From 60655a933bb786984afaf24904bf6e60f8e4cbf3 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Thu, 15 Oct 2026 03:27:43 +0000 Subject: [PATCH 001/239] fix(log-store): persist object store WAL region releases (#9579) * fix(log-store): persist object store WAL region releases A region's release lived only in memory: a restart after DROP TABLE and before the next object never re-established the dropped region's watermark, so its segments counted as live and pinned garbage collection of the whole prefix. A release now writes a release record (region id and the entry id the region is released through) into an object sealed at once, and returns once that object is durable and indexed, in both acknowledgement modes. Recovery seeds the obsolete watermarks from the release records of the indexed objects before it computes any reclaim boundary. obsolete_all releases in full; the new LogStore::release_region, a no-op by default, releases a dropped region through the largest entry id the store handed out for it, and Mito calls it after obsolete when it drops a region. Objects with release records use format version 2, the version 1 layout with the records at the end of the footer; every other object is still written as version 1. A counter of release records made durable is added. Signed-off-by: jeremyhi * docs(rfc): describe persisted object store WAL region releases Signed-off-by: jeremyhi * fix(log-store): carry a failed region release into the next object A release whose batch failed left its record only in an orphan object, so a drop that failed there could still leave the region unreleased after a restart. The store now keeps the release records of failed batches and seals them with its next batch, on a timer tick when no entry is admitted, so the release becomes durable once a create succeeds again. The version 1 and version 2 object fixtures are frozen under tests/compatibility/wal/object_store and loaded by the format tests, and the RFC no longer says that state is persisted only in headers. Signed-off-by: jeremyhi * fix(mito2): complete a region drop whose WAL release fails A drop request is not sent again once the region was removed, so a drop whose WAL release failed used to leave the region half dropped. The drop now completes, cleaning up the region's runtime state and scheduling its files for deletion, and reports the release failure. In the durable mode the object store WAL now creates one more object at stop with the release records of failed and dropped batches, so a graceful stop does not lose a release that failed. The release record counter is described as counting the records of indexed objects. Signed-off-by: jeremyhi * docs(log-store): describe the stop-time release create of the WAL actor Signed-off-by: jeremyhi * docs(rfc): state the accepted window of unreleased dropped regions Signed-off-by: jeremyhi * docs(log-store): describe release waiters in the WAL actor's outcome table Signed-off-by: jeremyhi * fix(log-store): fail region releases while the object store WAL is failing A failing store keeps its backlog and rejects appends until a create succeeds, so a release would otherwise wait for an upload that is not coming. A release that arrives in that state now fails at once with the retryable error, and one whose batch is still pending when the store enters it fails as well. Neither poisons the store: the record of the first waits for the next sealed batch, like that of a batch that failed, and the record of the second stays in the backlog, so both become durable once a create succeeds. Rejected releases count in rejected_appends_total. Signed-off-by: jeremyhi * docs(rfc): describe region releases while the object store WAL is failing Signed-off-by: jeremyhi * docs(rfc): place the drop crash window in the drop procedure The crash window between the deletion of a table's metadata and a durable release record is a gap of the drop procedure, not of the WAL. Describe it as far as the code shows, name the remedy that belongs in the procedure, cleaning up offline a region a resumed region drop finds not open, and list it under future work. Signed-off-by: jeremyhi * fix(log-store): keep rejected release records coalesced while failing A release rejected while the store is failing kept its record for the next sealed batch, but every timer tick then sealed that record alone behind the refused backlog, so a release repeated across ticks grew the backlog without bound. A failing store now keeps such records coalesced and seals them on the first tick after a create succeeds. The fault-matrix case of a release while the backlog is refused now runs with a 10 ms flush interval, repeats the release, and checks that the backlog does not grow and that recovery reads the record back. Signed-off-by: jeremyhi * docs(rfc): keep rejected release records unsealed while failing Signed-off-by: jeremyhi * docs(rfc): cite the drop procedure issue for the crash window Signed-off-by: jeremyhi * fix(log-store): let a stop seal rejected release records of a failing store Keeping the records of rejected releases out of sealed batches while the store is failing also kept them out of the seal a stop requests, so a stop whose last attempt of the backlog succeeded dropped them and the released region pinned collection after the restart. A stop now seals them with the backlog. A fault-matrix case rejects a release while the store is failing, stops, lets both creates succeed, and checks the record after a reopen. Signed-off-by: jeremyhi * docs(rfc): describe the stop of a failing store with rejected releases Signed-off-by: jeremyhi * fix(log-store): keep release records readable by older object store WAL readers Release records were written as a version 2 footer, which a build without it refuses to read, so rolling a node back after a drop failed to open the whole WAL prefix. They are now the single entry of a segment of the reserved region u64::MAX at the first position of a version 1 object: an older reader indexes that segment as a region it never opens, and recovery reads the segment back with one range read per object that holds one. The segment never holds the reclaim boundary and names no namespace. A durable stop whose last create of release records fails now reports the error instead of returning Ok and losing the records silently. The tests are trimmed: the release unit tests and fault-matrix cases that repeated each other are merged, and the Mito drop tests share one parameterized case. Signed-off-by: jeremyhi * docs(rfc): describe release records as a segment of a reserved region Signed-off-by: jeremyhi * fix(log-store): report release records a poisoning stop drops A durable stop whose last create of release records met a conflict poisoned the store and returned Ok, losing the records silently. Stop now reports the error whenever release records are dropped after it began, by a rollback or by poisoning. The test of a failed release covers that conflict, and seals its record on a tick after a one-shot create failure, so the next tick cannot fail too. The documentation now says that recovery reads the segments of release records, and that one that does not decode fails recovery. Signed-off-by: jeremyhi * docs(rfc): describe how recovery reads segments of release records Signed-off-by: jeremyhi * docs(rfc): say that stop reports release records a conflict drops Signed-off-by: jeremyhi * test(log-store): restore the transiently failing release create in enqueued mode The fault matrix again covers a release whose create fails transiently in the enqueued mode: the release waits for the repeat under the same sequence and returns once it succeeds. The encode_batch doc now says the release segment is the last segment of the object. Signed-off-by: jeremyhi * docs(rfc): place the release segment and count the restored case Signed-off-by: jeremyhi * fix(mito2): do not fail a region drop on a failed WAL release Returning the release error from the drop left a metric physical region half dropped: the metric engine stops after the data region and never drops the metadata region or forgets the physical region, and a retried drop finds the data region gone and counts it as done. The release error gave the caller nothing to retry, since the region is already removed and the log store keeps the record for its next object. The drop now logs the error and succeeds. A metric engine test drops a physical region while the WAL releases of both mito regions fail, and checks that both regions and the physical region state are gone. Signed-off-by: jeremyhi * docs(rfc): a failed release no longer fails the drop Signed-off-by: jeremyhi * fix(mito2): do not fail a region drop on a failed WAL obsolete A store that the release of another region poisoned fails obsolete as well, which left the metric engine with its metadata region undropped. Signed-off-by: jeremyhi * test(metric-engine): drop a physical region whose WAL obsolete or release fails Cover a failed data release, two failed releases and a conflicting release that poisons the store, after the regions wrote to the WAL, and check that a store opened later keeps no segment of either region. Signed-off-by: jeremyhi * docs(rfc): a failed obsolete does not fail the drop either Signed-off-by: jeremyhi * fix(mito2): fail the drop on an obsolete error of other WAL providers Only the object store WAL is shared by the regions of a store that the release of one can poison; Raft Engine and Kafka keep failing the drop when obsolete fails. Signed-off-by: jeremyhi * docs(rfc): a stopped or poisoned store keeps no failed release Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- Cargo.lock | 1 + docs/rfcs/2026-09-06-object-store-wal.md | 46 +- src/log-store/src/metrics.rs | 10 +- src/log-store/src/object_store_wal.rs | 12 +- src/log-store/src/object_store_wal/catalog.rs | 19 +- src/log-store/src/object_store_wal/format.rs | 139 +++ src/log-store/src/object_store_wal/store.rs | 1054 +++++++++++++++-- src/metric-engine/Cargo.toml | 1 + src/metric-engine/src/engine/drop.rs | 156 +++ src/metric-engine/src/test_util.rs | 11 + .../engine/object_store_wal_recovery_test.rs | 130 +- src/mito2/src/test_util.rs | 2 +- src/mito2/src/wal.rs | 12 + src/mito2/src/worker/handle_drop.rs | 22 +- src/store-api/src/logstore.rs | 16 + tests/compatibility/README.md | 6 + .../compatibility/wal/object_store/README.md | 16 + .../wal/object_store/release_record.wal | Bin 0 -> 272 bytes 18 files changed, 1514 insertions(+), 139 deletions(-) create mode 100644 tests/compatibility/wal/object_store/README.md create mode 100644 tests/compatibility/wal/object_store/release_record.wal diff --git a/Cargo.lock b/Cargo.lock index b0c66cc8b5c..edf460e4679 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8591,6 +8591,7 @@ dependencies = [ "humantime-serde", "itertools 0.14.0", "lazy_static", + "log-store", "mito-codec", "mito2", "moka", diff --git a/docs/rfcs/2026-09-06-object-store-wal.md b/docs/rfcs/2026-09-06-object-store-wal.md index b49f97c330c..1d784dd09ae 100644 --- a/docs/rfcs/2026-09-06-object-store-wal.md +++ b/docs/rfcs/2026-09-06-object-store-wal.md @@ -31,7 +31,7 @@ Object storage is already provisioned, already durable across zones, and already # Non-goals - Distributed mode. The provider is rejected when a datanode is controlled by a metasrv; the metasrv side has no allocation logic for it yet. -- Persistence of per-region obsolete watermarks. Mito re-establishes every watermark when it opens a region, which is what garbage collection reasons from; persisting them is a follow-up. +- Persistence of the obsolete watermarks that flushes and truncates move. Mito re-establishes each of them when it opens a region, which is what garbage collection reasons from. Only the release of a region that is never opened again, after a drop or an offline cleanup, is persisted, as described under *Entry ids*. - Matching Raft Engine's write latency. The provider trades latency for the absence of local durable WAL state. - Making standalone metadata independent of local disk. The metadata store is unchanged by this RFC. @@ -65,11 +65,12 @@ header | segment(region 1) | ... | segment(region N) | footer | trailer - The header carries the magic `GTWALOBJ`, format version 1, the object sequence number, the epoch of the store instance that wrote the object, a link to the object this one extends, as that object's sequence number and epoch or none for the first object of a chain, the reclaim boundary described under *Entry ids*, and a CRC32 of the preceding header bytes; the header is 54 bytes. A reclaim boundary above the object's own sequence fails the header check. A sequence number and an epoch together identify an object: an instance never writes two objects under one sequence, and instances have distinct epochs, as described under *Log store*. Recovery selects the history it replays by the links and the epochs, as described under *Log store*. - An object may hold no segment. A store writes such an object when it opens, as described under *Log store*; it is a chain object like any other, and its segments are trivially all flushed. +- An object may carry release records, which a store writes when it releases a region, as described under *Entry ids*; each names a region and the entry id the region is released through. They form the single entry of a segment of the reserved region id `u64::MAX`, whose entry id is at position 1 of the object; segments are ordered by region id, so it is the last segment of the object. Its payload is their count, then per record, ordered by region id and at most one per region, the region id and the entry id. The object keeps format version 1, so a reader that predates release records reads it as any other object. An object written for a release alone holds only that segment. - Each segment holds the entries of exactly one region, ordered by entry id; segments are ordered by region id. Ordering makes the encoding deterministic, which the identical-retry rule below depends on. - The footer records, per segment, the region id, the entry id range, the entry count, the byte range and a CRC32 of the segment. - The fixed-size trailer carries the magic `GTWALTRL`, the footer location, a CRC32 of the footer and a CRC32 of the whole object. A reader locates the footer by reading the trailer at the end of the object and can decode a single region's segment without touching the others. -Corruption is handled at two levels; nothing else is checked in production. Structural damage fails recovery: recovery reads only the header, the trailer and the footer of every object, as described under *Log store*, and verifies the header and trailer magic, the format version, the header CRC, the header sequence against the key, the footer CRC and that the segment byte ranges tile the object body exactly; the first object that fails any of these checks fails construction of the store, since nothing can be trusted about such an object. A segment is verified by the read that decodes it, and a segment that does not decode, because its CRC does not match or its content disagrees with its footer entry, is a data-level failure confined to one region. The read fetches the segment once more, since the download rather than the object may be damaged, and decodes it again; if it still does not decode, the behaviour follows `on_corrupted_segment`. `skip`, the default, skips the segment, records a WAL hole for the region with the object key, the object sequence and the entry id range the footer holds for the segment, increments the counter `greptime_logstore_object_store_wal_skipped_segments_total` and logs a warning with those fields, and the read continues with the next object; every other region in the object is unaffected. The holes are kept in memory per region and can be listed from the store; nothing is persisted, so a store opened later on the same prefix records a hole again when a read meets the segment. `fail` fails the read with the decode error, so a region whose replay meets the segment does not open. Skipping loses only the unflushed entries of that region in that object: nothing already flushed is affected, and the loss is reported rather than hidden. A failed fetch is not corruption and fails the read in both modes, unless garbage collection deleted the object, as described under *Entry ids*, in which case the object is skipped; an object the catalog still holds once a delete of it has settled fails the read whether it is missing or unreadable, and a segment of such an object that does not decode follows `on_corrupted_segment` as it would without any collection. The whole-object CRC is not a recovery requirement; it stays in the trailer as a fixture-level check. Because recovery never checks it, the header carries its own CRC, so a damaged link fails recovery instead of silently cutting acknowledged objects off the chain. The format is frozen at version 1. +Corruption is handled at two levels; nothing else is checked in production. Structural damage fails recovery: recovery reads only the header, the trailer and the footer of every object, as described under *Log store*, and verifies the header and trailer magic, the format version, the header CRC, the header sequence against the key, the footer CRC and that the segment byte ranges tile the object body exactly; the first object that fails any of these checks fails construction of the store, since nothing can be trusted about such an object. The segment of release records of an indexed object, which recovery reads too, is held to the same rule rather than to `on_corrupted_segment`: one that cannot be fetched or does not decode fails construction, since the watermarks it seeds would otherwise be lost. A segment is verified by the read that decodes it, and a segment that does not decode, because its CRC does not match or its content disagrees with its footer entry, is a data-level failure confined to one region. The read fetches the segment once more, since the download rather than the object may be damaged, and decodes it again; if it still does not decode, the behaviour follows `on_corrupted_segment`. `skip`, the default, skips the segment, records a WAL hole for the region with the object key, the object sequence and the entry id range the footer holds for the segment, increments the counter `greptime_logstore_object_store_wal_skipped_segments_total` and logs a warning with those fields, and the read continues with the next object; every other region in the object is unaffected. The holes are kept in memory per region and can be listed from the store; nothing is persisted, so a store opened later on the same prefix records a hole again when a read meets the segment. `fail` fails the read with the decode error, so a region whose replay meets the segment does not open. Skipping loses only the unflushed entries of that region in that object: nothing already flushed is affected, and the loss is reported rather than hidden. A failed fetch is not corruption and fails the read in both modes, unless garbage collection deleted the object, as described under *Entry ids*, in which case the object is skipped; an object the catalog still holds once a delete of it has settled fails the read whether it is missing or unreadable, and a segment of such an object that does not decode follows `on_corrupted_segment` as it would without any collection. The whole-object CRC is not a recovery requirement; it stays in the trailer as a fixture-level check. Because recovery never checks it, the header carries its own CRC, so a damaged link fails recovery instead of silently cutting acknowledged objects off the chain. The operator configures only the prefix root. The store runs under a node prefix derived from it, `/datanodes//epochs/`, and object keys are `/objects/.wal`, where the sequence is zero-padded to 20 decimal digits. Standalone uses its datanode id, zero unless configured, and generation zero, so with prefix `wal` the first object is `wal/datanodes/0/epochs/0/objects/00000000000000000000.wal`. Every place that names the prefix, whether to open the store or to allocate the WAL options of a region, derives it from the same function, so the prefix a region persists matches the one its store runs under. Listing accepts only well-formed keys under `/objects/` and returns them in sequence order. This is the standalone form of the key layout under *Alignment with cluster mode*; the cluster assignment of node id and generation is proposed there. @@ -105,17 +106,19 @@ Two rules keep every new id of a region above every id it already has, on a pref The second rule is the sequence floor of `obsolete`. Mito calls `obsolete` with the manifest's `flushed_entry_id` when it opens a region, and the actor applies the watermark and a floor derived from that entry id together: the next object sequence is raised to `(entry_id >> 20) + 1` unless it is there already, while zero, the watermark of a region without entries, names no object and needs no floor, so that a watermark assigned under another prefix, or one whose object this prefix no longer holds, is never passed by a new id. The floor moves only the next sequence, never the chain: the next batch still links to its predecessor, so a skipped sequence leaves no gap in the history recovery replays. While the open batch holds ids under the next sequence, the call is refused with a retryable error and records neither the watermark nor the floor, because those ids would fall below the floor. In the `enqueued` mode an id the store itself handed out needs no floor, even before it is durable: it is never handed out again while the store runs, so every later id of the region is greater. In the current single-prefix layout the floor never needs to move, since every watermark Mito hands the store names an object below the next sequence; the rule exists for the cross-prefix cases under *Alignment with cluster mode*. No additional persisted state is introduced. -Garbage collection deletes the objects below a reclaim boundary that the store persists in every header, so that recovery can tell a collected predecessor from one that never landed. The boundary is computed from the catalog and the watermarks the store already keeps: each region with a recorded watermark contributes the sequence of its first indexed object holding a segment whose largest entry id is above the watermark; a region without a recorded watermark contributes its first indexed object; a region whose watermark is the maximum entry id, after `obsolete_all`, contributes nothing; the boundary is the smallest contribution, or the last indexed sequence plus one when nothing contributes, and it never decreases. Whether a segment is flushed is decided per segment from the footer, on the full entry id, so a watermark that stops in the middle of an object keeps the object. The boundary is recomputed after every successful `obsolete`, after `obsolete_all` and after every object is indexed; Mito calls `obsolete` after a flush, after a truncate, when it drops a region and when it opens one, so a restart re-establishes every watermark and the boundary advances on the same rule. Every sealed batch carries the boundary current at sealing and the start object carries the boundary recovery found, so a boundary is always at or below the sequence of the object carrying it. The boundary that permits deletion is the one carried by the last indexed object; the boundary of a batch whose create failed or is in flight permits nothing. +Garbage collection deletes the objects below a reclaim boundary that the store persists in every header, so that recovery can tell a collected predecessor from one that never landed. The boundary is computed from the catalog and the watermarks the store already keeps: each region with a recorded watermark contributes the sequence of its first indexed object holding a segment whose largest entry id is above the watermark; a region without a recorded watermark contributes its first indexed object; a region whose watermark is the maximum entry id, after a release in full, contributes nothing; the boundary is the smallest contribution, or the last indexed sequence plus one when nothing contributes, and it never decreases. Whether a segment is flushed is decided per segment from the footer, on the full entry id, so a watermark that stops in the middle of an object keeps the object. The boundary is recomputed after every successful `obsolete`, after every release and after every object is indexed; Mito calls `obsolete` after a flush, after a truncate, when it drops a region and when it opens one, so a restart re-establishes the watermark of every region Mito opens again, and the boundary advances on the same rule. A region that is never opened again needs the release described below. Every sealed batch carries the boundary current at sealing and the start object carries the boundary recovery found, so a boundary is always at or below the sequence of the object carrying it. The boundary that permits deletion is the one carried by the last indexed object; the boundary of a batch whose create failed or is in flight permits nothing. -Every present object with a sequence below the durable boundary is deletable, on the chain or not: the objects the catalog indexes, the objects recovery listed but did not index, the objects of an earlier epoch that the start object moved past, and the sequences of batches that failed in this run, whose objects may or may not exist. A collection starts after a successful `obsolete`, after `obsolete_all` and after an indexing that raised the durable boundary. Deletes run in the background, at most four at a time and in no particular order, and never hold up admission, sealing, uploads or acknowledgements; an indexed object is unindexed only after its delete succeeded, and the WAL holes recorded for it are dropped with it. Every deleted object is counted in `greptime_logstore_object_store_wal_deleted_objects_total`; a failed delete is logged, counted in `greptime_logstore_object_store_wal_failed_deletes_total` and retried at the next collection, and the object stays indexed meanwhile. An object deleted but not yet unindexed when the process crashes is simply absent at the next recovery, which lists the prefix. Nothing is collected once `stop` has begun or while the store is poisoned; a delete in flight when `stop` begins runs to completion before `stop` returns, and the actor settles every attempt it abandons when it exits. +Every present object with a sequence below the durable boundary is deletable, on the chain or not: the objects the catalog indexes, the objects recovery listed but did not index, the objects of an earlier epoch that the start object moved past, and the sequences of batches that failed in this run, whose objects may or may not exist. A collection starts after a successful `obsolete`, after a release and after an indexing that raised the durable boundary. Deletes run in the background, at most four at a time and in no particular order, and never hold up admission, sealing, uploads or acknowledgements; an indexed object is unindexed only after its delete succeeded, and the WAL holes recorded for it are dropped with it. Every deleted object is counted in `greptime_logstore_object_store_wal_deleted_objects_total`; a failed delete is logged, counted in `greptime_logstore_object_store_wal_failed_deletes_total` and retried at the next collection, and the object stays indexed meanwhile. An object deleted but not yet unindexed when the process crashes is simply absent at the next recovery, which lists the prefix. Nothing is collected once `stop` has begun or while the store is poisoned; a delete in flight when `stop` begins runs to completion before `stop` returns, and the actor settles every attempt it abandons when it exits. -A read never needs a deleted object, since entries at or below a region's watermark are hidden anyway, but a read that listed an object before it was collected skips the object when it gets there. A read that meets an error on such an object makes two observations, in the reverse of the order the actor writes them: it looks at the set of objects whose delete is in flight, and then at the catalog. The actor registers an object in that set before its delete starts, unindexes it if the delete succeeded, and only then takes it out, so an object absent from the set was either never collected, in which case the catalog still holds it, or is already unindexed by the time the catalog is read. Observing the catalog first would let a collection complete between the two observations and leave the read with an object that still looked indexed and no longer looked deleting. An object that is still in the set has an undecided delete, so the read waits for that one delete, a single object store request, rather than assume it succeeds: a delete that fails leaves the object present and indexed, and the error the read met is then reported as any other, `on_corrupted_segment` included. What the read waits for is the delete attempt its observation found, taken together with the lookup, not the object: a retry of the same object is a separate attempt, so the attempt the read holds releases it even while that retry is in flight, and a read is never handed from one attempt to the next. A read stream outlives the store it came from, so the actor settles every attempt it abandons when it exits, whether the store was stopped or dropped; the objects of those attempts stay present and indexed, so the reads waiting for them report the error they met. A collected object is skipped silently: it is not recorded as a hole and does not count as a skipped segment. An object the catalog still holds after those observations fails the read like any other fetch failure, `on_corrupted_segment` included. Within one run the catalog keeps the largest entry id of a region after the objects holding it were collected, so `latest_entry_id`, the durability wait and the `enqueued` clamp keep their reference point; a store opened later learns only what is listed, so after a restart `latest_entry_id` of a region whose later objects were collected is the largest id among the objects left, which only the pruning hint of Mito consumes. No state is persisted outside the headers and no configuration option is added. +A region that Mito drops is never opened again, so no restart re-establishes its watermark, and a region without a watermark holds the boundary at its first indexed object. The store therefore persists the release of such a region in the WAL. `obsolete_all`, which Mito calls when it cleans a region up offline, releases the region in full, through the maximum entry id. `release_region`, a `LogStore` method that Mito calls after `obsolete` when it drops a region, releases it through the largest entry id the store has handed out for the region, so entries of a region that reuses the id later stay above the release. A release moves the region's watermark up in memory, recomputes the boundary, starts a collection and seals the open batch at once, empty or not, with a release record of the region. The call returns once that object is durable and indexed, in both acknowledgement modes, and fails whenever its batch fails, as an append in the `durable` mode does. In the `enqueued` mode the watermark of a release is not clamped to the durable entry id: a released region is never replayed, and its record becomes durable only after every object sealed before it. The object of a release carries the boundary its watermark allows, so the record and that boundary become durable together. Recovery seeds the watermarks from the release records of the indexed objects before it computes any boundary, and a later `obsolete` never moves a watermark down. A release record needs no bookkeeping of its own. It is sealed once every entry it covers was admitted, so it lies in the object holding the last segment it covers or in a later one; collection deletes, in any order, only objects below the durable boundary; and recovery indexes only objects at or above the boundary of the tip, which every start object carries forward. So recovery indexes a release record whenever it indexes a segment the record covers, and once the record is below the boundary, so is every segment it covers, which recovery never indexes again, whether collection deleted it yet or not. A release in full also covers the entries a region would be assigned afterwards, so it is made only for a region that writes nothing more. When the batch of a release fails, the call fails, and the store carries the record again in the next batch it seals; while such a record waits, a timer tick seals it even when no entry is admitted, and a stop in the `durable` mode gives it one more create and reports the record lost when that create fails, so the release becomes durable once a create succeeds again. While the store is failing, as described under *Log store*, a release fails at once with the retryable error rejected appends carry, and so does a release whose batch is still pending when the store becomes failing; neither poisons the store. The record of the first is kept like that of a failed batch, but no tick seals it while the store is failing, so repeated releases add no batch to the backlog; the first tick after a create succeeds seals it, and a `stop` seals it with the backlog for its last attempt. The record of the second stays in the backlog. Both become durable once a create succeeds. A crash before that loses the record. Two windows therefore remain in which a dropped region is never released: a crash between the deletion of the table metadata and a durable release record, which includes a crash during the drop itself, and a crash after a failed release before any later create succeeds. Nothing retries the drop of such a region. The first window is a gap of the drop procedure rather than of the WAL, issue #9577 (https://github.com/GreptimeTeam/greptimedb/issues/9577): a crash during the drop leaves the region's data behind whatever the WAL provider. In the second window the drop completed except for its release, and a standalone region drop that fails is not retried. An operator sees `greptime_logstore_object_store_wal_indexed_objects` only grow after the restart while `deleted_objects_total` stays put, because the region's segments hold the reclaim boundary. Its remedy belongs in the drop procedure, for example an offline cleanup of a region that a resumed region drop finds not open, whose `obsolete_all` writes a durable release. Reclamation that consults the metadata for live regions, as Kafka's pruning by the metasrv does, would be a second safety net for both windows; it is left to garbage collection across nodes, see *Alignment with cluster mode*. Neither is part of this change. The object of the failed batch is seeded only when nothing was acknowledged after it, as the entries of a failed append are replayed. Every release record in an object the store indexes is counted in `greptime_logstore_object_store_wal_release_records_total`, records carried over from a failed batch included. + +A read never needs a deleted object, since entries at or below a region's watermark are hidden anyway, but a read that listed an object before it was collected skips the object when it gets there. A read that meets an error on such an object makes two observations, in the reverse of the order the actor writes them: it looks at the set of objects whose delete is in flight, and then at the catalog. The actor registers an object in that set before its delete starts, unindexes it if the delete succeeded, and only then takes it out, so an object absent from the set was either never collected, in which case the catalog still holds it, or is already unindexed by the time the catalog is read. Observing the catalog first would let a collection complete between the two observations and leave the read with an object that still looked indexed and no longer looked deleting. An object that is still in the set has an undecided delete, so the read waits for that one delete, a single object store request, rather than assume it succeeds: a delete that fails leaves the object present and indexed, and the error the read met is then reported as any other, `on_corrupted_segment` included. What the read waits for is the delete attempt its observation found, taken together with the lookup, not the object: a retry of the same object is a separate attempt, so the attempt the read holds releases it even while that retry is in flight, and a read is never handed from one attempt to the next. A read stream outlives the store it came from, so the actor settles every attempt it abandons when it exits, whether the store was stopped or dropped; the objects of those attempts stay present and indexed, so the reads waiting for them report the error they met. A collected object is skipped silently: it is not recorded as a hole and does not count as a skipped segment. An object the catalog still holds after those observations fails the read like any other fetch failure, `on_corrupted_segment` included. Within one run the catalog keeps the largest entry id of a region after the objects holding it were collected, so `latest_entry_id`, the durability wait and the `enqueued` clamp keep their reference point; a store opened later learns only what is listed, so after a restart `latest_entry_id` of a region whose later objects were collected is the largest id among the objects left, which only the pruning hint of Mito consumes. Reclaim boundaries are persisted in the headers and release records in the segments of the reserved region, no other state is persisted, and no configuration option is added. ## Log store `ObjectStoreLogStore` implements the `LogStore` trait for `Provider::ObjectStore`. -**Construction and recovery.** The prefix, flush interval, batch size and backlog thresholds are validated, then the store lists the prefix and fetches, for up to eight objects at a time, the header and a window at the end of each object with range reads; a short object is read whole, and the footer is read separately only when it starts before the window. Each header and footer is verified as described under *Object format*; the first corrupted object fails construction; a fetch that finds the object gone is judged once the history is selected, as described below, and any other failed fetch abandons the rest. Segments are not read, so recovery costs a few small reads per object however large the objects are. +**Construction and recovery.** The prefix, flush interval, batch size and backlog thresholds are validated, then the store lists the prefix and fetches, for up to eight objects at a time, the header and a window at the end of each object with range reads; a short object is read whole, and the footer is read separately only when it starts before the window. Each header and footer is verified as described under *Object format*; the first corrupted object fails construction; a fetch that finds the object gone is judged once the history is selected, as described below, and any other failed fetch abandons the rest. Segments are not read, except the segment of release records of each indexed object that holds one, with one range read, so recovery costs a few small reads per object however large the objects are. The store then selects the history it replays from the headers. An object's link holds when the object has no predecessor, when its predecessor is present and carries the epoch the link records, or when the predecessor's sequence is below the reclaim boundary the object carries, in which case the chain ends there whether or not the predecessor is present; above that boundary a missing predecessor breaks the link. An object is complete when every link on its chain holds, each judged with the boundary of the object carrying the link. The tip is the complete object with the largest epoch and, among those, the largest sequence. Only the objects on the tip's chain are indexed and replayed, in sequence order whatever the order the fetches completed in, so the catalog checks the entry ranges of every object against its predecessors like a sequential replay, and a conflicting object fails construction. Every other present object is an orphan: it is neither indexed nor read, but its sequence counts, so the next sequence is above every listed object, and the objects below the tip's boundary are collected. A listed object whose header or footer is gone by the time it is fetched is ignored when it lies below the tip's boundary and fails construction when it lies at or above it, since collection never deletes one there. Construction also fails when a listed object carries an epoch above the next sequence, which no instance can have produced, since an epoch is a start sequence plus one. When objects are present but none is complete, construction fails rather than start an empty history. Construction then records the largest durable entry id per region. @@ -142,28 +145,28 @@ The mode also needs a barrier that the durable mode gets for free. An entry id t | create succeeds, or identical retry | advances | acknowledged in sequence order | healthy | | transient I/O error, `durable` mode | consumed and never reused, nor are the sequences of the other failed batches | this and every batch that is not indexed fail, including those whose objects were created; a create still in flight completes as an orphan | healthy: the next batch links to the last indexed object | | an object of an earlier epoch holds the sequence, `durable` mode | consumed and never reused | as for a transient I/O error | healthy | -| transient I/O error, `enqueued` mode | unchanged | already acknowledged | healthy: the create is repeated after a short delay | -| I/O error that requires intervention, `enqueued` mode | unchanged | already acknowledged; later appends are rejected with a retryable error until a create succeeds | failing: the lowest pending create is repeated alone with backoff; the first create that succeeds ends it | -| any I/O error, `enqueued` mode after `stop` began | unchanged | already acknowledged | the backlog is dropped and `stop` reports the loss | +| transient I/O error, `enqueued` mode | unchanged | already acknowledged, except a release, which waits for the repeat | healthy: the create is repeated after a short delay | +| I/O error that requires intervention, `enqueued` mode | unchanged | already acknowledged, except a release, which fails with a retryable error while its record stays in the backlog; later appends and releases are rejected with that error until a create succeeds | failing: the lowest pending create is repeated alone with backoff; the first create that succeeds ends it | +| any I/O error, `enqueued` mode after `stop` began | unchanged | already acknowledged, except a release, which fails with the stopped error | the backlog is dropped and `stop` reports the loss | | an object of the same or a later epoch holds the sequence, or any conflicting object in the `enqueued` mode | unchanged | every unacknowledged waiter fails, including those of batches whose objects were created but not yet indexed; an `enqueued` acknowledgement already returned cannot be revoked | poisoned | | encoding or catalog error | unchanged | every unacknowledged waiter fails, including those of batches whose objects were created but not yet indexed; an `enqueued` acknowledgement already returned cannot be revoked | poisoned | | create succeeds at the last representable sequence | cannot advance | acknowledged | poisoned: no later batch can be allocated a sequence, every later append fails | | the start object meets an object of an earlier epoch | moves to the next sequence | none admitted yet | construction continues | | the start object meets an object of the same or a later epoch, or its create has an unknown outcome | unchanged | none admitted yet | construction fails; an unknown outcome is retryable | -A store cannot be constructed on a prefix whose largest object already carries the maximum sequence. A poisoned store fails every `LogStore` operation except `stop` with the same terminal error; in particular `obsolete` can no longer move a watermark. In the `durable` mode transient errors are not retried inside the store, but a failed create is reconciled once: the store immediately reads the key back, accepts the batch if the object exists with identical bytes, treats it as a conflict if the bytes differ, and reports the transient error only when the key is absent or that read fails too. So a lost response alone does not fail the append. When the append does fail, neither the store nor Mito keeps the batch: it surfaces through Mito to the caller of the write, and a retry is a new write by the caller under a new sequence. If the create had in fact succeeded and both its response and the read-back were lost, or its request lands only later, the object is an orphan as soon as a later object is acknowledged. Retrying the batch itself inside the store in the `durable` mode would only hide the latency and would hold writes up for as long as an outage lasts; the `enqueued` mode repeats the create with the same bytes, so a create that succeeded without a response is accepted as an identical retry. +A release is a waiter of the batch that carries its record in both modes: it returns when the batch is indexed and fails whenever the batch fails. A store cannot be constructed on a prefix whose largest object already carries the maximum sequence. A poisoned store fails every `LogStore` operation except `stop` with the same terminal error; in particular `obsolete` can no longer move a watermark. In the `durable` mode transient errors are not retried inside the store, but a failed create is reconciled once: the store immediately reads the key back, accepts the batch if the object exists with identical bytes, treats it as a conflict if the bytes differ, and reports the transient error only when the key is absent or that read fails too. So a lost response alone does not fail the append. When the append does fail, neither the store nor Mito keeps the batch: it surfaces through Mito to the caller of the write, and a retry is a new write by the caller under a new sequence. If the create had in fact succeeded and both its response and the read-back were lost, or its request lands only later, the object is an orphan as soon as a later object is acknowledged. Retrying the batch itself inside the store in the `durable` mode would only hide the latency and would hold writes up for as long as an outage lasts; the `enqueued` mode repeats the create with the same bytes, so a create that succeeded without a response is accepted as an identical retry. -**Stopping.** `stop` is idempotent and awaits the actor. Once stop has begun, queued appends are not admitted, no timer or size triggered seal starts, and appends after stop, including empty ones, fail with the stopped error. In the `durable` mode the open batch and every sealed batch whose create has not started are dropped: their waiters receive the stopped error. A create that is already in flight runs to completion; if it succeeds while its batch is still pending, the batch is indexed in order and acknowledged. A transient failure after stop began fails every batch that is not indexed, as it does while the store runs, and those batches receive the stopped error; a create of such a batch that completes later is an orphan and acknowledges nothing. A create that finds a conflicting object of the same or a later epoch after stop began poisons the store as it would while the store runs, but the waiters of the batches that are not indexed receive the stopped error, the conflict is kept only as the terminal error, and `stop` returns successfully once the creates in flight have completed. In the `enqueued` mode the held-back appends receive the stopped error, the open batch is sealed and the whole backlog is uploaded before `stop` returns, by a failing store too; any object store failure in that phase drops the backlog and `stop` returns the error, as described under *Acknowledgement*. A durability wait still pending when the actor exits receives the stopped error, and a delete already in flight runs to completion before `stop` returns. +**Stopping.** `stop` is idempotent and awaits the actor. Once stop has begun, queued appends are not admitted, no timer or size triggered seal starts, and appends after stop, including empty ones, and releases fail with the stopped error. In the `durable` mode the release records of failed batches and of the batches dropped here are sealed into one more object, which is created before `stop` returns; `stop` reports the error when that create fails too. In the `durable` mode the open batch and every sealed batch whose create has not started are dropped: their waiters receive the stopped error. A create that is already in flight runs to completion; if it succeeds while its batch is still pending, the batch is indexed in order and acknowledged. A transient failure after stop began fails every batch that is not indexed, as it does while the store runs, and those batches receive the stopped error; a create of such a batch that completes later is an orphan and acknowledges nothing. A create that finds a conflicting object of the same or a later epoch after stop began poisons the store as it would while the store runs, but the waiters of the batches that are not indexed receive the stopped error, the conflict is kept only as the terminal error, and `stop` returns successfully once the creates in flight have completed, unless the poisoning drops release records, which `stop` reports by returning the conflict. In the `enqueued` mode the held-back appends receive the stopped error, the open batch is sealed and the whole backlog is uploaded before `stop` returns, by a failing store too; any object store failure in that phase drops the backlog and `stop` returns the error, as described under *Acknowledgement*. A durability wait still pending when the actor exits receives the stopped error, and a delete already in flight runs to completion before `stop` returns. -**Reading.** `read(provider, entry_id)` returns the region's entries with id at least `entry_id`, in order, fetching only that region's segments with range reads and decoding nothing else, and hides entries at or below the region's obsolete watermark. `obsolete` records the watermark in memory, clamped to the durable entry id in the `enqueued` mode; the watermark is re-established after a restart because Mito calls `obsolete` with the manifest's `flushed_entry_id` when it opens a region, so no WAL-owned watermark state is persisted. `obsolete_all` hides the whole region. `latest_entry_id` is not affected by the watermark. Every method taking a provider requires `Provider::ObjectStore` with the store's prefix and the caller's region. +**Reading.** `read(provider, entry_id)` returns the region's entries with id at least `entry_id`, in order, fetching only that region's segments with range reads and decoding nothing else, and hides entries at or below the region's obsolete watermark. `obsolete` records the watermark in memory, clamped to the durable entry id in the `enqueued` mode; the watermark is re-established after a restart because Mito calls `obsolete` with the manifest's `flushed_entry_id` when it opens a region, so the store persists no watermark of a region Mito opens again. The release of a region, which `obsolete_all` and `release_region` make, is persisted as described under *Entry ids*. `obsolete_all` hides the whole region. `latest_entry_id` is not affected by the watermark. Every method taking a provider requires `Provider::ObjectStore` with the store's prefix and the caller's region. -**Metrics.** The store exports metrics under `greptime_logstore_object_store_wal_`. The write path reports `seal_to_durable_seconds`, the time from sealing a batch until its object is durable, `append_ack_seconds`, the time an append waits for its acknowledgement in the `durable` mode, `object_bytes` and `object_entries`, the size and the entry count of every object sealed, and `created_objects_total`, the creates that succeeded, an identical retry included; a create is counted and timed whether or not its batch is still there to be indexed, since an object a poisoned store gave up on is as durable as any other. `indexed_objects` and `indexed_bytes` are the objects and the bytes the catalog holds, set by a recovery that succeeded and moved by every object indexed or unindexed afterwards, so a collection takes the object it deleted off both, and a recovery that is abandoned publishes neither. A recovery that succeeds reports `recovery_seconds` and the objects it read footers from in `recovered_objects_total`, both of the same recoveries, so one that is abandoned at a fetch or rejected by the catalog contributes to neither; a read stream reports `read_seconds` once it has been read to its end, and one that fails or is dropped part way reports nothing. The failure paths are counted: `create_conflicts_total` for a create that found a conflicting object, `create_failures_total` for one the object store did not confirm, `poisoned_total` for a store that recorded a terminal error, `skipped_segments_total` for a segment a read skipped, and `stalled_appends_total` with `stalled_append_seconds` for an append held back because the sealed batches were full, in either mode, or by the `enqueued` backlog thresholds, timed until it is admitted or refused, a stop, a terminal error, a failing store or the teardown of the actor releasing it included. The gauge `failing` counts the stores whose backlog the object store refuses with an error that requires intervention, and `rejected_appends_total` the appends such a store rejected; the attempts of a failing store are counted in `create_failures_total`. The delete counters of a collection are described under *Entry ids*. The histogram buckets are chosen for the default flush interval of 100 ms and the default `max_batch_bytes` of 8 MiB. +**Metrics.** The store exports metrics under `greptime_logstore_object_store_wal_`. The write path reports `seal_to_durable_seconds`, the time from sealing a batch until its object is durable, `append_ack_seconds`, the time an append waits for its acknowledgement in the `durable` mode, `object_bytes` and `object_entries`, the size and the entry count of every object sealed, and `created_objects_total`, the creates that succeeded, an identical retry included; a create is counted and timed whether or not its batch is still there to be indexed, since an object a poisoned store gave up on is as durable as any other. `indexed_objects` and `indexed_bytes` are the objects and the bytes the catalog holds, set by a recovery that succeeded and moved by every object indexed or unindexed afterwards, so a collection takes the object it deleted off both, and a recovery that is abandoned publishes neither. A recovery that succeeds reports `recovery_seconds` and the objects it read footers from in `recovered_objects_total`, both of the same recoveries, so one that is abandoned at a fetch or rejected by the catalog contributes to neither; a read stream reports `read_seconds` once it has been read to its end, and one that fails or is dropped part way reports nothing. The failure paths are counted: `create_conflicts_total` for a create that found a conflicting object, `create_failures_total` for one the object store did not confirm, `poisoned_total` for a store that recorded a terminal error, `skipped_segments_total` for a segment a read skipped, and `stalled_appends_total` with `stalled_append_seconds` for an append held back because the sealed batches were full, in either mode, or by the `enqueued` backlog thresholds, timed until it is admitted or refused, a stop, a terminal error, a failing store or the teardown of the actor releasing it included. The gauge `failing` counts the stores whose backlog the object store refuses with an error that requires intervention, and `rejected_appends_total` the appends and releases such a store rejected; the attempts of a failing store are counted in `create_failures_total`. The delete counters of a collection and the release record counter are described under *Entry ids*. The histogram buckets are chosen for the default flush interval of 100 ms and the default `max_batch_bytes` of 8 MiB. A `testing` cargo feature exposes hooks to wait for admitted appends, to seal the open batch, to hold and release creates and to make creates fail, so engine tests can flush deterministically and observe entries that are admitted but not durable. ## Mito integration -Mito maps `WalOptions::ObjectStore { prefix }` to `Provider::object_store_provider(region_id, prefix)`; resolving object store options against a Raft Engine or Kafka log store reports the existing incompatible-provider error. Reads go through the region-scoped reader used for Kafka, so a namespace shared by many regions never leaks entries across regions even if a log store returned them. Object store regions are opened individually rather than grouped by topic. After replay, `topic_latest_entry_id` follows the existing remote-WAL rule: it is taken from the store's latest entry id only when the region's memtables are empty, otherwise it stays at the replay start, because that value feeds pruning and must never advance past unflushed data. Before a flush, a full truncate or a discard of unflushed data publishes a manifest watermark, Mito waits on the durability barrier described under *Acknowledgement*; with the `durable` mode and with every other provider the wait returns at once. +Mito maps `WalOptions::ObjectStore { prefix }` to `Provider::object_store_provider(region_id, prefix)`; resolving object store options against a Raft Engine or Kafka log store reports the existing incompatible-provider error. Reads go through the region-scoped reader used for Kafka, so a namespace shared by many regions never leaks entries across regions even if a log store returned them. Object store regions are opened individually rather than grouped by topic. After replay, `topic_latest_entry_id` follows the existing remote-WAL rule: it is taken from the store's latest entry id only when the region's memtables are empty, otherwise it stays at the replay start, because that value feeds pruning and must never advance past unflushed data. Before a flush, a full truncate or a discard of unflushed data publishes a manifest watermark, Mito waits on the durability barrier described under *Acknowledgement*; with the `durable` mode and with every other provider the wait returns at once. When Mito drops a region it calls `release_region` after `obsolete`, as described under *Entry ids*; the default of that method does nothing, so the other providers handle a dropped region as before. Neither a release nor an `obsolete` of an object store region that fails, the latter for instance on a store that the release of another region poisoned, fails the drop, since nothing sends the drop of a removed region again and a caller that drops several regions, such as the metric engine with its data and metadata regions, would leave the others undone: the region's runtime state is cleaned up and its files are scheduled for deletion as for any drop, and the failure is logged; an `obsolete` that fails for another provider fails the drop as before. A store that is neither stopped nor poisoned keeps the record and writes it with its next object; a stopped or poisoned store keeps none, and the region stays unreleased after a restart. The store carries the record in a later object, as described under *Entry ids*; a crash before any later create succeeds loses it, and the dropped region's segments then hold the reclaim boundary after the restart. ## Datanode wiring and lifecycle @@ -215,9 +218,10 @@ on_corrupted_segment = "skip" ## Compatibility - `WalOptions::ObjectStore` is a new persisted variant. Existing `raft_engine`, `kafka` and `noop` options decode unchanged; a region's options never change provider implicitly. -- The object format is version 1 and frozen once the provider is released; a version bump after that is a new decoder, not an in-place change. The version 1 header layout was revised before release to carry the epoch, the link, the reclaim boundary and the header CRC. No object in an earlier layout exists, since the provider has not been enabled anywhere. +- The object format is version 1 and frozen once the provider is released; a version bump after that is a new decoder, not an in-place change. The version 1 header layout was revised before release to carry the epoch, the link, the reclaim boundary and the header CRC. No object in an earlier layout exists, since the provider has not been enabled anywhere. Release records add no format version: they are the segment of a reserved region in a version 1 object, so a build that predates them still opens the prefix. Such a build indexes that segment as a region it never opens, which holds its reclaim boundary at the first release until the release records are understood again. A frozen object with release records under `tests/compatibility/wal/object_store` is checked byte for byte by the format tests; the SQL compatibility runner owns `wal.provider` and does not cover this provider. - The configuration tag carries `experimental_`; the persisted tag does not, so promoting the provider to stable is a configuration rename with an alias, not a metadata migration. - `LogStore::wait_durable` has a default implementation that returns at once, so Raft Engine, Kafka and the noop provider keep their own acknowledgement and persistence semantics; the durability barrier binds only the object store provider. +- `LogStore::release_region` has a default implementation that does nothing, so Raft Engine, Kafka and the noop provider keep their handling of a dropped region. - The object-sequence-major id scheme needs no migration of existing objects. A prefix written under the contiguous scheme keeps its objects, and the sequence allocation rule under *Entry ids* keeps new ids above the old ones. Garbage collection reasons from footers and from the reclaim boundary, not from the id encoding, so it applies to objects of either scheme. - `on_corrupted_segment` is a new configuration option with a default, so existing configurations decode unchanged; it changes what a read does with a damaged segment, not what is written. @@ -238,12 +242,13 @@ The reference implementation was validated on the following evidence; the tests | Recovery never silently loses data | Unit coverage of every corruption class at recovery and on read; footer-only recovery compared with a full decode; a footer longer than the tail window. | Structural damage fails construction of the store; a corrupted segment is skipped, recorded as a hole of its region and counted, or fails the read under `fail`; the footer-only catalog equals the full decode. | | A region cannot be silently pointed at another WAL | Reopen on a mismatched prefix; reopen on an empty prefix. | The mismatch is rejected with both prefixes named and the region is not opened; the empty prefix opens cleanly with the expected row counts. | | Garbage collection deletes everything below the persisted reclaim boundary, and recovery stays correct while objects disappear | Unit tests of the boundary computation (a region without a watermark holds it at its first object, a watermark inside an object keeps the object, `obsolete_all` releases it, last indexed plus one when everything is flushed), of stamping and durability (headers carry the boundary current at sealing, a failed create's boundary permits nothing, the start object carries the recovered boundary), of deletion (chain objects, listed orphans and failed batches below the durable boundary, nothing at or above it, a failed delete retried, `stop` waiting for deletes in flight, holes dropped), of recovery (two restarts with a collection between them, a late object of an earlier epoch below the boundary collected, a missing predecessor above the boundary still failing) and of a read that listed an object before it was collected; the Mito flush-then-collect test in both acknowledgement modes; the standalone integration test's assertion after a flush. | Objects below the durable boundary are gone after a flush, the start object included; a restart replays everything above the boundary; exactly one WAL object remains after a flush and one more insert in the integration test. | -| Every row of the *Failure matrix* holds, and every metric moves as documented | A table-driven fault matrix, `test_store_fault_matrix`, that drives every failure the object store access can produce through the store and asserts, per case, which waiters fail and with what error, whether the store is poisoned, whether the sequence is consumed, the indexed gauges at the failure itself, the exact movement of every metric, and what a read of the store sees afterwards; every timer is observed twice around two forced waits whose intervals cannot overlap. | 36 cases: 19 for the eleven rows of the matrix and 17 for the failures listed under *Failures driven through the store*. The test fails if a row of the matrix is left uncovered. | +| A dropped region holds no object after a restart that never opens it | Unit tests of the release record an object carries and the boundary it allows; of recovery seeding the watermark from the release record before any boundary is computed, with collection proceeding once the remaining region is flushed; of a failed release carried by a timer tick on an idle prefix and by a `durable` stop, which reports the record lost when its create fails too; of the release of a region without segments, and of a region id used again after its release; of a release record that outlives a restart together with the segment it covers and goes once both are below the boundary, deleted out of order; the fault-matrix cases of a release, which return only once its object is durable and indexed, fail with its create, and fail at once while the store is failing; the Mito test that drops one of two regions, restarts gracefully, abruptly, or after a release that failed, with nothing written in between, then flushes the other region and writes once more; the metric engine test of a physical region whose data and metadata regions wrote to the WAL, dropped while the release of the data region fails, while both releases fail, and while the release of the data region meets a conflicting object that poisons the store, which still drops both regions and forgets the physical region, and, in the first two cases, after which a store opened later keeps no segment of either; the frozen object with release records. | Recovery seeds the release, and after the flush and the next write every object written before the restart is deleted. | +| Every row of the *Failure matrix* holds, and every metric moves as documented | A table-driven fault matrix, `test_store_fault_matrix`, that drives every failure the object store access can produce through the store and asserts, per case, which waiters fail and with what error, whether the store is poisoned, whether the sequence is consumed, the indexed gauges at the failure itself, the exact movement of every metric, and what a read of the store sees afterwards; every timer is observed twice around two forced waits whose intervals cannot overlap. | 42 cases: 24 for the eleven rows of the matrix, five of them for the create of a release record, and 18 for the failures listed under *Failures driven through the store*. The test fails if a row of the matrix is left uncovered. | | Recovery cost does not grow with object size | Footer-only recovery over 2,000 objects; recovery concurrency bounded and ordered by sequence. | Construction of the store fell from 3.6 s to 0.9 s. | ### Failures driven through the store -- Conditional create outcomes: success, identical retry, conflicting content from an earlier epoch and from the same or a later one, lost response, transient failure in `durable` mode with later objects already created and with creates still in flight (every batch that is not indexed fails, no sequence is reused, late completions are orphans), transient failure in `enqueued` mode and after `stop` began, a failure that requires intervention in `enqueued` mode and a `stop` while it lasts, completion after a conflict had already poisoned the store, a create at the last representable sequence. +- Conditional create outcomes: success, identical retry, conflicting content from an earlier epoch and from the same or a later one, lost response, transient failure in `durable` mode with later objects already created and with creates still in flight (every batch that is not indexed fails, no sequence is reused, late completions are orphans), transient failure in `enqueued` mode and after `stop` began, a failure that requires intervention in `enqueued` mode and a `stop` and releases while it lasts, completion after a conflict had already poisoned the store, a create at the last representable sequence. - Object store access: an unwritable store; a listing, a trailer fetch, a footer fetch, a segment fetch and a delete that fail; a delete that fails and is retried by the next collection; an object collected between the listing and its fetch, below and at the tip's reclaim boundary; an object store that answers slowly. - Store state: an object the catalog rejects, per-region entry id exhaustion, the catalog refusing a next sequence past the maximum, every stop path with several creates in flight that succeed, fail transiently or conflict after stop began, a stop with a delete in flight, an append held back by full sealed batches, and the appends the backlog thresholds were holding back released by admission, by a stop, by a terminal error and by the teardown of the actor. - Pipelined uploader: concurrent creates acknowledged in sequence order, every sealed batch linked to its pipeline predecessor, the open batch sealing at the threshold while creates are in flight, a transient failure failing every batch that is not indexed, created ones included, a permanent failure poisoning the store. @@ -322,7 +327,9 @@ Simpler catalog, but every region then produces its own timer-driven object, so - Write latency under the `durable` mode is the time until the batch seals (up to `flush_interval`, 100 ms by default), plus waiting for every earlier object, plus one PUT. A serial writer waits for the next seal, so its extra latency over Raft Engine is the interval minus the base insert cost modulo the interval, and concurrency does not close the gap: measured within one session on a laptop it was about 23% of throughput at concurrency 1 and about 28% at concurrency 8, and on S3 against Raft Engine on a local NVMe store the `durable` mode reached 31% of Raft Engine's throughput at concurrency 1 and 54% at concurrency 8, a serial writer being bound to one batch per interval, while the `enqueued` mode matched it, see *Write path measurements* and *Cloud write path measurements*. - Every object costs a request. Timer-triggered sealing creates at most one object per interval per node, ten per second at the default; the `enqueued` mode can seal earlier when the unpersisted backlog reaches a threshold, and the batch seals on size only when the admitted bytes reach `max_batch_bytes` before the next tick. Under sustained load on S3 the timer ceiling was reached, about 600 creates per minute per node, about $0.18 per node-hour at the S3 standard list price. -- Garbage collection keeps every object at or above the reclaim boundary, so a region that is written but rarely flushed, or never flushed, holds the boundary at its first object and keeps every later object the prefix shares with other regions; recovery time grows with the number of retained objects. A prefix that receives no appends cannot persist a higher boundary, so an idle prefix is collected only after its next object. +- Garbage collection keeps every object at or above the reclaim boundary, so a region that is written but rarely flushed, or never flushed, holds the boundary at its first object and keeps every later object the prefix shares with other regions; recovery time grows with the number of retained objects. A prefix that receives no appends cannot persist a higher boundary, so an idle prefix is collected only after its next object; a release seals that object itself. +- A dropped region stays unreleased, and holds the reclaim boundary for good, after a crash between the deletion of its table's metadata and a durable release record, or after a crash that follows a failed release before any later create succeeds; `indexed_objects` then only grows. The first window is a gap of the drop procedure for every WAL provider, issue #9577. See *Entry ids*. +- A release costs one object of its own and waits for every object sealed before it, in the `enqueued` mode as well, so dropping a region waits for the uploads ahead of its release record; when they fail, or the store is stopped or poisoned, the drop succeeds all the same and logs the failure; a stopped or poisoned store leaves the region unreleased. - Every open of the store writes one start object without segments, one extra PUT per open, and every object carries a 54-byte header. Orphans, the objects of failed batches that were stored anyway, hold storage together with the rows Mito wrote again after those failures until garbage collection deletes them. - Object-sequence-major entry ids consume 20 bits per object for positions, which caps a region at about a million entries per object; the batch seals when the cap is reached. Sequence numbers are 44 bits wide as a result, enough for one object per second for half a million years. @@ -344,12 +351,13 @@ Sealing may only exclude objects whose writes were never acknowledged, which tak **Corruption.** Segment-level skipping with the hole recorded per region is implemented, see *Object format*. *Proposed*: cluster mode routes the WAL hole event to the metasrv so that a follower or a takeover knows the region needs a fresh replica rather than a replay. -**Garbage collection across nodes.** Within one prefix, garbage collection is implemented as described under *Entry ids*: watermark from each region's manifest, per-segment footer comparison, a reclaim boundary persisted in every header. A source prefix that still holds segments of regions which have moved away cannot be trimmed by the source alone, because it no longer learns those regions' watermarks. Either the target reports to the metasrv, after each flush, the source sequence below which it no longer needs the region, and the metasrv drives deletion, or a cluster-level janitor collects the watermarks; which of the two is the last of the *Unresolved questions*. +**Garbage collection across nodes.** Within one prefix, garbage collection is implemented as described under *Entry ids*: watermark from each region's manifest, per-segment footer comparison, a reclaim boundary persisted in every header. A source prefix that still holds segments of regions which have moved away cannot be trimmed by the source alone, because it no longer learns those regions' watermarks. Either the target reports to the metasrv, after each flush, the source sequence below which it no longer needs the region, and the metasrv drives deletion, or a cluster-level janitor collects the watermarks; which of the two is the last of the *Unresolved questions*. Either way reclamation consults the metadata for the regions that still exist, as Kafka's pruning by the metasrv does, which also lets go of a dropped region whose release never became durable, see *Entry ids*. # Future work 1. Probe the next sequence with a GET from the last known sequence instead of listing the whole prefix at recovery; on a prefix with many objects the probe is cheaper than the listing. 2. Distributed mode, which needs metasrv-side allocation of node ids and generations, the takeover chain in `WalOptions`, read-only replay of a foreign prefix, and cross-node garbage collection, as sketched under *Alignment with cluster mode*. +3. Close the crash window of the drop procedure, issue #9577 (https://github.com/GreptimeTeam/greptimedb/issues/9577), which also leaves the region's WAL unreleased, see *Entry ids*. # Unresolved questions diff --git a/src/log-store/src/metrics.rs b/src/log-store/src/metrics.rs index c787820aaeb..8f923615964 100644 --- a/src/log-store/src/metrics.rs +++ b/src/log-store/src/metrics.rs @@ -111,6 +111,13 @@ lazy_static! { "object store logstore failed object deletes total", ) .unwrap(); + /// Counter of release records in the objects the object store logstore + /// indexed, including records carried over from batches that failed. + pub static ref METRIC_OBJECT_STORE_WAL_RELEASE_RECORDS_TOTAL: IntCounter = register_int_counter!( + "greptime_logstore_object_store_wal_release_records_total", + "object store logstore release records made durable total", + ) + .unwrap(); /// How long an object of the object store logstore takes to become durable /// once its batch is sealed, sampled for every conditional create that @@ -294,7 +301,8 @@ lazy_static! { /// Counter of appends the object store logstore rejected while its /// backlog is refused, see the gauge above. They take no entry id and - /// fail with a retryable error. + /// fail with a retryable error. Releases of dropped regions failed in + /// that state count too; their records wait for the next object. pub static ref METRIC_OBJECT_STORE_WAL_REJECTED_APPENDS_TOTAL: IntCounter = register_int_counter!( "greptime_logstore_object_store_wal_rejected_appends_total", "object store logstore appends rejected while the object store refuses the backlog total", diff --git a/src/log-store/src/object_store_wal.rs b/src/log-store/src/object_store_wal.rs index 6a25afb3e77..68046487873 100644 --- a/src/log-store/src/object_store_wal.rs +++ b/src/log-store/src/object_store_wal.rs @@ -51,11 +51,19 @@ //! The reclaim boundary of an object states that every object with a lower //! sequence was fully flushed when the object was sealed, so the chain need //! not reach below it. Every present object below the boundary of the last -//! indexed object is deleted, on the chain or not. +//! indexed object is deleted, on the chain or not. A region that is dropped, +//! or whose entries are all made obsolete, is released: a release record of +//! the region goes into an object sealed at once, as the single entry of a +//! segment of a reserved region id, so the object keeps format version 1. +//! Recovery reads those segments of the chain and seeds the obsolete +//! watermarks from them, so a store opened later collects the objects of a +//! region that is never opened again. //! //! Recovery lists the objects, reads and verifies only the header, trailer and //! footer of each, and indexes the footers of the chain in sequence order to -//! rebuild the object catalog, which rejects a sequence it already holds. +//! rebuild the object catalog, which rejects a sequence it already holds. Of +//! the segments it reads only those of release records, and one that does not +//! decode fails recovery. //! Segments are read and checksummed only when a read decodes them; a segment //! that does not decode is fetched once more and then skipped as a hole of its //! region or fails the read, as `on_corrupted_segment` says. diff --git a/src/log-store/src/object_store_wal/catalog.rs b/src/log-store/src/object_store_wal/catalog.rs index 70f7aa5b329..61b4d2c09cc 100644 --- a/src/log-store/src/object_store_wal/catalog.rs +++ b/src/log-store/src/object_store_wal/catalog.rs @@ -25,7 +25,7 @@ use crate::error::{ CorruptedWalObjectSnafu, InvalidWalEntryRangeSnafu, Result, WalObjectSequenceExhaustedSnafu, }; use crate::object_store_wal::batch::{OBJECT_SEQ_LIMIT, sequence_floor}; -use crate::object_store_wal::format::FooterEntry; +use crate::object_store_wal::format::{FooterEntry, RELEASE_REGION_ID}; /// Indexes objects by sequence and, per region, the objects that hold entries /// of that region. @@ -165,7 +165,9 @@ impl ObjectCatalog { /// /// A segment is live when its maximum entry id is above the watermark of /// its region, or its region has no watermark; a region whose watermark - /// is [`EntryId::MAX`] holds no live segment. Watermarks never move down + /// is [`EntryId::MAX`] holds no live segment, and neither does + /// [`RELEASE_REGION_ID`], whose records recovery reads only from the + /// objects it indexes. Watermarks never move down /// and a region first appears in an object above the boundary, so the /// scan starts at the previous boundary and never reads an object twice /// once the boundary has passed it. @@ -176,9 +178,10 @@ impl ObjectCatalog { ) -> u64 { let live = self.objects.range(from..).find(|(_, footer)| { footer.iter().any(|entry| { - obsolete_entry_ids - .get(&entry.region_id) - .is_none_or(|obsolete| entry.max_entry_id > *obsolete) + entry.region_id != RELEASE_REGION_ID + && obsolete_entry_ids + .get(&entry.region_id) + .is_none_or(|obsolete| entry.max_entry_id > *obsolete) }) }); match live { @@ -511,7 +514,10 @@ mod tests { catalog .insert_object( 2, - vec![footer_entry(region_one, entry_id(2, 1), entry_id(2, 1))], + vec![ + footer_entry(region_one, entry_id(2, 1), entry_id(2, 1)), + footer_entry(RELEASE_REGION_ID, entry_id(2, 1), entry_id(2, 1)), + ], ) .unwrap(); catalog @@ -542,6 +548,7 @@ mod tests { &[(region_one, entry_id(1, 2)), (region_two, entry_id(1, 1))] ) ); + // The segment of release records never holds the boundary. assert_eq!( 3, boundary( diff --git a/src/log-store/src/object_store_wal/format.rs b/src/log-store/src/object_store_wal/format.rs index 7a255711784..cbd35287050 100644 --- a/src/log-store/src/object_store_wal/format.rs +++ b/src/log-store/src/object_store_wal/format.rs @@ -77,6 +77,20 @@ pub(crate) struct Record { pub(crate) payload: Bytes, } +/// The region whose segment holds the release records of an object, as a +/// single entry whose payload [`encode_releases`] writes. It is an ordinary +/// segment, so a reader that predates release records reads the object as +/// any other; no region is ever assigned this id. +pub(crate) const RELEASE_REGION_ID: RegionId = RegionId::from_u64(u64::MAX); + +/// A release record: no entry of the region with an id at or below +/// `entry_id` is needed any more, `u64::MAX` releasing the region in full. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct Release { + pub(crate) region_id: RegionId, + pub(crate) entry_id: u64, +} + /// Footer entry describing the segment of one region. #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) struct FooterEntry { @@ -470,6 +484,51 @@ pub(crate) fn decode_segment(bytes: &[u8], entry: &FooterEntry) -> Result) -> Bytes { + let mut output = BytesMut::with_capacity(4 + releases.len() * 16); + output.put_u32(releases.len() as u32); + for (region_id, entry_id) in releases { + output.put_u64(region_id.as_u64()); + output.put_u64(*entry_id); + } + output.freeze() +} + +/// Decodes the payload [`encode_releases`] writes. +pub(crate) fn decode_releases(payload: &[u8]) -> Result> { + let mut reader = Reader::new(payload); + let count = reader.u32("release records")? as usize; + ensure!( + count > 0 && reader.remaining_len() == count * 16, + CorruptedWalObjectSnafu { + reason: format!( + "release records declare {count} records in {} bytes", + reader.remaining_len() + ), + } + ); + let mut releases = Vec::::with_capacity(count); + for _ in 0..count { + let region_id = RegionId::from_u64(reader.u64("release records")?); + ensure!( + releases + .last() + .is_none_or(|last| last.region_id < region_id), + CorruptedWalObjectSnafu { + reason: format!("release records are not strictly ordered at region {region_id}"), + } + ); + releases.push(Release { + region_id, + entry_id: reader.u64("release records")?, + }); + } + Ok(releases) +} + /// Lower bound on the length of an object: header, a footer holding only its /// entry count and the trailer. pub(crate) const MIN_OBJECT_LEN: usize = HEADER_LEN + 4 + TRAILER_LEN; @@ -1138,6 +1197,86 @@ mod tests { } } + /// An object with an entry of region (1, 1) and the release records of + /// regions (2, 1) and (3, 0), frozen under `tests/compatibility`. + const FIXTURE_RELEASE: &[u8] = + include_bytes!("../../../../tests/compatibility/wal/object_store/release_record.wal"); + + fn release_fixture_records() -> Vec { + let releases = BTreeMap::from([ + (RegionId::new(2, 1), (2 << 20) | 1), + (RegionId::new(3, 0), u64::MAX), + ]); + [ + (RegionId::new(1, 1), Bytes::from_static(b"a1")), + (RELEASE_REGION_ID, encode_releases(&releases)), + ] + .into_iter() + .map(|(region_id, payload)| Record { + region_id, + entry_id: (3 << 20) | 1, + payload, + }) + .collect() + } + + #[test] + fn test_format_matches_release_record_fixture() { + let header = Header { + object_seq: 3, + epoch: 2, + prev: Some(ChainLink { + object_seq: 2, + epoch: 2, + }), + reclaim_boundary: 1, + }; + let encoded = encode_object(header.clone(), &release_fixture_records()).unwrap(); + assert_eq!(FIXTURE_RELEASE, encoded.bytes.as_ref()); + + // The object is a version 1 object whose last segment holds the + // release records, so a reader that predates them decodes it. + assert_eq!([0, 1], FIXTURE_RELEASE[8..10]); + let decoded = decode_object(FIXTURE_RELEASE).unwrap(); + assert_eq!(header, decoded.header); + assert_eq!(release_fixture_records(), decoded.records); + assert_eq!( + vec![ + Release { + region_id: RegionId::new(2, 1), + entry_id: (2 << 20) | 1, + }, + Release { + region_id: RegionId::new(3, 0), + entry_id: u64::MAX, + }, + ], + decode_releases(&decoded.records[1].payload).unwrap() + ); + } + + #[test] + fn test_format_rejects_invalid_release_records() { + let one = BTreeMap::from([(RegionId::new(1, 1), 1)]); + let encoded = encode_releases(&one); + for (payload, reason) in [ + (&[0, 0, 0, 0][..], "declare 0 records in 0 bytes"), + ( + &encoded[..encoded.len() - 1], + "declare 1 records in 15 bytes", + ), + ] { + assert_corrupted(decode_releases(payload), reason); + } + let mut unordered = encode_releases(&BTreeMap::from([ + (RegionId::new(1, 1), 1), + (RegionId::new(1, 2), 1), + ])) + .to_vec(); + unordered[4..12].copy_from_slice(&RegionId::new(1, 2).as_u64().to_be_bytes()); + assert_corrupted(decode_releases(&unordered), "not strictly ordered"); + } + fn footer_ids(entry: &FooterEntry) -> (u64, u64, u32) { (entry.min_entry_id, entry.max_entry_id, entry.entry_count) } diff --git a/src/log-store/src/object_store_wal/store.rs b/src/log-store/src/object_store_wal/store.rs index f233f73bdea..fc32bef707e 100644 --- a/src/log-store/src/object_store_wal/store.rs +++ b/src/log-store/src/object_store_wal/store.rs @@ -55,7 +55,7 @@ use crate::metrics::{ METRIC_OBJECT_STORE_WAL_OBJECT_ENTRIES, METRIC_OBJECT_STORE_WAL_POISONED_TOTAL, METRIC_OBJECT_STORE_WAL_READ_SECONDS, METRIC_OBJECT_STORE_WAL_RECOVERED_OBJECTS_TOTAL, METRIC_OBJECT_STORE_WAL_RECOVERY_SECONDS, METRIC_OBJECT_STORE_WAL_REJECTED_APPENDS_TOTAL, - METRIC_OBJECT_STORE_WAL_SEAL_TO_DURABLE_SECONDS, + METRIC_OBJECT_STORE_WAL_RELEASE_RECORDS_TOTAL, METRIC_OBJECT_STORE_WAL_SEAL_TO_DURABLE_SECONDS, METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL, METRIC_OBJECT_STORE_WAL_STALLED_APPEND_SECONDS, METRIC_OBJECT_STORE_WAL_STALLED_APPENDS_TOTAL, }; @@ -63,8 +63,9 @@ use crate::object_store_wal::batch::{OBJECT_SEQ_LIMIT, OpenBatch, entry_id, sequ use crate::object_store_wal::catalog::ObjectCatalog; use crate::object_store_wal::format::{ ChainLink, EncodedObject, FixedTrailer, FooterEntry, HEADER_LEN, Header, MIN_OBJECT_LEN, - Record, TRAILER_LEN, decode_footer, decode_header, decode_segment, decode_trailer, - encode_object, footer_range, object_len, verify_segment_ranges, + RELEASE_REGION_ID, Record, Release, TRAILER_LEN, decode_footer, decode_header, decode_releases, + decode_segment, decode_trailer, encode_object, encode_releases, footer_range, object_len, + verify_segment_ranges, }; use crate::object_store_wal::io::{ListedObject, ObjectStoreIo, PutResult}; @@ -138,7 +139,9 @@ const RECOVERY_TAIL_WINDOW: usize = 64 * 1024; /// appends are rejected while the object store refuses the backlog with an /// error that requires intervention. A segment that a read cannot /// decode is handled as `on_corrupted_segment` says. Objects below the reclaim -/// boundary of the last indexed object are deleted in the background. +/// boundary of the last indexed object are deleted in the background. A +/// release of a region seals an object with a release record of the region at +/// once and returns, in both modes, once that object is durable and indexed. pub struct ObjectStoreLogStore { prefix: String, ack_mode: AckMode, @@ -315,6 +318,7 @@ impl ObjectStoreLogStore { max_epoch, reclaim_boundary, unindexed, + released_entry_ids, recovered_objects, } = recover(io.as_ref()).await?; // Every epoch is one above the sequence of the start object its @@ -349,10 +353,11 @@ impl ObjectStoreLogStore { })?; METRIC_OBJECT_STORE_WAL_INDEXED_OBJECTS.inc(); METRIC_OBJECT_STORE_WAL_INDEXED_BYTES.add(object_len(&[]) as i64); - // No watermark is recorded yet. - let writer_reclaim_boundary = catalog.reclaim_boundary(reclaim_boundary, &HashMap::new()); + // The release records are the only watermarks recorded yet. + let writer_reclaim_boundary = + catalog.reclaim_boundary(reclaim_boundary, &released_entry_ids); let catalog = Arc::new(RwLock::new(catalog)); - let obsolete_entry_ids = ObsoleteEntryIds::default(); + let obsolete_entry_ids = Arc::new(Mutex::new(released_entry_ids)); let wal_holes = WalHoles::default(); let deleting = Arc::new(DeletingObjects::default()); let terminal_error = TerminalError::default(); @@ -403,6 +408,7 @@ impl ObjectStoreLogStore { reclaim_boundary: writer_reclaim_boundary, durable_reclaim_boundary: reclaim_boundary, unindexed: unindexed.into_iter().chain(moved_past).collect(), + failed_releases: BTreeMap::new(), creating: HashSet::new(), deleting: deleting.clone(), deletes: FuturesUnordered::new(), @@ -521,6 +527,30 @@ impl ObjectStoreLogStore { None => Ok(()), } } + + /// Releases the region: its watermark moves up to the release in memory + /// and a release record of it goes into an object sealed at once, and the + /// call returns once that object is durable and indexed, in both + /// acknowledgement modes, or with the error that failed it. A `full` + /// release reaches [`EntryId::MAX`]; any other the largest id the store + /// handed out for the region, which every entry of the region it holds is + /// at or below. When the object fails, the next sealed batch carries the + /// record again. Recovery seeds the watermark from the record, so a store + /// opened later collects the objects of a region that is never opened + /// again. + async fn release(&self, region_id: RegionId, full: bool) -> Result<()> { + let (response_tx, response_rx) = oneshot::channel(); + self.command_tx + .send(Command::Release { + region_id, + full, + response: response_tx, + }) + .await + .ok() + .context(ObjectStoreWalStoppedSnafu)?; + response_rx.await.ok().context(ObjectStoreWalStoppedSnafu)? + } } /// The transient object store error of a create that a testing hook fails. @@ -607,6 +637,12 @@ impl ObjectStoreLogStore { self.creates_fail.store(true, Ordering::Release); } + /// Lets the creates that run from now on write again after + /// [`fail_creates`](Self::fail_creates). + pub fn allow_creates(&self) { + self.creates_fail.store(false, Ordering::Release); + } + /// Makes the next create that writes its object report a transient object /// store error afterwards, so the object exists although its create /// failed. @@ -660,8 +696,8 @@ impl ObjectStoreLogStore { /// Waits until no object delete is in flight. The deletes a collection /// starts are in flight by the time the `obsolete`, `obsolete_all`, - /// [`seal_open_batch`](Self::seal_open_batch) or `durable` append that - /// started it returns. + /// `release_region`, [`seal_open_batch`](Self::seal_open_batch) or + /// `durable` append that started it returns. pub async fn wait_for_garbage_collection(&self) { self.deleting.wait_until_empty().await; } @@ -888,6 +924,7 @@ impl LogStore for ObjectStoreLogStore { let regions = catalog .objects_in_order() .flat_map(|(_, footer)| footer.iter().map(|entry| entry.region_id)) + .filter(|region_id| *region_id != RELEASE_REGION_ID) .collect::>(); Ok(regions .into_iter() @@ -913,7 +950,9 @@ impl LogStore for ObjectStoreLogStore { /// ids under the next sequence the floor cannot move, and the call fails /// with [`Error::WalObjectSequenceUnsettled`] and records neither. The /// watermark is kept in memory; the caller re-establishes it after a - /// restart. In the `enqueued` mode an `entry_id` this store handed out + /// restart unless it released the region, see + /// [`release_region`](LogStore::release_region). In the `enqueued` mode + /// an `entry_id` this store handed out /// needs no floor: it is never handed out again while the store runs. async fn obsolete( &self, @@ -949,23 +988,19 @@ impl LogStore for ObjectStoreLogStore { } } - /// Hides every entry of the region in memory and lets collection delete - /// its objects. The caller re-establishes the watermark after a restart. + /// Hides every entry of the region and lets collection delete its + /// objects, by a release in full that is durable once this returns. async fn obsolete_all(&self, provider: &Provider, region_id: RegionId) -> Result<()> { self.check_region(provider, region_id)?; - record_obsolete(&self.obsolete_entry_ids, region_id, EntryId::MAX); - let (response_tx, response_rx) = oneshot::channel(); - if self - .command_tx - .send(Command::Collect { - response: response_tx, - }) - .await - .is_ok() - { - let _ = response_rx.await; - } - Ok(()) + self.release(region_id, true).await + } + + /// Releases the region through the largest id handed out for it, durably + /// once this returns. Should the region id be used again, its new entries + /// have greater ids and stay readable. + async fn release_region(&self, provider: &Provider, region_id: RegionId) -> Result<()> { + self.check_region(provider, region_id)?; + self.release(region_id, false).await } fn entry( @@ -1031,8 +1066,15 @@ enum Command { watermark: EntryId, response: oneshot::Sender>, }, - /// Answered once the collection the recorded watermarks allow started. - Collect { response: oneshot::Sender<()> }, + /// Answered once the object holding the release record of the region is + /// durable and indexed, or with the reason it is not. A `full` release + /// reaches [`EntryId::MAX`], any other the largest id handed out for the + /// region. + Release { + region_id: RegionId, + full: bool, + response: oneshot::Sender>, + }, Stop { response: oneshot::Sender>, }, @@ -1074,11 +1116,15 @@ struct SealedBatch { reclaim_boundary: u64, bytes: Bytes, footer: Vec, + releases: Vec, first_admitted_at: Instant, sealed_at: Instant, /// Number of creates that were attempted for the batch. attempts: u32, waiters: Vec, + /// Callers of a release whose record the batch carries, answered once the + /// object is durable and indexed. + release_waiters: Vec>>, #[cfg(any(test, feature = "testing"))] seal_waiters: Vec>>, state: CreateState, @@ -1093,6 +1139,9 @@ impl SealedBatch { for waiter in self.waiters { let _ = waiter.response.send(Err(error())); } + for waiter in self.release_waiters { + let _ = waiter.send(Err(error())); + } #[cfg(any(test, feature = "testing"))] for waiter in self.seal_waiters { let _ = waiter.send(Err(error())); @@ -1113,7 +1162,8 @@ type CreateOutcome = (u64, Result); /// flight, after a delay that doubles from [`CREATE_RETRY_DELAY`] up to /// [`MAX_FAILING_RETRY_DELAY`]. Appends are rejected at once with a retryable /// error that carries the last error: they take no entry id and never wait -/// behind the backlog. The first create that succeeds ends the state. +/// behind the backlog. Releases fail with it too, and their records wait for +/// a create. The first create that succeeds ends the state. struct Failing { /// The last error a create met, which rejected appends carry. error: Arc, @@ -1159,9 +1209,9 @@ impl Failing { /// | --- | --- | --- | --- | /// | created, or identical retry | advances | acknowledged in order | healthy | /// | transient error, `durable` mode | never reused: a create that reported an error may have written its object | every batch that is not indexed fails, created ones included; retries are assigned new ids | healthy: the next batch links to the last indexed object, so no chain reaches a failed batch | -/// | transient error, `enqueued` mode | unchanged | already acknowledged | healthy: the create is repeated with the same bytes after [`CREATE_RETRY_DELAY`], which an identical retry accepts | -/// | error that requires intervention, `enqueued` mode | unchanged | already acknowledged; later appends are rejected with a retryable error | failing, see [`Failing`]: the lowest pending create is repeated alone with backoff, and the first create that succeeds ends it | -/// | object store error, `enqueued` mode after `stop` began | never reused | already acknowledged | the backlog is dropped and `stop` reports the error | +/// | transient error, `enqueued` mode | unchanged | already acknowledged, except a release, which waits for the repeat | healthy: the create is repeated with the same bytes after [`CREATE_RETRY_DELAY`], which an identical retry accepts | +/// | error that requires intervention, `enqueued` mode | unchanged | already acknowledged, except a release, which fails while its record stays in the backlog; later appends and releases are rejected with a retryable error | failing, see [`Failing`]: the lowest pending create is repeated alone with backoff, and the first create that succeeds ends it | +/// | object store error, `enqueued` mode after `stop` began | never reused | already acknowledged, except a release, which fails with the stopped error | the backlog is dropped and `stop` reports the error | /// | an object of an earlier epoch holds the sequence, `durable` mode | never reused | as for a transient error | healthy: that object can never be on the chain | /// | an object of the same or a later epoch holds the sequence, any conflicting object in the `enqueued` mode, encoding or catalog error | unchanged | every batch that is not indexed fails | poisoned | /// | created at the last representable sequence | cannot advance | acknowledged | poisoned: no later batch can be allocated a sequence | @@ -1169,8 +1219,9 @@ impl Failing { /// A create still in flight when its batch failed runs to completion and its /// outcome is ignored: whatever it stored is off the chain. After `stop` began /// nothing is admitted and no create starts, except that the `enqueued` mode -/// uploads its backlog; creates in flight run to completion and acknowledge -/// if they succeed. +/// uploads its backlog and the `durable` mode creates one object with the +/// release records of failed batches; creates in flight run to completion and +/// acknowledge if they succeed. /// /// Appends arrive on their own channel. The actor holds one back while the /// sealed batches leave no room for the two an admission can seal, or at a @@ -1180,12 +1231,13 @@ impl Failing { /// /// The actor also collects garbage. Every batch carries the reclaim boundary /// current when it was sealed, see [`ObjectCatalog::reclaim_boundary`], which -/// is recomputed after every `obsolete` and every indexed object and never -/// moves down. The boundary of the last indexed object is durable: every -/// present object below it, indexed or not, is deleted in the background, up -/// to [`MAX_IN_FLIGHT_DELETES`] at a time, after an `obsolete` or an indexing -/// that moved it. An indexed object is unindexed once its delete succeeded; a -/// failed delete is attempted again at the next collection. Nothing is +/// is recomputed after every `obsolete`, every release and every indexed +/// object and never moves down. The boundary of the last indexed object is +/// durable: every present object below it, indexed or not, is deleted in the +/// background, up to [`MAX_IN_FLIGHT_DELETES`] at a time, after an +/// `obsolete`, a release or an indexing that moved it. An indexed object is +/// unindexed once its delete succeeded; a failed delete is attempted again at +/// the next collection. Nothing is /// collected once stop began or the store is poisoned, and deletes in flight /// complete before `stop` returns. struct Actor { @@ -1240,6 +1292,11 @@ struct Actor { /// Sequences of objects that may be present but are not indexed: the /// objects recovery listed off the chain and the batches that failed. unindexed: BTreeSet, + /// Release records of batches that failed, which the next sealed batch + /// carries again, so a release whose call failed still becomes durable + /// while the store runs. A timer tick seals them when no entry is + /// admitted. + failed_releases: BTreeMap, /// Sequences whose create is in flight. creating: HashSet, deleting: Arc, @@ -1297,10 +1354,8 @@ impl Actor { Some(Command::Obsolete { region_id, entry_id, watermark, response }) => { self.handle_obsolete(region_id, entry_id, watermark, response); } - Some(Command::Collect { response }) => { - self.update_reclaim_boundary(); - self.collect(); - let _ = response.send(()); + Some(Command::Release { region_id, full, response }) => { + self.handle_release(region_id, full, response); } Some(Command::Stop { response }) => { self.handle_stop(response); @@ -1474,7 +1529,22 @@ impl Actor { /// Seals the open batch as the object `next_object_seq` and starts its /// create when a slot is free. Returns whether a batch was sealed. fn flush_open_batch(&mut self) -> bool { - if self.open_batch.is_empty() { + self.seal(Vec::new()) + } + + /// Seals the open batch together with `releases` and the release records + /// of failed batches as the object `next_object_seq`, even when the open + /// batch holds no entry, and starts its create when a slot is free. + /// Returns whether a batch was sealed. + /// + /// A failing store keeps the records of failed batches coalesced instead + /// of sealing them alone behind its backlog on every tick, until `stop` + /// seals them for its last attempt of the backlog. + fn seal(&mut self, releases: Vec) -> bool { + if self.open_batch.is_empty() + && releases.is_empty() + && (self.failed_releases.is_empty() || (self.failing.is_some() && self.stop.is_empty())) + { return false; } if let Some(error) = terminal(&self.terminal_error) { @@ -1506,7 +1576,12 @@ impl Actor { ), reclaim_boundary: self.reclaim_boundary, }; - let encoded = match encode_batch(header, entries) { + let mut records = std::mem::take(&mut self.failed_releases); + for release in releases { + let entry_id = records.entry(release.region_id).or_insert(release.entry_id); + *entry_id = (*entry_id).max(release.entry_id); + } + let encoded = match encode_batch(header, entries, &records) { Ok(encoded) => encoded, Err(error) => { self.poison(error); @@ -1534,10 +1609,18 @@ impl Actor { reclaim_boundary: self.reclaim_boundary, bytes: encoded.bytes, footer: encoded.footer, + releases: records + .into_iter() + .map(|(region_id, entry_id)| Release { + region_id, + entry_id, + }) + .collect(), first_admitted_at, sealed_at: Instant::now(), attempts: 0, waiters: std::mem::take(&mut self.pending), + release_waiters: Vec::new(), #[cfg(any(test, feature = "testing"))] seal_waiters: Vec::new(), state: CreateState::Pending, @@ -1549,11 +1632,13 @@ impl Actor { /// Starts the creates of pending batches in sequence order while fewer /// than [`MAX_IN_FLIGHT_CREATES`] are in flight, counting the creates of /// batches that already failed. Nothing starts once stop began, except - /// the backlog of the `enqueued` mode. A failing store only schedules the + /// the backlog of the `enqueued` mode and, once the stop command was + /// handled, the release records [`handle_stop`](Self::handle_stop) + /// seals in the `durable` mode. A failing store only schedules the /// next attempt once nothing is in flight, see [`Failing`]; after stop /// began it uploads its backlog at once. fn start_creates(&mut self) { - if self.is_stopped() && self.ack_mode == AckMode::Durable { + if self.is_stopped() && self.ack_mode == AckMode::Durable && self.stop.is_empty() { return; } if !self.is_stopped() @@ -1760,6 +1845,13 @@ impl Actor { if let Some((_, response)) = stalled { let _ = response.send(Err(failing.reject())); } + // So would a release. Its record stays in the backlog and becomes + // durable with it. + for batch in &mut self.sealed { + for waiter in batch.release_waiters.drain(..) { + let _ = waiter.send(Err(failing.reject())); + } + } } /// Ends the failing state, if the store is in it, because of `reason`. @@ -1808,6 +1900,7 @@ impl Actor { }; METRIC_OBJECT_STORE_WAL_INDEXED_OBJECTS.inc(); METRIC_OBJECT_STORE_WAL_INDEXED_BYTES.add(batch.bytes.len() as i64); + METRIC_OBJECT_STORE_WAL_RELEASE_RECORDS_TOTAL.inc_by(batch.releases.len() as u64); self.last_indexed = ChainLink { object_seq: batch.object_seq, epoch: self.epoch, @@ -1824,6 +1917,9 @@ impl Actor { last_entry_ids: waiter.last_entry_ids, })); } + for waiter in batch.release_waiters { + let _ = waiter.send(Ok(())); + } #[cfg(any(test, feature = "testing"))] for waiter in batch.seal_waiters { let _ = waiter.send(Ok(())); @@ -1884,8 +1980,9 @@ impl Actor { shared(&error) } }; - for batch in self.sealed.drain(..) { + for mut batch in std::mem::take(&mut self.sealed) { self.unindexed.insert(batch.object_seq); + self.keep_releases(&mut batch); batch.fail(failure); } self.reset_open_batch(); @@ -1893,13 +1990,32 @@ impl Actor { // so `settle` admits it, or refuses it once stop began. self.fail_waiters(failure); // An acknowledged backlog was dropped: `stop` reports it, whether - // its caller has arrived yet or not. - if self.ack_mode == AckMode::Enqueued { + // its caller has arrived yet or not. So it does release records once + // stop began, which no later create carries. + if self.ack_mode == AckMode::Enqueued + || (!self.stop.is_empty() && !self.failed_releases.is_empty()) + { self.stop_error.get_or_insert(error); } self.leave_failing("its backlog was dropped"); } + /// Keeps the release records of a batch that failed for the next sealed + /// batch. + fn keep_releases(&mut self, batch: &mut SealedBatch) { + for release in batch.releases.drain(..) { + self.keep_release(release); + } + } + + fn keep_release(&mut self, release: Release) { + let entry_id = self + .failed_releases + .entry(release.region_id) + .or_insert(release.entry_id); + *entry_id = (*entry_id).max(release.entry_id); + } + /// Records `error` as terminal and fails every waiter that is not /// acknowledged with it, or with the stopped error if the store was /// stopped meanwhile. Creates in flight run to completion, but their @@ -1910,6 +2026,8 @@ impl Actor { fn poison(&mut self, error: Error) -> Arc { let error = set_terminal(&self.terminal_error, error); let stopped = self.is_stopped(); + let releases_lost = !self.failed_releases.is_empty() + || self.sealed.iter().any(|batch| !batch.releases.is_empty()); let failure = || { if stopped { ObjectStoreWalStoppedSnafu.build() @@ -1922,7 +2040,9 @@ impl Actor { } self.reset_open_batch(); self.fail_unacknowledged(failure); - if self.ack_mode == AckMode::Enqueued { + // As in `roll_back`, release records lost once stop began are + // reported by `stop` in both modes. + if self.ack_mode == AckMode::Enqueued || (!self.stop.is_empty() && releases_lost) { self.stop_error.get_or_insert(error.clone()); } self.leave_failing("the store is poisoned"); @@ -2014,6 +2134,63 @@ impl Actor { let _ = response.send(result); } + /// Moves the watermark of the region up to the release and seals the + /// open batch with a release record of the region; `response` is + /// answered once that object is durable and indexed. The batch carries + /// the reclaim boundary the watermark allows, so the object makes both + /// durable at once. + /// + /// The watermark is not clamped to the durable entry id in the `enqueued` + /// mode: a released region is never replayed, and the record is durable + /// only once every object sealed before it is. A release arrives as a + /// command, so [`MAX_SEALED_BATCHES`] does not hold it back; every caller + /// waits for its own object, so releases add at most one batch per caller. + fn handle_release( + &mut self, + region_id: RegionId, + full: bool, + response: oneshot::Sender>, + ) { + if self.is_stopped() { + let _ = response.send(Err(ObjectStoreWalStoppedSnafu.build())); + return; + } + if let Some(error) = terminal(&self.terminal_error) { + let _ = response.send(Err(shared(&error))); + return; + } + let entry_id = if full { + EntryId::MAX + } else { + self.issued_entry_ids.get(®ion_id).copied().unwrap_or(0) + }; + record_obsolete(&self.obsolete_entry_ids, region_id, entry_id); + self.update_reclaim_boundary(); + self.collect(); + let release = Release { + region_id, + entry_id, + }; + // A failing store admits nothing behind its backlog: the release + // fails at once, and its record is kept like that of a batch that + // failed, sealed once a create succeeds. + if let Some(failing) = &mut self.failing { + let _ = response.send(Err(failing.reject())); + self.keep_release(release); + return; + } + if self.seal(vec![release]) + && let Some(batch) = self.sealed.back_mut() + { + batch.release_waiters.push(response); + return; + } + // Sealing failed and poisoned the store. + if let Some(error) = terminal(&self.terminal_error) { + let _ = response.send(Err(shared(&error))); + } + } + /// Raises the reclaim boundary to what the catalog and the obsolete /// watermarks allow. fn update_reclaim_boundary(&mut self) { @@ -2162,7 +2339,8 @@ impl Actor { /// Begins stopping. Nothing is admitted from now on and the append held /// back is refused; the `durable` mode drops the open batch and the - /// batches whose create has not started, the `enqueued` mode seals its + /// batches whose create has not started, and seals the release records of + /// failed and dropped batches once more; the `enqueued` mode seals its /// backlog and uploads it, a failing store included. Stop is answered by /// [`finish_stop`](Self::finish_stop) once nothing is in flight. fn handle_stop(&mut self, response: oneshot::Sender>) { @@ -2183,10 +2361,12 @@ impl Actor { .iter() .position(|batch| matches!(batch.state, CreateState::Pending)) { - for batch in self.sealed.drain(index..) { + for mut batch in self.sealed.split_off(index) { + self.keep_releases(&mut batch); batch.fail(|| ObjectStoreWalStoppedSnafu.build()); } } + self.seal(Vec::new()); } AckMode::Enqueued => { self.flush_open_batch(); @@ -2486,7 +2666,19 @@ fn record_skipped_segment( .or_insert(hole); } -fn encode_batch(header: Header, entries: Vec) -> Result { +/// Encodes the entries of a batch, and its release records as the single +/// entry of [`RELEASE_REGION_ID`], whose id is at position 1 of the object; +/// segments are ordered by region id, so it is the last segment. +fn encode_batch( + header: Header, + entries: Vec, + releases: &BTreeMap, +) -> Result { + let release_record = (!releases.is_empty()).then(|| Record { + region_id: RELEASE_REGION_ID, + entry_id: entry_id(header.object_seq, 1), + payload: encode_releases(releases), + }); let records = entries .into_iter() .map(|entry| Record { @@ -2494,6 +2686,7 @@ fn encode_batch(header: Header, entries: Vec) -> Result { entry_id: entry.entry_id(), payload: Bytes::from(entry.into_bytes()), }) + .chain(release_record) .collect::>(); encode_object(header, &records) } @@ -2513,6 +2706,9 @@ struct Recovered { reclaim_boundary: u64, /// Present objects that are not indexed, in sequence order. unindexed: Vec, + /// The entry id each region is released through, from the release + /// records of the indexed objects, see [`read_releases`]. + released_entry_ids: HashMap, /// Number of objects whose header and footer were read. recovered_objects: u64, } @@ -2535,14 +2731,71 @@ struct FetchedObject { /// Rebuilds the catalog from object headers and footers, /// so recovery costs a few small reads per object however large the objects -/// are. Segments are not read; a segment checksum is verified by the read that -/// decodes it. Footers are fetched for up to [`RECOVERY_CONCURRENCY`] objects +/// are. Segments are not read, except those of release records, see +/// [`read_releases`]; a segment checksum is verified by the read that decodes +/// it. Footers are fetched for up to [`RECOVERY_CONCURRENCY`] objects /// at a time and indexed in sequence order, so the catalog checks the entry /// ranges of every object against its predecessors like a sequential replay. /// Only the objects on the chain [`select_chain`] picks are indexed. async fn recover(io: &dyn WalObjectIo) -> Result { let objects = io.list().await?; - finish_recovery(fetch_footers(io, objects, RECOVERY_CONCURRENCY).await?) + let mut recovered = finish_recovery(fetch_footers(io, objects, RECOVERY_CONCURRENCY).await?)?; + recovered.released_entry_ids = read_releases(io, &recovered.catalog).await?; + Ok(recovered) +} + +/// Reads the release records of the indexed objects, which seed the +/// watermarks of their regions: one range read of the segment of +/// [`RELEASE_REGION_ID`] per object that holds one. A segment that cannot be +/// fetched or does not decode fails recovery, whatever `on_corrupted_segment` +/// says: skipping it would lose the watermarks it seeds. +/// +/// A release record needs no bookkeeping of its own. It is sealed once every +/// entry it covers was admitted, so it lies in the object holding the last +/// segment it covers or in a later one. Collection deletes, in any order, +/// only objects below the durable reclaim boundary, and recovery indexes only +/// objects at or above the boundary of the tip, which a start object carries +/// forward. So recovery indexes a release record whenever it indexes a +/// segment the record covers, and once the record is below the boundary so is +/// every segment it covers, which recovery never indexes again, whether +/// collection deleted it yet or not. A release in full also covers entries a +/// region would be assigned afterwards, so it is made only for a region that +/// writes nothing more. +async fn read_releases( + io: &dyn WalObjectIo, + catalog: &ObjectCatalog, +) -> Result> { + let segments = catalog + .objects_in_order() + .filter_map(|(object_seq, footer)| { + footer + .iter() + .find(|entry| entry.region_id == RELEASE_REGION_ID) + .map(|entry| (object_seq, entry.clone())) + }) + .collect::>(); + let releases = futures::stream::iter(segments) + .map(|(object_seq, entry)| async move { + let bytes = io + .get_range(object_seq, entry.segment_offset, entry.segment_len) + .await?; + decode_segment(&bytes, &entry) + .and_then(|records| decode_releases(&records[0].payload)) + .with_context(|_| InvalidWalObjectSnafu { + path: io.object_path(object_seq), + }) + }) + .buffered(RECOVERY_CONCURRENCY) + .try_collect::>() + .await?; + let mut released_entry_ids = HashMap::new(); + for release in releases.into_iter().flatten() { + let entry_id = released_entry_ids + .entry(release.region_id) + .or_insert(release.entry_id); + *entry_id = (*entry_id).max(release.entry_id); + } + Ok(released_entry_ids) } /// Indexes the chain of the fetched objects. Objects off the chain, orphans @@ -2627,6 +2880,7 @@ fn finish_recovery(fetched: FetchedObjects) -> Result { max_epoch, reclaim_boundary, unindexed, + released_entry_ids: HashMap::new(), recovered_objects, }) } @@ -3417,6 +3671,7 @@ mod tests { assert_eq!(expected.max_epoch, actual.max_epoch); assert_eq!(expected.reclaim_boundary, actual.reclaim_boundary); assert_eq!(expected.unindexed, actual.unindexed); + assert_eq!(expected.released_entry_ids, actual.released_entry_ids); } fn catalog_contents(catalog: &ObjectCatalog) -> Vec<(u64, Vec)> { @@ -3600,8 +3855,8 @@ mod tests { let last = bytes.len() - 1; bytes[last] ^= 1; }), - ("unsupported format version 2", |bytes| { - bytes[8..10].copy_from_slice(&2u16.to_be_bytes()); + ("unsupported format version 3", |bytes| { + bytes[8..10].copy_from_slice(&3u16.to_be_bytes()); }), ("footer checksum mismatch", |bytes| { let (trailer, _) = footer_of(bytes); @@ -4320,19 +4575,10 @@ mod tests { assert_eq!(20, latest(&store, region(1))); store.obsolete(&p, region(1), 10).await.unwrap(); assert_eq!(remaining, read_entries(&store, region(1), 0).await); - store.obsolete_all(&p, region(1)).await.unwrap(); - store.obsolete(&p, region(1), 0).await.unwrap(); - assert_eq!( - Vec::::new(), - read_entries(&store, region(1), 0).await - ); - assert_eq!(20, latest(&store, region(1))); - assert_eq!( - expected_entries(region(2), &[(10, "e10")]), - read_entries(&store, region(2), 0).await - ); store.stop().await.unwrap(); - let reopened = open(object_store, &eager()).await; + + // A watermark is kept in memory only. + let reopened = open(object_store.clone(), &eager()).await; assert_eq!(all, read_entries(&reopened, region(1), 0).await); reopened.obsolete(&p, region(1), 20).await.unwrap(); assert_eq!( @@ -4340,6 +4586,31 @@ mod tests { read_entries(&reopened, region(1), 0).await ); assert_eq!(20, latest(&reopened, region(1))); + + // A release in full survives the restart, and a later watermark does + // not move it down. + reopened.obsolete_all(&p, region(1)).await.unwrap(); + reopened.obsolete(&p, region(1), 0).await.unwrap(); + assert_eq!( + Vec::::new(), + read_entries(&reopened, region(1), 0).await + ); + assert_eq!(20, latest(&reopened, region(1))); + reopened.stop().await.unwrap(); + let reopened = open(object_store, &eager()).await; + assert_eq!( + Vec::::new(), + read_entries(&reopened, region(1), 0).await + ); + reopened.obsolete(&p, region(1), 0).await.unwrap(); + assert_eq!( + Vec::::new(), + read_entries(&reopened, region(1), 0).await + ); + assert_eq!( + expected_entries(region(2), &[(10, "e10")]), + read_entries(&reopened, region(2), 0).await + ); reopened.stop().await.unwrap(); } @@ -7477,22 +7748,25 @@ mod tests { assert_eq!(2, reclaim_boundary_of(io.as_ref(), 3).await); assert_eq!(vec![2, 3], collected_object_seqs(&store).await); - // Region one dropped and region two flushed through object 2: object - // 4 carries the sequence after the last indexed object. + // Region two flushed through object 2 and region one released: the + // object of the release record, 4, carries the sequence after the + // last indexed object, and so does the next one. store.obsolete(&provider(two), two, id(2, 1)).await.unwrap(); store.obsolete_all(&provider(one), one).await.unwrap(); - append(&store, two, "b2").await.unwrap(); assert_eq!(4, reclaim_boundary_of(io.as_ref(), 4).await); assert_eq!(vec![4], collected_object_seqs(&store).await); + append(&store, two, "b2").await.unwrap(); + assert_eq!(5, reclaim_boundary_of(io.as_ref(), 5).await); + assert_eq!(vec![5], collected_object_seqs(&store).await); // A region that appears later does not move the boundary back. append(&store, three, "c1").await.unwrap(); - assert_eq!(4, reclaim_boundary_of(io.as_ref(), 5).await); - assert_eq!(vec![4, 5], collected_object_seqs(&store).await); + assert_eq!(5, reclaim_boundary_of(io.as_ref(), 6).await); + assert_eq!(vec![5, 6], collected_object_seqs(&store).await); // Other stores of the process count their deletes as well. - assert!(METRIC_OBJECT_STORE_WAL_DELETED_OBJECTS_TOTAL.get() >= deleted + 4); + assert!(METRIC_OBJECT_STORE_WAL_DELETED_OBJECTS_TOTAL.get() >= deleted + 5); assert_eq!( - expected_entries(two, &[(id(4, 1), "b2")]), + expected_entries(two, &[(id(5, 1), "b2")]), read_entries(&store, two, 0).await ); store.stop().await.unwrap(); @@ -7608,15 +7882,17 @@ mod tests { put_foreign(&object_store, 7, 100).await; append(&store, region_id, "a7").await.unwrap_err(); } else { - // The first poll queues the stop command; the actor answers - // the command of `obsolete_all` queued after it only once it + // The first poll queues the stop command; the actor refuses + // the release of `obsolete_all` queued after it only once it // handled the stop. let mut stopping = Box::pin(store.stop()); assert!((&mut stopping).now_or_never().is_none()); - store - .obsolete_all(&provider(region_id), region_id) - .await - .unwrap(); + assert_stopped( + &store + .obsolete_all(&provider(region_id), region_id) + .await + .unwrap_err(), + ); assert!((&mut stopping).now_or_never().is_none()); stop = Some(stopping); } @@ -7967,6 +8243,258 @@ mod tests { store.stop().await.unwrap(); } + fn spawn_release( + store: &Arc, + region_id: RegionId, + ) -> tokio::task::JoinHandle> { + let store = store.clone(); + tokio::spawn(async move { store.release_region(&provider(region_id), region_id).await }) + } + + fn release(region_id: RegionId, entry_id: EntryId) -> Release { + Release { + region_id, + entry_id, + } + } + + /// The release records of the object `object_seq`. + async fn releases_of(io: &dyn WalObjectIo, object_seq: u64) -> Vec { + let decoded = decode_object(&io.get(object_seq).await.unwrap()).unwrap(); + let record = decoded + .records + .iter() + .find(|record| record.region_id == RELEASE_REGION_ID) + .unwrap(); + assert_eq!(entry_id(object_seq, 1), record.entry_id); + decode_releases(&record.payload).unwrap() + } + + fn watermark(store: &ObjectStoreLogStore, region_id: RegionId) -> Option { + store + .obsolete_entry_ids + .lock() + .unwrap() + .get(®ion_id) + .copied() + } + + #[tokio::test] + async fn test_store_release_writes_a_record_of_the_largest_issued_id() { + let _serialized = SERIALIZED.lock().await; + let store = open(memory_store(), &eager()).await; + let (one, two) = (region(1), region(2)); + store + .append_batch(vec![entry(&store, one, "a1"), entry(&store, two, "b1")]) + .await + .unwrap(); + + // The release seals object 2 without any entry; its record releases + // region one through the largest id handed out for it, and it carries + // the boundary region two holds at object 1. + store.release_region(&provider(one), one).await.unwrap(); + assert_eq!( + vec![release(one, id(1, 1))], + releases_of(store.io.as_ref(), 2).await + ); + assert_eq!(1, reclaim_boundary_of(store.io.as_ref(), 2).await); + assert_eq!(vec![1, 2], collected_object_seqs(&store).await); + assert!(read_entries(&store, one, 0).await.is_empty()); + // The record is the segment of a reserved region, which names no + // namespace. + assert_eq!( + vec![provider(one), provider(two)], + store.list_namespaces().await.unwrap() + ); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_recovery_seeds_watermarks_from_release_records() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + let (one, two) = (region(1), region(2)); + let store = open(object_store.clone(), &eager()).await; + // Object 1 holds entries of both regions and object 2 of region one. + // The release of region one is the last object written, and region + // two holds the boundary at object 1. + store + .append_batch(vec![entry(&store, one, "a1"), entry(&store, two, "b1")]) + .await + .unwrap(); + append(&store, one, "a2").await.unwrap(); + store.release_region(&provider(one), one).await.unwrap(); + assert_eq!(1, reclaim_boundary_of(store.io.as_ref(), 3).await); + store.stop().await.unwrap(); + + // Recovery seeds the watermark of region one, which nothing opens + // again, and a later watermark never moves it down. + let store = open(object_store, &eager()).await; + assert_eq!(vec![1, 2, 3, 4], collected_object_seqs(&store).await); + assert_eq!(Some(id(2, 1)), watermark(&store, one)); + assert!(read_entries(&store, one, 0).await.is_empty()); + store.obsolete(&provider(one), one, id(1, 1)).await.unwrap(); + assert_eq!(Some(id(2, 1)), watermark(&store, one)); + + // Once region two is flushed through object 1, no object holds a + // live segment: object 5 carries the boundary past all of them. + store.obsolete(&provider(two), two, id(1, 1)).await.unwrap(); + append(&store, two, "b2").await.unwrap(); + assert_eq!(5, reclaim_boundary_of(store.io.as_ref(), 5).await); + assert_eq!(vec![5], collected_object_seqs(&store).await); + assert_eq!( + expected_entries(two, &[(id(5, 1), "b2")]), + read_entries(&store, two, 0).await + ); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_carries_the_record_of_a_failed_release() { + let _serialized = SERIALIZED.lock().await; + let region_id = region(1); + + // The release fails with its create, and a tick seals the record + // again although nothing is appended. + let (io, _) = RecordingIo::over(memory_store()); + let store = open_over(io.clone(), &config(MIN_FLUSH_INTERVAL, u64::MAX)).await; + io.fail_next_put_persistently(); + store + .release_region(&provider(region_id), region_id) + .await + .unwrap_err(); + wait_until(|| store.catalog.read().unwrap().contains_object(2)).await; + assert_eq!( + vec![release(region_id, 0)], + releases_of(io.as_ref(), 2).await + ); + store.stop().await.unwrap(); + + // A `durable` stop creates the record once more, and reports it lost + // when that create fails or conflicts too. + for last_create in ["succeeds", "fails", "conflicts"] { + let object_store = memory_store(); + let store = open(object_store.clone(), &eager()).await; + append(&store, region_id, "a1").await.unwrap(); + store.fail_creates(); + store + .release_region(&provider(region_id), region_id) + .await + .unwrap_err(); + match last_create { + "succeeds" => store.allow_creates(), + "conflicts" => { + store.allow_creates(); + put_foreign(&object_store, 3, 2).await; + } + _ => {} + } + let stopped = store.stop().await; + let store = open(object_store, &eager()).await; + if last_create == "succeeds" { + stopped.unwrap(); + assert_eq!(Some(id(1, 1)), watermark(&store, region_id)); + } else { + let error = stopped.unwrap_err(); + let conflict = matches!(unwrap_shared(&error), Error::WalObjectConflict { .. }); + assert_eq!(last_create == "conflicts", conflict, "{error:?}"); + assert_eq!(None, watermark(&store, region_id)); + } + store.stop().await.unwrap(); + } + } + + #[tokio::test] + async fn test_store_releases_a_region_without_segments() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + let store = open(object_store.clone(), &eager()).await; + let io = store.io.clone(); + let (one, two, three) = (region(1), region(2), region(3)); + append(&store, two, "b1").await.unwrap(); + + // A region the store holds no entry of is released through zero and + // a release in full reaches the largest id, each in an object of its + // own; neither moves the boundary region two holds. + store.release_region(&provider(one), one).await.unwrap(); + store.obsolete_all(&provider(three), three).await.unwrap(); + assert_eq!(vec![release(one, 0)], releases_of(io.as_ref(), 2).await); + assert_eq!( + vec![release(three, EntryId::MAX)], + releases_of(io.as_ref(), 3).await + ); + assert_eq!(1, reclaim_boundary_of(io.as_ref(), 3).await); + assert_eq!(vec![1, 2, 3], collected_object_seqs(&store).await); + + // Region one written again stays readable across a restart. + append(&store, one, "a1").await.unwrap(); + store.stop().await.unwrap(); + let store = open(object_store, &eager()).await; + assert_eq!(Some(0), watermark(&store, one)); + assert_eq!(Some(EntryId::MAX), watermark(&store, three)); + assert_eq!( + expected_entries(one, &[(id(4, 1), "a1")]), + read_entries(&store, one, 0).await + ); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_release_record_lives_as_long_as_the_segments_it_covers() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + let (one, two) = (region(1), region(2)); + let store = open(object_store.clone(), &eager()).await; + // Object 1 holds entries of both regions, object 2 the release of + // region one and object 3 an entry of region two, which holds the + // boundary at object 1. + store + .append_batch(vec![entry(&store, one, "a1"), entry(&store, two, "b1")]) + .await + .unwrap(); + store.release_region(&provider(one), one).await.unwrap(); + append(&store, two, "b2").await.unwrap(); + assert_eq!(vec![1, 2, 3], collected_object_seqs(&store).await); + store.stop().await.unwrap(); + + // The record outlives a restart together with the segment it covers. + let (io, mut parked) = ParkedIo::parking_creates(object_store.clone()); + let store = open_over(io.clone(), &eager()).await; + assert_eq!(Some(id(1, 1)), watermark(&store, one)); + assert!(read_entries(&store, one, 0).await.is_empty()); + + // Region two flushed through object 1 lets object 5 carry boundary + // 3, below which the record and the segment go in any order: the + // record is deleted, the delete of the segment fails. + io.deletes_parked.store(true, Ordering::SeqCst); + store.obsolete(&provider(two), two, id(1, 1)).await.unwrap(); + append(&store, two, "b3").await.unwrap(); + assert_eq!(3, reclaim_boundary_of(io.as_ref(), 5).await); + let mut deletes = parked_operations(&mut parked, 2).await; + deletes.remove(&2).unwrap().send(true).unwrap(); + deletes.remove(&1).unwrap().send(false).unwrap(); + timeout(WAIT, store.wait_for_garbage_collection()) + .await + .unwrap(); + assert_eq!(vec![1, 3, 4, 5], object_seqs(io.as_ref()).await); + store.stop().await.unwrap(); + + // Recovery indexes nothing below boundary 3, which the start object + // carries forward, so the segment the record covered is never live + // again, and the first collection deletes it. + let plain = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); + let recovered = recover(&plain).await.unwrap(); + assert_eq!(3, recovered.reclaim_boundary); + assert_eq!(vec![1], recovered.unindexed); + assert!(recovered.released_entry_ids.is_empty()); + let store = open(object_store, &eager()).await; + assert_eq!(3, reclaim_boundary_of(&plain, 6).await); + assert_eq!(None, watermark(&store, one)); + store.obsolete(&provider(two), two, id(3, 1)).await.unwrap(); + assert_eq!(vec![3, 4, 5, 6], collected_object_seqs(&store).await); + store.stop().await.unwrap(); + } + /// Object access whose next create stores the object but reports a /// transient failure, as a create whose response is lost. struct LostResponseIo { @@ -8430,6 +8958,7 @@ mod tests { reads: METRIC_OBJECT_STORE_WAL_READ_SECONDS.get_sample_count(), deleted_objects: METRIC_OBJECT_STORE_WAL_DELETED_OBJECTS_TOTAL.get(), failed_deletes: METRIC_OBJECT_STORE_WAL_FAILED_DELETES_TOTAL.get(), + release_records: METRIC_OBJECT_STORE_WAL_RELEASE_RECORDS_TOTAL.get(), /// No read of the matrix meets a corrupted segment. skipped_segments: METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.get(), stalled_appends: METRIC_OBJECT_STORE_WAL_STALLED_APPENDS_TOTAL.get(), @@ -8488,6 +9017,15 @@ mod tests { object_bytes_of(&[(region(1), "a1")]) } + /// The length of an object that holds one release record alone. + fn release_object() -> u64 { + let releases = BTreeMap::from([(region(1), 0)]); + encode_batch(chain_header(1, 1, None), Vec::new(), &releases) + .unwrap() + .bytes + .len() as u64 + } + /// The sum and the count of a timing histogram before a forced wait. fn timer_before(histogram: &Histogram) -> (f64, u64) { (histogram.get_sample_sum(), histogram.get_sample_count()) @@ -8680,6 +9218,36 @@ mod tests { row: Some(Row::LastSequence), run: || Box::pin(create_takes_the_last_object_sequence()), }, + Case { + fault: "a release whose create succeeds", + row: Some(Row::Created), + run: || Box::pin(release_create_succeeds()), + }, + Case { + fault: "a release whose create fails transiently, `durable` mode", + row: Some(Row::TransientDurable), + run: || Box::pin(release_create_fails_transiently_durable()), + }, + Case { + fault: "a release whose create fails transiently, `enqueued` mode", + row: Some(Row::TransientEnqueued), + run: || Box::pin(release_create_fails_transiently_enqueued()), + }, + Case { + fault: "a release whose create the object store refuses with an error that requires intervention, `enqueued` mode", + row: Some(Row::InterventionEnqueued), + run: || Box::pin(release_create_refused_enqueued()), + }, + Case { + fault: "a release while the object store refuses the backlog, `enqueued` mode", + row: None, + run: || Box::pin(release_while_the_backlog_is_refused()), + }, + Case { + fault: "a release whose create finds an object of the same or a later epoch", + row: Some(Row::Conflict), + run: || Box::pin(release_create_finds_another_writer()), + }, Case { fault: "a start object that meets an object of an earlier epoch", row: Some(Row::StartMovesPast), @@ -8877,6 +9445,310 @@ mod tests { store.stop().await.unwrap(); } + /// A release seals an object of its own and returns once it is + /// indexed, in both modes. + async fn release_create_succeeds() { + for config in [eager(), enqueued(eager())] { + let (store, io, mut parked) = open_parking_creates(memory_store(), &config).await; + let before = Counters::sample(); + let region_id = region(1); + let released = spawn_release(&store, region_id); + let (object_seq, create) = next_parked_operation(&mut parked).await; + assert_eq!(1, object_seq); + assert!(!released.is_finished()); + assert_indexed(io.as_ref(), &[0]).await; + create.send(true).unwrap(); + timeout(WAIT, released).await.unwrap().unwrap().unwrap(); + // Object 1 carries the boundary past the start object. + assert_eq!(vec![1], collected_object_seqs(&store).await); + assert_eq!( + vec![release(region_id, 0)], + releases_of(io.as_ref(), 1).await + ); + assert!(read_entries(&store, region_id, 1).await.is_empty()); + assert_indexed(io.as_ref(), &[1]).await; + assert_eq!( + Counters { + created_objects: 1, + seal_to_durable: 1, + sealed_objects: 1, + object_bytes: release_object(), + reads: 1, + deleted_objects: 1, + release_records: 1, + ..Counters::default() + }, + Counters::since(before) + ); + store.stop().await.unwrap(); + } + } + + /// The release fails with the create, its sequence is not reused, and + /// the release of the region again lands in the next object. + async fn release_create_fails_transiently_durable() { + let (store, io, mut parked) = open_parking_creates(memory_store(), &eager()).await; + let before = Counters::sample(); + let region_id = region(1); + let released = spawn_release(&store, region_id); + let (object_seq, create) = next_parked_operation(&mut parked).await; + assert_eq!(1, object_seq); + create.send(false).unwrap(); + let error = timeout(WAIT, released).await.unwrap().unwrap().unwrap_err(); + assert!( + matches!(unwrap_shared(&error), Error::WalObjectStore { .. }), + "unexpected error: {error:?}" + ); + assert!(terminal(&store.terminal_error).is_none()); + assert_eq!(vec![0], object_seqs(io.as_ref()).await); + assert_indexed(io.as_ref(), &[0]).await; + assert_eq!( + Counters { + create_failures: 1, + sealed_objects: 1, + object_bytes: release_object(), + ..Counters::default() + }, + Counters::since(before) + ); + + let released = spawn_release(&store, region_id); + let (object_seq, create) = next_parked_operation(&mut parked).await; + assert_eq!(2, object_seq); + create.send(true).unwrap(); + timeout(WAIT, released).await.unwrap().unwrap().unwrap(); + assert_eq!(vec![2], collected_object_seqs(&store).await); + assert_indexed(io.as_ref(), &[2]).await; + assert_eq!( + Counters { + created_objects: 1, + seal_to_durable: 1, + create_failures: 1, + sealed_objects: 2, + object_bytes: 2 * release_object(), + deleted_objects: 1, + release_records: 1, + ..Counters::default() + }, + Counters::since(before) + ); + store.stop().await.unwrap(); + } + + /// The store repeats the create, and the release waits for the repeat. + async fn release_create_fails_transiently_enqueued() { + let (store, io, mut parked) = + open_parking_creates(memory_store(), &enqueued(eager())).await; + let before = Counters::sample(); + let region_id = region(1); + let released = spawn_release(&store, region_id); + let (object_seq, create) = next_parked_operation(&mut parked).await; + assert_eq!(1, object_seq); + create.send(false).unwrap(); + let (object_seq, repeat) = next_parked_operation(&mut parked).await; + assert_eq!(1, object_seq); + assert!(!released.is_finished()); + assert_eq!(0, failing_stores()); + assert_eq!(vec![0], object_seqs(io.as_ref()).await); + assert_indexed(io.as_ref(), &[0]).await; + assert_eq!( + Counters { + create_failures: 1, + sealed_objects: 1, + object_bytes: release_object(), + ..Counters::default() + }, + Counters::since(before) + ); + + repeat.send(true).unwrap(); + timeout(WAIT, released).await.unwrap().unwrap().unwrap(); + assert_eq!(vec![1], collected_object_seqs(&store).await); + assert_indexed(io.as_ref(), &[1]).await; + assert_eq!( + Counters { + created_objects: 1, + seal_to_durable: 1, + create_failures: 1, + sealed_objects: 1, + object_bytes: release_object(), + deleted_objects: 1, + release_records: 1, + ..Counters::default() + }, + Counters::since(before) + ); + store.stop().await.unwrap(); + } + + /// The object store denies access to the object of a release: the + /// store becomes failing and the release fails at once with the + /// retryable error; its record stays in the backlog and becomes + /// durable with the next create that succeeds. + async fn release_create_refused_enqueued() { + let (store, io, mut parked) = + open_parking_creates(memory_store(), &enqueued(eager())).await; + let before = Counters::sample(); + let region_id = region(1); + let released = spawn_release(&store, region_id); + let (object_seq, create) = next_parked_operation(&mut parked).await; + assert_eq!(1, object_seq); + io.creates_refused.store(true, Ordering::SeqCst); + create.send(false).unwrap(); + let error = timeout(WAIT, released).await.unwrap().unwrap().unwrap_err(); + assert_rejected(&error, ErrorKind::PermissionDenied); + assert_eq!(1, failing_stores()); + assert!(terminal(&store.terminal_error).is_none()); + assert_indexed(io.as_ref(), &[0]).await; + assert_eq!( + Counters { + create_failures: 1, + sealed_objects: 1, + object_bytes: release_object(), + rejected_appends: 1, + ..Counters::default() + }, + Counters::since(before) + ); + + let (object_seq, create) = next_parked_operation(&mut parked).await; + assert_eq!(1, object_seq); + create.send(true).unwrap(); + wait_until(|| failing_stores() == 0).await; + assert_eq!(vec![1], collected_object_seqs(&store).await); + assert_eq!( + vec![release(region_id, 0)], + releases_of(io.as_ref(), 1).await + ); + assert_indexed(io.as_ref(), &[1]).await; + assert_eq!( + Counters { + created_objects: 1, + seal_to_durable: 1, + create_failures: 1, + sealed_objects: 1, + object_bytes: release_object(), + deleted_objects: 1, + release_records: 1, + rejected_appends: 1, + ..Counters::default() + }, + Counters::since(before) + ); + store.stop().await.unwrap(); + } + + /// Releases that arrive while the store is failing fail at once with + /// the retryable error, and their record is kept coalesced: the + /// ticks that pass seal nothing behind the backlog. A stop seals the + /// record with the backlog, and recovery reads it back. + async fn release_while_the_backlog_is_refused() { + let object_store = memory_store(); + let (store, io, mut parked) = open_parking_creates( + object_store.clone(), + &enqueued(config(MIN_FLUSH_INTERVAL, 1)), + ) + .await; + let before = Counters::sample(); + let region_id = region(1); + append(&store, region_id, "a1").await.unwrap(); + io.creates_refused.store(true, Ordering::SeqCst); + next_parked_operation(&mut parked) + .await + .1 + .send(false) + .unwrap(); + wait_until(|| failing_stores() == 1).await; + for _ in 0..3 { + let error = store + .release_region(&provider(region_id), region_id) + .await + .unwrap_err(); + assert_rejected(&error, ErrorKind::PermissionDenied); + } + // The next attempt of the backlog comes after several ticks. + let (object_seq, probe) = next_parked_operation(&mut parked).await; + assert_eq!(1, object_seq); + assert!(parked.try_recv().is_err()); + probe.send(false).unwrap(); + wait_until(|| Counters::since(before).create_failures == 2).await; + assert_eq!(1, Counters::since(before).sealed_objects); + + let stop = { + let store = store.clone(); + tokio::spawn(async move { store.stop().await }) + }; + for create in parked_operations(&mut parked, 2).await.into_values() { + create.send(true).unwrap(); + } + timeout(WAIT, stop).await.unwrap().unwrap().unwrap(); + assert_eq!(0, failing_stores()); + assert_eq!( + vec![release(region_id, id(1, 1))], + releases_of(io.as_ref(), 2).await + ); + assert_eq!( + Counters { + created_objects: 2, + seal_to_durable: 2, + create_failures: 2, + sealed_objects: 2, + object_bytes: one_entry_object() + release_object(), + object_entries: 1, + release_records: 1, + rejected_appends: 3, + ..Counters::default() + }, + Counters::since(before) + ); + let store = open(object_store, &eager()).await; + assert_eq!(Some(id(1, 1)), watermark(&store, region_id)); + store.stop().await.unwrap(); + } + + /// Another writer holds the sequence: the release fails with the + /// conflict, which poisons the store, in both modes. + async fn release_create_finds_another_writer() { + for config in [eager(), enqueued(eager())] { + let object_store = memory_store(); + let (store, io, mut parked) = + open_parking_creates(object_store.clone(), &config).await; + put_foreign(&object_store, 1, 2).await; + let before = Counters::sample(); + let region_id = region(1); + let released = spawn_release(&store, region_id); + let (_, create) = next_parked_operation(&mut parked).await; + create.send(true).unwrap(); + let error = timeout(WAIT, released).await.unwrap().unwrap().unwrap_err(); + let path = io.object_path(1); + assert!( + matches!(unwrap_shared(&error), Error::WalObjectConflict { path: actual, .. } if *actual == path), + "unexpected error: {error:?}" + ); + assert_read_is_terminal(&store, region_id).await; + assert_indexed(io.as_ref(), &[0]).await; + assert_eq!( + Counters { + create_conflicts: 1, + poisoned: 1, + sealed_objects: 1, + object_bytes: release_object(), + ..Counters::default() + }, + Counters::since(before) + ); + let stopped = store.stop().await; + if config.ack_mode == AckMode::Durable { + stopped.unwrap(); + } else { + assert!(matches!( + unwrap_shared(&stopped.unwrap_err()), + Error::WalObjectConflict { .. } + )); + } + } + } + /// Appends over `io` after `arm` made the create meet the very object /// it writes, and asserts it was accepted as an identical retry. async fn assert_identical_retry(io: Arc, arm: impl FnOnce()) { @@ -10254,15 +11126,17 @@ mod tests { } /// Round-trips a command through the actor once `stop` began, when - /// the seal command is refused: the collection that `obsolete_all` of - /// a region without entries requests is answered by the actor, and + /// the seal command is refused: the actor refuses the release that + /// `obsolete_all` of a region without entries requests as well, which /// hides nothing of the regions a case reads. async fn round_trip_stopping_actor(store: &ObjectStoreLogStore) { let region_id = region(9); - store - .obsolete_all(&provider(region_id), region_id) - .await - .unwrap(); + assert_stopped( + &store + .obsolete_all(&provider(region_id), region_id) + .await + .unwrap_err(), + ); } /// Four creates are in flight and a fifth batch waits for a slot when diff --git a/src/metric-engine/Cargo.toml b/src/metric-engine/Cargo.toml index 92000549832..0cc39dc5add 100644 --- a/src/metric-engine/Cargo.toml +++ b/src/metric-engine/Cargo.toml @@ -49,6 +49,7 @@ common-function = { workspace = true } common-meta = { workspace = true, features = ["testing"] } common-wal = { workspace = true } criterion = { workspace = true, features = ["async", "async_tokio"] } +log-store = { workspace = true, features = ["testing"] } mito2 = { workspace = true, features = ["test"] } mur3 = "0.1" partition = { workspace = true } diff --git a/src/metric-engine/src/engine/drop.rs b/src/metric-engine/src/engine/drop.rs index 0b3d876fd84..7e98effdeee 100644 --- a/src/metric-engine/src/engine/drop.rs +++ b/src/metric-engine/src/engine/drop.rs @@ -171,3 +171,159 @@ impl MetricEngineInner { Ok(0) } } + +#[cfg(test)] +mod tests { + use api::v1::Rows; + use common_wal::config::object_store::{ObjectStoreWalConfig, STANDALONE_GENERATION}; + use common_wal::options::{ObjectStoreWalOptions, WAL_OPTIONS_KEY, WalOptions}; + use log_store::ObjectStoreLogStore; + use mito2::config::MitoConfig; + use mito2::test_util::TestEnv as MitoTestEnv; + use object_store::ObjectStore; + use object_store::services::Memory; + use store_api::logstore::LogStore; + use store_api::region_engine::RegionEngine; + use store_api::region_request::{RegionDropRequest, RegionPutRequest, RegionRequest}; + use store_api::storage::RegionId; + + use crate::test_util::{TestEnv, build_rows, row_schema_with_tags}; + use crate::utils::{to_data_region_id, to_metadata_region_id}; + + fn drop_request() -> RegionRequest { + RegionRequest::Drop(RegionDropRequest { + fast_path: false, + force: false, + partial_drop: false, + }) + } + + /// How the WAL releases of the dropped regions fail. + #[derive(Clone, Copy, PartialEq)] + enum ReleaseFailure { + /// The release of the data region fails transiently. + Data, + /// The releases of the data and the metadata region fail transiently. + Both, + /// The release of the data region meets a conflicting object, which + /// poisons the WAL, so the cleanup of the metadata region fails too. + Conflict, + } + + #[tokio::test] + async fn test_drop_physical_region_whose_wal_release_fails() { + for failure in [ + ReleaseFailure::Data, + ReleaseFailure::Both, + ReleaseFailure::Conflict, + ] { + let object_store = ObjectStore::new(Memory::default()).unwrap(); + let config = ObjectStoreWalConfig::default(); + let open_wal = || { + ObjectStoreLogStore::try_new( + object_store.clone(), + &config, + 0, + STANDALONE_GENERATION, + ) + }; + let wal = open_wal().await.unwrap(); + let mut mito_env = MitoTestEnv::with_prefix("metric-drop-failed-release").await; + let mito = mito_env + .create_engine_with_log_store(MitoConfig::default(), wal.clone()) + .await; + let env = TestEnv::with_mito_engine(mito_env, mito); + let region_id = RegionId::new(1, 1); + let wal_options = WalOptions::ObjectStore(ObjectStoreWalOptions::new( + "wal/datanodes/0/epochs/0".to_string(), + )); + env.create_physical_region( + region_id, + "test_metric_table", + vec![( + WAL_OPTIONS_KEY.to_string(), + serde_json::to_string(&wal_options).unwrap(), + )], + ) + .await; + // Both mito regions write to the WAL: the data region the rows of a + // logical region, the metadata region its creation and its drop. + let logical_region_id = RegionId::new(1024, 1); + env.create_logical_region(region_id, logical_region_id) + .await; + let metric = env.metric(); + metric + .handle_request( + logical_region_id, + RegionRequest::Put(RegionPutRequest { + skip_wal: false, + rows: Rows { + schema: row_schema_with_tags(&["job"]), + rows: build_rows(1, 10), + }, + hint: None, + partition_expr_version: None, + }), + ) + .await + .unwrap(); + metric + .handle_request(logical_region_id, drop_request()) + .await + .unwrap(); + + // The drop still removes both regions and the physical region state. + match failure { + ReleaseFailure::Data => wal.fail_next_create_after_write(), + ReleaseFailure::Both => wal.fail_creates(), + ReleaseFailure::Conflict => { + // A copy of the last object, of the same epoch, takes the + // sequence of the next one. + let objects = "wal/datanodes/0/epochs/0/objects/"; + let last = object_store + .list(objects) + .await + .unwrap() + .iter() + .filter_map(|entry| entry.name().strip_suffix(".wal")?.parse::().ok()) + .max() + .unwrap(); + let path = |object_seq: u64| format!("{objects}{object_seq:020}.wal"); + let copy = object_store.read(&path(last)).await.unwrap(); + object_store.write(&path(last + 1), copy).await.unwrap(); + } + } + metric + .handle_request(region_id, drop_request()) + .await + .unwrap(); + assert!(!env.mito().is_region_exists(to_data_region_id(region_id))); + assert!( + !env.mito() + .is_region_exists(to_metadata_region_id(region_id)) + ); + assert!( + !metric + .inner + .state + .read() + .unwrap() + .exist_physical_region(to_data_region_id(region_id)) + ); + + // The stop writes the release records of failed releases. A store + // opened later finds both regions released and keeps no segment of + // either. A poisoned store writes nothing more. + if failure == ReleaseFailure::Conflict { + assert!(wal.list_namespaces().await.is_err()); + continue; + } + wal.allow_creates(); + wal.stop().await.unwrap(); + let wal = open_wal().await.unwrap(); + wal.wait_for_garbage_collection().await; + assert!(wal.list_namespaces().await.unwrap().is_empty()); + wal.stop().await.unwrap(); + } + } +} diff --git a/src/metric-engine/src/test_util.rs b/src/metric-engine/src/test_util.rs index cf64d968182..3c3762aeea7 100644 --- a/src/metric-engine/src/test_util.rs +++ b/src/metric-engine/src/test_util.rs @@ -92,6 +92,17 @@ impl TestEnv { } } + /// Returns a new env whose metric engine runs over `mito`, an engine of + /// `mito_env`. + pub fn with_mito_engine(mito_env: MitoTestEnv, mito: MitoEngine) -> Self { + let metric = MetricEngine::try_new(mito.clone(), EngineConfig::default()).unwrap(); + Self { + mito_env, + mito, + metric, + } + } + /// Returns a new env with specific `prefix` and `mito_env` for test. pub async fn with_mito_env(mut mito_env: MitoTestEnv) -> Self { let mito = mito_env.create_engine(MitoConfig::default()).await; diff --git a/src/mito2/src/engine/object_store_wal_recovery_test.rs b/src/mito2/src/engine/object_store_wal_recovery_test.rs index 54626b24432..9ecfc998697 100644 --- a/src/mito2/src/engine/object_store_wal_recovery_test.rs +++ b/src/mito2/src/engine/object_store_wal_recovery_test.rs @@ -1126,23 +1126,10 @@ async fn test_drop_cancels_a_flush_waiting_for_wal_durability() { wait_for(store.wait_for_durability_waits(1)).await; assert!(!flush.is_finished()); - // The drop cancels the flush instead of waiting behind the upload. - tokio::time::timeout( - WAIT, - engine.handle_request( - REGION_A, - RegionRequest::Drop(RegionDropRequest { - fast_path: false, - force: false, - partial_drop: false, - }), - ), - ) - .await - .unwrap() - .unwrap(); - assert!(!engine.is_region_exists(REGION_A)); - assert_eq!(vec![0], wal_object_seqs(&object_store).await); + // The drop cancels the flush instead of waiting behind the upload. Its + // release of the region then waits for its own object, which is created + // after the upload. + let dropped = spawn_drop(&engine, REGION_A); assert!( tokio::time::timeout(WAIT, flush) .await @@ -1150,8 +1137,17 @@ async fn test_drop_cancels_a_flush_waiting_for_wal_durability() { .unwrap() .is_err() ); + wait_for(store.wait_for_parked_creates(2)).await; + assert!(!dropped.is_finished()); + assert_eq!(vec![0], wal_object_seqs(&object_store).await); store.release_creates(); + tokio::time::timeout(WAIT, dropped) + .await + .unwrap() + .unwrap() + .unwrap(); + assert!(!engine.is_region_exists(REGION_A)); tokio::time::timeout(WAIT, seal) .await .unwrap() @@ -1159,6 +1155,106 @@ async fn test_drop_cancels_a_flush_waiting_for_wal_durability() { .unwrap(); } +fn spawn_drop( + engine: &MitoEngine, + region_id: RegionId, +) -> tokio::task::JoinHandle> { + let engine = engine.clone(); + tokio::spawn(async move { + engine + .handle_request( + region_id, + RegionRequest::Drop(RegionDropRequest { + fast_path: false, + force: false, + partial_drop: false, + }), + ) + .await + .map(|_| ()) + }) +} + +#[rstest] +#[case::graceful(false, false)] +#[case::abrupt(true, false)] +#[case::failed_release(false, true)] +#[tokio::test] +async fn test_restart_after_a_drop_collects_the_objects_of_the_dropped_region( + #[case] abrupt: bool, + #[case] release_fails: bool, +) { + let mut env = TestEnv::with_prefix("object-store-wal-drop-restart").await; + let object_store = memory_store(); + let store = open_store(&object_store, AckMode::Durable).await; + let engine = new_engine(&mut env, store.clone()).await; + let (_, schema_a) = create_region(&engine, REGION_A, &[]).await; + let (table_dir_b, schema_b) = create_region(&engine, REGION_B, &[]).await; + + // Object 1 holds entries of both regions and object 2 of region A. The + // drop releases region A in object 3, and region B, which is not + // flushed, holds the boundary at object 1. + let mut writer = SealedWriter::new(&store); + writer + .put_and_seal( + &engine, + vec![ + (REGION_A, rows(&schema_a, 0, 2)), + (REGION_B, rows(&schema_b, 0, 3)), + ], + ) + .await; + writer + .put_and_seal(&engine, vec![(REGION_A, rows(&schema_a, 2, 4))]) + .await; + if release_fails { + store.fail_creates(); + } + let dropped = tokio::time::timeout(WAIT, spawn_drop(&engine, REGION_A)) + .await + .unwrap() + .unwrap(); + // A drop whose release fails succeeds all the same, and the stop creates + // the record in object 4. + dropped.unwrap(); + assert!(!engine.is_region_exists(REGION_A)); + store.allow_creates(); + wait_for_collection(&store).await; + let released_in = if release_fails { 4 } else { 3 }; + drop(writer); + if abrupt { + crash(engine, store).await; + } else { + engine.stop().await.unwrap(); + drop(engine); + store.stop().await.unwrap(); + drop(store); + } + assert_eq!( + vec![1, 2, released_in], + wal_object_seqs(&object_store).await + ); + + // Nothing is written before the restart, and region A is never opened + // again. + let store = open_store(&object_store, AckMode::Durable).await; + let engine = new_engine(&mut env, store.clone()).await; + open_region(&engine, REGION_B, &table_dir_b, &[]) + .await + .unwrap(); + + // Once region B is flushed no object holds a live segment, and the next + // object carries the boundary past all of them. + flush_region(&engine, REGION_B, None).await; + let mut writer = SealedWriter::new(&store); + writer + .put_and_seal(&engine, vec![(REGION_B, rows(&schema_b, 3, 5))]) + .await; + wait_for_collection(&store).await; + assert_eq!(vec![released_in + 2], wal_object_seqs(&object_store).await); + assert_eq!(5, engine.get_region_statistic(REGION_B).unwrap().num_rows); +} + #[rstest] #[case(RegionTruncateRequest::All)] #[case(RegionTruncateRequest::Unflushed)] diff --git a/src/mito2/src/test_util.rs b/src/mito2/src/test_util.rs index d5104bb16b7..f265494058a 100644 --- a/src/mito2/src/test_util.rs +++ b/src/mito2/src/test_util.rs @@ -495,7 +495,7 @@ impl TestEnv { } /// Creates a new engine with specific config on `log_store` under this env. - pub(crate) async fn create_engine_with_log_store( + pub async fn create_engine_with_log_store( &mut self, config: MitoConfig, log_store: Arc, diff --git a/src/mito2/src/wal.rs b/src/mito2/src/wal.rs index 76b3e3a2cdf..9d00a76d87f 100644 --- a/src/mito2/src/wal.rs +++ b/src/mito2/src/wal.rs @@ -220,6 +220,18 @@ impl Wal { .context(DeleteWalSnafu { region_id }) } + /// Releases every WAL entry of a dropped region, see [LogStore::release_region]. + pub async fn release_region(&self, region_id: RegionId, provider: &Provider) -> Result<()> { + if let Provider::Noop = provider { + return Ok(()); + } + self.store + .release_region(provider, region_id) + .await + .map_err(BoxedError::new) + .context(DeleteWalSnafu { region_id }) + } + /// Marks all WAL entries of a region as obsolete and removes its dedicated namespace when /// supported by the backend. pub async fn obsolete_all(&self, region_id: RegionId, provider: &Provider) -> Result<()> { diff --git a/src/mito2/src/worker/handle_drop.rs b/src/mito2/src/worker/handle_drop.rs index a62e8f0068e..192106110eb 100644 --- a/src/mito2/src/worker/handle_drop.rs +++ b/src/mito2/src/worker/handle_drop.rs @@ -121,14 +121,30 @@ where // this cleanup and close/drop remove. self.cleanup_dropped_region_runtime_state(region_id).await; - // Delete region data in WAL. - self.wal + // Deletes region data in WAL, and releases it so that this outlives a restart, which + // never opens the region again. A failure of the object store WAL does not fail the + // drop: the region has left the map, a drop request for it is not sent again, and a + // caller that drops several regions on one store, such as the metric engine, would + // leave the others undone, for instance once the release of one poisoned the store. + // The store keeps a failed release and writes it with its next object unless it is + // stopped or poisoned. + if let Err(e) = self + .wal .obsolete( region_id, region.version_control.current().last_entry_id, ®ion.provider, ) - .await?; + .await + { + if region.provider.as_object_store_provider().is_none() { + return Err(e); + } + error!(e; "Failed to obsolete the WAL of dropped region {}", region_id); + } + if let Err(e) = self.wal.release_region(region_id, ®ion.provider).await { + error!(e; "Failed to release the WAL of dropped region {}", region_id); + } if let Some(store) = &self.series_index_store { crate::series_index::delete_catalogs(store, region_id).await; diff --git a/src/store-api/src/logstore.rs b/src/store-api/src/logstore.rs index dd71142def5..ef0fbe49f01 100644 --- a/src/store-api/src/logstore.rs +++ b/src/store-api/src/logstore.rs @@ -106,6 +106,22 @@ pub trait LogStore: Send + Sync + 'static + std::fmt::Debug { ) -> Result<(), Self::Error> { Ok(()) } + + /// Releases every entry of a region the caller dropped, after `obsolete` + /// marked them obsolete, so that the log store keeps none of them even + /// across a restart, which never opens the region again to mark its + /// entries obsolete. + /// + /// The default does nothing and keeps each log store's existing handling + /// of a dropped region. A log store that keeps obsolete marks only in + /// memory overrides this method and returns once the release is durable. + async fn release_region( + &self, + _provider: &Provider, + _region_id: RegionId, + ) -> Result<(), Self::Error> { + Ok(()) + } } /// The response of an `append` operation. diff --git a/tests/compatibility/README.md b/tests/compatibility/README.md index 404b6332b94..1daf42cc736 100644 --- a/tests/compatibility/README.md +++ b/tests/compatibility/README.md @@ -15,6 +15,12 @@ read-preference defaulting and unrelated metadata preservation: cargo nextest run -p session -E 'test(test_legacy_query_context_wire_defaults_to_leader)' ``` +## Object store WAL fixtures + +Frozen objects of the object store WAL are checked byte for byte by the format tests of +`log-store`. The SQL runner cannot cover the object store WAL: it owns `wal.provider`. See the +[fixture notes](wal/object_store/README.md). + ## Quick Start ```shell diff --git a/tests/compatibility/wal/object_store/README.md b/tests/compatibility/wal/object_store/README.md new file mode 100644 index 00000000000..411900197f4 --- /dev/null +++ b/tests/compatibility/wal/object_store/README.md @@ -0,0 +1,16 @@ +# Object store WAL object fixtures + +Frozen objects of the object store WAL, checked byte for byte by the format tests in +`src/log-store/src/object_store_wal/format.rs`, which never rewrite them: + +```shell +cargo nextest run -p log-store -E 'test(test_format_matches_release_record_fixture)' +``` + +## `release_record.wal` + +A version 1 object, sequence 3, holding entry `(3 << 20) | 1` of region 1 and, as the single +entry of the reserved region `u64::MAX` at the same position, the release records of region +`(2, 1)` through `(2 << 20) | 1` and of region `(3, 0)` in full. A reader that predates +release records decodes it as an ordinary version 1 object. 272 bytes, SHA-256 +`7e8de4d50142f3b48c97a24f20744ee9c0d56f843daf3a36d00ba9b465a5d4b5`. diff --git a/tests/compatibility/wal/object_store/release_record.wal b/tests/compatibility/wal/object_store/release_record.wal new file mode 100644 index 0000000000000000000000000000000000000000..0318e3630d15d860f40d67c50e3fb7113446155b GIT binary patch literal 272 zcmZ<{33v4Ick*IjWPku>D9wa}W|Us#t__p}sfS?(1_LmgDbesh6hLKRGAckm$V3o^ z*rWia!L}jkfyqPEKum?10TE+hFaw$a0#bi!yqb{o!qmZNnEC*Sat51R$@Q5~b`O+J T*(Cqi%fox2nmZ^QLV|n%u)Qe_ literal 0 HcmV?d00001 From ef1298ce026924bbd8f80cf279eedfed06718635 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Thu, 15 Oct 2026 03:28:25 +0000 Subject: [PATCH 002/239] fix(query): accept narrowed nullability in MergeScan remote schema check (#9592) On append-mode tables a region reports an exact row count. When it is zero, DataFusion's AggregateStatistics replaces the region-side __min_state/__max_state aggregate with a non-null struct literal, so the remote stream advertises the state column as non-nullable. The MergeScan schema check required equal nullability and failed the query. Accept a remote field that is non-nullable where the expected field is nullable, matching DataFusion's own physical-plan schema invariant. Batches are already rebuilt with the expected schema. Widening is still rejected, and nested field nullability is still compared exactly as part of the data type. Signed-off-by: Dennis Zhuang --- src/query/src/dist_plan/merge_scan.rs | 87 ++++++++++-- .../aggregate/empty_region_stats.result | 126 ++++++++++++++++++ .../common/aggregate/empty_region_stats.sql | 42 ++++++ 3 files changed, 242 insertions(+), 13 deletions(-) create mode 100644 tests/cases/standalone/common/aggregate/empty_region_stats.result create mode 100644 tests/cases/standalone/common/aggregate/empty_region_stats.sql diff --git a/src/query/src/dist_plan/merge_scan.rs b/src/query/src/dist_plan/merge_scan.rs index 041fea160cf..9717c436e0b 100644 --- a/src/query/src/dist_plan/merge_scan.rs +++ b/src/query/src/dist_plan/merge_scan.rs @@ -167,8 +167,20 @@ fn is_json_field(field: &Field) -> bool { == Some("Json") } +/// Returns true when the remote field's nullability fits the expected field. +/// +/// The remote physical optimizer may narrow a nullable column to non-nullable, +/// as DataFusion's own schema invariant permits: `AggregateStatistics` +/// replaces an aggregate over an empty region with a non-null literal. Batches +/// are rebuilt with the expected schema, so the narrowing does not leak. +/// Widening is rejected. Nested field nullability is part of the data type and +/// still has to match exactly. +fn nullability_compatible(expected_field: &Field, actual_field: &Field) -> bool { + expected_field.is_nullable() || !actual_field.is_nullable() +} + /// Returns true when two fields describe the same column semantics: same -/// name, same data type, and same nullability. +/// name, same data type, and compatible nullability. /// /// Arrow `Field` metadata is auxiliary information (indexing, encoding, /// compression hints, ...) and does not participate in the comparison. This @@ -179,7 +191,7 @@ fn is_json_field(field: &Field) -> bool { fn fields_semantically_equal(expected_field: &Field, actual_field: &Field) -> bool { expected_field.name() == actual_field.name() && expected_field.data_type() == actual_field.data_type() - && expected_field.is_nullable() == actual_field.is_nullable() + && nullability_compatible(expected_field, actual_field) } /// Returns true when two fields are semantically equivalent JSON columns. @@ -196,7 +208,7 @@ fn json_fields_compatible(expected_field: &Field, actual_field: &Field) -> bool is_json_field(expected_field) && is_json_field(actual_field) && expected_field.name() == actual_field.name() - && expected_field.is_nullable() == actual_field.is_nullable() + && nullability_compatible(expected_field, actual_field) // Both must carry the same JSON marker. && expected_field.metadata().get(datatypes::schema::TYPE_KEY) == actual_field.metadata().get(datatypes::schema::TYPE_KEY) @@ -212,14 +224,15 @@ fn json_fields_compatible(expected_field: &Field, actual_field: &Field) -> bool /// Validates the remote schema before positional column handling. /// /// Field metadata (indexing/encoding/compression hints) is non-semantic at -/// this boundary and is ignored: two fields are equal when their name, data -/// type, and nullability match (see [`fields_semantically_equal`]). A -/// timestamp timezone difference is the only intentional exception, accepted -/// when the fields share the same timestamp unit, distinct timezones, and -/// equal name and nullability. Top-level Arrow schema metadata is -/// non-semantic at this boundary as well. JSON columns are compared -/// semantically (see [`json_fields_compatible`]) because their wire and -/// decoded representations use different physical arrow types. +/// this boundary and is ignored: two fields are equal when their name and +/// data type match and their nullability is compatible (see +/// [`fields_semantically_equal`]). A timestamp timezone difference is the only +/// intentional exception, accepted when the fields share the same timestamp +/// unit, distinct timezones, equal name, and compatible nullability. +/// Top-level Arrow schema metadata is non-semantic at this boundary as well. +/// JSON columns are compared semantically (see [`json_fields_compatible`]) +/// because their wire and decoded representations use different physical +/// arrow types. fn validate_remote_schema( expected: &ArrowSchema, actual: &ArrowSchema, @@ -260,7 +273,7 @@ fn validate_remote_schema( ) if expected_unit == actual_unit && expected_timezone != actual_timezone && expected_field.name() == actual_field.name() - && expected_field.is_nullable() == actual_field.is_nullable() + && nullability_compatible(expected_field, actual_field) ); if !timezone_only_difference { return Err(remote_schema_mismatch(format!( @@ -2691,6 +2704,30 @@ mod tests { assert_eq!(batches[0].schema().as_ref(), &expected_schema); } + #[tokio::test] + async fn qbs_merge_scan_remote_schema_identity_emits_expected_schema_on_nullability_narrowing() + { + let batch = record_batch( + int64_schema(&["a", "b"]), + vec![ + Arc::new(Int64Vector::from_slice([11])) as _, + Arc::new(Int64Vector::from_slice([12])) as _, + ], + ); + let expected_schema = ArrowSchema::new(vec![ + Field::new("a", TestArrowDataType::Int64, true), + Field::new("b", TestArrowDataType::Int64, false), + ]); + let batches = collect_merge_scan(merge_scan_exec( + vec![(RegionId::new(1024, 1), batch)], + expected_schema.clone(), + 1, + )) + .await + .unwrap(); + assert_eq!(batches[0].schema().as_ref(), &expected_schema); + } + #[tokio::test] async fn qbs_merge_scan_remote_schema_identity_rejects_incompatible_type() { let batch = record_batch( @@ -2897,7 +2934,7 @@ mod tests { } #[test] - fn merge_scan_remote_schema_identity_ignores_field_metadata_but_rejects_nullability_mismatch() { + fn merge_scan_remote_schema_identity_ignores_field_metadata_but_rejects_nullability_widening() { let expected = expected_int64_schema(); let metadata_mismatch = ArrowSchema::new_with_metadata( vec![ @@ -2926,6 +2963,30 @@ mod tests { assert!(validate_remote_schema(&expected, &nullability_mismatch, "test").is_err()); } + #[test] + fn merge_scan_remote_schema_identity_accepts_top_level_nullability_narrowing_only() { + let state_type = |child_nullable| { + TestArrowDataType::Struct(arrow_schema::Fields::from(vec![Field::new( + "min[value]", + TestArrowDataType::Int64, + child_nullable, + )])) + }; + let schema = |nullable, child_nullable| { + ArrowSchema::new(vec![Field::new( + "__min_state(t.v)", + state_type(child_nullable), + nullable, + )]) + }; + + // `AggregateStatistics` turns the state aggregate over an empty region + // into a non-null struct literal. + assert!(validate_remote_schema(&schema(true, true), &schema(false, true), "test").is_ok()); + assert!(validate_remote_schema(&schema(false, true), &schema(true, true), "test").is_err()); + assert!(validate_remote_schema(&schema(true, true), &schema(true, false), "test").is_err()); + } + #[test] fn merge_scan_remote_schema_identity_ignores_skipping_index_field_metadata() { // Regression: a remote stream may advertise auxiliary field metadata diff --git a/tests/cases/standalone/common/aggregate/empty_region_stats.result b/tests/cases/standalone/common/aggregate/empty_region_stats.result new file mode 100644 index 00000000000..5f4a1bc59b0 --- /dev/null +++ b/tests/cases/standalone/common/aggregate/empty_region_stats.result @@ -0,0 +1,126 @@ +-- Append-mode regions report exact row counts, so the datanode can answer +-- count/min/max over a region from statistics. This only happens when every +-- aggregate in the query can be answered that way. +CREATE TABLE a1 (ts TIMESTAMP(9) TIME INDEX, k STRING, v BIGINT) +PARTITION ON COLUMNS (k) (k < 'b', k >= 'b' AND k < 'c', k >= 'c') +WITH (append_mode = 'true'); + +Affected Rows: 0 + +-- All regions are empty. +SELECT min(v), max(v) FROM a1; + ++-----------+-----------+ +| min(a1.v) | max(a1.v) | ++-----------+-----------+ +| | | ++-----------+-----------+ + +SELECT count(*), count(v) FROM a1; + ++----------+-------------+ +| count(*) | count(a1.v) | ++----------+-------------+ +| 0 | 0 | ++----------+-------------+ + +INSERT INTO a1 VALUES (1000, 'a', 1), (2000, 'a', 2); + +Affected Rows: 2 + +-- Two of three regions are empty. +SELECT min(v) FROM a1; + ++-----------+ +| min(a1.v) | ++-----------+ +| 1 | ++-----------+ + +SELECT max(v) FROM a1; + ++-----------+ +| max(a1.v) | ++-----------+ +| 2 | ++-----------+ + +SELECT min(ts), max(k) FROM a1; + ++----------------------------+-----------+ +| min(a1.ts) | max(a1.k) | ++----------------------------+-----------+ +| 1970-01-01T00:00:00.000001 | a | ++----------------------------+-----------+ + +SELECT count(*), count(v) FROM a1; + ++----------+-------------+ +| count(*) | count(a1.v) | ++----------+-------------+ +| 2 | 2 | ++----------+-------------+ + +SELECT min(v), count(*) FROM a1; + ++-----------+----------+ +| min(a1.v) | count(*) | ++-----------+----------+ +| 1 | 2 | ++-----------+----------+ + +SELECT sum(v), avg(v) FROM a1; + ++-----------+-----------+ +| sum(a1.v) | avg(a1.v) | ++-----------+-----------+ +| 3 | 1.5 | ++-----------+-----------+ + +SELECT k, min(v), max(v), count(*) FROM a1 GROUP BY k ORDER BY k; + ++---+-----------+-----------+----------+ +| k | min(a1.v) | max(a1.v) | count(*) | ++---+-----------+-----------+----------+ +| a | 1 | 2 | 2 | ++---+-----------+-----------+----------+ + +INSERT INTO a1 VALUES (3000, 'b', 3); + +Affected Rows: 1 + +-- One of three regions is empty. +SELECT min(v), max(v) FROM a1; + ++-----------+-----------+ +| min(a1.v) | max(a1.v) | ++-----------+-----------+ +| 1 | 3 | ++-----------+-----------+ + +SELECT k, min(v), max(v), count(*) FROM a1 GROUP BY k ORDER BY k; + ++---+-----------+-----------+----------+ +| k | min(a1.v) | max(a1.v) | count(*) | ++---+-----------+-----------+----------+ +| a | 1 | 2 | 2 | +| b | 3 | 3 | 1 | ++---+-----------+-----------+----------+ + +INSERT INTO a1 VALUES (4000, 'c', 4); + +Affected Rows: 1 + +-- No region is empty. +SELECT min(v), max(v), count(*) FROM a1; + ++-----------+-----------+----------+ +| min(a1.v) | max(a1.v) | count(*) | ++-----------+-----------+----------+ +| 1 | 4 | 4 | ++-----------+-----------+----------+ + +DROP TABLE a1; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/aggregate/empty_region_stats.sql b/tests/cases/standalone/common/aggregate/empty_region_stats.sql new file mode 100644 index 00000000000..a2d03e6b0e0 --- /dev/null +++ b/tests/cases/standalone/common/aggregate/empty_region_stats.sql @@ -0,0 +1,42 @@ +-- Append-mode regions report exact row counts, so the datanode can answer +-- count/min/max over a region from statistics. This only happens when every +-- aggregate in the query can be answered that way. +CREATE TABLE a1 (ts TIMESTAMP(9) TIME INDEX, k STRING, v BIGINT) +PARTITION ON COLUMNS (k) (k < 'b', k >= 'b' AND k < 'c', k >= 'c') +WITH (append_mode = 'true'); + +-- All regions are empty. +SELECT min(v), max(v) FROM a1; + +SELECT count(*), count(v) FROM a1; + +INSERT INTO a1 VALUES (1000, 'a', 1), (2000, 'a', 2); + +-- Two of three regions are empty. +SELECT min(v) FROM a1; + +SELECT max(v) FROM a1; + +SELECT min(ts), max(k) FROM a1; + +SELECT count(*), count(v) FROM a1; + +SELECT min(v), count(*) FROM a1; + +SELECT sum(v), avg(v) FROM a1; + +SELECT k, min(v), max(v), count(*) FROM a1 GROUP BY k ORDER BY k; + +INSERT INTO a1 VALUES (3000, 'b', 3); + +-- One of three regions is empty. +SELECT min(v), max(v) FROM a1; + +SELECT k, min(v), max(v), count(*) FROM a1 GROUP BY k ORDER BY k; + +INSERT INTO a1 VALUES (4000, 'c', 4); + +-- No region is empty. +SELECT min(v), max(v), count(*) FROM a1; + +DROP TABLE a1; From 5cfdc30a696b864fff75a078dcdecfb9ea74d212 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Thu, 15 Oct 2026 03:45:48 +0000 Subject: [PATCH 003/239] fix(table): stop claiming exact null and distinct counts for append scans (#9593) RegionScanExec reported null_count = 0 and distinct_count = total_rows for every column of an append-only table. DataFusion answers count(col) and count(DISTINCT col) from these statistics, so both returned the row count. Keep the exact row count and an exact zero null count for non-nullable columns, so count(*) is still answered from statistics. Signed-off-by: Dennis Zhuang --- src/table/src/table/scan.rs | 37 ++-- .../aggregate/append_count_nulls.result | 165 ++++++++++++++++++ .../common/aggregate/append_count_nulls.sql | 67 +++++++ 3 files changed, 258 insertions(+), 11 deletions(-) create mode 100644 tests/cases/standalone/common/aggregate/append_count_nulls.result create mode 100644 tests/cases/standalone/common/aggregate/append_count_nulls.sql diff --git a/src/table/src/table/scan.rs b/src/table/src/table/scan.rs index 1bfaea0d1df..fc7aaa24dc6 100644 --- a/src/table/src/table/scan.rs +++ b/src/table/src/table/scan.rs @@ -522,13 +522,21 @@ impl ExecutionPlan for RegionScanExec { let statistics = if scanner.properties().total_rows_is_exact() && !scanner.has_predicate_without_region() { + // Only the row count is known. DataFusion answers `count(col)` as + // `num_rows - null_count` and `count(DISTINCT col)` as `distinct_count` + // without scanning, so these must stay absent unless the schema + // guarantees them. Non-nullable columns (including the time index that + // `count(*)` is rewritten to) keep an exact zero null count. let column_statistics = self .arrow_schema .fields .iter() - .map(|_| ColumnStatistics { - distinct_count: Precision::Exact(self.total_rows), - null_count: Precision::Exact(0), // all null rows are counted for append-only table + .map(|field| ColumnStatistics { + null_count: if field.is_nullable() { + Precision::Absent + } else { + Precision::Exact(0) + }, ..Default::default() }) .collect(); @@ -1218,7 +1226,8 @@ mod test { (true, true, Precision::Exact(3)), (false, true, Precision::Absent), ] { - let (schema, batch, metadata) = dynamic_filter_test_data(5685, false); + // Column `a` is nullable and `ts` is not. + let (schema, batch, metadata) = dynamic_filter_test_data(5685, true); let scanner = Box::new(RepeatableScanner { batches: RecordBatches::try_new(schema, vec![batch]).unwrap(), properties: ScannerProperties::default() @@ -1229,13 +1238,19 @@ mod test { dynamic_filters: Arc::new(Mutex::new(Vec::new())), }); let plan = RegionScanExec::new(scanner, ScanRequest::default(), None).unwrap(); - assert_eq!( - StatisticsContext::new() - .compute(&plan, &StatisticsArgs::new()) - .unwrap() - .num_rows, - expected, - ); + let stats = StatisticsContext::new() + .compute(&plan, &StatisticsArgs::new()) + .unwrap(); + assert_eq!(stats.num_rows, expected); + if expected != Precision::Absent { + let [a, ts] = stats.column_statistics.as_slice() else { + panic!("unexpected column statistics: {stats:?}"); + }; + assert_eq!(a.null_count, Precision::Absent); + assert_eq!(ts.null_count, Precision::Exact(0)); + assert_eq!(a.distinct_count, Precision::Absent); + assert_eq!(ts.distinct_count, Precision::Absent); + } assert_eq!( StatisticsContext::new() .compute(&plan, &StatisticsArgs::new().with_partition(Some(0))) diff --git a/tests/cases/standalone/common/aggregate/append_count_nulls.result b/tests/cases/standalone/common/aggregate/append_count_nulls.result new file mode 100644 index 00000000000..acb4714224b --- /dev/null +++ b/tests/cases/standalone/common/aggregate/append_count_nulls.result @@ -0,0 +1,165 @@ +-- count(col) and count(DISTINCT col) on append-only tables must not be answered +-- from row-count statistics, while count(*) still is. +CREATE TABLE n1 (ts TIMESTAMP TIME INDEX, v BIGINT, s STRING) WITH (append_mode='true'); + +Affected Rows: 0 + +INSERT INTO n1 VALUES (1000, 1, 'a'), (2000, NULL, 'a'), (3000, NULL, NULL), (4000, 4, 'b'); + +Affected Rows: 4 + +SELECT count(*), count(v), count(s), count(DISTINCT s) FROM n1; + ++----------+-------------+-------------+----------------------+ +| count(*) | count(n1.v) | count(n1.s) | count(DISTINCT n1.s) | ++----------+-------------+-------------+----------------------+ +| 4 | 2 | 3 | 2 | ++----------+-------------+-------------+----------------------+ + +SELECT count(v) FROM n1; + ++-------------+ +| count(n1.v) | ++-------------+ +| 2 | ++-------------+ + +-- SQLNESS REPLACE (metrics.*) REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +EXPLAIN ANALYZE SELECT count(*) FROM n1; + ++-+-+-+ +| stage | node | plan_| ++-+-+-+ +| 0_| 0_|_CooperativeExec REDACTED +|_|_|_MergeScanExec: REDACTED +|_|_|_| +| 1_| 0_|_ProjectionExec: expr=[4 as count(*)] REDACTED +|_|_|_PlaceholderRowExec REDACTED +|_|_|_| +|_|_| Total rows: 1_| ++-+-+-+ + +ADMIN FLUSH_TABLE('n1'); + ++-------------------------+ +| ADMIN FLUSH_TABLE('n1') | ++-------------------------+ +| 0 | ++-------------------------+ + +SELECT count(*), count(v), count(s), count(DISTINCT s) FROM n1; + ++----------+-------------+-------------+----------------------+ +| count(*) | count(n1.v) | count(n1.s) | count(DISTINCT n1.s) | ++----------+-------------+-------------+----------------------+ +| 4 | 2 | 3 | 2 | ++----------+-------------+-------------+----------------------+ + +-- SQLNESS REPLACE (metrics.*) REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +EXPLAIN ANALYZE SELECT count(*) FROM n1; + ++-+-+-+ +| stage | node | plan_| ++-+-+-+ +| 0_| 0_|_CooperativeExec REDACTED +|_|_|_MergeScanExec: REDACTED +|_|_|_| +| 1_| 0_|_ProjectionExec: expr=[4 as count(*)] REDACTED +|_|_|_PlaceholderRowExec REDACTED +|_|_|_| +|_|_| Total rows: 1_| ++-+-+-+ + +DROP TABLE n1; + +Affected Rows: 0 + +CREATE TABLE n2 ( + host STRING, + ts TIMESTAMP TIME INDEX, + v BIGINT, + s STRING, + PRIMARY KEY (host) +) +PARTITION ON COLUMNS (host) ( + host < 'm', + host >= 'm' +) +WITH (append_mode='true'); + +Affected Rows: 0 + +INSERT INTO n2 VALUES + ('a', 1000, 1, 'a'), + ('b', 2000, NULL, 'a'), + ('x', 3000, NULL, NULL), + ('y', 4000, 4, 'b'); + +Affected Rows: 4 + +-- count is pushed down to each region as a state aggregate. +SELECT count(v), count(s) FROM n2; + ++-------------+-------------+ +| count(n2.v) | count(n2.s) | ++-------------+-------------+ +| 2 | 3 | ++-------------+-------------+ + +ADMIN FLUSH_TABLE('n2'); + ++-------------------------+ +| ADMIN FLUSH_TABLE('n2') | ++-------------------------+ +| 0 | ++-------------------------+ + +SELECT count(v), count(s) FROM n2; + ++-------------+-------------+ +| count(n2.v) | count(n2.s) | ++-------------+-------------+ +| 2 | 3 | ++-------------+-------------+ + +-- SQLNESS REPLACE (metrics.*) REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +EXPLAIN ANALYZE SELECT count(*) FROM n2; + ++-+-+-+ +| stage | node | plan_| ++-+-+-+ +| 0_| 0_|_ProjectionExec: expr=[count(Int64(1))@0 as count(*)] REDACTED +|_|_|_AggregateExec: mode=Final, gby=[], aggr=[__count_merge(__count_state(n2.ts)) as count(Int64(1))] REDACTED +|_|_|_CoalescePartitionsExec REDACTED +|_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__count_merge(__count_state(n2.ts)) as count(Int64(1))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED +|_|_|_MergeScanExec: REDACTED +|_|_|_| +| 1_| 0_|_ProjectionExec: expr=[{count[count]:2} as __count_state(n2.ts)] REDACTED +|_|_|_PlaceholderRowExec REDACTED +|_|_|_| +| 1_| 1_|_ProjectionExec: expr=[{count[count]:2} as __count_state(n2.ts)] REDACTED +|_|_|_PlaceholderRowExec REDACTED +|_|_|_| +|_|_| Total rows: 1_| ++-+-+-+ + +DROP TABLE n2; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/aggregate/append_count_nulls.sql b/tests/cases/standalone/common/aggregate/append_count_nulls.sql new file mode 100644 index 00000000000..8f2e1f5c2f7 --- /dev/null +++ b/tests/cases/standalone/common/aggregate/append_count_nulls.sql @@ -0,0 +1,67 @@ +-- count(col) and count(DISTINCT col) on append-only tables must not be answered +-- from row-count statistics, while count(*) still is. +CREATE TABLE n1 (ts TIMESTAMP TIME INDEX, v BIGINT, s STRING) WITH (append_mode='true'); + +INSERT INTO n1 VALUES (1000, 1, 'a'), (2000, NULL, 'a'), (3000, NULL, NULL), (4000, 4, 'b'); + +SELECT count(*), count(v), count(s), count(DISTINCT s) FROM n1; + +SELECT count(v) FROM n1; + +-- SQLNESS REPLACE (metrics.*) REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +EXPLAIN ANALYZE SELECT count(*) FROM n1; + +ADMIN FLUSH_TABLE('n1'); + +SELECT count(*), count(v), count(s), count(DISTINCT s) FROM n1; + +-- SQLNESS REPLACE (metrics.*) REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +EXPLAIN ANALYZE SELECT count(*) FROM n1; + +DROP TABLE n1; + +CREATE TABLE n2 ( + host STRING, + ts TIMESTAMP TIME INDEX, + v BIGINT, + s STRING, + PRIMARY KEY (host) +) +PARTITION ON COLUMNS (host) ( + host < 'm', + host >= 'm' +) +WITH (append_mode='true'); + +INSERT INTO n2 VALUES + ('a', 1000, 1, 'a'), + ('b', 2000, NULL, 'a'), + ('x', 3000, NULL, NULL), + ('y', 4000, 4, 'b'); + +-- count is pushed down to each region as a state aggregate. +SELECT count(v), count(s) FROM n2; + +ADMIN FLUSH_TABLE('n2'); + +SELECT count(v), count(s) FROM n2; + +-- SQLNESS REPLACE (metrics.*) REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +EXPLAIN ANALYZE SELECT count(*) FROM n2; + +DROP TABLE n2; From ce66d8d4ac0dbd53e50c1780c3b5729ab6f29c5b Mon Sep 17 00:00:00 2001 From: Yingwen Date: Thu, 15 Oct 2026 03:51:05 +0000 Subject: [PATCH 004/239] fix: avoid pinning memtables during series-index reconciliation (#9588) Signed-off-by: evenyag --- src/mito2/src/series_index/builder.rs | 36 ++++++----- src/mito2/src/series_index/maintenance.rs | 78 ++++++++++++++--------- src/mito2/src/series_index/tests.rs | 4 +- 3 files changed, 70 insertions(+), 48 deletions(-) diff --git a/src/mito2/src/series_index/builder.rs b/src/mito2/src/series_index/builder.rs index 63a51badaeb..447f3d409e0 100644 --- a/src/mito2/src/series_index/builder.rs +++ b/src/mito2/src/series_index/builder.rs @@ -21,6 +21,7 @@ use common_telemetry::warn; use futures::TryStreamExt; use object_store::ObjectStore; use snafu::{OptionExt, ensure}; +use store_api::metadata::RegionMetadataRef; use crate::error::{Result, UnexpectedSnafu}; use crate::read::BoxedRecordBatchStream; @@ -30,7 +31,6 @@ use crate::read::flat_projection::FlatProjectionMapper; use crate::read::prune::FlatPruneReader; use crate::read::read_columns::ReadColumns; use crate::region::MitoRegionRef; -use crate::region::version::VersionRef; use crate::series_index::bucket::SeriesBucket; use crate::series_index::catalog::{ RangeIndexEntry, SeriesIndexEntry, range_index_path, series_index_path, series_metadata, @@ -65,18 +65,18 @@ async fn reader_input( pub(crate) async fn build_range_index( store: &ObjectStore, region: &MitoRegionRef, - version: &VersionRef, + metadata: &RegionMetadataRef, file: FileHandle, ) -> Result> { let file_id = file.file_id().file_id(); let Some((context, mut selection)) = reader_input(region, file).await? else { return Ok(None); }; - let mapper = FlatProjectionMapper::new(&version.metadata, [])?; + let mapper = FlatProjectionMapper::new(metadata, [])?; let compat = FlatCompatBatch::try_new(&mapper, context.read_format(), false)?; let path = range_index_path(region.region_id, file_id); let mut writer = SstRangeIndexWriter::try_new( - version.metadata.clone(), + metadata.clone(), store.clone(), &path, SstRangeIndexWriterOptions::default(), @@ -127,13 +127,13 @@ pub(crate) async fn build_range_index( pub(crate) async fn build_series_index( store: &ObjectStore, region: &MitoRegionRef, - version: &VersionRef, + metadata: &RegionMetadataRef, bucket: &SeriesBucket, entry: &SeriesIndexEntry, purger: &IndexFilePurger, ) -> Result { let mut sources = Vec::::new(); - let mapper = FlatProjectionMapper::new(&version.metadata, [])?; + let mapper = FlatProjectionMapper::new(metadata, [])?; let schema = mapper.input_arrow_schema(false); for file in &bucket.files { let Some((context, mut selection)) = reader_input(region, file.clone()).await? else { @@ -182,7 +182,7 @@ pub(crate) async fn build_series_index( // TODO(yingwen): Deduplicate update-mode rows before series indexes are used by queries. let path = series_index_path(region.region_id, entry.index_uuid); let mut writer = SeriesIndexWriter::try_new( - version.metadata.clone(), + metadata.clone(), store.clone(), &path, SeriesIndexWriterOptions::default(), @@ -226,6 +226,7 @@ mod tests { use store_api::region_engine::RegionEngine; use super::*; + use crate::region::version::VersionRef; use crate::series_index::bucket::{group_files_into_series_buckets, plan_series_indexes}; use crate::series_index::catalog::{ SeriesIndexCatalog, load_version_control, series_catalog_path, store_catalog, @@ -280,7 +281,7 @@ mod tests { let mut range_bytes = HashMap::new(); if build_ranges { for file in files { - let completed = build_range_index(&store, ®ion, &version, file.clone()) + let completed = build_range_index(&store, ®ion, &version.metadata, file.clone()) .await .unwrap() .unwrap(); @@ -318,7 +319,7 @@ mod tests { let result = build_series_index( &store.clone().layer(layer), ®ion, - &version, + &version.metadata, &bucket, &entry, &purger, @@ -466,14 +467,15 @@ mod tests { let store = ObjectStore::new(Memory::default()).unwrap(); let (purger, _receiver) = series_index_channel(store.clone()); for file in &bucket.files { - build_range_index(&store, ®ion, &version, file.clone()) + build_range_index(&store, ®ion, &version.metadata, file.clone()) .await .unwrap() .unwrap(); } - let _handle = build_series_index(&store, ®ion, &version, &bucket, &entry, &purger) - .await - .unwrap(); + let _handle = + build_series_index(&store, ®ion, &version.metadata, &bucket, &entry, &purger) + .await + .unwrap(); let bytes = store .read(&series_index_path(region.region_id, entry.index_uuid)) .await @@ -533,14 +535,14 @@ mod tests { let result: Result> = async { for file in files { let Some(entry) = - build_range_index(&store, ®ion, &version, file.clone()).await? + build_range_index(&store, ®ion, &version.metadata, file.clone()).await? else { // Defer series construction if a needed range is not ready. return Ok(None); }; completed.push(entry.file_id); } - build_series_index(&store, ®ion, &version, &bucket, &entry, &purger) + build_series_index(&store, ®ion, &version.metadata, &bucket, &entry, &purger) .await .map(Some) } @@ -741,12 +743,12 @@ mod tests { .layer(write_layer); fail.store(fail_before_open, Ordering::Relaxed); let result = if !series { - build_range_index(&store, ®ion, &version, bucket.files[0].clone()) + build_range_index(&store, ®ion, &version.metadata, bucket.files[0].clone()) .await .map(|_| ()) } else { let (purger, _receiver) = series_index_channel(store.clone()); - build_series_index(&store, ®ion, &version, &bucket, &entry, &purger) + build_series_index(&store, ®ion, &version.metadata, &bucket, &entry, &purger) .await .map(|_| ()) }; diff --git a/src/mito2/src/series_index/maintenance.rs b/src/mito2/src/series_index/maintenance.rs index f5221bb1ebe..2520e32422e 100644 --- a/src/mito2/src/series_index/maintenance.rs +++ b/src/mito2/src/series_index/maintenance.rs @@ -19,13 +19,14 @@ use std::sync::Arc; use std::time::{Duration, Instant}; use common_telemetry::{debug, info}; +use common_time::TimeToLive; use object_store::ObjectStore; +use store_api::metadata::RegionMetadataRef; use store_api::storage::{FileId, RegionId}; use crate::error::Result; use crate::metrics::{SERIES_INDEX_RECONCILE_ELAPSED, SERIES_INDEX_RECONCILE_TOTAL}; use crate::read::series_candidate::is_sparse_metric_metadata; -use crate::region::version::VersionRef; use crate::region::{MitoRegionRef, RegionLeaderState, RegionRoleState}; use crate::series_index::bucket::{ group_files_into_series_buckets, plan_series_indexes, rounded_bucket_width, @@ -37,6 +38,15 @@ use crate::series_index::catalog::{ }; use crate::series_index::purger::IndexFilePurger; use crate::series_index::version::{SeriesIndexFileHandle, SeriesIndexVersion}; +use crate::sst::file::FileHandle; + +/// Captured index-build inputs that retain SSTs without retaining memtables. +struct ReconcileSnapshot { + metadata: RegionMetadataRef, + files: Vec, + compaction_time_window: Option, + ttl: Option, +} /// Retires newly completed series files unless their snapshot is published. #[derive(Default)] @@ -86,21 +96,36 @@ pub(crate) async fn reconcile_series_indexes( allow_builds: bool, ) -> Result { let total_start = Instant::now(); - // Use this snapshot throughout reconciliation, even if the region version advances. - let version = region.version_control.current().version; - if !is_sparse_metric_metadata(&version.metadata) { - SERIES_INDEX_RECONCILE_TOTAL - .with_label_values(&["noop"]) - .inc(); - return Ok(ReconcileStats::default()); - } + // Retaining the full version during I/O would pin flushed memtables and delay + // write-buffer reclamation. Capture consistent build inputs without the memtables. + let snapshot = { + let version = region.version_control.current().version; + if !is_sparse_metric_metadata(&version.metadata) { + SERIES_INDEX_RECONCILE_TOTAL + .with_label_values(&["noop"]) + .inc(); + return Ok(ReconcileStats::default()); + } + ReconcileSnapshot { + metadata: version.metadata.clone(), + files: version + .ssts + .levels() + .iter() + .flat_map(|level| level.files()) + .cloned() + .collect(), + compaction_time_window: version.compaction_time_window, + ttl: version.options.ttl, + } + }; let build_start = Instant::now(); let mut unpublished = UnpublishedSeriesFiles::default(); let (next, mut stats) = build_index_version( worker_id, &store, ®ion, - &version, + &snapshot, requested_bucket_width, now_ms, &purger, @@ -161,7 +186,7 @@ async fn build_index_version( worker_id: u32, store: &ObjectStore, region: &MitoRegionRef, - version: &VersionRef, + snapshot: &ReconcileSnapshot, requested_bucket_width: Duration, now_ms: i64, purger: &IndexFilePurger, @@ -170,13 +195,12 @@ async fn build_index_version( allow_builds: bool, ) -> Result<(Option, ReconcileStats)> { let mut stats = ReconcileStats::default(); - let files = version - .ssts - .levels() - .iter() - .flat_map(|level| level.files()) - .cloned() - .collect::>(); + let ReconcileSnapshot { + metadata, + files, + compaction_time_window, + ttl, + } = snapshot; stats.source_files = files.len(); let visible = files .iter() @@ -184,12 +208,12 @@ async fn build_index_version( .collect::>(); let current = region.series_index_version(); // Empty build inputs still let the planner expire established bucket coverage. - let buckets = match (allow_builds, version.compaction_time_window) { + let buckets = match (allow_builds, *compaction_time_window) { (true, Some(window)) => rounded_bucket_width(requested_bucket_width, window) // Successful rounding guarantees the window fits in i64; subsecond windows // use the same one-second minimum as rounded_bucket_width. .map(|width| { - group_files_into_series_buckets(&files, width, (window.as_secs() as i64).max(1)) + group_files_into_series_buckets(files, width, (window.as_secs() as i64).max(1)) }) .unwrap_or_default(), (true, None) => { @@ -201,12 +225,7 @@ async fn build_index_version( } (false, _) => Vec::new(), }; - let plan = plan_series_indexes( - buckets, - current.index_buckets.clone(), - version.options.ttl, - now_ms, - ); + let plan = plan_series_indexes(buckets, current.index_buckets.clone(), *ttl, now_ms); stats.computed_buckets = plan.computed_buckets; stats.skipped_buckets = plan.skipped_buckets; let mut next = prune_index_version(¤t, &visible, &plan.expired_index_ids, &mut stats); @@ -240,14 +259,15 @@ async fn build_index_version( for file in &bucket.files { if enable_range_index && !range_indexes.contains_key(&file.file_id().file_id()) - && let Some(entry) = build_range_index(store, region, version, file.clone()).await? + && let Some(entry) = + build_range_index(store, region, metadata, file.clone()).await? { stats.built_range += 1; range_indexes.insert(entry.file_id, entry); } } let series_handle = - build_series_index(store, region, version, &bucket, &expected, purger).await?; + build_series_index(store, region, metadata, &bucket, &expected, purger).await?; unpublished.0.push(series_handle.clone()); stats.built_series += 1; series_indexes.insert(expected.index_uuid, series_handle); @@ -259,7 +279,7 @@ async fn build_index_version( if range_indexes.contains_key(&file_id) { continue; } - if let Some(entry) = build_range_index(store, region, version, file).await? { + if let Some(entry) = build_range_index(store, region, metadata, file.clone()).await? { stats.built_range += 1; range_indexes.insert(entry.file_id, entry); } diff --git a/src/mito2/src/series_index/tests.rs b/src/mito2/src/series_index/tests.rs index 631bdc75b87..5f13d0b4b2d 100644 --- a/src/mito2/src/series_index/tests.rs +++ b/src/mito2/src/series_index/tests.rs @@ -720,13 +720,13 @@ async fn test_failed_reconcile_retires_only_unpublished_series() { let (bucket, entry) = &plan.builds[0]; let mut ranges = HashMap::new(); for file in &bucket.files { - let entry = build_range_index(store, ®ion, &version, file.clone()) + let entry = build_range_index(store, ®ion, &version.metadata, file.clone()) .await .unwrap() .unwrap(); ranges.insert(entry.file_id, entry); } - let handle = build_series_index(store, ®ion, &version, bucket, entry, purger) + let handle = build_series_index(store, ®ion, &version.metadata, bucket, entry, purger) .await .unwrap(); let reused_path = series_index_path(region.region_id, entry.index_uuid); From af396e60a380595e1459db9138b5561b0f9e102f Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Thu, 15 Oct 2026 04:46:37 +0000 Subject: [PATCH 005/239] perf(query): infer safe aggregate input bounds from HAVING (#9582) * perf(query): infer safe aggregate input bounds from HAVING Signed-off-by: Lei, HUANG * test(query): refresh explain snapshots for aggregate bound inference Signed-off-by: Lei, HUANG * fix(query): avoid reinferring bounds across MergeScan Signed-off-by: Lei, HUANG * fix(query): validate aggregate bound column provenance Signed-off-by: Lei, HUANG * test(query): retain zero-match groups for nonpositive HAVING thresholds Signed-off-by: Lei, HUANG * test(query): assert inferred bounds in remote scan plans Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- src/query/src/dist_plan/analyzer.rs | 2 + src/query/src/dist_plan/analyzer/test.rs | 110 +++ src/query/src/optimizer.rs | 1 + .../optimizer/infer_aggregate_input_bounds.rs | 683 ++++++++++++++++++ src/query/src/query_engine/state.rs | 14 + .../aggregate/having_input_bounds.result | 229 ++++++ .../common/aggregate/having_input_bounds.sql | 126 ++++ .../common/tql-explain-analyze/explain.result | 6 + 8 files changed, 1171 insertions(+) create mode 100644 src/query/src/optimizer/infer_aggregate_input_bounds.rs create mode 100644 tests/cases/standalone/common/aggregate/having_input_bounds.result create mode 100644 tests/cases/standalone/common/aggregate/having_input_bounds.sql diff --git a/src/query/src/dist_plan/analyzer.rs b/src/query/src/dist_plan/analyzer.rs index a39514a0e43..d90e5d31f9f 100644 --- a/src/query/src/dist_plan/analyzer.rs +++ b/src/query/src/dist_plan/analyzer.rs @@ -54,6 +54,7 @@ use crate::dist_plan::merge_scan::MergeScanLogicalPlan; use crate::dist_plan::merge_sort::MergeSortLogicalPlan; use crate::metrics::PUSH_DOWN_FALLBACK_ERRORS_TOTAL; use crate::optimizer::expr_depth; +use crate::optimizer::infer_aggregate_input_bounds::InferAggregateInputBoundsRule; use crate::options::ScheduledTimeExtension; use crate::plan::ExtractExpr; use crate::query_engine::DefaultSerializer; @@ -225,6 +226,7 @@ fn pre_merge_scan_optimizer() -> Optimizer { Arc::new(EliminateFilter::new()), Arc::new(PropagateEmptyRelation::new()), Arc::new(FilterNullJoinKeys::default()), + Arc::new(InferAggregateInputBoundsRule), Arc::new(PushDownFilter::new()), Arc::new(SimplifyExpressions::new()), ]) diff --git a/src/query/src/dist_plan/analyzer/test.rs b/src/query/src/dist_plan/analyzer/test.rs index 694f2bcde59..24f8a9bc43a 100644 --- a/src/query/src/dist_plan/analyzer/test.rs +++ b/src/query/src/dist_plan/analyzer/test.rs @@ -1633,6 +1633,116 @@ fn sibling_merge_scans_have_unique_remote_dyn_filter_producer_ids() { ); } +#[tokio::test] +async fn having_bound_reaches_remote_scan_before_aggregate_splitting() { + for pushdown in [FilterPushDownType::Inexact, FilterPushDownType::Exact] { + let table = TestTable::table_with_filter_pushdown(0, "t".to_string(), pushdown); + let ctx = SessionContext::new(); + ctx.register_table("t", Arc::new(DfTableProviderAdapter::new(table))) + .unwrap(); + // Both region-local groups and groups requiring coordinator aggregation + // must receive the bound before MIN/SUM become state aggregates. + for group_by in ["pk1, pk2", "pk3"] { + let sql = format!( + "SELECT {group_by}, MIN(CAST(ts AS BIGINT)) FROM t GROUP BY {group_by} + HAVING SUM(CASE WHEN number > 0 AND CAST(ts AS BIGINT) >= 10 + AND CAST(ts AS BIGINT) < 20 THEN 1 ELSE 0 END) > 0" + ); + let state = ctx.state(); + let plan = state.create_logical_plan(&sql).await.unwrap(); + let plan = state + .analyzer() + .execute_and_check(plan, state.config_options(), |_, _| {}) + .unwrap(); + let planned = DistPlannerAnalyzer + .analyze(plan, state.config_options()) + .unwrap(); + let remote = find_merge_scan_for_table(&planned, "t").unwrap().input(); + let scan = find_table_scan(remote, "t").unwrap(); + let bound = datafusion_expr::cast(col("t.ts"), DataType::Int64).lt(lit(20_i64)); + assert!(scan.filters.contains(&bound), "{planned}"); + assert!( + planned.to_string().contains("> Int64(0)"), + "HAVING must remain: {planned}" + ); + assert_remote_table_scan_filters_are_safe(&planned); + } + } +} + +#[tokio::test] +async fn frontend_having_does_not_reinfer_remote_bounds() { + let ctx = SessionContext::new(); + for (id, name) in [(0, "t"), (1, "u")] { + let table = TestTable::table_with_filter_pushdown( + id, + name.to_string(), + FilterPushDownType::Inexact, + ); + ctx.register_table(name, Arc::new(DfTableProviderAdapter::new(table))) + .unwrap(); + } + for source in [ + "SELECT t.pk3, t.ts, t.number FROM t JOIN u ON t.pk1 = u.pk1", + "SELECT pk3, ts, number FROM t UNION ALL SELECT pk3, ts, number FROM u", + ] { + let sql = format!( + "SELECT s.pk3, MIN(CAST(s.ts AS BIGINT)) FROM ({source}) s + GROUP BY s.pk3 HAVING SUM(CASE WHEN s.number > 0 AND CAST(s.ts AS BIGINT) < 20 + THEN 1 ELSE 0 END) > 0" + ); + let state = ctx.state(); + let plan = state.create_logical_plan(&sql).await.unwrap(); + let plan = state + .analyzer() + .execute_and_check(plan, state.config_options(), |_, _| {}) + .unwrap(); + let planned = DistPlannerAnalyzer + .analyze(plan, state.config_options()) + .unwrap(); + let remote = find_merge_scan_for_table(&planned, "t").unwrap().input(); + let scan = find_table_scan(remote, "t").unwrap(); + let bound = datafusion_expr::cast(col("t.ts"), DataType::Int64).lt(lit(20_i64)); + assert!(scan.filters.contains(&bound), "{planned}"); + + let config = datafusion_optimizer::OptimizerContext::new(); + let rule_only = Optimizer::with_rules(vec![Arc::new(InferAggregateInputBoundsRule)]); + assert_eq!( + planned, + rule_only + .optimize(planned.clone(), &config, |_, _| {}) + .unwrap() + ); + + let mut optimizer = Optimizer::new(); + let index = optimizer + .rules + .iter() + .position(|rule| rule.name() == "push_down_filter") + .unwrap(); + optimizer + .rules + .insert(index, Arc::new(InferAggregateInputBoundsRule)); + let optimized = optimizer.optimize(planned, &config, |_, _| {}).unwrap(); + // MergeScan hides remote filters from this walk. Only HAVING should + // remain on the frontend, not another copy of the input bound. + let mut filters = 0; + optimized + .apply(|plan| { + filters += usize::from(matches!(plan, LogicalPlan::Filter(_))); + Ok(TreeNodeRecursion::Continue) + }) + .unwrap(); + assert_eq!(1, filters, "{optimized}"); + assert_eq!( + optimized, + optimizer + .optimize(optimized.clone(), &config, |_, _| {}) + .unwrap() + ); + } +} + #[test] fn pre_merge_scan_optimizer_eliminates_projected_false_filter() { init_default_ut_logging(); diff --git a/src/query/src/optimizer.rs b/src/query/src/optimizer.rs index 35a3cac2639..1fc0730c4fd 100644 --- a/src/query/src/optimizer.rs +++ b/src/query/src/optimizer.rs @@ -19,6 +19,7 @@ pub mod count_wildcard; pub mod enforce_sorting; pub(crate) mod expr_depth; pub mod global_limit; +pub(crate) mod infer_aggregate_input_bounds; pub(crate) mod insert_assignment; pub(crate) mod json_schema_concretize; pub(crate) mod json_type_concretize; diff --git a/src/query/src/optimizer/infer_aggregate_input_bounds.rs b/src/query/src/optimizer/infer_aggregate_input_bounds.rs new file mode 100644 index 00000000000..87d8dc3df69 --- /dev/null +++ b/src/query/src/optimizer/infer_aggregate_input_bounds.rs @@ -0,0 +1,683 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use arrow_schema::DataType; +use datafusion::functions_aggregate::min_max::{max_udaf, min_udaf}; +use datafusion::functions_aggregate::sum::sum_udaf; +use datafusion_common::tree_node::{Transformed, TreeNode}; +use datafusion_common::{Column, DFSchema, Result, ScalarValue}; +use datafusion_expr::expr::AggregateFunction; +use datafusion_expr::utils::split_conjunction; +use datafusion_expr::{Aggregate, Expr, ExprSchemable, Filter, LogicalPlan, Operator}; +use datafusion_optimizer::optimizer::ApplyOrder; +use datafusion_optimizer::{OptimizerConfig, OptimizerRule}; + +use crate::dist_plan::MergeScanLogicalPlan; + +/// Derives an input bound from an existential HAVING condition and MIN/MAX. +/// +/// For example, `MIN(x), SUM(CASE WHEN p AND x < u THEN 1 ELSE 0 END)` with +/// HAVING `SUM(...) > 0` can discard rows outside `x < u`: they contribute zero +/// to SUM, and every surviving group has a witness below u, preserving its MIN. +/// MAX permits the symmetric lower bound. The original HAVING must remain. +/// +/// This deliberately requires exactly these two aggregates and ordinary column +/// grouping. Extra aggregates, grouping sets, floating-point extrema and arbitrary +/// expressions need separate preservation proofs. It must run before MergeScan +/// hides remote inputs, as well as in the non-distributed logical optimizer. +#[derive(Debug)] +pub(crate) struct InferAggregateInputBoundsRule; + +impl OptimizerRule for InferAggregateInputBoundsRule { + fn name(&self) -> &str { + "infer_aggregate_input_bounds" + } + + fn apply_order(&self) -> Option { + Some(ApplyOrder::TopDown) + } + + fn rewrite( + &self, + plan: LogicalPlan, + _config: &dyn OptimizerConfig, + ) -> Result> { + let LogicalPlan::Filter(mut having) = plan else { + return Ok(Transformed::no(plan)); + }; + let LogicalPlan::Aggregate(aggregate) = having.input.as_ref() else { + return Ok(Transformed::no(LogicalPlan::Filter(having))); + }; + // The pre-MergeScan pass already handled remote inputs. Their filters + // are hidden from this pass, so re-inference can add useless local work. + if aggregate.input.exists(|plan| { + Ok(matches!(plan, LogicalPlan::Extension(extension) + if extension.node.as_any().is::())) + })? { + return Ok(Transformed::no(LogicalPlan::Filter(having))); + } + let Some(bound) = infer_bound(aggregate, &having.predicate)? else { + return Ok(Transformed::no(LogicalPlan::Filter(having))); + }; + if already_filtered(&aggregate.input, &bound) { + return Ok(Transformed::no(LogicalPlan::Filter(having))); + } + + let mut aggregate = aggregate.clone(); + aggregate.input = Arc::new(LogicalPlan::Filter(Filter::try_new( + bound, + aggregate.input, + )?)); + having.input = Arc::new(LogicalPlan::Aggregate(aggregate)); + Ok(Transformed::yes(LogicalPlan::Filter(having))) + } +} + +fn infer_bound(aggregate: &Aggregate, having: &Expr) -> Result> { + if aggregate.aggr_expr.len() != 2 + || aggregate.group_expr.is_empty() + || !aggregate + .group_expr + .iter() + .all(|e| matches!(e, Expr::Column(_))) + || !safe_predicate(having, aggregate.schema.as_ref())? + { + return Ok(None); + } + + for conjunct in split_conjunction(having) { + let Some(sum_index) = positive_sum_index(conjunct, aggregate) else { + continue; + }; + let Some(sum) = plain_aggregate(&aggregate.aggr_expr[sum_index]) else { + continue; + }; + let Some(extremum) = plain_aggregate(&aggregate.aggr_expr[1 - sum_index]) else { + continue; + }; + if let Some(bound) = bound_from_aggregates(sum, extremum, &aggregate.input)? { + return Ok(Some(bound)); + } + } + Ok(None) +} + +fn positive_sum_index(predicate: &Expr, aggregate: &Aggregate) -> Option { + let Expr::BinaryExpr(comparison) = predicate else { + return None; + }; + let (output, threshold, operator) = match (comparison.left.as_ref(), comparison.right.as_ref()) + { + (Expr::Column(column), literal) => (column, literal, comparison.op), + (literal, Expr::Column(column)) => (column, literal, comparison.op.swap()?), + _ => return None, + }; + if !((operator == Operator::Gt && integer_literal(threshold) == Some(0)) + || (operator == Operator::GtEq && integer_literal(threshold) == Some(1))) + { + return None; + } + let index = aggregate.schema.index_of_column(output).ok()?; + index + .checked_sub(aggregate.group_expr.len()) + .filter(|&i| i < 2) +} + +fn plain_aggregate(expr: &Expr) -> Option<&AggregateFunction> { + let expr = match expr { + Expr::Alias(alias) => alias.expr.as_ref(), + _ => expr, + }; + let Expr::AggregateFunction(function) = expr else { + return None; + }; + let params = &function.params; + (params.args.len() == 1 + && !params.distinct + && params.filter.is_none() + && params.order_by.is_empty() + && params.null_treatment.is_none()) + .then_some(function) +} + +fn bound_from_aggregates( + sum: &AggregateFunction, + extremum: &AggregateFunction, + input: &LogicalPlan, +) -> Result> { + let schema = input.schema(); + // Compare function identities, not names: a user-defined "sum" need not + // have the built-in aggregate's zero-contribution semantics. + if sum.func != sum_udaf() || (extremum.func != min_udaf() && extremum.func != max_udaf()) { + return Ok(None); + } + let argument = &extremum.params.args[0]; + let argument_type = argument.get_type(schema)?; + if !(argument_type.is_integer() || matches!(argument_type, DataType::Timestamp(..))) + || !safe_operand(argument, schema)? + { + return Ok(None); + } + let Some(condition) = indicator_condition(&sum.params.args[0]) else { + return Ok(None); + }; + if !safe_predicate(condition, schema)? + || !columns_are_from_scan(argument, input)? + || !columns_are_from_scan(condition, input)? + { + return Ok(None); + } + Ok(split_conjunction(condition) + .into_iter() + .find_map(|predicate| extremum_bound(predicate, argument, extremum, &argument_type))) +} + +fn indicator_condition(expr: &Expr) -> Option<&Expr> { + let Expr::Case(case) = expr else { + return None; + }; + let [(condition, contribution)] = case.when_then_expr.as_slice() else { + return None; + }; + (case.expr.is_none() + && integer_literal(contribution) == Some(1) + && case.else_expr.as_deref().and_then(integer_literal) == Some(0)) + .then_some(condition) +} + +fn extremum_bound( + predicate: &Expr, + argument: &Expr, + extremum: &AggregateFunction, + argument_type: &DataType, +) -> Option { + let Expr::BinaryExpr(comparison) = predicate else { + return None; + }; + let (literal, operator) = if comparison.left.as_ref() == argument { + (comparison.right.as_ref(), comparison.op) + } else if comparison.right.as_ref() == argument { + (comparison.left.as_ref(), comparison.op.swap()?) + } else { + return None; + }; + let Expr::Literal(bound, _) = literal else { + return None; + }; + if bound.is_null() || &bound.data_type() != argument_type { + return None; + } + let preserves_extremum = if extremum.func == min_udaf() { + matches!(operator, Operator::Lt | Operator::LtEq) + } else { + matches!(operator, Operator::Gt | Operator::GtEq) + }; + preserves_extremum.then(|| predicate.clone()) +} + +fn integer_literal(expr: &Expr) -> Option { + match expr { + Expr::Literal(ScalarValue::Int64(integer), _) => *integer, + Expr::Literal(ScalarValue::Int32(integer), _) => integer.map(i64::from), + _ => None, + } +} + +/// Rejects explicitly volatile, fallible or subquery expressions. Column +/// references additionally require input-plan provenance checks. +fn safe_predicate(expr: &Expr, schema: &DFSchema) -> Result { + match expr { + Expr::BinaryExpr(binary) => Ok(matches!( + binary.op, + Operator::And + | Operator::Or + | Operator::Eq + | Operator::NotEq + | Operator::Lt + | Operator::LtEq + | Operator::Gt + | Operator::GtEq + ) && safe_predicate(&binary.left, schema)? + && safe_predicate(&binary.right, schema)?), + Expr::IsNull(expr) | Expr::IsNotNull(expr) => safe_operand(expr, schema), + _ => safe_operand(expr, schema), + } +} + +fn safe_operand(expr: &Expr, schema: &DFSchema) -> Result { + match expr { + Expr::Column(_) | Expr::Literal(..) => Ok(true), + // Timestamp-to-Int64 preserves its native unit and cannot fail. Keep + // the cast intact rather than assuming timestamps are nanoseconds. + Expr::Cast(cast) => Ok(cast.field.data_type() == &DataType::Int64 + && matches!(cast.expr.as_ref(), Expr::Column(_)) + && matches!(cast.expr.get_type(schema)?, DataType::Timestamp(..))), + _ => Ok(false), + } +} + +/// A projected column can hide division, a fallible cast or a volatile UDF. +/// PushDownFilter may inline it below the WHERE that made evaluation safe. +fn columns_are_from_scan(expr: &Expr, input: &LogicalPlan) -> Result { + for column in expr.column_refs() { + if !column_is_from_scan(column, input)? { + return Ok(false); + } + } + Ok(true) +} + +fn column_is_from_scan(column: &Column, input: &LogicalPlan) -> Result { + let index = input.schema().index_of_column(column)?; + match input { + LogicalPlan::TableScan(_) => Ok(true), + LogicalPlan::Filter(filter) => column_is_from_scan(column, &filter.input), + LogicalPlan::SubqueryAlias(alias) => column_at_is_from_scan(index, &alias.input), + LogicalPlan::Projection(projection) => { + let expr = match &projection.expr[index] { + Expr::Alias(alias) => alias.expr.as_ref(), + expr => expr, + }; + Ok(safe_operand(expr, projection.input.schema())? + && columns_are_from_scan(expr, &projection.input)?) + } + LogicalPlan::Join(join) => { + match ( + join.left.schema().has_column(column), + join.right.schema().has_column(column), + ) { + (true, false) => column_is_from_scan(column, &join.left), + (false, true) => column_is_from_scan(column, &join.right), + _ => Ok(false), + } + } + LogicalPlan::Union(union) => { + for input in &union.inputs { + if !column_at_is_from_scan(index, input)? { + return Ok(false); + } + } + Ok(true) + } + // Unknown plan nodes may synthesize columns or hide their origin. + _ => Ok(false), + } +} + +fn column_at_is_from_scan(index: usize, input: &LogicalPlan) -> Result { + let (qualifier, field) = input.schema().qualified_field(index); + column_is_from_scan(&Column::new(qualifier.cloned(), field.name()), input) +} + +fn already_filtered(input: &LogicalPlan, bound: &Expr) -> bool { + match input { + LogicalPlan::Filter(filter) => { + split_conjunction(&filter.predicate).contains(&bound) + || already_filtered(&filter.input, bound) + } + LogicalPlan::TableScan(scan) => scan + .filters + .iter() + .any(|predicate| split_conjunction(predicate).contains(&bound)), + _ => false, + } +} + +#[cfg(test)] +mod tests { + use arrow::array::{ + BooleanArray, DictionaryArray, Float64Array, Int64Array, StringArray, + TimestampMillisecondArray, TimestampNanosecondArray, UInt32Array, + }; + use arrow::datatypes::UInt32Type; + use arrow::record_batch::RecordBatch; + use arrow::util::pretty::pretty_format_batches; + use datafusion::prelude::SessionContext; + use datafusion_optimizer::OptimizerContext; + use datafusion_optimizer::optimizer::Optimizer; + use datafusion_optimizer::push_down_filter::PushDownFilter; + use datafusion_optimizer::simplify_expressions::SimplifyExpressions; + + use super::*; + + fn context() -> SessionContext { + // Covers an earlier nonmatching row, both interval endpoints, NULL + // values/predicates/group keys, and groups with no matching rows. + let keys = vec![ + Some("a"), + Some("a"), + Some("a"), + Some("a"), + Some("a"), + Some("b"), + Some("c"), + Some("c"), + None, + None, + None, + ]; + let numbers = vec![ + Some(5), + Some(10), + Some(20), + Some(25), + None, + Some(25), + Some(15), + Some(4), + Some(2), + Some(12), + Some(22), + ]; + let flags = vec![ + Some(false), + Some(true), + Some(true), + Some(true), + Some(true), + Some(true), + None, + Some(false), + Some(false), + Some(true), + Some(false), + ]; + let services = DictionaryArray::::try_new( + UInt32Array::from( + flags + .iter() + .map(|flag| flag.map(|f| u32::from(!f))) + .collect::>(), + ), + Arc::new(StringArray::from(vec!["cart", "other"])), + ) + .unwrap(); + let batch = RecordBatch::try_from_iter(vec![ + ("k", Arc::new(StringArray::from(keys)) as _), + ("service", Arc::new(services) as _), + ("v", Arc::new(Int64Array::from(numbers.clone())) as _), + ( + "f", + Arc::new(Float64Array::from(vec![1.0; numbers.len()])) as _, + ), + ( + "ts", + Arc::new(TimestampNanosecondArray::from(numbers.clone())) as _, + ), + ( + "ts_ms", + Arc::new(TimestampMillisecondArray::from(numbers)) as _, + ), + ("flag", Arc::new(BooleanArray::from(flags)) as _), + ]) + .unwrap(); + let ctx = SessionContext::new(); + ctx.register_batch("t", batch).unwrap(); + ctx + } + + async fn analyzed(ctx: &SessionContext, sql: &str) -> LogicalPlan { + let state = ctx.state(); + let plan = state.create_logical_plan(sql).await.unwrap(); + let plan = state + .analyzer() + .execute_and_check(plan, state.config_options(), |_, _| {}) + .unwrap(); + Optimizer::with_rules(vec![Arc::new(SimplifyExpressions::new())]) + .optimize(plan, &OptimizerContext::new(), |_, _| {}) + .unwrap() + } + + fn optimize(plan: LogicalPlan) -> LogicalPlan { + Optimizer::with_rules(vec![Arc::new(InferAggregateInputBoundsRule)]) + .optimize(plan, &OptimizerContext::new(), |_, _| {}) + .unwrap() + } + + fn aggregate_input_bound(plan: &LogicalPlan) -> Option<&Expr> { + if let LogicalPlan::Aggregate(aggregate) = plan + && let LogicalPlan::Filter(filter) = aggregate.input.as_ref() + { + return Some(&filter.predicate); + } + plan.inputs().into_iter().find_map(aggregate_input_bound) + } + + async fn rows(ctx: &SessionContext, plan: LogicalPlan) -> String { + let batches = ctx + .execute_logical_plan(plan) + .await + .unwrap() + .collect() + .await + .unwrap(); + pretty_format_batches(&batches).unwrap().to_string() + } + + #[tokio::test] + async fn existential_having_preserves_extrema_and_conditional_sum() { + let ctx = context(); + for (extremum, argument, condition, having, expected_bound) in [ + ( + "MIN", + "v", + "service = 'cart' AND v >= 10 AND v < 20", + "n > 0", + "t.v < Int64(20)", + ), + ("MIN", "v", "flag AND v <= 20", "n >= 1", "t.v <= Int64(20)"), + ( + "MAX", + "v", + "flag AND v >= 10 AND v < 20", + "0 < n", + "t.v >= Int64(10)", + ), + ("MAX", "v", "flag AND 10 < v", "1 <= n", "t.v > Int64(10)"), + ( + "MIN", + "CAST(ts AS BIGINT)", + "flag AND CAST(ts AS BIGINT) < 20", + "n > 0", + "CAST(t.ts AS Int64) < Int64(20)", + ), + ( + "MIN", + "CAST(ts_ms AS BIGINT)", + "flag AND CAST(ts_ms AS BIGINT) < 20", + "n > 0", + "CAST(t.ts_ms AS Int64) < Int64(20)", + ), + ( + "MIN", + "ts", + "flag AND ts < TIMESTAMP '1970-01-01 00:00:00.000000020'", + "n > 0", + "t.ts < TimestampNanosecond(20, None)", + ), + ] { + let sql = format!( + "SELECT k, {extremum}({argument}) AS endpoint, + SUM(CASE WHEN {condition} THEN 1 ELSE 0 END) AS n + FROM t GROUP BY k HAVING {having} ORDER BY k NULLS FIRST" + ); + let original = analyzed(&ctx, &sql).await; + let optimized = optimize(original.clone()); + assert_eq!( + aggregate_input_bound(&optimized).unwrap().to_string(), + expected_bound, + "{sql}" + ); + assert_eq!( + optimized, + optimize(optimized.clone()), + "rule must be idempotent: {sql}" + ); + assert_eq!( + rows(&ctx, original).await, + rows(&ctx, optimized).await, + "{sql}" + ); + } + } + + #[tokio::test] + async fn guarded_derived_operands_are_not_pruned() { + let ctx = SessionContext::new(); + let batch = RecordBatch::try_from_iter(vec![ + ("k", Arc::new(StringArray::from(vec!["a", "a", "b"])) as _), + ("v", Arc::new(Int64Array::from(vec![12, 5, 25])) as _), + ("d", Arc::new(Int64Array::from(vec![0, 1, 1])) as _), + ( + "raw", + Arc::new(StringArray::from(vec!["bad", "5", "25"])) as _, + ), + ]) + .unwrap(); + ctx.register_batch("guarded", batch).unwrap(); + ctx.sql("CREATE VIEW guarded_view AS SELECT k, v / d AS r FROM guarded WHERE abs(d) > 0") + .await + .unwrap() + .collect() + .await + .unwrap(); + for source in [ + "SELECT k, v / d AS r FROM guarded WHERE abs(d) > 0", + "SELECT k, quotient AS r FROM (SELECT k, v / d AS quotient FROM guarded WHERE abs(d) > 0) q", + "SELECT k, CAST(raw AS BIGINT) AS r FROM guarded WHERE abs(d) > 0", + "SELECT k, r FROM guarded_view", + "SELECT k, v AS r FROM guarded UNION ALL SELECT k, v / d AS r FROM guarded WHERE abs(d) > 0", + "SELECT l.k, l.r FROM guarded_view l JOIN guarded q ON l.k = q.k", + ] { + let sql = format!( + "SELECT k, MIN(r) FROM ({source}) s GROUP BY k + HAVING SUM(CASE WHEN r < 20 THEN 1 ELSE 0 END) > 0 ORDER BY k" + ); + let original = analyzed(&ctx, &sql).await; + let optimized = optimize(original.clone()); + assert_eq!( + rows(&ctx, original.clone()).await, + rows(&ctx, optimized.clone()).await, + "{sql}" + ); + assert_eq!(original, optimized, "{sql}"); + } + } + + #[tokio::test] + async fn scan_column_aliases_still_allow_bounds() { + let ctx = context(); + for source in [ + "SELECT k, v AS r FROM t", + "SELECT k, renamed AS r FROM (SELECT k, v AS renamed FROM t) q", + "SELECT k, CAST(ts AS BIGINT) AS r FROM t", + "SELECT k, v AS r FROM t UNION ALL SELECT k, v AS r FROM t", + "SELECT l.k, l.v AS r FROM t l JOIN t r ON l.k = r.k", + ] { + let sql = format!( + "SELECT k, MIN(r) FROM ({source}) s GROUP BY k + HAVING SUM(CASE WHEN r < 20 THEN 1 ELSE 0 END) > 0 ORDER BY k NULLS FIRST" + ); + let original = analyzed(&ctx, &sql).await; + let optimized = optimize(original.clone()); + assert!(aggregate_input_bound(&optimized).is_some(), "{sql}"); + assert_eq!( + rows(&ctx, original).await, + rows(&ctx, optimized).await, + "{sql}" + ); + } + } + + #[tokio::test] + async fn unsafe_or_unsupported_aggregates_are_not_pruned() { + let ctx = context(); + let indicator = "SUM(CASE WHEN flag AND v < 20 THEN 1 ELSE 0 END)"; + for sql in [ + format!("SELECT k, MIN(v), COUNT(*) FROM t GROUP BY k HAVING {indicator}>0"), + format!("SELECT k, MIN(v), MAX(v) FROM t GROUP BY k HAVING {indicator}>0"), + format!("SELECT k, MIN(f) FROM t GROUP BY k HAVING {indicator}>0"), + format!("SELECT k, MAX(v) FROM t GROUP BY k HAVING {indicator}>0"), + format!("SELECT MIN(v) FROM t HAVING {indicator}>0"), + format!("SELECT k, MIN(v) FROM t GROUP BY ROLLUP(k) HAVING {indicator}>0"), + format!("SELECT k, MIN(v) FROM t GROUP BY k HAVING {indicator}=0"), + format!("SELECT k, MIN(v) FROM t GROUP BY k HAVING {indicator}>0 OR MIN(v)>0"), + format!("SELECT k, MIN(v) FILTER (WHERE flag) FROM t GROUP BY k HAVING {indicator}>0"), + "SELECT k, MIN(v) FROM t GROUP BY k HAVING SUM(CASE WHEN v<20 THEN 1 ELSE 2 END)>0".into(), + "SELECT k, MIN(v) FROM t GROUP BY k HAVING SUM(CASE WHEN v<20 THEN 1 END)>0".into(), + "SELECT k, MIN(v) FROM t GROUP BY k HAVING SUM(DISTINCT CASE WHEN v<20 THEN 1 ELSE 0 END)>0".into(), + "SELECT k, MIN(v) FROM t GROUP BY k HAVING SUM(CASE WHEN v<20 OR flag THEN 1 ELSE 0 END)>0".into(), + "SELECT k, MIN(v) FROM t GROUP BY k HAVING SUM(CASE WHEN random()>0.5 AND v<20 THEN 1 ELSE 0 END)>0".into(), + "SELECT k, MIN(CAST(k AS BIGINT)) FROM t GROUP BY k HAVING SUM(CASE WHEN CAST(k AS BIGINT)<20 THEN 1 ELSE 0 END)>0".into(), + "SELECT k, MIN(v) FROM t GROUP BY k HAVING SUM(CASE WHEN v<20 THEN 1 WHEN flag THEN 1 ELSE 0 END)>0".into(), + "SELECT k, MIN(f) FROM t GROUP BY k HAVING SUM(CASE WHEN f<20.0 THEN 1 ELSE 0 END)>0".into(), + "SELECT k, MIN(v) FROM t GROUP BY k HAVING SUM(CASE WHEN v<20 THEN 1 ELSE 0 END)>0 AND SUM(CASE WHEN v<100 THEN 1 ELSE 0 END)>0".into(), + ] { + let original = analyzed(&ctx, &sql).await; + assert_eq!(original, optimize(original.clone()), "{sql}"); + } + } + + #[tokio::test] + async fn nonpositive_thresholds_preserve_zero_match_groups() { + let ctx = context(); + let expected = [ + "+---+----------+", + "| k | endpoint |", + "+---+----------+", + "| | 2 |", + "| a | 5 |", + "| b | 25 |", + "| c | 4 |", + "+---+----------+", + ] + .join("\n"); + for threshold in [">= 0", "> -1"] { + let sql = format!( + "SELECT k, MIN(v) AS endpoint FROM t GROUP BY k + HAVING SUM(CASE WHEN flag AND v < 20 THEN 1 ELSE 0 END) {threshold} + ORDER BY k NULLS FIRST" + ); + let original = analyzed(&ctx, &sql).await; + let optimized = optimize(original.clone()); + assert_eq!(original, optimized, "{sql}"); + assert_eq!(expected, rows(&ctx, optimized).await, "{sql}"); + } + } + + #[tokio::test] + async fn bound_survives_filter_pushdown_without_being_added_again() { + let ctx = context(); + let original = analyzed( + &ctx, + "SELECT k, MIN(v) FROM t GROUP BY k + HAVING SUM(CASE WHEN flag AND v<20 THEN 1 ELSE 0 END)>0", + ) + .await; + let optimizer = Optimizer::with_rules(vec![ + Arc::new(InferAggregateInputBoundsRule), + Arc::new(PushDownFilter::new()), + ]); + let optimized = optimizer + .optimize(original, &OptimizerContext::new(), |_, _| {}) + .unwrap(); + let repeated = optimizer + .optimize(optimized.clone(), &OptimizerContext::new(), |_, _| {}) + .unwrap(); + assert_eq!(optimized, repeated); + assert!(aggregate_input_bound(&optimized).is_some()); + } +} diff --git a/src/query/src/query_engine/state.rs b/src/query/src/query_engine/state.rs index be681cfcf57..361a0d2991d 100644 --- a/src/query/src/query_engine/state.rs +++ b/src/query/src/query_engine/state.rs @@ -68,6 +68,7 @@ use crate::optimizer::count_wildcard::CountWildcardToTimeIndexRule; use crate::optimizer::enforce_sorting::EnforceSorting; use crate::optimizer::expr_depth; use crate::optimizer::global_limit::EnsureGlobalLimitForFetch; +use crate::optimizer::infer_aggregate_input_bounds::InferAggregateInputBoundsRule; use crate::optimizer::json_schema_concretize::JsonSchemaConcretizeRule; use crate::optimizer::json_type_concretize::JsonTypeConcretizeRule; use crate::optimizer::parallelize_scan::ParallelizeScan; @@ -224,6 +225,19 @@ impl QueryEngineState { // handled by the parser's blacklist. let mut optimizer = Optimizer::new(); + // Constants must be simplified, but CASE arguments must not yet be + // hidden by common-subexpression extraction. + let filter_pushdown_index = optimizer + .rules + .iter() + .position(|rule| rule.name() == "push_down_filter") + .ok_or_else(|| { + datafusion_common::internal_datafusion_err!("Missing PushDownFilter rule") + })?; + optimizer.rules.insert( + filter_pushdown_index, + Arc::new(InferAggregateInputBoundsRule), + ); optimizer.rules.push(Arc::new(ScanHintRule)); optimizer.rules.push(Arc::new(JsonTypeConcretizeRule)); diff --git a/tests/cases/standalone/common/aggregate/having_input_bounds.result b/tests/cases/standalone/common/aggregate/having_input_bounds.result new file mode 100644 index 00000000000..0f24ab70b8f --- /dev/null +++ b/tests/cases/standalone/common/aggregate/having_input_bounds.result @@ -0,0 +1,229 @@ +-- Matching spans and the earliest/latest spans deliberately live in different regions. +CREATE TABLE having_input_bounds ( + trace_id STRING, + service_name STRING, + span_name STRING, + product STRING, + duration BIGINT, + ts TIMESTAMP(9) TIME INDEX, + PRIMARY KEY (service_name) +) +PARTITION ON COLUMNS (service_name) ( + service_name < 'm', + service_name >= 'm' +) +WITH ('append_mode' = 'true'); + +Affected Rows: 0 + +INSERT INTO having_input_bounds VALUES + ('a', 'other', 'parent', NULL, 30, 5), + ('a', 'cart', 'GetCart', NULL, 6, 10), + ('a', 'frontend', 'GetProduct', 'product-7', 8, 15), + ('a', 'cart', 'GetCart', NULL, 6, 20), + ('a', 'other', 'parent', NULL, 30, 25), + ('b', 'other', 'parent', NULL, 1, 2), + ('b', 'frontend', 'GetProduct', 'product-x', 18, 12), + ('b', 'other', 'parent', NULL, 30, 22), + ('c', 'cart', 'GetCart', NULL, 5, 30), + ('d', 'frontend', NULL, 'product-7', 8, 15), + (NULL, 'other', 'parent', NULL, 2, 3), + (NULL, 'cart', 'GetCart', NULL, 6, 11), + (NULL, 'other', 'parent', NULL, 30, 25); + +Affected Rows: 13 + +-- T02 shape: only the upper bound may be inferred; preserve earlier spans and NULL keys. +SELECT trace_id, MIN(CAST(ts AS BIGINT)) AS trace_start +FROM having_input_bounds +GROUP BY trace_id +HAVING SUM(CASE WHEN service_name = 'cart' AND span_name = 'GetCart' + AND CAST(ts AS BIGINT) >= 10 AND CAST(ts AS BIGINT) < 20 THEN 1 ELSE 0 END) > 0 +ORDER BY trace_start, trace_id +LIMIT 20; + ++----------+-------------+ +| trace_id | trace_start | ++----------+-------------+ +| | 3 | +| a | 5 | ++----------+-------------+ + +-- Pin the optimization itself: the remote scan has only the inferred time upper +-- bound, while HAVING stays above aggregation. Result equality alone cannot prove this. +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED +-- SQLNESS REPLACE (MergeScanExec:.*) MergeScanExec: REDACTED +EXPLAIN SELECT trace_id, MIN(CAST(ts AS BIGINT)) AS trace_start +FROM having_input_bounds +GROUP BY trace_id +HAVING SUM(CASE WHEN service_name = 'cart' AND span_name = 'GetCart' + AND CAST(ts AS BIGINT) >= 10 AND CAST(ts AS BIGINT) < 20 THEN 1 ELSE 0 END) > 0 +ORDER BY trace_start, trace_id +LIMIT 20; + ++-+-+ +| plan_type_| plan_| ++-+-+ +| logical_plan_| Sort: trace_start ASC NULLS LAST, having_input_bounds.trace_id ASC NULLS LAST, fetch=20_| +|_|_Projection: having_input_bounds.trace_id, min(having_input_bounds.ts) AS trace_start_| +|_|_Filter: sum(CASE WHEN having_input_bounds.service_name = Utf8("cart") AND having_input_bounds.span_name = Utf8("GetCart") AND having_input_bounds.ts >= Int64(10) AND having_input_bounds.ts < Int64(20) THEN Int64(1) ELSE Int64(0) END) > Int64(0)_| +|_|_Aggregate: groupBy=[[having_input_bounds.trace_id]], aggr=[[__min_merge(__min_state(having_input_bounds.ts)) AS min(having_input_bounds.ts), __sum_merge(__sum_state(CASE WHEN having_input_bounds.service_name = Utf8("cart") AND having_input_bounds.span_name = Utf8("GetCart") AND having_input_bounds.ts >= Int64(10) AND having_input_bounds.ts < Int64(20) THEN Int64(1) ELSE Int64(0) END)) AS sum(CASE WHEN having_input_bounds.service_name = Utf8("cart") AND having_input_bounds.span_name = Utf8("GetCart") AND having_input_bounds.ts >= Int64(10) AND having_input_bounds.ts < Int64(20) THEN Int64(1) ELSE Int64(0) END)]]_| +|_|_MergeScan [is_placeholder=false, remote_input=[_| +|_| Aggregate: groupBy=[[having_input_bounds.trace_id]], aggr=[[__min_state(CAST(having_input_bounds.ts AS Int64)), __sum_state(CASE WHEN having_input_bounds.service_name = CAST(Utf8("cart") AS Dictionary(UInt32, Utf8)) AND having_input_bounds.span_name = Utf8("GetCart") AND CAST(having_input_bounds.ts AS Int64) >= Int64(10) AND CAST(having_input_bounds.ts AS Int64) < Int64(20) THEN Int64(1) ELSE Int64(0) END)]]_| +|_|_Filter: CAST(having_input_bounds.ts AS Int64) < Int64(20)_| +|_|_TableScan: having_input_bounds, partial_filters=[CAST(having_input_bounds.ts AS Int64) < Int64(20)]_| +|_| ]]_| +| physical_plan | SortPreservingMergeExec: [trace_start@1 ASC NULLS LAST, trace_id@0 ASC NULLS LAST], fetch=20_| +|_|_SortExec: TopK(fetch=20), expr=[trace_start@1 ASC NULLS LAST, trace_id@0 ASC NULLS LAST], preserve_partitioning=[true]_| +|_|_ProjectionExec: expr=[trace_id@0 as trace_id, min(having_input_bounds.ts)@1 as trace_start]_| +|_|_FilterExec: sum(CASE WHEN having_input_bounds.service_name = Utf8("cart") AND having_input_bounds.span_name = Utf8("GetCart") AND having_input_bounds.ts >= Int64(10) AND having_input_bounds.ts < Int64(20) THEN Int64(1) ELSE Int64(0) END)@2 > 0, projection=[trace_id@0, min(having_input_bounds.ts)@1]_| +|_|_AggregateExec: mode=FinalPartitioned, gby=[trace_id@0 as trace_id], aggr=[__min_merge(__min_state(having_input_bounds.ts)) as min(having_input_bounds.ts), __sum_merge(__sum_state(CASE WHEN having_input_bounds.service_name = Utf8("cart") AND having_input_bounds.span_name = Utf8("GetCart") AND having_input_bounds.ts >= Int64(10) AND having_input_bounds.ts < Int64(20) THEN Int64(1) ELSE Int64(0) END)) as sum(CASE WHEN having_input_bounds.service_name = Utf8("cart") AND having_input_bounds.span_name = Utf8("GetCart") AND having_input_bounds.ts >= Int64(10) AND having_input_bounds.ts < Int64(20) THEN Int64(1) ELSE Int64(0) END)] | +|_|_RepartitionExec: REDACTED +|_|_AggregateExec: mode=Partial, gby=[trace_id@0 as trace_id], aggr=[__min_merge(__min_state(having_input_bounds.ts)) as min(having_input_bounds.ts), __sum_merge(__sum_state(CASE WHEN having_input_bounds.service_name = Utf8("cart") AND having_input_bounds.span_name = Utf8("GetCart") AND having_input_bounds.ts >= Int64(10) AND having_input_bounds.ts < Int64(20) THEN Int64(1) ELSE Int64(0) END)) as sum(CASE WHEN having_input_bounds.service_name = Utf8("cart") AND having_input_bounds.span_name = Utf8("GetCart") AND having_input_bounds.ts >= Int64(10) AND having_input_bounds.ts < Int64(20) THEN Int64(1) ELSE Int64(0) END)]_| +|_|_RepartitionExec: REDACTED +|_|_MergeScanExec: REDACTED +|_|_| ++-+-+ + +-- T03 shape: a span attribute is part of the existence condition, not an input filter. +SELECT trace_id, MIN(CAST(ts AS BIGINT)) AS trace_start +FROM having_input_bounds +GROUP BY trace_id +HAVING SUM(CASE WHEN service_name = 'frontend' AND span_name = 'GetProduct' + AND product = 'product-7' AND CAST(ts AS BIGINT) >= 10 + AND CAST(ts AS BIGINT) < 20 THEN 1 ELSE 0 END) > 0 +ORDER BY trace_start, trace_id; + ++----------+-------------+ +| trace_id | trace_start | ++----------+-------------+ +| a | 5 | ++----------+-------------+ + +-- T04 shape: bounds on a different expression must not become aggregate input bounds. +SELECT trace_id, MIN(CAST(ts AS BIGINT)) AS trace_start +FROM having_input_bounds +GROUP BY trace_id +HAVING SUM(CASE WHEN service_name = 'frontend' AND span_name = 'GetProduct' + AND duration >= 5 AND duration < 20 AND CAST(ts AS BIGINT) >= 10 + AND CAST(ts AS BIGINT) < 20 THEN 1 ELSE 0 END) > 0 +ORDER BY trace_start, trace_id; + ++----------+-------------+ +| trace_id | trace_start | ++----------+-------------+ +| b | 2 | +| a | 5 | ++----------+-------------+ + +-- MAX is symmetric: preserve later spans by inferring only the lower bound. +-- Also check that the conditional SUM itself remains unchanged. +SELECT trace_id, MAX(CAST(ts AS BIGINT)) AS trace_end, + SUM(CASE WHEN service_name = 'cart' AND CAST(ts AS BIGINT) >= 10 + AND CAST(ts AS BIGINT) < 20 THEN 1 ELSE 0 END) AS matches +FROM having_input_bounds +GROUP BY trace_id +HAVING matches >= 1 +ORDER BY trace_end, trace_id; + ++----------+-----------+---------+ +| trace_id | trace_end | matches | ++----------+-----------+---------+ +| a | 25 | 1 | +| | 25 | 1 | ++----------+-----------+---------+ + +-- Full-group aggregates forbid pruning, even when the MIN would be preserved. +SELECT trace_id, MIN(CAST(ts AS BIGINT)) AS trace_start, + MAX(CAST(ts AS BIGINT)) AS trace_end, COUNT(*) AS spans +FROM having_input_bounds +GROUP BY trace_id +HAVING SUM(CASE WHEN service_name = 'cart' AND CAST(ts AS BIGINT) >= 10 + AND CAST(ts AS BIGINT) < 20 THEN 1 ELSE 0 END) > 0 +ORDER BY trace_start, trace_id; + ++----------+-------------+-----------+-------+ +| trace_id | trace_start | trace_end | spans | ++----------+-------------+-----------+-------+ +| | 3 | 25 | 3 | +| a | 5 | 25 | 5 | ++----------+-------------+-----------+-------+ + +-- OR does not guarantee a matching span in every surviving group. +SELECT trace_id, MIN(CAST(ts AS BIGINT)) AS trace_start +FROM having_input_bounds +GROUP BY trace_id +HAVING SUM(CASE WHEN service_name = 'cart' AND CAST(ts AS BIGINT) < 20 THEN 1 ELSE 0 END) > 0 + OR MIN(CAST(ts AS BIGINT)) > 20 +ORDER BY trace_start, trace_id; + ++----------+-------------+ +| trace_id | trace_start | ++----------+-------------+ +| | 3 | +| a | 5 | +| c | 30 | ++----------+-------------+ + +DROP TABLE having_input_bounds; + +Affected Rows: 0 + +-- A derived column must not move fallible arithmetic/casts ahead of its WHERE guard. +CREATE TABLE having_guarded(k STRING, v BIGINT, d BIGINT, text_value STRING, ts TIMESTAMP TIME INDEX); + +Affected Rows: 0 + +INSERT INTO having_guarded VALUES ('a', 12, 0, 'bad', 1), ('a', 5, 1, '5', 2), ('b', 25, 1, '25', 3); + +Affected Rows: 3 + +SELECT k, MIN(r) FROM ( + SELECT k, v / d AS r FROM having_guarded WHERE abs(d) > 0 +) s GROUP BY k +HAVING SUM(CASE WHEN r < 20 THEN 1 ELSE 0 END) > 0 +ORDER BY k; + ++---+----------+ +| k | min(s.r) | ++---+----------+ +| a | 5 | ++---+----------+ + +SELECT k, MIN(r) FROM ( + SELECT k, CAST(text_value AS BIGINT) AS r FROM having_guarded WHERE abs(d) > 0 +) s GROUP BY k +HAVING SUM(CASE WHEN r < 20 THEN 1 ELSE 0 END) > 0 +ORDER BY k; + ++---+----------+ +| k | min(s.r) | ++---+----------+ +| a | 5 | ++---+----------+ + +CREATE VIEW having_guarded_view AS +SELECT k, v / d AS r FROM having_guarded WHERE abs(d) > 0; + +Affected Rows: 0 + +SELECT k, MIN(r) FROM having_guarded_view GROUP BY k +HAVING SUM(CASE WHEN r < 20 THEN 1 ELSE 0 END) > 0 +ORDER BY k; + ++---+----------------------------+ +| k | min(having_guarded_view.r) | ++---+----------------------------+ +| a | 5 | ++---+----------------------------+ + +DROP VIEW having_guarded_view; + +Affected Rows: 0 + +DROP TABLE having_guarded; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/aggregate/having_input_bounds.sql b/tests/cases/standalone/common/aggregate/having_input_bounds.sql new file mode 100644 index 00000000000..b18b1e2a1e8 --- /dev/null +++ b/tests/cases/standalone/common/aggregate/having_input_bounds.sql @@ -0,0 +1,126 @@ +-- Matching spans and the earliest/latest spans deliberately live in different regions. +CREATE TABLE having_input_bounds ( + trace_id STRING, + service_name STRING, + span_name STRING, + product STRING, + duration BIGINT, + ts TIMESTAMP(9) TIME INDEX, + PRIMARY KEY (service_name) +) +PARTITION ON COLUMNS (service_name) ( + service_name < 'm', + service_name >= 'm' +) +WITH ('append_mode' = 'true'); + +INSERT INTO having_input_bounds VALUES + ('a', 'other', 'parent', NULL, 30, 5), + ('a', 'cart', 'GetCart', NULL, 6, 10), + ('a', 'frontend', 'GetProduct', 'product-7', 8, 15), + ('a', 'cart', 'GetCart', NULL, 6, 20), + ('a', 'other', 'parent', NULL, 30, 25), + ('b', 'other', 'parent', NULL, 1, 2), + ('b', 'frontend', 'GetProduct', 'product-x', 18, 12), + ('b', 'other', 'parent', NULL, 30, 22), + ('c', 'cart', 'GetCart', NULL, 5, 30), + ('d', 'frontend', NULL, 'product-7', 8, 15), + (NULL, 'other', 'parent', NULL, 2, 3), + (NULL, 'cart', 'GetCart', NULL, 6, 11), + (NULL, 'other', 'parent', NULL, 30, 25); + +-- T02 shape: only the upper bound may be inferred; preserve earlier spans and NULL keys. +SELECT trace_id, MIN(CAST(ts AS BIGINT)) AS trace_start +FROM having_input_bounds +GROUP BY trace_id +HAVING SUM(CASE WHEN service_name = 'cart' AND span_name = 'GetCart' + AND CAST(ts AS BIGINT) >= 10 AND CAST(ts AS BIGINT) < 20 THEN 1 ELSE 0 END) > 0 +ORDER BY trace_start, trace_id +LIMIT 20; + +-- Pin the optimization itself: the remote scan has only the inferred time upper +-- bound, while HAVING stays above aggregation. Result equality alone cannot prove this. +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED +-- SQLNESS REPLACE (MergeScanExec:.*) MergeScanExec: REDACTED +EXPLAIN SELECT trace_id, MIN(CAST(ts AS BIGINT)) AS trace_start +FROM having_input_bounds +GROUP BY trace_id +HAVING SUM(CASE WHEN service_name = 'cart' AND span_name = 'GetCart' + AND CAST(ts AS BIGINT) >= 10 AND CAST(ts AS BIGINT) < 20 THEN 1 ELSE 0 END) > 0 +ORDER BY trace_start, trace_id +LIMIT 20; + +-- T03 shape: a span attribute is part of the existence condition, not an input filter. +SELECT trace_id, MIN(CAST(ts AS BIGINT)) AS trace_start +FROM having_input_bounds +GROUP BY trace_id +HAVING SUM(CASE WHEN service_name = 'frontend' AND span_name = 'GetProduct' + AND product = 'product-7' AND CAST(ts AS BIGINT) >= 10 + AND CAST(ts AS BIGINT) < 20 THEN 1 ELSE 0 END) > 0 +ORDER BY trace_start, trace_id; + +-- T04 shape: bounds on a different expression must not become aggregate input bounds. +SELECT trace_id, MIN(CAST(ts AS BIGINT)) AS trace_start +FROM having_input_bounds +GROUP BY trace_id +HAVING SUM(CASE WHEN service_name = 'frontend' AND span_name = 'GetProduct' + AND duration >= 5 AND duration < 20 AND CAST(ts AS BIGINT) >= 10 + AND CAST(ts AS BIGINT) < 20 THEN 1 ELSE 0 END) > 0 +ORDER BY trace_start, trace_id; + +-- MAX is symmetric: preserve later spans by inferring only the lower bound. +-- Also check that the conditional SUM itself remains unchanged. +SELECT trace_id, MAX(CAST(ts AS BIGINT)) AS trace_end, + SUM(CASE WHEN service_name = 'cart' AND CAST(ts AS BIGINT) >= 10 + AND CAST(ts AS BIGINT) < 20 THEN 1 ELSE 0 END) AS matches +FROM having_input_bounds +GROUP BY trace_id +HAVING matches >= 1 +ORDER BY trace_end, trace_id; + +-- Full-group aggregates forbid pruning, even when the MIN would be preserved. +SELECT trace_id, MIN(CAST(ts AS BIGINT)) AS trace_start, + MAX(CAST(ts AS BIGINT)) AS trace_end, COUNT(*) AS spans +FROM having_input_bounds +GROUP BY trace_id +HAVING SUM(CASE WHEN service_name = 'cart' AND CAST(ts AS BIGINT) >= 10 + AND CAST(ts AS BIGINT) < 20 THEN 1 ELSE 0 END) > 0 +ORDER BY trace_start, trace_id; + +-- OR does not guarantee a matching span in every surviving group. +SELECT trace_id, MIN(CAST(ts AS BIGINT)) AS trace_start +FROM having_input_bounds +GROUP BY trace_id +HAVING SUM(CASE WHEN service_name = 'cart' AND CAST(ts AS BIGINT) < 20 THEN 1 ELSE 0 END) > 0 + OR MIN(CAST(ts AS BIGINT)) > 20 +ORDER BY trace_start, trace_id; + +DROP TABLE having_input_bounds; + +-- A derived column must not move fallible arithmetic/casts ahead of its WHERE guard. +CREATE TABLE having_guarded(k STRING, v BIGINT, d BIGINT, text_value STRING, ts TIMESTAMP TIME INDEX); +INSERT INTO having_guarded VALUES ('a', 12, 0, 'bad', 1), ('a', 5, 1, '5', 2), ('b', 25, 1, '25', 3); + +SELECT k, MIN(r) FROM ( + SELECT k, v / d AS r FROM having_guarded WHERE abs(d) > 0 +) s GROUP BY k +HAVING SUM(CASE WHEN r < 20 THEN 1 ELSE 0 END) > 0 +ORDER BY k; + +SELECT k, MIN(r) FROM ( + SELECT k, CAST(text_value AS BIGINT) AS r FROM having_guarded WHERE abs(d) > 0 +) s GROUP BY k +HAVING SUM(CASE WHEN r < 20 THEN 1 ELSE 0 END) > 0 +ORDER BY k; + +CREATE VIEW having_guarded_view AS +SELECT k, v / d AS r FROM having_guarded WHERE abs(d) > 0; + +SELECT k, MIN(r) FROM having_guarded_view GROUP BY k +HAVING SUM(CASE WHEN r < 20 THEN 1 ELSE 0 END) > 0 +ORDER BY k; + +DROP VIEW having_guarded_view; +DROP TABLE having_guarded; diff --git a/tests/cases/standalone/common/tql-explain-analyze/explain.result b/tests/cases/standalone/common/tql-explain-analyze/explain.result index 465072da81e..3184bb261cc 100644 --- a/tests/cases/standalone/common/tql-explain-analyze/explain.result +++ b/tests/cases/standalone/common/tql-explain-analyze/explain.result @@ -122,6 +122,7 @@ TQL EXPLAIN VERBOSE (0, 10, '5s') test; | logical_plan after filter_null_join_keys_| SAME TEXT AS ABOVE_| | logical_plan after eliminate_outer_join_| SAME TEXT AS ABOVE_| | logical_plan after push_down_limit_| SAME TEXT AS ABOVE_| +| logical_plan after infer_aggregate_input_bounds_| SAME TEXT AS ABOVE_| | logical_plan after push_down_filter_| SAME TEXT AS ABOVE_| | logical_plan after single_distinct_aggregation_to_group_by | SAME TEXT AS ABOVE_| | logical_plan after eliminate_group_by_constant_| SAME TEXT AS ABOVE_| @@ -155,6 +156,7 @@ TQL EXPLAIN VERBOSE (0, 10, '5s') test; | logical_plan after filter_null_join_keys_| SAME TEXT AS ABOVE_| | logical_plan after eliminate_outer_join_| SAME TEXT AS ABOVE_| | logical_plan after push_down_limit_| SAME TEXT AS ABOVE_| +| logical_plan after infer_aggregate_input_bounds_| SAME TEXT AS ABOVE_| | logical_plan after push_down_filter_| SAME TEXT AS ABOVE_| | logical_plan after single_distinct_aggregation_to_group_by | SAME TEXT AS ABOVE_| | logical_plan after eliminate_group_by_constant_| SAME TEXT AS ABOVE_| @@ -277,6 +279,7 @@ TQL EXPLAIN VERBOSE (0, 10, '5s') test AS series; | logical_plan after filter_null_join_keys_| SAME TEXT AS ABOVE_| | logical_plan after eliminate_outer_join_| SAME TEXT AS ABOVE_| | logical_plan after push_down_limit_| SAME TEXT AS ABOVE_| +| logical_plan after infer_aggregate_input_bounds_| SAME TEXT AS ABOVE_| | logical_plan after push_down_filter_| SAME TEXT AS ABOVE_| | logical_plan after single_distinct_aggregation_to_group_by | SAME TEXT AS ABOVE_| | logical_plan after eliminate_group_by_constant_| SAME TEXT AS ABOVE_| @@ -311,6 +314,7 @@ TQL EXPLAIN VERBOSE (0, 10, '5s') test AS series; | logical_plan after filter_null_join_keys_| SAME TEXT AS ABOVE_| | logical_plan after eliminate_outer_join_| SAME TEXT AS ABOVE_| | logical_plan after push_down_limit_| SAME TEXT AS ABOVE_| +| logical_plan after infer_aggregate_input_bounds_| SAME TEXT AS ABOVE_| | logical_plan after push_down_filter_| SAME TEXT AS ABOVE_| | logical_plan after single_distinct_aggregation_to_group_by | SAME TEXT AS ABOVE_| | logical_plan after eliminate_group_by_constant_| SAME TEXT AS ABOVE_| @@ -464,6 +468,7 @@ TQL EXPLAIN VERBOSE (0, 10, '5s') test_nano; | logical_plan after filter_null_join_keys_| SAME TEXT AS ABOVE_| | logical_plan after eliminate_outer_join_| SAME TEXT AS ABOVE_| | logical_plan after push_down_limit_| SAME TEXT AS ABOVE_| +| logical_plan after infer_aggregate_input_bounds_| SAME TEXT AS ABOVE_| | logical_plan after push_down_filter_| SAME TEXT AS ABOVE_| | logical_plan after single_distinct_aggregation_to_group_by | SAME TEXT AS ABOVE_| | logical_plan after eliminate_group_by_constant_| SAME TEXT AS ABOVE_| @@ -498,6 +503,7 @@ TQL EXPLAIN VERBOSE (0, 10, '5s') test_nano; | logical_plan after filter_null_join_keys_| SAME TEXT AS ABOVE_| | logical_plan after eliminate_outer_join_| SAME TEXT AS ABOVE_| | logical_plan after push_down_limit_| SAME TEXT AS ABOVE_| +| logical_plan after infer_aggregate_input_bounds_| SAME TEXT AS ABOVE_| | logical_plan after push_down_filter_| SAME TEXT AS ABOVE_| | logical_plan after single_distinct_aggregation_to_group_by | SAME TEXT AS ABOVE_| | logical_plan after eliminate_group_by_constant_| SAME TEXT AS ABOVE_| From cf16c7313bec7cbfc818c2250762e19b7a349e6e Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Thu, 15 Oct 2026 07:47:36 +0000 Subject: [PATCH 006/239] fix(jaeger): return the most recent traces when a search hits its limit (#9595) * fix(jaeger): return the most recent traces when a search hits its limit Trace search applied LIMIT to an unordered DISTINCT trace_id, so a search matching more traces than the limit returned an arbitrary subset that changed between identical requests. Order traces by their latest matching span before applying the limit. Signed-off-by: Dennis Zhuang * refactor(jaeger): name the trace dataframe builder with a verb and document execution Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/frontend/src/instance/jaeger.rs | 118 ++++++++++++++++++++-------- tests-integration/tests/http.rs | 83 +++++++++++++++++++ 2 files changed, 170 insertions(+), 31 deletions(-) diff --git a/src/frontend/src/instance/jaeger.rs b/src/frontend/src/instance/jaeger.rs index ac59698678b..7a7a2c8b44d 100644 --- a/src/frontend/src/instance/jaeger.rs +++ b/src/frontend/src/instance/jaeger.rs @@ -35,6 +35,7 @@ use datafusion::dataframe::DataFrame; use datafusion::execution::SessionStateBuilder; use datafusion::execution::context::SessionContext; use datafusion::functions::core::expr_fn::coalesce; +use datafusion::functions_aggregate::expr_fn::max; use datafusion::functions_window::expr_fn::row_number; use datafusion_expr::select_expr::SelectExpr; use datafusion_expr::{Expr, ExprFunctionExt, SortExpr, col, lit, lit_timestamp_nano, wildcard}; @@ -59,6 +60,7 @@ use crate::instance::Instance; const DEFAULT_LIMIT: usize = 2000; const KEY_RN: &str = "greptime_rn"; +const KEY_LATEST_MATCH: &str = "greptime_latest_match"; impl Instance { async fn check_jaeger_query_permission(&self, ctx: &QueryContextRef) -> ServerResult<()> { @@ -226,11 +228,11 @@ impl JaegerQueryHandler for Instance { filters.push(col(DURATION_NANO_COLUMN).lt_eq(lit(max_duration))); } - // Get all distinct trace ids that match the filters. + // Get the ids of the most recent traces that match the filters. // It's equivalent to the following SQL query: // // ``` - // SELECT DISTINCT trace_id + // SELECT trace_id // FROM // {db}.{trace_table} // WHERE @@ -240,17 +242,16 @@ impl JaegerQueryHandler for Instance { // timestamp <= {end_time} AND // duration >= {min_duration} AND // duration <= {max_duration} + // GROUP BY trace_id + // ORDER BY max(timestamp) DESC, trace_id // LIMIT {limit} // ```. - let output = query_trace_table( + let output = find_latest_trace_ids( ctx.clone(), self, - vec![wildcard()], filters, - vec![], - Some(query_params.limit.unwrap_or(DEFAULT_LIMIT)), + query_params.limit.unwrap_or(DEFAULT_LIMIT), query_params.tags, - vec![col(TRACE_ID_COLUMN)], ) .await?; @@ -338,6 +339,79 @@ async fn query_trace_table( tags: Option>, distincts: Vec, ) -> ServerResult { + let dataframe = + build_trace_dataframe(&ctx, instance, selects, filters, tags, &distincts).await?; + + // Apply the distinct if needed. + let dataframe = if !distincts.is_empty() { + dataframe + .distinct_on(distincts.clone(), distincts, None) + .context(DataFusionSnafu)? + } else { + dataframe + }; + + // Apply the sorts if needed. + let dataframe = if !sorts.is_empty() { + dataframe.sort(sorts).context(DataFusionSnafu)? + } else { + dataframe + }; + + // Apply the limit if needed. + let dataframe = if let Some(limit) = limit { + dataframe.limit(0, Some(limit)).context(DataFusionSnafu)? + } else { + dataframe + }; + + execute_trace_dataframe(dataframe).await +} + +/// Returns the ids of the `limit` traces whose latest span matching `filters` +/// and `tags` is the most recent, newest first. +/// +/// Ordering before the limit keeps a search that matches more traces than the +/// limit deterministic. +async fn find_latest_trace_ids( + ctx: QueryContextRef, + instance: &Instance, + filters: Vec, + limit: usize, + tags: Option>, +) -> ServerResult { + let dataframe = + build_trace_dataframe(&ctx, instance, vec![wildcard()], filters, tags, &[]).await?; + let dataframe = dataframe + .aggregate( + vec![col(TRACE_ID_COLUMN)], + vec![max(col(TIMESTAMP_COLUMN)).alias(KEY_LATEST_MATCH)], + ) + .context(DataFusionSnafu)? + .sort(vec![ + col(KEY_LATEST_MATCH).sort(false, false), + col(TRACE_ID_COLUMN).sort(true, false), + ]) + .context(DataFusionSnafu)? + .limit(0, Some(limit)) + .context(DataFusionSnafu)? + .select(vec![col(TRACE_ID_COLUMN)]) + .context(DataFusionSnafu)?; + + execute_trace_dataframe(dataframe).await +} + +/// Reads the trace table, or the services or operations lookup table that +/// answers `selects` and `distincts`, then projects `selects` and applies +/// `filters` and the Jaeger tag filters. +async fn build_trace_dataframe( + ctx: &QueryContextRef, + instance: &Instance, + selects: Vec, + filters: Vec, + tags: Option>, + distincts: &[Expr], +) -> ServerResult { let trace_table_name = ctx .extension(JAEGER_QUERY_TABLE_NAME_KEY) .unwrap_or(TRACE_TABLE_NAME); @@ -366,7 +440,7 @@ async fn query_trace_table( ctx.current_catalog(), &ctx.current_schema(), table_name, - Some(&ctx), + Some(ctx), ) .await? .with_context(|| TableNotFoundSnafu { @@ -403,30 +477,12 @@ async fn query_trace_table( df.filter(expr).context(DataFusionSnafu) })?; - // Apply the distinct if needed. - let dataframe = if !distincts.is_empty() { - dataframe - .distinct_on(distincts.clone(), distincts, None) - .context(DataFusionSnafu)? - } else { - dataframe - }; + Ok(dataframe) +} - // Apply the sorts if needed. - let dataframe = if !sorts.is_empty() { - dataframe.sort(sorts).context(DataFusionSnafu)? - } else { - dataframe - }; - - // Apply the limit if needed. - let dataframe = if let Some(limit) = limit { - dataframe.limit(0, Some(limit)).context(DataFusionSnafu)? - } else { - dataframe - }; - - // Execute the query and collect the result. +/// Executes `dataframe` and wraps its stream into an [`Output`] with dictionary +/// columns decoded to values. +async fn execute_trace_dataframe(dataframe: DataFrame) -> ServerResult { let stream = dataframe.execute_stream().await.context(DataFusionSnafu)?; let output = Output::new_with_stream(Box::pin( diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index f3b439e19bc..1bb4af7839c 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -230,6 +230,7 @@ macro_rules! http_tests { test_jaeger_v3_query_api, test_jaeger_query_api_for_trace_v1, test_jaeger_query_api_for_trace_v2, + test_jaeger_search_limit_returns_latest_traces, test_influxdb_write, test_influxdb_write_with_hints, @@ -12480,6 +12481,88 @@ pub async fn test_jaeger_query_api_for_trace_v2( Ok(()) } +pub async fn test_jaeger_search_limit_returns_latest_traces(store_type: StorageType) { + common_telemetry::init_default_ut_logging(); + let (app, mut guard) = + setup_test_http_app_with_frontend(store_type, "test_jaeger_search_limit").await; + let client = TestClient::new(app).await; + + // Twelve single-span traces spread over all trace partitions. The oldest + // traces get the smallest ids and are written first, so neither id nor + // write order points at the newest ones. + let now = Utc::now().timestamp_nanos_opt().unwrap() as u64; + let traces = (0..12u64) + .map(|i| { + let start = now - (60 - i) * 1_000_000_000; + (format!("{:x}{:031x}", i, i + 1), start) + }) + .collect::>(); + let spans = traces + .iter() + .map(|(trace_id, start)| { + json!({ + "traceId": trace_id, + "spanId": &trace_id[16..], + "name": "op", + "kind": 2, + "startTimeUnixNano": start.to_string(), + "endTimeUnixNano": (start + 1_000).to_string(), + }) + }) + .collect::>(); + let content = json!({ + "resourceSpans": [{ + "resource": {"attributes": [{"key": "service.name", "value": {"stringValue": "svc"}}]}, + "scopeSpans": [{"scope": {"name": "s"}, "spans": spans}], + }] + }); + let req: ExportTraceServiceRequest = serde_json::from_value(content).unwrap(); + let res = send_req( + &client, + vec![ + ( + HeaderName::from_static("content-type"), + HeaderValue::from_static("application/x-protobuf"), + ), + ( + HeaderName::from_static("x-greptime-pipeline-name"), + HeaderValue::from_static(GREPTIME_INTERNAL_TRACE_PIPELINE_V1_NAME), + ), + ], + "/v1/otlp/v1/traces", + req.encode_to_vec(), + false, + ) + .await; + assert_eq!(StatusCode::OK, res.status()); + + let mut expected = traces[9..] + .iter() + .map(|(trace_id, _)| trace_id.clone()) + .collect::>(); + expected.sort(); + let path = format!( + "/v1/jaeger/api/traces?service=svc&operation=op&start={}&end={}&limit=3", + (now - 120_000_000_000) / 1_000, + now / 1_000 + ); + for _ in 0..5 { + let res = client.get(&path).send().await; + assert_eq!(StatusCode::OK, res.status()); + let resp: Value = serde_json::from_str(&res.text().await).unwrap(); + let mut found = resp["data"] + .as_array() + .unwrap() + .iter() + .map(|trace| trace["traceID"].as_str().unwrap().to_string()) + .collect::>(); + found.sort(); + assert_eq!(found, expected); + } + + guard.remove_all().await; +} + pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { common_telemetry::init_default_ut_logging(); let (app, mut guard) = From f7a3a7033c0e494ca3fe8d0af5bb2697421ab96d Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Thu, 15 Oct 2026 07:51:17 +0000 Subject: [PATCH 007/239] fix(query): keep state column names unique when splitting aggregates (#9594) * fix(query): keep state column names unique when splitting aggregates Signed-off-by: Dennis Zhuang * fix(query): probe allocated names when renaming colliding state columns Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/common/function/src/aggrs/aggr_wrapper.rs | 43 ++++++++- src/query/src/dist_plan/analyzer/test.rs | 93 +++++++++++++++++++ .../common/aggregate/multi_regions.result | 26 ++++++ .../common/aggregate/multi_regions.sql | 8 ++ 4 files changed, 166 insertions(+), 4 deletions(-) diff --git a/src/common/function/src/aggrs/aggr_wrapper.rs b/src/common/function/src/aggrs/aggr_wrapper.rs index 710c383a834..8c0c475c265 100644 --- a/src/common/function/src/aggrs/aggr_wrapper.rs +++ b/src/common/function/src/aggrs/aggr_wrapper.rs @@ -22,6 +22,7 @@ //! `foo_merge`'s input arg is the same as `foo_state`'s output, and its output is the same as `foo`'s input. //! +use std::collections::HashSet; use std::hash::{Hash, Hasher}; use std::sync::Arc; @@ -183,7 +184,13 @@ impl StateMergeHelper { ))); } }; - let mut lower_aggr_exprs = vec![]; + // (unaliased state expression, output column name) + let mut lower_states: Vec<(Expr, String)> = vec![]; + let mut lower_output_names: HashSet = aggr + .group_expr + .iter() + .map(|e| e.qualified_name().1) + .collect(); let mut upper_aggr_exprs = vec![]; // group exprs for upper plan should refer to the output group expr as column from lower plan @@ -218,9 +225,28 @@ impl StateMergeHelper { params: aggr_func.params.clone(), }; let expr = Expr::AggregateFunction(expr); - let lower_state_output_col_name = expr.schema_name().to_string(); - lower_aggr_exprs.push(expr); + // The state column is named after the unaliased state expression, which can + // repeat even though the original output names are unique: a name-preserving + // rewrite such as `count(*)` -> `count(ts) AS count(*)` can make two aggregates + // share one state expression, and schema names hide casts, so + // `count(CAST(ts AS Int64))` is named `count(t.ts)` too. Identical states are + // computed once; a different state with a taken name gets a numeric suffix. + let lower_state_output_col_name = + if let Some((_, name)) = lower_states.iter().find(|(e, _)| *e == expr) { + name.clone() + } else { + let base = expr.schema_name().to_string(); + let mut name = base.clone(); + let mut suffix = 0; + while lower_output_names.contains(&name) { + suffix += 1; + name = format!("{base}_{suffix}"); + } + lower_output_names.insert(name.clone()); + lower_states.push((expr, name.clone())); + name + }; // then create the merge function using the physical expression of the original aggregate function let (name, human_display) = match aggr_expr { @@ -270,7 +296,16 @@ impl StateMergeHelper { } let mut lower = aggr.clone(); - lower.aggr_expr = lower_aggr_exprs; + lower.aggr_expr = lower_states + .into_iter() + .map(|(expr, name)| { + if expr.schema_name().to_string() == name { + expr + } else { + expr.alias(name) + } + }) + .collect(); let lower_plan = LogicalPlan::Aggregate(lower); // update aggregate's output schema diff --git a/src/query/src/dist_plan/analyzer/test.rs b/src/query/src/dist_plan/analyzer/test.rs index 24f8a9bc43a..a689d94c282 100644 --- a/src/query/src/dist_plan/analyzer/test.rs +++ b/src/query/src/dist_plan/analyzer/test.rs @@ -58,6 +58,7 @@ use table::table::numbers::NumbersTable; use table::{Table, TableRef}; use super::*; +use crate::optimizer::count_wildcard::CountWildcardToTimeIndexRule; fn collect_merge_scan_remote_dyn_filter_producer_ids( plan: &LogicalPlan, @@ -1997,6 +1998,98 @@ fn expand_step_aggr_avg_limit() { assert_eq!(expected, result.to_string()); } +/// `count(*)` and `count(1)` are rewritten to `count(ts)` with their names preserved, so +/// all three share one state expression and must share one state column. +#[test] +fn expand_step_aggr_count_wildcard_and_time_index() { + init_default_ut_logging(); + let test_table = TestTable::table_with_name(0, "t".to_string()); + let table_source = Arc::new(DefaultTableSource::new(Arc::new( + DfTableProviderAdapter::new(test_table), + ))); + let plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![]) + .unwrap() + .aggregate( + Vec::::new(), + vec![ + datafusion::functions_aggregate::count::count_all(), + datafusion::functions_aggregate::expr_fn::count(lit(1i64)), + datafusion::functions_aggregate::expr_fn::count(col("ts")), + ], + ) + .unwrap() + .project(vec![ + col("count(*)").alias("a"), + col("count(Int64(1))"), + col("count(t.ts)").alias("b"), + ]) + .unwrap() + .build() + .unwrap(); + + let config = ConfigOptions::default(); + let plan = CountWildcardToTimeIndexRule.analyze(plan, &config).unwrap(); + let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let expected = [ + "Projection: count(*) AS a, count(Int64(1)), count(t.ts) AS b", + " Projection: count(*), count(Int64(1)), count(t.ts)", + " Aggregate: groupBy=[[]], aggr=[[__count_merge(__count_state(t.ts)) AS count(*), __count_merge(__count_state(t.ts)) AS count(Int64(1)), __count_merge(__count_state(t.ts)) AS count(t.ts)]]", + " MergeScan [is_placeholder=false, remote_input=[", + "Aggregate: groupBy=[[]], aggr=[[__count_state(t.ts)]]", + " TableScan: t", + "]]", + ] + .join("\n"); + assert_eq!(expected, result.to_string()); +} + +/// After the `count(*)` rewrite, `count(ts)` and the casts of `ts` are different states with +/// the same schema name, as schema names hide casts. The output alias `t.ts` and the second +/// cast make sure a renamed state never takes a name that is already allocated. +#[test] +fn expand_step_aggr_count_wildcard_and_cast_time_index() { + init_default_ut_logging(); + let test_table = TestTable::table_with_name(0, "t".to_string()); + let table_source = Arc::new(DefaultTableSource::new(Arc::new( + DfTableProviderAdapter::new(test_table), + ))); + let plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![]) + .unwrap() + .aggregate( + Vec::::new(), + vec![ + datafusion::functions_aggregate::count::count_all(), + datafusion::functions_aggregate::expr_fn::count(datafusion_expr::cast( + col("ts"), + DataType::Int64, + )) + .alias("t.ts"), + datafusion::functions_aggregate::expr_fn::count(datafusion_expr::cast( + col("ts"), + DataType::Utf8, + )) + .alias("c"), + ], + ) + .unwrap() + .build() + .unwrap(); + + let config = ConfigOptions::default(); + let plan = CountWildcardToTimeIndexRule.analyze(plan, &config).unwrap(); + let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let expected = [ + "Projection: count(*), t.ts, c", + " Aggregate: groupBy=[[]], aggr=[[__count_merge(__count_state(t.ts)) AS count(*), __count_merge(__count_state(t.ts)_1) AS t.ts, __count_merge(__count_state(t.ts)_2) AS c]]", + " MergeScan [is_placeholder=false, remote_input=[", + "Aggregate: groupBy=[[]], aggr=[[__count_state(t.ts), __count_state(CAST(t.ts AS Int64)) AS __count_state(t.ts)_1, __count_state(CAST(t.ts AS Utf8)) AS __count_state(t.ts)_2]]", + " TableScan: t", + "]]", + ] + .join("\n"); + assert_eq!(expected, result.to_string()); +} + /// notice how `Limit` can still get expanded #[test] fn expand_part_col_aggr_limit() { diff --git a/tests/cases/standalone/common/aggregate/multi_regions.result b/tests/cases/standalone/common/aggregate/multi_regions.result index 79ae85fbc04..bdfc5ed2dd1 100644 --- a/tests/cases/standalone/common/aggregate/multi_regions.result +++ b/tests/cases/standalone/common/aggregate/multi_regions.result @@ -131,6 +131,32 @@ select substr(host, 6, 1) as g, count(*), sum(val) from t group by g order by g; | y | 2 | 6.0 | +---+----------+------------+ +-- count(*) and count(1) are rewritten to count(ts), sharing the same state column +select count(*), count(1), count(ts) from t; + ++----------+-----------------+-------------+ +| count(*) | count(Int64(1)) | count(t.ts) | ++----------+-----------------+-------------+ +| 4 | 4 | 4 | ++----------+-----------------+-------------+ + +select count(*) as a, count(ts) as b from t; + ++---+---+ +| a | b | ++---+---+ +| 4 | 4 | ++---+---+ + +-- the cast is hidden from the state column name +select count(*), count(ts::bigint) from t; + ++----------+-------------+ +| count(*) | count(t.ts) | ++----------+-------------+ +| 4 | 4 | ++----------+-------------+ + -- grouping sets are computed on the frontend, only the PostgreSQL dialect parses them -- SQLNESS PROTOCOL POSTGRES select host, idc, sum(val) from t group by grouping sets ((host, idc), (host)) order by host, idc; diff --git a/tests/cases/standalone/common/aggregate/multi_regions.sql b/tests/cases/standalone/common/aggregate/multi_regions.sql index 4f99e1ed465..9c8d8af4490 100644 --- a/tests/cases/standalone/common/aggregate/multi_regions.sql +++ b/tests/cases/standalone/common/aggregate/multi_regions.sql @@ -54,6 +54,14 @@ insert into t values -- group keys derived from the partition column span regions select substr(host, 6, 1) as g, count(*), sum(val) from t group by g order by g; +-- count(*) and count(1) are rewritten to count(ts), sharing the same state column +select count(*), count(1), count(ts) from t; + +select count(*) as a, count(ts) as b from t; + +-- the cast is hidden from the state column name +select count(*), count(ts::bigint) from t; + -- grouping sets are computed on the frontend, only the PostgreSQL dialect parses them -- SQLNESS PROTOCOL POSTGRES select host, idc, sum(val) from t group by grouping sets ((host, idc), (host)) order by host, idc; From cbfbad469f579c9cd031d44dbad1f730303156ad Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Thu, 15 Oct 2026 09:41:31 +0000 Subject: [PATCH 008/239] feat(meta): add affinity hooks (#9596) * feat(meta): add affinity hooks Signed-off-by: WenyXu * test(meta): cover load-based affinity fallback queries Signed-off-by: WenyXu * refactor(meta): borrow affinity requirements during selection Signed-off-by: WenyXu * fix(meta): remove unused affinity error variant Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- src/cmd/src/standalone.rs | 2 + src/common/meta/AGENTS.md | 5 + src/common/meta/src/affinity.rs | 254 +++++++++ src/common/meta/src/cluster.rs | 12 +- src/common/meta/src/ddl.rs | 5 + src/common/meta/src/ddl/create_table.rs | 13 +- src/common/meta/src/ddl/tests/create_table.rs | 140 +++++ src/common/meta/src/ddl_manager.rs | 14 + src/common/meta/src/lib.rs | 1 + src/common/meta/src/peer.rs | 5 +- src/common/meta/src/test_util.rs | 2 + src/meta-srv/AGENTS.md | 6 + src/meta-srv/src/metasrv.rs | 13 + src/meta-srv/src/metasrv/builder.rs | 109 +++- src/meta-srv/src/peer.rs | 217 ++++++- src/meta-srv/src/procedure/utils.rs | 2 + src/meta-srv/src/region/supervisor.rs | 2 +- src/meta-srv/src/selector.rs | 17 +- src/meta-srv/src/selector/common.rs | 539 +++++++++++++++++- src/meta-srv/src/selector/lease_based.rs | 48 +- src/meta-srv/src/selector/load_based.rs | 51 +- src/meta-srv/src/selector/round_robin.rs | 98 +++- src/meta-srv/src/selector/test_utils.rs | 9 +- tests-integration/src/standalone.rs | 2 + tests-integration/tests/region_migration.rs | 2 +- 25 files changed, 1475 insertions(+), 93 deletions(-) create mode 100644 src/common/meta/src/affinity.rs diff --git a/src/cmd/src/standalone.rs b/src/cmd/src/standalone.rs index ca0f3b74057..278afb7f1da 100644 --- a/src/cmd/src/standalone.rs +++ b/src/cmd/src/standalone.rs @@ -681,6 +681,8 @@ impl StartCommand { )); let ddl_context = DdlContext { + affinity_resolver: None, + peer_alloc_context_provider: None, node_manager: node_manager.clone(), cache_invalidator: layered_cache_registry.clone(), memory_region_keeper: Arc::new(MemoryRegionKeeper::default()), diff --git a/src/common/meta/AGENTS.md b/src/common/meta/AGENTS.md index 6c9710b9e45..69d76d3c5ca 100644 --- a/src/common/meta/AGENTS.md +++ b/src/common/meta/AGENTS.md @@ -11,6 +11,7 @@ and durable DDL procedures. Metasrv process and service wiring live in | Area | Path | Entry point | | --- | --- | --- | +| Affinity | `src/common/meta/src/affinity.rs` | Requirements, reference resolution, node matching, and allocation-context provider | | Metadata model | `src/common/meta/src/key.rs`, `src/common/meta/src/key/` | Typed keys, values, and `TableMetadataManager` | | KV storage | `src/common/meta/src/kv_backend.rs`, `src/common/meta/src/kv_backend/` | `KvBackend`, transactions, memory/etcd/RDS backends | | DDL procedures | `src/common/meta/src/ddl.rs`, `src/common/meta/src/ddl/`, `src/common/meta/src/ddl_manager.rs` | Durable DDL state machines and task dispatch | @@ -31,6 +32,10 @@ and durable DDL procedures. Metasrv process and service wiring live in `src/table/src/requests.rs` — classification, validation, logical-table eligibility, and the metadata-only alter flow come with it. Do not add per-option special cases to the alter conversion or DDL routing. +- Create-table preparation obtains allocation inputs from the optional query + provider in `DdlContext`. Install providers before registering procedure + loaders so restored procedures use the same runtime dependencies. Affinity + requirements and fallback policy are runtime inputs, not persisted fields. - `KvBackend` behavior changes should extend the shared tests in `src/common/meta/src/kv_backend/test.rs`. diff --git a/src/common/meta/src/affinity.rs b/src/common/meta/src/affinity.rs new file mode 100644 index 00000000000..2bcb0e934f0 --- /dev/null +++ b/src/common/meta/src/affinity.rs @@ -0,0 +1,254 @@ +// 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. + +//! Resolve reference attributes into requirements, then evaluate candidates. + +use std::collections::HashMap; +use std::sync::Arc; + +use crate::DatanodeId; +use crate::cluster::NodeInfo; +use crate::error::Result; +use crate::peer::PeerAllocContext; +use crate::rpc::ddl::QueryContext; + +/// Named attributes used to express affinity requirements. +pub type AffinityAttributes = HashMap; + +/// Identifies the node whose attributes define affinity requirements. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub enum AffinityReference { + /// Uses the datanode's last-known attributes, even if its lease has expired. + Datanode(DatanodeId), + /// Uses the frontend identified by its advertised address. + Frontend(String), +} + +/// Behavior when affinity cannot be satisfied. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub enum AffinityFallback { + /// Returns an error instead of selecting without affinity. + #[default] + Error, + /// Allows selection among eligible nodes without affinity. + Unconstrained, +} + +/// Builds allocation inputs from the query that initiated an operation. +#[async_trait::async_trait] +pub trait PeerAllocContextProvider: Send + Sync { + /// Resolves query-specific allocation inputs using metadata, without selecting peers. + async fn build(&self, query_context: &QueryContext) -> Result; +} + +/// Shared allocation-context provider. +pub type PeerAllocContextProviderRef = Arc; + +/// Nonempty attributes a candidate must match. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AffinityRequirements { + attributes: AffinityAttributes, +} + +impl AffinityRequirements { + /// Returns no requirements when the supplied attributes are empty. + pub fn new(attributes: AffinityAttributes) -> Option { + (!attributes.is_empty()).then_some(Self { attributes }) + } + + /// Returns the attributes that candidates are evaluated against. + pub fn attributes(&self) -> &AffinityAttributes { + &self.attributes + } +} + +/// Resolves references using current metadata, without selecting candidates. +#[async_trait::async_trait] +pub trait AffinityResolver: Send + Sync { + /// Returns one result per reference, in order. + /// None means requirements are unavailable; Err means resolution failed. + async fn resolve( + &self, + references: &[AffinityReference], + ) -> Result>>; +} + +/// Shared requirement resolver. +pub type AffinityResolverRef = Arc; + +/// Evaluates candidate nodes without I/O or scheduling decisions. +pub trait Affinity: Send + Sync { + /// Evaluates a node snapshot; implementations choose which fields to inspect. + fn matches(&self, requirements: &AffinityRequirements, candidate: &NodeInfo) -> bool; +} + +/// Accepts every candidate without imposing affinity constraints. +pub struct NoopAffinity; + +impl Affinity for NoopAffinity { + fn matches(&self, _: &AffinityRequirements, _: &NodeInfo) -> bool { + true + } +} + +/// Shared affinity evaluator. +pub type AffinityRef = Arc; + +#[cfg(test)] +mod tests { + use std::sync::RwLock; + + use super::*; + use crate::cluster::{DatanodeStatus, NodeStatus}; + use crate::peer::Peer; + + struct NodeAffinityResolver { + attributes: RwLock>, + } + + #[async_trait::async_trait] + impl AffinityResolver for NodeAffinityResolver { + async fn resolve( + &self, + references: &[AffinityReference], + ) -> Result>> { + let snapshot = self.attributes.read().unwrap(); + Ok(references + .iter() + .map(|reference| { + snapshot + .get(reference) + .cloned() + .and_then(AffinityRequirements::new) + }) + .collect()) + } + } + + struct ExactAffinity; + + impl Affinity for ExactAffinity { + fn matches(&self, requirements: &AffinityRequirements, candidate: &NodeInfo) -> bool { + requirements + .attributes() + .iter() + .all(|(key, expected)| candidate.env_vars.get(key) == Some(expected)) + } + } + + fn zone(name: &str) -> AffinityAttributes { + HashMap::from([("AZ".to_string(), name.to_string())]) + } + + fn node(id: DatanodeId, env_vars: AffinityAttributes) -> NodeInfo { + NodeInfo { + peer: Peer::empty(id), + last_activity_ts: 0, + status: NodeStatus::Datanode(DatanodeStatus { + rcus: 0, + wcus: 0, + leader_regions: 0, + follower_regions: 0, + workloads: Default::default(), + }), + version: String::new(), + git_commit: String::new(), + start_time_ms: 0, + total_cpu_millicores: 0, + total_memory_bytes: 0, + cpu_usage_millicores: 0, + memory_usage_bytes: 0, + hostname: String::new(), + env_vars, + } + } + + #[tokio::test] + async fn test_resolve_reference_and_match_candidates() { + let resolver = NodeAffinityResolver { + attributes: RwLock::new(HashMap::from([(AffinityReference::Datanode(1), zone("a"))])), + }; + let references = [AffinityReference::Datanode(1)]; + let resolved = resolver.resolve(&references).await.unwrap(); + assert_eq!(resolved, vec![AffinityRequirements::new(zone("a"))]); + let requirements = resolved[0].as_ref().unwrap(); + let candidates = [ + node(2, zone("a")), + node(3, zone("b")), + node(4, HashMap::new()), + ]; + assert_eq!( + candidates + .iter() + .map(|node| (node.peer.id, ExactAffinity.matches(requirements, node))) + .collect::>(), + vec![(2, true), (3, false), (4, false),] + ); + + assert!( + candidates + .iter() + .all(|node| NoopAffinity.matches(requirements, node)) + ); + + // A new allocation attempt resolves current attributes rather than a saved requirement. + resolver + .attributes + .write() + .unwrap() + .insert(AffinityReference::Datanode(1), zone("b")); + let resolved = resolver.resolve(&references).await.unwrap(); + let requirements = resolved[0].as_ref().unwrap(); + let matched: Vec<_> = candidates + .iter() + .filter(|node| ExactAffinity.matches(requirements, node)) + .map(|node| node.peer.id) + .collect(); + assert_eq!(matched, vec![3]); + } + + #[tokio::test] + async fn test_resolution_preserves_order_and_unavailable_requirements() { + let resolver = NodeAffinityResolver { + attributes: RwLock::new(HashMap::from([ + (AffinityReference::Datanode(1), zone("a")), + (AffinityReference::Datanode(2), HashMap::new()), + ( + AffinityReference::Frontend("frontend:4001".into()), + zone("b"), + ), + ])), + }; + let references = [ + AffinityReference::Frontend("missing:4001".into()), + AffinityReference::Datanode(1), + AffinityReference::Frontend("frontend:4001".into()), + AffinityReference::Datanode(2), + AffinityReference::Datanode(3), + AffinityReference::Frontend("frontend:4001".into()), + ]; + assert_eq!( + resolver.resolve(&references).await.unwrap(), + vec![ + None, + AffinityRequirements::new(zone("a")), + AffinityRequirements::new(zone("b")), + None, + None, + AffinityRequirements::new(zone("b")), + ] + ); + assert!(AffinityRequirements::new(HashMap::new()).is_none()); + } +} diff --git a/src/common/meta/src/cluster.rs b/src/common/meta/src/cluster.rs index 149e5a059d5..4e7ffb7f920 100644 --- a/src/common/meta/src/cluster.rs +++ b/src/common/meta/src/cluster.rs @@ -106,7 +106,7 @@ fn calculate_node_id(addr: &str) -> u64 { } /// The information of a node in the cluster. -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Clone, Serialize, Deserialize)] pub struct NodeInfo { /// The peer information. [node_id, address] pub peer: Peer, @@ -148,7 +148,7 @@ pub enum Role { Metasrv, } -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Clone, Serialize, Deserialize)] pub enum NodeStatus { Datanode(DatanodeStatus), Frontend(FrontendStatus), @@ -171,7 +171,7 @@ impl NodeStatus { } /// The status of a datanode. -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Clone, Serialize, Deserialize)] pub struct DatanodeStatus { /// The read capacity units during this period. pub rcus: i64, @@ -186,7 +186,7 @@ pub struct DatanodeStatus { } /// The status of a frontend. -#[derive(Debug, Default, Serialize, Deserialize)] +#[derive(Debug, Clone, Default, Serialize, Deserialize)] pub struct FrontendStatus { /// The workloads of the frontend. #[serde(default)] @@ -194,7 +194,7 @@ pub struct FrontendStatus { } /// The status of a flownode. -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Clone, Serialize, Deserialize)] pub struct FlownodeStatus { /// The workloads of the flownode. #[serde(default)] @@ -202,7 +202,7 @@ pub struct FlownodeStatus { } /// The status of a metasrv. -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Clone, Serialize, Deserialize)] pub struct MetasrvStatus { pub is_leader: bool, } diff --git a/src/common/meta/src/ddl.rs b/src/common/meta/src/ddl.rs index f70ba2b3861..3eea16fcfa5 100644 --- a/src/common/meta/src/ddl.rs +++ b/src/common/meta/src/ddl.rs @@ -18,6 +18,7 @@ use std::time::Duration; use store_api::storage::{RegionId, TableId}; use crate::DatanodeId; +use crate::affinity::{AffinityResolverRef, PeerAllocContextProviderRef}; use crate::cache_invalidator::CacheInvalidatorRef; use crate::ddl::flow_meta::FlowMetadataAllocatorRef; use crate::ddl::table_meta::TableMetadataAllocatorRef; @@ -105,6 +106,10 @@ impl RegionFailureDetectorController for NoopRegionFailureDetectorControl { /// The context of ddl. #[derive(Clone)] pub struct DdlContext { + /// Optional resolver for automatic allocation affinity. + pub affinity_resolver: Option, + /// Builds allocation inputs from the initiating query. + pub peer_alloc_context_provider: Option, /// Sends querying and requests to nodes. pub node_manager: NodeManagerRef, /// Cache invalidation. diff --git a/src/common/meta/src/ddl/create_table.rs b/src/common/meta/src/ddl/create_table.rs index ce3e1aa060a..122b0f24030 100644 --- a/src/common/meta/src/ddl/create_table.rs +++ b/src/common/meta/src/ddl/create_table.rs @@ -165,7 +165,10 @@ impl CreateTableProcedure { return Ok(Status::done_with_output(table_id)); } - self.data.state = CreateTableState::DatanodeCreateRegions; + let alloc_context = match &self.context.peer_alloc_context_provider { + Some(provider) => provider.build(&self.data.query_context).await?, + None => PeerAllocContext::default(), + }; let TableMetadata { table_id, table_route, @@ -173,14 +176,10 @@ impl CreateTableProcedure { } = self .context .table_metadata_allocator - .create_with_context( - &self.data.task, - &PeerAllocContext { - extensions: self.data.query_context.extensions.clone(), - }, - ) + .create_with_context(&self.data.task, &alloc_context) .await?; self.set_allocated_metadata(table_id, table_route, region_wal_options); + self.data.state = CreateTableState::DatanodeCreateRegions; Ok(Status::executing(true)) } diff --git a/src/common/meta/src/ddl/tests/create_table.rs b/src/common/meta/src/ddl/tests/create_table.rs index c97e61926c0..d39346e02da 100644 --- a/src/common/meta/src/ddl/tests/create_table.rs +++ b/src/common/meta/src/ddl/tests/create_table.rs @@ -398,3 +398,143 @@ async fn test_memory_region_keeper_guard_dropped_on_procedure_done() { .contains(datanode_id, region_id) ); } + +#[tokio::test] +async fn test_on_prepare_uses_peer_alloc_context_provider() { + use std::sync::Mutex; + + use crate::affinity::{AffinityRequirements, PeerAllocContextProvider}; + use crate::ddl::table_meta::TableMetadataAllocator; + use crate::error; + use crate::peer::{PeerAllocContext, PeerAllocator}; + use crate::rpc::ddl::QueryContext; + + #[derive(Clone, Copy, PartialEq, Eq)] + enum ProviderMode { + Absent, + Available, + Error, + } + + struct Provider { + mode: ProviderMode, + calls: Mutex>, + } + + #[async_trait::async_trait] + impl PeerAllocContextProvider for Provider { + async fn build(&self, query_context: &QueryContext) -> Result { + self.calls.lock().unwrap().push(query_context.clone()); + if self.mode == ProviderMode::Error { + return error::UnexpectedSnafu { + err_msg: "allocation context lookup failed", + } + .fail(); + } + Ok(PeerAllocContext { + affinity: AffinityRequirements::new(HashMap::from([("pool".into(), "a".into())])), + }) + } + } + + #[derive(Default)] + struct Allocator(Mutex>); + + #[async_trait::async_trait] + impl PeerAllocator for Allocator { + async fn alloc(&self, num: usize, ctx: &PeerAllocContext) -> Result> { + self.0.lock().unwrap().push(ctx.clone()); + Ok(vec![Peer::empty(2); num]) + } + } + + for (mode, table_exists) in [ + (ProviderMode::Available, false), + (ProviderMode::Absent, false), + (ProviderMode::Error, false), + (ProviderMode::Error, true), + ] { + let provider = Arc::new(Provider { + mode, + calls: Mutex::new(Vec::new()), + }); + let allocator = Arc::new(Allocator::default()); + let mut context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); + if mode != ProviderMode::Absent { + context.peer_alloc_context_provider = Some(provider.clone()); + } + context.table_metadata_allocator = Arc::new(TableMetadataAllocator::with_peer_allocator( + context.table_metadata_allocator.table_id_allocator(), + context.table_metadata_allocator.wal_options_allocator(), + allocator.clone(), + )); + let mut task = test_create_table_task("provider_table"); + task.create_table.create_if_not_exists = true; + if table_exists { + context + .table_metadata_manager + .create_table_metadata( + task.table_info.clone(), + TableRouteValue::physical(vec![]), + HashMap::new(), + ) + .await + .unwrap(); + } + let query_context = QueryContext { + current_catalog: "test_catalog".into(), + current_schema: "test_schema".into(), + timezone: "UTC".into(), + extensions: HashMap::from([("custom-key".into(), "custom-value".into())]), + channel: 1, + snapshot_seqs: HashMap::from([(1, 10)]), + sst_min_sequences: HashMap::from([(1, 2)]), + }; + let mut procedure = + CreateTableProcedure::new_with_query_context(task, query_context.clone(), context) + .unwrap(); + let result = procedure.on_prepare().await; + if table_exists { + assert_matches!(result.unwrap(), Status::Done { .. }); + assert!(provider.calls.lock().unwrap().is_empty()); + assert!(allocator.0.lock().unwrap().is_empty()); + continue; + } + let expected_calls = if mode == ProviderMode::Absent { + vec![] + } else { + vec![query_context] + }; + assert_eq!(*provider.calls.lock().unwrap(), expected_calls); + if mode == ProviderMode::Error { + let error = result.unwrap_err(); + assert_matches!(error, Error::Unexpected { .. }); + assert!( + error + .to_string() + .contains("allocation context lookup failed") + ); + assert_matches!(procedure.data.state, CreateTableState::Prepare); + assert!(allocator.0.lock().unwrap().is_empty()); + } else { + result.unwrap(); + assert_matches!( + procedure.data.state, + CreateTableState::DatanodeCreateRegions + ); + let calls = allocator.0.lock().unwrap(); + assert_eq!(calls.len(), 1); + let expected = if mode == ProviderMode::Available { + PeerAllocContext { + affinity: AffinityRequirements::new(HashMap::from([( + "pool".into(), + "a".into(), + )])), + } + } else { + PeerAllocContext::default() + }; + assert_eq!(calls[0].affinity, expected.affinity); + } + } +} diff --git a/src/common/meta/src/ddl_manager.rs b/src/common/meta/src/ddl_manager.rs index eddb4148767..0a1abde30f7 100644 --- a/src/common/meta/src/ddl_manager.rs +++ b/src/common/meta/src/ddl_manager.rs @@ -33,6 +33,7 @@ use snafu::{OptionExt, ResultExt, ensure}; use store_api::storage::{RegionId, TableId}; use table::table_name::TableName; +use crate::affinity::PeerAllocContextProviderRef; use crate::ddl::alter_database::AlterDatabaseProcedure; use crate::ddl::alter_logical_tables::AlterLogicalTablesProcedure; use crate::ddl::alter_table::{AlterTableProcedure, RegionRouteChanged, only_sets_skip_wal}; @@ -300,6 +301,15 @@ impl DdlManager { self } + /// Sets the provider used to build allocation inputs from queries. + pub fn with_peer_alloc_context_provider( + mut self, + provider: PeerAllocContextProviderRef, + ) -> Self { + self.ddl_context.peer_alloc_context_provider = Some(provider); + self + } + pub fn with_create_database_metadata_committer( mut self, committer: CreateDatabaseMetadataCommitterRef, @@ -1875,6 +1885,8 @@ mod tests { let ddl_manager = DdlManager::new( DdlContext { + affinity_resolver: None, + peer_alloc_context_provider: None, node_manager: Arc::new(DummyDatanodeManager), cache_invalidator: Arc::new(DummyCacheInvalidator), table_metadata_manager, @@ -1941,6 +1953,8 @@ mod tests { DdlManager::new( DdlContext { + affinity_resolver: None, + peer_alloc_context_provider: None, node_manager: Arc::new(DummyDatanodeManager), cache_invalidator: Arc::new(DummyCacheInvalidator), table_metadata_manager, diff --git a/src/common/meta/src/lib.rs b/src/common/meta/src/lib.rs index ef5109dc03e..22fea83563e 100644 --- a/src/common/meta/src/lib.rs +++ b/src/common/meta/src/lib.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +pub mod affinity; pub mod cache; pub mod cache_invalidator; pub mod cluster; diff --git a/src/common/meta/src/peer.rs b/src/common/meta/src/peer.rs index 18ffd92f8bb..0e1f53ac462 100644 --- a/src/common/meta/src/peer.rs +++ b/src/common/meta/src/peer.rs @@ -12,12 +12,12 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::HashMap; use std::sync::Arc; pub use api::v1::meta::Peer; use api::v1::meta::heartbeat_request::NodeWorkloads; +use crate::affinity::AffinityRequirements; use crate::cluster::NodeInfo; use crate::error::Error; use crate::{DatanodeId, FlownodeId}; @@ -80,7 +80,8 @@ pub type PeerDiscoveryRef = Arc; #[derive(Debug, Clone, Default)] pub struct PeerAllocContext { - pub extensions: HashMap, + /// Resolved requirements passed to node selection. + pub affinity: Option, } /// [`PeerAllocator`] allocates [`Peer`]s for creating region or flow. diff --git a/src/common/meta/src/test_util.rs b/src/common/meta/src/test_util.rs index 618603c13e5..e574161d367 100644 --- a/src/common/meta/src/test_util.rs +++ b/src/common/meta/src/test_util.rs @@ -198,6 +198,8 @@ pub fn new_ddl_context_with_kv_backend( .build(), ))); DdlContext { + affinity_resolver: None, + peer_alloc_context_provider: None, node_manager, cache_invalidator: Arc::new(DummyCacheInvalidator), memory_region_keeper: Arc::new(MemoryRegionKeeper::new()), diff --git a/src/meta-srv/AGENTS.md b/src/meta-srv/AGENTS.md index 5f9f573978c..863b2a2aab4 100644 --- a/src/meta-srv/AGENTS.md +++ b/src/meta-srv/AGENTS.md @@ -64,6 +64,12 @@ API under `src/meta-srv/src/service/admin/`. must be idempotent for crash recovery. - **Cache invalidation**: new metadata types need matching invalidation in `cache_invalidator.rs` and the heartbeat publish handler. +- **Affinity**: `common-meta/affinity.rs` defines requirements, matching, batch + resolution, and the query allocation-context provider. The selector factory + optionally supplies a resolver for allocation inputs. DDL configurators can + install a query provider before procedure loaders are registered. Each selector + allows one unconstrained retry on candidate shortage only when explicitly + configured. Procedure state is unchanged. - **Selector stats fields**: keep `selector/` implementations in sync with the stats they consume. diff --git a/src/meta-srv/src/metasrv.rs b/src/meta-srv/src/metasrv.rs index 86aee9f8b82..f9b7a8dd5d7 100644 --- a/src/meta-srv/src/metasrv.rs +++ b/src/meta-srv/src/metasrv.rs @@ -26,6 +26,7 @@ use common_base::readable_size::ReadableSize; use common_config::{Configurable, DEFAULT_DATA_HOME}; use common_event_recorder::EventRecorderOptions; use common_greptimedb_telemetry::GreptimeDBTelemetryTask; +use common_meta::affinity::{AffinityFallback, AffinityRef, AffinityResolverRef}; use common_meta::cache_invalidator::CacheInvalidatorRef; use common_meta::ddl::allocator::resource_id::ResourceIdAllocatorRef; use common_meta::ddl_manager::DdlManagerRef; @@ -486,6 +487,10 @@ impl Display for SelectTarget { #[derive(Clone)] pub struct SelectorContext { pub peer_discovery: PeerDiscoveryRef, + /// Evaluator shared by datanode and flownode selection. + pub affinity: AffinityRef, + /// Whether selection may retry without affinity when candidates are insufficient. + pub affinity_fallback: AffinityFallback, } pub type SelectorRef = Arc>>; @@ -507,6 +512,14 @@ pub struct SelectorFactoryContext { /// Builds the final datanode selector metasrv should use. pub trait SelectorFactory: Send + Sync { fn build(&self, ctx: SelectorFactoryContext) -> SelectorRef; + + /// Builds the resolver shared by allocation and failover callers. + fn build_affinity_resolver( + &self, + _meta_peer_client: MetaPeerClientRef, + ) -> Option { + None + } } /// Shared selector factory plugin registered through [`common_base::Plugins`]. diff --git a/src/meta-srv/src/metasrv/builder.rs b/src/meta-srv/src/metasrv/builder.rs index 57e7deb3db6..96cd2a8adb5 100644 --- a/src/meta-srv/src/metasrv/builder.rs +++ b/src/meta-srv/src/metasrv/builder.rs @@ -22,6 +22,7 @@ use client::inserter::InsertOptions; use common_base::Plugins; use common_catalog::consts::{MIN_USER_FLOW_ID, MIN_USER_TABLE_ID}; use common_event_recorder::{DEFAULT_COMPACTION_TIME_WINDOW, EventRecorderImpl, EventRecorderRef}; +use common_meta::affinity::{AffinityFallback, AffinityRef, NoopAffinity}; use common_meta::ddl::flow_meta::FlowMetadataAllocator; use common_meta::ddl::table_meta::{TableMetadataAllocator, TableMetadataAllocatorRef}; use common_meta::ddl::{ @@ -66,7 +67,8 @@ use crate::handler::region_lease_handler::{CustomizedRegionLeaseRenewerRef, Regi use crate::handler::{HeartbeatHandlerGroupBuilder, HeartbeatMailbox, Pushers}; use crate::metasrv::{ ElectionRef, FLOW_ID_SEQ, METASRV_DATA_DIR, Metasrv, MetasrvInfo, MetasrvOptions, - RegionStatAwareSelectorRef, SelectTarget, SelectorContext, SelectorRef, TABLE_ID_SEQ, + RegionStatAwareSelectorRef, SelectTarget, SelectorContext, SelectorFactoryRef, SelectorRef, + TABLE_ID_SEQ, }; use crate::peer::MetasrvPeerAllocator; use crate::procedure::region_migration::DefaultContextFactory; @@ -274,6 +276,14 @@ impl MetasrvBuilder { )); let selector_ctx = SelectorContext { + affinity_fallback: plugins + .as_ref() + .and_then(|plugins| plugins.get::()) + .unwrap_or_default(), + affinity: plugins + .as_ref() + .and_then(|plugins| plugins.get::()) + .unwrap_or_else(|| Arc::new(NoopAffinity)), peer_discovery: meta_peer_client.clone(), }; @@ -385,6 +395,10 @@ impl MetasrvBuilder { ), )); region_migration_manager.try_start()?; + let affinity_resolver = plugins + .as_ref() + .and_then(|plugins| plugins.get::()) + .and_then(|factory| factory.build_affinity_resolver(meta_peer_client.clone())); let region_supervisor_selector = plugins .as_ref() .and_then(|plugins| plugins.get::()); @@ -422,6 +436,8 @@ impl MetasrvBuilder { let topic_stats_registry = Arc::new(TopicStatsRegistry::default()); let ddl_context = DdlContext { + peer_alloc_context_provider: None, + affinity_resolver, node_manager: node_manager.clone(), cache_invalidator: cache_invalidator.clone(), memory_region_keeper: memory_region_keeper.clone(), @@ -795,10 +811,99 @@ mod tests { options.gc.enable = false; options.gc.max_concurrent_tables = 0; - MetasrvBuilder::new() + let metasrv = MetasrvBuilder::new() .options(options) .build() .await .unwrap(); + let context = metasrv.ddl_manager().create_context(); + assert!(context.affinity_resolver.is_none()); + assert!(context.peer_alloc_context_provider.is_none()); + } + + #[tokio::test] + async fn test_builder_injects_affinity_components() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use common_meta::affinity::{ + Affinity, AffinityRef, AffinityReference, AffinityRequirements, AffinityResolver, + AffinityResolverRef, + }; + use common_meta::cluster::NodeInfo; + use common_meta::error::Result as MetaResult; + + use crate::metasrv::{SelectorFactory, SelectorFactoryContext}; + + struct TestAffinity; + + impl Affinity for TestAffinity { + fn matches(&self, _: &AffinityRequirements, _: &NodeInfo) -> bool { + true + } + } + + struct TestResolver; + + #[async_trait::async_trait] + impl AffinityResolver for TestResolver { + async fn resolve( + &self, + references: &[AffinityReference], + ) -> MetaResult>> { + Ok(vec![None; references.len()]) + } + } + + struct RecordingFactory { + calls: Arc, + resolver: AffinityResolverRef, + } + + impl SelectorFactory for RecordingFactory { + fn build_affinity_resolver(&self, _: MetaPeerClientRef) -> Option { + self.calls.fetch_add(1, Ordering::Relaxed); + Some(self.resolver.clone()) + } + + fn build(&self, ctx: SelectorFactoryContext) -> SelectorRef { + ctx.base_selector + } + } + + let calls = Arc::new(AtomicUsize::new(0)); + let resolver: AffinityResolverRef = Arc::new(TestResolver); + let affinity: AffinityRef = Arc::new(TestAffinity); + let plugins = Plugins::new(); + plugins.insert(affinity.clone()); + plugins.insert(AffinityFallback::Unconstrained); + plugins.insert(Arc::new(RecordingFactory { + calls: calls.clone(), + resolver: resolver.clone(), + }) as SelectorFactoryRef); + let metasrv = MetasrvBuilder::new() + .options(MetasrvOptions { + enable_telemetry: false, + ..Default::default() + }) + .plugins(plugins) + .build() + .await + .unwrap(); + + assert_eq!( + metasrv.selector_ctx().affinity_fallback, + AffinityFallback::Unconstrained + ); + assert_eq!(calls.load(Ordering::Relaxed), 1); + assert!(Arc::ptr_eq(&metasrv.selector_ctx().affinity, &affinity)); + assert!(Arc::ptr_eq( + metasrv + .ddl_manager() + .create_context() + .affinity_resolver + .as_ref() + .unwrap(), + &resolver + )); } } diff --git a/src/meta-srv/src/peer.rs b/src/meta-srv/src/peer.rs index 9572006182b..8843d58ad87 100644 --- a/src/meta-srv/src/peer.rs +++ b/src/meta-srv/src/peer.rs @@ -21,8 +21,8 @@ use common_meta::peer::{Peer, PeerAllocContext, PeerAllocator}; use snafu::{ResultExt, ensure}; use crate::discovery::utils::accept_ingest_workload; -use crate::error::{Result, TooManyPartitionsSnafu}; -use crate::metasrv::{SelectorContext, SelectorRef}; +use crate::error::{self, Result, TooManyPartitionsSnafu}; +use crate::metasrv::{SelectTarget, SelectorContext, SelectorRef}; use crate::selector::SelectorOptions; pub struct MetasrvPeerAllocator { @@ -67,18 +67,23 @@ impl MetasrvPeerAllocator { ); } - self.selector - .select( - &self.ctx, - SelectorOptions { - min_required_items, - allow_duplication: true, - exclude_peer_ids: HashSet::new(), - workload_filter: Some(accept_ingest_workload), - extensions: alloc_context.extensions.clone(), - }, - ) - .await + let options = SelectorOptions { + affinity: alloc_context.affinity.as_ref(), + min_required_items, + allow_duplication: true, + exclude_peer_ids: HashSet::new(), + workload_filter: Some(accept_ingest_workload), + }; + let peers = self.selector.select(&self.ctx, options).await?; + ensure!( + peers.len() >= min_required_items, + error::NoEnoughAvailableNodeSnafu { + required: min_required_items, + available: peers.len(), + select_target: SelectTarget::Datanode, + } + ); + Ok(peers) } } @@ -91,3 +96,187 @@ impl PeerAllocator for MetasrvPeerAllocator { .context(ExternalSnafu) } } + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::ptr; + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + + use common_meta::affinity::{ + Affinity, AffinityFallback, AffinityRef, AffinityRequirements, NoopAffinity, + }; + use common_meta::cluster::NodeInfo; + + use super::*; + use crate::selector::Selector; + use crate::selector::round_robin::RoundRobinSelector; + use crate::test_util::{create_meta_peer_client, put_datanodes}; + + struct OnlySecondNode; + + impl Affinity for OnlySecondNode { + fn matches(&self, _: &AffinityRequirements, node: &NodeInfo) -> bool { + node.peer.id == 2 + } + } + + #[derive(Clone, Copy)] + enum Outcome { + Base, + BackendError, + Short, + } + + struct CountingSelector { + base: RoundRobinSelector, + calls: AtomicUsize, + requirements_address: AtomicUsize, + outcome: Outcome, + } + + #[async_trait] + impl Selector for CountingSelector { + type Context = SelectorContext; + type Output = Vec; + + async fn select( + &self, + ctx: &SelectorContext, + opts: SelectorOptions<'_>, + ) -> Result { + self.calls.fetch_add(1, Ordering::Relaxed); + self.requirements_address.store( + opts.affinity + .map_or(0, |requirements| ptr::from_ref(requirements) as usize), + Ordering::Relaxed, + ); + assert!(opts.allow_duplication); + assert!(opts.exclude_peer_ids.is_empty()); + assert!(opts.workload_filter.is_some()); + assert_eq!(opts.min_required_items, 2); + match self.outcome { + Outcome::Base => self.base.select(ctx, opts).await, + Outcome::BackendError => error::UnexpectedSnafu { + violated: "selection backend failed", + } + .fail(), + Outcome::Short => Ok(vec![Peer::empty(1)]), + } + } + } + + #[tokio::test] + async fn test_allocator_delegates_fallback_and_validates_results() { + let discovery = create_meta_peer_client(); + put_datanodes(&discovery, vec![Peer::empty(1)]).await; + for (has_requirements, fallback, noop, outcome, expected_calls, succeeds) in [ + ( + true, + AffinityFallback::Error, + false, + Outcome::Base, + 1, + false, + ), + ( + true, + AffinityFallback::Unconstrained, + false, + Outcome::Base, + 1, + true, + ), + (true, AffinityFallback::Error, true, Outcome::Base, 1, true), + ( + false, + AffinityFallback::Unconstrained, + false, + Outcome::Base, + 1, + true, + ), + ( + true, + AffinityFallback::Unconstrained, + false, + Outcome::BackendError, + 1, + false, + ), + ( + true, + AffinityFallback::Error, + false, + Outcome::Short, + 1, + false, + ), + ( + true, + AffinityFallback::Unconstrained, + false, + Outcome::Short, + 1, + false, + ), + ( + false, + AffinityFallback::Unconstrained, + false, + Outcome::Short, + 1, + false, + ), + ] { + let selector = Arc::new(CountingSelector { + base: RoundRobinSelector::default(), + calls: AtomicUsize::new(0), + requirements_address: AtomicUsize::new(0), + outcome, + }); + let affinity: AffinityRef = if noop { + Arc::new(NoopAffinity) + } else { + Arc::new(OnlySecondNode) + }; + let allocator = MetasrvPeerAllocator::new( + SelectorContext { + affinity_fallback: fallback, + peer_discovery: discovery.clone(), + affinity, + }, + selector.clone(), + ); + let alloc_context = PeerAllocContext { + affinity: has_requirements.then(|| { + AffinityRequirements::new(HashMap::from([("pool".into(), "a".into())])).unwrap() + }), + }; + let result = allocator.alloc(2, &alloc_context).await; + assert_eq!( + selector.requirements_address.load(Ordering::Relaxed), + alloc_context + .affinity + .as_ref() + .map_or(0, |requirements| { ptr::from_ref(requirements) as usize }), + ); + assert_eq!(selector.calls.load(Ordering::Relaxed), expected_calls); + if succeeds { + assert_eq!(result.unwrap(), vec![Peer::empty(1); 2]); + } else { + match outcome { + Outcome::BackendError => assert!(matches!( + result.unwrap_err(), + error::Error::Unexpected { .. } + )), + _ => assert!(matches!( + result.unwrap_err(), + error::Error::NoEnoughAvailableNode { .. } + )), + } + } + } + } +} diff --git a/src/meta-srv/src/procedure/utils.rs b/src/meta-srv/src/procedure/utils.rs index 49964de24b7..1f975b2f9c0 100644 --- a/src/meta-srv/src/procedure/utils.rs +++ b/src/meta-srv/src/procedure/utils.rs @@ -535,6 +535,8 @@ pub mod test_data { Arc::new(SequenceBuilder::new("test", kv_backend).build()), )); DdlContext { + peer_alloc_context_provider: None, + affinity_resolver: None, node_manager, cache_invalidator: Arc::new(MetasrvCacheInvalidator::new( mailbox, diff --git a/src/meta-srv/src/region/supervisor.rs b/src/meta-srv/src/region/supervisor.rs index 99c68ac23c8..93624a61892 100644 --- a/src/meta-srv/src/region/supervisor.rs +++ b/src/meta-srv/src/region/supervisor.rs @@ -667,7 +667,7 @@ impl RegionSupervisor { allow_duplication: true, exclude_peer_ids, workload_filter: Some(accept_ingest_workload), - extensions: Default::default(), + affinity: None, }; let peers = selector.select(&self.selector_context, opt).await?; ensure!( diff --git a/src/meta-srv/src/selector.rs b/src/meta-srv/src/selector.rs index a02d294d822..e5c643b7b7c 100644 --- a/src/meta-srv/src/selector.rs +++ b/src/meta-srv/src/selector.rs @@ -20,9 +20,10 @@ pub mod round_robin; pub(crate) mod test_utils; pub mod weight_compute; pub mod weighted_choose; -use std::collections::{HashMap, HashSet}; +use std::collections::HashSet; use api::v1::meta::heartbeat_request::NodeWorkloads; +use common_meta::affinity::AffinityRequirements; use serde::{Deserialize, Serialize}; use store_api::storage::RegionId; use strum::AsRefStr; @@ -35,7 +36,7 @@ pub trait Selector: Send + Sync { type Context; type Output; - async fn select(&self, ctx: &Self::Context, opts: SelectorOptions) -> Result; + async fn select(&self, ctx: &Self::Context, opts: SelectorOptions<'_>) -> Result; } /// A selector that aware of region statistics @@ -56,8 +57,10 @@ pub trait RegionStatAwareSelector: Send + Sync { ) -> Result; } -#[derive(Debug)] -pub struct SelectorOptions { +#[derive(Debug, Clone)] +pub struct SelectorOptions<'a> { + /// Resolved requirements used to filter candidate nodes. + pub affinity: Option<&'a AffinityRequirements>, /// Minimum number of selected results. pub min_required_items: usize, /// Whether duplicates are allowed in the selected result, default false. @@ -66,18 +69,16 @@ pub struct SelectorOptions { pub exclude_peer_ids: HashSet, /// The filter to select the peers based on their workloads. pub workload_filter: Option bool>, - /// Additional placement context, such as the frontend identity that originated the request. - pub extensions: HashMap, } -impl Default for SelectorOptions { +impl Default for SelectorOptions<'_> { fn default() -> Self { Self { + affinity: None, min_required_items: 1, allow_duplication: false, exclude_peer_ids: HashSet::new(), workload_filter: None, - extensions: HashMap::new(), } } } diff --git a/src/meta-srv/src/selector/common.rs b/src/meta-srv/src/selector/common.rs index 2a31c1ce86d..00e72b23096 100644 --- a/src/meta-srv/src/selector/common.rs +++ b/src/meta-srv/src/selector/common.rs @@ -32,7 +32,7 @@ pub fn filter_out_excluded_peers( } /// According to the `opts`, choose peers from the `weight_array` through `weighted_choose`. -pub fn choose_items(opts: &SelectorOptions, weighted_choose: &mut W) -> Result> +pub fn choose_items(opts: &SelectorOptions<'_>, weighted_choose: &mut W) -> Result> where W: WeightedChoose, { @@ -86,13 +86,36 @@ where #[cfg(test)] mod tests { + use std::any::Any; use std::collections::{HashMap, HashSet}; + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; - use common_meta::peer::Peer; + use api::v1::meta::DatanodeWorkloads; + use api::v1::meta::heartbeat_request::NodeWorkloads; + use common_meta::affinity::{Affinity, AffinityFallback, AffinityRequirements, NoopAffinity}; + use common_meta::cluster::{DatanodeStatus, NodeInfo, NodeStatus}; + use common_meta::datanode::{DatanodeStatKey, DatanodeStatValue, Stat}; + use common_meta::error::{Error as MetaError, Result as MetaResult}; + use common_meta::kv_backend::memory::MemoryKvBackend; + use common_meta::kv_backend::{KvBackend, KvBackendRef, ResettableKvBackend, TxnService}; + use common_meta::peer::{Peer, PeerDiscovery}; + use common_meta::rpc::store::{ + BatchDeleteRequest, BatchDeleteResponse, BatchGetRequest, BatchGetResponse, + BatchPutRequest, BatchPutResponse, DeleteRangeRequest, DeleteRangeResponse, PutRequest, + PutResponse, RangeRequest, RangeResponse, + }; - use crate::selector::SelectorOptions; + use crate::cluster::MetaPeerClientBuilder; + use crate::discovery::utils::accept_ingest_workload; + use crate::metasrv::{SelectTarget, SelectorContext, SelectorRef}; use crate::selector::common::{choose_items, filter_out_excluded_peers}; + use crate::selector::lease_based::LeaseBasedSelector; + use crate::selector::load_based::LoadBasedSelector; + use crate::selector::round_robin::RoundRobinSelector; + use crate::selector::weight_compute::RegionNumsBasedWeightCompute; use crate::selector::weighted_choose::{RandomWeightedChoose, WeightedItem}; + use crate::selector::{Selector, SelectorOptions}; #[test] fn test_choose_peers() { @@ -136,11 +159,11 @@ mod tests { for i in 1..=5 { let opts = SelectorOptions { + affinity: None, min_required_items: i, allow_duplication: false, exclude_peer_ids: HashSet::new(), workload_filter: None, - extensions: HashMap::new(), }; let selected_peers: HashSet<_> = @@ -153,11 +176,11 @@ mod tests { } let opts = SelectorOptions { + affinity: None, min_required_items: 6, allow_duplication: false, exclude_peer_ids: HashSet::new(), workload_filter: None, - extensions: HashMap::new(), }; let selected_result = @@ -166,11 +189,11 @@ mod tests { for i in 1..=50 { let opts = SelectorOptions { + affinity: None, min_required_items: i, allow_duplication: true, exclude_peer_ids: HashSet::new(), workload_filter: None, - extensions: HashMap::new(), }; let selected_peers = @@ -205,4 +228,508 @@ mod tests { assert_eq!(weight_array.len(), 1); assert_eq!(weight_array[0].item.id, 2); } + + struct ExactAffinity; + + impl Affinity for ExactAffinity { + fn matches(&self, requirements: &AffinityRequirements, candidate: &NodeInfo) -> bool { + requirements + .attributes() + .iter() + .all(|(key, value)| candidate.env_vars.get(key) == Some(value)) + } + } + + fn requirements() -> AffinityRequirements { + AffinityRequirements::new(HashMap::from([("pool".into(), "a".into())])).unwrap() + } + + fn node(id: u64, pool: Option<&str>, accepts_ingest: bool) -> NodeInfo { + NodeInfo { + peer: Peer::empty(id), + last_activity_ts: 0, + status: NodeStatus::Datanode(DatanodeStatus { + rcus: 0, + wcus: 0, + leader_regions: 0, + follower_regions: 0, + workloads: DatanodeWorkloads { + types: if accepts_ingest { vec![0] } else { vec![] }, + }, + }), + version: String::new(), + git_commit: String::new(), + start_time_ms: 0, + total_cpu_millicores: 0, + total_memory_bytes: 0, + cpu_usage_millicores: 0, + memory_usage_bytes: 0, + hostname: String::new(), + env_vars: pool + .map(|pool| HashMap::from([("pool".into(), pool.into())])) + .unwrap_or_default(), + } + } + + struct Discovery { + snapshots: Vec>, + fail: bool, + reads: AtomicUsize, + } + + #[async_trait::async_trait] + impl PeerDiscovery for Discovery { + async fn active_frontends(&self) -> MetaResult> { + Ok(vec![]) + } + + async fn active_datanodes( + &self, + filter: Option fn(&'a NodeWorkloads) -> bool>, + ) -> MetaResult> { + let index = self.reads.fetch_add(1, Ordering::Relaxed); + if self.fail { + return common_meta::error::UnexpectedSnafu { + err_msg: "discovery unavailable", + } + .fail(); + } + let index = index.min(self.snapshots.len() - 1); + Ok(self.snapshots[index] + .iter() + .filter(|node| { + let NodeStatus::Datanode(status) = &node.status else { + unreachable!() + }; + filter.is_none_or(|filter| { + filter(&NodeWorkloads::Datanode(status.workloads.clone())) + }) + }) + .cloned() + .collect()) + } + + async fn active_flownodes( + &self, + filter: Option fn(&'a NodeWorkloads) -> bool>, + ) -> MetaResult> { + self.active_datanodes(filter).await + } + } + + #[tokio::test] + async fn test_affinity_constrains_base_selectors() { + let affinity_requirements = requirements(); + for lease_based in [false, true] { + for (use_exact_affinity, has_requirements) in + [(true, false), (true, true), (false, true)] + { + let enabled = use_exact_affinity && has_requirements; + let selector: SelectorRef = if lease_based { + Arc::new(LeaseBasedSelector) + } else { + Arc::new(RoundRobinSelector::new(SelectTarget::Datanode)) + }; + let ctx = SelectorContext { + affinity_fallback: AffinityFallback::Error, + affinity: if use_exact_affinity { + Arc::new(ExactAffinity) + } else { + Arc::new(NoopAffinity) + }, + peer_discovery: Arc::new(Discovery { + snapshots: vec![vec![ + node(1, Some("b"), true), + node(2, Some("a"), true), + node(3, Some("a"), true), + node(4, None, true), + node(5, Some("a"), false), + ]], + reads: AtomicUsize::new(0), + fail: false, + }), + }; + // Selecting a complete round makes the lease-based result deterministic as a set. + let peers = selector + .select( + &ctx, + SelectorOptions { + affinity: has_requirements.then_some(&affinity_requirements), + min_required_items: if enabled { 1 } else { 3 }, + exclude_peer_ids: HashSet::from([3]), + workload_filter: Some(accept_ingest_workload), + ..Default::default() + }, + ) + .await + .unwrap(); + let ids = peers + .into_iter() + .map(|peer| peer.id) + .collect::>(); + assert_eq!( + ids, + if enabled { + HashSet::from([2]) + } else { + HashSet::from([1, 2, 4]) + } + ); + } + } + } + + #[tokio::test] + async fn test_round_robin_filters_each_discovery_read_for_both_targets() { + let affinity_requirements = requirements(); + for target in [SelectTarget::Datanode, SelectTarget::Flownode] { + let discovery = Arc::new(Discovery { + snapshots: vec![ + vec![node(1, Some("b"), true), node(2, Some("a"), true)], + vec![ + node(1, Some("b"), true), + node(2, Some("a"), true), + node(3, None, true), + ], + ], + reads: AtomicUsize::new(0), + fail: false, + }); + let ctx = SelectorContext { + affinity_fallback: AffinityFallback::Error, + peer_discovery: discovery.clone(), + affinity: Arc::new(ExactAffinity), + }; + let selector = RoundRobinSelector::new(target); + let opts = SelectorOptions { + affinity: Some(&affinity_requirements), + ..Default::default() + }; + for _ in 0..2 { + assert_eq!( + selector.select(&ctx, opts.clone()).await.unwrap(), + vec![Peer::empty(2)] + ); + } + assert_eq!(discovery.reads.load(Ordering::Relaxed), 2); + } + } + + #[tokio::test] + async fn test_round_robin_excludes_peers_for_both_targets() { + let affinity_requirements = requirements(); + for target in [SelectTarget::Datanode, SelectTarget::Flownode] { + let ctx = SelectorContext { + affinity_fallback: AffinityFallback::Error, + peer_discovery: Arc::new(Discovery { + snapshots: vec![vec![node(1, Some("a"), true), node(2, Some("a"), true)]], + reads: AtomicUsize::new(0), + fail: false, + }), + affinity: Arc::new(ExactAffinity), + }; + for affinity in [None, Some(&affinity_requirements)] { + // Keep the topology fixed and vary only the exclusion set. + for (excluded, expected) in [(vec![], vec![1, 2]), (vec![1], vec![2, 2])] { + let selector = RoundRobinSelector::new(target); + let peers = selector + .select( + &ctx, + SelectorOptions { + affinity, + min_required_items: 2, + allow_duplication: true, + exclude_peer_ids: excluded.into_iter().collect(), + ..Default::default() + }, + ) + .await + .unwrap(); + assert_eq!( + peers.into_iter().map(|peer| peer.id).collect::>(), + expected + ); + } + let selector = RoundRobinSelector::new(target); + let error = selector + .select( + &ctx, + SelectorOptions { + affinity, + exclude_peer_ids: HashSet::from([1, 2]), + ..Default::default() + }, + ) + .await + .unwrap_err(); + assert!(matches!( + error, + crate::error::Error::NoEnoughAvailableNode { available: 0, .. } + )); + } + } + } + + #[tokio::test] + async fn test_base_selector_affinity_fallback() { + let affinity_requirements = requirements(); + for lease_based in [false, true] { + for fallback in [AffinityFallback::Error, AffinityFallback::Unconstrained] { + for has_requirements in [false, true] { + for required in [2, 3] { + let discovery = Arc::new(Discovery { + snapshots: vec![vec![ + node(1, Some("a"), true), + node(2, Some("b"), true), + node(3, Some("a"), true), + node(4, Some("a"), false), + ]], + reads: AtomicUsize::new(0), + fail: false, + }); + let ctx = SelectorContext { + peer_discovery: discovery.clone(), + affinity: Arc::new(ExactAffinity), + affinity_fallback: fallback, + }; + let selector: SelectorRef = if lease_based { + Arc::new(LeaseBasedSelector) + } else { + Arc::new(RoundRobinSelector::new(SelectTarget::Datanode)) + }; + let result = selector + .select( + &ctx, + SelectorOptions { + affinity: has_requirements.then_some(&affinity_requirements), + min_required_items: required, + allow_duplication: false, + exclude_peer_ids: HashSet::from([3]), + workload_filter: Some(accept_ingest_workload), + }, + ) + .await; + let retries = + has_requirements && fallback == AffinityFallback::Unconstrained; + assert_eq!( + discovery.reads.load(Ordering::Relaxed), + 1 + usize::from(retries) + ); + let constrained = has_requirements && !retries; + if required == 2 && !constrained { + assert_eq!( + result + .unwrap() + .into_iter() + .map(|peer| peer.id) + .collect::>(), + HashSet::from([1, 2]) + ); + } else { + let available = if constrained { 1 } else { 2 }; + assert!( + matches!(result.unwrap_err(), crate::error::Error::NoEnoughAvailableNode { + required: actual_required, available: actual_available, .. + } if actual_required == required && actual_available == available) + ); + } + } + } + } + } + } + + #[tokio::test] + async fn test_base_selector_discovery_error_does_not_fallback() { + let affinity_requirements = requirements(); + for lease_based in [false, true] { + let discovery = Arc::new(Discovery { + snapshots: vec![], + reads: AtomicUsize::new(0), + fail: true, + }); + let ctx = SelectorContext { + peer_discovery: discovery.clone(), + affinity: Arc::new(ExactAffinity), + affinity_fallback: AffinityFallback::Unconstrained, + }; + let selector: SelectorRef = if lease_based { + Arc::new(LeaseBasedSelector) + } else { + Arc::new(RoundRobinSelector::new(SelectTarget::Datanode)) + }; + let error = selector + .select( + &ctx, + SelectorOptions { + affinity: Some(&affinity_requirements), + ..Default::default() + }, + ) + .await + .unwrap_err(); + assert!(matches!( + error, + crate::error::Error::ListActiveDatanodes { .. } + )); + assert_eq!(discovery.reads.load(Ordering::Relaxed), 1); + } + } + + struct CountingStatBackend { + inner: MemoryKvBackend, + reads: AtomicUsize, + fail: bool, + } + + #[async_trait::async_trait] + impl TxnService for CountingStatBackend { + type Error = MetaError; + } + + #[async_trait::async_trait] + impl KvBackend for CountingStatBackend { + fn name(&self) -> &str { + "counting-stat-backend" + } + + fn as_any(&self) -> &dyn Any { + self + } + + async fn batch_get(&self, req: BatchGetRequest) -> MetaResult { + self.reads.fetch_add(1, Ordering::Relaxed); + if self.fail { + return common_meta::error::UnexpectedSnafu { + err_msg: "stat backend unavailable", + } + .fail(); + } + self.inner.batch_get(req).await + } + + async fn range(&self, req: RangeRequest) -> MetaResult { + self.inner.range(req).await + } + + async fn put(&self, req: PutRequest) -> MetaResult { + self.inner.put(req).await + } + + async fn batch_put(&self, req: BatchPutRequest) -> MetaResult { + self.inner.batch_put(req).await + } + + async fn delete_range(&self, req: DeleteRangeRequest) -> MetaResult { + self.inner.delete_range(req).await + } + + async fn batch_delete(&self, req: BatchDeleteRequest) -> MetaResult { + self.inner.batch_delete(req).await + } + } + + impl ResettableKvBackend for CountingStatBackend { + fn reset(&self) { + self.inner.reset(); + } + + fn as_kv_backend_ref(self: Arc) -> KvBackendRef { + self + } + } + + #[tokio::test] + async fn test_load_based_affinity_fallback_queries() { + let affinity_requirements = requirements(); + for (has_requirements, fallback, required, fail_stats, expected_reads, available) in [ + (false, AffinityFallback::Unconstrained, 2, false, 1, 2), + (true, AffinityFallback::Error, 2, false, 1, 1), + (true, AffinityFallback::Unconstrained, 2, false, 2, 2), + (true, AffinityFallback::Unconstrained, 3, false, 2, 2), + (true, AffinityFallback::Unconstrained, 2, true, 1, 0), + ] { + let backend = Arc::new(CountingStatBackend { + inner: MemoryKvBackend::new(), + reads: AtomicUsize::new(0), + fail: fail_stats, + }); + for id in 1..=4 { + let value = DatanodeStatValue { + stats: vec![Stat { + id, + addr: format!("node-{id}"), + region_num: id, + ..Default::default() + }], + }; + backend + .put( + PutRequest::new() + .with_key(Vec::::from(DatanodeStatKey { node_id: id })) + .with_value(Vec::::try_from(value).unwrap()), + ) + .await + .unwrap(); + } + let client = Arc::new( + MetaPeerClientBuilder::default() + .election(None) + .in_memory(backend.clone()) + .build() + .unwrap(), + ); + let selector = LoadBasedSelector::new(RegionNumsBasedWeightCompute, client); + let discovery = Arc::new(Discovery { + snapshots: vec![vec![ + node(1, Some("a"), true), + node(2, Some("b"), true), + node(3, Some("a"), true), + node(4, Some("a"), false), + ]], + reads: AtomicUsize::new(0), + fail: false, + }); + let ctx = SelectorContext { + peer_discovery: discovery.clone(), + affinity: Arc::new(ExactAffinity), + affinity_fallback: fallback, + }; + let result = selector + .select( + &ctx, + SelectorOptions { + affinity: has_requirements.then_some(&affinity_requirements), + min_required_items: required, + allow_duplication: false, + exclude_peer_ids: HashSet::from([3]), + workload_filter: Some(accept_ingest_workload), + }, + ) + .await; + assert_eq!(discovery.reads.load(Ordering::Relaxed), expected_reads); + assert_eq!(backend.reads.load(Ordering::Relaxed), expected_reads); + if fail_stats { + assert!(matches!( + result.unwrap_err(), + crate::error::Error::KvBackend { .. } + )); + } else if available < required { + assert!( + matches!(result.unwrap_err(), crate::error::Error::NoEnoughAvailableNode { + required: actual_required, available: actual_available, .. + } if actual_required == required && actual_available == available) + ); + } else { + assert_eq!( + result + .unwrap() + .into_iter() + .map(|peer| peer.id) + .collect::>(), + HashSet::from([1, 2]) + ); + } + } + } } diff --git a/src/meta-srv/src/selector/lease_based.rs b/src/meta-srv/src/selector/lease_based.rs index 4a44bfca1de..b9ddd852dc4 100644 --- a/src/meta-srv/src/selector/lease_based.rs +++ b/src/meta-srv/src/selector/lease_based.rs @@ -12,10 +12,11 @@ // See the License for the specific language governing permissions and // limitations under the License. +use common_meta::affinity::AffinityFallback; use common_meta::peer::Peer; use snafu::ResultExt; -use crate::error::{ListActiveDatanodesSnafu, Result}; +use crate::error::{self, ListActiveDatanodesSnafu, Result}; use crate::metasrv::SelectorContext; use crate::selector::common::{choose_items, filter_out_excluded_peers}; use crate::selector::weighted_choose::{RandomWeightedChoose, WeightedItem}; @@ -25,12 +26,12 @@ use crate::selector::{Selector, SelectorOptions}; #[derive(Default)] pub struct LeaseBasedSelector; -#[async_trait::async_trait] -impl Selector for LeaseBasedSelector { - type Context = SelectorContext; - type Output = Vec; - - async fn select(&self, ctx: &Self::Context, opts: SelectorOptions) -> Result { +impl LeaseBasedSelector { + async fn select_once( + &self, + ctx: &SelectorContext, + opts: &SelectorOptions<'_>, + ) -> Result> { // 1. get alive datanodes. let alive_datanodes = ctx .peer_discovery @@ -38,6 +39,10 @@ impl Selector for LeaseBasedSelector { .await .context(ListActiveDatanodesSnafu)? .into_iter() + .filter(|node| { + opts.affinity + .is_none_or(|requirements| ctx.affinity.matches(requirements, node)) + }) .map(|node| node.peer) .collect::>(); @@ -53,8 +58,35 @@ impl Selector for LeaseBasedSelector { // 3. choose peers by weight_array. filter_out_excluded_peers(&mut weight_array, &opts.exclude_peer_ids); let mut weighted_choose = RandomWeightedChoose::new(weight_array); - let selected = choose_items(&opts, &mut weighted_choose)?; + let selected = choose_items(opts, &mut weighted_choose)?; Ok(selected) } } + +#[async_trait::async_trait] +impl Selector for LeaseBasedSelector { + type Context = SelectorContext; + type Output = Vec; + + async fn select( + &self, + ctx: &Self::Context, + mut opts: SelectorOptions<'_>, + ) -> Result { + match self.select_once(ctx, &opts).await { + Err(error::Error::NoEnoughAvailableNode { + required, + available, + .. + }) if opts.affinity.is_some() + && ctx.affinity_fallback == AffinityFallback::Unconstrained + && available < required => + { + opts.affinity = None; + self.select_once(ctx, &opts).await + } + result => result, + } + } +} diff --git a/src/meta-srv/src/selector/load_based.rs b/src/meta-srv/src/selector/load_based.rs index 9f36ba051f9..84196f3973e 100644 --- a/src/meta-srv/src/selector/load_based.rs +++ b/src/meta-srv/src/selector/load_based.rs @@ -14,13 +14,14 @@ use std::collections::HashMap; +use common_meta::affinity::AffinityFallback; use common_meta::datanode::{DatanodeStatKey, DatanodeStatValue}; use common_meta::peer::Peer; use common_telemetry::debug; use snafu::ResultExt; use crate::cluster::MetaPeerClientRef; -use crate::error::{ListActiveDatanodesSnafu, Result}; +use crate::error::{self, ListActiveDatanodesSnafu, Result}; use crate::metasrv::SelectorContext; use crate::selector::common::{choose_items, filter_out_excluded_peers}; use crate::selector::weight_compute::WeightCompute; @@ -41,15 +42,15 @@ impl LoadBasedSelector { } } -#[async_trait::async_trait] -impl Selector for LoadBasedSelector +impl LoadBasedSelector where C: WeightCompute>, { - type Context = SelectorContext; - type Output = Vec; - - async fn select(&self, ctx: &Self::Context, opts: SelectorOptions) -> Result { + async fn select_once( + &self, + ctx: &SelectorContext, + opts: &SelectorOptions<'_>, + ) -> Result> { // 1. get alive datanodes. let alive_datanodes = ctx .peer_discovery @@ -57,6 +58,10 @@ where .await .context(ListActiveDatanodesSnafu)? .into_iter() + .filter(|node| { + opts.affinity + .is_none_or(|requirements| ctx.affinity.matches(requirements, node)) + }) .map(|node| node.peer) .collect::>(); @@ -77,7 +82,7 @@ where filter_out_excluded_peers(&mut weight_array, &opts.exclude_peer_ids); // 5. choose peers by weight_array. let mut weighted_choose = RandomWeightedChoose::new(weight_array); - let selected = choose_items(&opts, &mut weighted_choose)?; + let selected = choose_items(opts, &mut weighted_choose)?; debug!( "LoadBasedSelector select peers: {:?}, opts: {:?}.", @@ -88,6 +93,36 @@ where } } +#[async_trait::async_trait] +impl Selector for LoadBasedSelector +where + C: WeightCompute>, +{ + type Context = SelectorContext; + type Output = Vec; + + async fn select( + &self, + ctx: &Self::Context, + mut opts: SelectorOptions<'_>, + ) -> Result { + match self.select_once(ctx, &opts).await { + Err(error::Error::NoEnoughAvailableNode { + required, + available, + .. + }) if opts.affinity.is_some() + && ctx.affinity_fallback == AffinityFallback::Unconstrained + && available < required => + { + opts.affinity = None; + self.select_once(ctx, &opts).await + } + result => result, + } + } +} + fn filter_out_expired_datanode( mut stat_kvs: HashMap, datanodes: &[Peer], diff --git a/src/meta-srv/src/selector/round_robin.rs b/src/meta-srv/src/selector/round_robin.rs index f9cf71e949b..05921921145 100644 --- a/src/meta-srv/src/selector/round_robin.rs +++ b/src/meta-srv/src/selector/round_robin.rs @@ -14,11 +14,12 @@ use std::sync::atomic::AtomicUsize; +use common_meta::affinity::AffinityFallback; use common_meta::peer::Peer; use snafu::{ResultExt, ensure}; use crate::error::{ - ListActiveDatanodesSnafu, ListActiveFlownodesSnafu, NoEnoughAvailableNodeSnafu, Result, + self, ListActiveDatanodesSnafu, ListActiveFlownodesSnafu, NoEnoughAvailableNodeSnafu, Result, }; use crate::metasrv::{SelectTarget, SelectorContext}; use crate::selector::{Selector, SelectorOptions}; @@ -52,35 +53,42 @@ impl RoundRobinSelector { } } - async fn get_peers(&self, opts: &SelectorOptions, ctx: &SelectorContext) -> Result> { + async fn get_peers( + &self, + opts: &SelectorOptions<'_>, + ctx: &SelectorContext, + ) -> Result> { let mut peers = match self.select_target { - SelectTarget::Datanode => { - // 1. get alive datanodes. - let alive_datanodes = ctx - .peer_discovery - .active_datanodes(opts.workload_filter) - .await - .context(ListActiveDatanodesSnafu)?; - - // 2. filter out excluded datanodes. - alive_datanodes - .into_iter() - .map(|node| node.peer) - .filter(|p| !opts.exclude_peer_ids.contains(&p.id)) - .collect::>() - } + SelectTarget::Datanode => ctx + .peer_discovery + .active_datanodes(opts.workload_filter) + .await + .context(ListActiveDatanodesSnafu)? + .into_iter() + .filter(|node| { + opts.affinity + .is_none_or(|requirements| ctx.affinity.matches(requirements, node)) + }) + .map(|node| node.peer) + .filter(|peer| !opts.exclude_peer_ids.contains(&peer.id)) + .collect::>(), SelectTarget::Flownode => ctx .peer_discovery .active_flownodes(opts.workload_filter) .await .context(ListActiveFlownodesSnafu)? .into_iter() + .filter(|node| { + opts.affinity + .is_none_or(|requirements| ctx.affinity.matches(requirements, node)) + }) .map(|node| node.peer) + .filter(|peer| !opts.exclude_peer_ids.contains(&peer.id)) .collect::>(), }; ensure!( - !peers.is_empty(), + !peers.is_empty() && (opts.allow_duplication || peers.len() >= opts.min_required_items), NoEnoughAvailableNodeSnafu { required: opts.min_required_items, available: peers.len(), @@ -93,15 +101,13 @@ impl RoundRobinSelector { Ok(peers) } -} -#[async_trait::async_trait] -impl Selector for RoundRobinSelector { - type Context = SelectorContext; - type Output = Vec; - - async fn select(&self, ctx: &Self::Context, opts: SelectorOptions) -> Result> { - let peers = self.get_peers(&opts, ctx).await?; + async fn select_once( + &self, + ctx: &SelectorContext, + opts: &SelectorOptions<'_>, + ) -> Result> { + let peers = self.get_peers(opts, ctx).await?; // choose peers let mut selected = Vec::with_capacity(opts.min_required_items); for _ in 0..opts.min_required_items { @@ -116,9 +122,39 @@ impl Selector for RoundRobinSelector { } } +#[async_trait::async_trait] +impl Selector for RoundRobinSelector { + type Context = SelectorContext; + type Output = Vec; + + async fn select( + &self, + ctx: &Self::Context, + mut opts: SelectorOptions<'_>, + ) -> Result { + match self.select_once(ctx, &opts).await { + Err(error::Error::NoEnoughAvailableNode { + required, + available, + .. + }) if opts.affinity.is_some() + && ctx.affinity_fallback == AffinityFallback::Unconstrained + && available < required => + { + opts.affinity = None; + self.select_once(ctx, &opts).await + } + result => result, + } + } +} + #[cfg(test)] mod test { use std::collections::HashSet; + use std::sync::Arc; + + use common_meta::affinity::NoopAffinity; use super::*; use crate::test_util::{create_meta_peer_client, put_datanodes}; @@ -128,6 +164,8 @@ mod test { let selector = RoundRobinSelector::default(); let meta_peer_client = create_meta_peer_client(); let ctx = SelectorContext { + affinity_fallback: Default::default(), + affinity: Arc::new(NoopAffinity), peer_discovery: meta_peer_client.clone(), }; // add three nodes @@ -150,11 +188,11 @@ mod test { .select( &ctx, SelectorOptions { + affinity: None, min_required_items: 4, allow_duplication: true, exclude_peer_ids: HashSet::new(), workload_filter: None, - extensions: Default::default(), }, ) .await @@ -169,11 +207,11 @@ mod test { .select( &ctx, SelectorOptions { + affinity: None, min_required_items: 2, allow_duplication: true, exclude_peer_ids: HashSet::new(), workload_filter: None, - extensions: Default::default(), }, ) .await @@ -187,6 +225,8 @@ mod test { let selector = RoundRobinSelector::new(SelectTarget::Datanode); let meta_peer_client = create_meta_peer_client(); let ctx = SelectorContext { + affinity_fallback: Default::default(), + affinity: Arc::new(NoopAffinity), peer_discovery: meta_peer_client.clone(), }; // add three nodes @@ -212,11 +252,11 @@ mod test { .select( &ctx, SelectorOptions { + affinity: None, min_required_items: 1, allow_duplication: true, exclude_peer_ids: HashSet::from([2, 5]), workload_filter: None, - extensions: Default::default(), }, ) .await diff --git a/src/meta-srv/src/selector/test_utils.rs b/src/meta-srv/src/selector/test_utils.rs index 6a992e967fc..9b79a435e90 100644 --- a/src/meta-srv/src/selector/test_utils.rs +++ b/src/meta-srv/src/selector/test_utils.rs @@ -14,6 +14,7 @@ use std::sync::Arc; +use common_meta::affinity::NoopAffinity; use common_meta::kv_backend::memory::MemoryKvBackend; use common_meta::peer::Peer; use rand::prelude::SliceRandom; @@ -34,6 +35,8 @@ pub fn new_test_selector_context() -> SelectorContext { .unwrap(); SelectorContext { + affinity_fallback: Default::default(), + affinity: Arc::new(NoopAffinity), peer_discovery: meta_peer_client, } } @@ -54,7 +57,11 @@ impl Selector for RandomNodeSelector { type Context = SelectorContext; type Output = Vec; - async fn select(&self, _ctx: &Self::Context, _opts: SelectorOptions) -> Result { + async fn select( + &self, + _ctx: &Self::Context, + _opts: SelectorOptions<'_>, + ) -> Result { let mut rng = rand::rng(); let mut nodes = self.nodes.clone(); nodes.shuffle(&mut rng); diff --git a/tests-integration/src/standalone.rs b/tests-integration/src/standalone.rs index 074be282543..f8e24cb001d 100644 --- a/tests-integration/src/standalone.rs +++ b/tests-integration/src/standalone.rs @@ -281,6 +281,8 @@ impl GreptimeDbStandaloneBuilder { let repartition_procedure_factory = Arc::new(StandaloneRepartitionProcedureFactory); let ddl_manager = Arc::new(DdlManager::new( DdlContext { + affinity_resolver: None, + peer_alloc_context_provider: None, node_manager: node_manager.clone(), cache_invalidator: cache_registry.clone(), memory_region_keeper: Arc::new(MemoryRegionKeeper::default()), diff --git a/tests-integration/tests/region_migration.rs b/tests-integration/tests/region_migration.rs index 0df42ee1e88..d50905637db 100644 --- a/tests-integration/tests/region_migration.rs +++ b/tests-integration/tests/region_migration.rs @@ -1224,7 +1224,7 @@ impl Selector for ConstNodeSelector { async fn select( &self, _ctx: &Self::Context, - _opts: SelectorOptions, + _opts: SelectorOptions<'_>, ) -> MetaResult { Ok(self.peers.clone()) } From 7df99f946e8c35f565705905734076121c89131d Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Thu, 15 Oct 2026 09:52:28 +0000 Subject: [PATCH 009/239] test: add manual S3 metric snapshot benchmark (#9573) * test: add repeatable metric snapshot S3 benchmark Signed-off-by: jeremyhi * test: use only existing CI S3 permissions in the metric snapshot benchmark The CI S3 user is not allowed GetBucketLocation or ListBucketMultipartUploads, so the bucket probe and the prefix cleanup failed with AccessDenied. Use only list, put, get and delete under the run prefix, as the object store WAL benchmark does, and probe S3 before allocating the EC2 runner. Signed-off-by: jeremyhi * test: scope metric snapshot inventory to the current round The full round's S3 inventory listed the whole run prefix, so it also counted the smoke snapshot. List only the round's snapshot prefix while cleanup still covers the run, and assert that the teardown's S3 cleanup step runs always() with the validated run prefix. Signed-off-by: jeremyhi * docs: record the AWS validation run of the metric snapshot benchmark Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- .github/scripts/metric-snapshot-fixture.py | 294 +++++++++++++ .github/scripts/metric-snapshot-s3-bench.py | 392 ++++++++++++++++++ .github/scripts/metric-snapshot-s3-ops.py | 217 ++++++++++ .../tests/test_metric_snapshot_fixture.py | 101 +++++ .../tests/test_metric_snapshot_s3_bench.py | 74 ++++ .../tests/test_metric_snapshot_s3_ops.py | 203 +++++++++ .github/workflows/checks.yml | 2 + .../workflows/metric-snapshot-s3-bench.yml | 230 ++++++++++ docs/how-to/metric-snapshot-s3-bench.md | 79 ++++ 9 files changed, 1592 insertions(+) create mode 100644 .github/scripts/metric-snapshot-fixture.py create mode 100644 .github/scripts/metric-snapshot-s3-bench.py create mode 100644 .github/scripts/metric-snapshot-s3-ops.py create mode 100644 .github/scripts/tests/test_metric_snapshot_fixture.py create mode 100644 .github/scripts/tests/test_metric_snapshot_s3_bench.py create mode 100644 .github/scripts/tests/test_metric_snapshot_s3_ops.py create mode 100644 .github/workflows/metric-snapshot-s3-bench.yml create mode 100644 docs/how-to/metric-snapshot-s3-bench.md diff --git a/.github/scripts/metric-snapshot-fixture.py b/.github/scripts/metric-snapshot-fixture.py new file mode 100644 index 00000000000..4341bc6f579 --- /dev/null +++ b/.github/scripts/metric-snapshot-fixture.py @@ -0,0 +1,294 @@ +#!/usr/bin/env python3 +# 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. + +"""Deterministic wide, heterogeneous Metric tables for snapshot benchmarks.""" + +import argparse +import base64 +import concurrent.futures +import hashlib +import http.client +import json +import os +from pathlib import Path +import threading +import time +import urllib.parse + +GROUPS = 10 +TAG_POOL = 298 +ROWS = 32 +DATABASE = "metrics" + + +class Client: + def __init__(self, addr): + self.addr = addr + self.connections = threading.local() + self.headers = {"Content-Type": "application/x-www-form-urlencoded"} + if auth := os.environ.get("GREPTIME_BENCH_AUTH"): + self.headers["Authorization"] = "Basic " + base64.b64encode( + auth.encode() + ).decode() + + def sql(self, query, database=DATABASE, batch=False): + connection = getattr(self.connections, "connection", None) + if connection is None: + connection = http.client.HTTPConnection(self.addr, timeout=600) + self.connections.connection = connection + body = urllib.parse.urlencode({"db": database, "sql": query}) + try: + connection.request( + "POST", "/v1/ddl/logical-tables" if batch else "/v1/sql", + body, self.headers, + ) + response = connection.getresponse() + raw = response.read() + except (OSError, http.client.HTTPException): + connection.close() + self.connections.connection = None + raise + data = json.loads(raw) + if response.status != 200 or "error" in data: + raise RuntimeError(f"SQL failed (HTTP {response.status}): {data}") + return data["output"] + + def values(self, query, database=DATABASE): + return self.sql(query, database)[0]["records"]["rows"] + + +def tags(index, per_group): + return sorted((index % per_group * 7 + offset) % TAG_POOL for offset in range(10)) + + +def column_names(tag_ids): + return ["ts", *[f"tag_{tag:03d}" for tag in tag_ids], "val"] + + +def columns(tag_ids): + names = [f"tag_{tag:03d}" for tag in tag_ids] + return ( + "ts TIMESTAMP TIME INDEX, " + + ", ".join(name + " STRING" for name in names) + + ", val DOUBLE, PRIMARY KEY(" + ",".join(names) + ")" + ) + + +def table_ddl(index, per_group): + return ( + f"CREATE TABLE m{index:06d} ({columns(tags(index, per_group))}) " + f"ENGINE=metric WITH(on_physical_table='phy_{index // per_group:02d}')" + ) + + +def row(index, row_index, per_group): + values = [] + for tag in tags(index, per_group): + if (row_index + tag) % 29 == 0: + value = None + elif (row_index + tag) % 17 == 0: + value = "" + elif tag % 10 == 0: + value = f"instance-{index:06d}-{row_index % 4}" + elif tag % 10 == 1: + value = "区域-" + str(row_index % 3) + elif tag % 10 == 2: + value = "label-" + "x" * 48 + "-" + str(row_index % 2) + else: + value = f"tag{tag:03d}-v{row_index % (2 + tag % 7)}" + values.append(value) + return [row_index, *values, None if row_index % 7 == 0 else index % 97 + row_index / 4] + + +def sql_value(value): + if value is None: + return "NULL" + if isinstance(value, str): + return "'" + value.replace("'", "''") + "'" + return str(value) + + +def canonical(value): + return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":")) + + +def update_digest(digest, value): + digest.update((canonical(value) + "\n").encode()) + + +def progress(phase, completed, total): + if completed == total or completed % 4096 == 0: + print(f"{phase}: {completed}/{total}", flush=True) + + +def prepare(client, per_group): + stages = {} + started = time.monotonic() + client.sql(f"CREATE DATABASE {DATABASE}", "public") + for group in range(GROUPS): + client.sql( + f"CREATE TABLE phy_{group:02d} ({columns(range(TAG_POOL))}) " + "ENGINE=metric WITH(physical_metric_table='')" + ) + total = GROUPS * per_group + for group in range(GROUPS): + for base in range(group * per_group, (group + 1) * per_group, 128): + end = min((group + 1) * per_group, base + 128) + client.sql(";".join(table_ddl(i, per_group) for i in range(base, end)), batch=True) + progress("create", end, total) + stages["create_seconds"] = time.monotonic() - started + + def insert(index): + rows = ",".join( + "(" + ",".join(sql_value(v) for v in row(index, j, per_group)) + ")" + for j in range(ROWS) + ) + names = ",".join(column_names(tags(index, per_group))) + client.sql(f"INSERT INTO m{index:06d} ({names}) VALUES {rows}") + + started = time.monotonic() + with concurrent.futures.ThreadPoolExecutor(max_workers=4) as pool: + for base in range(0, total, 128): + end = min(total, base + 128) + list(pool.map(insert, range(base, end))) + progress("insert", end, total) + stages["insert_seconds"] = time.monotonic() - started + started = time.monotonic() + for group in range(GROUPS): + client.sql(f"ADMIN FLUSH_TABLE('{DATABASE}.phy_{group:02d}')") + stages["flush_seconds"] = time.monotonic() - started + return stages + + +def check_records(index, per_group, records): + names = column_names(tags(index, per_group)) + types = ["TimestampMillisecond", *["String"] * 10, "Float64"] + expected_columns = [dict(name=name, data_type=kind) for name, kind in zip(names, types)] + actual_columns = records["schema"]["column_schemas"] + if sorted(actual_columns, key=lambda c: c["name"]) != sorted( + expected_columns, key=lambda c: c["name"] + ): + raise RuntimeError(f"logical schema mismatch: m{index:06d}") + # Metric tables canonicalize column order; SELECT * still checks every column. + positions = [names.index(column["name"]) for column in actual_columns] + expected_rows = [ + [values[position] for position in positions] + for values in (row(index, j, per_group) for j in range(ROWS)) + ] + # Canonical JSON also distinguishes integers, floats, booleans, and NULL. + if canonical(records["rows"]) != canonical(expected_rows): + raise RuntimeError(f"typed values/NULL mismatch: m{index:06d}") + + +def validate(client, per_group): + total = GROUPS * per_group + expected_names = {f"m{i:06d}" for i in range(total)} | { + f"phy_{group:02d}" for group in range(GROUPS) + } + table_rows = client.values( + "SELECT table_name, table_id FROM information_schema.tables " + f"WHERE table_schema='{DATABASE}' ORDER BY table_name" + ) + table_ids = dict(table_rows) + if len(table_rows) != len(expected_names) or set(table_ids) != expected_names: + raise RuntimeError("table membership mismatch") + if len(set(table_ids.values())) != len(table_ids): + raise RuntimeError("duplicate table IDs") + digests = {name: hashlib.sha256() for name in ("schema", "data", "ddl")} + + def check_ddl(names): + outputs = client.sql(";".join(f"SHOW CREATE TABLE {name}" for name in names)) + if len(outputs) != len(names): + raise RuntimeError("incomplete SHOW CREATE TABLE response") + for name, output in zip(names, outputs): + records = output["records"]["rows"] + if len(records) != 1 or records[0][0] != name or not records[0][1]: + raise RuntimeError(f"invalid SHOW CREATE TABLE response: {name}") + update_digest(digests["ddl"], records) + + for base in range(0, total, 64): + end = min(total, base + 64) + names = [f"m{i:06d}" for i in range(base, end)] + outputs = client.sql(";".join(f"SELECT * FROM {name} ORDER BY ts" for name in names)) + if len(outputs) != len(names): + raise RuntimeError("incomplete SELECT response") + for index, output in enumerate(outputs, base): + records = output["records"] + check_records(index, per_group, records) + update_digest(digests["schema"], [f"m{index:06d}", records["schema"]]) + update_digest(digests["data"], [f"m{index:06d}", records["rows"]]) + check_ddl(names) + progress("validate", end, total) + check_ddl([f"phy_{group:02d}" for group in range(GROUPS)]) + return {**{name + "_sha256": digest.hexdigest() for name, digest in digests.items()}, + "table_ids": table_ids} + + +def compare_source(result, source): + for field in ("fixture", "schema_sha256", "data_sha256", "ddl_sha256"): + if result[field] != source[field]: + raise RuntimeError(f"source/target mismatch: {field}") + source_ids, target_ids = source["table_ids"], result["table_ids"] + if source_ids.keys() != target_ids.keys(): + raise RuntimeError("source/target table membership mismatch") + changed = sum(source_ids[name] != target_ids[name] for name in source_ids) + if changed != len(source_ids): + raise RuntimeError(f"table IDs not all remapped: {changed}/{len(source_ids)}") + result["remapped_table_ids"] = changed + + +def placeholders(client): + database = "metric_bench_guard" + client.sql(f"CREATE DATABASE {database}", "public") + for index in range(17): + client.sql(f"CREATE TABLE placeholder_{index} (ts TIMESTAMP TIME INDEX)", database) + return {"placeholder_tables": 17} + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("action", choices=("prepare", "validate", "placeholders")) + parser.add_argument("--addr", default="127.0.0.1:4000") + parser.add_argument("--tables-per-group", type=int, default=10000) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--source", type=Path, help="Source validation JSON, for target comparison") + args = parser.parse_args() + if not 50 <= args.tables_per_group <= 10000: + parser.error("--tables-per-group must be between 50 and 10000") + if args.source is not None and args.action != "validate": + parser.error("--source only applies to validate") + coverage = {tag for i in range(args.tables_per_group) for tag in tags(i, args.tables_per_group)} + if coverage != set(range(TAG_POOL)): + parser.error("fixture does not cover all 298 tags") + client = Client(args.addr) + started = time.monotonic() + result = {"fixture": {"groups": GROUPS, "tables_per_group": args.tables_per_group, + "rows_per_table": ROWS, "tag_pool": TAG_POOL}} + if args.action == "placeholders": + result.update(placeholders(client)) + elif args.action == "prepare": + result.update(prepare(client, args.tables_per_group)) + else: + result.update(validate(client, args.tables_per_group)) + result["wall_seconds"] = time.monotonic() - started + if args.source is not None: + compare_source(result, json.loads(args.source.read_text())) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(result, ensure_ascii=False, indent=2) + "\n") + print(json.dumps({key: value for key, value in result.items() if key != "table_ids"})) + + +if __name__ == "__main__": + main() diff --git a/.github/scripts/metric-snapshot-s3-bench.py b/.github/scripts/metric-snapshot-s3-bench.py new file mode 100644 index 00000000000..7bb01a410b1 --- /dev/null +++ b/.github/scripts/metric-snapshot-s3-bench.py @@ -0,0 +1,392 @@ +#!/usr/bin/env python3 +# 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. + +"""Run one packed Metric snapshot smoke, then one full S3 round trip.""" + +import argparse +from datetime import datetime +import json +import os +from pathlib import Path +import re +import signal +import subprocess +import sys +import threading +import time +import urllib.request + +HERE = Path(__file__).resolve().parent + + +def command(args, **kwargs): + result = subprocess.run(args, text=True, stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, **kwargs) + if result.returncode: + # CLI argv and error output can contain storage credentials. + raise RuntimeError(f"{args[0]} exited with status {result.returncode}") + return result.stdout + + +def write_json(path, value): + path.write_text(json.dumps(value, indent=2, sort_keys=True) + "\n") + + +def metric_values(text): + values = {} + for line in text.splitlines(): + if line and not line.startswith("#"): + key, value = line.rsplit(" ", 1) + values[key] = float(value) + return values + + +def docker_timestamp(value): + # Docker emits nanoseconds; Python 3.9 accepts at most microseconds. + return datetime.fromisoformat(re.sub(r"(\.\d{6})\d+", r"\1", value).replace("Z", "+00:00")) + + +def host_memory(): + if not Path("/proc/vmstat").exists(): + return {"available": False} + vm = dict(line.split() for line in Path("/proc/vmstat").read_text().splitlines()) + mem = dict(line.split(":", 1) for line in Path("/proc/meminfo").read_text().splitlines()) + return {"available": True, "pswpin_pages": int(vm["pswpin"]), + "pswpout_pages": int(vm["pswpout"]), "oom_kill": int(vm.get("oom_kill", 0)), + **{key + "_kib": int(mem[key].split()[0]) + for key in ("MemTotal", "MemAvailable", "SwapTotal", "SwapFree")}} + + +def process_usage(pid): + try: + stat = Path(f"/proc/{pid}/stat").read_text().rsplit(")", 1)[1].split() + return {"cpu_seconds": (int(stat[11]) + int(stat[12])) / os.sysconf("SC_CLK_TCK"), + "rss_bytes": int(stat[21]) * os.sysconf("SC_PAGE_SIZE")} + except (OSError, IndexError): + return None + + +class Benchmark: + def __init__(self, local_smoke): + self.local_smoke = local_smoke + self.image = os.environ["GREPTIMEDB_IMAGE"] + self.commit = os.environ["EXPECTED_COMMIT"] + if not re.fullmatch(r"greptime/greptimedb(?:-nightly)?@sha256:[0-9a-f]{64}", self.image): + raise ValueError("GREPTIMEDB_IMAGE must be an immutable official image digest") + if not re.fullmatch(r"[0-9a-f]{40}", self.commit): + raise ValueError("EXPECTED_COMMIT must be a full commit SHA") + self.per_group = int(os.environ.get("LOGICAL_TABLES_PER_GROUP", "10000")) + if not 50 <= self.per_group <= 10000: + raise ValueError("LOGICAL_TABLES_PER_GROUP must be in 50..10000") + self.prefix = os.environ["RUN_PREFIX"] + if not re.fullmatch(r"metric-snapshot-bench/[0-9]+-[0-9]+/", self.prefix): + raise ValueError("RUN_PREFIX must be metric-snapshot-bench/-/") + self.results = Path(os.environ["RESULTS"]).resolve() + self.data = Path(os.environ["DATA_ROOT"]).resolve() + self.results.mkdir(parents=True, exist_ok=False) + self.data.mkdir(parents=True, exist_ok=False) + self.network = "metric-snapshot-" + self.prefix.split("/")[1] + self.server = self.network + "-db" + self.cli_name = self.network + "-cli" + self.port = int(os.environ.get("BENCH_HTTP_PORT", "19410")) + self.base_url = f"http://127.0.0.1:{self.port}" + self.endpoint = os.environ.get("S3_ENDPOINT", "") + if self.endpoint and not local_smoke: + raise ValueError("Custom S3 endpoints are restricted to local smoke") + if not local_smoke and os.environ["S3_REGION"] != "ap-southeast-1": + raise ValueError("Cloud benchmark requires ap-southeast-1") + self.containers = set() + self.network_created = False + self.phases = [] + self.server_pid = 0 + self.round_dir = None + + def docker(self, *args): + return command(["docker", *args]) + + def metrics(self): + try: + with urllib.request.urlopen(self.base_url + "/metrics", timeout=5) as response: + return response.read().decode() + except OSError: + return "" + + def phase(self, name, action): + print(f"Starting {self.round_dir.name}/{name}", flush=True) + before = host_memory() + metrics_before = self.metrics() + samples = [] + done = threading.Event() + + def sample(): + while not done.is_set(): + samples.append({"time": time.time(), "server": process_usage(self.server_pid), + "cli": process_usage(getattr(self, "cli_pid", 0)), + "query_memory": {k: v for k, v in metric_values(self.metrics()).items() + if "query_memory" in k} if self.server_pid else {}}) + done.wait(0.5) + + monitor = threading.Thread(target=sample, daemon=True) + monitor.start() + start = time.monotonic() + outcome = "failed" + result = None + try: + result = action() + outcome = "passed" + return result + finally: + elapsed = time.monotonic() - start + done.set() + monitor.join() + after = host_memory() + interfered = before["available"] and after["available"] and ( + after["pswpout_pages"] > before["pswpout_pages"] or after["oom_kill"] > before["oom_kill"]) + if interfered: + outcome = "resource-invalid" + record = {"round": self.round_dir.name, "phase": name, "outcome": outcome, + "wall_seconds": elapsed, "host_before": before, "host_after": after, + "process_samples": samples} + if isinstance(result, dict): + record.update(result) + if before["available"] and after["available"]: + record["new_swapout_pages"] = after["pswpout_pages"] - before["pswpout_pages"] + record["new_oom_kills"] = after["oom_kill"] - before["oom_kill"] + metrics_after = self.metrics() + pre, post = metric_values(metrics_before), metric_values(metrics_after) + record["server_counter_deltas"] = { + key: value - pre.get(key, 0) for key, value in post.items() + if ("opendal" in key and ("_count{" in key or "_sum{" in key)) + or key.startswith("greptime_servers_http_requests_total{")} + write_json(self.round_dir / f"{name}-resources.json", record) + (self.round_dir / f"{name}-before.prom").write_text(metrics_before) + (self.round_dir / f"{name}-after.prom").write_text(metrics_after) + self.phases.append({k: v for k, v in record.items() if k != "process_samples"}) + self.summary() + print(f"Finished {name}: {outcome}, {elapsed:.3f}s", flush=True) + if interfered: + raise RuntimeError(f"Resource interference during {name}; evidence retained") + + def fixture(self, operation, per_group, output, source=None): + args = [sys.executable, str(HERE / "metric-snapshot-fixture.py"), operation, + "--addr", f"127.0.0.1:{self.port}", "--tables-per-group", str(per_group), + "--output", str(output)] + if source: + args += ["--source", str(source)] + with (self.round_dir / f"{operation}-{output.stem}.log").open("w") as log: + status = subprocess.run(args, stdout=log, stderr=subprocess.STDOUT).returncode + if status: + raise RuntimeError(f"Fixture {operation} failed; see its artifact log") + + def start_server(self, role): + root = self.data / self.round_dir.name / role + root.mkdir(parents=True) + evidence = self.round_dir / role + evidence.mkdir() + config = root / "config.toml" + config.write_text('''enable_telemetry = false +[http] +addr = "0.0.0.0:4000" +[grpc] +bind_addr = "0.0.0.0:4001" +[mysql] +enable = false +[postgres] +enable = false +[storage] +type = "File" +data_home = "/data/db" +[query] +parallelism = 4 +memory_pool_size = "2GiB" +[logging] +dir = "/evidence/logs" +level = "info" +''') + self.containers.add(self.server) + self.docker("run", "-d", "--name", self.server, "--network", self.network, + "--network-alias", "metric-db", "--user", f"{os.getuid()}:{os.getgid()}", + "--add-host", "host.docker.internal:host-gateway", + "-p", f"127.0.0.1:{self.port}:4000", + "-v", f"{root}:/data", "-v", f"{evidence}:/evidence", + self.image, "standalone", "start", "-c", "/data/config.toml") + self.server_pid = int(self.docker("inspect", "-f", "{{.State.Pid}}", self.server)) + deadline = time.monotonic() + 120 + while time.monotonic() < deadline: + try: + with urllib.request.urlopen(self.base_url + "/v1/capabilities", timeout=2) as response: + caps = json.load(response) + if any(caps.get(key) != 1 for key in + ("metric_packed_export", "metric_packed_import", "metric_batch_ddl")): + raise RuntimeError("Required packed snapshot capabilities are absent") + write_json(evidence / "capabilities.json", caps) + return + except OSError: + if self.docker("inspect", "-f", "{{.State.Running}}", self.server).strip() != "true": + break + time.sleep(1) + raise RuntimeError("Server did not become ready") + + def stop_server(self, role): + self.docker("stop", "-t", "60", self.server) + state = json.loads(self.docker("inspect", "-f", "{{json .State}}", self.server)) + write_json(self.round_dir / role / "exit.json", state) + (self.round_dir / role / "stdout.log").write_text(self.docker("logs", self.server)) + self.docker("rm", self.server) + self.containers.remove(self.server) + self.server_pid = 0 + if state["OOMKilled"] or state["ExitCode"] != 0: + raise RuntimeError(f"{role} server did not shut down cleanly") + + def cli(self, args, name): + # Expand credentials only inside the container. Never persist argv or env. + script = 'exec greptime "$@" --s3 --s3-region "$S3_REGION" --s3-access-key-id "$AWS_ACCESS_KEY_ID" --s3-secret-access-key "$AWS_SECRET_ACCESS_KEY"' + if self.endpoint: + args += ["--s3-endpoint", self.endpoint] + self.containers.add(self.cli_name) + self.docker("create", "--name", self.cli_name, "--network", self.network, + "--user", f"{os.getuid()}:{os.getgid()}", + "--add-host", "host.docker.internal:host-gateway", + "-e", "AWS_ACCESS_KEY_ID", "-e", "AWS_SECRET_ACCESS_KEY", "-e", "S3_REGION", + "-v", f"{self.round_dir}:/evidence", "--entrypoint", "/bin/sh", self.image, + "-c", script, "sh", "--log-dir", f"/evidence/{name}-logs", "cli", "data", *args) + + def execute(): + self.docker("start", self.cli_name) + self.cli_pid = int(self.docker("inspect", "-f", "{{.State.Pid}}", self.cli_name)) + status = int(self.docker("wait", self.cli_name)) + state = json.loads(self.docker("inspect", "-f", "{{json .State}}", self.cli_name)) + (self.round_dir / f"{name}.log").write_text(self.docker("logs", self.cli_name)) + write_json(self.round_dir / f"{name}-exit.json", state) + if status or state["ExitCode"] or state["OOMKilled"]: + raise RuntimeError(f"{name} failed; see artifact log") + return {"cli_wall_seconds": (docker_timestamp(state["FinishedAt"]) - + docker_timestamp(state["StartedAt"])).total_seconds()} + + try: + self.phase(name, execute) + finally: + self.cli_pid = 0 + if not (self.round_dir / f"{name}-exit.json").exists(): + self.save_container(self.cli_name, name + "-failure") + self.docker("rm", "-f", self.cli_name) + self.containers.remove(self.cli_name) + + def round(self, label, per_group): + self.round_dir = self.results / label + self.round_dir.mkdir() + self.start_server("source") + source = self.round_dir / "source-validation.json" + self.phase("prepare", lambda: self.fixture("prepare", per_group, self.round_dir / "prepare.json")) + self.phase("source-validation", lambda: self.fixture("validate", per_group, source)) + location = f"s3://{os.environ['S3_BUCKET']}/{self.prefix}{label}" + self.cli(["export-v2", "create", "--addr", "metric-db:4000", "--to", location, + "--schemas", "metrics", "--pack-metric-data", "--parallelism", "4", + "--chunk-parallelism", "1", "--start-time", "1970-01-01T00:00:00Z", + "--end-time", "1970-01-01T00:00:01Z", "--timeout", "1h", + "--no-proxy", "--progress", "always"], "export") + if not self.local_smoke: + command(["uv", "run", str(HERE / "metric-snapshot-s3-ops.py"), "inventory", "--round", label, + "--output", str(self.round_dir / "s3-inventory.json")]) + self.stop_server("source") + self.cli(["export-v2", "verify", "--snapshot", location], "verify") + self.start_server("target") + self.fixture("placeholders", per_group, self.round_dir / "placeholders.json") + self.cli(["import-v2", "--addr", "metric-db:4000", "--from", location, + "--task-parallelism", "1", "--state-path", "/evidence/import-state.json", + "--timeout", "1h", "--no-proxy", "--progress", "always"], "import") + imported = json.loads((self.round_dir / "import-resources.json").read_text()) + batch_requests = sum(value for key, value in imported["server_counter_deltas"].items() + if '/v1/ddl/logical-tables' in key and 'code="200"' in key) + if batch_requests <= 0: + raise RuntimeError("Import did not complete any successful batch DDL HTTP requests") + write_json(self.round_dir / "import-completion.json", {"batch_ddl_requests": batch_requests, + "resume_state_removed": not (self.round_dir / "import-state.json").exists()}) + if (self.round_dir / "import-state.json").exists(): + raise RuntimeError("Successful import left its resume state file") + self.phase("target-validation", lambda: self.fixture( + "validate", per_group, self.round_dir / "target-validation.json", source)) + self.stop_server("target") + + def summary(self): + lines = ["# Metric packed snapshot benchmark", "", f"Image: `{self.image}`", + f"Source commit: `{self.commit}`", "", + "Local smoke only; timings are not AWS results." if self.local_smoke else + "AWS S3 in ap-southeast-1; source and target use local File storage.", "", + "| Round | Phase | Seconds | Result |", "| --- | --- | ---: | --- |"] + for phase in self.phases: + lines.append(f"| {phase['round']} | {phase['phase']} | {phase['wall_seconds']:.3f} | {phase['outcome']} |") + lines += ["", "Each logical table has 32 rows, 10 overlapping string tags, a timestamp and a DOUBLE.", + "Phase durations include Docker start, wait and log collection. cli_wall_seconds records Docker StartedAt to FinishedAt separately; setup, loading and full validation are separate.", + "Process CPU/RSS are sampled every 0.5 s on Linux, not exact process peaks. Linux host swap/OOM counters are checked per phase; configured swap is not RAM.", + "Raw server Prometheus snapshots and CLI progress/logs are retained. OpenDAL operation counts are not AWS HTTP request counts: multipart requests, retries and CLI operations are not fully covered. Exact AWS request totals are unavailable."] + (self.results / "summary.md").write_text("\n".join(lines) + "\n") + + def run(self): + self.docker("pull", self.image) + version = self.docker("run", "--rm", self.image, "--version") + (self.results / "version.txt").write_text(version) + if self.commit not in version or "clean: true" not in version: + raise RuntimeError("Image version does not match the expected clean source commit") + for args, required in ((["export-v2", "create"], "--pack-metric-data"), + (["import-v2"], "--task-parallelism")): + help_text = self.docker("run", "--rm", self.image, "cli", "data", *args, "--help") + (self.results / (args[0] + "-help.txt")).write_text(help_text) + if required not in help_text: + raise RuntimeError(f"Image CLI lacks {required}") + write_json(self.results / "host.json", {"memory": host_memory(), "local_smoke": self.local_smoke, + "platform": self.docker("image", "inspect", "-f", "{{.Os}}/{{.Architecture}}", self.image).strip()}) + self.docker("network", "create", self.network) + self.network_created = True + self.round("smoke", 50) + if not self.local_smoke and self.per_group != 50: + self.round("full", self.per_group) + + def save_container(self, name, label): + state = subprocess.run(["docker", "inspect", "-f", "{{json .State}}", name], + capture_output=True, text=True) + (self.round_dir / (label + "-state.json")).write_text(state.stdout) + result = subprocess.run(["docker", "logs", name], capture_output=True, text=True) + (self.round_dir / (label + ".log")).write_text(result.stdout + result.stderr) + + def cleanup(self): + for name in self.containers: + if self.round_dir: + self.save_container(name, name + "-failure") + subprocess.run(["docker", "rm", "-f", name], stdout=subprocess.DEVNULL, check=False) + if self.network_created: + subprocess.run(["docker", "network", "rm", self.network], stdout=subprocess.DEVNULL, check=False) + self.summary() + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("command", choices=["run"]) + parser.add_argument("--local-smoke", action="store_true") + args = parser.parse_args() + bench = Benchmark(args.local_smoke) + + def cancelled(signum, frame): + raise RuntimeError(f"Received signal {signum}") + + signal.signal(signal.SIGTERM, cancelled) + try: + bench.run() + finally: + bench.cleanup() + + +if __name__ == "__main__": + main() diff --git a/.github/scripts/metric-snapshot-s3-ops.py b/.github/scripts/metric-snapshot-s3-ops.py new file mode 100644 index 00000000000..f1e34defb89 --- /dev/null +++ b/.github/scripts/metric-snapshot-s3-ops.py @@ -0,0 +1,217 @@ +#!/usr/bin/env python3 +# 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. + +# /// script +# requires-python = ">=3.10" +# dependencies = ["boto3>=1.35,<2"] +# /// + +"""Budget preflight and strictly scoped S3 operations for the metric snapshot benchmark.""" + +import argparse +from decimal import Decimal +import json +import os +from pathlib import Path +import re +import time +import urllib.request + +REGION = "ap-southeast-1" +HOURLY_USD = Decimal("0.9408") +OTHER_RESERVE_USD = Decimal("1.50") + + +def owned_prefix(prefix): + if not re.fullmatch(r"metric-snapshot-bench/[0-9]+-[0-9]+/", prefix): + raise ValueError("prefix must be exactly metric-snapshot-bench/-/") + return prefix + + +def plan(env): + image = env["GREPTIMEDB_IMAGE"] + if not re.fullmatch(r"greptime/greptimedb(?:-nightly)?@sha256:[0-9a-f]{64}", image): + raise ValueError("image must be an immutable greptime/greptimedb sha256 reference") + if not re.fullmatch(r"[0-9a-f]{40}", env["EXPECTED_COMMIT"]): + raise ValueError("expected commit must be a full lowercase source SHA") + tables, minutes = int(env["LOGICAL_TABLES_PER_GROUP"]), int(env["BENCHMARK_TIMEOUT_MINUTES"]) + if not 50 <= tables <= 10_000 or not 30 <= minutes <= 240: + raise ValueError("logical tables per group must be 50..10000; timeout must be 30..240 minutes") + if env["S3_REGION"] != REGION or env["RUNNER_REGION"] != REGION: + raise ValueError("bucket and runner must both use ap-southeast-1") + if not env["S3_BUCKET"]: + raise ValueError("S3 bucket is required") + prior = Decimal(env["PRIOR_SPEND_USD"]) + # Reserve 45m provision, 10m queue allowance and 20m cleanup beyond the job timeout. + reservation = HOURLY_USD * Decimal(minutes + 75) / 60 + OTHER_RESERVE_USD + if not prior.is_finite() or prior < 0 or prior >= 15 or prior + reservation > 20: + raise ValueError("campaign budget exhausted: prior spend must be below $15 and leave the full attempt reserve within $20") + return {"image": image, "expected_commit": env["EXPECTED_COMMIT"], + "logical_tables_per_group": tables, "smoke_logical_tables": 500, + "timeout_minutes": minutes, "prefix": owned_prefix(env["RUN_PREFIX"]), + "instance_type": "c6id.4xlarge", "region": REGION, + "prior_estimated_usd": str(prior), "attempt_reserve_usd": str(reservation), + "ec2_hourly_usd": str(HOURLY_USD), "other_reserve_usd": str(OTHER_RESERVE_USD), + "cost_note": "Estimate, not an AWS billing total or a hard spending cap."} + + +def client(region): + import boto3 + from botocore.config import Config + + if region != REGION: + raise ValueError("S3 region must be ap-southeast-1") + return boto3.client("s3", region_name=region, config=Config( + connect_timeout=10, read_timeout=30, retries={"mode": "standard", "max_attempts": 3})) + + +def entries(s3, operation, bucket, prefix, fields, scope=""): + listed = owned_prefix(prefix) + scope + for page in s3.get_paginator(operation).paginate(Bucket=bucket, Prefix=listed): + for field in fields: + for item in page.get(field, []): + if not item["Key"].startswith(listed): + raise ValueError("S3 returned an object outside the owned prefix") + yield item + + +def inventory(s3, bucket, prefix, round=None): + # A round's snapshot lives under /; earlier rounds stay until cleanup. + if round not in (None, "smoke", "full"): + raise ValueError("round must be smoke or full") + scope = round + "/" if round else "" + objects = list(entries(s3, "list_objects_v2", bucket, prefix, ("Contents",), scope)) + return {"prefix": prefix + scope, "objects": len(objects), "bytes": sum(item["Size"] for item in objects), + "measurement": "S3 listing inventory; not request counts or transferred bytes."} + + +def probe(s3, bucket, prefix): + if inventory(s3, bucket, prefix)["objects"]: + raise ValueError("run prefix is already occupied; use a fresh workflow attempt") + key = prefix + "_probe" + s3.put_object(Bucket=bucket, Key=key, Body=b"metric snapshot benchmark") + try: + body = s3.get_object(Bucket=bucket, Key=key)["Body"] + try: + if body.read() != b"metric snapshot benchmark": + raise ValueError("S3 probe returned unexpected content") + finally: + body.close() + finally: + s3.delete_object(Bucket=bucket, Key=key) + return inventory(s3, bucket, prefix) + + +def cleanup(s3, bucket, prefix): + owned_prefix(prefix) + for item in entries(s3, "list_objects_v2", bucket, prefix, ("Contents",)): + s3.delete_object(Bucket=bucket, Key=item["Key"]) + after = inventory(s3, bucket, prefix) + if after["objects"]: + raise RuntimeError("owned S3 prefix is not empty after cleanup") + return after + + +def scrub(root, env): + secrets = [env[name].encode() for name in ("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_SESSION_TOKEN") + if env.get(name)] + if not secrets: + raise ValueError("credential scan requires the CI credential environment") + overlap = max(map(len, secrets)) - 1 + for result in Path(root).rglob("*"): + if result.is_file(): + with result.open("rb") as stream: + previous = b"" + while chunk := stream.read(1024 * 1024): + data = previous + chunk + if any(secret in data for secret in secrets): + raise ValueError("credential detected in results; artifact upload is blocked") + previous = data[-overlap:] if overlap else b"" + + +def wait_for_benchmark(env, started_at): + """Bound runner allocation even when the benchmark job never leaves its queue.""" + run_id, attempt = owned_prefix(env["RUN_PREFIX"]).split("/")[1].split("-") + repo = env["GITHUB_REPOSITORY"] + if run_id != env["GITHUB_RUN_ID"] or not re.fullmatch(r"[\w.-]+/[\w.-]+", repo): + raise ValueError("workflow identity does not match its owned resource prefix") + minutes = int(env["BENCHMARK_TIMEOUT_MINUTES"]) + if not 30 <= minutes <= 240 or started_at <= 0: + raise ValueError("invalid benchmark allocation deadline") + # Use the original allocation timestamp; a delayed or retried watcher cannot extend it. + deadline = started_at + (45 + minutes + 10) * 60 + request = urllib.request.Request( + f"https://api.github.com/repos/{repo}/actions/runs/{run_id}/attempts/{attempt}/jobs?per_page=100", + headers={"Authorization": f"Bearer {env['GH_TOKEN']}", "Accept": "application/vnd.github+json", + "X-GitHub-Api-Version": "2022-11-28"}) + while True: + remaining = deadline - time.time() + if remaining <= 0: + raise TimeoutError("benchmark allocation deadline reached; proceeding to resource cleanup") + # API errors deliberately fail this step: subsequent always() steps clean up immediately. + with urllib.request.urlopen(request, timeout=min(15, remaining)) as response: + jobs = json.load(response)["jobs"] + benchmark = next((job for job in jobs if job["name"] == "benchmark"), None) + if benchmark and benchmark["status"] == "completed": + return {"benchmark_conclusion": benchmark["conclusion"], "owned_attempt": attempt} + remaining = deadline - time.time() + if remaining > 0: + time.sleep(min(30, remaining)) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("command", choices=("plan", "probe", "inventory", "cleanup", "scrub", "cost", "wait-for-benchmark")) + parser.add_argument("--output") + parser.add_argument("--started-at", type=int) + parser.add_argument("--round") + args = parser.parse_args() + if args.round is not None and args.command != "inventory": + parser.error("--round applies only to inventory") + if args.command == "plan": + result = plan(os.environ) + elif args.command == "scrub": + scrub(os.environ["RESULTS"], os.environ) + return + elif args.command == "wait-for-benchmark": + if args.started_at is None: + parser.error("wait-for-benchmark requires --started-at") + result = wait_for_benchmark(os.environ, args.started_at) + elif args.command == "cost": + if args.started_at is None: + parser.error("cost requires --started-at") + seconds = max(0, int(time.time()) - args.started_at) + estimate = HOURLY_USD * Decimal(seconds) / 3600 + result = {"allocation_window_seconds": seconds, "ec2_estimate_usd": str(round(estimate, 4)), + "campaign_estimate_with_other_reserve_usd": str(round( + Decimal(os.environ["PRIOR_SPEND_USD"]) + estimate + OTHER_RESERVE_USD, 4)), + "note": "Allocation wall-time estimate plus $1.50 reserve for S3, disk and IPv4; not an AWS bill."} + else: + prefix = owned_prefix(os.environ["RUN_PREFIX"]) + s3, bucket = client(os.environ["S3_REGION"]), os.environ["S3_BUCKET"] + if args.command == "inventory": + result = inventory(s3, bucket, prefix, args.round) + else: + result = globals()[args.command](s3, bucket, prefix) + rendered = json.dumps(result, indent=2) + "\n" + if args.output: + destination = Path(args.output) + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_text(rendered) + print(rendered, end="") + + +if __name__ == "__main__": + main() diff --git a/.github/scripts/tests/test_metric_snapshot_fixture.py b/.github/scripts/tests/test_metric_snapshot_fixture.py new file mode 100644 index 00000000000..f4c16a11af3 --- /dev/null +++ b/.github/scripts/tests/test_metric_snapshot_fixture.py @@ -0,0 +1,101 @@ +# 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. + +import copy +import importlib.util +from pathlib import Path +import unittest + +spec = importlib.util.spec_from_file_location( + "fixture", Path(__file__).resolve().parents[1] / "metric-snapshot-fixture.py" +) +fixture = importlib.util.module_from_spec(spec) +spec.loader.exec_module(fixture) + + +class ValidatorTests(unittest.TestCase): + def records(self): + return { + "schema": {"column_schemas": [ + {"name": "ts", "data_type": "TimestampMillisecond"}, + *[{"name": f"tag_{i:03d}", "data_type": "String"} for i in range(10)], + {"name": "val", "data_type": "Float64"}, + ]}, + "rows": [fixture.row(0, j, 50) for j in range(32)], + } + + def test_full_tag_coverage_and_heterogeneous_schemas(self): + for per_group in (50, 10000): + variants = {tuple(fixture.tags(i, per_group)) for i in range(per_group)} + self.assertEqual(set().union(*map(set, variants)), set(range(298))) + self.assertGreater(len(variants), 1) + self.assertTrue(all(len(tags) == 10 for tags in variants)) + + def test_accepts_engine_column_order_with_matching_values(self): + records = self.records() + order = [*range(1, 11), 0, 11] # The engine returns tags, timestamp, value. + records["schema"]["column_schemas"] = [ + records["schema"]["column_schemas"][index] for index in order + ] + records["rows"] = [[values[index] for index in order] for values in records["rows"]] + fixture.check_records(0, 50, records) + records["rows"][1][10] = 2 + with self.assertRaisesRegex(RuntimeError, "typed values"): + fixture.check_records(0, 50, records) + + def test_rejects_column_type_and_membership_changes(self): + fixture.check_records(0, 50, self.records()) + for mutation in ("type", "extra", "missing"): + records = self.records() + columns = records["schema"]["column_schemas"] + if mutation == "type": + columns[0]["data_type"] = "TimestampMicrosecond" + elif mutation == "extra": + columns.append({"name": "unexpected", "data_type": "String"}) + else: + del columns[1] + with self.subTest(mutation=mutation), self.assertRaisesRegex(RuntimeError, "schema"): + fixture.check_records(0, 50, records) + + def test_rejects_null_and_numeric_type_changes(self): + for value in ("", 0, False): + records = self.records() + records["rows"][0][-1] = value # Expected SQL NULL. + with self.subTest(value=value), self.assertRaisesRegex(RuntimeError, "typed values"): + fixture.check_records(0, 50, records) + records = self.records() + records["rows"][4][-1] = 1 # Numeric equality to 1.0 must not hide a type change. + with self.assertRaisesRegex(RuntimeError, "typed values"): + fixture.check_records(0, 50, records) + + def test_source_comparison_requires_same_ddl_and_remapped_ids(self): + source = {"fixture": {"tables_per_group": 50}, "schema_sha256": "schema", + "data_sha256": "data", "ddl_sha256": "ddl", + "table_ids": {"m000000": 1024, "phy_00": 1025}} + target = copy.deepcopy(source) + target["table_ids"] = {"m000000": 1041, "phy_00": 1042} + fixture.compare_source(target, source) + self.assertEqual(target["remapped_table_ids"], 2) + for field in ("schema_sha256", "data_sha256", "ddl_sha256"): + changed = copy.deepcopy(target) + changed[field] = "changed" + with self.subTest(field=field), self.assertRaisesRegex(RuntimeError, field): + fixture.compare_source(changed, source) + target["table_ids"]["phy_00"] = 1025 + with self.assertRaisesRegex(RuntimeError, "not all remapped"): + fixture.compare_source(target, source) + + +if __name__ == "__main__": + unittest.main() diff --git a/.github/scripts/tests/test_metric_snapshot_s3_bench.py b/.github/scripts/tests/test_metric_snapshot_s3_bench.py new file mode 100644 index 00000000000..c662ecbff76 --- /dev/null +++ b/.github/scripts/tests/test_metric_snapshot_s3_bench.py @@ -0,0 +1,74 @@ +#!/usr/bin/env python3 +# 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. + +"""Failure accounting for the metric snapshot runner.""" +import importlib.util +import json +from pathlib import Path +import tempfile +import unittest +from unittest.mock import patch + +SPEC = importlib.util.spec_from_file_location( + "metric_snapshot_bench", Path(__file__).resolve().parents[1] / "metric-snapshot-s3-bench.py") +bench = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(bench) + + +class PhaseTests(unittest.TestCase): + def test_docker_nanosecond_timestamps(self): + start = bench.docker_timestamp("2026-10-14T06:06:44.72000818Z") + end = bench.docker_timestamp("2026-10-14T06:06:46.123456789Z") + self.assertAlmostEqual((end - start).total_seconds(), 1.403448) + + def test_resource_interference_invalidates_success_and_preserves_fixture(self): + with tempfile.TemporaryDirectory() as directory: + runner = bench.Benchmark.__new__(bench.Benchmark) + runner.round_dir = Path(directory) + runner.server_pid = 0 + runner.phases = [] + runner.metrics = lambda: "" + runner.summary = lambda: None + fixture = runner.round_dir / "source-validation.json" + fixture.write_text('{"data_sha256":"fixture"}') + before = {"available": True, "pswpout_pages": 0, "oom_kill": 0} + after = {"available": True, "pswpout_pages": 1, "oom_kill": 0} + with patch.object(bench, "host_memory", side_effect=[before, after]): + with self.assertRaisesRegex(RuntimeError, "Resource interference"): + runner.phase("source-validation", lambda: None) + self.assertEqual(json.loads(fixture.read_text()), {"data_sha256": "fixture"}) + evidence = json.loads((runner.round_dir / "source-validation-resources.json").read_text()) + self.assertEqual(evidence["outcome"], "resource-invalid") + self.assertEqual(evidence["new_swapout_pages"], 1) + + def test_action_failure_keeps_failed_resource_evidence(self): + with tempfile.TemporaryDirectory() as directory: + runner = bench.Benchmark.__new__(bench.Benchmark) + runner.round_dir = Path(directory) + runner.server_pid = 0 + runner.phases = [] + runner.metrics = lambda: "" + runner.summary = lambda: None + def fail(): + raise RuntimeError("export failed") + with patch.object(bench, "host_memory", return_value={"available": False}): + with self.assertRaisesRegex(RuntimeError, "export failed"): + runner.phase("export", fail) + evidence = json.loads((runner.round_dir / "export-resources.json").read_text()) + self.assertEqual(evidence["outcome"], "failed") + + +if __name__ == "__main__": + unittest.main() diff --git a/.github/scripts/tests/test_metric_snapshot_s3_ops.py b/.github/scripts/tests/test_metric_snapshot_s3_ops.py new file mode 100644 index 00000000000..d22c09044b2 --- /dev/null +++ b/.github/scripts/tests/test_metric_snapshot_s3_ops.py @@ -0,0 +1,203 @@ +#!/usr/bin/env python3 +# 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. + +"""Offline safety checks for the metric snapshot workflow's budget and cleanup.""" + +import importlib.util +import io +import json +from pathlib import Path +import tempfile +import unittest +from unittest import mock + +SPEC = importlib.util.spec_from_file_location("snapshot_ops", Path(__file__).parents[1] / "metric-snapshot-s3-ops.py") +ops = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(ops) +PREFIX = "metric-snapshot-bench/123-1/" + + +class FakeS3: + def __init__(self): + self.objects = [{"Key": PREFIX + "current", "Size": 11}, {"Key": PREFIX + "data", "Size": 7}] + self.deleted = [] + + def get_paginator(self, operation): + assert operation == "list_objects_v2" + parent = self + + class Paginator: + def paginate(self, **kwargs): + assert kwargs["Bucket"] == "bucket" and kwargs["Prefix"].startswith(PREFIX) + # Keys outside the run prefix model a misbehaving listing and are always returned. + return iter([{"Contents": [item for item in parent.objects if not item["Key"].startswith(PREFIX) + or item["Key"].startswith(kwargs["Prefix"])]}]) + + return Paginator() + + def put_object(self, **kwargs): + self.objects.append({"Key": kwargs["Key"], "Size": len(kwargs["Body"])}) + return {} + + def get_object(self, **kwargs): + return {"Body": io.BytesIO(b"metric snapshot benchmark")} + + def delete_object(self, **kwargs): + self.deleted.append(kwargs) + self.objects = [item for item in self.objects if item["Key"] != kwargs["Key"]] + + +class OperationsTest(unittest.TestCase): + def env(self, **changes): + return dict(GREPTIMEDB_IMAGE="greptime/greptimedb-nightly@sha256:" + "a" * 64, + EXPECTED_COMMIT="b" * 40, LOGICAL_TABLES_PER_GROUP="10000", + BENCHMARK_TIMEOUT_MINUTES="240", PRIOR_SPEND_USD="0", + S3_REGION="ap-southeast-1", RUNNER_REGION="ap-southeast-1", S3_BUCKET="bucket", + RUN_PREFIX=PREFIX) | changes + + def test_budget_includes_provision_cleanup_and_other_cost_reserve(self): + plan = ops.plan(self.env()) + self.assertEqual(plan["attempt_reserve_usd"], "6.4392") + ops.plan(self.env(PRIOR_SPEND_USD="13.5608")) + for changes in ({"PRIOR_SPEND_USD": "13.5609"}, {"PRIOR_SPEND_USD": "15"}, + {"PRIOR_SPEND_USD": "NaN"}, {"PRIOR_SPEND_USD": "-1"}, + {"BENCHMARK_TIMEOUT_MINUTES": "241"}, {"LOGICAL_TABLES_PER_GROUP": "49"}, + {"LOGICAL_TABLES_PER_GROUP": "10001"}, {"S3_REGION": "us-east-1"}, + {"GREPTIMEDB_IMAGE": "greptime/greptimedb:latest"}, + {"GREPTIMEDB_IMAGE": "docker.io/greptime/greptimedb@sha256:" + "a" * 64}, + {"GREPTIMEDB_IMAGE": "greptime/greptimedb:nightly@sha256:" + "a" * 64}, + {"EXPECTED_COMMIT": "abcdef"}): + with self.subTest(changes=changes), self.assertRaises(ValueError): + ops.plan(self.env(**changes)) + + def watcher_env(self): + return self.env(GITHUB_REPOSITORY="GreptimeTeam/greptimedb", GITHUB_RUN_ID="123", + GITHUB_RUN_ATTEMPT="2", GH_TOKEN="test-token", BENCHMARK_TIMEOUT_MINUTES="30") + + def test_watcher_completed_job_uses_original_attempt_and_does_not_wait(self): + for conclusion in ("success", "failure", "cancelled"): + response = io.BytesIO(json.dumps({"jobs": [{"name": "benchmark", "status": "completed", + "conclusion": conclusion}]}).encode()) + with self.subTest(conclusion=conclusion), \ + mock.patch.object(ops.urllib.request, "urlopen", return_value=response) as fetch, \ + mock.patch.object(ops.time, "time", return_value=1005), \ + mock.patch.object(ops.time, "sleep") as sleep: + self.assertEqual(ops.wait_for_benchmark(self.watcher_env(), 1000), + {"benchmark_conclusion": conclusion, "owned_attempt": "1"}) + self.assertIn("/runs/123/attempts/1/jobs?", fetch.call_args.args[0].full_url) + sleep.assert_not_called() + + def test_watcher_queue_reaches_durable_deadline(self): + deadline = 1000 + (45 + 30 + 10) * 60 + clock = [deadline - 31] + + def advance(seconds): + self.assertLessEqual(seconds, 30) + clock[0] += seconds + + def queued(*args, **kwargs): + return io.BytesIO(b'{"jobs": [{"name": "benchmark", "status": "queued"}]}') + + with mock.patch.object(ops.urllib.request, "urlopen", side_effect=queued) as fetch, \ + mock.patch.object(ops.time, "time", side_effect=lambda: clock[0]), \ + mock.patch.object(ops.time, "sleep", side_effect=advance), \ + self.assertRaisesRegex(TimeoutError, "proceeding to resource cleanup"): + ops.wait_for_benchmark(self.watcher_env(), 1000) + self.assertEqual(clock[0], deadline) + self.assertEqual(fetch.call_count, 2) + + def test_watcher_api_error_fails_immediately_for_cleanup(self): + with mock.patch.object(ops.urllib.request, "urlopen", side_effect=OSError("API unavailable")), \ + mock.patch.object(ops.time, "time", return_value=1005), \ + mock.patch.object(ops.time, "sleep") as sleep, \ + self.assertRaises(OSError): + ops.wait_for_benchmark(self.watcher_env(), 1000) + sleep.assert_not_called() + + def test_cleanup_deletes_every_listed_object(self): + s3 = FakeS3() + result = ops.cleanup(s3, "bucket", PREFIX) + self.assertEqual(result["objects"], 0) + self.assertEqual(s3.deleted, [{"Bucket": "bucket", "Key": PREFIX + "current"}, + {"Bucket": "bucket", "Key": PREFIX + "data"}]) + + def test_probe_refuses_occupied_prefix_and_removes_its_object(self): + with self.assertRaisesRegex(ValueError, "already occupied"): + ops.probe(FakeS3(), "bucket", PREFIX) + s3 = FakeS3() + s3.objects = [] + self.assertEqual(ops.probe(s3, "bucket", PREFIX)["objects"], 0) + self.assertEqual(s3.deleted, [{"Bucket": "bucket", "Key": PREFIX + "_probe"}]) + + def test_failed_probe_delete_still_has_hosted_cleanup_path(self): + s3 = FakeS3() + s3.objects = [] + with mock.patch.object(s3, "delete_object", side_effect=OSError("temporary delete failure")), \ + self.assertRaises(OSError): + ops.probe(s3, "bucket", PREFIX) + self.assertEqual([item["Key"] for item in s3.objects], [PREFIX + "_probe"]) + workflow = (Path(__file__).parents[2] / "workflows/metric-snapshot-s3-bench.yml").read_text() + teardown = workflow.split(" teardown:\n", 1)[1] + job_if = next(line.strip() for line in teardown.splitlines() if line.startswith(" if:")) + self.assertEqual(job_if, "if: ${{ always() && needs.validate.outputs.run_id != '' }}") + ec2 = teardown.split(" - name: Delete EC2, its temporary disks and runner registration\n", 1)[1] + self.assertEqual(ec2.splitlines()[0].strip(), "if: ${{ always() && needs.provision.result != 'skipped' }}") + self.assertIn(" RUN_PREFIX: ${{ needs.validate.outputs.run_prefix }}\n", teardown.split(" steps:", 1)[0]) + s3_cleanup = teardown.split(" - name: Remove and verify only this run's S3 prefix\n", 1)[1].split(" - ", 1)[0] + self.assertEqual(s3_cleanup.splitlines()[0].strip(), "if: always()") + self.assertIn("run: uv run greptimedb/.github/scripts/metric-snapshot-s3-ops.py cleanup\n", s3_cleanup) + self.assertEqual(ops.cleanup(s3, "bucket", PREFIX)["objects"], 0) + + def test_round_inventory_excludes_other_rounds_while_cleanup_covers_the_run(self): + s3 = FakeS3() + s3.objects = [{"Key": PREFIX + "smoke/pack-000000.bin", "Size": 3}, + {"Key": PREFIX + "full/pack-000000.bin", "Size": 5}, + {"Key": PREFIX + "full/index.json", "Size": 2}] + full = ops.inventory(s3, "bucket", PREFIX, "full") + self.assertEqual((full["prefix"], full["objects"], full["bytes"]), (PREFIX + "full/", 2, 7)) + self.assertEqual(ops.inventory(s3, "bucket", PREFIX, "smoke")["bytes"], 3) + self.assertEqual(ops.inventory(s3, "bucket", PREFIX)["objects"], 3) + for round in ("", "../full", "full/"): + with self.subTest(round=round), self.assertRaises(ValueError): + ops.inventory(s3, "bucket", PREFIX, round) + self.assertEqual(ops.cleanup(s3, "bucket", PREFIX)["objects"], 0) + self.assertEqual(len(s3.deleted), 3) + + def test_bad_prefix_and_out_of_scope_response_cannot_delete(self): + for prefix in ("", "metric-snapshot-bench/", "metric-snapshot-bench/123-1", PREFIX + "../", "oswal-bench/123-1/"): + s3 = FakeS3() + with self.subTest(prefix=prefix), self.assertRaises(ValueError): + ops.cleanup(s3, "bucket", prefix) + self.assertEqual(s3.deleted, []) + s3 = FakeS3() + s3.objects[0]["Key"] = "other-run/current" + with self.assertRaises(ValueError): + ops.inventory(s3, "bucket", PREFIX, "full") + with self.assertRaises(ValueError): + ops.cleanup(s3, "bucket", PREFIX) + self.assertEqual(s3.deleted, []) + + def test_credential_crossing_read_boundary_blocks_artifact_upload(self): + with tempfile.TemporaryDirectory() as directory: + logfile = Path(directory) / "server.log" + logfile.write_bytes(b"a" * (1024 * 1024 - 4) + b"private-key" + b"tail") + with self.assertRaisesRegex(ValueError, "upload is blocked"): + ops.scrub(directory, {"AWS_SECRET_ACCESS_KEY": "private-key"}) + logfile.write_text("safe log\n") + ops.scrub(directory, {"AWS_SECRET_ACCESS_KEY": "private-key"}) + + +if __name__ == "__main__": + unittest.main() diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml index deba5010c80..2aa460ec0ca 100644 --- a/.github/workflows/checks.yml +++ b/.github/workflows/checks.yml @@ -88,6 +88,8 @@ jobs: run: .github/scripts/run-fuzz-targets-test.sh - name: Run CI comment-command tests run: python3 .github/scripts/tests/test_ci_slash.py + - name: Run metric snapshot benchmark tests + run: python3 -m unittest discover -s .github/scripts/tests -p 'test_metric_snapshot*.py' - name: Run query-regression Python tooling tests run: | python3 tests/perf/test_query_regression_runner_compaction_toctou.py diff --git a/.github/workflows/metric-snapshot-s3-bench.yml b/.github/workflows/metric-snapshot-s3-bench.yml new file mode 100644 index 00000000000..e4a12dde69f --- /dev/null +++ b/.github/workflows/metric-snapshot-s3-bench.yml @@ -0,0 +1,230 @@ +name: Metric Snapshot S3 Benchmark +on: + workflow_dispatch: + inputs: + greptimedb_image: + description: 'Immutable greptime/greptimedb or greptimedb-nightly image (repository@sha256:digest)' + required: true + type: string + expected_commit: + description: 'Full source commit SHA reported by the image' + required: true + type: string + logical_tables_per_group: + description: 'Logical tables per physical table, 50..10000; 10 physical tables, 32 rows per logical table' + type: number + default: 10000 + benchmark_timeout_minutes: + description: 'Setup, 500-table smoke and full run timeout, 30..240 minutes' + type: number + default: 240 + prior_spend_usd: + description: 'Cumulative estimated cloud spend for all earlier attempts in this $20 campaign' + required: true + type: string + default: '0' +permissions: + contents: read +concurrency: + group: metric-snapshot-s3-bench + cancel-in-progress: false +env: + GREPTIMEDB_IMAGE: ${{ inputs.greptimedb_image }} + EXPECTED_COMMIT: ${{ inputs.expected_commit }} + LOGICAL_TABLES_PER_GROUP: ${{ inputs.logical_tables_per_group }} + BENCHMARK_TIMEOUT_MINUTES: ${{ inputs.benchmark_timeout_minutes }} + PRIOR_SPEND_USD: ${{ inputs.prior_spend_usd }} + RUN_PREFIX: metric-snapshot-bench/${{ github.run_id }}-${{ github.run_attempt }}/ + S3_BUCKET: ${{ vars.AWS_CI_TEST_BUCKET }} + S3_REGION: ${{ vars.AWS_CI_TEST_BUCKET_REGION }} + RUNNER_REGION: ${{ vars.EC2_RUNNER_REGION }} +jobs: + validate: + runs-on: ubuntu-latest + timeout-minutes: 15 + outputs: + instance_type: ${{ steps.runner.outputs.instance_type }} + run_id: ${{ steps.identity.outputs.run_id }} + run_prefix: ${{ steps.identity.outputs.run_prefix }} + steps: + - uses: actions/checkout@v4 + with: + path: greptimedb + persist-credentials: false + - name: Test benchmark boundaries + run: | + python3 greptimedb/.github/scripts/tests/test_metric_snapshot_s3_ops.py + python3 greptimedb/.github/scripts/tests/test_metric_snapshot_fixture.py + python3 greptimedb/.github/scripts/tests/test_metric_snapshot_s3_bench.py + - name: Validate pinned version, scale, region and cumulative budget + run: python3 greptimedb/.github/scripts/metric-snapshot-s3-ops.py plan + - name: Bind resource ownership to this validated attempt + id: identity + run: | + echo "run_id=$GITHUB_RUN_ID-$GITHUB_RUN_ATTEMPT" >> "$GITHUB_OUTPUT" + echo "run_prefix=$RUN_PREFIX" >> "$GITHUB_OUTPUT" + - uses: astral-sh/setup-uv@v6 + - name: Check S3 access before allocating EC2 + env: + AWS_ACCESS_KEY_ID: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} + AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} + run: uv run greptimedb/.github/scripts/metric-snapshot-s3-ops.py probe + - name: Resolve the fixed benchmark runner + id: runner + run: python3 greptimedb/.github/scripts/ci-runner-config.py --provider AWS --instance-type c6id.4xlarge + provision: + needs: validate + runs-on: ubuntu-latest + timeout-minutes: 45 + outputs: + label: ${{ steps.aws.outputs.label }} + instance_id: ${{ steps.aws.outputs.instance_id }} + runner_name: ${{ steps.aws.outputs.runner_name }} + started_at: ${{ steps.start.outputs.started_at }} + steps: + - uses: actions/checkout@v4 + with: + path: greptimedb + persist-credentials: false + - uses: aws-actions/configure-aws-credentials@v4 + with: + aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }} + aws-region: ${{ env.RUNNER_REGION }} + - uses: astral-sh/setup-uv@v6 + - name: Record allocation start + id: start + run: echo "started_at=$(date +%s)" >> "$GITHUB_OUTPUT" + - name: Create ephemeral EC2 runner + id: aws + env: + GH_PERSONAL_ACCESS_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + AWS_REGION: ${{ env.RUNNER_REGION }} + AWS_EC2_IMAGE_ID: ${{ vars.BENCHMARK_EC2_IMAGE_ID || 'auto' }} + AWS_EC2_INSTANCE_TYPE: ${{ needs.validate.outputs.instance_type }} + AWS_EC2_SUBNET_ID: ${{ vars.EC2_RUNNER_SUBNET_ID }} + AWS_EC2_SECURITY_GROUP_ID: ${{ vars.EC2_RUNNER_SECURITY_GROUP_ID }} + AWS_EC2_SYSTEM_DISK_GIB: '60' + AWS_EC2_TTL_HOURS: '6' + AWS_EC2_RUN_ID: ${{ needs.validate.outputs.run_id }} + run: uv run greptimedb/.github/scripts/aws-ec2-runner-provision.py + benchmark: + needs: + - validate + - provision + runs-on: ${{ needs.provision.outputs.label }} + timeout-minutes: ${{ fromJSON(format('{0}', inputs.benchmark_timeout_minutes)) }} + env: + RUN_PREFIX: ${{ needs.validate.outputs.run_prefix }} + RESULTS: ${{ github.workspace }}/metric-snapshot-results + DISK_ROOT: /mnt/metric-snapshot-disk + DATA_ROOT: /mnt/metric-snapshot-disk/data + steps: + - uses: actions/checkout@v4 + with: + path: greptimedb + persist-credentials: false + - name: Set up benchmark dependencies + uses: ./greptimedb/.github/actions/setup-benchmark + - uses: astral-sh/setup-uv@v6 + - name: Mount the temporary local File data disk + run: | + device=$(lsblk -dnpo NAME,MODEL | awk '/Instance Storage/ {print $1; exit}') + test -n "$device" + if [[ "$(lsblk -nro TYPE "$device" | wc -l)" -ne 1 ]] || findmnt --source "$device" >/dev/null; then + echo 'Refusing partitioned or mounted instance storage' >&2 + exit 1 + fi + probe_status=0 + sudo -n blkid -p "$device" || probe_status=$? + if ((probe_status != 2)); then + echo 'Refusing a disk with an existing signature or an unsuccessful signature probe' >&2 + exit 1 + fi + sudo -n mkfs.ext4 -q "$device" + sudo -n mkdir -p "$DISK_ROOT" + sudo -n mount -o noatime "$device" "$DISK_ROOT" + sudo -n chown "$(id -u):$(id -g)" "$DISK_ROOT" + findmnt -no SOURCE,FSTYPE -T "$DISK_ROOT" + - name: Probe the real same-region S3 bucket + env: + AWS_ACCESS_KEY_ID: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} + AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} + run: uv run greptimedb/.github/scripts/metric-snapshot-s3-ops.py probe + - name: Run the 500-table smoke, then one full packed export and restore + env: + AWS_ACCESS_KEY_ID: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} + AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} + run: python3 greptimedb/.github/scripts/metric-snapshot-s3-bench.py run + - name: Check results for credentials + id: scrub + if: always() + env: + AWS_ACCESS_KEY_ID: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} + AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} + run: python3 greptimedb/.github/scripts/metric-snapshot-s3-ops.py scrub + - name: Publish the readable result + if: ${{ always() && steps.scrub.outcome == 'success' }} + run: | + if test -f "$RESULTS/summary.md"; then + cat "$RESULTS/summary.md" >> "$GITHUB_STEP_SUMMARY" + fi + - name: Upload bounded-retention evidence + if: ${{ always() && steps.scrub.outcome == 'success' }} + uses: actions/upload-artifact@v4 + with: + name: metric-snapshot-s3-${{ github.run_id }}-${{ github.run_attempt }} + path: ${{ env.RESULTS }} + if-no-files-found: warn + retention-days: 14 + teardown: + needs: + - validate + - provision + # Start alongside the benchmark so an offline runner cannot hold cleanup in its queue. + if: ${{ always() && needs.validate.outputs.run_id != '' }} + runs-on: ubuntu-latest + timeout-minutes: 330 + permissions: + contents: read + actions: read + env: + RUN_PREFIX: ${{ needs.validate.outputs.run_prefix }} + steps: + - uses: actions/checkout@v4 + with: + path: greptimedb + persist-credentials: false + - uses: astral-sh/setup-uv@v6 + - name: Wait for completion within the allocation deadline + if: ${{ !cancelled() && needs.provision.result == 'success' && needs.provision.outputs.started_at != '' }} + env: + GH_TOKEN: ${{ github.token }} + STARTED_AT: ${{ needs.provision.outputs.started_at }} + run: python3 greptimedb/.github/scripts/metric-snapshot-s3-ops.py wait-for-benchmark --started-at "$STARTED_AT" + - uses: aws-actions/configure-aws-credentials@v4 + if: ${{ always() && needs.provision.result != 'skipped' }} + with: + aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }} + aws-region: ${{ env.RUNNER_REGION }} + - name: Delete EC2, its temporary disks and runner registration + if: ${{ always() && needs.provision.result != 'skipped' }} + env: + GH_PERSONAL_ACCESS_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + AWS_REGION: ${{ env.RUNNER_REGION }} + AWS_EC2_INSTANCE_ID: ${{ needs.provision.outputs.instance_id }} + AWS_EC2_RUNNER_NAME: ${{ needs.provision.outputs.runner_name }} + AWS_EC2_RUN_ID: ${{ needs.validate.outputs.run_id }} + run: uv run greptimedb/.github/scripts/aws-ec2-runner-teardown.py + - name: Remove and verify only this run's S3 prefix + if: always() + env: + AWS_ACCESS_KEY_ID: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} + AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} + run: uv run greptimedb/.github/scripts/metric-snapshot-s3-ops.py cleanup + - name: Estimate the cumulative cloud cost + if: ${{ always() && needs.provision.outputs.started_at != '' }} + env: + STARTED_AT: ${{ needs.provision.outputs.started_at }} + run: python3 greptimedb/.github/scripts/metric-snapshot-s3-ops.py cost --started-at "$STARTED_AT" >> "$GITHUB_STEP_SUMMARY" diff --git a/docs/how-to/metric-snapshot-s3-bench.md b/docs/how-to/metric-snapshot-s3-bench.md new file mode 100644 index 00000000000..86468bf3ffd --- /dev/null +++ b/docs/how-to/metric-snapshot-s3-bench.md @@ -0,0 +1,79 @@ +# Metric snapshot benchmark on AWS S3 + +The manual `Metric Snapshot S3 Benchmark` workflow tests packed Metric export and +restore against real S3 in `ap-southeast-1`. It must first exist on the repository's +default branch for GitHub to register `workflow_dispatch`. It was validated on AWS +by run [38544464731](https://github.com/GreptimeTeam/greptimedb/actions/runs/38544464731); +[#9573](https://github.com/GreptimeTeam/greptimedb/pull/9573) records its conditions and results. + +## Run + +Use the repository's existing EC2 runner configuration and CI S3 credentials. +The workflow passes credentials only through CI; no local AWS credentials are +needed. Source, target and CLI use the same immutable image. This image's source +SHA and required CLI capabilities have been checked: + +```bash +gh workflow run metric-snapshot-s3-bench.yml \ + --repo GreptimeTeam/greptimedb --ref main \ + -f greptimedb_image=greptime/greptimedb-nightly@sha256:20a1182bed0f452e6d573ca110b313e0b5cb76de7aedfa41a0c339963d4683d3 \ + -f expected_commit=f690bbb2040865c0d22b0b57d51612b60ae27cd0 \ + -f logical_tables_per_group=10000 \ + -f benchmark_timeout_minutes=240 \ + -f prior_spend_usd=0 +``` + +Start a new dispatch for each attempt and update `prior_spend_usd` to include all +earlier cloud attempts in the campaign, including failed setup and cleanup. +Do not reset the estimate to zero when retrying. The workflow accepts 50–10,000 +logical tables per group and a 30–240 minute benchmark job timeout. + +The runner is fixed at `c6id.4xlarge` (16 vCPU, 32 GiB RAM). At the recorded +Singapore Linux on-demand estimate of $0.9408/hour, the default attempt reserves +$6.4392: four benchmark hours, 45 minutes of provisioning, 10 minutes for queueing, +20 minutes of cleanup, and $1.50 for S3, disk and IPv4 costs. Preflight rejects prior spend at or above +$15, or an attempt whose reserve would exceed the cumulative $20 budget. +These estimates and checks are not an AWS hard spending cap or an AWS bill. + +## Workload and validation + +- A 10-physical-table × 50-logical-table smoke runs first. Its 500 logical tables + cover all 298 string tag columns before the larger workload starts. +- The default full run uses 10 physical tables and 100,000 logical tables, with + 32 rows per logical table: 3.2 million rows. Each logical schema has 10 + overlapping tags, a timestamp and a DOUBLE; schemas vary between tables. +- Source and target use local File storage on the temporary NVMe disk. Only + snapshots use S3. Source validation finishes before the source process stops; + restoration then starts a separate target process. +- Export uses `--pack-metric-data`, parallelism 4, chunk parallelism 1, and one + complete time chunk. Import uses task parallelism 1 and automatic batch DDL. +- Every source and restored logical table is checked for membership, schema, + DDL and typed values including NULLs. Target placeholders force table ID + remapping; successful batch DDL and removal of import resume state are checked. + +## Results and cleanup + +The job summary reports complete export, structural verify and import durations; +data generation and full data validation are timed separately. Evidence artifacts +are credential-scanned and retained for 14 days. Keep conclusions limited to +this version's absolute AWS results; this workflow does not compare old versions. + +Artifacts include sampled CPU/RSS, host memory, swap/OOM counters, server metrics, +CLI logs and S3 object/byte inventories. Configured swap is not RAM. A phase with +new swap-out or OOM fails and retains its evidence. OpenDAL operation counters +are not complete AWS HTTP request counts: multipart calls, retries and CLI traffic +are not fully covered. Object inventory bytes are not transferred bytes. + +A separate GitHub-hosted teardown job starts when provisioning finishes and polls +the benchmark job. Its deadline includes queue time: allocation start plus 45 +minutes, the benchmark timeout and 10 minutes. Completion, cancellation, an API +error or the deadline leads to cleanup. Failed provisioning skips the wait. +It terminates the EC2 instance and its disks, unregisters the runner, deletes every +object listed under the run's unique prefix and verifies the listing is empty. Like +the object store WAL benchmark, it needs only list, put, get and delete on the CI +bucket. It does not abort incomplete multipart uploads or remove noncurrent object +versions; those depend on the bucket's lifecycle configuration. Check teardown completion before another +attempt. A deadline may leave the GitHub benchmark job queued after its resources +are removed; cancel that workflow manually. If cleanup fails, rerun the teardown +job only; it keeps the original attempt's resource identity and deadline. A new benchmark needs a new dispatch because teardown +removes its runner. The daily EC2 janitor is only a fallback and does not clean S3. From 0dc81777e3de532a711fac26c53b78ca2bee8ab2 Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Thu, 15 Oct 2026 10:07:35 +0000 Subject: [PATCH 010/239] fix(pprof): initialize framehop before signal sampling (#9612) fix(pprof): pin upstream framehop initialization fix Pin pprof workspace-wide to 01cff82d containing tikv/pprof-rs#282. Initialize framehop before sampling instead of allocating and discovering shared libraries in the first SIGPROF handler. Regenerate the lockfile without unrelated dependency changes. Signed-off-by: Lei, HUANG --- Cargo.lock | 3 +-- Cargo.toml | 4 ++++ 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index edf460e4679..f9410e292ba 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11106,8 +11106,7 @@ checksum = "439ee305def115ba05938db6eb1644ff94165c5ab5e9420d1c1bcedbba909391" [[package]] name = "pprof" version = "0.15.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "38a01da47675efa7673b032bf8efd8214f1917d89685e07e395ab125ea42b187" +source = "git+https://github.com/tikv/pprof-rs?rev=01cff82dbe6fe110a707bf2b38d8ebb1d14a18f8#01cff82dbe6fe110a707bf2b38d8ebb1d14a18f8" dependencies = [ "aligned-vec", "backtrace", diff --git a/Cargo.toml b/Cargo.toml index 702498ee5e5..9fa5d013272 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -384,6 +384,10 @@ datafusion-pg-catalog = { git = "https://github.com/datafusion-contrib/datafusio arrow-pg = { git = "https://github.com/datafusion-contrib/datafusion-postgres.git", rev = "ffb14a52b3d7c7489559812e1ec38096f225ce93" } sqlparser = { git = "https://github.com/GreptimeTeam/sqlparser-rs.git", rev = "9e9019bb1c7040ed956f654e39378dd42ab17884" } +# pprof 0.15.0 initializes framehop inside SIGPROF; pin the unreleased fix +# that initializes it before sampling. Remove once tikv/pprof-rs#282 is released. +pprof = { git = "https://github.com/tikv/pprof-rs", rev = "01cff82dbe6fe110a707bf2b38d8ebb1d14a18f8" } + # Temporary: use the GreptimeTeam fork of tikv-jemalloc-sys embedding # jemalloc 5.3.1 + backport of 54f22c83 ("Initialize TSD tcache before # enabling it"), fixing the TSD/tcache enable-before-init reentrancy From ad823db020cc532cfd94e466d9c1be8037efe182 Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Thu, 15 Oct 2026 10:10:50 +0000 Subject: [PATCH 011/239] fix(pipeline)!: disable VRL host access by default (#9610) fix(pipeline): disable VRL host access by default Signed-off-by: shuiyisong --- Cargo.lock | 2 + config/config.md | 2 + config/frontend.example.toml | 4 + config/standalone.example.toml | 4 + src/frontend/src/frontend.rs | 12 + src/frontend/src/instance/log_handler.rs | 5 +- src/pipeline/Cargo.toml | 2 + src/pipeline/src/etl.rs | 11 +- src/pipeline/src/etl/README.md | 21 ++ src/pipeline/src/etl/processor.rs | 17 +- .../src/etl/processor/vrl_processor.rs | 266 +++++++++++++++++- src/pipeline/src/lib.rs | 2 +- src/pipeline/src/manager/pipeline_cache.rs | 8 +- src/pipeline/src/manager/pipeline_operator.rs | 12 +- src/pipeline/src/manager/table.rs | 18 +- src/pipeline/src/options.rs | 5 + src/standalone/src/options.rs | 13 + tests-integration/Cargo.toml | 2 +- tests-integration/src/standalone.rs | 10 + tests-integration/tests/http.rs | 68 +++++ 20 files changed, 437 insertions(+), 47 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index f9410e292ba..b696ee8224d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10898,6 +10898,8 @@ dependencies = [ "snafu 0.8.6", "sql", "table", + "temp-env", + "tempfile", "tokio", "urlencoding", "vrl", diff --git a/config/config.md b/config/config.md index 6faae09c347..74a94891d7d 100644 --- a/config/config.md +++ b/config/config.md @@ -268,6 +268,7 @@ | `tracing.tokio_console_addr` | String | Unset | The tokio console address. | | `pipeline` | -- | -- | The pipeline options. | | `pipeline.cache_ttl` | String | `10s` | Time to live of the local pipeline cache. Default is `10s`. | +| `pipeline.allow_vrl_host_access` | Bool | `false` | Allow VRL to access the server environment, filesystem, hostname, local timezone,
network, and logs. Default is false. Enable only for trusted pipeline authors.
Applies to pipeline creation, execution, and dry-runs. Requires a restart. | | `event_recorder` | -- | -- | Configuration options for the event recorder. | | `event_recorder.ttl` | String | `90d` | TTL for the events table that will be used to store the events. Default is `90d`. | | `event_recorder.event_types` | Array | All | Event types to record. Current available event types: `create_database`,
`alter_database`, `drop_database`, `create_flow`, `drop_flow`,
`create_table`, `create_logical_tables`, `alter_table`, `alter_logical_tables`,
`drop_table`, `undrop_table`, `purge_dropped_table`, `truncate_table`,
`create_view`, `drop_view`, `admin_function`, `reconcile_catalog`,
`reconcile_database`, `reconcile_table`, `reconcile_logical_tables`.
When omitted, all current and future event types are recorded.
Set to an empty array to disable event recording. | @@ -429,6 +430,7 @@ | `memory.enable_heap_profiling` | Bool | `true` | Whether to enable heap profiling activation during startup.
When enabled, heap profiling will be activated if the `MALLOC_CONF` environment variable
is set to "prof:true,prof_active:false". The official image adds this env variable.
Default is true. | | `pipeline` | -- | -- | The pipeline options. | | `pipeline.cache_ttl` | String | `10s` | Time to live of the frontend-local pipeline cache. A pipeline created or deleted on
another frontend takes effect on this one after at most this duration. | +| `pipeline.allow_vrl_host_access` | Bool | `false` | Allow VRL to access the server environment, filesystem, hostname, local timezone,
network, and logs. Default is false. Enable only for trusted pipeline authors.
Applies to pipeline creation, execution, and dry-runs. Requires a restart. | | `event_recorder` | -- | -- | Configuration options for the event recorder. | | `event_recorder.ttl` | String | `90d` | TTL for the events table that will be used to store the events. Default is `90d`. | | `event_recorder.event_types` | Array | All | Event types to record. Current available event type: `admin_function`.
When omitted, all current and future event types are recorded.
Set to an empty array to disable event recording. | diff --git a/config/frontend.example.toml b/config/frontend.example.toml index c42d254671c..24aef8bd796 100644 --- a/config/frontend.example.toml +++ b/config/frontend.example.toml @@ -512,6 +512,10 @@ enable_heap_profiling = true ## Time to live of the frontend-local pipeline cache. A pipeline created or deleted on ## another frontend takes effect on this one after at most this duration. cache_ttl = "10s" +## Allow VRL to access the server environment, filesystem, hostname, local timezone, +## network, and logs. Default is false. Enable only for trusted pipeline authors. +## Applies to pipeline creation, execution, and dry-runs. Requires a restart. +allow_vrl_host_access = false ## Configuration options for the event recorder. [event_recorder] diff --git a/config/standalone.example.toml b/config/standalone.example.toml index 82fb3d2c489..aa5d09c730c 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -1095,6 +1095,10 @@ ttl = "90d" [pipeline] ## Time to live of the local pipeline cache. Default is `10s`. cache_ttl = "10s" +## Allow VRL to access the server environment, filesystem, hostname, local timezone, +## network, and logs. Default is false. Enable only for trusted pipeline authors. +## Applies to pipeline creation, execution, and dry-runs. Requires a restart. +allow_vrl_host_access = false ## Configuration options for the event recorder. [event_recorder] diff --git a/src/frontend/src/frontend.rs b/src/frontend/src/frontend.rs index 2ccbc3915e5..15e4ad66ca5 100644 --- a/src/frontend/src/frontend.rs +++ b/src/frontend/src/frontend.rs @@ -229,6 +229,18 @@ mod tests { type GrpcStream = Pin> + Send + Sync + 'static>>; + #[test] + fn test_pipeline_host_access_config() { + let defaults: FrontendOptions = toml::from_str("").unwrap(); + assert!(!defaults.pipeline.allow_vrl_host_access); + let options: FrontendOptions = + toml::from_str("[pipeline]\nallow_vrl_host_access = true").unwrap(); + assert!(options.pipeline.allow_vrl_host_access); + let serialized = toml::to_string(&options).unwrap(); + let restored: FrontendOptions = toml::from_str(&serialized).unwrap(); + assert_eq!(options.pipeline, restored.pipeline); + } + #[test] fn test_write_retry_config() { let defaults: FrontendOptions = toml::from_str("").unwrap(); diff --git a/src/frontend/src/instance/log_handler.rs b/src/frontend/src/instance/log_handler.rs index 44738b9dcc1..7b6a8fcb8ed 100644 --- a/src/frontend/src/instance/log_handler.rs +++ b/src/frontend/src/instance/log_handler.rs @@ -23,7 +23,6 @@ use auth::{ use client::Output; use common_error::ext::BoxedError; use datatypes::timestamp::TimestampNanosecond; -use pipeline::pipeline_operator::PipelineOperator; use pipeline::{Pipeline, PipelineInfo, PipelineVersion}; use servers::error::{ AuthSnafu, Error as ServerError, ExecuteGrpcRequestSnafu, PipelineSnafu, Result as ServerResult, @@ -148,7 +147,9 @@ impl PipelineHandler for Instance { } fn build_pipeline(&self, pipeline: &str) -> ServerResult { - PipelineOperator::build_pipeline(pipeline).context(PipelineSnafu) + self.pipeline_operator + .build_pipeline(pipeline) + .context(PipelineSnafu) } async fn get_pipeline_str( diff --git a/src/pipeline/Cargo.toml b/src/pipeline/Cargo.toml index 9ebe51ecf24..43083d07e70 100644 --- a/src/pipeline/Cargo.toml +++ b/src/pipeline/Cargo.toml @@ -61,6 +61,8 @@ yaml-rust = "0.4" [dev-dependencies] catalog = { workspace = true, features = ["testing"] } criterion = { workspace = true, features = ["html_reports"] } +temp-env.workspace = true +tempfile.workspace = true [[bench]] name = "processor" diff --git a/src/pipeline/src/etl.rs b/src/pipeline/src/etl.rs index f78a515fb59..ba68fde79ee 100644 --- a/src/pipeline/src/etl.rs +++ b/src/pipeline/src/etl.rs @@ -42,8 +42,8 @@ use crate::etl::transform::transformer::greptime::{ }; use crate::tablesuffix::TableSuffixTemplate; use crate::{ - ContextOpt, GreptimeTransformer, IdentityTimeIndex, PipelineContext, SchemaInfo, - unwrap_or_continue_if_err, + ContextOpt, GreptimeTransformer, IdentityTimeIndex, PipelineContext, PipelineOptions, + SchemaInfo, unwrap_or_continue_if_err, }; const DESCRIPTION: &str = "description"; @@ -60,6 +60,11 @@ pub enum Content<'a> { } pub fn parse(input: &Content) -> Result { + parse_with_options(input, &PipelineOptions::default()) +} + +/// Parses a pipeline using the server's compilation policy. +pub fn parse_with_options(input: &Content, options: &PipelineOptions) -> Result { match input { Content::Yaml(str) => { let docs = YamlLoader::load_from_str(str).context(YamlLoadSnafu)?; @@ -73,7 +78,7 @@ pub fn parse(input: &Content) -> Result { let doc_version = (&doc[DOC_VERSION]).try_into()?; let processors = if let Some(v) = doc[PROCESSORS].as_vec() { - v.try_into()? + Processors::from_yaml(v, options.allow_vrl_host_access)? } else { Processors::default() }; diff --git a/src/pipeline/src/etl/README.md b/src/pipeline/src/etl/README.md index e512ea31f7b..020bcdc304d 100644 --- a/src/pipeline/src/etl/README.md +++ b/src/pipeline/src/etl/README.md @@ -2,6 +2,27 @@ ETL capability +## VRL host access + +VRL host access is disabled by default. The server configuration option +`[pipeline].allow_vrl_host_access = true` restores the full VRL standard library +for trusted pipeline authors. It applies to both stored pipelines and dry-runs, +requires a server restart, and cannot be overridden in a pipeline definition. + +With host access disabled, environment, hostname, system timezone, filesystem, +network, and logging functions are unavailable. This includes `get_env_var`, +`get_hostname`, `get_timezone_name`, `http_request`, `dns_lookup`, `reverse_dns`, +`parse_proto`, `encode_proto`, `parse_etld`, `validate_json_schema`, and `log`. +Grok parsing supports inline +aliases but rejects nonempty `alias_sources`. Timestamp parsing and formatting +use UTC by default; an explicit timezone must be a constant named timezone such +as `Asia/Shanghai`, rather than `local`, an empty string, or an event-derived +value. Ordinary time and random-value functions remain available. + +Existing pipelines using restricted functions or arguments fail compilation +under the default policy, including when loaded from storage. This option +restricts VRL capabilities; it does not impose CPU or memory limits. + ## processors refer [elastic ingest processor][elastic-ingest-processor] for detail diff --git a/src/pipeline/src/etl/processor.rs b/src/pipeline/src/etl/processor.rs index 761bcd6cba4..8c999cd13ee 100644 --- a/src/pipeline/src/etl/processor.rs +++ b/src/pipeline/src/etl/processor.rs @@ -51,7 +51,7 @@ use urlencoding::UrlEncodingProcessor; use vrl::value::Value as VrlValue; use crate::error::{ - Error, FailedParseFieldFromStringSnafu, FieldMustBeTypeSnafu, InvalidFieldRenameSnafu, + FailedParseFieldFromStringSnafu, FieldMustBeTypeSnafu, InvalidFieldRenameSnafu, ProcessorKeyMustBeStringSnafu, ProcessorMustBeMapSnafu, ProcessorMustHaveStringKeySnafu, Result, UnsupportedProcessorSnafu, }; @@ -173,13 +173,12 @@ impl std::ops::DerefMut for Processors { } } -impl TryFrom<&Vec> for Processors { - type Error = Error; - - fn try_from(vec: &Vec) -> Result { +impl Processors { + /// Parses processors using the server's VRL host access policy. + pub fn from_yaml(vec: &[yaml_rust::Yaml], allow_vrl_host_access: bool) -> Result { let mut processors_builders = vec![]; for doc in vec { - let processor = parse_processor(doc)?; + let processor = parse_processor(doc, allow_vrl_host_access)?; processors_builders.push(processor); } Ok(Processors { @@ -188,7 +187,7 @@ impl TryFrom<&Vec> for Processors { } } -fn parse_processor(doc: &yaml_rust::Yaml) -> Result { +fn parse_processor(doc: &yaml_rust::Yaml, allow_vrl_host_access: bool) -> Result { let map = doc.as_hash().context(ProcessorMustBeMapSnafu)?; let key = map.keys().next().context(ProcessorMustHaveStringKeySnafu)?; @@ -227,7 +226,9 @@ fn parse_processor(doc: &yaml_rust::Yaml) -> Result { json_parse::PROCESSOR_JSON_PARSE => { ProcessorKind::JsonParse(JsonParseProcessor::try_from(value)?) } - vrl_processor::PROCESSOR_VRL => ProcessorKind::Vrl(VrlProcessor::try_from(value)?), + vrl_processor::PROCESSOR_VRL => { + ProcessorKind::Vrl(VrlProcessor::from_yaml(value, allow_vrl_host_access)?) + } select::PROCESSOR_SELECT => ProcessorKind::Select(SelectProcessor::try_from(value)?), filter::PROCESSOR_FILTER => ProcessorKind::Filter(FilterProcessor::try_from(value)?), _ => return UnsupportedProcessorSnafu { processor: str_key }.fail(), diff --git a/src/pipeline/src/etl/processor/vrl_processor.rs b/src/pipeline/src/etl/processor/vrl_processor.rs index ee3452523d7..d950834a2a8 100644 --- a/src/pipeline/src/etl/processor/vrl_processor.rs +++ b/src/pipeline/src/etl/processor/vrl_processor.rs @@ -21,11 +21,14 @@ use snafu::{OptionExt, ensure}; use vrl::compiler::runtime::Runtime; use vrl::compiler::{Program, TargetValue, compile}; use vrl::diagnostic::Formatter; -use vrl::prelude::TimeZone; +use vrl::prelude::{ + ArgumentList, Compiled, Example, Function, FunctionCompileContext, Parameter, TimeZone, + TypeState, function, +}; use vrl::value::{Kind, Secrets, Value as VrlValue}; use crate::error::{ - CompileVrlSnafu, Error, ExecuteVrlSnafu, KeyMustBeStringSnafu, Result, VrlRegexValueSnafu, + CompileVrlSnafu, ExecuteVrlSnafu, KeyMustBeStringSnafu, Result, VrlRegexValueSnafu, VrlReturnValueSnafu, }; use crate::etl::processor::yaml_string; @@ -46,8 +49,8 @@ pub struct VrlProcessor { } impl VrlProcessor { - pub fn new(source: String) -> Result { - let fns = vrl::stdlib::all(); + pub fn new(source: String, allow_host_access: bool) -> Result { + let fns = functions(allow_host_access); let compile_result = compile(&source, &fns).map_err(|e| { CompileVrlSnafu { @@ -101,10 +104,9 @@ impl VrlProcessor { } } -impl TryFrom<&yaml_rust::yaml::Hash> for VrlProcessor { - type Error = Error; - - fn try_from(value: &yaml_rust::yaml::Hash) -> Result { +impl VrlProcessor { + /// Parses VRL source using the server's host access policy. + pub fn from_yaml(value: &yaml_rust::yaml::Hash, allow_host_access: bool) -> Result { let mut source = String::new(); for (k, v) in value.iter() { let key = k @@ -114,8 +116,111 @@ impl TryFrom<&yaml_rust::yaml::Hash> for VrlProcessor { source = yaml_string(v, SOURCE)?; } } - let processor = VrlProcessor::new(source)?; - Ok(processor) + Self::new(source, allow_host_access) + } +} + +fn functions(allow_host_access: bool) -> Vec> { + let mut functions = vrl::stdlib::all(); + if allow_host_access { + return functions; + } + + // These functions bypass VRL's event target, including file reads during compilation. + // When upgrading VRL, re-audit this denylist and the argument restrictions below; + // new host-access functions are otherwise enabled by default. + functions.retain(|function| { + !matches!( + function.identifier(), + "get_env_var" + | "get_hostname" + | "get_timezone_name" + | "http_request" + | "dns_lookup" + | "reverse_dns" + | "encode_proto" + | "parse_proto" + | "parse_etld" + | "validate_json_schema" + | "log" + ) + }); + functions + .into_iter() + .map(|function| match function.identifier() { + "parse_groks" | "parse_timestamp" | "format_timestamp" => { + Box::new(RestrictedFunction(function)) as Box + } + _ => function, + }) + .collect() +} + +/// Retains data-only uses of functions that also accept host-access arguments. +#[derive(Debug)] +struct RestrictedFunction(Box); + +impl Function for RestrictedFunction { + fn identifier(&self) -> &'static str { + self.0.identifier() + } + + fn usage(&self) -> &'static str { + self.0.usage() + } + + fn category(&self) -> &'static str { + self.0.category() + } + + fn return_kind(&self) -> u16 { + self.0.return_kind() + } + + fn examples(&self) -> &'static [Example] { + self.0.examples() + } + + fn parameters(&self) -> &'static [Parameter] { + self.0.parameters() + } + + fn pure(&self) -> bool { + self.0.pure() + } + + fn compile( + &self, + state: &TypeState, + ctx: &mut FunctionCompileContext, + arguments: ArgumentList, + ) -> Compiled { + if self.identifier() == "parse_groks" { + if arguments + .optional_array("alias_sources")? + .is_some_and(|sources| !sources.is_empty()) + { + return Err(function::Error::InvalidArgument { + keyword: "alias_sources", + value: "alias_sources".into(), + error: "file imports require pipeline.allow_vrl_host_access", + } + .into()); + } + } else if let Some(timezone) = arguments.optional_literal("timezone", state)? { + // A static named timezone cannot resolve to the host's local timezone at runtime. + if !matches!(&timezone, VrlValue::Bytes(bytes) + if matches!(TimeZone::parse(&String::from_utf8_lossy(bytes)), Some(TimeZone::Named(_)))) + { + return Err(function::Error::InvalidArgument { + keyword: "timezone", + value: timezone, + error: "a named timezone is required when pipeline.allow_vrl_host_access is false", + } + .into()); + } + } + self.0.compile(state, ctx, arguments) } } @@ -163,6 +268,141 @@ mod tests { use super::*; + #[test] + fn test_host_functions_require_opt_in() { + let calls = [ + ("get_env_var", r#"get_env_var!("GREPTIME_VRL_TEST_SECRET")"#), + ("get_hostname", "get_hostname!()"), + ("get_timezone_name", "get_timezone_name!()"), + ("http_request", r#"http_request!("http://127.0.0.1:9/")"#), + ("dns_lookup", r#"dns_lookup!("localhost")"#), + ("reverse_dns", r#"reverse_dns!("127.0.0.1")"#), + ( + "parse_proto", + r#"parse_proto!("", "missing.desc", "Message")"#, + ), + ( + "encode_proto", + r#"encode_proto!({}, "missing.desc", "Message")"#, + ), + ("parse_etld", r#"parse_etld!("example.com")"#), + ( + "validate_json_schema", + r#"validate_json_schema!("{}", "missing.json")"#, + ), + ("log", r#"log("host log write")"#), + ]; + let unrestricted = functions(true); + for (name, call) in calls { + assert!( + unrestricted + .iter() + .any(|function| function.identifier() == name) + ); + let source = format!(".result = {call}\n."); + let error = VrlProcessor::new(source, false).unwrap_err().to_string(); + assert!(error.contains("undefined function"), "{name}: {error}"); + } + + temp_env::with_var("GREPTIME_VRL_TEST_SECRET", Some("decoy"), || { + let source = ".secret = get_env_var!(\"GREPTIME_VRL_TEST_SECRET\")\n."; + let processor = VrlProcessor::new(source.to_string(), true).unwrap(); + let result = processor + .resolve(VrlValue::Object(BTreeMap::new())) + .unwrap(); + assert_eq!( + result.as_object().unwrap()["secret"], + VrlValue::from("decoy") + ); + assert!(VrlProcessor::new(source.to_string(), false).is_err()); + }); + } + + #[test] + fn test_grok_file_import_requires_opt_in() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("aliases.json"); + std::fs::write(&path, r#"{"custom_word":"%{word}"}"#).unwrap(); + let path = serde_json::to_string(&path).unwrap(); + let source = format!( + r#".parsed = parse_groks!("hello", patterns: ["%{{custom_word:message}}"], alias_sources: [{path}]) +."# + ); + let error = VrlProcessor::new(source.clone(), false) + .unwrap_err() + .to_string(); + assert!(error.contains("file imports require pipeline.allow_vrl_host_access")); + let result = VrlProcessor::new(source, true) + .unwrap() + .resolve(VrlValue::Object(BTreeMap::new())) + .unwrap(); + assert_eq!( + result.as_object().unwrap()["parsed"].as_object().unwrap()["message"], + VrlValue::from("hello") + ); + + let source = r#".parsed = parse_groks!("hello", patterns: ["%{custom_word:message}"], aliases: {"custom_word": "%{word}"}, alias_sources: []) +."#; + assert!(VrlProcessor::new(source.to_string(), false).is_ok()); + } + + #[test] + fn test_local_and_dynamic_timezones_require_opt_in() { + for timezone in [r#""local""#, r#""""#, "string!(.timezone)"] { + for call in [ + format!( + r#"format_timestamp!(t'2020-01-01T00:00:00Z', format: "%z", timezone: {timezone})"# + ), + format!( + r#"parse_timestamp!("2020-01-01", format: "%Y-%m-%d", timezone: {timezone})"# + ), + ] { + let source = format!(".result = {call}\n."); + assert!( + VrlProcessor::new(source.clone(), false).is_err(), + "{source}" + ); + assert!(VrlProcessor::new(source.clone(), true).is_ok(), "{source}"); + } + } + + let source = r#".parsed = parse_timestamp!("2020-01-01 00:00:00", format: "%Y-%m-%d %H:%M:%S") +.utc = format_timestamp!(.parsed, format: "%z") +.named = format_timestamp!(.parsed, format: "%z", timezone: "Asia/Shanghai") +."#; + let result = VrlProcessor::new(source.to_string(), false) + .unwrap() + .resolve(VrlValue::Object(BTreeMap::new())) + .unwrap(); + assert_eq!(result.as_object().unwrap()["utc"], VrlValue::from("+0000")); + assert_eq!( + result.as_object().unwrap()["named"], + VrlValue::from("+0800") + ); + } + + #[test] + fn test_pipeline_host_policy_is_server_controlled() { + let yaml = r#" +allow_vrl_host_access: true +processors: + - vrl: + allow_vrl_host_access: true + source: | + .secret = get_env_var!("GREPTIME_VRL_TEST_SECRET") + . +transform: + - field: secret + type: string +"#; + assert!(crate::parse(&crate::Content::Yaml(yaml)).is_err()); + let options = crate::PipelineOptions { + allow_vrl_host_access: true, + ..Default::default() + }; + assert!(crate::parse_with_options(&crate::Content::Yaml(yaml), &options).is_ok()); + } + #[test] fn test_vrl() { let source = r#" @@ -173,7 +413,7 @@ del(.user_info) . "#; - let v = VrlProcessor::new(source.to_string()); + let v = VrlProcessor::new(source.to_string(), false); assert!(v.is_ok()); let v = v.unwrap(); @@ -226,7 +466,7 @@ processors: .and_then(|x| x.as_hash()) .unwrap(); - let vrl = VrlProcessor::try_from(vrl_processor_yaml); + let vrl = VrlProcessor::from_yaml(vrl_processor_yaml, false); assert!(vrl.is_ok()); let vrl = vrl.unwrap(); @@ -245,7 +485,7 @@ del(.re) . "#; - let v = VrlProcessor::new(source.to_string()); + let v = VrlProcessor::new(source.to_string(), false); assert!(v.is_err()); } } diff --git a/src/pipeline/src/lib.rs b/src/pipeline/src/lib.rs index 15954b18f98..086758fa9cc 100644 --- a/src/pipeline/src/lib.rs +++ b/src/pipeline/src/lib.rs @@ -27,7 +27,7 @@ pub use etl::transform::transformer::greptime::{GreptimePipelineParams, SchemaIn pub use etl::transform::transformer::identity_pipeline; pub use etl::{ Content, DispatchedTo, Pipeline, PipelineExecOutput, PipelineProcessOutput, TransformedOutput, - TransformerMode, parse, + TransformerMode, parse, parse_with_options, }; pub use manager::{ GREPTIME_INTERNAL_IDENTITY_PIPELINE_NAME, GREPTIME_INTERNAL_TRACE_PIPELINE_V1_NAME, diff --git a/src/pipeline/src/manager/pipeline_cache.rs b/src/pipeline/src/manager/pipeline_cache.rs index f0e9ca52970..8f6a91df2c6 100644 --- a/src/pipeline/src/manager/pipeline_cache.rs +++ b/src/pipeline/src/manager/pipeline_cache.rs @@ -282,11 +282,9 @@ mod tests { } .build()); } - Ok(Arc::new( - crate::manager::table::PipelineTable::compile_pipeline( - "transform:\n - fields: [message]\n type: string", - )?, - )) + Ok(Arc::new(crate::parse(&crate::Content::Yaml( + "transform:\n - fields: [message]\n type: string", + ))?)) }) .await .map(|_| ()) diff --git a/src/pipeline/src/manager/pipeline_operator.rs b/src/pipeline/src/manager/pipeline_operator.rs index cace74171c2..7568fc088bd 100644 --- a/src/pipeline/src/manager/pipeline_operator.rs +++ b/src/pipeline/src/manager/pipeline_operator.rs @@ -14,7 +14,7 @@ use std::collections::HashMap; use std::sync::{Arc, RwLock}; -use std::time::{Duration, Instant}; +use std::time::Instant; use api::v1::CreateTableExpr; use catalog::{CatalogManagerRef, RegisterSystemTableRequest}; @@ -57,7 +57,7 @@ pub struct PipelineOperator { catalog_manager: CatalogManagerRef, query_engine: QueryEngineRef, tables: RwLock>, - cache_ttl: Duration, + options: PipelineOptions, } impl PipelineOperator { @@ -100,7 +100,7 @@ impl PipelineOperator { self.statement_executor.clone(), table, self.query_engine.clone(), - self.cache_ttl, + &self.options, )), ); } @@ -186,7 +186,7 @@ impl PipelineOperator { catalog_manager, tables: RwLock::new(HashMap::new()), query_engine, - cache_ttl: options.cache_ttl, + options: options.clone(), } } @@ -284,7 +284,7 @@ impl PipelineOperator { } /// Compile a pipeline. - pub fn build_pipeline(pipeline: &str) -> Result { - PipelineTable::compile_pipeline(pipeline) + pub fn build_pipeline(&self, pipeline: &str) -> Result { + crate::parse_with_options(&crate::Content::Yaml(pipeline), &self.options) } } diff --git a/src/pipeline/src/manager/table.rs b/src/pipeline/src/manager/table.rs index 0b81493b14c..ab675b6cea5 100644 --- a/src/pipeline/src/manager/table.rs +++ b/src/pipeline/src/manager/table.rs @@ -13,7 +13,6 @@ // limitations under the License. use std::sync::Arc; -use std::time::Duration; use api::v1::value::ValueData; use api::v1::{ @@ -46,10 +45,11 @@ use crate::error::{ ExecuteInternalStatementSnafu, InsertPipelineSnafu, InvalidPipelineVersionSnafu, MultiPipelineWithDiffSchemaSnafu, PipelineNotFoundSnafu, RecordBatchLenNotMatchSnafu, Result, }; -use crate::etl::{Content, Pipeline, parse}; +use crate::etl::{Content, Pipeline, parse_with_options}; use crate::manager::pipeline_cache::{PipelineCache, PipelineContent}; use crate::manager::{PipelineInfo, PipelineVersion}; use crate::metrics::METRIC_PIPELINE_TABLE_FIND_COUNT; +use crate::options::PipelineOptions; use crate::util::prepare_dataframe_conditions; pub(crate) const PIPELINE_TABLE_NAME: &str = "pipelines"; @@ -68,6 +68,7 @@ pub struct PipelineTable { table: TableRef, query_engine: QueryEngineRef, cache: PipelineCache, + options: PipelineOptions, } impl PipelineTable { @@ -77,14 +78,15 @@ impl PipelineTable { statement_executor: StatementExecutorRef, table: TableRef, query_engine: QueryEngineRef, - cache_ttl: Duration, + options: &PipelineOptions, ) -> Self { Self { inserter, statement_executor, table, query_engine, - cache: PipelineCache::new(cache_ttl), + cache: PipelineCache::new(options.cache_ttl), + options: options.clone(), } } @@ -198,9 +200,9 @@ impl PipelineTable { } /// Compile a pipeline from a string. - pub fn compile_pipeline(pipeline: &str) -> Result { + pub fn compile_pipeline(&self, pipeline: &str) -> Result { let yaml_content = Content::Yaml(pipeline); - parse(&yaml_content) + parse_with_options(&yaml_content, &self.options) } /// Insert a pipeline into the pipeline table. @@ -267,7 +269,7 @@ impl PipelineTable { self.cache .get_pipeline_with(schema, name, input_version, async { let pipeline_content = self.get_pipeline_str(schema, name, input_version).await?; - Ok(Arc::new(Self::compile_pipeline(&pipeline_content.content)?)) + Ok(Arc::new(self.compile_pipeline(&pipeline_content.content)?)) }) .await } @@ -373,7 +375,7 @@ impl PipelineTable { content_type: &str, pipeline: &str, ) -> Result { - let compiled_pipeline = Arc::new(Self::compile_pipeline(pipeline)?); + let compiled_pipeline = Arc::new(self.compile_pipeline(pipeline)?); // we will use the version in the future let version = self .insert_pipeline_to_pipeline_table(name, content_type, pipeline) diff --git a/src/pipeline/src/options.rs b/src/pipeline/src/options.rs index 6e34f0ce933..ad1c72a430b 100644 --- a/src/pipeline/src/options.rs +++ b/src/pipeline/src/options.rs @@ -24,12 +24,17 @@ pub struct PipelineOptions { /// duration. Default: "10s". #[serde(with = "humantime_serde")] pub cache_ttl: Duration, + /// Allows VRL programs to access the server environment, filesystem, hostname, + /// local timezone, network, and logs. Enable only for trusted pipeline authors. + /// Defaults to false. Changes require a restart. + pub allow_vrl_host_access: bool, } impl Default for PipelineOptions { fn default() -> Self { Self { cache_ttl: Duration::from_secs(10), + allow_vrl_host_access: false, } } } diff --git a/src/standalone/src/options.rs b/src/standalone/src/options.rs index 05d929cc7da..657961e9298 100644 --- a/src/standalone/src/options.rs +++ b/src/standalone/src/options.rs @@ -222,6 +222,19 @@ mod tests { use crate::options::*; + #[test] + fn test_pipeline_host_access_config_forwarding() { + let defaults: StandaloneOptions = toml::from_str("").unwrap(); + assert!(!defaults.pipeline.allow_vrl_host_access); + assert!(!defaults.frontend_options().pipeline.allow_vrl_host_access); + let options: StandaloneOptions = + toml::from_str("[pipeline]\nallow_vrl_host_access = true").unwrap(); + assert!(options.frontend_options().pipeline.allow_vrl_host_access); + let serialized = toml::to_string(&options).unwrap(); + let restored: StandaloneOptions = toml::from_str(&serialized).unwrap(); + assert_eq!(options.pipeline, restored.pipeline); + } + #[test] fn test_frontend_write_retries_disabled() { assert!(FrontendOptions::default().write_retry.max_retries > 0); diff --git a/tests-integration/Cargo.toml b/tests-integration/Cargo.toml index 79df85475d9..31a238f7bd5 100644 --- a/tests-integration/Cargo.toml +++ b/tests-integration/Cargo.toml @@ -56,6 +56,7 @@ meta-srv = { workspace = true, features = ["mock"] } mito2.workspace = true object-store.workspace = true operator.workspace = true +pipeline.workspace = true prost.workspace = true query.workspace = true rand.workspace = true @@ -107,7 +108,6 @@ opentelemetry-proto.workspace = true operator = { workspace = true, features = ["testing"] } otel-arrow-rust.workspace = true paste.workspace = true -pipeline.workspace = true plugins.workspace = true prost.workspace = true rand.workspace = true diff --git a/tests-integration/src/standalone.rs b/tests-integration/src/standalone.rs index f8e24cb001d..e0113f90425 100644 --- a/tests-integration/src/standalone.rs +++ b/tests-integration/src/standalone.rs @@ -91,6 +91,7 @@ pub struct GreptimeDbStandaloneBuilder { logical_batcher: Option, table_batcher: BatcherOptions, mito_config: Option, + pipeline_options: pipeline::PipelineOptions, } impl GreptimeDbStandaloneBuilder { @@ -112,6 +113,7 @@ impl GreptimeDbStandaloneBuilder { logical_batcher: None, table_batcher: BatcherOptions::default(), mito_config: None, + pipeline_options: pipeline::PipelineOptions::default(), } } @@ -122,6 +124,13 @@ impl GreptimeDbStandaloneBuilder { self } + /// Configures pipeline compilation and caching for this test instance. + #[must_use] + pub fn with_pipeline_options(mut self, options: pipeline::PipelineOptions) -> Self { + self.pipeline_options = options; + self + } + /// Configures ordinary-table batching for protocol integration tests. #[must_use] pub fn with_table_batcher(mut self, options: BatcherOptions) -> Self { @@ -391,6 +400,7 @@ impl GreptimeDbStandaloneBuilder { slow_query: self.slow_query_options.clone(), event_recorder: self.event_recorder_options.clone(), auto_create_table: self.auto_create_table, + pipeline: self.pipeline_options.clone(), pending_rows_batcher: PendingRowsBatcherOptions { logical_table: self.logical_batcher.clone(), table: self.table_batcher.clone(), diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 1bb4af7839c..5e006c5d80d 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -200,6 +200,7 @@ macro_rules! http_tests { test_pipeline_suffix_template, test_pipeline_context, test_pipeline_with_vrl, + test_pipeline_vrl_host_access, test_pipeline_with_hint_vrl, test_pipeline_one_to_many_vrl, test_pipeline_2, @@ -3519,6 +3520,7 @@ enable_heap_profiling = true [pipeline] cache_ttl = "10s" +allow_vrl_host_access = false [event_recorder] ttl = "2months 29days 2h 52m 48s" @@ -7012,6 +7014,72 @@ table_suffix: _${type} guard.remove_all().await; } +pub async fn test_pipeline_vrl_host_access(storage_type: StorageType) { + let pipeline = r#" +processors: + - vrl: + source: | + .value = is_string(get_env_var("GREPTIME_VRL_HOST_ACCESS_TEST") ?? "missing") + . +transform: + - field: value + type: boolean +"#; + + for allow_vrl_host_access in [false, true] { + let mut instance = tests_integration::standalone::GreptimeDbStandaloneBuilder::new( + "test_pipeline_vrl_host_access", + ) + .with_default_store_type(storage_type) + .with_pipeline_options(pipeline::PipelineOptions { + // Force named dry-runs to reload and compile the stored definition. + cache_ttl: Duration::ZERO, + allow_vrl_host_access, + }) + .build() + .await; + let server = HttpServerBuilder::new(HttpOptions::default()) + .with_log_ingest_handler(instance.fe_instance().clone(), None, None) + .build(); + let client = TestClient::new(server.build(server.make_app()).unwrap()).await; + + let response = client + .post("/v1/pipelines/_dryrun") + .header("Content-Type", "application/json") + .body(json!({"pipeline": pipeline, "data": "[{}]"}).to_string()) + .send() + .await; + let expected_status = if allow_vrl_host_access { + StatusCode::OK + } else { + StatusCode::BAD_REQUEST + }; + assert_eq!(response.status(), expected_status); + if !allow_vrl_host_access { + assert!(response.text().await.contains("undefined function")); + } + + let response = client + .post("/v1/pipelines/host_policy") + .header("Content-Type", "application/x-yaml") + .body(pipeline) + .send() + .await; + assert_eq!(response.status(), expected_status); + + if allow_vrl_host_access { + let response = client + .post("/v1/pipelines/_dryrun") + .header("Content-Type", "application/json") + .body(json!({"pipeline_name": "host_policy", "data": "[{}]"}).to_string()) + .send() + .await; + assert_eq!(response.status(), StatusCode::OK); + } + instance.guard.remove_all().await; + } +} + pub async fn test_pipeline_with_vrl(storage_type: StorageType) { common_telemetry::init_default_ut_logging(); let (app, mut guard) = From 33879b07c861b0a7e757172de9b52c687f914259 Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Thu, 15 Oct 2026 10:47:29 +0000 Subject: [PATCH 012/239] chore: bump rustls to 0.23.46 (#9614) rustls 0.23.46 fixes rustls/rustls#3301: on TLS 1.2 with AES-GCM, reaching the 2^24 record limit dropped a fragment while reporting the write as complete, then kept sending records after close_notify. This stalled Kafka WAL produce requests until the request timeout. --- Cargo.lock | 17 +++++++++-------- Cargo.toml | 2 +- 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index b696ee8224d..22f3921a133 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1002,9 +1002,9 @@ checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" [[package]] name = "aws-lc-rs" -version = "1.17.0" +version = "1.18.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5ec2f1fc3ec205783a5da9a7e6c1509cc69dedf09a1949e412c1e18469326d00" +checksum = "faac5829c2b74c28f830747e7818ccfb684261b5f48a1118b1e2a13d36dfab13" dependencies = [ "aws-lc-sys", "untrusted 0.7.1", @@ -1013,14 +1013,15 @@ dependencies = [ [[package]] name = "aws-lc-sys" -version = "0.41.0" +version = "0.46.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1a2f9779ce85b93ab6170dd940ad0169b5766ff848247aff13bb788b832fe3f4" +checksum = "f1622d8446a2d4b2ce0c7eefc73dd43a99779028d5ee5c2dd8073a658ba8a2bc" dependencies = [ "cc", "cmake", "dunce", "fs_extra", + "pkg-config", ] [[package]] @@ -12972,9 +12973,9 @@ dependencies = [ [[package]] name = "rustls" -version = "0.23.28" +version = "0.23.46" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7160e3e10bf4535308537f3c4e1641468cd0e485175d6163087c0393c7d46643" +checksum = "48e13bd8c0e9365c43cfa5c9e8f9ad49d3c8444926c9aac819e0e4dc503c8fdf" dependencies = [ "aws-lc-rs", "log", @@ -13059,9 +13060,9 @@ checksum = "f87165f0995f63a9fbeea62b64d10b4d9d8e78ec6d7d51fb2125fda7bb36788f" [[package]] name = "rustls-webpki" -version = "0.103.13" +version = "0.103.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "61c429a8649f110dddef65e2a5ad240f747e85f7758a6bccc7e5777bd33f756e" +checksum = "f3c3cf1d8b1e7d4927e2d154c3fcb02979afb9939629c62cd9048d4f07b60ac2" dependencies = [ "aws-lc-rs", "ring", diff --git a/Cargo.toml b/Cargo.toml index 9fa5d013272..e2f91a54336 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -237,7 +237,7 @@ rskafka = { git = "https://github.com/influxdata/rskafka.git", rev = "80f42eb15e rstest = "0.25" rstest_reuse = "0.7" rust_decimal = "1.33" -rustls = { version = "0.23.25", default-features = false } +rustls = { version = "0.23.46", default-features = false } ryu = "1.0" sea-query = "0.32" serde = { version = "1.0", features = ["derive"] } From c3469fad51a1f47ba9df567a5019412968d0f9e2 Mon Sep 17 00:00:00 2001 From: discord9 Date: Thu, 15 Oct 2026 10:49:09 +0000 Subject: [PATCH 013/239] fix(flow): retain tracked write completion across transient failures (#9525) * fix(flow): retain tracked write completion across transient failures Use receiver-issued single-write capabilities at the frontend and datanode. Rejoin exact admitted writes across unknown transport outcomes, then close and drain before allowing subsequent rounds. Keep ordinary clients and reads on their existing paths. Separate schema map equality from encoded row identity so fresh protobuf decodes replay consistently even with NaN values. Cover receiver admission, owner cancellation, all-child drain, and retired capabilities. Local feature-worktree checkpoint only. Main verified core enterprise check and 484 tests, plus DN enterprise check and all four tracked owner tests. Corrected seven-case timer acceptance remains to be rerun; this checkpoint does not authorize publishing or integration. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(flow): preserve upstream dependencies after rebase Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(flow): fail fast on definitive tracked-write open rejection A tracked write treats every application-level error from the frontend as an ambiguous transport failure and retries the same pinned frontend forever. During a rolling upgrade an older frontend does not recognize the tracked-write command extension and processes Open as an ordinary empty SQL query, rejecting it with InvalidSql (InvalidArguments); a non-internal endpoint rejects the extension outright. Both are definitive rejections that no same-frontend retry can ever resolve, so the flow stalls. Classify the Open-phase errors by status code: a definitive rejection (Unsupported / InvalidArguments / PermissionDenied) fails fast with a terminal pre-write error, while genuinely ambiguous transport failures (Unavailable, timeout, connection reset, ...) keep the same-capability retry. Write/Close rejections are deliberately NOT failed fast: once a write may have been admitted, a same-code rejection (e.g. WriterLost, or UnknownCapability after a VIP moved Open->A/Write->B) leaves the outcome ambiguous, and failing fast would let the caller replay and duplicate rows. Add a regression test that a legacy empty-SQL Open rejection fails fast with a single attempt and surfaces as Error::External carrying a TerminalError. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(flow): preflight tracked writes without submitting data Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(flow): retain tracked command extensions while decoding Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- Cargo.lock | 1 + licenserc-enterprise.toml | 5 + licenserc.toml | 5 + src/common/grpc/Cargo.toml | 4 + src/common/grpc/src/flow_write.rs | 477 +++++++++ src/common/grpc/src/lib.rs | 2 + src/datanode/Cargo.toml | 2 +- src/datanode/src/region_server.rs | 32 +- src/datanode/src/region_server/flow_write.rs | 959 ++++++++++++++++++ src/flow/Cargo.toml | 3 + src/flow/src/batching_mode/frontend_client.rs | 102 ++ .../frontend_client/flow_write.rs | 758 ++++++++++++++ src/frontend/Cargo.toml | 2 + src/frontend/src/instance.rs | 83 +- src/frontend/src/instance/builder.rs | 2 + src/frontend/src/instance/flow_write.rs | 227 +++++ src/frontend/src/instance/grpc.rs | 14 + src/operator/Cargo.toml | 1 + src/operator/src/insert.rs | 44 + src/operator/src/insert/flow_write.rs | 391 +++++++ src/servers/src/grpc/region_server.rs | 14 +- src/session/Cargo.toml | 3 + src/session/src/context.rs | 5 + src/standalone/src/datanode_manager.rs | 8 +- 24 files changed, 3104 insertions(+), 40 deletions(-) create mode 100644 src/common/grpc/src/flow_write.rs create mode 100644 src/datanode/src/region_server/flow_write.rs create mode 100644 src/flow/src/batching_mode/frontend_client/flow_write.rs create mode 100644 src/frontend/src/instance/flow_write.rs create mode 100644 src/operator/src/insert/flow_write.rs diff --git a/Cargo.lock b/Cargo.lock index 22f3921a133..658180353e3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2581,6 +2581,7 @@ dependencies = [ "tokio-util", "tonic 0.14.2", "tower 0.5.2", + "uuid", "vec1", ] diff --git a/licenserc-enterprise.toml b/licenserc-enterprise.toml index d11d344b64d..805358f1f1e 100644 --- a/licenserc-enterprise.toml +++ b/licenserc-enterprise.toml @@ -15,6 +15,11 @@ headerPath = "licenses/enterprise-header.txt" # with the "# enterprise" block in `licenserc.toml`'s `excludes` — both are # enforced by `make check-enterprise-license`. includes = [ + "src/common/grpc/src/flow_write.rs", + "src/frontend/src/instance/flow_write.rs", + "src/datanode/src/region_server/flow_write.rs", + "src/operator/src/insert/flow_write.rs", + "src/flow/src/batching_mode/frontend_client/flow_write.rs", "src/common/meta/src/rpc/ddl/trigger.rs", "src/operator/src/expr_helper/trigger.rs", "src/sql/src/statements/alter/trigger.rs", diff --git a/licenserc.toml b/licenserc.toml index 42e6464017c..e84d3a1f8ff 100644 --- a/licenserc.toml +++ b/licenserc.toml @@ -31,6 +31,11 @@ excludes = [ # Keep this list in sync with that file's `includes`; enforced by # `make check-enterprise-license`. # enterprise:start + "src/common/grpc/src/flow_write.rs", + "src/frontend/src/instance/flow_write.rs", + "src/datanode/src/region_server/flow_write.rs", + "src/operator/src/insert/flow_write.rs", + "src/flow/src/batching_mode/frontend_client/flow_write.rs", "src/common/meta/src/rpc/ddl/trigger.rs", "src/operator/src/expr_helper/trigger.rs", "src/sql/src/statements/alter/trigger.rs", diff --git a/src/common/grpc/Cargo.toml b/src/common/grpc/Cargo.toml index 58bd04db093..fbfab1b6318 100644 --- a/src/common/grpc/Cargo.toml +++ b/src/common/grpc/Cargo.toml @@ -7,6 +7,9 @@ license.workspace = true [lints] workspace = true +[features] +enterprise = [] + [dependencies] api.workspace = true arrow-flight.workspace = true @@ -32,6 +35,7 @@ tokio.workspace = true tokio-util.workspace = true tonic.workspace = true tower.workspace = true +uuid.workspace = true vec1 = "1.12" [dev-dependencies] diff --git a/src/common/grpc/src/flow_write.rs b/src/common/grpc/src/flow_write.rs new file mode 100644 index 00000000000..6d085a500bd --- /dev/null +++ b/src/common/grpc/src/flow_write.rs @@ -0,0 +1,477 @@ +// Copyright 2023-2026 GrepTime Inc. +// +// This file is part of the GreptimeDB Enterprise Edition and is licensed under +// the GreptimeDB Enterprise License. You may not use this file except in +// compliance with that license. A copy of the license is available at the root +// of this repository in the file LICENSE-ENTERPRISE. +// +// Unless required by applicable law or agreed to in writing, this software is +// distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, +// either express or implied. + +use std::collections::HashMap; +use std::sync::{Arc, Mutex}; + +use common_error::ext::{ErrorExt, StackError}; +use common_error::status_code::StatusCode; + +use common_recordbatch::{RecordBatch, RecordBatches}; +use datatypes::arrow::array::{Array, StringArray}; +use datatypes::prelude::{ConcreteDataType, VectorRef}; +use datatypes::schema::{ColumnSchema, Schema}; +use datatypes::vectors::StringVector; +use serde::{Deserialize, Serialize}; +use tokio::sync::watch; +use uuid::Uuid; + +pub const COMMAND_KEY: &str = "flow_write_command"; +pub const CHILD_KEY: &str = "flow_write_child"; + +/// The control command carried by a Flow write request. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub enum Command { + Open { nonce: String }, + Write { capability: String }, + Close { capability: String }, +} + +impl Command { + /// Serializes this command as JSON. + pub fn encode(&self) -> serde_json::Result { + serde_json::to_string(self) + } + + /// Parses a serialized command. + pub fn decode(value: &str) -> serde_json::Result { + serde_json::from_str(value) + } +} + +/// A safe terminal application error. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct TerminalError { + pub code: u32, + pub message: String, +} + +impl std::fmt::Display for TerminalError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.message) + } +} + +impl std::error::Error for TerminalError {} + +impl StackError for TerminalError { + fn debug_fmt(&self, layer: usize, buf: &mut Vec) { + buf.push(format!("{layer}: {self:?}")); + } + + fn next(&self) -> Option<&dyn StackError> { + None + } +} + +impl ErrorExt for TerminalError { + fn status_code(&self) -> StatusCode { + StatusCode::from_u32(self.code).unwrap_or(StatusCode::Internal) + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } +} + +/// A finite response for a tracked Flow write. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum Reply { + Opened { + capability: String, + }, + Closed, + Complete { + rows: usize, + cost: usize, + metrics: Option, + error: Option, + }, +} + +fn envelope_schema() -> Arc { + Arc::new(Schema::new(vec![ColumnSchema::new( + "envelope", + ConcreteDataType::string_datatype(), + false, + )])) +} + +fn response_error(reason: impl Into) -> common_recordbatch::error::Error { + common_recordbatch::error::Error::CreateRecordBatches { + reason: reason.into(), + location: snafu::Location::new(file!(), line!(), column!()), + } +} + +/// Encodes a reply in the finite one-string-column record-batch envelope. +pub fn reply_batches(reply: &Reply) -> common_recordbatch::error::Result { + let value = serde_json::to_string(reply).map_err(|e| response_error(e.to_string()))?; + let batch = RecordBatch::new( + envelope_schema(), + vec![Arc::new(StringVector::from(vec![value])) as VectorRef], + )?; + RecordBatches::try_new(batch.schema.clone(), vec![batch]) +} + +/// Decodes a reply from its finite one-string-column record-batch envelope. +pub fn decode_reply(batches: &RecordBatches) -> common_recordbatch::error::Result { + let batch = batches + .iter() + .next() + .ok_or_else(|| response_error("missing reply batch"))?; + if batches.iter().count() != 1 || batch.num_columns() != 1 || batch.num_rows() != 1 { + return Err(response_error("reply must be exactly one string value")); + } + let values = batch.columns()[0] + .as_any() + .downcast_ref::() + .ok_or_else(|| response_error("reply value is not a string"))?; + if values.is_null(0) { + return Err(response_error("reply value is null")); + } + let value = values.value(0); + serde_json::from_str(value).map_err(|e| response_error(e.to_string())) +} + +/// Errors that reject admission or indicate that an admitted writer was lost. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum RegistryError { + UnknownCapability, + Closing, + WriterLost, + PayloadMismatch, +} + +impl std::fmt::Display for RegistryError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{self:?}") + } +} + +impl std::error::Error for RegistryError {} + +impl StackError for RegistryError { + fn debug_fmt(&self, layer: usize, buf: &mut Vec) { + buf.push(format!("{layer}: {self:?}")); + } + + fn next(&self) -> Option<&dyn StackError> { + None + } +} + +impl ErrorExt for RegistryError { + fn status_code(&self) -> StatusCode { + StatusCode::InvalidArguments + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } +} + +/// Result of atomically admitting or joining a payload. +#[derive(Debug)] +pub enum Admission { + Start(watch::Sender>>), + Join(watch::Receiver>>), +} + +#[derive(Debug)] +struct Entry { + payload: Option

, + sender: Option>>>, + receiver: watch::Receiver>>, + closing: bool, +} + +#[derive(Debug)] +struct State { + openers: HashMap, + entries: HashMap>, +} + +/// Receiver-issued capabilities for single-write admission and completion sharing. +#[derive(Debug)] +pub struct FlowWriteRegistry { + state: Mutex>, +} + +impl Default for FlowWriteRegistry { + fn default() -> Self { + Self { + state: Mutex::new(State { + openers: HashMap::new(), + entries: HashMap::new(), + }), + } + } +} + +impl FlowWriteRegistry { + /// Returns the active capability for this opener nonce, creating one if needed. + pub fn open(&self, nonce: &str) -> Result { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(capability) = state.openers.get(nonce) { + if state + .entries + .get(capability) + .is_some_and(|entry| entry.closing) + { + return Err(RegistryError::Closing); + } + return Ok(capability.clone()); + } + let capability = Uuid::new_v4().to_string(); + let (sender, receiver) = watch::channel(None); + state.openers.insert(nonce.to_string(), capability.clone()); + state.entries.insert( + capability.clone(), + Entry { + payload: None, + sender: Some(sender), + receiver, + closing: false, + }, + ); + Ok(capability) + } + + /// Atomically starts one payload or joins its identical in-flight completion. + pub fn admit(&self, capability: &str, payload: P) -> Result, RegistryError> { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let entry = state + .entries + .get_mut(capability) + .ok_or(RegistryError::UnknownCapability)?; + if entry.closing { + return Err(RegistryError::Closing); + } + match &entry.payload { + Some(existing) if existing != &payload => Err(RegistryError::PayloadMismatch), + Some(_) => Ok(Admission::Join(entry.receiver.clone())), + None => { + entry.payload = Some(payload); + let sender = entry.sender.take().ok_or(RegistryError::Closing)?; + Ok(Admission::Start(sender)) + } + } + } + + /// Closes admission, drains an admitted writer, then deletes the capability. + pub async fn close(&self, capability: &str) -> Result<(), RegistryError> { + let receiver = { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let Some(entry) = state.entries.get_mut(capability) else { + return Ok(()); + }; + entry.closing = true; + if entry.payload.is_none() { + state.entries.remove(capability); + state.openers.retain(|_, value| value != capability); + return Ok(()); + } + entry.receiver.clone() + }; + wait_completion(receiver).await?; + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + state.entries.remove(capability); + state.openers.retain(|_, value| value != capability); + Ok(()) + } +} + +/// Waits for an admitted writer's positive terminal completion. +pub async fn wait_completion( + mut receiver: watch::Receiver>>, +) -> Result, RegistryError> { + loop { + if let Some(value) = receiver.borrow().clone() { + return Ok(value); + } + if receiver.changed().await.is_err() { + return Err(RegistryError::WriterLost); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Barrier; + use std::thread; + + #[tokio::test] + async fn concurrent_admission_has_one_owner_and_same_payload_joins() { + let registry = Arc::new(FlowWriteRegistry::::default()); + let cap = registry.open("nonce").unwrap(); + let barrier = Arc::new(Barrier::new(8)); + let mut workers = Vec::new(); + for _ in 0..8 { + let registry = registry.clone(); + let cap = cap.clone(); + let barrier = barrier.clone(); + workers.push(thread::spawn(move || { + barrier.wait(); + registry.admit(&cap, "payload".into()).unwrap() + })); + } + let admissions: Vec<_> = workers.into_iter().map(|w| w.join().unwrap()).collect(); + assert_eq!( + admissions + .iter() + .filter(|a| matches!(a, Admission::Start(_))) + .count(), + 1 + ); + let sender = admissions + .into_iter() + .find_map(|a| match a { + Admission::Start(s) => Some(s), + _ => None, + }) + .unwrap(); + sender.send(Some(Arc::new(7))).unwrap(); + registry.close(&cap).await.unwrap(); + } + + #[tokio::test] + async fn close_drains_active_writer_and_deletes_old_capability() { + let registry = Arc::new(FlowWriteRegistry::::default()); + let cap = registry.open("first").unwrap(); + let Admission::Start(sender) = registry.admit(&cap, "payload".into()).unwrap() else { + panic!() + }; + let close_registry = registry.clone(); + let close_cap = cap.clone(); + let close = tokio::spawn(async move { close_registry.close(&close_cap).await }); + tokio::task::yield_now().await; + assert!(matches!( + registry.admit(&cap, "payload".into()), + Err(RegistryError::Closing) + )); + assert!(!close.is_finished()); + sender.send(Some(Arc::new(9))).unwrap(); + close.await.unwrap().unwrap(); + assert!(matches!( + registry.admit(&cap, "payload".into()), + Err(RegistryError::UnknownCapability) + )); + registry.close(&cap).await.unwrap(); + let new_cap = registry.open("first").unwrap(); + assert_ne!(cap, new_cap); + assert!(matches!( + registry.admit(&cap, "late".into()), + Err(RegistryError::UnknownCapability) + )); + assert!(matches!( + registry.admit(&new_cap, "new".into()).unwrap(), + Admission::Start(_) + )); + } + + #[tokio::test] + async fn writer_loss_is_not_completion() { + let registry = FlowWriteRegistry::::default(); + let cap = registry.open("nonce").unwrap(); + let Admission::Start(sender) = registry.admit(&cap, "payload".into()).unwrap() else { + panic!() + }; + let Admission::Join(receiver) = registry.admit(&cap, "payload".into()).unwrap() else { + panic!() + }; + drop(sender); + assert_eq!( + wait_completion(receiver).await.unwrap_err(), + RegistryError::WriterLost + ); + assert_eq!( + registry.close(&cap).await.unwrap_err(), + RegistryError::WriterLost + ); + } + + #[test] + fn command_and_reply_envelope_round_trip_and_malformed_envelopes() { + let command = Command::Open { nonce: "n".into() }; + assert_eq!( + Command::decode(&command.encode().unwrap()).unwrap(), + command + ); + let reply = Reply::Complete { + rows: 2, + cost: 3, + metrics: Some(common_recordbatch::adapter::RecordBatchMetrics::default()), + error: Some(TerminalError { + code: 4, + message: "terminal".into(), + }), + }; + let batches = reply_batches(&reply).unwrap(); + let decoded = decode_reply(&batches).unwrap(); + assert!(matches!( + decoded, + Reply::Complete { + rows: 2, + cost: 3, + error: Some(_), + .. + } + )); + + let schema = envelope_schema(); + let two_rows = RecordBatch::new( + schema.clone(), + vec![Arc::new(StringVector::from(vec!["a", "b"])) as VectorRef], + ) + .unwrap(); + assert!( + decode_reply(&RecordBatches::try_new(schema.clone(), vec![two_rows]).unwrap()).is_err() + ); + let int_schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "envelope", + ConcreteDataType::int32_datatype(), + false, + )])); + let non_string = RecordBatch::new( + int_schema.clone(), + vec![Arc::new(datatypes::vectors::Int32Vector::from(vec![Some(1)])) as VectorRef], + ) + .unwrap(); + assert!( + decode_reply(&RecordBatches::try_new(int_schema, vec![non_string]).unwrap()).is_err() + ); + let invalid_json = RecordBatch::new( + envelope_schema(), + vec![Arc::new(StringVector::from(vec!["not json"])) as VectorRef], + ) + .unwrap(); + assert!( + decode_reply( + &RecordBatches::try_new(invalid_json.schema.clone(), vec![invalid_json]).unwrap() + ) + .is_err() + ); + } +} diff --git a/src/common/grpc/src/lib.rs b/src/common/grpc/src/lib.rs index 8527dd079b8..0ac05863f3e 100644 --- a/src/common/grpc/src/lib.rs +++ b/src/common/grpc/src/lib.rs @@ -15,6 +15,8 @@ pub mod channel_manager; pub mod error; pub mod flight; +#[cfg(feature = "enterprise")] +pub mod flow_write; pub mod precision; pub mod reloadable_tls; pub mod select; diff --git a/src/datanode/Cargo.toml b/src/datanode/Cargo.toml index 4d4e470e665..7fedd2fc263 100644 --- a/src/datanode/Cargo.toml +++ b/src/datanode/Cargo.toml @@ -6,7 +6,7 @@ license.workspace = true [features] testing = [] -enterprise = ["mito2/enterprise"] +enterprise = ["mito2/enterprise", "common-grpc/enterprise"] [lints] workspace = true diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index a1ccc382ef1..185183ce308 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -13,6 +13,10 @@ // limitations under the License. mod catalog; +#[cfg(feature = "enterprise")] +mod flow_write; +#[cfg(feature = "enterprise")] +use flow_write::TrackedPayload; mod registrations; mod remote_dyn_filter; @@ -29,7 +33,8 @@ use api::region::RegionResponse; use api::v1::meta::TopicStat; use api::v1::region::sync_request::ManifestInfo; use api::v1::region::{ - ListMetadataRequest, RegionResponse as RegionResponseV1, SyncRequest, region_request, + ListMetadataRequest, RegionRequest as RegionRequestV1, RegionResponse as RegionResponseV1, + SyncRequest, region_request, }; use api::v1::{ResponseHeader, Status}; use arrow_flight::{FlightData, Ticket}; @@ -37,6 +42,8 @@ use async_trait::async_trait; use bytes::Bytes; use common_error::ext::{BoxedError, ErrorExt}; use common_error::status_code::StatusCode; +#[cfg(feature = "enterprise")] +use common_grpc::flow_write::{COMMAND_KEY, FlowWriteRegistry, Reply}; use common_meta::datanode::TopicStatsReporter; use common_query::OutputData; use common_query::request::QueryRequest; @@ -115,6 +122,8 @@ pub struct RegionServer { inner: Arc, flight_compression: FlightCompression, suspend: Arc, + #[cfg(feature = "enterprise")] + flow_write_registry: Arc>, } pub struct RegionStat { @@ -163,6 +172,8 @@ impl RegionServer { )), flight_compression, suspend: Arc::new(AtomicBool::new(false)), + #[cfg(feature = "enterprise")] + flow_write_registry: Arc::new(FlowWriteRegistry::default()), } } @@ -925,6 +936,25 @@ impl Stream for RegionWatermarkStream { #[async_trait] impl RegionServerHandler for RegionServer { + async fn handle_region_request( + &self, + request: RegionRequestV1, + ) -> ServerResult { + #[cfg(feature = "enterprise")] + if request + .header + .as_ref() + .and_then(|header| header.query_context.as_ref()) + .is_some_and(|context| context.extensions.contains_key(COMMAND_KEY)) + { + return self.handle_flow_write_request(request).await; + } + let body = request.body.context(servers_error::InvalidQuerySnafu { + reason: "Expecting non-empty region request body.", + })?; + self.handle(body).await + } + async fn handle(&self, request: region_request::Body) -> ServerResult { let failed_requests_cnt = crate::metrics::REGION_SERVER_REQUEST_FAILURE_COUNT .with_label_values(&[request.as_ref()]); diff --git a/src/datanode/src/region_server/flow_write.rs b/src/datanode/src/region_server/flow_write.rs new file mode 100644 index 00000000000..adb9d88568a --- /dev/null +++ b/src/datanode/src/region_server/flow_write.rs @@ -0,0 +1,959 @@ +// Copyright 2023-2026 GrepTime Inc. +// +// This file is part of the GreptimeDB Enterprise Edition and is licensed under +// the GreptimeDB Enterprise License. You may not use this file except in +// compliance with that license. A copy of the license is available at the root +// of this repository in the file LICENSE-ENTERPRISE. +// +// Unless required by applicable law or agreed to in writing, this software is +// distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, +// either express or implied. + +//! Datanode receiver for tracked Flow writes. +//! +//! A single immutable operation per receiver-issued capability: the first +//! admitted `Write` becomes the owner and executes the Inserts exactly once on +//! the region engines; concurrent or retried identical writers join the +//! owner's completion. Dropping a waiter (for example a lost RPC response) +//! never cancels the engine work, and closing a capability rejects later +//! writes with that capability. If the owner's outcome becomes unknown (a +//! child ingest worker was lost), no positive terminal completion is +//! published: waiters observe `WriterLost`; the capability admits no second +//! owner, so no overlapping engine write can be produced. + +use std::sync::Arc; + +use api::v1::region::RegionRequest as RegionRequestV1; +use api::v1::region::region_request; +use common_error::ext::ErrorExt; +use common_grpc::flow_write::{Admission, COMMAND_KEY, Command, Reply, TerminalError}; +use common_runtime::runtime::RuntimeTrait; +use common_telemetry::tracing_context::{FutureExt, TracingContext}; +use futures_util::future::join_all; +use prost::Message; +use servers::error::{InvalidQuerySnafu, Result as ServerResult}; +use snafu::OptionExt; + +use crate::region_server::RegionServer; + +/// Payload of a tracked write: encoded Inserts rows (exact bytes, so floats +/// keep bit equality even for NaN), schemas compared by protobuf equality, and +/// semantic header identity (dbname and command-free query context). Retrying +/// with different business data or context is rejected instead of joining. +pub(crate) type TrackedPayload = ( + Vec, + Vec>>, + String, + Option, +); + +fn tracked_payload(request: &RegionRequestV1) -> ServerResult { + let body = request.body.as_ref().context(InvalidQuerySnafu { + reason: "Expecting non-empty region request body.", + })?; + if !matches!(body, region_request::Body::Inserts(_)) { + return Err(servers::error::Error::InvalidQuery { + reason: "Tracked Flow write supports Inserts requests only.".to_string(), + location: snafu::Location::new(file!(), line!(), column!()), + }); + } + // Oneof is not `prost::Message`; the variant message itself is. Schema + // options contain unordered maps, so compare schemas separately rather than + // letting their wire ordering affect retries; retain encoded rows for NaN. + let (bytes, schemas) = match body { + region_request::Body::Inserts(inserts) => { + let mut inserts = inserts.clone(); + let schemas = inserts + .requests + .iter_mut() + .map(|request| { + request + .rows + .as_mut() + .map(|rows| std::mem::take(&mut rows.schema)) + }) + .collect(); + (inserts.encode_to_vec(), schemas) + } + _ => unreachable!("checked above"), + }; + let header = request.header.as_ref(); + let mut query_context = header.and_then(|header| header.query_context.clone()); + if let Some(context) = query_context.as_mut() { + // The command extension differs per control request, not per business. + context.extensions.remove(COMMAND_KEY); + } + Ok(( + bytes, + schemas, + header + .map(|header| header.dbname.clone()) + .unwrap_or_default(), + query_context, + )) +} + +fn ensure_empty_control_body(request: &RegionRequestV1) -> ServerResult<()> { + match request.body.as_ref() { + Some(region_request::Body::Inserts(inserts)) if inserts.requests.is_empty() => Ok(()), + _ => Err(servers::error::Error::InvalidQuery { + reason: "Tracked Flow write control requests must carry an empty Inserts body." + .to_string(), + location: snafu::Location::new(file!(), line!(), column!()), + }), + } +} + +fn registry_error(error: common_grpc::flow_write::RegistryError) -> servers::error::Error { + servers::error::Error::UnexpectedResult { + reason: format!("Tracked Flow write registry rejected request: {error:?}"), + location: snafu::Location::new(file!(), line!(), column!()), + } +} + +fn response(reply: Reply) -> ServerResult { + let encoded = + serde_json::to_vec(&reply).map_err(|error| servers::error::Error::UnexpectedResult { + reason: format!("Failed to encode tracked Flow write reply: {error}"), + location: snafu::Location::new(file!(), line!(), column!()), + })?; + let mut extensions = std::collections::HashMap::new(); + extensions.insert(COMMAND_KEY.to_string(), encoded); + Ok(api::v1::region::RegionResponse { + header: Some(api::v1::ResponseHeader { + status: Some(api::v1::Status { + status_code: common_error::status_code::StatusCode::Success as _, + ..Default::default() + }), + }), + affected_rows: 0, + extensions, + metadata: Vec::new(), + }) +} + +impl RegionServer { + pub(crate) async fn handle_flow_write_request( + &self, + request: RegionRequestV1, + ) -> ServerResult { + let command = request + .header + .as_ref() + .and_then(|header| header.query_context.as_ref()) + .and_then(|context| context.extensions.get(COMMAND_KEY)) + .and_then(|value| Command::decode(value).ok()); + let command = match command { + Some(command) => command, + None => { + // The extension exists but is not a decodable command: reject + // before any ordinary DML dispatch. + return Err(InvalidQuerySnafu { + reason: "Malformed tracked Flow write command.", + } + .build()); + } + }; + + match command { + Command::Open { nonce } => { + ensure_empty_control_body(&request)?; + let capability = self + .flow_write_registry + .open(&nonce) + .map_err(registry_error)?; + response(Reply::Opened { capability }) + } + Command::Close { capability } => { + ensure_empty_control_body(&request)?; + self.flow_write_registry + .close(&capability) + .await + .map_err(registry_error)?; + response(Reply::Closed) + } + Command::Write { capability } => { + let payload = tracked_payload(&request)?; + let admission = self + .flow_write_registry + .admit(&capability, payload) + .map_err(registry_error)?; + let reply = match admission { + Admission::Start(owner) => { + // Subscribe before the owner sender moves into the + // spawned task: this caller stays a plain waiter whose + // drop cannot cancel the engine work. + let receiver = owner.subscribe(); + let server = self.clone(); + let body = request.body; + let header = request.header; + self.inner.runtime.spawn(async move { + let completion = server + .run_tracked_write(header, body) + .trace(TracingContext::from_current_span().attach( + common_telemetry::tracing::info_span!( + "RegionServer::tracked_flow_write" + ), + )) + .await; + match completion { + Some(reply) => { + let _ = owner.send(Some(Arc::new(reply))); + } + None => { + // Unknown worker outcome: no positive terminal + // result is published. Dropping the sender lets + // waiters observe `WriterLost`. + drop(owner); + } + } + }); + common_grpc::flow_write::wait_completion(receiver) + .await + .map_err(registry_error)? + } + Admission::Join(waiter) => common_grpc::flow_write::wait_completion(waiter) + .await + .map_err(registry_error)?, + }; + response((*reply).clone()) + } + } + } + + /// Executes one admitted tracked write. Returns `None` only when the + /// outcome is unknown (a child worker was lost); malformed request bodies + /// are terminal because no child was ever dispatched. + async fn run_tracked_write( + &self, + header: Option, + body: Option, + ) -> Option { + let requests = + match body.map(store_api::region_request::RegionRequest::try_from_request_body) { + Some(Ok(requests)) => requests, + Some(Err(error)) => { + return Some(Reply::Complete { + rows: 0, + cost: 0, + metrics: None, + error: Some(TerminalError { + code: error.status_code() as u32, + message: error.output_msg(), + }), + }); + } + // `None` is unreachable here: the Write path already rejected an + // empty body before admission. Treat it as unknown rather than + // inventing a terminal result. + None => return None, + }; + let tracing = header + .as_ref() + .map(|header| TracingContext::from_w3c(&header.tracing_context)) + .unwrap_or_default(); + + let mut unknown_worker = false; + let mut first_error: Option = None; + let mut affected_rows = 0usize; + + let tasks = requests.into_iter().map(|(region_id, request)| { + let span = tracing.attach(common_telemetry::tracing::info_span!( + "RegionServer::tracked_region_write", + region_id = region_id.to_string() + )); + async move { self.handle_request(region_id, request).trace(span).await } + }); + // Drain every per-region child before any terminal decision. A lost + // ingest worker (`RuntimeJoin`) is unknown and outranks any ordinary + // application error observed during the same drain. + for result in join_all(tasks).await { + match result { + Ok(response) => affected_rows += response.affected_rows, + Err(crate::error::Error::RuntimeJoin { .. }) => { + unknown_worker = true; + } + Err(error) if first_error.is_none() => first_error = Some(error), + Err(_) => {} + } + } + + if unknown_worker { + return None; + } + + Some(match first_error { + None => Reply::Complete { + rows: affected_rows, + cost: affected_rows, + metrics: None, + error: None, + }, + Some(error) => Reply::Complete { + rows: 0, + cost: 0, + metrics: None, + error: Some(TerminalError { + code: error.status_code() as u32, + message: error.output_msg(), + }), + }, + }) + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + + use api::region::RegionResponse; + use api::v1::region::region_request::Body; + use api::v1::region::{ + InsertRequest, InsertRequests, RegionRequest as Request, RegionRequestHeader, + }; + use api::v1::{ColumnDataType, ColumnSchema, QueryContext, Row, Rows, SemanticType, Value}; + use common_grpc::flow_write::{self, Command, Reply}; + use servers::grpc::region_server::RegionServerHandler; + use store_api::region_engine::{RegionEngine, RegionEngineRef}; + use store_api::storage::RegionId; + use tokio::sync::Notify; + use tokio::time::{Duration, timeout}; + + use super::*; + use crate::tests::{MockRegionEngine, mock_region_server}; + + #[derive(Clone, Default)] + struct Gate { + entered: Arc, + notified: Arc, + open: Arc, + fail: bool, + } + + impl Gate { + async fn wait_entered(&self, count: usize) { + timeout(Duration::from_secs(5), async { + loop { + let notified = self.notified.notified(); + tokio::pin!(notified); + notified.as_mut().enable(); + if self.entered.load(Ordering::SeqCst) >= count { + break; + } + notified.await; + } + }) + .await + .unwrap(); + } + fn release(&self) { + self.open.store(1, Ordering::SeqCst); + self.notified.notify_waiters(); + } + async fn wait_open(&self) { + loop { + let notified = self.notified.notified(); + tokio::pin!(notified); + notified.as_mut().enable(); + if self.open.load(Ordering::SeqCst) != 0 { + break; + } + notified.await; + } + } + } + + struct GatedEngine { + gate: Gate, + inner: Arc, + panic_id: Option, + } + #[async_trait::async_trait] + impl RegionEngine for GatedEngine { + fn name(&self) -> &str { + self.inner.name() + } + async fn handle_request( + &self, + id: RegionId, + _request: store_api::region_request::RegionRequest, + ) -> Result { + self.gate.entered.fetch_add(1, Ordering::SeqCst); + self.gate.notified.notify_waiters(); + if self.panic_id == Some(id.as_u64()) { + panic!("test ingest panic"); + } + if self.gate.fail && id.as_u64() == 1 { + return Err(common_error::ext::BoxedError::new( + crate::error::RegionNotFoundSnafu { region_id: id }.build(), + )); + } + self.gate.wait_open().await; + Ok(RegionResponse::new(1)) + } + async fn handle_query( + &self, + id: RegionId, + req: store_api::storage::ScanRequest, + ) -> Result + { + self.inner.handle_query(id, req).await + } + async fn get_committed_sequence( + &self, + id: RegionId, + ) -> Result { + self.inner.get_committed_sequence(id).await + } + async fn get_metadata( + &self, + id: RegionId, + ) -> Result { + self.inner.get_metadata(id).await + } + fn region_statistic( + &self, + id: RegionId, + ) -> Option { + self.inner.region_statistic(id) + } + async fn stop(&self) -> Result<(), common_error::ext::BoxedError> { + self.inner.stop().await + } + fn set_region_role( + &self, + id: RegionId, + role: store_api::region_engine::RegionRole, + ) -> Result<(), common_error::ext::BoxedError> { + self.inner.set_region_role(id, role) + } + async fn sync_region( + &self, + id: RegionId, + req: store_api::region_engine::SyncRegionFromRequest, + ) -> Result + { + self.inner.sync_region(id, req).await + } + async fn remap_manifests( + &self, + req: store_api::region_engine::RemapManifestsRequest, + ) -> Result + { + self.inner.remap_manifests(req).await + } + async fn set_region_role_state_gracefully( + &self, + id: RegionId, + state: store_api::region_engine::SettableRegionRoleState, + ) -> Result< + store_api::region_engine::SetRegionRoleStateResponse, + common_error::ext::BoxedError, + > { + self.inner.set_region_role_state_gracefully(id, state).await + } + fn role(&self, id: RegionId) -> Option { + self.inner.role(id) + } + fn as_any(&self) -> &dyn std::any::Any { + self + } + } + + fn body(ids: &[u64], nan: bool) -> Body { + Body::Inserts(InsertRequests { + requests: ids + .iter() + .map(|id| InsertRequest { + region_id: *id, + rows: Some(Rows { + schema: vec![ + ColumnSchema { + column_name: "ts".into(), + datatype: ColumnDataType::TimestampMillisecond as i32, + semantic_type: SemanticType::Timestamp as i32, + datatype_extension: None, + options: None, + }, + ColumnSchema { + column_name: "v".into(), + datatype: ColumnDataType::Float64 as i32, + semantic_type: SemanticType::Field as i32, + datatype_extension: None, + options: None, + }, + ], + rows: vec![Row { + values: vec![ + Value { + value_data: Some( + api::v1::value::ValueData::TimestampMillisecondValue(1), + ), + }, + Value { + value_data: Some(api::v1::value::ValueData::F64Value(if nan { + f64::NAN + } else { + 7. + })), + }, + ], + }], + }), + partition_expr_version: None, + skip_wal: false, + }) + .collect(), + }) + } + + fn request(cmd: Command, body: Option) -> Request { + let mut extensions = HashMap::new(); + extensions.insert(flow_write::COMMAND_KEY.into(), cmd.encode().unwrap()); + Request { + header: Some(RegionRequestHeader { + dbname: "db".into(), + tracing_context: Default::default(), + query_context: Some(QueryContext { + extensions, + ..Default::default() + }), + }), + body, + } + } + fn decode(resp: &api::v1::region::RegionResponse) -> Reply { + assert_eq!( + resp.header + .as_ref() + .unwrap() + .status + .as_ref() + .unwrap() + .status_code, + common_error::status_code::StatusCode::Success as u32 + ); + serde_json::from_slice(resp.extensions.get(flow_write::COMMAND_KEY).unwrap()).unwrap() + } + async fn open(s: &RegionServer) -> String { + match decode( + &s.handle_region_request(request( + Command::Open { nonce: "n".into() }, + Some(Body::Inserts(InsertRequests { requests: vec![] })), + )) + .await + .unwrap(), + ) { + Reply::Opened { capability } => capability, + x => panic!("{x:?}"), + } + } + fn setup(gate: Gate, panic_id: Option, ids: &[u64]) -> RegionServer { + let s = mock_region_server(); + let (inner, _) = MockRegionEngine::new("mock"); + let engine: RegionEngineRef = Arc::new(GatedEngine { + gate, + inner, + panic_id, + }); + for id in ids { + s.inner.region_map.insert( + RegionId::from_u64(*id), + super::super::RegionEngineWithStatus::Ready(engine.clone()), + ); + } + s + } + fn write(cap: &str, ids: &[u64], nan: bool) -> Request { + request( + Command::Write { + capability: cap.into(), + }, + Some(body(ids, nan)), + ) + } + + #[tokio::test] + async fn dropped_rpc_retry_and_close_drain_exactly_once() { + let gate = Gate::default(); + let s = setup(gate.clone(), None, &[1, 2]); + let cap = open(&s).await; + let w = write(&cap, &[1, 2], false); + let first = tokio::spawn({ + let s = s.clone(); + let w = w.clone(); + async move { s.handle_region_request(w).await } + }); + gate.wait_entered(2).await; + first.abort(); + // Exercise the same actual registry admission used by retry, before release. + let payload = tracked_payload(&w).unwrap(); + let Admission::Join(waiter) = s.flow_write_registry.admit(&cap, payload).unwrap() else { + panic!("retry did not join") + }; + let mut retry_waiter = tokio::spawn(common_grpc::flow_write::wait_completion(waiter)); + assert!( + timeout(Duration::from_millis(20), &mut retry_waiter) + .await + .is_err() + ); + let mut closer = tokio::spawn({ + let s = s.clone(); + let cap = cap.clone(); + async move { + s.handle_region_request(request( + Command::Close { capability: cap }, + Some(Body::Inserts(InsertRequests { requests: vec![] })), + )) + .await + } + }); + assert!( + timeout(Duration::from_millis(20), &mut closer) + .await + .is_err() + ); + gate.release(); + let joined = timeout(Duration::from_secs(5), &mut retry_waiter) + .await + .unwrap() + .unwrap() + .unwrap(); + assert!(matches!( + *joined, + Reply::Complete { + rows: 2, + error: None, + .. + } + )); + assert!(matches!( + decode( + &timeout(Duration::from_secs(5), &mut closer) + .await + .unwrap() + .unwrap() + .unwrap() + ), + Reply::Closed + )); + assert!( + s.handle_region_request(write(&cap, &[1], false)) + .await + .is_err() + ); + let cap2 = open(&s).await; + gate.release(); + let r = s + .handle_region_request(write(&cap2, &[1], false)) + .await + .unwrap(); + assert!(matches!(decode(&r), Reply::Complete { rows: 1, .. })); + assert!( + s.handle_region_request(write(&cap, &[1], false)) + .await + .is_err() + ); + assert_eq!(gate.entered.load(Ordering::SeqCst), 3); + let close2 = request( + Command::Close { capability: cap2 }, + Some(Body::Inserts(InsertRequests { requests: vec![] })), + ); + assert!(matches!( + decode(&s.handle_region_request(close2).await.unwrap()), + Reply::Closed + )); + } + + #[tokio::test] + async fn application_error_waits_for_sibling_and_panic_overrides_it() { + let gate = Gate { + fail: true, + ..Default::default() + }; + let s = setup(gate.clone(), None, &[1, 2]); + let cap = open(&s).await; + let task = tokio::spawn({ + let s = s.clone(); + async move { s.handle_region_request(write(&cap, &[1, 2], false)).await } + }); + gate.wait_entered(2).await; + // The ordinary error is immediate for region 1, but region 2 remains gated. + assert!(!task.is_finished()); + gate.release(); + let r = timeout(Duration::from_secs(5), task) + .await + .unwrap() + .unwrap() + .unwrap(); + assert!(matches!(decode(&r), Reply::Complete { error: Some(_), .. })); + + let gate = Gate { + fail: true, + ..Default::default() + }; + let s = setup(gate.clone(), Some(2), &[1, 2, 3]); + let cap = open(&s).await; + let task = tokio::spawn({ + let s = s.clone(); + let cap = cap.clone(); + async move { + s.handle_region_request(write(&cap, &[1, 2, 3], false)) + .await + } + }); + gate.wait_entered(3).await; + assert!(!task.is_finished()); + gate.release(); + assert!( + timeout(Duration::from_secs(5), task) + .await + .unwrap() + .unwrap() + .is_err() + ); + assert!( + s.handle_region_request(write(&cap, &[1, 2, 3], false)) + .await + .is_err() + ); + assert_eq!(gate.entered.load(Ordering::SeqCst), 3); + } + + #[tokio::test] + async fn decoded_schema_option_maps_join_live_nan_write() { + let gate = Gate::default(); + let s = setup(gate.clone(), None, &[1]); + let cap = open(&s).await; + let mut original = write(&cap, &[1], true); + let Body::Inserts(inserts) = original.body.as_mut().unwrap() else { + unreachable!() + }; + inserts.requests[0].rows.as_mut().unwrap().schema[1].options = + Some(api::v1::ColumnOptions { + options: [ + ("alpha".to_string(), "one".to_string()), + ("beta".to_string(), "two".to_string()), + ("gamma".to_string(), "three".to_string()), + ] + .into_iter() + .collect(), + }); + let encoded = original.encode_to_vec(); + let Body::Inserts(original_inserts) = original.body.as_ref().unwrap() else { + unreachable!() + }; + let original_rows = original_inserts.requests[0].rows.clone(); + let original_body_bytes = original_inserts.encode_to_vec(); + let mut retries = Vec::new(); + let mut found_different_wire_order = false; + for _ in 0..128 { + let decoded = Request::decode(encoded.as_slice()).unwrap(); + let Body::Inserts(decoded_inserts) = decoded.body.as_ref().unwrap() else { + unreachable!() + }; + assert_eq!( + decoded_inserts.requests[0].rows.as_ref().unwrap().schema, + original_rows.as_ref().unwrap().schema + ); + assert_eq!( + tracked_payload(&decoded).unwrap(), + tracked_payload(&original).unwrap() + ); + found_different_wire_order |= decoded_inserts.encode_to_vec() != original_body_bytes; + retries.push(decoded); + if found_different_wire_order && retries.len() >= 16 { + break; + } + } + assert!(found_different_wire_order); + + let owner = tokio::spawn({ + let s = s.clone(); + async move { s.handle_region_request(original).await } + }); + gate.wait_entered(1).await; + let retry_tasks: Vec<_> = retries + .into_iter() + .map(|retry| { + let s = s.clone(); + tokio::spawn(async move { s.handle_region_request(retry).await }) + }) + .collect(); + tokio::task::yield_now().await; + assert!(retry_tasks.iter().all(|retry| !retry.is_finished())); + assert_eq!(gate.entered.load(Ordering::SeqCst), 1); + gate.release(); + let owner_reply = decode( + &timeout(Duration::from_secs(5), owner) + .await + .unwrap() + .unwrap() + .unwrap(), + ); + assert!(matches!( + &owner_reply, + Reply::Complete { + rows: 1, + error: None, + .. + } + )); + let owner_reply_json = serde_json::to_string(&owner_reply).unwrap(); + for retry in retry_tasks { + let retry_reply = decode( + &timeout(Duration::from_secs(5), retry) + .await + .unwrap() + .unwrap() + .unwrap(), + ); + assert!(matches!( + &retry_reply, + Reply::Complete { + rows: 1, + error: None, + .. + } + )); + assert_eq!( + serde_json::to_string(&retry_reply).unwrap(), + owner_reply_json + ); + } + assert_eq!(gate.entered.load(Ordering::SeqCst), 1); + + let replay = Request::decode(encoded.as_slice()).unwrap(); + assert!(matches!( + decode(&s.handle_region_request(replay.clone()).await.unwrap()), + Reply::Complete { rows: 1, .. } + )); + let mut changed_schema = replay.clone(); + let Body::Inserts(inserts) = changed_schema.body.as_mut().unwrap() else { + unreachable!() + }; + inserts.requests[0].rows.as_mut().unwrap().schema[1] + .options + .as_mut() + .unwrap() + .options + .insert("gamma".into(), "changed".into()); + assert!(s.handle_region_request(changed_schema).await.is_err()); + let close = request( + Command::Close { + capability: cap.clone(), + }, + Some(Body::Inserts(InsertRequests { requests: vec![] })), + ); + assert!(matches!( + decode(&s.handle_region_request(close).await.unwrap()), + Reply::Closed + )); + assert!(s.handle_region_request(replay).await.is_err()); + assert_eq!(gate.entered.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn payload_identity_nan_semantics_and_pre_admission_rejection() { + let s = mock_region_server(); + let cap = open(&s).await; + let a = write(&cap, &[1], true); + let b = write(&cap, &[1], true); + assert_eq!(tracked_payload(&a).unwrap(), tracked_payload(&b).unwrap()); + let owner = match s + .flow_write_registry + .admit(&cap, tracked_payload(&a).unwrap()) + .unwrap() + { + Admission::Start(owner) => owner, + _ => panic!(), + }; + let mut semantic_mismatch = b.clone(); + semantic_mismatch.header.as_mut().unwrap().dbname = "other_db".into(); + assert_eq!( + s.flow_write_registry + .admit(&cap, tracked_payload(&semantic_mismatch).unwrap()) + .unwrap_err(), + flow_write::RegistryError::PayloadMismatch + ); + semantic_mismatch = b.clone(); + semantic_mismatch + .header + .as_mut() + .unwrap() + .query_context + .as_mut() + .unwrap() + .extensions + .insert("flow_lane".into(), "changed".into()); + assert_eq!( + s.flow_write_registry + .admit(&cap, tracked_payload(&semantic_mismatch).unwrap()) + .unwrap_err(), + flow_write::RegistryError::PayloadMismatch + ); + let Admission::Join(waiter) = s + .flow_write_registry + .admit(&cap, tracked_payload(&b).unwrap()) + .unwrap() + else { + panic!() + }; + let mut joined = tokio::spawn(common_grpc::flow_write::wait_completion(waiter)); + assert!( + timeout(Duration::from_millis(20), &mut joined) + .await + .is_err() + ); + owner + .send(Some(Arc::new(Reply::Complete { + rows: 1, + cost: 1, + metrics: None, + error: None, + }))) + .unwrap(); + assert!(matches!( + *timeout(Duration::from_secs(5), &mut joined) + .await + .unwrap() + .unwrap() + .unwrap(), + Reply::Complete { rows: 1, .. } + )); + s.flow_write_registry.close(&cap).await.unwrap(); + assert_ne!( + tracked_payload(&a).unwrap(), + tracked_payload(&write(&cap, &[1], false)).unwrap() + ); + assert!( + s.handle_region_request(request( + Command::Write { + capability: cap.clone() + }, + Some(Body::Deletes(Default::default())) + )) + .await + .is_err() + ); + let fresh = open(&s).await; + let mut bad = request( + Command::Write { + capability: fresh.clone(), + }, + Some(body(&[1], false)), + ); + bad.header + .as_mut() + .unwrap() + .query_context + .as_mut() + .unwrap() + .extensions + .insert(flow_write::COMMAND_KEY.into(), "not-json".into()); + assert!(s.handle_region_request(bad).await.is_err()); + s.flow_write_registry.close(&fresh).await.unwrap(); + } +} diff --git a/src/flow/Cargo.toml b/src/flow/Cargo.toml index fea7b1eeda2..2cf316ad245 100644 --- a/src/flow/Cargo.toml +++ b/src/flow/Cargo.toml @@ -4,6 +4,9 @@ version.workspace = true edition.workspace = true license.workspace = true +[features] +enterprise = ["common-grpc/enterprise", "session/enterprise"] + [lints] workspace = true diff --git a/src/flow/src/batching_mode/frontend_client.rs b/src/flow/src/batching_mode/frontend_client.rs index 822d8dea45b..a06055298d9 100644 --- a/src/flow/src/batching_mode/frontend_client.rs +++ b/src/flow/src/batching_mode/frontend_client.rs @@ -95,8 +95,12 @@ pub enum FrontendClient { control_channel_manager: ChannelManager, query: QueryOptions, batch_opts: BatchingModeOptions, + #[cfg(feature = "enterprise")] + tracked_writes: bool, }, Standalone { + #[cfg(feature = "enterprise")] + tracked_writes: bool, /// for the sake of simplicity still use grpc even in standalone mode /// notice the client here should all be lazy, so that can wait after frontend is booted then make conn database_client: HandlerMutable, @@ -105,6 +109,26 @@ pub enum FrontendClient { } impl FrontendClient { + #[cfg(feature = "enterprise")] + pub fn with_tracked_writes(&self) -> Self { + let mut client = self.clone(); + match &mut client { + Self::Distributed { tracked_writes, .. } | Self::Standalone { tracked_writes, .. } => { + *tracked_writes = true; + } + } + client + } + + #[cfg(feature = "enterprise")] + fn tracked_writes_enabled(&self) -> bool { + match self { + Self::Distributed { tracked_writes, .. } | Self::Standalone { tracked_writes, .. } => { + *tracked_writes + } + } + } + /// Create a new empty frontend client, with a `HandlerMutable` to set the grpc handler later pub fn from_empty_grpc_handler(query: QueryOptions) -> (Self, HandlerMutable) { let is_initialized = Arc::new(SetOnce::new()); @@ -116,6 +140,8 @@ impl FrontendClient { Self::Standalone { database_client: handler.clone(), query, + #[cfg(feature = "enterprise")] + tracked_writes: false, }, handler, ) @@ -152,6 +178,8 @@ impl FrontendClient { control_channel_manager, query, batch_opts, + #[cfg(feature = "enterprise")] + tracked_writes: false, }) } @@ -168,6 +196,8 @@ impl FrontendClient { Self::Standalone { database_client: handler, query, + #[cfg(feature = "enterprise")] + tracked_writes: false, } } } @@ -205,7 +235,39 @@ impl DatabaseWithPeer { } } +#[cfg(feature = "enterprise")] +pub(crate) mod flow_write; + impl FrontendClient { + /// Preflight tracked-write support by opening and closing a capability without writing. + #[cfg(feature = "enterprise")] + pub async fn check_tracked_writes(&self, catalog: &str, schema: &str) -> Result<(), Error> { + if !self.tracked_writes_enabled() { + return UnexpectedSnafu { + reason: "Tracked writes are not enabled on this frontend client", + } + .fail(); + } + flow_write::check_tracked_writes(self, catalog, schema).await + } + + /// Execute a finite SQL write with receiver-issued completion tracking. + #[cfg(feature = "enterprise")] + pub async fn tracked_sql_write( + &self, + catalog: &str, + schema: &str, + sql: &str, + ) -> Result { + if !self.tracked_writes_enabled() { + return UnexpectedSnafu { + reason: "Tracked writes are not enabled on this frontend client", + } + .fail(); + } + flow_write::tracked_sql_write(self, catalog, schema, sql).await + } + /// scan for available frontend from metadata pub(crate) async fn scan_for_frontend(&self) -> Result, Error> { let Self::Distributed { meta_client, .. } = self else { @@ -295,6 +357,7 @@ impl FrontendClient { control_channel_manager, query: _, batch_opts, + .. } = self else { return UnexpectedSnafu { @@ -495,6 +558,33 @@ impl FrontendClient { peer_desc: &mut Option, ) -> Result { let flow_extensions = build_flow_extensions(extensions)?; + #[cfg(feature = "enterprise")] + if self.tracked_writes_enabled() && is_tracked_insert_request(&request) { + let effective_extensions: HashMap = extensions + .iter() + .map(|(key, value)| ((*key).to_string(), (*value).to_string())) + .collect(); + if effective_extensions.contains_key(common_grpc::flow_write::COMMAND_KEY) { + return UnexpectedSnafu { + reason: "Client-supplied tracked Flow command is forbidden", + } + .fail(); + } + let effective_extension_refs: Vec<(&str, &str)> = effective_extensions + .iter() + .map(|(key, value)| (key.as_str(), value.as_str())) + .collect(); + return flow_write::tracked_query( + self, + catalog, + schema, + Some(request), + &effective_extension_refs, + snapshot_seqs, + peer_desc, + ) + .await; + } match self { FrontendClient::Distributed { query, batch_opts, .. @@ -522,6 +612,7 @@ impl FrontendClient { FrontendClient::Standalone { database_client, query, + .. } => { *peer_desc = Some(PeerDesc::Standalone); let mut extensions_map = HashMap::from([( @@ -600,6 +691,7 @@ impl FrontendClient { FrontendClient::Standalone { database_client, query, + .. } => { let ctx = QueryContextBuilder::default() .current_catalog(catalog.to_string()) @@ -681,6 +773,7 @@ impl FrontendClient { FrontendClient::Standalone { database_client, query, + .. } => { let ctx = QueryContextBuilder::default() .current_catalog(catalog.to_string()) @@ -729,6 +822,15 @@ impl FrontendClient { } } +#[cfg(feature = "enterprise")] +fn is_tracked_insert_request(request: &QueryRequest) -> bool { + match request.query.as_ref() { + Some(Query::InsertIntoPlan(_)) => true, + Some(Query::Sql(sql)) => sql.trim_start().to_ascii_uppercase().starts_with("INSERT "), + _ => false, + } +} + fn build_flow_extensions(extensions: &[(&str, &str)]) -> Result { let flow_extensions = HashMap::from_iter( extensions diff --git a/src/flow/src/batching_mode/frontend_client/flow_write.rs b/src/flow/src/batching_mode/frontend_client/flow_write.rs new file mode 100644 index 00000000000..0d7663ccd4d --- /dev/null +++ b/src/flow/src/batching_mode/frontend_client/flow_write.rs @@ -0,0 +1,758 @@ +// Copyright 2023-2026 GrepTime Inc. +// +// This file is part of the GreptimeDB Enterprise Edition and is licensed under +// the GreptimeDB Enterprise License. You may not use this file except in +// compliance with that license. A copy of the license is available at the root +// of this repository in the file LICENSE-ENTERPRISE. +// +// Unless required by applicable law or agreed to in writing, this software is +// distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, +// either express or implied. + +use std::collections::HashMap; +use std::sync::Arc; +use std::time::Duration; + +use api::v1::query_request::Query; +use api::v1::{QueryRequest, greptime_request::Request}; +use client::OutputWithMetrics; +use common_error::ext::{BoxedError, ErrorExt}; +use common_error::status_code::StatusCode; +use common_grpc::flow_write::{self, COMMAND_KEY, Command, Reply, TerminalError}; +use common_query::{Output, OutputData, OutputMeta}; +use common_telemetry::warn; +use session::context::{Channel, QueryContextBuilder, generate_remote_query_id}; +use snafu::OptionExt; +use tokio::time::sleep; + +use crate::Error; +use crate::batching_mode::frontend_client::{DatabaseWithPeer, FrontendClient, PeerDesc}; +use crate::error::UnexpectedSnafu; + +const RETRY_PACE: Duration = Duration::from_secs(1); + +pub(crate) async fn tracked_sql_write( + frontend: &FrontendClient, + catalog: &str, + schema: &str, + sql: &str, +) -> Result { + tracked_query( + frontend, + catalog, + schema, + Some(QueryRequest { + query: Some(Query::Sql(sql.to_string())), + }), + &[], + &HashMap::new(), + &mut None, + ) + .await +} + +pub(crate) async fn check_tracked_writes( + frontend: &FrontendClient, + catalog: &str, + schema: &str, +) -> Result<(), Error> { + tracked_query( + frontend, + catalog, + schema, + None, + &[], + &HashMap::new(), + &mut None, + ) + .await + .map(|_| ()) +} + +pub(crate) async fn tracked_query( + frontend: &FrontendClient, + catalog: &str, + schema: &str, + request: Option, + extensions: &[(&str, &str)], + snapshot_seqs: &HashMap, + peer_desc: &mut Option, +) -> Result { + let (selected, handler, parallelism, read_preference, timeout) = match frontend { + FrontendClient::Distributed { + query, batch_opts, .. + } => ( + Some( + frontend + .get_random_active_frontend(catalog, schema) + .await + .map_err(prewrite_error)?, + ), + None, + query.parallelism.to_string(), + Some(batch_opts.read_preference.as_ref().to_string()), + batch_opts.experimental_flight_do_get_timeout, + ), + FrontendClient::Standalone { + database_client, + query, + .. + } => { + let handler = database_client + .handler + .lock() + .map_err(|e| { + prewrite_reason( + format!("Failed to lock database client: {e}"), + StatusCode::Internal, + ) + })? + .as_ref() + .context(UnexpectedSnafu { + reason: "Standalone's frontend instance is not set", + }) + .map_err(prewrite_error)? + .upgrade() + .context(UnexpectedSnafu { + reason: "Failed to upgrade database client", + }) + .map_err(prewrite_error)?; + ( + None, + Some(handler), + query.parallelism.to_string(), + None, + RETRY_PACE, + ) + } + }; + *peer_desc = Some(match &selected { + Some(DatabaseWithPeer { peer, .. }) => PeerDesc::Dist { peer: peer.clone() }, + None => PeerDesc::Standalone, + }); + + let nonce = generate_remote_query_id(); + let mut command = Command::Open { nonce }; + let mut completion: Option<( + usize, + usize, + Option, + Option, + )> = None; + + loop { + let encoded = command.encode().map_err(|error| { + UnexpectedSnafu { + reason: format!("Failed to encode tracked Flow command: {error}"), + } + .build() + })?; + let mut operation_extensions: Vec<(String, String)> = extensions + .iter() + .filter(|(key, _)| *key != COMMAND_KEY) + .map(|(key, value)| ((*key).to_string(), (*value).to_string())) + .collect(); + operation_extensions.push((COMMAND_KEY.to_string(), encoded)); + let extension_refs: Vec<(&str, &str)> = operation_extensions + .iter() + .map(|(k, v)| (k.as_str(), v.as_str())) + .collect(); + let phase_request = match &command { + Command::Write { .. } => request + .as_ref() + .context(UnexpectedSnafu { + reason: "Tracked Flow write command has no request", + })? + .clone(), + Command::Open { .. } | Command::Close { .. } => QueryRequest { + query: Some(Query::Sql(String::new())), + }, + }; + let response = if let Some(DatabaseWithPeer { database, .. }) = &selected { + let mut hints = vec![( + query::datafusion::QUERY_PARALLELISM_HINT, + parallelism.as_str(), + )]; + if let Some(preference) = read_preference.as_deref() { + hints.push((session::hints::READ_PREFERENCE_HINT, preference)); + } + database + .flight_request() + .with_hints(&hints) + .with_flow_extensions(&extension_refs) + .with_snapshot_seqs(snapshot_seqs) + .with_timeout(timeout) + .query_with_terminal_metrics(phase_request) + .await + .map(|output| output.output) + .map_err(|e| (e.status_code(), e.to_string())) + } else { + let mut ctx_extensions: HashMap = + operation_extensions.into_iter().collect(); + ctx_extensions.insert( + query::datafusion::QUERY_PARALLELISM_HINT.to_string(), + parallelism.clone(), + ); + let ctx = QueryContextBuilder::default() + .current_catalog(catalog.to_string()) + .current_schema(schema.to_string()) + .channel(Channel::Internal) + .extensions(ctx_extensions) + .snapshot_seqs(Arc::new(std::sync::RwLock::new(snapshot_seqs.clone()))) + .build(); + let handler = handler.as_ref().ok_or_else(|| { + UnexpectedSnafu { + reason: "Standalone handler was not pinned", + } + .build() + })?; + handler + .do_query(Request::Query(phase_request), Arc::new(ctx)) + .await + .map_err(|e| (e.status_code(), e.to_string())) + }; + let reply = match response { + Ok(output) => match decode_output(output).await { + Ok(reply) => Some(reply), + Err(e) if matches!(&command, Command::Open { .. }) => { + return Err(prewrite_reason( + format!("Frontend does not support tracked Flow writes: {e}"), + StatusCode::Unsupported, + )); + } + Err(e) => { + warn!("Malformed tracked Flow {:?} reply: {}", command, e); + None + } + }, + Err((code, msg)) => { + // Only Open is pre-admission; definitive rejection there is safe to surface. + if matches!(&command, Command::Open { .. }) && is_definitive_rejection(code) { + return Err(prewrite_reason( + format!("Frontend rejected tracked Flow open: {msg}"), + code, + )); + } + warn!( + "Tracked Flow {:?} transport failure (code={:?}): {}", + command, code, msg + ); + None + } + }; + match (command.clone(), reply) { + (Command::Open { .. }, Some(Reply::Opened { capability: cap })) if !cap.is_empty() => { + command = if request.is_some() { + Command::Write { capability: cap } + } else { + Command::Close { capability: cap } + }; + continue; + } + (Command::Open { .. }, Some(other)) => { + return Err(prewrite_reason( + format!("Invalid tracked Flow open reply: {other:?}"), + StatusCode::InvalidArguments, + )); + } + ( + Command::Write { capability }, + Some(Reply::Complete { + rows, + cost, + metrics, + error, + }), + ) => { + completion = Some((rows, cost, metrics, error)); + command = Command::Close { capability }; + continue; + } + (Command::Write { .. }, Some(other)) => { + warn!("Unexpected tracked Flow write reply: {other:?}") + } + (Command::Close { .. }, Some(Reply::Closed)) => { + let Some((rows, cost, metrics, error)) = completion.take() else { + if request.is_none() { + return Ok(OutputWithMetrics::from_output(Output::new( + OutputData::AffectedRows(0), + OutputMeta::default(), + ))); + } + return UnexpectedSnafu { + reason: "Close completed without positive write result", + } + .fail(); + }; + if let Some(error) = error { + return Err(Error::External { + source: BoxedError::new(error), + location: snafu::Location::new(file!(), line!(), column!()), + }); + } + let output = OutputWithMetrics::from_output(Output::new( + OutputData::AffectedRows(rows), + OutputMeta { + cost, + ..Default::default() + }, + )); + output.metrics.update(metrics); + return Ok(output); + } + (Command::Close { .. }, Some(other)) => { + warn!("Unexpected tracked Flow close reply: {other:?}") + } + (_, None) => {} + } + sleep(RETRY_PACE).await; + } +} + +fn prewrite_error(error: Error) -> Error { + let status = error.status_code(); + prewrite_reason(error.to_string(), status) +} + +/// Returns true when a tracked Flow command error is a definitive, non-retryable +/// rejection from the frontend rather than an ambiguous transport failure. These +/// codes indicate the frontend itself refused the command (unsupported protocol, +/// malformed/empty command, or wrong channel), so retrying the same pinned +/// frontend can never succeed; the caller must fail fast instead of looping. +fn is_definitive_rejection(code: StatusCode) -> bool { + matches!( + code, + StatusCode::Unsupported | StatusCode::InvalidArguments | StatusCode::PermissionDenied + ) +} + +fn prewrite_reason(message: String, status: StatusCode) -> Error { + Error::External { + source: BoxedError::new(TerminalError { + code: status as u32, + message, + }), + location: snafu::Location::new(file!(), line!(), column!()), + } +} + +async fn decode_output(output: Output) -> Result { + let batches = match output.data { + OutputData::RecordBatches(batches) => batches, + OutputData::Stream(stream) => common_recordbatch::util::collect_batches(stream) + .await + .map_err(|e| e.to_string())?, + _ => return Err("tracked control request returned a non-record-batch result".to_string()), + }; + flow_write::decode_reply(&batches).map_err(|e| e.to_string()) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::atomic::{AtomicUsize, Ordering}; + + use common_error::ext::{BoxedError, PlainError}; + use common_error::status_code::StatusCode; + use session::context::QueryContextRef; + + #[derive(Debug, Default)] + struct LostReplyHandler { + opens: AtomicUsize, + writes: AtomicUsize, + closes: AtomicUsize, + executions: AtomicUsize, + registry: common_grpc::flow_write::FlowWriteRegistry, + capabilities: std::sync::Mutex>, + } + + #[async_trait::async_trait] + impl super::super::GrpcQueryHandlerWithBoxedError for LostReplyHandler { + async fn do_query( + &self, + request: Request, + ctx: QueryContextRef, + ) -> Result { + let encoded = ctx.extensions().get(COMMAND_KEY).cloned().ok_or_else(|| { + BoxedError::new(PlainError::new( + "missing command".into(), + StatusCode::Unexpected, + )) + })?; + match Command::decode(&encoded).map_err(|error| { + BoxedError::new(PlainError::new( + error.to_string(), + StatusCode::InvalidArguments, + )) + })? { + Command::Open { nonce } => { + let call = self.opens.fetch_add(1, Ordering::SeqCst); + let capability = self.registry.open(&nonce).unwrap(); + let previous = self + .capabilities + .lock() + .unwrap() + .insert(nonce, capability.clone()); + if let Some(previous) = previous { + assert_eq!(previous, capability); + } + if call == 0 { + return Err(transport_error("lost Open reply")); + } + assert_empty_sql(&request); + response(Reply::Opened { capability }) + } + Command::Write { capability } => { + let sql = request_sql(&request); + let call = self.writes.fetch_add(1, Ordering::SeqCst); + let admission = self.registry.admit(&capability, sql.clone()).unwrap(); + let result = match admission { + common_grpc::flow_write::Admission::Start(sender) => { + self.executions.fetch_add(1, Ordering::SeqCst); + let reply = Reply::Complete { + rows: 1, + cost: 9, + metrics: Some(common_recordbatch::adapter::RecordBatchMetrics { + region_watermarks: vec![ + common_recordbatch::adapter::RegionWatermarkEntry { + region_id: 41, + watermark: Some(17), + }, + ], + ..Default::default() + }), + error: None, + }; + sender.send(Some(Arc::new(reply.clone()))).unwrap(); + reply + } + common_grpc::flow_write::Admission::Join(receiver) => { + common_grpc::flow_write::wait_completion(receiver) + .await + .unwrap() + .as_ref() + .clone() + } + }; + if call == 0 { + return Err(transport_error("lost Write reply")); + } + response(result) + } + Command::Close { capability } => { + let call = self.closes.fetch_add(1, Ordering::SeqCst); + self.registry.close(&capability).await.unwrap(); + if call == 0 { + return Err(transport_error("lost Close reply")); + } + response(Reply::Closed) + } + } + } + } + + fn request_sql(request: &Request) -> String { + match request { + Request::Query(query) => match query.query.as_ref() { + Some(Query::Sql(sql)) => sql.clone(), + other => panic!("expected SQL INSERT, got {other:?}"), + }, + other => panic!("expected Query request, got {other:?}"), + } + } + + fn assert_empty_sql(request: &Request) { + assert!( + matches!(request, Request::Query(QueryRequest { query: Some(Query::Sql(sql)) }) if sql.is_empty()) + ); + } + + fn transport_error(message: &str) -> BoxedError { + BoxedError::new(PlainError::new( + message.to_string(), + StatusCode::StorageUnavailable, + )) + } + + fn response(reply: Reply) -> Result { + let batches = flow_write::reply_batches(&reply).map_err(BoxedError::new)?; + Ok(Output::new_with_record_batches(batches)) + } + + #[tokio::test] + async fn retries_same_open_write_and_close_capability_and_executes_once() { + let handler = Arc::new(LostReplyHandler::default()); + let trait_handler: Arc = handler.clone(); + let weak = Arc::downgrade(&trait_handler); + let frontend = + FrontendClient::from_grpc_handler(weak, query::options::QueryOptions::default()) + .with_tracked_writes(); + let output = frontend + .tracked_sql_write("greptime", "public", "INSERT INTO t VALUES (7)") + .await + .unwrap(); + assert!(matches!(output.output.data, OutputData::AffectedRows(1))); + assert_eq!(output.output.meta.cost, 9); + assert_eq!( + output.metrics.region_watermark_map(), + Some(HashMap::from([(41, 17)])) + ); + assert_eq!(handler.opens.load(Ordering::SeqCst), 2); + assert_eq!(handler.writes.load(Ordering::SeqCst), 2); + assert_eq!(handler.closes.load(Ordering::SeqCst), 2); + assert_eq!(handler.executions.load(Ordering::SeqCst), 1); + let second = frontend + .tracked_sql_write("greptime", "public", "INSERT INTO t VALUES (8)") + .await + .unwrap(); + assert!(matches!(second.output.data, OutputData::AffectedRows(1))); + assert_eq!(handler.executions.load(Ordering::SeqCst), 2); + } + + /// A legacy frontend that does not understand the tracked-write extension + /// processes `Open` as an ordinary empty-SQL query and rejects it with + /// `InvalidArguments`. That is a definitive application rejection, not an + /// ambiguous transport failure: the client must fail fast instead of + /// retrying the same pinned frontend forever. + #[derive(Debug, Default)] + struct LegacyRejectHandler { + opens: AtomicUsize, + } + + #[async_trait::async_trait] + impl super::super::GrpcQueryHandlerWithBoxedError for LegacyRejectHandler { + async fn do_query( + &self, + _request: Request, + _ctx: QueryContextRef, + ) -> Result { + self.opens.fetch_add(1, Ordering::SeqCst); + Err(BoxedError::new(PlainError::new( + "empty statements".to_string(), + StatusCode::InvalidArguments, + ))) + } + } + + #[derive(Debug)] + struct PreflightHandler { + opens: AtomicUsize, + closes: AtomicUsize, + reject_open: bool, + } + + #[async_trait::async_trait] + impl super::super::GrpcQueryHandlerWithBoxedError for PreflightHandler { + async fn do_query( + &self, + request: Request, + ctx: QueryContextRef, + ) -> Result { + let encoded = ctx.extensions().get(COMMAND_KEY).cloned().unwrap(); + assert_empty_sql(&request); + match Command::decode(&encoded).unwrap() { + Command::Open { .. } => { + self.opens.fetch_add(1, Ordering::SeqCst); + if self.reject_open { + Err(BoxedError::new(PlainError::new( + "preflight rejected".into(), + StatusCode::PermissionDenied, + ))) + } else { + response(Reply::Opened { + capability: "preflight-capability".into(), + }) + } + } + Command::Close { capability } => { + assert_eq!(capability, "preflight-capability"); + self.closes.fetch_add(1, Ordering::SeqCst); + response(Reply::Closed) + } + Command::Write { .. } => panic!("preflight must not write"), + } + } + } + + #[tokio::test] + async fn preflight_opens_and_closes_without_writing() { + let handler = Arc::new(PreflightHandler { + opens: AtomicUsize::default(), + closes: AtomicUsize::default(), + reject_open: false, + }); + let trait_handler: Arc = handler.clone(); + let frontend = FrontendClient::from_grpc_handler( + Arc::downgrade(&trait_handler), + query::options::QueryOptions::default(), + ) + .with_tracked_writes(); + tokio::time::timeout( + Duration::from_secs(10), + frontend.check_tracked_writes("greptime", "public"), + ) + .await + .expect("tracked-write preflight should finish within 10 seconds") + .unwrap(); + assert_eq!(handler.opens.load(Ordering::SeqCst), 1); + assert_eq!(handler.closes.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn preflight_preserves_definitive_open_error() { + let handler = Arc::new(PreflightHandler { + opens: AtomicUsize::default(), + closes: AtomicUsize::default(), + reject_open: true, + }); + let trait_handler: Arc = handler.clone(); + let frontend = FrontendClient::from_grpc_handler( + Arc::downgrade(&trait_handler), + query::options::QueryOptions::default(), + ) + .with_tracked_writes(); + let error = tokio::time::timeout( + Duration::from_secs(10), + frontend.check_tracked_writes("greptime", "public"), + ) + .await + .expect("tracked-write preflight rejection should finish within 10 seconds") + .unwrap_err(); + assert_eq!(error.status_code(), StatusCode::PermissionDenied); + assert!(matches!( + &error, + Error::External { source, .. } + if source.as_any().downcast_ref::().is_some_and(|terminal| { + terminal.code == StatusCode::PermissionDenied as u32 + && terminal.message.contains("preflight rejected") + }) + )); + assert_eq!(handler.opens.load(Ordering::SeqCst), 1); + assert_eq!(handler.closes.load(Ordering::SeqCst), 0); + } + + #[derive(Debug)] + struct RejectionRetryHandler { + status: StatusCode, + opens: AtomicUsize, + writes: AtomicUsize, + closes: AtomicUsize, + capability: std::sync::Mutex>, + } + + #[async_trait::async_trait] + impl super::super::GrpcQueryHandlerWithBoxedError for RejectionRetryHandler { + async fn do_query( + &self, + request: Request, + ctx: QueryContextRef, + ) -> Result { + let encoded = ctx.extensions().get(COMMAND_KEY).cloned().unwrap(); + match Command::decode(&encoded).unwrap() { + Command::Open { nonce } => { + assert_empty_sql(&request); + self.opens.fetch_add(1, Ordering::SeqCst); + let mut capability = self.capability.lock().unwrap(); + let value = capability.get_or_insert_with(|| nonce); + response(Reply::Opened { + capability: value.clone(), + }) + } + Command::Write { capability } => { + assert_eq!(request_sql(&request), "INSERT INTO t VALUES (11)"); + assert_eq!(self.capability.lock().unwrap().as_ref(), Some(&capability)); + if self.writes.fetch_add(1, Ordering::SeqCst) == 0 { + Err(BoxedError::new(PlainError::new( + "write rejection".into(), + self.status, + ))) + } else { + response(Reply::Complete { + rows: 2, + cost: 0, + metrics: None, + error: None, + }) + } + } + Command::Close { capability } => { + assert_empty_sql(&request); + assert_eq!(self.capability.lock().unwrap().as_ref(), Some(&capability)); + if self.closes.fetch_add(1, Ordering::SeqCst) == 0 { + Err(BoxedError::new(PlainError::new( + "close rejection".into(), + self.status, + ))) + } else { + response(Reply::Closed) + } + } + } + } + } + + #[tokio::test] + async fn post_admission_rejections_retry_same_capability_until_close() { + for status in [ + StatusCode::Unsupported, + StatusCode::InvalidArguments, + StatusCode::PermissionDenied, + ] { + let handler = Arc::new(RejectionRetryHandler { + status, + opens: AtomicUsize::default(), + writes: AtomicUsize::default(), + closes: AtomicUsize::default(), + capability: std::sync::Mutex::new(None), + }); + let trait_handler: Arc = + handler.clone(); + let frontend = FrontendClient::from_grpc_handler( + Arc::downgrade(&trait_handler), + query::options::QueryOptions::default(), + ) + .with_tracked_writes(); + let output = tokio::time::timeout( + Duration::from_secs(10), + frontend.tracked_sql_write("greptime", "public", "INSERT INTO t VALUES (11)"), + ) + .await + .expect("post-admission retries should finish within 10 seconds") + .unwrap(); + assert!(matches!(output.output.data, OutputData::AffectedRows(2))); + assert_eq!(handler.opens.load(Ordering::SeqCst), 1); + assert_eq!(handler.writes.load(Ordering::SeqCst), 2); + assert_eq!(handler.closes.load(Ordering::SeqCst), 2); + } + } + + #[tokio::test] + async fn fails_fast_on_definitive_open_rejection() { + let handler = Arc::new(LegacyRejectHandler::default()); + let trait_handler: Arc = handler.clone(); + let weak = Arc::downgrade(&trait_handler); + let frontend = + FrontendClient::from_grpc_handler(weak, query::options::QueryOptions::default()) + .with_tracked_writes(); + let result = tokio::time::timeout( + Duration::from_secs(10), + frontend.tracked_sql_write("greptime", "public", "INSERT INTO t VALUES (7)"), + ) + .await + .expect("definitive Open rejection should finish within 10 seconds"); + let err = result.expect_err("definitive Open rejection must not be retried forever"); + assert_eq!(err.status_code(), StatusCode::InvalidArguments); + // The durable-execution layer keys repair off a TerminalError inside + // Error::External; pin that exact shape so a refactor cannot silently + // downgrade the rejection to an unresolved (permanently stalled) error. + assert!( + matches!(&err, Error::External { source, .. } if source.as_any().is::()), + "expected Error::External carrying TerminalError, got {err:?}" + ); + assert_eq!( + handler.opens.load(Ordering::SeqCst), + 1, + "Open must not be retried after a definitive rejection" + ); + } +} diff --git a/src/frontend/Cargo.toml b/src/frontend/Cargo.toml index 92803933a8b..44094e662a8 100644 --- a/src/frontend/Cargo.toml +++ b/src/frontend/Cargo.toml @@ -7,6 +7,8 @@ license.workspace = true [features] testing = [] enterprise = [ + "common-grpc/enterprise", + "session/enterprise", "common-meta/enterprise", "operator/enterprise", "sql/enterprise", diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index 59559950aef..d6fe5f49208 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -16,6 +16,8 @@ pub mod builder; mod dashboard; mod entity_graph; mod export_database; +#[cfg(feature = "enterprise")] +mod flow_write; mod grpc; mod import_packed; mod influxdb; @@ -146,6 +148,13 @@ pub struct Instance { trace_aux_limiter: Option, otlp_resource_info: bool, suspend: Arc, + #[cfg(feature = "enterprise")] + flow_write_registry: Arc< + common_grpc::flow_write::FlowWriteRegistry< + flow_write::WritePayload, + common_grpc::flow_write::Reply, + >, + >, // cache for otlp metrics // first layer key: db-string @@ -321,25 +330,35 @@ impl Instance { let query_fut = self.exec_statement_with_timeout( stmt, - query_ctx, + query_ctx.clone(), query_interceptor, timeout_recorder, ); - CancellableFuture::new(query_fut, ticket.cancellation_handle.clone()) + #[cfg(feature = "enterprise")] + let result = if query_ctx.channel() == Channel::Internal + && query_ctx.extension(common_grpc::flow_write::CHILD_KEY) == Some("1") + { + query_fut.await + } else { + CancellableFuture::new(query_fut, ticket.cancellation_handle.clone()) + .await + .map_err(|_| error::CancelledSnafu.build())? + }; + #[cfg(not(feature = "enterprise"))] + let result = CancellableFuture::new(query_fut, ticket.cancellation_handle.clone()) .await - .map_err(|_| error::CancelledSnafu.build())? - .map(|output| { - let Output { meta, data } = output; - - let data = match data { - OutputData::Stream(stream) => OutputData::Stream(Box::pin( - CancellableStreamWrapper::new(stream, ticket), - )), - other => other, - }; - Output { data, meta } - }) + .map_err(|_| error::CancelledSnafu.build())?; + result.map(|output| { + let Output { meta, data } = output; + let data = match data { + OutputData::Stream(stream) => { + OutputData::Stream(Box::pin(CancellableStreamWrapper::new(stream, ticket))) + } + other => other, + }; + Output { data, meta } + }) } else { self.exec_statement_with_timeout(stmt, query_ctx, query_interceptor, None) .await @@ -1033,20 +1052,30 @@ impl Instance { let query_fut = self.exec_plan_with_timeout(plan, query_ctx.clone(), timeout_recorder); - CancellableFuture::new(query_fut, ticket.cancellation_handle.clone()) + #[cfg(feature = "enterprise")] + let result = if query_ctx.channel() == Channel::Internal + && query_ctx.extension(common_grpc::flow_write::CHILD_KEY) == Some("1") + { + query_fut.await + } else { + CancellableFuture::new(query_fut, ticket.cancellation_handle.clone()) + .await + .map_err(|_| error::CancelledSnafu.build())? + }; + #[cfg(not(feature = "enterprise"))] + let result = CancellableFuture::new(query_fut, ticket.cancellation_handle.clone()) .await - .map_err(|_| error::CancelledSnafu.build())? - .map(|output| { - let Output { meta, data } = output; - - let data = match data { - OutputData::Stream(stream) => OutputData::Stream(Box::pin( - CancellableStreamWrapper::new(stream, ticket), - )), - other => other, - }; - Output { data, meta } - }) + .map_err(|_| error::CancelledSnafu.build())?; + result.map(|output| { + let Output { meta, data } = output; + let data = match data { + OutputData::Stream(stream) => { + OutputData::Stream(Box::pin(CancellableStreamWrapper::new(stream, ticket))) + } + other => other, + }; + Output { data, meta } + }) } else { self.exec_plan_with_timeout(plan, query_ctx.clone(), None) .await diff --git a/src/frontend/src/instance/builder.rs b/src/frontend/src/instance/builder.rs index 624ccb2f9cd..0a03163199b 100644 --- a/src/frontend/src/instance/builder.rs +++ b/src/frontend/src/instance/builder.rs @@ -440,6 +440,8 @@ impl FrontendBuilder { trace_aux_limiter, otlp_resource_info: self.options.otlp.experimental_enable_resource_info, suspend: Arc::new(AtomicBool::new(false)), + #[cfg(feature = "enterprise")] + flow_write_registry: Default::default(), }) } } diff --git a/src/frontend/src/instance/flow_write.rs b/src/frontend/src/instance/flow_write.rs new file mode 100644 index 00000000000..f908d438d4e --- /dev/null +++ b/src/frontend/src/instance/flow_write.rs @@ -0,0 +1,227 @@ +// Copyright 2023-2026 GrepTime Inc. +// +// This file is part of the GreptimeDB Enterprise Edition and is licensed under +// the GreptimeDB Enterprise License. You may not use this file except in +// compliance with that license. A copy of the license is available at the root +// of this repository in the file LICENSE-ENTERPRISE. +// +// Unless required by applicable law or agreed to in writing, this software is +// distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, +// either express or implied. + +use std::sync::Arc; + +use api::v1::{greptime_request::Request, query_request::Query}; +use common_error::ext::ErrorExt; +use common_grpc::flow_write::{ + Admission, CHILD_KEY, COMMAND_KEY, Command, Reply, TerminalError, wait_completion, +}; +use common_query::{Output, OutputData}; +use servers::error::{self as server_error, Error as ServerError}; +use servers::interceptor::GrpcQueryInterceptor; +use servers::query_handler::grpc::GrpcQueryHandler; +use session::context::QueryContextRef; + +use crate::instance::Instance; + +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct WritePayload { + request: Request, + catalog: String, + schema: String, + snapshots: Vec<(u64, u64)>, + sst_min_sequences: Vec<(u64, u64)>, + extensions: Vec<(String, String)>, +} + +impl WritePayload { + fn new(request: &Request, ctx: &QueryContextRef) -> Result { + if !is_finite_insert(request) { + return Err("tracked flow writes require INSERT or InsertIntoPlan".to_string()); + } + let mut snapshots: Vec<_> = ctx.snapshots().into_iter().collect(); + snapshots.sort_unstable(); + let mut sst_min_sequences: Vec<_> = ctx.sst_min_sequences().into_iter().collect(); + sst_min_sequences.sort_unstable(); + let mut extensions: Vec<_> = ctx + .extensions() + .into_iter() + .filter(|(key, _)| { + key != COMMAND_KEY + && key != CHILD_KEY + && key != session::context::REMOTE_QUERY_ID_EXTENSION_KEY + }) + .collect(); + extensions.sort_unstable(); + Ok(Self { + request: request.clone(), + catalog: ctx.current_catalog().to_string(), + schema: ctx.current_schema(), + snapshots, + sst_min_sequences, + extensions, + }) + } +} + +fn is_finite_insert(request: &Request) -> bool { + match request { + Request::Query(query) => match query.query.as_ref() { + Some(Query::InsertIntoPlan(_)) => true, + Some(Query::Sql(sql)) => sql.split_whitespace().next().is_some_and(|statement| { + statement + .trim_end_matches(';') + .eq_ignore_ascii_case("INSERT") + }), + _ => false, + }, + _ => false, + } +} + +impl Instance { + fn intercept_control( + &self, + request: &Request, + ctx: &QueryContextRef, + ) -> server_error::Result<()> { + let interceptor = self + .plugins + .get::>(); + interceptor + .as_ref() + .pre_execute(request, ctx.clone()) + .map_err(|error| ServerError::ExecuteQuery { + source: common_error::ext::BoxedError::new(error), + location: snafu::Location::new(file!(), line!(), column!()), + }) + } + + pub(crate) async fn handle_flow_write_request( + &self, + request: Request, + ctx: QueryContextRef, + encoded_command: String, + ) -> server_error::Result { + let command = Command::decode(&encoded_command).map_err(|e| ServerError::InvalidQuery { + reason: e.to_string(), + location: snafu::Location::new(file!(), line!(), column!()), + })?; + let reply = match command { + Command::Open { nonce } => { + self.intercept_control(&request, &ctx)?; + if !is_empty_control(&request) { + return Err(ServerError::InvalidQuery { + reason: "Open requires an empty SQL request".into(), + location: snafu::Location::new(file!(), line!(), column!()), + }); + } + Reply::Opened { + capability: self + .flow_write_registry + .open(&nonce) + .map_err(registry_error)?, + } + } + Command::Close { capability } => { + self.intercept_control(&request, &ctx)?; + if !is_empty_control(&request) { + return Err(ServerError::InvalidQuery { + reason: "Close requires an empty SQL request".into(), + location: snafu::Location::new(file!(), line!(), column!()), + }); + } + self.flow_write_registry + .close(&capability) + .await + .map_err(registry_error)?; + Reply::Closed + } + Command::Write { capability } => { + let payload = WritePayload::new(&request, &ctx).map_err(|reason| { + ServerError::InvalidQuery { + reason, + location: snafu::Location::new(file!(), line!(), column!()), + } + })?; + match self + .flow_write_registry + .admit(&capability, payload) + .map_err(registry_error)? + { + Admission::Join(receiver) => { + (*wait_completion(receiver).await.map_err(registry_error)?).clone() + } + Admission::Start(sender) => { + let receiver = sender.subscribe(); + let instance = self.clone(); + let mut child_ctx = (*ctx).clone(); + child_ctx.remove_extension(COMMAND_KEY); + child_ctx.set_extension(CHILD_KEY, "1"); + let child_ctx = Arc::new(child_ctx); + common_runtime::spawn_global(async move { + let terminal = match instance.do_query(request, child_ctx).await { + Ok(output) => terminal_reply(output), + Err(error) => Reply::Complete { + rows: 0, + cost: 0, + metrics: None, + error: Some(TerminalError { + code: error.status_code() as u32, + message: error.output_msg(), + }), + }, + }; + let _ = sender.send(Some(Arc::new(terminal))); + }); + (*wait_completion(receiver).await.map_err(registry_error)?).clone() + } + } + } + }; + let batches = common_grpc::flow_write::reply_batches(&reply).map_err(|e| { + ServerError::InvalidQuery { + reason: e.to_string(), + location: snafu::Location::new(file!(), line!(), column!()), + } + })?; + Ok(Output::new_with_record_batches(batches)) + } +} + +fn is_empty_control(request: &Request) -> bool { + matches!(request, Request::Query(query) if matches!(&query.query, Some(Query::Sql(sql)) if sql.trim().is_empty())) +} + +fn terminal_reply(output: Output) -> Reply { + let (rows, cost) = output.extract_rows_and_cost(); + let metrics = output + .meta + .plan + .clone() + .and_then(query::metrics::terminal_recordbatch_metrics_from_plan); + match output.data { + OutputData::AffectedRows(_) => Reply::Complete { + rows, + cost, + metrics, + error: None, + }, + _ => Reply::Complete { + rows: 0, + cost, + metrics, + error: Some(TerminalError { + code: 1001, + message: "tracked INSERT did not produce finite affected rows".into(), + }), + }, + } +} + +fn registry_error(error: common_grpc::flow_write::RegistryError) -> ServerError { + ServerError::InvalidQuery { + reason: format!("flow write capability rejected: {error:?}"), + location: snafu::Location::new(file!(), line!(), column!()), + } +} diff --git a/src/frontend/src/instance/grpc.rs b/src/frontend/src/instance/grpc.rs index 5cb546485bc..725c3667738 100644 --- a/src/frontend/src/instance/grpc.rs +++ b/src/frontend/src/instance/grpc.rs @@ -66,6 +66,20 @@ impl GrpcQueryHandler for Instance { request: Request, ctx: QueryContextRef, ) -> server_error::Result { + #[cfg(feature = "enterprise")] + if let Some(command) = ctx + .extension(common_grpc::flow_write::COMMAND_KEY) + .map(str::to_string) + { + if ctx.channel() != session::context::Channel::Internal { + return Err(server_error::Error::InvalidQuery { + reason: "flow write commands are only accepted on the internal channel".into(), + location: snafu::Location::new(file!(), line!(), column!()), + }); + } + return self.handle_flow_write_request(request, ctx, command).await; + } + let result: Result = async { let interceptor_ref = self.plugins.get::>(); let interceptor = interceptor_ref.as_ref(); diff --git a/src/operator/Cargo.toml b/src/operator/Cargo.toml index b89bf4b13d3..225f1bb6dec 100644 --- a/src/operator/Cargo.toml +++ b/src/operator/Cargo.toml @@ -7,6 +7,7 @@ license.workspace = true [features] testing = [] enterprise = [ + "common-grpc/enterprise", "common-meta/enterprise", "sql/enterprise", "query/enterprise", diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index 3f67d4fc576..e8f183bf7cc 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -12,6 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. +#[cfg(feature = "enterprise")] +mod flow_write; + use std::collections::BTreeMap; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, LazyLock}; @@ -839,6 +842,47 @@ impl Inserter { has_instant_rows, )?; + // Tracked Flow children use a same-capability write protocol. Keep the + // ordinary dispatch path unchanged for every other request. + #[cfg(feature = "enterprise")] + if flow_write::is_child(ctx) { + let dispatcher = self.write_dispatcher.clone(); + let request = request_factory.build_insert(normal_requests); + let writes = dispatcher.prepare(vec![request]).await?; + let tracked_tasks = writes.into_iter().map(|(peer, mut request)| { + let mut query_context = crate::utils::to_meta_query_context(ctx.clone()); + query_context.extensions.insert( + common_grpc::flow_write::CHILD_KEY.to_string(), + "1".to_string(), + ); + if let Some(header) = request.header.as_mut() { + header.query_context = Some(query_context.into()); + } + let node_manager = self.node_manager.clone(); + async move { flow_write::dispatch(&node_manager, &peer, request).await } + }); + let results = future::join_all(tracked_tasks).await; + let mut affected_rows = 0; + let mut first_error = None; + for result in results { + match result { + Ok(response) => affected_rows += response.affected_rows, + Err(error) if first_error.is_none() => first_error = Some(error), + Err(_) => {} + } + } + if let Some(error) = first_error { + return Err(error); + } + crate::metrics::DIST_INGEST_ROW_COUNT + .with_label_values(&[ctx.get_db_string().as_str()]) + .inc_by(affected_rows as u64); + return Ok(Output::new( + OutputData::AffectedRows(affected_rows), + OutputMeta::new_with_cost(write_cost as _), + )); + } + let dispatcher = self.write_dispatcher.clone(); let request = request_factory.build_insert(normal_requests); let writes = dispatcher.prepare(vec![request]).await?; diff --git a/src/operator/src/insert/flow_write.rs b/src/operator/src/insert/flow_write.rs new file mode 100644 index 00000000000..62935e46981 --- /dev/null +++ b/src/operator/src/insert/flow_write.rs @@ -0,0 +1,391 @@ +// Copyright 2023-2026 GrepTime Inc. +// +// This file is part of the GreptimeDB Enterprise Edition and is licensed under +// the GreptimeDB Enterprise License. You may not use this file except in +// compliance with that license. A copy of the license is available at the root +// of this repository in the file LICENSE-ENTERPRISE. +// +// Unless required by applicable law or agreed to in writing, this software is +// distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, +// either express or implied. + +use std::time::Duration; + +use api::v1::region::RegionRequest; +use common_error::ext::BoxedError; +use common_grpc::flow_write::{self, Command, Reply}; +use common_meta::node_manager::{DatanodeRef, NodeManagerRef}; +use common_meta::peer::Peer; +use session::context::generate_remote_query_id; + +use crate::error::{Error, Result}; + +const RETRY_DELAY: Duration = Duration::from_secs(1); + +fn error(message: impl Into) -> Error { + Error::Unexpected { + violated: message.into(), + location: snafu::Location::new(file!(), line!(), column!()), + } +} + +fn command_request( + request: &RegionRequest, + command: &Command, + control: bool, +) -> Result { + let mut request = request.clone(); + if control { + request.body = Some(api::v1::region::region_request::Body::Inserts( + Default::default(), + )); + } + let header = request.header.get_or_insert_with(Default::default); + let query_context = header + .query_context + .get_or_insert_with(api::v1::QueryContext::default); + query_context.extensions.insert( + flow_write::COMMAND_KEY.to_string(), + command.encode().map_err(|e| error(e.to_string()))?, + ); + Ok(request) +} + +fn reply(response: &api::region::RegionResponse) -> Option { + response + .extensions + .get(flow_write::COMMAND_KEY) + .and_then(|value| std::str::from_utf8(value).ok()) + .and_then(|value| serde_json::from_str(value).ok()) +} + +pub(crate) async fn dispatch( + node_manager: &NodeManagerRef, + peer: &Peer, + request: RegionRequest, +) -> Result { + dispatch_datanode(node_manager.datanode(peer).await, request).await +} + +async fn dispatch_datanode( + datanode: DatanodeRef, + request: RegionRequest, +) -> Result { + let nonce = generate_remote_query_id(); + let open = command_request(&request, &Command::Open { nonce }, true)?; + let capability = loop { + match datanode.handle(open.clone()).await { + Ok(response) => match reply(&response) { + Some(Reply::Opened { capability }) if !capability.is_empty() => break capability, + _ => return Err(error("Datanode does not support tracked Flow writes")), + }, + Err(_) => tokio::time::sleep(RETRY_DELAY).await, + } + }; + + let write = command_request( + &request, + &Command::Write { + capability: capability.clone(), + }, + false, + )?; + let complete = loop { + match datanode.handle(write.clone()).await { + Ok(response) => match reply(&response) { + Some(Reply::Complete { + rows, + cost: _, + metrics: _, + error, + }) => break (response, rows, error), + _ => tokio::time::sleep(RETRY_DELAY).await, + }, + Err(_) => tokio::time::sleep(RETRY_DELAY).await, + } + }; + + let close = command_request( + &request, + &Command::Close { + capability: capability.clone(), + }, + true, + )?; + loop { + match datanode.handle(close.clone()).await { + Ok(response) if matches!(reply(&response), Some(Reply::Closed)) => break, + _ => tokio::time::sleep(RETRY_DELAY).await, + } + } + + let (mut response, rows, terminal_error) = complete; + if let Some(terminal_error) = terminal_error { + return Err(Error::External { + source: BoxedError::new(terminal_error), + location: snafu::location!(), + }); + } + response.affected_rows = rows as _; + Ok(response) +} + +pub(crate) fn is_child(ctx: &session::context::QueryContextRef) -> bool { + ctx.channel() == common_base::protocol::Channel::Internal + && ctx.extension(flow_write::CHILD_KEY) == Some("1") +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::{Arc, Mutex}; + + use api::region::RegionResponse; + use api::v1::region::{RegionRequest, region_request::Body}; + use common_error::ext::BoxedError; + use common_error::ext::ErrorExt; + use common_grpc::flow_write::{self, Command, Reply}; + use common_meta::node_manager::{Datanode, DatanodeRef}; + use common_query::request::QueryRequest; + use common_recordbatch::SendableRecordBatchStream; + use futures_util::future::join_all; + use tokio::sync::Notify; + + use super::dispatch_datanode; + + #[derive(Default)] + struct MockDatanode { + capabilities: Mutex>, + business_write_attempts: AtomicUsize, + write_commands: Mutex>, + requests: Mutex>, + drop_open_reply: AtomicUsize, + drop_write_reply: AtomicUsize, + drop_close_reply: AtomicUsize, + latch_write: bool, + terminal_error: bool, + write_started: Notify, + release_write: Notify, + close_started: Notify, + } + + impl MockDatanode { + fn response(reply: Reply) -> RegionResponse { + let mut response = RegionResponse::new(0); + response.extensions.insert( + flow_write::COMMAND_KEY.to_string(), + serde_json::to_vec(&reply).unwrap(), + ); + response + } + + fn command(request: &RegionRequest) -> Command { + let serialized = request + .header + .as_ref() + .and_then(|header| header.query_context.as_ref()) + .and_then(|context| context.extensions.get(flow_write::COMMAND_KEY)) + .unwrap(); + Command::decode(serialized).unwrap() + } + } + + #[async_trait::async_trait] + impl Datanode for MockDatanode { + async fn handle_request( + &self, + request: Arc, + ) -> common_meta::error::Result { + self.requests.lock().unwrap().push(request.as_ref().clone()); + match Self::command(&request) { + Command::Open { nonce } => { + if !matches!(request.body, Some(Body::Inserts(ref inserts)) if inserts.requests.is_empty()) + { + return Err(common_meta::error::Error::External { + source: BoxedError::new(common_error::mock::MockError::new( + common_error::status_code::StatusCode::InvalidArguments, + )), + location: snafu::location!(), + }); + } + let mut capabilities = self.capabilities.lock().unwrap(); + let capability = capabilities + .entry(nonce) + .or_insert_with(|| "cap-1".into()) + .clone(); + if self + .drop_open_reply + .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |n| n.checked_sub(1)) + .is_ok() + { + return Err(common_meta::error::Error::External { + source: BoxedError::new(common_error::mock::MockError::new( + common_error::status_code::StatusCode::TableUnavailable, + )), + location: snafu::location!(), + }); + } + Ok(Self::response(Reply::Opened { capability })) + } + Command::Write { capability } => { + self.write_commands.lock().unwrap().push(capability.clone()); + if matches!(request.body, Some(Body::Inserts(ref inserts)) if !inserts.requests.is_empty()) + && self.business_write_attempts.fetch_add(1, Ordering::SeqCst) == 0 + && self.latch_write + { + self.write_started.notify_one(); + self.release_write.notified().await; + } + if self + .drop_write_reply + .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |n| n.checked_sub(1)) + .is_ok() + { + return Err(common_meta::error::Error::External { + source: BoxedError::new(common_error::mock::MockError::new( + common_error::status_code::StatusCode::TableUnavailable, + )), + location: snafu::location!(), + }); + } + Ok(Self::response(Reply::Complete { + rows: 7, + cost: 0, + metrics: None, + error: self.terminal_error.then(|| flow_write::TerminalError { + code: common_error::status_code::StatusCode::TableNotFound as u32, + message: "terminal child error".to_string(), + }), + })) + } + Command::Close { .. } => { + self.close_started.notify_one(); + if self + .drop_close_reply + .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |n| n.checked_sub(1)) + .is_ok() + { + return Err(common_meta::error::Error::External { + source: BoxedError::new(common_error::mock::MockError::new( + common_error::status_code::StatusCode::TableUnavailable, + )), + location: snafu::location!(), + }); + } + Ok(Self::response(Reply::Closed)) + } + } + } + + async fn handle_query( + &self, + _: QueryRequest, + ) -> common_meta::error::Result { + unreachable!() + } + } + + fn request() -> RegionRequest { + RegionRequest { + header: None, + body: Some(Body::Inserts(api::v1::region::InsertRequests { + requests: vec![Default::default()], + })), + } + } + + #[tokio::test] + async fn lost_control_and_write_replies_reuse_capability_and_exact_request() { + let datanode = Arc::new(MockDatanode::default()); + datanode.drop_open_reply.store(1, Ordering::SeqCst); + datanode.drop_write_reply.store(1, Ordering::SeqCst); + datanode.drop_close_reply.store(1, Ordering::SeqCst); + let response = dispatch_datanode(datanode.clone(), request()) + .await + .unwrap(); + assert_eq!(response.affected_rows, 7); + assert_eq!(datanode.business_write_attempts.load(Ordering::SeqCst), 2); + let writes = datanode.write_commands.lock().unwrap(); + assert_eq!(writes.len(), 2); + assert!(writes.iter().all(|capability| capability == "cap-1")); + let requests = datanode.requests.lock().unwrap(); + let writes: Vec<_> = requests + .iter() + .filter(|request| matches!(MockDatanode::command(request), Command::Write { .. })) + .collect(); + assert_eq!(writes.len(), 2); + assert_eq!(writes[0].body, writes[1].body); + assert_eq!(writes[0].header, writes[1].header); + assert_eq!( + requests + .iter() + .filter(|request| matches!( + MockDatanode::command(request), + Command::Open { .. } | Command::Close { .. } + )) + .count(), + 4 + ); + assert!(requests + .iter() + .filter(|request| matches!(MockDatanode::command(request), Command::Open { .. } | Command::Close { .. })) + .all(|request| matches!(request.body, Some(Body::Inserts(ref inserts)) if inserts.requests.is_empty()))); + } + + #[tokio::test] + async fn unsupported_open_never_submits_business_request() { + struct Unsupported; + #[async_trait::async_trait] + impl Datanode for Unsupported { + async fn handle_request( + &self, + _: Arc, + ) -> common_meta::error::Result { + Ok(RegionResponse::new(0)) + } + async fn handle_query( + &self, + _: QueryRequest, + ) -> common_meta::error::Result { + unreachable!() + } + } + let datanode: DatanodeRef = Arc::new(Unsupported); + let error = dispatch_datanode(datanode, request()).await.unwrap_err(); + assert_eq!( + error.status_code(), + common_error::status_code::StatusCode::Unexpected + ); + } + + #[tokio::test] + async fn all_peer_writes_drain_before_terminal_error_returns() { + let failing = Arc::new(MockDatanode { + terminal_error: true, + ..Default::default() + }); + let latched = Arc::new(MockDatanode { + latch_write: true, + ..Default::default() + }); + let first_datanode: DatanodeRef = failing.clone(); + let second_datanode: DatanodeRef = latched.clone(); + let aggregate = tokio::spawn(async move { + join_all([ + dispatch_datanode(first_datanode, request()), + dispatch_datanode(second_datanode, request()), + ]) + .await + }); + latched.write_started.notified().await; + failing.close_started.notified().await; + assert!(!aggregate.is_finished()); + latched.release_write.notify_one(); + let results = aggregate.await.unwrap(); + assert_eq!(results.len(), 2); + assert_eq!(failing.business_write_attempts.load(Ordering::SeqCst), 1); + assert_eq!(latched.business_write_attempts.load(Ordering::SeqCst), 1); + assert!(results.into_iter().any(|result| result.is_err())); + } +} diff --git a/src/servers/src/grpc/region_server.rs b/src/servers/src/grpc/region_server.rs index cb1cbee7891..aa2bdf3a74b 100644 --- a/src/servers/src/grpc/region_server.rs +++ b/src/servers/src/grpc/region_server.rs @@ -32,6 +32,13 @@ use crate::grpc::{TonicResult, cancellation}; #[async_trait] pub trait RegionServerHandler: Send + Sync { async fn handle(&self, request: region_request::Body) -> Result; + + async fn handle_region_request(&self, request: RegionRequest) -> Result { + let body = request.body.context(InvalidQuerySnafu { + reason: "Expecting non-empty region request body.", + })?; + self.handle(body).await + } } pub type RegionServerHandlerRef = Arc; @@ -51,15 +58,12 @@ impl RegionServerRequestHandler { let tracing_context = TracingContext::from_w3c( &request .header + .as_ref() .context(InvalidQuerySnafu { reason: "Expecting non-empty region request header.", })? .tracing_context, ); - let query = request.body.context(InvalidQuerySnafu { - reason: "Expecting non-empty region request body.", - })?; - let handler = self.handler.clone(); // Executes requests in another runtime to @@ -71,7 +75,7 @@ impl RegionServerRequestHandler { // 2. avoid the handler blocks the gRPC runtime incidentally. let handle = self.runtime.spawn(async move { handler - .handle(query) + .handle_region_request(request) .trace(tracing_context.attach(info_span!("RegionServerRequestHandler::handle"))) .await .map_err(|e| { diff --git a/src/session/Cargo.toml b/src/session/Cargo.toml index 93c46ff7ff7..630d0713a62 100644 --- a/src/session/Cargo.toml +++ b/src/session/Cargo.toml @@ -4,6 +4,9 @@ version.workspace = true edition.workspace = true license.workspace = true +[features] +enterprise = [] + [lints] workspace = true diff --git a/src/session/src/context.rs b/src/session/src/context.rs index 20462933a39..040d33f2405 100644 --- a/src/session/src/context.rs +++ b/src/session/src/context.rs @@ -447,6 +447,11 @@ impl QueryContext { self.extensions.insert(key.into(), value.into()); } + #[cfg(feature = "enterprise")] + pub fn remove_extension(&mut self, key: &str) { + self.extensions.remove(key); + } + pub fn extension>(&self, key: S) -> Option<&str> { self.extensions.get(key.as_ref()).map(|v| v.as_str()) } diff --git a/src/standalone/src/datanode_manager.rs b/src/standalone/src/datanode_manager.rs index a27a94b3c8d..14104eacc2f 100644 --- a/src/standalone/src/datanode_manager.rs +++ b/src/standalone/src/datanode_manager.rs @@ -30,7 +30,7 @@ use common_telemetry::tracing; use common_telemetry::tracing_context::{FutureExt, TracingContext}; use datanode::region_server::RegionServer; use servers::grpc::region_server::RegionServerHandler; -use snafu::{OptionExt, ResultExt}; +use snafu::ResultExt; use crate::error::{InvalidRegionRequestSnafu, InvokeRegionServerSnafu, Result}; @@ -64,12 +64,8 @@ impl RegionInvoker { } async fn handle_inner(&self, request: RegionRequest) -> Result { - let body = request.body.with_context(|| InvalidRegionRequestSnafu { - reason: "body not found", - })?; - self.region_server - .handle(body) + .handle_region_request(request) .await .context(InvokeRegionServerSnafu) } From 60097a062872f0fe5cfbc8341c4c600d824476c7 Mon Sep 17 00:00:00 2001 From: discord9 Date: Thu, 15 Oct 2026 11:17:18 +0000 Subject: [PATCH 014/239] ci(coverage): set RUST_MIN_STACK=8MB in merge_group coverage job (#9615) The merge_group coverage job did not set RUST_MIN_STACK, so test threads ran with the default 2MiB stack while the regular test job (and the production runtime in debug) use 8MiB. The deeper PromQL planning chain from #9494 overflows 2MiB under llvm-cov instrumentation, aborting tests::promql_test::anon_promql_ratio_repro::case_1_test_with_standalone with SIGABRT (12/12 failures across three merge-queue candidates). Align the coverage job with the existing 8MiB test policy. Signed-off-by: discord9 --- .github/workflows/rust.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index e29b69678c1..e9c77554e07 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -294,6 +294,7 @@ jobs: env: CARGO_BUILD_RUSTFLAGS: "-C link-arg=-fuse-ld=mold" RUST_BACKTRACE: 1 + RUST_MIN_STACK: 8388608 # 8MB, aligned with the test job CARGO_INCREMENTAL: 0 # Runs the object-store matrix against the local MinIO from # tests-integration/fixtures. Real AWS S3 is covered by the nightly From cf0f6c308ae092ccedf1f505c04d4ca1a747a2f7 Mon Sep 17 00:00:00 2001 From: fys <40801205+fengys1996@users.noreply.github.com> Date: Thu, 15 Oct 2026 12:16:14 +0000 Subject: [PATCH 015/239] fix(prom): return empty remote read results for missing metrics (#9613) * fix: return empty remote read results for missing metrics Signed-off-by: fys * test(prom): cover empty reads for missing databases Signed-off-by: fys * test(prom): simplify remote read regression cases Signed-off-by: fys * fix(prom): log missing tables during remote read Signed-off-by: fys --------- Signed-off-by: fys --- src/frontend/src/instance/prom_store.rs | 32 ++++-- src/servers/src/prom_store.rs | 26 ++++- tests-integration/src/prom_store.rs | 127 ++++++++++++++++++++++++ tests-integration/tests/http.rs | 6 +- 4 files changed, 180 insertions(+), 11 deletions(-) diff --git a/src/frontend/src/instance/prom_store.rs b/src/frontend/src/instance/prom_store.rs index d972abc9112..875cbaa873d 100644 --- a/src/frontend/src/instance/prom_store.rs +++ b/src/frontend/src/instance/prom_store.rs @@ -61,7 +61,7 @@ use tracing::instrument; use crate::error::{ AmbiguousValueColumnSnafu, CatalogSnafu, ColumnNotFoundSnafu, ExecLogicalPlanSnafu, - PromStoreRemoteQueryPlanSnafu, ReadTableSnafu, Result, TableNotFoundSnafu, + PromStoreRemoteQueryPlanSnafu, ReadTableSnafu, Result, }; use crate::instance::Instance; @@ -207,15 +207,21 @@ impl Instance { schema_name: &str, table_name: &str, query: &Query, - ) -> Result { - let table = self + ) -> Result> { + let Some(table) = self .catalog_manager .table(catalog_name, schema_name, table_name, Some(ctx)) .await .context(CatalogSnafu)? - .with_context(|| TableNotFoundSnafu { - table_name: format_full_table_name(catalog_name, schema_name, table_name), - })?; + else { + debug!( + catalog = catalog_name, + schema = schema_name, + table = table_name, + "Table not found during remote read, returning empty result" + ); + return Ok(None); + }; let timestamp_column_name = table .schema() @@ -250,12 +256,12 @@ impl Instance { .await .context(ExecLogicalPlanSnafu)?; - Ok(RemoteQueryOutput { + Ok(Some(RemoteQueryOutput { table_name: table_name.to_string(), timestamp_column_name, value_column_name, output, - }) + })) } #[tracing::instrument(skip_all)] @@ -264,7 +270,7 @@ impl Instance { ctx: QueryContextRef, queries: &[Query], query_targets: &[PermissionTableTarget], - ) -> ServerResult> { + ) -> ServerResult>> { let mut results = Vec::with_capacity(queries.len()); for (query, target) in queries.iter().zip(query_targets) { @@ -599,6 +605,10 @@ impl PromStoreProtocolHandler for Instance { let response_type = negotiate_response_type(&request.accepted_response_types)?; + for query in &request.queries { + prom_store::validate_matchers(&query.matchers)?; + } + // TODO(dennis): use read_hints to speedup query if possible let results = self .handle_remote_queries(ctx, &request.queries, &query_targets) @@ -609,6 +619,10 @@ impl PromStoreProtocolHandler for Instance { let mut query_results = Vec::with_capacity(results.len()); let mut map = HashMap::new(); for result in results { + let Some(result) = result else { + query_results.push(QueryResult { timeseries: vec![] }); + continue; + }; let RemoteQueryOutput { table_name, timestamp_column_name, diff --git a/src/servers/src/prom_store.rs b/src/servers/src/prom_store.rs index ff82f8f933b..0568d967010 100644 --- a/src/servers/src/prom_store.rs +++ b/src/servers/src/prom_store.rs @@ -21,7 +21,7 @@ use std::hash::{Hash, Hasher}; use std::sync::Arc; use api::prom_store::remote::label_matcher::Type as MatcherType; -use api::prom_store::remote::{Label, Query, Sample, TimeSeries, WriteRequest}; +use api::prom_store::remote::{Label, LabelMatcher, Query, Sample, TimeSeries, WriteRequest}; use api::v1::RowInsertRequests; use arrow::array::{ Array, ArrayRef, AsArray, DictionaryArray, LargeStringArray, StringArray, StringViewArray, @@ -138,6 +138,30 @@ pub fn extract_schema_from_query(query: &Query) -> Option { .map(|matcher| matcher.value.clone()) } +/// Validate matcher syntax independently of whether the metric exists. +pub fn validate_matchers(matchers: &[LabelMatcher]) -> Result<()> { + for matcher in matchers { + let matcher_type = MatcherType::try_from(matcher.r#type).map_err(|e| { + error::InvalidPromRemoteRequestSnafu { + msg: format!("invalid LabelMatcher type, decode error: {e}"), + } + .build() + })?; + if matches!(matcher_type, MatcherType::Re | MatcherType::Nre) { + // Validate using the same regex engine as query execution. + // Its syntax differs from Prometheus's; e.g. `(?x) a` is accepted here + // but rejected by Prometheus. + regex::Regex::new(&matcher.value).map_err(|e| { + error::InvalidPromRemoteRequestSnafu { + msg: format!("invalid regex for label '{}': {e}", matcher.name), + } + .build() + })?; + } + } + Ok(()) +} + /// Create a DataFrame from a remote Query #[tracing::instrument(skip_all)] pub fn query_to_plan( diff --git a/tests-integration/src/prom_store.rs b/tests-integration/src/prom_store.rs index af901120454..80da82c0e89 100644 --- a/tests-integration/src/prom_store.rs +++ b/tests-integration/src/prom_store.rs @@ -34,6 +34,8 @@ mod tests { use frontend::instance::Instance; use prost::Message; use servers::http::prom_store::PHYSICAL_TABLE_PARAM; + use servers::http::test_helpers::{TestClient, TestResponse}; + use servers::http::{HttpOptions, HttpServerBuilder}; use servers::prom_store; use servers::prom_store::to_grpc_row_insert_requests; use servers::query_handler::PromStoreProtocolHandler; @@ -663,10 +665,135 @@ mod tests { ] ); + instance + .do_query( + "CREATE TABLE invalid_metric (ts TIMESTAMP TIME INDEX, a DOUBLE, b DOUBLE)", + ctx.clone(), + ) + .await[0] + .as_ref() + .unwrap(); + let server = HttpServerBuilder::new(HttpOptions::default()) + .with_prom_handler(instance.clone(), None, true, Default::default(), None) + .build(); + let client = TestClient::new(server.build(server.make_app()).unwrap()).await; + + // Cover missing schemas/catalogs and each database label alias. + for (selector, database) in [ + ("__schema__", "unknown_schema"), + ("__database__", "unknown_catalog-prometheus"), + ("x_greptime_database", "unknown_catalog-information_schema"), + ] { + let mut query = remote_query("missing_metric"); + query.matchers.push(LabelMatcher { + name: selector.to_string(), + value: database.to_string(), + r#type: MatcherType::Eq as i32, + }); + let response = post_remote_read(&client, vec![query]).await; + assert_eq!(response.status(), 200, "{selector}={database}"); + let body = prom_store::snappy_decompress(&response.bytes().await).unwrap(); + let response = ReadResponse::decode(&body[..]).unwrap(); + assert_eq!(response.results.len(), 1); + assert!(response.results[0].timeseries.is_empty()); + } + + // Matcher syntax must be checked even when the metric does not exist. + for metric in ["metric1", "missing_metric"] { + for (matcher_type, value, expected_status) in [ + (99, "spark", 400), + (MatcherType::Re as i32, "[", 400), + (MatcherType::Nre as i32, "[", 400), + (MatcherType::Eq as i32, "[", 200), + (MatcherType::Neq as i32, "[", 200), + (MatcherType::Re as i32, "spark", 200), + (MatcherType::Nre as i32, "api.*", 200), + ] { + let mut query = remote_query(metric); + query.matchers.push(LabelMatcher { + name: "job".to_string(), + value: value.to_string(), + r#type: matcher_type, + }); + let response = post_remote_read(&client, vec![query]).await; + assert_eq!( + response.status().as_u16(), + expected_status, + "metric={metric}, type={matcher_type}, value={value}" + ); + } + } + + // Missing metrics must retain their result slots in multi-query responses. + for metrics in [ + vec!["missing_metric"], + vec!["missing_metric", "another_missing_metric"], + vec![ + "missing_metric", + "metric1", + "missing_metric", + "metric3", + "missing_metric", + ], + vec!["missing_metric", "invalid_metric"], + ] { + let queries = metrics.iter().map(|metric| remote_query(metric)).collect(); + let response = post_remote_read(&client, queries).await; + if metrics.contains(&"invalid_metric") { + assert!(!response.status().is_success()); + assert!(response.text().await.contains("Ambiguous value column")); + continue; + } + assert_eq!(response.status(), 200, "{}", response.text().await); + assert_eq!(response.headers()["content-type"], "application/x-protobuf"); + assert_eq!(response.headers()["content-encoding"], "snappy"); + let body = prom_store::snappy_decompress(&response.bytes().await).unwrap(); + let response = ReadResponse::decode(&body[..]).unwrap(); + assert_eq!(response.results.len(), metrics.len()); + for (result, metric) in response.results.iter().zip(metrics) { + if metric.ends_with("missing_metric") { + assert!(result.timeseries.is_empty()); + continue; + } + assert!(!result.timeseries.is_empty()); + for series in &result.timeseries { + assert!(series.labels.iter().any(|label| { + label.name == prom_store::METRIC_NAME_LABEL && label.value == metric + })); + assert_eq!(series.samples.len(), 2); + } + } + } + // check physical table if provided if let Some(physical_table) = physical_table { let sql = format!("DESC TABLE {physical_table};"); instance.do_query(&sql, ctx).await[0].as_ref().unwrap(); } } + + fn remote_query(metric: &str) -> Query { + Query { + start_timestamp_ms: 1000, + end_timestamp_ms: 2000, + matchers: vec![LabelMatcher { + name: prom_store::METRIC_NAME_LABEL.to_string(), + value: metric.to_string(), + r#type: MatcherType::Eq as i32, + }], + ..Default::default() + } + } + + async fn post_remote_read(client: &TestClient, queries: Vec) -> TestResponse { + let request = ReadRequest { + queries, + ..Default::default() + }; + client + .post("/v1/prometheus/read?db=prometheus") + .body(prom_store::snappy_compress(&request.encode_to_vec()).unwrap()) + .send() + .await + } } diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 5e006c5d80d..20d6f6dde7f 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -5463,7 +5463,11 @@ pub async fn test_prometheus_remote_schema_labels(store_type: StorageType) { .body(compressed_unknown_read_request) .send() .await; - assert_eq!(unknown_result.status(), StatusCode::BAD_REQUEST); + assert_eq!(unknown_result.status(), StatusCode::OK); + let body = prom_store::snappy_decompress(&unknown_result.bytes().await).unwrap(); + let response = ReadResponse::decode(&body[..]).unwrap(); + assert_eq!(response.results.len(), 1); + assert!(response.results[0].timeseries.is_empty()); guard.remove_all().await; } From deae75ffee02f7a89b6a8d9482cef378b8f95597 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Thu, 15 Oct 2026 12:38:01 +0000 Subject: [PATCH 016/239] fix(log-store): refuse to open object store WAL regions whose history is missing (#9561) * fix(log-store): refuse to open object store WAL regions whose history is missing A region persists only its WAL prefix, not the object store that holds it. After `[wal].storage_provider`, or the bucket or root of the provider it names, changes, or after the objects under the prefix are deleted, every region reopened against a prefix without its history, replayed nothing and silently lost its acknowledged writes that were not flushed. Add `LogStore::check_region_history` and `LogStore::start_region_history`, whose defaults accept and do nothing, so Raft Engine and Kafka are unchanged. The object store WAL reports the history of a region as missing when recovery found no object under the prefix, or when the region's watermark names an object above the newest object holding entries that recovery indexed, unless the catalog holds the region's entries through the watermark or the region's history starts in the store. Both rest on invariants of the prefix: an opened prefix is never empty, and a persisted watermark names an object at or below the newest object holding entries, which collection never deletes. The segment of release records counts as entries, since an object holding only release records can be the one collection keeps; start objects hold no segment and do not move the bound. When a store that found its prefix empty refuses a region, the actor gives up its start object, so the next start refuses the region again. It agrees only while nothing was handed an id, no batch, append, release or create waits, no failed release is kept, no other object was listed, the catalog holds only the start object and no region's history started in the store, deciding under the lock region registration takes. It then records a terminal error, deletes the object as one of its tracked deletes, which `stop` waits for, and answers every refusal once that delete settled. Mito runs the check before it replays an existing region on open, on catch-up and when a create finds the region present, with the larger of the flushed and truncated entry ids, and records the history of a region it creates; regions that skip the WAL and noop regions are not checked. A missing history fails the open with an error naming the region, the prefix, the configured storage and how to move the WAL by copying its objects. The new `[wal].accept_missing_history` option, off by default, opens such a region without its unflushed writes instead, with a warning and the counter `greptime_logstore_object_store_wal_missing_history_accepted_total`. Signed-off-by: jeremyhi * docs(rfc): describe the object store WAL history check Describe the open-time history check, the prefix invariants it relies on, how release records bound the history, the start object withdrawal of a refused open and the window that remains, the accept_missing_history option and how to move a WAL by copying its objects, and update the metrics, Mito integration, evidence, compatibility and cluster alignment sections. Signed-off-by: jeremyhi * fix(log-store): unindex the start object a refused open withdrew A withdrawal that deleted the start object left it in the catalog and in the indexed object and byte gauges for the rest of the process. Remove it from both once the delete succeeded, and keep it when the delete failed. Signed-off-by: jeremyhi * docs(rfc): separate the withdrawal windows of start objects and entries An append or a release admitted before a refusal leaves an object holding entries, which may have collected the start object, so the next start also passes watermarks at or below that object. Signed-off-by: jeremyhi * fix(log-store): keep a successor's chain complete across a late withdrawal A start object withdrawn after another instance recovered it was linked by that instance's start object, which carried reclaim boundary 0, so once the delete landed the next open found no complete chain. A start object now reclaims every object below it when the recovered chain holds no segment, so such a delete stays below its boundary. Signed-off-by: jeremyhi * fix(mito2): record a created region's WAL history before its manifest A create recorded the region's WAL history after it persisted the manifest, so a withdrawal in between failed the create but left a manifest that the next start refused as a region without its history. Record the history first, so such a create fails before it writes anything. Signed-off-by: jeremyhi * docs(rfc): describe the start boundary of a chain without segments and the create order Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- config/config.md | 2 + config/datanode.example.toml | 7 + config/standalone.example.toml | 7 + docs/rfcs/2026-09-06-object-store-wal.md | 31 +- src/common/wal/src/config.rs | 4 + src/common/wal/src/config/object_store.rs | 6 + src/log-store/src/error.rs | 33 + src/log-store/src/metrics.rs | 11 + src/log-store/src/object_store_wal/io.rs | 12 + src/log-store/src/object_store_wal/store.rs | 852 +++++++++++++++++- .../engine/object_store_wal_recovery_test.rs | 409 ++++++++- src/mito2/src/error.rs | 12 +- src/mito2/src/region/catchup.rs | 8 +- src/mito2/src/region/opener.rs | 32 +- src/mito2/src/wal.rs | 40 +- src/store-api/src/logstore.rs | 26 + 16 files changed, 1464 insertions(+), 28 deletions(-) diff --git a/config/config.md b/config/config.md index 74a94891d7d..82e98fe96ec 100644 --- a/config/config.md +++ b/config/config.md @@ -131,6 +131,7 @@ | `wal.max_unpersisted_bytes` | String | `64MB` | The size of the unpersisted backlog at which new appends stall until an upload completes, in `enqueued` mode.
Nothing is dropped and nothing is rejected; the threshold does not bound what a crash during an outage can lose.
**It's only used when the provider is `experimental_object_store` and `ack_mode` is `enqueued`**. | | `wal.max_unpersisted_age` | String | `8s` | The age of the oldest unpersisted entry at which new appends stall until an upload completes, in `enqueued` mode.
**It's only used when the provider is `experimental_object_store` and `ack_mode` is `enqueued`**. | | `wal.on_corrupted_segment` | String | `skip` | What a read does with a segment that still does not decode after a second fetch, because its checksum does not match or its content disagrees with its footer entry.
- `skip`: the segment is skipped and recorded as a WAL hole of its region, a metric is incremented and a warning is logged; the other regions of the object are unaffected (the default).
- `fail`: the read fails, so the region does not open.
**It's only used when the provider is `experimental_object_store`**. | +| `wal.accept_missing_history` | Bool | `false` | Whether a region opens although the WAL objects under the prefix do not hold its history.
An existing region fails to open when the prefix held no object when the WAL opened, or when the flushed entry id of the region names an object above every object of the prefix that holds entries. Both happen after `storage_provider`, or the bucket or root of the provider it names, changed without the WAL objects being copied, or after the objects were deleted.
To move the WAL, stop the instance, copy every object under the prefix to the new location and then change the configuration; the regions open with all their writes.
When the old WAL is lost, set this to `true` for one start: the regions open without their writes that were not flushed, which are lost, a warning names each region, and `greptime_logstore_object_store_wal_missing_history_accepted_total` counts them. Turn it off again once the instance has written to the WAL after that start; until then a later start refuses a region whose flushed entry id names an object above every object of the prefix that holds entries.
**It's only used when the provider is `experimental_object_store`**. | | `metadata_store` | -- | -- | Metadata storage options. | | `metadata_store.file_size` | String | `64MB` | The size of the metadata store log file. | | `metadata_store.purge_threshold` | String | `256MB` | The threshold of the metadata store size to trigger a purge. | @@ -614,6 +615,7 @@ | `wal.max_unpersisted_bytes` | String | `64MB` | The size of the unpersisted backlog at which new appends stall until an upload completes, in `enqueued` mode.
Nothing is dropped and nothing is rejected; the threshold does not bound what a crash during an outage can lose.
**It's only used when the provider is `experimental_object_store` and `ack_mode` is `enqueued`**. | | `wal.max_unpersisted_age` | String | `8s` | The age of the oldest unpersisted entry at which new appends stall until an upload completes, in `enqueued` mode.
**It's only used when the provider is `experimental_object_store` and `ack_mode` is `enqueued`**. | | `wal.on_corrupted_segment` | String | `skip` | What a read does with a segment that still does not decode after a second fetch, because its checksum does not match or its content disagrees with its footer entry.
- `skip`: the segment is skipped and recorded as a WAL hole of its region, a metric is incremented and a warning is logged; the other regions of the object are unaffected (the default).
- `fail`: the read fails, so the region does not open.
**It's only used when the provider is `experimental_object_store`**. | +| `wal.accept_missing_history` | Bool | `false` | Whether a region opens although the WAL objects under the prefix do not hold its history.
An existing region fails to open when the prefix held no object when the WAL opened, or when the flushed entry id of the region names an object above every object of the prefix that holds entries. Both happen after `storage_provider`, or the bucket or root of the provider it names, changed without the WAL objects being copied, or after the objects were deleted.
To move the WAL, stop the instance, copy every object under the prefix to the new location and then change the configuration; the regions open with all their writes.
When the old WAL is lost, set this to `true` for one start: the regions open without their writes that were not flushed, which are lost, a warning names each region, and `greptime_logstore_object_store_wal_missing_history_accepted_total` counts them. Turn it off again once the instance has written to the WAL after that start; until then a later start refuses a region whose flushed entry id names an object above every object of the prefix that holds entries.
**It's only used when the provider is `experimental_object_store`**. | | `query` | -- | -- | The query engine options. | | `query.parallelism` | Integer | `0` | Parallelism of the query engine.
Default to 0, which means the number of CPU cores. | | `query.memory_pool_size` | String | 0 | Memory pool size for query execution operators (aggregation, sorting, join).
Supports absolute size (e.g., "2GB", "4GB") or percentage of system memory (e.g., "20%").
Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`.
When this limit is reached, queries will fail with ResourceExhausted error.
NOTE: This does NOT limit memory used by table scans. | diff --git a/config/datanode.example.toml b/config/datanode.example.toml index 2d72f795189..d5f8e1e9be0 100644 --- a/config/datanode.example.toml +++ b/config/datanode.example.toml @@ -258,6 +258,13 @@ overwrite_entry_start_id = false ## **It's only used when the provider is `experimental_object_store`**. #+ on_corrupted_segment = "skip" +## Whether a region opens although the WAL objects under the prefix do not hold its history. +## An existing region fails to open when the prefix held no object when the WAL opened, or when the flushed entry id of the region names an object above every object of the prefix that holds entries. Both happen after `storage_provider`, or the bucket or root of the provider it names, changed without the WAL objects being copied, or after the objects were deleted. +## To move the WAL, stop the instance, copy every object under the prefix to the new location and then change the configuration; the regions open with all their writes. +## When the old WAL is lost, set this to `true` for one start: the regions open without their writes that were not flushed, which are lost, a warning names each region, and `greptime_logstore_object_store_wal_missing_history_accepted_total` counts them. Turn it off again once the instance has written to the WAL after that start; until then a later start refuses a region whose flushed entry id names an object above every object of the prefix that holds entries. +## **It's only used when the provider is `experimental_object_store`**. +#+ accept_missing_history = false + # The Kafka SASL configuration. # **It's only used when the provider is `kafka`**. # Available SASL mechanisms: diff --git a/config/standalone.example.toml b/config/standalone.example.toml index aa5d09c730c..6fba1d91398 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -449,6 +449,13 @@ overwrite_entry_start_id = false ## **It's only used when the provider is `experimental_object_store`**. #+ on_corrupted_segment = "skip" +## Whether a region opens although the WAL objects under the prefix do not hold its history. +## An existing region fails to open when the prefix held no object when the WAL opened, or when the flushed entry id of the region names an object above every object of the prefix that holds entries. Both happen after `storage_provider`, or the bucket or root of the provider it names, changed without the WAL objects being copied, or after the objects were deleted. +## To move the WAL, stop the instance, copy every object under the prefix to the new location and then change the configuration; the regions open with all their writes. +## When the old WAL is lost, set this to `true` for one start: the regions open without their writes that were not flushed, which are lost, a warning names each region, and `greptime_logstore_object_store_wal_missing_history_accepted_total` counts them. Turn it off again once the instance has written to the WAL after that start; until then a later start refuses a region whose flushed entry id names an object above every object of the prefix that holds entries. +## **It's only used when the provider is `experimental_object_store`**. +#+ accept_missing_history = false + # The Kafka SASL configuration. # **It's only used when the provider is `kafka`**. # Available SASL mechanisms: diff --git a/docs/rfcs/2026-09-06-object-store-wal.md b/docs/rfcs/2026-09-06-object-store-wal.md index 1d784dd09ae..d3e38bb5a26 100644 --- a/docs/rfcs/2026-09-06-object-store-wal.md +++ b/docs/rfcs/2026-09-06-object-store-wal.md @@ -50,7 +50,7 @@ The provider is added to the four parallel WAL enums so that every resolution pa Two choices are deliberate: -- The persisted `WalOptions` carry the derived node prefix `/datanodes//epochs/` described under *Object format*, not the configured root. On reopen the log store compares the persisted prefix with the one it runs under and fails fast on a mismatch, so an operator who changes the prefix or the node id cannot silently point a region at an empty WAL. Changing the prefix of an existing region requires an explicit migration, which is out of scope here. +- The persisted `WalOptions` carry the derived node prefix `/datanodes//epochs/` described under *Object format*, not the configured root. On reopen the log store compares the persisted prefix with the one it runs under and fails fast on a mismatch, so an operator who changes the prefix or the node id cannot silently point a region at an empty WAL. Changing the prefix of an existing region requires an explicit migration, which is out of scope here. The prefix does not name the object store that holds it, so changing `storage_provider`, or the bucket or root of the provider it names, keeps the prefix; the history check under *Log store* catches such a change when a region opens. - `Provider::ObjectStore` is classified as a remote WAL. Mito's remote-WAL branches (initial flushed entry id from `latest_entry_id`, `topic_latest_entry_id` after replay, region-scoped reads, no topic grouping on batch open) apply unchanged. This is a classification of replay semantics, not of physical location. A datanode that is configured with a metasrv client rejects the provider before any store is built: `DatanodeBuilder` returns an error stating that the object store WAL is standalone-only. Separately, `DatanodeWalConfig::ObjectStore` cannot be converted into `MetasrvWalConfig`, so a metasrv cannot be configured with it either. The standalone bootstrap matches the datanode WAL configuration before that conversion and constructs `WalProvider::ObjectStore` directly. @@ -102,11 +102,11 @@ Positions start at one, so entry id zero is never assigned: Mito treats zero as `latest_entry_id(provider)` is region-scoped: the highest entry id durable under the prefix for the provider's region, or zero when the region has none. A freshly created region therefore starts at zero even if other regions under the prefix have entries. This differs from Kafka, whose high watermark is per topic; it is sufficient here because ids are per region, and Mito consumes the value only for the region that asked. -Two rules keep every new id of a region above every id it already has, on a prefix that already holds objects and across the watermark Mito hands the store. On open, the next object sequence is the larger of the largest listed sequence plus one, orphans included, and, over every region, the object named by the region's largest existing entry id plus one, that is `(max_entry_id >> 20) + 1`, so that the catalog's strictly increasing ranges hold; sequence gaps are tolerated, and a sequence at or above `2^44` is rejected as exhausted. Objects written under an earlier contiguous scheme are not rewritten: they remain readable, their ids carry no object information, and the rule places them by their high bits just the same. +Two rules keep every new id of a region above every id it already has, on a prefix that already holds objects and across the watermark Mito hands the store. On open, the next object sequence is the larger of the largest listed sequence plus one, orphans included, and, over every region, the object named by the region's largest existing entry id plus one, that is `(max_entry_id >> 20) + 1`, so that the catalog's strictly increasing ranges hold; sequence gaps are tolerated, and a sequence at or above `2^44` is rejected as exhausted. Ids have been object-sequence-major since the first version of the provider; a contiguous scheme existed only in the reference implementation before release, so no deployed prefix holds ids without object information, and the history check under *Log store* relies on that. -The second rule is the sequence floor of `obsolete`. Mito calls `obsolete` with the manifest's `flushed_entry_id` when it opens a region, and the actor applies the watermark and a floor derived from that entry id together: the next object sequence is raised to `(entry_id >> 20) + 1` unless it is there already, while zero, the watermark of a region without entries, names no object and needs no floor, so that a watermark assigned under another prefix, or one whose object this prefix no longer holds, is never passed by a new id. The floor moves only the next sequence, never the chain: the next batch still links to its predecessor, so a skipped sequence leaves no gap in the history recovery replays. While the open batch holds ids under the next sequence, the call is refused with a retryable error and records neither the watermark nor the floor, because those ids would fall below the floor. In the `enqueued` mode an id the store itself handed out needs no floor, even before it is durable: it is never handed out again while the store runs, so every later id of the region is greater. In the current single-prefix layout the floor never needs to move, since every watermark Mito hands the store names an object below the next sequence; the rule exists for the cross-prefix cases under *Alignment with cluster mode*. No additional persisted state is introduced. +The second rule is the sequence floor of `obsolete`. Mito calls `obsolete` with the manifest's `flushed_entry_id` when it opens a region, and the actor applies the watermark and a floor derived from that entry id together: the next object sequence is raised to `(entry_id >> 20) + 1` unless it is there already, while zero, the watermark of a region without entries, names no object and needs no floor, so that a watermark assigned under another prefix, or one whose object this prefix no longer holds, is never passed by a new id. The floor moves only the next sequence, never the chain: the next batch still links to its predecessor, so a skipped sequence leaves no gap in the history recovery replays. While the open batch holds ids under the next sequence, the call is refused with a retryable error and records neither the watermark nor the floor, because those ids would fall below the floor. In the `enqueued` mode an id the store itself handed out needs no floor, even before it is durable: it is never handed out again while the store runs, so every later id of the region is greater. In the current single-prefix layout the floor never needs to move for a prefix that holds the region's history, since every watermark Mito hands the store then names an object below the next sequence. A watermark above the newest object holding entries that the store found means the prefix does not hold that history: the history check under *Log store* refuses such an open before Mito calls `obsolete`, and the floor moves only when `accept_missing_history` accepted the loss. Otherwise the rule exists for the cross-prefix cases under *Alignment with cluster mode*. No additional persisted state is introduced. -Garbage collection deletes the objects below a reclaim boundary that the store persists in every header, so that recovery can tell a collected predecessor from one that never landed. The boundary is computed from the catalog and the watermarks the store already keeps: each region with a recorded watermark contributes the sequence of its first indexed object holding a segment whose largest entry id is above the watermark; a region without a recorded watermark contributes its first indexed object; a region whose watermark is the maximum entry id, after a release in full, contributes nothing; the boundary is the smallest contribution, or the last indexed sequence plus one when nothing contributes, and it never decreases. Whether a segment is flushed is decided per segment from the footer, on the full entry id, so a watermark that stops in the middle of an object keeps the object. The boundary is recomputed after every successful `obsolete`, after every release and after every object is indexed; Mito calls `obsolete` after a flush, after a truncate, when it drops a region and when it opens one, so a restart re-establishes the watermark of every region Mito opens again, and the boundary advances on the same rule. A region that is never opened again needs the release described below. Every sealed batch carries the boundary current at sealing and the start object carries the boundary recovery found, so a boundary is always at or below the sequence of the object carrying it. The boundary that permits deletion is the one carried by the last indexed object; the boundary of a batch whose create failed or is in flight permits nothing. +Garbage collection deletes the objects below a reclaim boundary that the store persists in every header, so that recovery can tell a collected predecessor from one that never landed. The boundary is computed from the catalog and the watermarks the store already keeps: each region with a recorded watermark contributes the sequence of its first indexed object holding a segment whose largest entry id is above the watermark; a region without a recorded watermark contributes its first indexed object; a region whose watermark is the maximum entry id, after a release in full, contributes nothing; the boundary is the smallest contribution, or the last indexed sequence plus one when nothing contributes, and it never decreases. Whether a segment is flushed is decided per segment from the footer, on the full entry id, so a watermark that stops in the middle of an object keeps the object. The boundary is recomputed after every successful `obsolete`, after every release and after every object is indexed; Mito calls `obsolete` after a flush, after a truncate, when it drops a region and when it opens one, so a restart re-establishes the watermark of every region Mito opens again, and the boundary advances on the same rule. A region that is never opened again needs the release described below. Every sealed batch carries the boundary current at sealing and the start object carries the boundary recovery found, or its own sequence when the recovered chain holds no segment, so a boundary is always at or below the sequence of the object carrying it. The boundary that permits deletion is the one carried by the last indexed object; the boundary of a batch whose create failed or is in flight permits nothing. Every present object with a sequence below the durable boundary is deletable, on the chain or not: the objects the catalog indexes, the objects recovery listed but did not index, the objects of an earlier epoch that the start object moved past, and the sequences of batches that failed in this run, whose objects may or may not exist. A collection starts after a successful `obsolete`, after a release and after an indexing that raised the durable boundary. Deletes run in the background, at most four at a time and in no particular order, and never hold up admission, sealing, uploads or acknowledgements; an indexed object is unindexed only after its delete succeeded, and the WAL holes recorded for it are dropped with it. Every deleted object is counted in `greptime_logstore_object_store_wal_deleted_objects_total`; a failed delete is logged, counted in `greptime_logstore_object_store_wal_failed_deletes_total` and retried at the next collection, and the object stays indexed meanwhile. An object deleted but not yet unindexed when the process crashes is simply absent at the next recovery, which lists the prefix. Nothing is collected once `stop` has begun or while the store is poisoned; a delete in flight when `stop` begins runs to completion before `stop` returns, and the actor settles every attempt it abandons when it exits. @@ -160,13 +160,21 @@ A release is a waiter of the batch that carries its record in both modes: it ret **Reading.** `read(provider, entry_id)` returns the region's entries with id at least `entry_id`, in order, fetching only that region's segments with range reads and decoding nothing else, and hides entries at or below the region's obsolete watermark. `obsolete` records the watermark in memory, clamped to the durable entry id in the `enqueued` mode; the watermark is re-established after a restart because Mito calls `obsolete` with the manifest's `flushed_entry_id` when it opens a region, so the store persists no watermark of a region Mito opens again. The release of a region, which `obsolete_all` and `release_region` make, is persisted as described under *Entry ids*. `obsolete_all` hides the whole region. `latest_entry_id` is not affected by the watermark. Every method taking a provider requires `Provider::ObjectStore` with the store's prefix and the caller's region. -**Metrics.** The store exports metrics under `greptime_logstore_object_store_wal_`. The write path reports `seal_to_durable_seconds`, the time from sealing a batch until its object is durable, `append_ack_seconds`, the time an append waits for its acknowledgement in the `durable` mode, `object_bytes` and `object_entries`, the size and the entry count of every object sealed, and `created_objects_total`, the creates that succeeded, an identical retry included; a create is counted and timed whether or not its batch is still there to be indexed, since an object a poisoned store gave up on is as durable as any other. `indexed_objects` and `indexed_bytes` are the objects and the bytes the catalog holds, set by a recovery that succeeded and moved by every object indexed or unindexed afterwards, so a collection takes the object it deleted off both, and a recovery that is abandoned publishes neither. A recovery that succeeds reports `recovery_seconds` and the objects it read footers from in `recovered_objects_total`, both of the same recoveries, so one that is abandoned at a fetch or rejected by the catalog contributes to neither; a read stream reports `read_seconds` once it has been read to its end, and one that fails or is dropped part way reports nothing. The failure paths are counted: `create_conflicts_total` for a create that found a conflicting object, `create_failures_total` for one the object store did not confirm, `poisoned_total` for a store that recorded a terminal error, `skipped_segments_total` for a segment a read skipped, and `stalled_appends_total` with `stalled_append_seconds` for an append held back because the sealed batches were full, in either mode, or by the `enqueued` backlog thresholds, timed until it is admitted or refused, a stop, a terminal error, a failing store or the teardown of the actor releasing it included. The gauge `failing` counts the stores whose backlog the object store refuses with an error that requires intervention, and `rejected_appends_total` the appends and releases such a store rejected; the attempts of a failing store are counted in `create_failures_total`. The delete counters of a collection and the release record counter are described under *Entry ids*. The histogram buckets are chosen for the default flush interval of 100 ms and the default `max_batch_bytes` of 8 MiB. +**History check.** A region persists its prefix, not the object store that holds the prefix, so pointing `storage_provider`, or the bucket or root of the provider it names, at another location, or deleting the objects under the prefix, would let every region open against a prefix without its history, replay nothing and lose its acknowledged writes that were not flushed. Two invariants of a prefix expose such a location without any persisted state of their own. A prefix that a store opened is never empty afterwards, unless the store withdrew its start object as described below: the store writes its start object before it admits an append, and collection deletes only objects below the durable boundary, the boundary of the last indexed object, which is at or below that object's own sequence, since a larger one fails the header check. And a watermark a region persisted names an object at or below the newest object holding entries that any later recovery of the prefix indexes: the entry it names was durable and indexed when the watermark was persisted, as described under *Acknowledgement*; no durable boundary passes the newest object holding entries, since a batch carries at most the sequence after the last indexed object and a start object carries the boundary recovery found, so collection keeps that object and recovery indexes it on the chain of the tip. The segment of release records counts as entries: it never holds the reclaim boundary, so an object holding only release records can be the one that collection keeps while it deletes every older object. Objects without segments, start objects and the like, do not move it. `LogStore::check_region_history(provider, entry_id)` therefore reports the history of the provider's region as missing when recovery found no object under the prefix, or when `entry_id` names an object above the newest object holding entries that recovery indexed, unless the catalog holds the region's entries through `entry_id`, as it does for a region written since the store opened, or the region's history starts in the store: Mito calls `LogStore::start_region_history` when it creates a region, before it persists the region's manifest, and an open that accepted a missing history records the region the same way, in memory, so the region reopens on the same store without any entry and is counted once. Both methods default to accepting and to doing nothing, so Raft Engine and Kafka are unchanged. Mito calls it before it replays an existing region from its persisted state, on open, on catch-up and when a create finds the region already present, with the larger of the manifest's `flushed_entry_id` and `truncated_entry_id`; a region whose options skip the WAL keeps no history there and is not checked, nor is a region with the noop provider. A missing history fails the open, and such a create, which does not fall back to creating the region over its manifest, with an `InvalidArguments` error that names the region, the prefix, the configured storage provider with its scheme, bucket and root, and the two ways forward below; in standalone a region that fails to open stops startup. The check runs before Mito calls `obsolete` with the watermark, so the sequence floor never hides a missing history. + +A refused open must not make the next start pass. When the store found its prefix empty and refuses a region, it gives up the start object it wrote: the actor agrees only while nothing was ever handed an id, no batch, append, release or create waits, no failed release is kept for a later object, recovery listed no other object, the catalog holds only the start object and no region's history started in the store, and from then on the store records a terminal error and accepts nothing, since a later object would link to the given-up one. A release seals its own batch without handing an id, so once its object is indexed the catalog shows it, alongside the start object or alone after collection deleted the start object; the actor deletes the object like a collected one, so `stop` waits for that delete and a cancelled caller does not abandon it, and every refusal returns its error only once the delete settled. A region registers its history and the actor decides the withdrawal under one lock, so a region created at that moment either prevents the withdrawal or fails to register, and then fails its create before it leaves a manifest that a later start would refuse. Another instance may open the prefix while the delete is in flight, as when two processes of the same node overlap: its recovery finds a chain without any segment, so its start object reclaims every object below it, and the delete, which may land after that start object, leaves its chain complete. The terminal error is not counted in `poisoned_total`. In standalone, regions open before the instance serves, so at that point nothing else can have written to the prefix, and the next start, an operator retry or an orchestrator restart, finds the prefix empty again and refuses the same regions. A refused start on a prefix that held objects writes a start object, which holds no entry and so does not move the newest object holding entries: later starts refuse the same watermarks. + +To move the WAL to another location, stop the instance, copy every object under the prefix to the same keys in the new location, and then change the configuration: recovery finds the same objects and every region opens with all of its acknowledged writes. When the old objects are lost, `accept_missing_history = true` turns the failure into an accepted loss for that start: the region opens without its writes that were not flushed, a warning names the region, the prefix, the location and the watermark above which nothing was replayed, and `greptime_logstore_object_store_wal_missing_history_accepted_total` counts the region. The option is meant to be turned off again once the instance has written to the WAL after that start: until a write creates an object holding entries at or above the sequence the accepted watermark names, which the floor raised by the accepted open makes the next write do, a later start reports such a region again. + +The check does not catch every change of location. A location whose prefix already holds objects passes a region that never flushed, and a region whose watermark names an object at or below the newest object holding entries there, whatever history those objects hold. A refused start on an empty prefix leaves its start object behind when the process ends before the delete completes or when the delete fails, which is logged; the next start then finds a prefix with only start objects and passes a region that never flushed, while it still refuses every region whose watermark names an object. The same holds when the store created a region before the refusal. When it admitted an append or a release before the refusal, as it can when regions open while the instance already serves, the store withdraws nothing and the prefix keeps an object holding entries, which may already have collected the start object; the next start then also passes every region whose watermark names an object at or below that one. + +**Metrics.** The store exports metrics under `greptime_logstore_object_store_wal_`. The write path reports `seal_to_durable_seconds`, the time from sealing a batch until its object is durable, `append_ack_seconds`, the time an append waits for its acknowledgement in the `durable` mode, `object_bytes` and `object_entries`, the size and the entry count of every object sealed, and `created_objects_total`, the creates that succeeded, an identical retry included; a create is counted and timed whether or not its batch is still there to be indexed, since an object a poisoned store gave up on is as durable as any other. `indexed_objects` and `indexed_bytes` are the objects and the bytes the catalog holds, set by a recovery that succeeded and moved by every object indexed or unindexed afterwards, so a collection takes the object it deleted off both, and a recovery that is abandoned publishes neither. A recovery that succeeds reports `recovery_seconds` and the objects it read footers from in `recovered_objects_total`, both of the same recoveries, so one that is abandoned at a fetch or rejected by the catalog contributes to neither; a read stream reports `read_seconds` once it has been read to its end, and one that fails or is dropped part way reports nothing. The failure paths are counted: `create_conflicts_total` for a create that found a conflicting object, `create_failures_total` for one the object store did not confirm, `poisoned_total` for a store that recorded a terminal error, `skipped_segments_total` for a segment a read skipped, `missing_history_accepted_total` for a region that `accept_missing_history` let open without its history, and `stalled_appends_total` with `stalled_append_seconds` for an append held back because the sealed batches were full, in either mode, or by the `enqueued` backlog thresholds, timed until it is admitted or refused, a stop, a terminal error, a failing store or the teardown of the actor releasing it included. The gauge `failing` counts the stores whose backlog the object store refuses with an error that requires intervention, and `rejected_appends_total` the appends and releases such a store rejected; the attempts of a failing store are counted in `create_failures_total`. The delete counters of a collection and the release record counter are described under *Entry ids*. The histogram buckets are chosen for the default flush interval of 100 ms and the default `max_batch_bytes` of 8 MiB. A `testing` cargo feature exposes hooks to wait for admitted appends, to seal the open batch, to hold and release creates and to make creates fail, so engine tests can flush deterministically and observe entries that are admitted but not durable. ## Mito integration -Mito maps `WalOptions::ObjectStore { prefix }` to `Provider::object_store_provider(region_id, prefix)`; resolving object store options against a Raft Engine or Kafka log store reports the existing incompatible-provider error. Reads go through the region-scoped reader used for Kafka, so a namespace shared by many regions never leaks entries across regions even if a log store returned them. Object store regions are opened individually rather than grouped by topic. After replay, `topic_latest_entry_id` follows the existing remote-WAL rule: it is taken from the store's latest entry id only when the region's memtables are empty, otherwise it stays at the replay start, because that value feeds pruning and must never advance past unflushed data. Before a flush, a full truncate or a discard of unflushed data publishes a manifest watermark, Mito waits on the durability barrier described under *Acknowledgement*; with the `durable` mode and with every other provider the wait returns at once. When Mito drops a region it calls `release_region` after `obsolete`, as described under *Entry ids*; the default of that method does nothing, so the other providers handle a dropped region as before. Neither a release nor an `obsolete` of an object store region that fails, the latter for instance on a store that the release of another region poisoned, fails the drop, since nothing sends the drop of a removed region again and a caller that drops several regions, such as the metric engine with its data and metadata regions, would leave the others undone: the region's runtime state is cleaned up and its files are scheduled for deletion as for any drop, and the failure is logged; an `obsolete` that fails for another provider fails the drop as before. A store that is neither stopped nor poisoned keeps the record and writes it with its next object; a stopped or poisoned store keeps none, and the region stays unreleased after a restart. The store carries the record in a later object, as described under *Entry ids*; a crash before any later create succeeds loses it, and the dropped region's segments then hold the reclaim boundary after the restart. +Mito maps `WalOptions::ObjectStore { prefix }` to `Provider::object_store_provider(region_id, prefix)`; resolving object store options against a Raft Engine or Kafka log store reports the existing incompatible-provider error. Reads go through the region-scoped reader used for Kafka, so a namespace shared by many regions never leaks entries across regions even if a log store returned them. Object store regions are opened individually rather than grouped by topic. After replay, `topic_latest_entry_id` follows the existing remote-WAL rule: it is taken from the store's latest entry id only when the region's memtables are empty, otherwise it stays at the replay start, because that value feeds pruning and must never advance past unflushed data. Before a flush, a full truncate or a discard of unflushed data publishes a manifest watermark, Mito waits on the durability barrier described under *Acknowledgement*; with the `durable` mode and with every other provider the wait returns at once. When Mito drops a region it calls `release_region` after `obsolete`, as described under *Entry ids*; the default of that method does nothing, so the other providers handle a dropped region as before. Neither a release nor an `obsolete` of an object store region that fails, the latter for instance on a store that the release of another region poisoned, fails the drop, since nothing sends the drop of a removed region again and a caller that drops several regions, such as the metric engine with its data and metadata regions, would leave the others undone: the region's runtime state is cleaned up and its files are scheduled for deletion as for any drop, and the failure is logged; an `obsolete` that fails for another provider fails the drop as before. A store that is neither stopped nor poisoned keeps the record and writes it with its next object; a stopped or poisoned store keeps none, and the region stays unreleased after a restart. The store carries the record in a later object, as described under *Entry ids*; a crash before any later create succeeds loses it, and the dropped region's segments then hold the reclaim boundary after the restart. Before Mito replays an existing region on open, on catch-up or for a create that finds the region present, it runs the history check described under *Log store*. ## Datanode wiring and lifecycle @@ -211,6 +219,10 @@ max_unpersisted_age = "8s" # regions of the object are unaffected (the default). "fail": the read # fails, so the region does not open. on_corrupted_segment = "skip" +# Whether a region opens although the prefix does not hold its WAL +# history, losing its writes that were not flushed; see the history +# check under *Log store*. Meant for one start after the old WAL is lost. +accept_missing_history = false ``` `config/config.md` is generated from the example files. @@ -222,8 +234,9 @@ on_corrupted_segment = "skip" - The configuration tag carries `experimental_`; the persisted tag does not, so promoting the provider to stable is a configuration rename with an alias, not a metadata migration. - `LogStore::wait_durable` has a default implementation that returns at once, so Raft Engine, Kafka and the noop provider keep their own acknowledgement and persistence semantics; the durability barrier binds only the object store provider. - `LogStore::release_region` has a default implementation that does nothing, so Raft Engine, Kafka and the noop provider keep their handling of a dropped region. -- The object-sequence-major id scheme needs no migration of existing objects. A prefix written under the contiguous scheme keeps its objects, and the sequence allocation rule under *Entry ids* keeps new ids above the old ones. Garbage collection reasons from footers and from the reclaim boundary, not from the id encoding, so it applies to objects of either scheme. +- The object-sequence-major id scheme is the only one the provider has written, so no prefix needs a migration; the contiguous scheme of the pre-release reference implementation has no deployed prefixes, and the history check relies on every id naming its object. - `on_corrupted_segment` is a new configuration option with a default, so existing configurations decode unchanged; it changes what a read does with a damaged segment, not what is written. +- `accept_missing_history` is a new configuration option with a default of `false`, so existing configurations decode unchanged. `LogStore::check_region_history` has a default implementation that accepts and `LogStore::start_region_history` one that does nothing, so only the object store provider checks a region's history, and the check reads only the recovered catalog: no object format or region metadata changes. A region that previously opened against a prefix without its history, replaying nothing, now fails to open until its objects are restored or the loss is accepted. ## Evidence @@ -240,7 +253,7 @@ The reference implementation was validated on the following evidence; the tests | A manifest watermark never names an entry that is not durable | Durability barrier tests: a flush held until its entry is durable; a crash in `enqueued` mode after an acknowledgement and a flush attempt but before the object exists, followed by a restart, new writes and a second crash; a drop cancelling a flush that waits; a truncate that waits; a flush after a backlog was lost. | Every durable entry replays; the flush after the lost backlog publishes no frontier; `obsolete` never passes the durable id. | | A region replays only its own entries | The two-region recovery test above. | Per-region row-level scan equality with the other region unflushed. | | Recovery never silently loses data | Unit coverage of every corruption class at recovery and on read; footer-only recovery compared with a full decode; a footer longer than the tail window. | Structural damage fails construction of the store; a corrupted segment is skipped, recorded as a hole of its region and counted, or fails the read under `fail`; the footer-only catalog equals the full decode. | -| A region cannot be silently pointed at another WAL | Reopen on a mismatched prefix; reopen on an empty prefix. | The mismatch is rejected with both prefixes named and the region is not opened; the empty prefix opens cleanly with the expected row counts. | +| A region cannot be silently pointed at another WAL | Reopen on a mismatched prefix; open and catch-up against a location whose prefix is empty, and open against one whose prefix holds objects only below a flushed or a truncated watermark; a copy of every object on a new location; a restart after collection left only the newest object; the option accepting the loss; store tests of the same cases. | The mismatch is rejected with both prefixes named; a missing history fails the open and the catch-up with the region, the prefix and the location named, while a region that skips the WAL, a region whose history the location holds and a new region still open; the copy and the restart open with every row; the option opens the region without its unflushed rows and counts it once. | | Garbage collection deletes everything below the persisted reclaim boundary, and recovery stays correct while objects disappear | Unit tests of the boundary computation (a region without a watermark holds it at its first object, a watermark inside an object keeps the object, `obsolete_all` releases it, last indexed plus one when everything is flushed), of stamping and durability (headers carry the boundary current at sealing, a failed create's boundary permits nothing, the start object carries the recovered boundary), of deletion (chain objects, listed orphans and failed batches below the durable boundary, nothing at or above it, a failed delete retried, `stop` waiting for deletes in flight, holes dropped), of recovery (two restarts with a collection between them, a late object of an earlier epoch below the boundary collected, a missing predecessor above the boundary still failing) and of a read that listed an object before it was collected; the Mito flush-then-collect test in both acknowledgement modes; the standalone integration test's assertion after a flush. | Objects below the durable boundary are gone after a flush, the start object included; a restart replays everything above the boundary; exactly one WAL object remains after a flush and one more insert in the integration test. | | A dropped region holds no object after a restart that never opens it | Unit tests of the release record an object carries and the boundary it allows; of recovery seeding the watermark from the release record before any boundary is computed, with collection proceeding once the remaining region is flushed; of a failed release carried by a timer tick on an idle prefix and by a `durable` stop, which reports the record lost when its create fails too; of the release of a region without segments, and of a region id used again after its release; of a release record that outlives a restart together with the segment it covers and goes once both are below the boundary, deleted out of order; the fault-matrix cases of a release, which return only once its object is durable and indexed, fail with its create, and fail at once while the store is failing; the Mito test that drops one of two regions, restarts gracefully, abruptly, or after a release that failed, with nothing written in between, then flushes the other region and writes once more; the metric engine test of a physical region whose data and metadata regions wrote to the WAL, dropped while the release of the data region fails, while both releases fail, and while the release of the data region meets a conflicting object that poisons the store, which still drops both regions and forgets the physical region, and, in the first two cases, after which a store opened later keeps no segment of either; the frozen object with release records. | Recovery seeds the release, and after the flush and the next write every object written before the restart is deleted. | | Every row of the *Failure matrix* holds, and every metric moves as documented | A table-driven fault matrix, `test_store_fault_matrix`, that drives every failure the object store access can produce through the store and asserts, per case, which waiters fail and with what error, whether the store is poisoned, whether the sequence is consumed, the indexed gauges at the failure itself, the exact movement of every metric, and what a read of the store sees afterwards; every timer is observed twice around two forced waits whose intervals cannot overlap. | 42 cases: 24 for the eleven rows of the matrix, five of them for the create of a release record, and 18 for the failures listed under *Failures driven through the store*. The test fails if a row of the matrix is left uncovered. | @@ -345,7 +358,7 @@ Distributed mode is out of scope, but several decisions were taken so that the c Sealing may only exclude objects whose writes were never acknowledged, which takes two fences. Admission is fenced by the lease mechanism Mito already has: when a region's lease expires the datanode demotes it to a follower, which refuses writes, and the metasrv waits for that deadline before it upgrades the candidate. Completion is fenced by the store: every conditional create runs under an upload deadline, a create that has not completed by its deadline is reported failed to its waiters and is never acknowledged afterwards even if the object turns out to exist, and when a region's lease expires the store fails every waiter of that region that is not yet acknowledged, so that a create still in flight at that moment cannot be acknowledged later. The metasrv seals the segment only after the lease deadline plus the upload deadline have both passed, so an object that appears after sealing belongs to a create the old writer had already reported as failed. Under the durable acknowledgement mode this means a sealed set never excludes an acknowledged write; under the `enqueued` mode the entries in the old writer's unpersisted backlog at that moment are lost, which is that mode's stated recovery point objective. The late-completing create is part of the acceptance criteria for takeover, stated in terms of the sealing instant: an object that was absent at sealing, whether it fills a hole later or lies above the cutover, is excluded on every open of the region, and so is a visible object outside the selected chain; an object on the sealed chain replays on every open even if its write was never acknowledged, which is the same outcome as a crash between an object's creation and its acknowledgement; and in the durable mode every caller whose create had not completed when the lease expired must have been failed, so no acknowledged write is among the excluded objects. The generation in the key is an additional guarantee that such a late object lands in a directory the new writer never reads beyond the sealed set. Object-sequence-major entry ids are what make this cheap: `entry_id >> 20` locates the object under the right prefix without the source node's in-memory catalog. -**Entry id floor across prefixes.** A region that arrives under a new prefix, by migration or by failover, brings a watermark and replayed ids that were assigned under another prefix's sequence. Before the target assigns the region its first new id, the target store raises its next object sequence above `max_id >> 20`, where `max_id` is the largest of the region's inherited watermark and every id it replayed, so that every new id of that region is greater than every id it already has; sequence gaps are tolerated and the sequence space is wide enough for this. The rule is applied on every open from the region's persisted state, in the same place Mito already hands the store the manifest watermark, so it needs no persisted state of its own and survives a crash before the target's first flush: the restart re-derives the floor and the region's acknowledged-but-unflushed entries under the target prefix are above the inherited watermark and replay. +**Entry id floor across prefixes.** A region that arrives under a new prefix, by migration or by failover, brings a watermark and replayed ids that were assigned under another prefix's sequence. Before the target assigns the region its first new id, the target store raises its next object sequence above `max_id >> 20`, where `max_id` is the largest of the region's inherited watermark and every id it replayed, so that every new id of that region is greater than every id it already has; sequence gaps are tolerated and the sequence space is wide enough for this. The rule is applied on every open from the region's persisted state, in the same place Mito already hands the store the manifest watermark, so it needs no persisted state of its own and survives a crash before the target's first flush: the restart re-derives the floor and the region's acknowledged-but-unflushed entries under the target prefix are above the inherited watermark and replay. The history check under *Log store* stays compatible with this: it judges a watermark against the prefix a region persists, and a region taken over from another prefix carries that prefix as an earlier segment of its `WalOptions`, so its inherited watermark is judged against the segment that assigned it rather than against the target prefix, and the floor, not the check, applies to it on the target. **Acknowledgement modes.** Both are implemented. Cluster mode is expected to run `enqueued` for latency-sensitive workloads; the object layout, the flush durability barrier and the recovery path are identical in both modes. diff --git a/src/common/wal/src/config.rs b/src/common/wal/src/config.rs index b92c3d13d36..950f833dec9 100644 --- a/src/common/wal/src/config.rs +++ b/src/common/wal/src/config.rs @@ -298,6 +298,7 @@ mod tests { assert_eq!(config.max_unpersisted_bytes, ReadableSize::mb(64)); assert_eq!(config.max_unpersisted_age, Duration::from_secs(8)); assert_eq!(config.on_corrupted_segment, CorruptedSegmentAction::Skip); + assert!(!config.accept_missing_history); assert!(matches!( MetasrvWalConfig::try_from(datanode_wal_config).unwrap_err(), Error::UnsupportedWalProvider { .. } @@ -313,6 +314,7 @@ mod tests { max_unpersisted_bytes = "32MB" max_unpersisted_age = "4s" on_corrupted_segment = "fail" + accept_missing_history = true "#; let datanode_wal_config: DatanodeWalConfig = toml::from_str(toml_str).unwrap(); let expected = ObjectStoreWalConfig { @@ -324,6 +326,7 @@ mod tests { max_unpersisted_bytes: ReadableSize::mb(32), max_unpersisted_age: Duration::from_secs(4), on_corrupted_segment: CorruptedSegmentAction::Fail, + accept_missing_history: true, }; assert_eq!( datanode_wal_config, @@ -334,6 +337,7 @@ mod tests { assert_eq!(table["ack_mode"].as_str(), Some("enqueued")); assert_eq!(table["max_unpersisted_bytes"].as_str(), Some("32MiB")); assert_eq!(table["max_unpersisted_age"].as_str(), Some("4s")); + assert_eq!(table["accept_missing_history"].as_bool(), Some(true)); assert_eq!( datanode_wal_config, toml::from_str::(&serialized).unwrap() diff --git a/src/common/wal/src/config/object_store.rs b/src/common/wal/src/config/object_store.rs index dd2c4544129..9e4a371baae 100644 --- a/src/common/wal/src/config/object_store.rs +++ b/src/common/wal/src/config/object_store.rs @@ -81,6 +81,11 @@ pub struct ObjectStoreWalConfig { /// second fetch, because its checksum does not match or its content /// disagrees with its footer entry, defaults to `skip`. pub on_corrupted_segment: CorruptedSegmentAction, + /// Whether a region whose WAL history the configured location does not + /// hold opens anyway, with its writes that were not flushed lost, instead + /// of failing to open, defaults to false. Meant for one start after the + /// old WAL location is lost. + pub accept_missing_history: bool, } impl Default for ObjectStoreWalConfig { @@ -94,6 +99,7 @@ impl Default for ObjectStoreWalConfig { max_unpersisted_bytes: ReadableSize::mb(64), max_unpersisted_age: Duration::from_secs(8), on_corrupted_segment: CorruptedSegmentAction::Skip, + accept_missing_history: false, } } } diff --git a/src/log-store/src/error.rs b/src/log-store/src/error.rs index c7b9c981c44..51f4be7677e 100644 --- a/src/log-store/src/error.rs +++ b/src/log-store/src/error.rs @@ -424,6 +424,37 @@ pub enum Error { location: Location, }, + #[snafu(display( + "Object store WAL prefix {} on {} does not hold the history of region {}: {}. \ + If the WAL was moved, stop the instance, copy every object under the prefix from the \ + old location to this one and start again. If the old WAL is lost, start once with \ + `[wal].accept_missing_history = true`, which opens the region without its writes that \ + were not flushed, and turn the option off again once the instance has written to the WAL", + prefix, + storage, + region_id, + reason + ))] + MissingWalHistory { + region_id: RegionId, + prefix: String, + /// The storage the prefix is read from, as configured. + storage: String, + reason: String, + #[snafu(implicit)] + location: Location, + }, + + #[snafu(display( + "Object store WAL gave up its start object {} after refusing a region without its history, and accepts nothing more until it opens again", + path + ))] + WalStartWithdrawn { + path: String, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Object store WAL log store is stopped"))] ObjectStoreWalStopped { #[snafu(implicit)] @@ -478,6 +509,7 @@ impl ErrorExt for Error { | OverrideCompactedEntry { .. } | InvalidWalObjectStore { .. } | MismatchedWalPrefix { .. } + | MissingWalHistory { .. } | MismatchedWalRegion { .. } | IncompleteWalEntry { .. } | InvalidWalEntryRange { .. } => StatusCode::InvalidArguments, @@ -494,6 +526,7 @@ impl ErrorExt for Error { | WaitProduceResultReceiver { .. } | MetaLengthExceededLimit { .. } | ObjectStoreWalStopped { .. } + | WalStartWithdrawn { .. } | WalSegmentFetch { .. } => StatusCode::Internal, CorruptedWalObject { .. } diff --git a/src/log-store/src/metrics.rs b/src/log-store/src/metrics.rs index 8f923615964..bdb123759eb 100644 --- a/src/log-store/src/metrics.rs +++ b/src/log-store/src/metrics.rs @@ -97,6 +97,17 @@ lazy_static! { "object store logstore skipped corrupted segments total", ) .unwrap(); + + /// Counter of regions the object store logstore let open although its + /// prefix does not hold their WAL history, because + /// `accept_missing_history` is set; their writes that were not flushed + /// are lost. + pub static ref METRIC_OBJECT_STORE_WAL_MISSING_HISTORY_ACCEPTED_TOTAL: IntCounter = register_int_counter!( + "greptime_logstore_object_store_wal_missing_history_accepted_total", + "object store logstore regions opened without their WAL history total", + ) + .unwrap(); + /// Counter of objects the object store logstore deleted because no /// region needs them any more. pub static ref METRIC_OBJECT_STORE_WAL_DELETED_OBJECTS_TOTAL: IntCounter = register_int_counter!( diff --git a/src/log-store/src/object_store_wal/io.rs b/src/log-store/src/object_store_wal/io.rs index 4fc7e971535..f9ed5eee36b 100644 --- a/src/log-store/src/object_store_wal/io.rs +++ b/src/log-store/src/object_store_wal/io.rs @@ -181,6 +181,18 @@ impl ObjectStoreIo { ) } + /// Describes the object store the objects live in by its scheme, its + /// bucket when it has one and its root, none of which is a credential. + pub(crate) fn storage(&self) -> String { + let info = self.store.info(); + let name = info.name(); + if name.is_empty() { + format!("{}, root {:?}", info.scheme(), info.root()) + } else { + format!("{}, bucket {name:?}, root {:?}", info.scheme(), info.root()) + } + } + fn parse_object_seq(&self, path: &str) -> Option { let value = path .strip_prefix(&self.object_prefix)? diff --git a/src/log-store/src/object_store_wal/store.rs b/src/log-store/src/object_store_wal/store.rs index fc32bef707e..0a8c68c83e7 100644 --- a/src/log-store/src/object_store_wal/store.rs +++ b/src/log-store/src/object_store_wal/store.rs @@ -42,16 +42,18 @@ use tokio_util::task::AbortOnDropHandle; use crate::error::{ CorruptedWalObjectSnafu, Error, IncompleteWalEntrySnafu, InvalidProviderSnafu, InvalidWalObjectSnafu, InvalidWalObjectStoreSnafu, MismatchedWalPrefixSnafu, - MismatchedWalRegionSnafu, ObjectStoreWalFailingSnafu, ObjectStoreWalSnafu, - ObjectStoreWalStoppedSnafu, Result, StaleWalObjectSnafu, UnconfirmedWalEpochStartSnafu, - WalObjectSequenceExhaustedSnafu, WalObjectSequenceUnsettledSnafu, WalSegmentFetchSnafu, + MismatchedWalRegionSnafu, MissingWalHistorySnafu, ObjectStoreWalFailingSnafu, + ObjectStoreWalSnafu, ObjectStoreWalStoppedSnafu, Result, StaleWalObjectSnafu, + UnconfirmedWalEpochStartSnafu, WalObjectSequenceExhaustedSnafu, + WalObjectSequenceUnsettledSnafu, WalSegmentFetchSnafu, WalStartWithdrawnSnafu, }; use crate::metrics::{ METRIC_OBJECT_STORE_WAL_APPEND_ACK_SECONDS, METRIC_OBJECT_STORE_WAL_CREATE_CONFLICTS_TOTAL, METRIC_OBJECT_STORE_WAL_CREATE_FAILURES_TOTAL, METRIC_OBJECT_STORE_WAL_CREATED_OBJECTS_TOTAL, METRIC_OBJECT_STORE_WAL_DELETED_OBJECTS_TOTAL, METRIC_OBJECT_STORE_WAL_FAILED_DELETES_TOTAL, METRIC_OBJECT_STORE_WAL_FAILING, METRIC_OBJECT_STORE_WAL_INDEXED_BYTES, - METRIC_OBJECT_STORE_WAL_INDEXED_OBJECTS, METRIC_OBJECT_STORE_WAL_OBJECT_BYTES, + METRIC_OBJECT_STORE_WAL_INDEXED_OBJECTS, + METRIC_OBJECT_STORE_WAL_MISSING_HISTORY_ACCEPTED_TOTAL, METRIC_OBJECT_STORE_WAL_OBJECT_BYTES, METRIC_OBJECT_STORE_WAL_OBJECT_ENTRIES, METRIC_OBJECT_STORE_WAL_POISONED_TOTAL, METRIC_OBJECT_STORE_WAL_READ_SECONDS, METRIC_OBJECT_STORE_WAL_RECOVERED_OBJECTS_TOTAL, METRIC_OBJECT_STORE_WAL_RECOVERY_SECONDS, METRIC_OBJECT_STORE_WAL_REJECTED_APPENDS_TOTAL, @@ -144,8 +146,17 @@ const RECOVERY_TAIL_WINDOW: usize = 64 * 1024; /// once and returns, in both modes, once that object is durable and indexed. pub struct ObjectStoreLogStore { prefix: String, + /// The storage the prefix is read from, as configured, for messages. + storage: String, ack_mode: AckMode, on_corrupted_segment: CorruptedSegmentAction, + accept_missing_history: bool, + /// Whether the prefix held no object when the store opened. + found_empty: bool, + /// The sequence above the newest object holding entries that recovery + /// indexed, zero when it indexed none. + history_end: u64, + started_regions: StartedRegions, io: Arc, catalog: Arc>, obsolete_entry_ids: ObsoleteEntryIds, @@ -175,6 +186,12 @@ pub struct ObjectStoreLogStore { type ObsoleteEntryIds = Arc>>; +/// Regions whose history starts in the store: those created on it and those +/// it opened by accepting a missing history. The store registers a region and +/// the actor decides a withdrawal under this lock, so neither misses the +/// other. +type StartedRegions = Arc>>; + type WalHoles = Arc>>>; /// A segment a read skipped because it did not decode: the entries of one @@ -332,14 +349,29 @@ impl ObjectStoreLogStore { ), } ); + let found_empty = next_object_seq == 0; + let history_end = durable_entry_ids + .values() + .map(|entry_id| sequence_floor(*entry_id)) + .max() + .unwrap_or(0); // Only a recovery that succeeded is reported, so that the objects // counted and the time measured are those of the same recoveries. METRIC_OBJECT_STORE_WAL_RECOVERED_OBJECTS_TOTAL.inc_by(recovered_objects); METRIC_OBJECT_STORE_WAL_RECOVERY_SECONDS .observe(recovery_started_at.elapsed().as_secs_f64()); publish_indexed(&catalog); + // A chain without any segment holds no history, so the start object + // reclaims every object below it. A start object that another + // instance withdraws after this recovery listed it is then below the + // boundary, and its delete leaves this chain complete. + let start_boundary = if durable_entry_ids.is_empty() { + next_object_seq + } else { + reclaim_boundary + }; let (start, moved_past) = - start_epoch(io.as_ref(), next_object_seq, tip, reclaim_boundary).await?; + start_epoch(io.as_ref(), next_object_seq, tip, start_boundary).await?; let epoch = start.epoch; // The epoch identifies this instance in every object it writes. info!( @@ -377,8 +409,11 @@ impl ObjectStoreLogStore { #[cfg(any(test, feature = "testing"))] let (durability_waits_tx, durability_waits_rx) = watch::channel(0); + let started_regions = StartedRegions::default(); let actor = Actor { io: io.clone(), + started_regions: started_regions.clone(), + withdrawal: Withdrawal::Kept, catalog: catalog.clone(), obsolete_entry_ids: obsolete_entry_ids.clone(), wal_holes: wal_holes.clone(), @@ -430,10 +465,19 @@ impl ObjectStoreLogStore { #[cfg(test)] LIVE_ACTORS.fetch_add(1, Ordering::SeqCst); common_runtime::spawn_global(actor.run()); + let storage = match config.storage_provider.as_str() { + "" => format!("the default storage ({})", io.storage()), + name => format!("storage provider {name:?} ({})", io.storage()), + }; Ok(Arc::new(Self { prefix, + storage, ack_mode: config.ack_mode, on_corrupted_segment: config.on_corrupted_segment, + accept_missing_history: config.accept_missing_history, + found_empty, + history_end, + started_regions, io, catalog, obsolete_entry_ids, @@ -551,6 +595,97 @@ impl ObjectStoreLogStore { .context(ObjectStoreWalStoppedSnafu)?; response_rx.await.ok().context(ObjectStoreWalStoppedSnafu)? } + + /// Returns why the prefix cannot hold the WAL history of `region_id` + /// through the watermark `entry_id`, or `None` when it can. + /// + /// Two invariants of a prefix decide it. First, a prefix that any store + /// opened is never empty afterwards, unless that store withdrew its start + /// object, see [`withdraw_start`](Self::withdraw_start): every open writes + /// its start object before it admits an append, and collection deletes + /// only objects below the reclaim boundary of the last indexed object, + /// which carries a boundary at or below its own sequence (a larger one + /// fails its header check), so the newest indexed object always stays. + /// Second, a watermark a region persisted names an object at or below + /// the newest object holding entries that any later recovery of the + /// prefix indexes. The entry it names was durable and indexed when the + /// watermark was persisted, since an append of the `durable` mode returns + /// only then and Mito waits for durability before it publishes a + /// watermark in the `enqueued` mode. No durable boundary passes the + /// newest object holding entries: a batch carries at most the sequence + /// after the last indexed object, which is below the batch, and a start + /// object carries the boundary recovery found, or its own sequence on a + /// chain without segments. So that object stays, and + /// recovery indexes it on the chain of the tip. The segment of release + /// records counts as entries: it never holds the reclaim boundary, so an + /// object holding only releases can be the one that collection keeps + /// while it deletes every older object. Objects without segments, such as + /// the start objects of opens this check refused, do not move it. + /// + /// So the history is missing when this store found the prefix empty, or + /// when the watermark names an object above the newest one holding + /// entries that it found, unless the region's history starts in this + /// store or the catalog holds the region's entries through the + /// watermark: those of a region written since the store opened. + fn missing_history(&self, region_id: RegionId, entry_id: EntryId) -> Option { + if self + .started_regions + .lock() + .unwrap_or_else(PoisonError::into_inner) + .contains(®ion_id) + { + return None; + } + if self + .catalog + .read() + .unwrap_or_else(PoisonError::into_inner) + .region_max_entry_id(region_id) + .is_some_and(|max_entry_id| max_entry_id >= entry_id) + { + return None; + } + if self.found_empty { + return Some("the prefix held no object when the WAL opened".to_string()); + } + let sequence_floor = sequence_floor(entry_id); + (sequence_floor > self.history_end).then(|| { + let held = match self.history_end { + 0 => "no entry".to_string(), + history_end => format!("entries only in objects below {history_end}"), + }; + format!( + "its watermark, entry {entry_id}, names object {}, but the prefix held {held} when the WAL opened", + sequence_floor - 1 + ) + }) + } + + /// Deletes the start object of a store that found its prefix empty and + /// whose open of an existing region is refused, so that the next open + /// finds the prefix empty again and refuses the region as well, instead + /// of passing a region that never flushed. The actor gives the object up + /// only while the store has admitted nothing, has no batch, holds no + /// other object and no region's history started in it, and from then on + /// the store accepts nothing, since a later object would link to the + /// deleted one. The actor owns the delete like any collection, so `stop` + /// waits for it and a cancelled caller does not abandon it, and every + /// refusal returns only once it settled. A failed delete, or a crash + /// before the delete completes, leaves the start object behind, and the + /// next open passes a region that never flushed. + async fn withdraw_start(&self) { + let (response_tx, response_rx) = oneshot::channel(); + if self + .command_tx + .send(Command::Withdraw { + response: response_tx, + }) + .await + .is_ok() + { + let _ = response_rx.await; + } + } } /// The transient object store error of a create that a testing hook fails. @@ -942,7 +1077,12 @@ impl LogStore for ObjectStoreLogStore { /// object is raised above the object `entry_id` names unless it is there /// already, so a watermark assigned under another prefix, or one whose /// object this prefix no longer holds, is never passed by a new id. Zero - /// names no object and needs no floor. + /// names no object and needs no floor. Before Mito opens a region that + /// writes this WAL from its persisted state, it calls + /// [`check_region_history`](LogStore::check_region_history), which + /// refuses a watermark naming an object at or above the next sequence + /// the store found unless `accept_missing_history` is set, so the floor + /// moves for such a region only once that loss was accepted. /// /// The watermark and the floor are applied together by the actor: a call /// cancelled before its command is queued publishes neither, while a @@ -1044,6 +1184,51 @@ impl LogStore for ObjectStoreLogStore { .context(ObjectStoreWalStoppedSnafu)?; response_rx.await.ok().context(ObjectStoreWalStoppedSnafu)? } + + /// Fails with [`Error::MissingWalHistory`] when the prefix cannot hold + /// the history of the provider's region through `entry_id`, unless + /// `accept_missing_history` is set: then the region's writes that were + /// not flushed are lost, which is logged and counted. + async fn check_region_history(&self, provider: &Provider, entry_id: EntryId) -> Result<()> { + let region_id = self.region_of(provider)?; + let Some(reason) = self.missing_history(region_id, entry_id) else { + return self.check_terminal(); + }; + if !self.accept_missing_history { + if self.found_empty { + self.withdraw_start().await; + } + return MissingWalHistorySnafu { + region_id, + prefix: self.prefix.clone(), + storage: self.storage.clone(), + reason, + } + .fail(); + } + self.check_terminal()?; + warn!( + "Opening region {region_id} without its WAL history, as accept_missing_history allows: prefix {} on {} does not hold it, {reason}; the region's writes above entry {entry_id} that were not flushed are lost", + self.prefix, self.storage + ); + METRIC_OBJECT_STORE_WAL_MISSING_HISTORY_ACCEPTED_TOTAL.inc(); + self.start_region_history(provider) + } + + /// Records the provider's region as one whose history starts in this + /// store, which is kept in memory: a later store on the prefix judges + /// the region by its watermark and the objects it finds. + fn start_region_history(&self, provider: &Provider) -> Result<()> { + let region_id = self.region_of(provider)?; + let mut started_regions = self + .started_regions + .lock() + .unwrap_or_else(PoisonError::into_inner); + // A withdrawal records its terminal error under the same lock. + self.check_terminal()?; + started_regions.insert(region_id); + Ok(()) + } } type AppendResponse = oneshot::Sender>; @@ -1075,6 +1260,9 @@ enum Command { full: bool, response: oneshot::Sender>, }, + /// Answered once the store keeps its start object, or once the delete + /// of the start object it gave up settled. + Withdraw { response: oneshot::Sender<()> }, Stop { response: oneshot::Sender>, }, @@ -1151,6 +1339,19 @@ impl SealedBatch { type CreateOutcome = (u64, Result); +/// Whether the store gave up its start object, see +/// [`Actor::handle_withdraw`]. +enum Withdrawal { + Kept, + /// The delete of the start object `object_seq` is in flight; `waiters` + /// are answered once it settled. + Deleting { + object_seq: u64, + waiters: Vec>, + }, + Settled, +} + /// The state of an `enqueued` store whose backlog the object store refuses /// with an error that requires intervention, see [`requires_intervention`], /// such as denied @@ -1242,6 +1443,8 @@ impl Failing { /// complete before `stop` returns. struct Actor { io: Arc, + started_regions: StartedRegions, + withdrawal: Withdrawal, catalog: Arc>, obsolete_entry_ids: ObsoleteEntryIds, wal_holes: WalHoles, @@ -1357,6 +1560,9 @@ impl Actor { Some(Command::Release { region_id, full, response }) => { self.handle_release(region_id, full, response); } + Some(Command::Withdraw { response }) => { + self.handle_withdraw(response); + } Some(Command::Stop { response }) => { self.handle_stop(response); } @@ -2114,6 +2320,79 @@ impl Actor { self.durability_waits.send_modify(|count| *count += 1); } + /// Gives up the start object while it is the only object this store + /// knows of: the catalog indexes nothing else, nothing was ever handed an + /// id, no batch, append or release waits, no other object was listed or + /// created, no region's history started in the store, and the store is + /// neither failing nor terminal. A release seals its own batch without + /// handing an id, so the catalog decides once its object is indexed, and + /// the object it may then leave alone holds a segment. The store then + /// records a terminal error, so nothing is admitted that would link to + /// the object, and deletes it like a collected object; `response` is + /// answered once that delete settled, or at once when the store keeps + /// the object. A later request waits for the same delete. + fn handle_withdraw(&mut self, response: oneshot::Sender<()>) { + match &mut self.withdrawal { + Withdrawal::Kept => {} + Withdrawal::Deleting { waiters, .. } => { + waiters.push(response); + return; + } + Withdrawal::Settled => { + let _ = response.send(()); + return; + } + } + let started_regions = self + .started_regions + .lock() + .unwrap_or_else(PoisonError::into_inner); + if !started_regions.is_empty() + || terminal(&self.terminal_error).is_some() + || self.failing.is_some() + || !self.issued_entry_ids.is_empty() + || !self.open_batch.is_empty() + || !self.sealed.is_empty() + || !self.creates.is_empty() + || self.stalled.is_some() + || !self.unindexed.is_empty() + || !self.failed_releases.is_empty() + { + let _ = response.send(()); + return; + } + // The start object is the one object indexed, and the only one + // without segments: a release object can be indexed alone once it + // collected the start object. + let start_object = { + let catalog = self.catalog.read().unwrap_or_else(PoisonError::into_inner); + let mut objects = catalog.objects_in_order(); + match (objects.next(), objects.next()) { + (Some((object_seq, [])), None) => Some(object_seq), + _ => None, + } + }; + let Some(object_seq) = start_object else { + let _ = response.send(()); + return; + }; + let path = self.io.object_path(object_seq); + self.terminal_error + .lock() + .unwrap_or_else(PoisonError::into_inner) + .get_or_insert_with(|| Arc::new(WalStartWithdrawnSnafu { path }.build())); + drop(started_regions); + self.deleting.start(object_seq); + let io = self.io.clone(); + self.deletes.push(Box::pin(async move { + (object_seq, io.delete(object_seq).await) + })); + self.withdrawal = Withdrawal::Deleting { + object_seq, + waiters: vec![response], + }; + } + /// Makes sure no id of the region at or below `entry_id` is assigned /// from now on, then records the obsolete watermark of the region; /// neither is done when the sequence cannot be raised. @@ -2252,6 +2531,36 @@ impl Actor { /// succeeded, then settles its attempt, so a read that the attempt /// releases sees the outcome in the catalog. fn on_delete_completed(&mut self, object_seq: u64, result: Result<()>) { + if let Withdrawal::Deleting { + object_seq: withdrawn, + waiters, + } = &mut self.withdrawal + && *withdrawn == object_seq + { + let path = self.io.object_path(object_seq); + match result { + Ok(()) => { + let mut catalog = self.catalog.write().unwrap_or_else(PoisonError::into_inner); + if let Some(footer) = catalog.remove_object(object_seq) { + METRIC_OBJECT_STORE_WAL_INDEXED_OBJECTS.dec(); + METRIC_OBJECT_STORE_WAL_INDEXED_BYTES.sub(object_len(&footer) as i64); + } + info!( + "Deleted the start object {path} of the object store WAL, whose prefix held no object, after refusing a region without its history" + ); + } + Err(error) => warn!( + error; + "Failed to delete the start object {path} of the object store WAL after refusing a region without its history; the next open passes regions that never flushed" + ), + } + for waiter in waiters.drain(..) { + let _ = waiter.send(()); + } + self.withdrawal = Withdrawal::Settled; + self.deleting.finish(object_seq); + return; + } match result { Ok(()) => { self.unindexed.remove(&object_seq); @@ -3187,6 +3496,9 @@ pub(crate) trait WalObjectIo: Send + Sync { async fn list(&self) -> Result>; fn object_path(&self, object_seq: u64) -> String; + + /// Describes the object store the objects live in, for messages. + fn storage(&self) -> String; } #[async_trait::async_trait] @@ -3214,6 +3526,10 @@ impl WalObjectIo for ObjectStoreIo { fn object_path(&self, object_seq: u64) -> String { ObjectStoreIo::object_path(self, object_seq) } + + fn storage(&self) -> String { + ObjectStoreIo::storage(self) + } } #[cfg(test)] @@ -4116,6 +4432,500 @@ mod tests { store.stop().await.unwrap(); } + /// Returns the message of the missing history error the check of + /// `region_id` at `entry_id` fails with, or `None` when it passes. + async fn missing_history( + store: &ObjectStoreLogStore, + region_id: RegionId, + entry_id: EntryId, + ) -> Option { + match store + .check_region_history(&provider(region_id), entry_id) + .await + { + Ok(()) => None, + Err(error @ Error::MissingWalHistory { .. }) => { + assert_eq!( + common_error::status_code::StatusCode::InvalidArguments, + error.status_code() + ); + Some(error.to_string()) + } + Err(error) => panic!("unexpected error: {error:?}"), + } + } + + /// Copies every object under the prefix of `from` to the same key in `to`. + async fn copy_objects(from: &ObjectStore, to: &ObjectStore) { + let io = ObjectStoreIo::new(from.clone(), PREFIX).unwrap(); + for object in io.list().await.unwrap() { + let bytes = from.read(&object.path).await.unwrap().to_vec(); + to.write(&object.path, bytes).await.unwrap(); + } + } + + #[tokio::test] + async fn test_store_history_is_missing_from_a_prefix_found_empty() { + let _serialized = SERIALIZED.lock().await; + let config = ObjectStoreWalConfig { + storage_provider: "wal-store".to_string(), + ..eager() + }; + let object_store = memory_store(); + let store = open(object_store.clone(), &config).await; + + // A region written since the store opened holds its history through + // its entries, not past them. + let written = append(&store, region(2), "a").await.unwrap().last_entry_ids[®ion(2)]; + assert_eq!(None, missing_history(&store, region(2), 0).await); + assert_eq!(None, missing_history(&store, region(2), written).await); + assert!( + missing_history(&store, region(2), written + 1) + .await + .is_some() + ); + + // No other region has a history under a prefix that held no object, + // not even one that never flushed. The store admitted an append, so + // it keeps its objects. + for entry_id in [0, id(0, 1), id(7, 1)] { + let message = missing_history(&store, region(1), entry_id).await.unwrap(); + for expected in [ + ®ion(1).to_string(), + PREFIX, + // A memory store names itself by its address as its bucket. + "storage provider \"wal-store\" (memory, bucket ", + "root \"/\")", + "the prefix held no object when the WAL opened", + "copy every object under the prefix", + "`[wal].accept_missing_history = true`", + ] { + assert!(message.contains(expected), "{message}"); + } + } + assert_eq!(vec![1], object_seqs(store.io.as_ref()).await); + + // A region created on the store has its history here without any + // entry. + store.start_region_history(&provider(region(3))).unwrap(); + assert_eq!(None, missing_history(&store, region(3), 0).await); + + // The provider must select the store's prefix. + let foreign = Provider::object_store_provider(region(1), "wal/datanodes/9/epochs/2".into()); + assert!(matches!( + store.check_region_history(&foreign, 0).await.unwrap_err(), + Error::MismatchedWalPrefix { .. } + )); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_withdraws_its_start_object_when_it_refuses_a_region() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + + // Every open of the prefix refuses the region, since each refusal + // deletes the start object the open wrote and the next open finds + // the prefix empty again. + for _ in 0..2 { + let store = open(object_store.clone(), &eager()).await; + assert_eq!(vec![0], object_seqs(store.io.as_ref()).await); + let indexed_objects = METRIC_OBJECT_STORE_WAL_INDEXED_OBJECTS.get(); + let indexed_bytes = METRIC_OBJECT_STORE_WAL_INDEXED_BYTES.get(); + assert!(missing_history(&store, region(1), 0).await.is_some()); + assert!(object_seqs(store.io.as_ref()).await.is_empty()); + // The catalog and the indexed gauges no longer hold the object. + assert!(!store.catalog.read().unwrap().contains_object(0)); + assert_eq!( + (indexed_objects - 1, indexed_bytes - object_len(&[]) as i64), + ( + METRIC_OBJECT_STORE_WAL_INDEXED_OBJECTS.get(), + METRIC_OBJECT_STORE_WAL_INDEXED_BYTES.get() + ) + ); + + // The store accepts nothing more, and refuses again without a + // second delete. + assert!(missing_history(&store, region(2), 0).await.is_some()); + let entry = Entry::Naive(NaiveEntry { + provider: provider(region(2)), + region_id: region(2), + entry_id: 0, + data: b"a".to_vec(), + }); + for error in [ + store.append_batch(vec![entry]).await.unwrap_err(), + store + .start_region_history(&provider(region(3))) + .unwrap_err(), + ] { + assert!( + matches!(unwrap_shared(&error), Error::WalStartWithdrawn { .. }), + "{error:?}" + ); + } + store.stop().await.unwrap(); + } + + // A store whose region history started keeps its start object. + let store = open(object_store.clone(), &eager()).await; + store.start_region_history(&provider(region(3))).unwrap(); + assert!(missing_history(&store, region(1), 0).await.is_some()); + assert_eq!(vec![0], object_seqs(store.io.as_ref()).await); + assert_eq!(None, missing_history(&store, region(3), 0).await); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_withdrawal_outlives_its_caller_and_holds_up_stop() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + let (io, mut parked) = ParkedIo::parking_creates(object_store.clone()); + let store = open_over(io.clone(), &eager()).await; + io.deletes_parked.store(true, Ordering::SeqCst); + + // The refusal that withdraws the start object is cancelled once the + // actor started its delete. + let first = { + let store = store.clone(); + tokio::spawn(async move { store.check_region_history(&provider(region(1)), 0).await }) + }; + let (object_seq, release) = next_parked_operation(&mut parked).await; + assert_eq!(0, object_seq); + first.abort(); + + // A later withdrawal request and stop reach the actor, and a + // durability wait queued behind them is answered, so the actor + // handled both; neither is answered while the delete is in flight. + let (withdrawn_tx, mut withdrawn) = oneshot::channel(); + store + .command_tx + .send(Command::Withdraw { + response: withdrawn_tx, + }) + .await + .unwrap(); + store.begin_stop(); + let (stopped_tx, mut stopped) = oneshot::channel(); + store + .command_tx + .send(Command::Stop { + response: stopped_tx, + }) + .await + .unwrap(); + let (answered_tx, answered) = oneshot::channel(); + store + .command_tx + .send(Command::WaitDurable { + region_id: region(1), + entry_id: 1, + response: answered_tx, + }) + .await + .unwrap(); + timeout(WAIT, answered).await.unwrap().unwrap().unwrap_err(); + assert!(withdrawn.try_recv().is_err()); + assert!(stopped.try_recv().is_err()); + + release.send(true).unwrap(); + timeout(WAIT, withdrawn).await.unwrap().unwrap(); + timeout(WAIT, stopped).await.unwrap().unwrap().unwrap(); + assert!(object_seqs(io.as_ref()).await.is_empty()); + assert!(missing_history(&store, region(2), 0).await.is_some()); + + // The next open finds the prefix empty and refuses the region again. + let store = open(object_store, &eager()).await; + assert!(missing_history(&store, region(1), 0).await.is_some()); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_withdrawal_keeps_the_chain_of_a_successor() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + let (io, mut parked) = ParkedIo::parking_creates(object_store.clone()); + let store = open_over(io.clone(), &eager()).await; + io.deletes_parked.store(true, Ordering::SeqCst); + + // While the withdrawal of start object 0 waits on its delete, a + // successor opens the prefix, links its start object 1 to object 0 + // and stops before any append. + let refused = { + let store = store.clone(); + tokio::spawn(async move { store.check_region_history(&provider(region(1)), 0).await }) + }; + let (object_seq, release) = next_parked_operation(&mut parked).await; + assert_eq!(0, object_seq); + let successor = open(object_store.clone(), &eager()).await; + assert_eq!(vec![0, 1], object_seqs(successor.io.as_ref()).await); + successor.stop().await.unwrap(); + release.send(true).unwrap(); + timeout(WAIT, refused).await.unwrap().unwrap().unwrap_err(); + store.stop().await.unwrap(); + + // Object 1 reclaims the deleted object 0, so its chain is complete. + let store = open(object_store, &eager()).await; + assert_eq!(vec![1, 2], object_seqs(store.io.as_ref()).await); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_keeps_indexing_a_start_object_whose_withdrawal_failed() { + let _serialized = SERIALIZED.lock().await; + let (io, mut parked) = ParkedIo::parking_creates(memory_store()); + let store = open_over(io.clone(), &eager()).await; + io.deletes_parked.store(true, Ordering::SeqCst); + let indexed_objects = METRIC_OBJECT_STORE_WAL_INDEXED_OBJECTS.get(); + let indexed_bytes = METRIC_OBJECT_STORE_WAL_INDEXED_BYTES.get(); + + let refused = { + let store = store.clone(); + tokio::spawn(async move { store.check_region_history(&provider(region(1)), 0).await }) + }; + let (object_seq, release) = next_parked_operation(&mut parked).await; + assert_eq!(0, object_seq); + release.send(false).unwrap(); + assert!(matches!( + timeout(WAIT, refused).await.unwrap().unwrap().unwrap_err(), + Error::MissingWalHistory { .. } + )); + + // The object is still there, so the catalog and the indexed gauges + // keep it. + assert_eq!(vec![0], object_seqs(io.as_ref()).await); + assert!(store.catalog.read().unwrap().contains_object(0)); + assert_eq!( + (indexed_objects, indexed_bytes), + ( + METRIC_OBJECT_STORE_WAL_INDEXED_OBJECTS.get(), + METRIC_OBJECT_STORE_WAL_INDEXED_BYTES.get() + ) + ); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_history_is_missing_above_the_newest_object_holding_entries() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + // The start object 0, then objects 1 and 2 of region 1. + let store = open(object_store.clone(), &eager()).await; + append(&store, region(1), "a").await.unwrap(); + append(&store, region(1), "b").await.unwrap(); + store.stop().await.unwrap(); + + // Every restart writes one more start object, but the newest object + // holding entries stays object 2, so a refused open does not let + // the next one pass a watermark of object 3. + for start_object in [3, 4] { + let store = open(object_store.clone(), &eager()).await; + assert!( + store.storage.starts_with("the default storage (memory, ") + && store.storage.ends_with("root \"/\")"), + "{}", + store.storage + ); + for region_id in [region(1), region(2)] { + for entry_id in [0, id(1, 1), id(2, 1), id(2, POSITION_LIMIT - 1)] { + assert_eq!(None, missing_history(&store, region_id, entry_id).await); + } + for entry_id in [id(3, 1), id(9, 1)] { + let message = missing_history(&store, region_id, entry_id).await.unwrap(); + let reason = format!( + "its watermark, entry {entry_id}, names object {}, but the prefix held entries only in objects below 3 when the WAL opened", + entry_id >> 20 + ); + assert!(message.contains(&reason), "{message}"); + } + } + assert_eq!( + Some(&start_object), + object_seqs(store.io.as_ref()).await.last() + ); + store.stop().await.unwrap(); + } + + // An entry written since then lies above the objects found, and the + // catalog holds it. + let store = open(object_store, &eager()).await; + let written = append(&store, region(2), "c").await.unwrap().last_entry_ids[®ion(2)]; + assert_eq!(id(6, 1), written); + assert_eq!(None, missing_history(&store, region(2), written).await); + assert!(missing_history(&store, region(1), written).await.is_some()); + store.stop().await.unwrap(); + + // A prefix that holds only start objects holds no entry. + let object_store = memory_store(); + open(object_store.clone(), &eager()) + .await + .stop() + .await + .unwrap(); + let store = open(object_store, &eager()).await; + assert_eq!(None, missing_history(&store, region(1), 0).await); + let message = missing_history(&store, region(1), id(0, 1)).await.unwrap(); + assert!( + message.contains("names object 0, but the prefix held no entry when the WAL opened"), + "{message}" + ); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_history_survives_collection_down_to_the_newest_object() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + let store = open(object_store.clone(), &eager()).await; + // Region 1 flushes object 1; object 2 of region 2 carries the boundary + // past it, so only object 2 is left once region 2 flushes it too. + let flushed_1 = append(&store, region(1), "a").await.unwrap().last_entry_ids[®ion(1)]; + store + .obsolete(&provider(region(1)), region(1), flushed_1) + .await + .unwrap(); + let flushed_2 = append(&store, region(2), "b").await.unwrap().last_entry_ids[®ion(2)]; + store + .obsolete(&provider(region(2)), region(2), flushed_2) + .await + .unwrap(); + assert_eq!(vec![2], collected_object_seqs(&store).await); + store.stop().await.unwrap(); + + // The object region 1 flushed is gone, but the newest object stays + // and keeps the next sequence above it. + let store = open(object_store, &eager()).await; + assert_eq!(0, latest(&store, region(1))); + assert_eq!(None, missing_history(&store, region(1), flushed_1).await); + assert_eq!(None, missing_history(&store, region(2), flushed_2).await); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_history_survives_collection_down_to_a_release_object() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + let store = open(object_store.clone(), &eager()).await; + // Region 1 flushes object 1; the release of region 2 seals object 2, + // which holds only its release record and carries the boundary past + // object 1. + let flushed = append(&store, region(1), "a").await.unwrap().last_entry_ids[®ion(1)]; + store + .obsolete(&provider(region(1)), region(1), flushed) + .await + .unwrap(); + store + .release_region(&provider(region(2)), region(2)) + .await + .unwrap(); + assert_eq!(vec![2], collected_object_seqs(&store).await); + store.stop().await.unwrap(); + + // The object holding the release record bounds the history. + let store = open(object_store, &eager()).await; + assert_eq!(None, missing_history(&store, region(1), flushed).await); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_keeps_its_objects_after_a_release_when_it_refuses_a_region() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + let store = open(object_store.clone(), &eager()).await; + // The release seals object 1, which collects the start object. + store + .release_region(&provider(region(2)), region(2)) + .await + .unwrap(); + assert_eq!(vec![1], collected_object_seqs(&store).await); + + // The refusal leaves the object holding the release record, and the + // store keeps accepting appends. + assert!(missing_history(&store, region(1), 0).await.is_some()); + assert_eq!(vec![1], object_seqs(store.io.as_ref()).await); + append(&store, region(3), "a").await.unwrap(); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_history_moves_with_a_copy_of_the_objects() { + let _serialized = SERIALIZED.lock().await; + let old_location = memory_store(); + let store = open(old_location.clone(), &eager()).await; + let flushed = append(&store, region(1), "a").await.unwrap().last_entry_ids[®ion(1)]; + store + .obsolete(&provider(region(1)), region(1), flushed) + .await + .unwrap(); + let unflushed = append(&store, region(1), "b").await.unwrap().last_entry_ids[®ion(1)]; + assert_eq!(vec![2], collected_object_seqs(&store).await); + store.stop().await.unwrap(); + + // Another location without the objects does not hold the history. + let store = open(memory_store(), &eager()).await; + assert!(missing_history(&store, region(1), flushed).await.is_some()); + store.stop().await.unwrap(); + + // A copy of every object does, and replays the unflushed entry. + let new_location = memory_store(); + copy_objects(&old_location, &new_location).await; + let store = open(new_location, &eager()).await; + assert_eq!(None, missing_history(&store, region(1), flushed).await); + assert_eq!( + expected_entries(region(1), &[(unflushed, "b")]), + read_entries(&store, region(1), flushed + 1).await + ); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_accepts_a_missing_history_when_configured() { + let _serialized = SERIALIZED.lock().await; + let config = ObjectStoreWalConfig { + accept_missing_history: true, + ..eager() + }; + let store = open(memory_store(), &config).await; + let accepted = METRIC_OBJECT_STORE_WAL_MISSING_HISTORY_ACCEPTED_TOTAL.get(); + store + .check_region_history(&provider(region(1)), id(5, 1)) + .await + .unwrap(); + assert_eq!( + accepted + 1, + METRIC_OBJECT_STORE_WAL_MISSING_HISTORY_ACCEPTED_TOTAL.get() + ); + + // The accepted region's history starts in this store, so a second + // open of it is not counted again, and neither is a history the + // prefix holds. + store + .check_region_history(&provider(region(1)), id(5, 1)) + .await + .unwrap(); + let written = append(&store, region(2), "a").await.unwrap().last_entry_ids[®ion(2)]; + store + .check_region_history(&provider(region(2)), written) + .await + .unwrap(); + assert_eq!( + accepted + 1, + METRIC_OBJECT_STORE_WAL_MISSING_HISTORY_ACCEPTED_TOTAL.get() + ); + + // Once Mito records the accepted watermark, the region's next id is + // above it. + store + .obsolete(&provider(region(1)), region(1), id(5, 1)) + .await + .unwrap(); + assert_eq!( + HashMap::from([(region(1), id(6, 1))]), + append(&store, region(1), "b").await.unwrap().last_entry_ids + ); + store.stop().await.unwrap(); + } + #[tokio::test] async fn test_store_rejects_exhausted_sequence() { let _serialized = SERIALIZED.lock().await; @@ -4206,8 +5016,13 @@ mod tests { let (append_tx, _) = mpsc::channel(APPEND_BUFFER); let store = ObjectStoreLogStore { prefix: PREFIX.to_string(), + storage: String::new(), ack_mode: AckMode::Durable, on_corrupted_segment: CorruptedSegmentAction::Skip, + accept_missing_history: false, + found_empty: false, + history_end: 0, + started_regions: StartedRegions::default(), io: Arc::new(ObjectStoreIo::new(memory_store(), PREFIX).unwrap()), catalog: Arc::default(), obsolete_entry_ids: ObsoleteEntryIds::default(), @@ -6680,9 +7495,12 @@ mod tests { } else { (0, 1, None) }; - let same = encode_object(chain_header(start_seq, epoch, tip), &[]) - .unwrap() - .bytes; + // A prefix without segments is reclaimed up to the start object. + let header = Header { + reclaim_boundary: start_seq, + ..chain_header(start_seq, epoch, tip) + }; + let same = encode_object(header, &[]).unwrap().bytes; let io = Arc::new(RacingIo { inner: ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(), late: Mutex::new(Some((start_seq, same))), @@ -8538,6 +9356,10 @@ mod tests { fn object_path(&self, object_seq: u64) -> String { self.inner.object_path(object_seq) } + + fn storage(&self) -> String { + self.inner.storage() + } } /// Object access whose first create lets another object land at a given @@ -8576,6 +9398,10 @@ mod tests { fn object_path(&self, object_seq: u64) -> String { self.inner.object_path(object_seq) } + + fn storage(&self) -> String { + self.inner.storage() + } } /// Object access whose whole-object reads, or conditional creates, and @@ -8704,6 +9530,10 @@ mod tests { fn object_path(&self, object_seq: u64) -> String { self.inner.object_path(object_seq) } + + fn storage(&self) -> String { + self.inner.storage() + } } type RangeReads = Arc>>; @@ -8851,6 +9681,10 @@ mod tests { fn object_path(&self, object_seq: u64) -> String { self.inner.object_path(object_seq) } + + fn storage(&self) -> String { + self.inner.storage() + } } /// Drives every outcome of a conditional create and every failure of the diff --git a/src/mito2/src/engine/object_store_wal_recovery_test.rs b/src/mito2/src/engine/object_store_wal_recovery_test.rs index 9ecfc998697..d5da18bf521 100644 --- a/src/mito2/src/engine/object_store_wal_recovery_test.rs +++ b/src/mito2/src/engine/object_store_wal_recovery_test.rs @@ -35,7 +35,9 @@ use common_wal::config::object_store::{ }; use common_wal::options::{ObjectStoreWalOptions, WAL_OPTIONS_KEY, WalOptions}; use log_store::ObjectStoreLogStore; -use log_store::metrics::METRIC_OBJECT_STORE_WAL_FAILING; +use log_store::metrics::{ + METRIC_OBJECT_STORE_WAL_FAILING, METRIC_OBJECT_STORE_WAL_MISSING_HISTORY_ACCEPTED_TOTAL, +}; use log_store::object_store_wal::{WalHole, entry_id}; use object_store::ObjectStore; use object_store::layers::mock::{self, MockLayerBuilder, oio}; @@ -47,8 +49,8 @@ use store_api::logstore::provider::Provider; use store_api::mito_engine_options::SKIP_WAL_KEY; use store_api::region_engine::{RegionEngine, RegionRole}; use store_api::region_request::{ - PathType, RegionDropRequest, RegionFlushRequest, RegionOpenRequest, RegionPutRequest, - RegionRequest, RegionTruncateRequest, + PathType, RegionCatchupRequest, RegionCloseRequest, RegionDropRequest, RegionFlushRequest, + RegionOpenRequest, RegionPutRequest, RegionRequest, RegionTruncateRequest, }; use store_api::storage::{RegionId, ScanRequest}; @@ -652,8 +654,72 @@ async fn test_open_region_rejects_mismatched_wal_prefix() { assert!(!engine.is_region_exists(REGION_A)); } +/// Opens a store under [`PREFIX`] like [`open_store`] that lets a region +/// open although the prefix does not hold its WAL history. +async fn open_store_accepting_missing_history( + object_store: &ObjectStore, +) -> Arc { + let config = ObjectStoreWalConfig { + accept_missing_history: true, + ..store_config(AckMode::Durable) + }; + ObjectStoreLogStore::try_new(object_store.clone(), &config, 0, STANDALONE_GENERATION) + .await + .unwrap() +} + +/// Copies every WAL object under [`PREFIX`] from `from` to the same key in +/// `to`, which is how the WAL moves to another location. +async fn copy_wal_objects(from: &ObjectStore, to: &ObjectStore) { + for object_seq in wal_object_seqs(from).await { + let path = format!("{PREFIX}/objects/{object_seq:020}.wal"); + let bytes = from.read(&path).await.unwrap().to_vec(); + to.write(&path, bytes).await.unwrap(); + } +} + +/// The request that opens a region under [`PREFIX`] and replays its WAL. +fn open_request(table_dir: &str) -> RegionOpenRequest { + RegionOpenRequest { + engine: String::new(), + table_dir: table_dir.to_string(), + options: wal_options(), + skip_wal_replay: false, + path_type: PathType::Bare, + checkpoint: None, + requirements: Default::default(), + } +} + +async fn close_region(engine: &MitoEngine, region_id: RegionId) { + engine + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) + .await + .unwrap(); +} + +/// Asserts that an open of `region_id` failed because the WAL does not hold +/// its history for `reason`, with the ways forward named. +fn assert_missing_history(err: &BoxedError, region_id: RegionId, reason: &str) { + assert_eq!(StatusCode::InvalidArguments, err.status_code()); + let message = err.output_msg(); + for expected in [ + ®ion_id.to_string(), + PREFIX, + "the default storage (memory", + reason, + "copy every object under the prefix", + "`[wal].accept_missing_history = true`", + ] { + assert!(message.contains(expected), "unexpected error: {message}"); + } +} + #[tokio::test] -async fn test_reopen_on_empty_prefix_without_durable_entries() { +async fn test_open_on_an_empty_prefix_requires_accepting_the_missing_history() { let mut env = TestEnv::with_prefix("object-store-wal-empty-prefix").await; let store = open_store(&memory_store(), AckMode::Durable).await; let engine = new_engine(&mut env, store.clone()).await; @@ -668,16 +734,52 @@ async fn test_reopen_on_empty_prefix_without_durable_entries() { drop(engine); drop(store); - // A fresh object store holds no object under the prefix. + // A fresh object store holds no object under the prefix. Region B keeps + // no history in the WAL and opens from its SST alone. Nothing tells + // region A, which never wrote, from a region whose unflushed entries + // were left behind, so it does not open, and the refused open deletes + // the start object it wrote: a second start, as an orchestrator that + // restarts a failed process makes, finds the prefix empty again and + // refuses region A the same way. + let new_location = memory_store(); + for _ in 0..2 { + let store = open_store(&new_location, AckMode::Durable).await; + let engine = new_engine(&mut env, store.clone()).await; + open_region(&engine, REGION_B, &table_dir_b, &[(SKIP_WAL_KEY, "true")]) + .await + .unwrap(); + assert_eq!(rows_b, scan_rows(&engine, REGION_B).await); + let err = open_region(&engine, REGION_A, &table_dir_a, &[]) + .await + .unwrap_err(); + assert_missing_history( + &err, + REGION_A, + "the prefix held no object when the WAL opened", + ); + assert!(!engine.is_region_exists(REGION_A)); + assert!(wal_object_seqs(&new_location).await.is_empty()); + engine.stop().await.unwrap(); + drop(engine); + drop(store); + } + + // Accepting the missing history opens region A without it, counted + // once; region B is not checked. let object_store = memory_store(); - let store = open_store(&object_store, AckMode::Durable).await; + let store = open_store_accepting_missing_history(&object_store).await; let engine = new_engine(&mut env, store.clone()).await; + let accepted = METRIC_OBJECT_STORE_WAL_MISSING_HISTORY_ACCEPTED_TOTAL.get(); open_region(&engine, REGION_A, &table_dir_a, &[]) .await .unwrap(); open_region(&engine, REGION_B, &table_dir_b, &[(SKIP_WAL_KEY, "true")]) .await .unwrap(); + assert_eq!( + accepted + 1, + METRIC_OBJECT_STORE_WAL_MISSING_HISTORY_ACCEPTED_TOTAL.get() + ); assert_eq!(0, latest(&store, REGION_A)); assert_eq!(0, latest(&store, REGION_B)); @@ -711,6 +813,301 @@ async fn test_reopen_on_empty_prefix_without_durable_entries() { assert_eq!(2, engine.get_region_statistic(REGION_A).unwrap().num_rows); } +#[tokio::test] +async fn test_create_after_a_withdrawal_leaves_no_manifest() { + let mut env = TestEnv::with_prefix("object-store-wal-create-after-withdrawal").await; + let store = open_store(&memory_store(), AckMode::Durable).await; + let engine = new_engine(&mut env, store.clone()).await; + let (table_dir_a, _) = create_region(&engine, REGION_A, &[]).await; + engine.stop().await.unwrap(); + drop(engine); + drop(store); + + // On a new location the refused open of region A gives up the start + // object, and a create of region B after it fails before it writes a + // manifest. + let new_location = memory_store(); + let store = open_store(&new_location, AckMode::Durable).await; + let engine = new_engine(&mut env, store.clone()).await; + open_region(&engine, REGION_A, &table_dir_a, &[]) + .await + .unwrap_err(); + let request = CreateRequestBuilder::new() + .insert_option(WAL_OPTIONS_KEY, &wal_options()[WAL_OPTIONS_KEY]) + .build(); + let err = engine + .handle_request(REGION_B, RegionRequest::Create(request.clone())) + .await + .unwrap_err(); + assert!( + format!("{err:?}").contains("gave up its start object"), + "unexpected error: {err:?}" + ); + assert!(!engine.is_region_exists(REGION_B)); + engine.stop().await.unwrap(); + drop(engine); + drop(store); + + // The next start creates region B, which a manifest left behind would + // have turned into an existing region without its history. + let store = open_store(&new_location, AckMode::Durable).await; + let engine = new_engine(&mut env, store.clone()).await; + engine + .handle_request(REGION_B, RegionRequest::Create(request)) + .await + .unwrap(); + assert!(engine.is_region_exists(REGION_B)); + engine.stop().await.unwrap(); +} + +#[tokio::test] +async fn test_reopen_a_region_created_on_an_empty_prefix_without_writes() { + let mut env = TestEnv::with_prefix("object-store-wal-created-region").await; + let store = open_store(&memory_store(), AckMode::Durable).await; + let engine = new_engine(&mut env, store.clone()).await; + let (table_dir, _) = create_region(&engine, REGION_A, &[]).await; + + // The region's history starts in this store, so it opens again on it, + // directly and through a catch-up, although it holds no entry. + close_region(&engine, REGION_A).await; + open_region(&engine, REGION_A, &table_dir, &[]) + .await + .unwrap(); + close_region(&engine, REGION_A).await; + engine + .handle_request( + REGION_A, + RegionRequest::Open(RegionOpenRequest { + skip_wal_replay: true, + ..open_request(&table_dir) + }), + ) + .await + .unwrap(); + engine + .handle_request( + REGION_A, + RegionRequest::Catchup(RegionCatchupRequest { + set_writable: true, + ..Default::default() + }), + ) + .await + .unwrap(); + assert!(region(&engine, REGION_A).is_writable()); + assert_eq!( + EMPTY_RECOVERY_STATE, + region_recovery_state(&engine, REGION_A).await + ); +} + +#[tokio::test] +async fn test_moving_the_wal_keeps_the_history_only_with_its_objects() { + let mut env = TestEnv::with_prefix("object-store-wal-move").await; + let old_location = memory_store(); + let store = open_store(&old_location, AckMode::Durable).await; + let engine = new_engine(&mut env, store.clone()).await; + let (table_dir, schema) = create_region(&engine, REGION_A, &[]).await; + // Object 1 holds the region's rows, which are never flushed. + let mut writer = SealedWriter::new(&store); + writer + .put_and_seal(&engine, vec![(REGION_A, rows(&schema, 0, 10))]) + .await; + let rows_before = scan_rows(&engine, REGION_A).await; + engine.stop().await.unwrap(); + drop(engine); + drop(writer); + drop(store); + + // On a location without the objects the region opens neither directly + // nor through a catch-up after an open that skipped the replay. + let store = open_store(&memory_store(), AckMode::Durable).await; + let engine = new_engine(&mut env, store.clone()).await; + let reason = "the prefix held no object when the WAL opened"; + let err = open_region(&engine, REGION_A, &table_dir, &[]) + .await + .unwrap_err(); + assert_missing_history(&err, REGION_A, reason); + engine + .handle_request( + REGION_A, + RegionRequest::Open(RegionOpenRequest { + skip_wal_replay: true, + ..open_request(&table_dir) + }), + ) + .await + .unwrap(); + let err = engine + .handle_request( + REGION_A, + RegionRequest::Catchup(RegionCatchupRequest { + set_writable: true, + ..Default::default() + }), + ) + .await + .unwrap_err(); + assert_missing_history(&err, REGION_A, reason); + assert!(!region(&engine, REGION_A).is_writable()); + // A repeated create of the existing region fails the same way rather + // than creating the region again over its manifest. + close_region(&engine, REGION_A).await; + let request = CreateRequestBuilder::new() + .insert_option(WAL_OPTIONS_KEY, &wal_options()[WAL_OPTIONS_KEY]) + .build(); + assert_eq!(table_dir, request.table_dir); + let err = engine + .handle_request(REGION_A, RegionRequest::Create(request)) + .await + .unwrap_err(); + assert_missing_history(&err, REGION_A, reason); + assert!(!engine.is_region_exists(REGION_A)); + engine.stop().await.unwrap(); + drop(engine); + drop(store); + + // Moving the WAL means copying every object: the region replays all of + // its rows on the new location. + let new_location = memory_store(); + copy_wal_objects(&old_location, &new_location).await; + let store = open_store(&new_location, AckMode::Durable).await; + let engine = new_engine(&mut env, store.clone()).await; + open_region(&engine, REGION_A, &table_dir, &[]) + .await + .unwrap(); + assert_eq!(rows_before, scan_rows(&engine, REGION_A).await); + assert_eq!(10, engine.get_region_statistic(REGION_A).unwrap().num_rows); +} + +#[tokio::test] +async fn test_open_region_fails_when_its_watermark_is_above_the_wal_objects() { + let mut env = TestEnv::with_prefix("object-store-wal-watermark-above").await; + let store = open_store(&memory_store(), AckMode::Durable).await; + let engine = new_engine(&mut env, store.clone()).await; + let (table_dir_a, schema_a) = create_region(&engine, REGION_A, &[]).await; + let (table_dir_b, schema_b) = create_region(&engine, REGION_B, &[]).await; + // Region A flushes its entry of object 2, region B is truncated after its + // entry of object 3. + let mut writer = SealedWriter::new(&store); + writer + .put_and_seal(&engine, vec![(REGION_A, rows(&schema_a, 0, 2))]) + .await; + writer + .put_and_seal(&engine, vec![(REGION_A, rows(&schema_a, 2, 4))]) + .await; + flush_region(&engine, REGION_A, None).await; + writer + .put_and_seal(&engine, vec![(REGION_B, rows(&schema_b, 0, 2))]) + .await; + engine + .handle_request( + REGION_B, + RegionRequest::Truncate(RegionTruncateRequest::All), + ) + .await + .unwrap(); + let manifest_a = region(&engine, REGION_A).manifest_ctx.manifest().await; + assert_eq!(entry_id(2, 1), manifest_a.flushed_entry_id); + let manifest_b = region(&engine, REGION_B).manifest_ctx.manifest().await; + assert_eq!(Some(entry_id(3, 1)), manifest_b.truncated_entry_id); + engine.stop().await.unwrap(); + drop(engine); + drop(writer); + drop(store); + + // The other location holds the start object 0 and object 1 of region C. + let other_location = memory_store(); + let store = open_store(&other_location, AckMode::Durable).await; + let engine = new_engine(&mut env, store.clone()).await; + let region_c = RegionId::new(1, 3); + let (table_dir_c, schema_c) = create_region(&engine, region_c, &[]).await; + let mut writer = SealedWriter::new(&store); + writer + .put_and_seal(&engine, vec![(region_c, rows(&schema_c, 0, 3))]) + .await; + engine.stop().await.unwrap(); + drop(engine); + drop(writer); + drop(store); + + // Its prefix held entries only in objects below 2, so neither watermark + // can name one of its objects, while region C, whose history it holds, + // opens. Every start adds a start object, which holds no entry, so a + // second start after the refused one refuses the regions as well. + for _ in 0..2 { + let store = open_store(&other_location, AckMode::Durable).await; + let engine = new_engine(&mut env, store.clone()).await; + for (region_id, table_dir, watermark) in [ + (REGION_A, &table_dir_a, entry_id(2, 1)), + (REGION_B, &table_dir_b, entry_id(3, 1)), + ] { + let err = open_region(&engine, region_id, table_dir, &[]) + .await + .unwrap_err(); + let reason = format!( + "its watermark, entry {watermark}, names object {}, but the prefix held entries only in objects below 2 when the WAL opened", + watermark >> 20 + ); + assert_missing_history(&err, region_id, &reason); + assert!(!engine.is_region_exists(region_id)); + } + open_region(&engine, region_c, &table_dir_c, &[]) + .await + .unwrap(); + assert_eq!(3, engine.get_region_statistic(region_c).unwrap().num_rows); + engine.stop().await.unwrap(); + drop(engine); + drop(store); + } +} + +#[tokio::test] +async fn test_reopen_after_collection_left_only_the_newest_object() { + let mut env = TestEnv::with_prefix("object-store-wal-newest-object").await; + let object_store = memory_store(); + let store = open_store(&object_store, AckMode::Durable).await; + let engine = new_engine(&mut env, store.clone()).await; + let (table_dir_a, schema_a) = create_region(&engine, REGION_A, &[]).await; + let (table_dir_b, schema_b) = create_region(&engine, REGION_B, &[]).await; + // Region A flushes object 1; object 2 of region B carries the boundary + // past it, so once region B flushes too only object 2 is left. + let mut writer = SealedWriter::new(&store); + writer + .put_and_seal(&engine, vec![(REGION_A, rows(&schema_a, 0, 2))]) + .await; + flush_region(&engine, REGION_A, None).await; + writer + .put_and_seal(&engine, vec![(REGION_B, rows(&schema_b, 0, 3))]) + .await; + flush_region(&engine, REGION_B, None).await; + wait_for_collection(&store).await; + assert_eq!(vec![2], wal_object_seqs(&object_store).await); + engine.stop().await.unwrap(); + drop(engine); + drop(writer); + drop(store); + + // The object region A flushed is gone, but the newest object stays and + // keeps the next sequence above it, so both regions open. + let store = open_store(&object_store, AckMode::Durable).await; + let engine = new_engine(&mut env, store.clone()).await; + open_region(&engine, REGION_A, &table_dir_a, &[]) + .await + .unwrap(); + open_region(&engine, REGION_B, &table_dir_b, &[]) + .await + .unwrap(); + assert_eq!( + entry_id(1, 1), + region_recovery_state(&engine, REGION_A) + .await + .flushed_entry_id + ); + assert_eq!(2, engine.get_region_statistic(REGION_A).unwrap().num_rows); + assert_eq!(3, engine.get_region_statistic(REGION_B).unwrap().num_rows); +} + #[rstest] #[case(AckMode::Durable)] #[case(AckMode::Enqueued)] diff --git a/src/mito2/src/error.rs b/src/mito2/src/error.rs index 54bf1ba6b6e..ab781f790f4 100644 --- a/src/mito2/src/error.rs +++ b/src/mito2/src/error.rs @@ -461,6 +461,14 @@ pub enum Error { source: BoxedError, }, + #[snafu(display("Failed to check the WAL history, region_id: {}", region_id))] + CheckWalHistory { + region_id: RegionId, + #[snafu(implicit)] + location: Location, + source: BoxedError, + }, + // Shared error for each writer in the write group. #[snafu(display("Failed to write region"))] WriteGroup { source: Arc }, @@ -1465,7 +1473,8 @@ impl ErrorExt for Error { WriteWal { source, .. } | ReadWal { source, .. } | DeleteWal { source, .. } - | WaitWalDurable { source, .. } => source.status_code(), + | WaitWalDurable { source, .. } + | CheckWalHistory { source, .. } => source.status_code(), CompressObject { .. } | DecompressObject { .. } | SerdeJson { .. } @@ -1684,6 +1693,7 @@ impl ErrorExt for Error { | ReadWal { source, .. } | DeleteWal { source, .. } | WaitWalDurable { source, .. } + | CheckWalHistory { source, .. } | FetchManifests { source, .. } | External { source, .. } => source.retry_hint(), diff --git a/src/mito2/src/region/catchup.rs b/src/mito2/src/region/catchup.rs index e9f7f7ef43c..4b8f1288c4b 100644 --- a/src/mito2/src/region/catchup.rs +++ b/src/mito2/src/region/catchup.rs @@ -21,7 +21,7 @@ use store_api::logstore::LogStore; use crate::error::{self, Result}; use crate::region::MitoRegion; -use crate::region::opener::replay_memtable; +use crate::region::opener::{check_wal_history, replay_memtable}; use crate::wal::Wal; use crate::wal::entry_distributor::WalEntryReceiver; @@ -89,6 +89,12 @@ impl RegionCatchupTask { "Trying to replay memtable for region: {region_id}, provider: {:?}, replay from entry id: {replay_from_entry_id}, flushed entry id: {flushed_entry_id}", self.region.provider ); + check_wal_history( + &self.wal, + &self.region.provider, + &self.region.version_control.current().version, + ) + .await?; let timer = Instant::now(); let wal_entry_reader = self .entry_receiver diff --git a/src/mito2/src/region/opener.rs b/src/mito2/src/region/opener.rs index 37c91946db3..35ca3129fe0 100644 --- a/src/mito2/src/region/opener.rs +++ b/src/mito2/src/region/opener.rs @@ -63,7 +63,7 @@ use crate::memtable::{MemtableBuilderProvider, ensure_json2_not_use_time_series_ use crate::metrics::{CACHE_FILL_DOWNLOADED_FILES, CACHE_FILL_PENDING_FILES}; use crate::read::series_candidate::is_sparse_metric_metadata; use crate::region::options::RegionOptions; -use crate::region::version::{VersionBuilder, VersionControl, VersionControlRef}; +use crate::region::version::{Version, VersionBuilder, VersionControl, VersionControlRef}; use crate::region::{ ManifestContext, ManifestStats, MitoRegion, MitoRegionRef, RegionLeaderState, RegionRoleState, RegionStats, @@ -360,6 +360,9 @@ impl RegionOpener { region_dir, self.region_id ); } + // The region exists but the WAL does not hold its history: + // creating it again would replace its manifest. + Err(e @ error::Error::CheckWalHistory { .. }) => return Err(e), Err(e) => { warn!(e; "Failed to open region {} before creating it, region_dir: {}", @@ -385,6 +388,11 @@ impl RegionOpener { }; ensure_json2_not_use_time_series_memtable(&metadata, &options)?; options.validate_column_compression(&metadata)?; + // The region's WAL history starts in this log store, so reopening + // the region on it passes the history check without any entry. It is + // recorded before the manifest exists: a log store that gave up its + // history meanwhile fails the create before it leaves a manifest. + wal.start_region_history(region_id, &provider)?; // Create a manifest manager for this region and writes regions to the manifest file. let mut region_manifest_options = RegionManifestOptions::new(config, ®ion_dir, &object_store); @@ -502,6 +510,8 @@ impl RegionOpener { } /// Tries to open the region and returns `None` if the region directory is empty. + /// A region it replays must find its WAL history in the log store, see + /// [`check_wal_history`]. async fn maybe_open( &mut self, config: &MitoConfig, @@ -605,6 +615,7 @@ impl RegionOpener { .unwrap_or_default() .max(flushed_entry_id); let topic_latest_entry_id = if !self.skip_wal_replay { + check_wal_history(wal, &provider, &version_control.current().version).await?; info!( "Start replaying memtable at replay_from_entry_id: {} for region {}, manifest version: {}, flushed entry id: {}, elapsed: {:?}", replay_from_entry_id, @@ -909,6 +920,25 @@ pub(crate) fn check_recovered_region( Ok(()) } +/// Checks that the log store holds the WAL history of a region that writes +/// to the WAL through the largest replay watermark `version` persisted, +/// before the region is replayed. A region that skips the WAL keeps no +/// history there. +pub(crate) async fn check_wal_history( + wal: &Wal, + provider: &Provider, + version: &Version, +) -> Result<()> { + if version.options.skip_wal { + return Ok(()); + } + let watermark = version + .flushed_entry_id + .max(version.truncated_entry_id.unwrap_or_default()); + wal.check_region_history(version.metadata.region_id, provider, watermark) + .await +} + /// Replays the mutations from WAL and inserts mutations to memtable of given region. pub(crate) async fn replay_memtable( provider: &Provider, diff --git a/src/mito2/src/wal.rs b/src/mito2/src/wal.rs index 9d00a76d87f..bae0556a4fa 100644 --- a/src/mito2/src/wal.rs +++ b/src/mito2/src/wal.rs @@ -36,7 +36,10 @@ use store_api::logstore::provider::Provider; use store_api::logstore::{AppendBatchResponse, LogStore}; use store_api::storage::RegionId; -use crate::error::{BuildEntrySnafu, DeleteWalSnafu, Result, WaitWalDurableSnafu, WriteWalSnafu}; +use crate::error::{ + BuildEntrySnafu, CheckWalHistorySnafu, DeleteWalSnafu, Result, WaitWalDurableSnafu, + WriteWalSnafu, +}; use crate::wal::entry_reader::{LogStoreEntryReader, WalEntryReader}; use crate::wal::raw_entry_reader::{LogStoreRawEntryReader, RegionRawEntryReader}; @@ -158,6 +161,41 @@ impl Wal { })) } + /// Checks that the log store holds the WAL history of the region through + /// `entry_id`, the largest replay watermark the region persisted, before + /// the region is replayed from its persisted state. + pub(crate) async fn check_region_history( + &self, + region_id: RegionId, + provider: &Provider, + entry_id: EntryId, + ) -> Result<()> { + if let Provider::Noop = provider { + return Ok(()); + } + self.store + .check_region_history(provider, entry_id) + .await + .map_err(BoxedError::new) + .context(CheckWalHistorySnafu { region_id }) + } + + /// Records that the WAL history of the region starts in the log store, + /// as it does for a region created on it. + pub(crate) fn start_region_history( + &self, + region_id: RegionId, + provider: &Provider, + ) -> Result<()> { + if let Provider::Noop = provider { + return Ok(()); + } + self.store + .start_region_history(provider) + .map_err(BoxedError::new) + .context(CheckWalHistorySnafu { region_id }) + } + /// Returns a [WalEntryReader] pub(crate) fn wal_entry_reader( &self, diff --git a/src/store-api/src/logstore.rs b/src/store-api/src/logstore.rs index ef0fbe49f01..88abd9809a5 100644 --- a/src/store-api/src/logstore.rs +++ b/src/store-api/src/logstore.rs @@ -122,6 +122,32 @@ pub trait LogStore: Send + Sync + 'static + std::fmt::Debug { ) -> Result<(), Self::Error> { Ok(()) } + + /// Checks, before an existing region is replayed from its persisted + /// state, that the log store holds the region's WAL history through + /// `entry_id`, the largest replay watermark the region persisted, so that + /// an open against a WAL that lost that history fails instead of + /// silently replaying nothing. + /// + /// The default accepts, which keeps the behaviour of every log store that + /// does not override it. + async fn check_region_history( + &self, + _provider: &Provider, + _entry_id: EntryId, + ) -> Result<(), Self::Error> { + Ok(()) + } + + /// Records that the WAL history of the provider's region starts in this + /// log store, as it does for a region created on it, so that + /// [`check_region_history`](Self::check_region_history) accepts a later + /// open of the region by the same log store. + /// + /// The default does nothing. + fn start_region_history(&self, _provider: &Provider) -> Result<(), Self::Error> { + Ok(()) + } } /// The response of an `append` operation. From 4859e57e0bec44a871168be0f364fee262ddd827 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Thu, 15 Oct 2026 14:29:10 +0000 Subject: [PATCH 017/239] feat(meta): apply affinity to failover and repartition (#9604) * feat(meta): apply affinity to failover and repartition Signed-off-by: WenyXu * refactor(meta): share repartition affinity until group allocation Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- src/cmd/src/standalone.rs | 1 + src/common/meta/src/ddl.rs | 4 +- src/common/meta/src/ddl_manager.rs | 2 + src/common/meta/src/test_util.rs | 1 + src/meta-srv/AGENTS.md | 10 +- src/meta-srv/src/metasrv/builder.rs | 11 + src/meta-srv/src/procedure/repartition.rs | 5 + .../procedure/repartition/allocate_region.rs | 576 +++++++++++++++++- src/meta-srv/src/procedure/utils.rs | 1 + src/meta-srv/src/region/supervisor.rs | 318 +++++++++- src/meta-srv/src/selector.rs | 1 + tests-integration/src/standalone.rs | 1 + 12 files changed, 911 insertions(+), 20 deletions(-) diff --git a/src/cmd/src/standalone.rs b/src/cmd/src/standalone.rs index 278afb7f1da..3dca89f6c25 100644 --- a/src/cmd/src/standalone.rs +++ b/src/cmd/src/standalone.rs @@ -682,6 +682,7 @@ impl StartCommand { let ddl_context = DdlContext { affinity_resolver: None, + affinity_fallback: Default::default(), peer_alloc_context_provider: None, node_manager: node_manager.clone(), cache_invalidator: layered_cache_registry.clone(), diff --git a/src/common/meta/src/ddl.rs b/src/common/meta/src/ddl.rs index 3eea16fcfa5..82b43cd832c 100644 --- a/src/common/meta/src/ddl.rs +++ b/src/common/meta/src/ddl.rs @@ -18,7 +18,7 @@ use std::time::Duration; use store_api::storage::{RegionId, TableId}; use crate::DatanodeId; -use crate::affinity::{AffinityResolverRef, PeerAllocContextProviderRef}; +use crate::affinity::{AffinityFallback, AffinityResolverRef, PeerAllocContextProviderRef}; use crate::cache_invalidator::CacheInvalidatorRef; use crate::ddl::flow_meta::FlowMetadataAllocatorRef; use crate::ddl::table_meta::TableMetadataAllocatorRef; @@ -108,6 +108,8 @@ impl RegionFailureDetectorController for NoopRegionFailureDetectorControl { pub struct DdlContext { /// Optional resolver for automatic allocation affinity. pub affinity_resolver: Option, + /// Fallback policy for allocation without query-specific inputs. + pub affinity_fallback: AffinityFallback, /// Builds allocation inputs from the initiating query. pub peer_alloc_context_provider: Option, /// Sends querying and requests to nodes. diff --git a/src/common/meta/src/ddl_manager.rs b/src/common/meta/src/ddl_manager.rs index 0a1abde30f7..7b490895a4b 100644 --- a/src/common/meta/src/ddl_manager.rs +++ b/src/common/meta/src/ddl_manager.rs @@ -1886,6 +1886,7 @@ mod tests { let ddl_manager = DdlManager::new( DdlContext { affinity_resolver: None, + affinity_fallback: Default::default(), peer_alloc_context_provider: None, node_manager: Arc::new(DummyDatanodeManager), cache_invalidator: Arc::new(DummyCacheInvalidator), @@ -1954,6 +1955,7 @@ mod tests { DdlManager::new( DdlContext { affinity_resolver: None, + affinity_fallback: Default::default(), peer_alloc_context_provider: None, node_manager: Arc::new(DummyDatanodeManager), cache_invalidator: Arc::new(DummyCacheInvalidator), diff --git a/src/common/meta/src/test_util.rs b/src/common/meta/src/test_util.rs index e574161d367..c1c3af8f5e9 100644 --- a/src/common/meta/src/test_util.rs +++ b/src/common/meta/src/test_util.rs @@ -199,6 +199,7 @@ pub fn new_ddl_context_with_kv_backend( ))); DdlContext { affinity_resolver: None, + affinity_fallback: Default::default(), peer_alloc_context_provider: None, node_manager, cache_invalidator: Arc::new(DummyCacheInvalidator), diff --git a/src/meta-srv/AGENTS.md b/src/meta-srv/AGENTS.md index 863b2a2aab4..f1153b4cc35 100644 --- a/src/meta-srv/AGENTS.md +++ b/src/meta-srv/AGENTS.md @@ -66,10 +66,12 @@ API under `src/meta-srv/src/service/admin/`. `cache_invalidator.rs` and the heartbeat publish handler. - **Affinity**: `common-meta/affinity.rs` defines requirements, matching, batch resolution, and the query allocation-context provider. The selector factory - optionally supplies a resolver for allocation inputs. DDL configurators can - install a query provider before procedure loaders are registered. Each selector - allows one unconstrained retry on candidate shortage only when explicitly - configured. Procedure state is unchanged. + optionally supplies a resolver shared by DDL allocation and the region + supervisor; without one, allocation skips source lookup. DDL configurators can + install a query provider before procedure loaders are registered. Repartition + and failover resolve source requirements before selection. Each selector allows + one unconstrained retry on candidate shortage only when explicitly configured; + metadata errors propagate. Procedure state is unchanged. - **Selector stats fields**: keep `selector/` implementations in sync with the stats they consume. diff --git a/src/meta-srv/src/metasrv/builder.rs b/src/meta-srv/src/metasrv/builder.rs index 96cd2a8adb5..93ce228a156 100644 --- a/src/meta-srv/src/metasrv/builder.rs +++ b/src/meta-srv/src/metasrv/builder.rs @@ -423,6 +423,11 @@ impl MetasrvBuilder { leader_cached_kv_backend.clone(), ) .with_state(state.clone()); + let region_supervisor = if let Some(resolver) = &affinity_resolver { + region_supervisor.with_affinity_resolver(resolver.clone()) + } else { + region_supervisor + }; Some(RegionFailureHandler::new( region_supervisor, @@ -437,6 +442,7 @@ impl MetasrvBuilder { let ddl_context = DdlContext { peer_alloc_context_provider: None, + affinity_fallback: selector_ctx.affinity_fallback, affinity_resolver, node_manager: node_manager.clone(), cache_invalidator: cache_invalidator.clone(), @@ -817,6 +823,7 @@ mod tests { .await .unwrap(); let context = metasrv.ddl_manager().create_context(); + assert_eq!(context.affinity_fallback, AffinityFallback::Error); assert!(context.affinity_resolver.is_none()); assert!(context.peer_alloc_context_provider.is_none()); } @@ -894,6 +901,10 @@ mod tests { metasrv.selector_ctx().affinity_fallback, AffinityFallback::Unconstrained ); + assert_eq!( + metasrv.ddl_manager().create_context().affinity_fallback, + AffinityFallback::Unconstrained + ); assert_eq!(calls.load(Ordering::Relaxed), 1); assert!(Arc::ptr_eq(&metasrv.selector_ctx().affinity, &affinity)); assert!(Arc::ptr_eq( diff --git a/src/meta-srv/src/procedure/repartition.rs b/src/meta-srv/src/procedure/repartition.rs index 008c0ce05bb..b137eac2166 100644 --- a/src/meta-srv/src/procedure/repartition.rs +++ b/src/meta-srv/src/procedure/repartition.rs @@ -30,6 +30,7 @@ use std::fmt::{Debug, Display}; use std::time::{Duration, Instant}; use common_error::ext::BoxedError; +use common_meta::affinity::{AffinityFallback, AffinityResolverRef}; use common_meta::cache_invalidator::CacheInvalidatorRef; use common_meta::ddl::DdlContext; use common_meta::ddl::allocator::region_routes::RegionRoutesAllocatorRef; @@ -148,6 +149,8 @@ impl PersistentContext { #[derive(Clone)] pub struct Context { + pub affinity_resolver: Option, + pub affinity_fallback: AffinityFallback, pub persistent_ctx: PersistentContext, pub volatile_ctx: VolatileContext, pub table_metadata_manager: TableMetadataManagerRef, @@ -248,6 +251,8 @@ impl Context { persistent_ctx: PersistentContext, ) -> Self { Self { + affinity_resolver: ddl_ctx.affinity_resolver.clone(), + affinity_fallback: ddl_ctx.affinity_fallback, persistent_ctx, table_metadata_manager: ddl_ctx.table_metadata_manager.clone(), memory_region_keeper: ddl_ctx.memory_region_keeper.clone(), diff --git a/src/meta-srv/src/procedure/repartition/allocate_region.rs b/src/meta-srv/src/procedure/repartition/allocate_region.rs index 68000682142..edbae33c17d 100644 --- a/src/meta-srv/src/procedure/repartition/allocate_region.rs +++ b/src/meta-srv/src/procedure/repartition/allocate_region.rs @@ -15,6 +15,9 @@ use std::any::Any; use std::collections::{HashMap, HashSet}; +use common_meta::affinity::{ + AffinityFallback, AffinityReference, AffinityRequirements, AffinityResolver, +}; use common_meta::ddl::create_table::executor::CreateTableExecutor; use common_meta::ddl::create_table::template::{ CreateRequestBuilder, build_template_from_raw_table_info_for_physical_table, @@ -155,18 +158,37 @@ impl ExecutePlan { let region_number_and_partition_exprs = AllocateRegion::prepare_region_allocation_data(&allocate_regions)?; let table_info_value = ctx.get_table_info_value().await?; - let new_allocated_region_routes = ctx - .region_routes_allocator - .allocate( - table_id, - ®ion_number_and_partition_exprs + let affinity = if let Some(resolver) = &ctx.affinity_resolver { + Some(resolve_allocation_affinity(ctx, resolver.as_ref()).await?) + } else { + None + }; + let groups = if let Some(affinity) = &affinity { + group_allocations(table_id, ®ion_number_and_partition_exprs, affinity) + } else { + vec![AllocationGroup { + affinity: None, + regions: region_number_and_partition_exprs .iter() - .map(|(n, p)| (*n, p.as_str())) - .collect::>(), - &PeerAllocContext::default(), - ) - .await - .context(error::AllocateRegionRoutesSnafu { table_id })?; + .map(|(number, partition)| (*number, partition.as_str())) + .collect(), + }] + }; + let mut new_allocated_region_routes = Vec::with_capacity(allocate_regions.len()); + for AllocationGroup { affinity, regions } in groups { + let routes = ctx + .region_routes_allocator + .allocate( + table_id, + ®ions, + &PeerAllocContext { + affinity: affinity.cloned(), + }, + ) + .await + .context(error::AllocateRegionRoutesSnafu { table_id })?; + new_allocated_region_routes.extend(routes); + } let skip_wal = table_info_value.table_info.meta.options.skip_wal; let allocate_noop = if skip_wal { let table_route_value = ctx.get_table_route_value().await?; @@ -439,8 +461,146 @@ impl AllocateRegion { } } +/// Selects the lowest-ID contributing source for a target spanning multiple sources. +/// Legacy plans without a transition map fall back to the lowest-ID source in the group. +fn collect_allocation_sources( + plans: &[RepartitionPlanEntry], + routes: &[RegionRoute], +) -> Vec<(RegionId, Option)> { + let leaders: HashMap<_, _> = routes + .iter() + .filter_map(|route| { + route + .leader_peer + .as_ref() + .map(|peer| (route.region.id, peer.id)) + }) + .collect(); + let mut sources = Vec::new(); + for plan in plans { + let allocated: HashSet<_> = plan.allocated_region_ids.iter().copied().collect(); + let mut target_sources = vec![None; plan.target_regions.len()]; + if plan.transition_map.is_empty() { + let source = plan + .source_regions + .iter() + .map(|source| source.region_id()) + .min(); + target_sources.fill(source); + } else { + for (source, targets) in plan.source_regions.iter().zip(&plan.transition_map) { + for target_index in targets { + if let Some(target_source) = target_sources.get_mut(*target_index) { + let source_id = source.region_id(); + *target_source = + Some(target_source.map_or(source_id, |current| current.min(source_id))); + } + } + } + } + for (target, source) in plan.target_regions.iter().zip(target_sources) { + if allocated.contains(&target.region_id) { + sources.push(( + target.region_id, + source.and_then(|source| leaders.get(&source)).copied(), + )); + } + } + } + sources +} + +/// Owns requirements once per source; targets refer to them by index. +#[derive(Debug)] +struct AllocationAffinity { + requirements: Vec>, + target_sources: HashMap, +} + +async fn resolve_allocation_affinity( + ctx: &Context, + resolver: &dyn AffinityResolver, +) -> Result { + let table_id = ctx.persistent_ctx.table_id; + let routes = ctx.get_table_route_value().await?; + let sources = + collect_allocation_sources(&ctx.persistent_ctx.plans, routes.region_routes().unwrap()); + let mut source_indices = HashMap::new(); + let mut references = Vec::new(); + for (_, source_id) in &sources { + if let Some(source_id) = source_id { + source_indices.entry(*source_id).or_insert_with(|| { + let index = references.len(); + references.push(AffinityReference::Datanode(*source_id)); + index + }); + } + } + let resolved = resolver + .resolve(&references) + .await + .context(error::AllocateRegionRoutesSnafu { table_id })?; + snafu::ensure!( + resolved.len() == source_indices.len(), + error::UnexpectedSnafu { + violated: "affinity resolution must preserve the number and order of references", + } + ); + snafu::ensure!( + ctx.affinity_fallback == AffinityFallback::Unconstrained + || sources.iter().all(|(_, source)| { + source.is_some_and(|id| resolved[source_indices[&id]].is_some()) + }), + error::UnexpectedSnafu { + violated: "source affinity requirements are unavailable", + } + ); + Ok(AllocationAffinity { + requirements: resolved, + target_sources: sources + .into_iter() + .filter_map(|(region_id, source_id)| { + source_id.map(|source_id| (region_id, source_indices[&source_id])) + }) + .collect(), + }) +} + +#[derive(Debug, PartialEq, Eq)] +struct AllocationGroup<'a> { + affinity: Option<&'a AffinityRequirements>, + regions: Vec<(RegionNumber, &'a str)>, +} + +fn group_allocations<'a>( + table_id: TableId, + regions: &'a [(RegionNumber, String)], + affinity: &'a AllocationAffinity, +) -> Vec> { + let mut groups: Vec> = Vec::new(); + for (number, partition) in regions { + let requirements = affinity + .target_sources + .get(&RegionId::new(table_id, *number)) + .and_then(|index| affinity.requirements[*index].as_ref()); + if let Some(group) = groups + .iter_mut() + .find(|group| group.affinity == requirements) + { + group.regions.push((*number, partition.as_str())); + } else { + groups.push(AllocationGroup { + affinity: requirements, + regions: vec![(*number, partition.as_str())], + }); + } + } + groups +} + #[cfg(test)] mod tests { + use std::ptr; use std::sync::Arc; use api::v1::region::region_request::Body; @@ -1035,4 +1195,398 @@ mod tests { AllocateRegion::Build(_) => panic!("expected execute plan"), } } + #[test] + fn test_collect_allocation_sources_use_target_overlap() { + let table_id = 1024; + let mut plan = RepartitionPlanEntry { + group_id: Uuid::nil(), + source_regions: vec![ + create_region_descriptor(table_id, 2, "x", 100, 200), + create_region_descriptor(table_id, 1, "x", 0, 100), + ], + target_regions: vec![ + create_target_region_descriptor(table_id, 3, "x", 0, 50), + create_target_region_descriptor(table_id, 4, "x", 50, 150), + create_target_region_descriptor(table_id, 5, "x", 150, 200), + ], + allocated_region_ids: (3..=5).map(|n| RegionId::new(table_id, n)).collect(), + pending_deallocate_region_ids: vec![], + transition_map: vec![vec![1, 2], vec![0, 1]], + original_target_routes: vec![], + }; + let mut routes = create_current_region_routes(table_id, &[1, 2]); + routes[0].leader_peer = Some(Peer::empty(10)); + routes[1].leader_peer = Some(Peer::empty(20)); + assert_eq!( + collect_allocation_sources(&[plan.clone()], &routes), + vec![ + (RegionId::new(table_id, 3), Some(10)), + (RegionId::new(table_id, 4), Some(10)), + (RegionId::new(table_id, 5), Some(20)), + ] + ); + // Source enumeration does not change the tie-break for overlapping targets. + plan.source_regions.reverse(); + plan.transition_map.reverse(); + assert_eq!( + collect_allocation_sources(std::slice::from_ref(&plan), &routes)[1], + (RegionId::new(table_id, 4), Some(10)) + ); + plan.transition_map.clear(); + assert_eq!( + collect_allocation_sources(std::slice::from_ref(&plan), &routes), + (3..=5) + .map(|number| (RegionId::new(table_id, number), Some(10))) + .collect::>() + ); + plan.source_regions.clear(); + assert_eq!( + collect_allocation_sources(&[plan], &routes), + (3..=5) + .map(|number| (RegionId::new(table_id, number), None)) + .collect::>() + ); + } + + #[test] + fn test_group_allocations_preserves_region_partition_mapping() { + use common_meta::affinity::AffinityRequirements; + let table_id = 1024; + let requirements = |zone: &str| { + AffinityRequirements::new(HashMap::from([("ZONE".into(), zone.into())])).unwrap() + }; + let region_requirements = AllocationAffinity { + requirements: vec![ + Some(requirements("a")), + Some(requirements("b")), + Some(requirements("a")), + ], + target_sources: HashMap::from([ + (RegionId::new(table_id, 1), 0), + (RegionId::new(table_id, 2), 1), + (RegionId::new(table_id, 3), 2), + ]), + }; + let regions = vec![ + (1, "p1".into()), + (2, "p2".into()), + (3, "p3".into()), + (4, "p4".into()), + ]; + let groups = group_allocations(table_id, ®ions, ®ion_requirements); + assert!(ptr::eq( + groups[0].affinity.unwrap(), + region_requirements.requirements[0].as_ref().unwrap(), + )); + assert!(!ptr::eq( + region_requirements.requirements[0].as_ref().unwrap(), + region_requirements.requirements[2].as_ref().unwrap(), + )); + assert_eq!( + groups, + vec![ + AllocationGroup { + affinity: region_requirements.requirements[0].as_ref(), + regions: vec![(1, "p1"), (3, "p3")], + }, + AllocationGroup { + affinity: region_requirements.requirements[1].as_ref(), + regions: vec![(2, "p2")], + }, + AllocationGroup { + affinity: None, + regions: vec![(4, "p4")], + }, + ] + ); + } + + #[tokio::test] + async fn test_execute_refreshes_affinity_on_retry() { + use std::sync::Mutex; + use std::sync::atomic::{AtomicUsize, Ordering}; + + use common_meta::affinity::AffinityResolver; + use common_meta::error::{self as meta_error, Result as MetaResult}; + use common_meta::key::table_route::TableRouteKey; + use common_meta::peer::PeerAllocator; + + #[derive(Default)] + struct ChangingResolver { + calls: AtomicUsize, + } + + #[async_trait::async_trait] + impl AffinityResolver for ChangingResolver { + async fn resolve( + &self, + references: &[AffinityReference], + ) -> MetaResult>> { + assert_eq!( + references, + &[ + AffinityReference::Datanode(1), + AffinityReference::Datanode(2) + ] + ); + let generation = self.calls.fetch_add(1, Ordering::Relaxed); + Ok(references + .iter() + .map(|_| { + AffinityRequirements::new(HashMap::from([( + "ZONE".into(), + format!("zone-{generation}"), + )])) + }) + .collect()) + } + } + + #[derive(Default)] + struct RecordingAllocator(Mutex>>); + + #[async_trait::async_trait] + impl PeerAllocator for RecordingAllocator { + async fn alloc(&self, _num: usize, ctx: &PeerAllocContext) -> MetaResult> { + self.0.lock().unwrap().push(ctx.affinity.clone()); + meta_error::UnexpectedSnafu { + err_msg: "stop before creating regions", + } + .fail() + } + } + + let env = TestingEnv::new(); + let table_id = 1024; + let mut routes = create_current_region_routes(table_id, &[1, 2]); + routes[1].leader_peer = Some(Peer::empty(2)); + env.create_physical_table_metadata_for_repartition( + table_id, + routes, + test_region_wal_options(&[1, 2]), + ) + .await; + let (sender, _receiver) = mpsc::channel(1); + let mut ctx = new_parent_context( + &env, + Arc::new(MockDatanodeManager::new(DatanodeWatcher::new(sender))), + table_id, + ); + let resolver = Arc::new(ChangingResolver::default()); + let allocator = Arc::new(RecordingAllocator::default()); + ctx.affinity_resolver = Some(resolver.clone()); + ctx.region_routes_allocator = allocator.clone(); + ctx.persistent_ctx.plans = vec![RepartitionPlanEntry { + group_id: Uuid::nil(), + source_regions: vec![ + create_region_descriptor(table_id, 1, "x", 0, 100), + create_region_descriptor(table_id, 2, "x", 100, 200), + ], + target_regions: vec![ + create_target_region_descriptor(table_id, 3, "x", 0, 50), + create_target_region_descriptor(table_id, 4, "x", 50, 100), + create_target_region_descriptor(table_id, 5, "x", 100, 200), + ], + allocated_region_ids: (3..=5).map(|n| RegionId::new(table_id, n)).collect(), + pending_deallocate_region_ids: vec![], + transition_map: vec![vec![0, 1], vec![2]], + original_target_routes: vec![], + }]; + let procedure_ctx = ProcedureContext { + procedure_id: ProcedureId::random(), + provider: Arc::new(MockContextProvider::default()), + event_context: None, + }; + for _ in 0..2 { + assert!(ExecutePlan.next(&mut ctx, &procedure_ctx).await.is_err()); + } + assert_eq!(resolver.calls.load(Ordering::Relaxed), 2); + assert_eq!( + *allocator.0.lock().unwrap(), + (0..2) + .map(|generation| AffinityRequirements::new(HashMap::from([( + "ZONE".into(), + format!("zone-{generation}") + ),]))) + .collect::>() + ); + // Without a resolver, allocation does not need source route metadata. + // Removing it makes any accidental lookup fail before the recording allocator. + ctx.affinity_resolver = None; + env.kv_backend + .delete(TableRouteKey::new(table_id).to_string().as_bytes(), false) + .await + .unwrap(); + assert!(ExecutePlan.next(&mut ctx, &procedure_ctx).await.is_err()); + assert_eq!(resolver.calls.load(Ordering::Relaxed), 2); + let allocations = allocator.0.lock().unwrap(); + assert_eq!(allocations.len(), 3); + assert_eq!(allocations.last(), Some(&None)); + } + + #[tokio::test] + async fn test_resolved_affinity_preserves_source_associations() { + use common_meta::affinity::AffinityResolver; + use common_meta::error::Result as MetaResult; + + struct Resolver(Vec>); + + #[async_trait::async_trait] + impl AffinityResolver for Resolver { + async fn resolve( + &self, + references: &[AffinityReference], + ) -> MetaResult>> { + assert_eq!( + references, + &[ + AffinityReference::Datanode(2), + AffinityReference::Datanode(1) + ] + ); + Ok(self.0.clone()) + } + } + + let env = TestingEnv::new(); + let table_id = 1024; + let mut routes = create_current_region_routes(table_id, &[1, 2]); + routes[1].leader_peer = Some(Peer::empty(2)); + env.create_physical_table_metadata_for_repartition( + table_id, + routes, + test_region_wal_options(&[1, 2]), + ) + .await; + let (sender, _receiver) = mpsc::channel(1); + let mut ctx = new_parent_context( + &env, + Arc::new(MockDatanodeManager::new(DatanodeWatcher::new(sender))), + table_id, + ); + ctx.persistent_ctx.plans = vec![RepartitionPlanEntry { + group_id: Uuid::nil(), + source_regions: vec![ + create_region_descriptor(table_id, 1, "x", 0, 100), + create_region_descriptor(table_id, 2, "x", 100, 200), + ], + target_regions: vec![ + create_target_region_descriptor(table_id, 3, "x", 100, 200), + create_target_region_descriptor(table_id, 4, "x", 0, 50), + create_target_region_descriptor(table_id, 5, "x", 50, 100), + ], + allocated_region_ids: (3..=5).map(|n| RegionId::new(table_id, n)).collect(), + pending_deallocate_region_ids: vec![], + transition_map: vec![vec![1, 2], vec![0]], + original_target_routes: vec![], + }]; + let a = AffinityRequirements::new(HashMap::from([("ZONE".into(), "a".into())])).unwrap(); + let b = AffinityRequirements::new(HashMap::from([("ZONE".into(), "b".into())])).unwrap(); + assert!(matches!( + resolve_allocation_affinity(&ctx, &Resolver(vec![None, Some(a.clone())])).await, + Err(error::Error::Unexpected { .. }) + )); + ctx.affinity_fallback = AffinityFallback::Unconstrained; + for (resolved, expected) in [ + ( + vec![Some(b.clone()), Some(a.clone())], + HashMap::from([ + (RegionId::new(table_id, 3), b.clone()), + (RegionId::new(table_id, 4), a.clone()), + (RegionId::new(table_id, 5), a.clone()), + ]), + ), + ( + vec![None, Some(a.clone())], + HashMap::from([ + (RegionId::new(table_id, 4), a.clone()), + (RegionId::new(table_id, 5), a), + ]), + ), + ( + vec![Some(b.clone()), None], + HashMap::from([(RegionId::new(table_id, 3), b)]), + ), + ] { + let affinity = resolve_allocation_affinity(&ctx, &Resolver(resolved)) + .await + .unwrap(); + assert_eq!(affinity.requirements.len(), 2); + assert_eq!( + affinity.target_sources[&RegionId::new(table_id, 4)], + affinity.target_sources[&RegionId::new(table_id, 5)] + ); + let actual = affinity + .target_sources + .iter() + .filter_map(|(target, index)| { + affinity.requirements[*index] + .as_ref() + .map(|requirements| (*target, requirements)) + }) + .collect::>(); + if let Some(first) = actual.get(&RegionId::new(table_id, 4)) { + let second = actual[&RegionId::new(table_id, 5)]; + assert!(ptr::eq(*first, second)); + } + assert_eq!( + actual, + expected + .iter() + .map(|(target, requirements)| (*target, requirements)) + .collect() + ); + } + } + + #[tokio::test] + async fn test_missing_source_leader_requires_explicit_fallback() { + use common_meta::error::Result as MetaResult; + + struct Resolver; + + #[async_trait::async_trait] + impl AffinityResolver for Resolver { + async fn resolve( + &self, + references: &[AffinityReference], + ) -> MetaResult>> { + assert!(references.is_empty()); + Ok(vec![]) + } + } + + let env = TestingEnv::new(); + let table_id = 1024; + let mut routes = create_current_region_routes(table_id, &[1]); + routes[0].leader_peer = None; + env.create_physical_table_metadata_for_repartition( + table_id, + routes, + test_region_wal_options(&[1]), + ) + .await; + let mut ctx = new_parent_context(&env, Arc::new(MockDatanodeManager::new(())), table_id); + ctx.persistent_ctx.plans = vec![RepartitionPlanEntry { + group_id: Uuid::nil(), + source_regions: vec![create_region_descriptor(table_id, 1, "x", 0, 100)], + target_regions: vec![create_target_region_descriptor(table_id, 2, "x", 0, 100)], + allocated_region_ids: vec![RegionId::new(table_id, 2)], + pending_deallocate_region_ids: vec![], + transition_map: vec![vec![0]], + original_target_routes: vec![], + }]; + assert!(matches!( + resolve_allocation_affinity(&ctx, &Resolver).await, + Err(error::Error::Unexpected { .. }) + )); + ctx.affinity_fallback = AffinityFallback::Unconstrained; + assert!( + resolve_allocation_affinity(&ctx, &Resolver) + .await + .unwrap() + .target_sources + .is_empty() + ); + } } diff --git a/src/meta-srv/src/procedure/utils.rs b/src/meta-srv/src/procedure/utils.rs index 1f975b2f9c0..ff88d1b38ba 100644 --- a/src/meta-srv/src/procedure/utils.rs +++ b/src/meta-srv/src/procedure/utils.rs @@ -536,6 +536,7 @@ pub mod test_data { )); DdlContext { peer_alloc_context_provider: None, + affinity_fallback: Default::default(), affinity_resolver: None, node_manager, cache_invalidator: Arc::new(MetasrvCacheInvalidator::new( diff --git a/src/meta-srv/src/region/supervisor.rs b/src/meta-srv/src/region/supervisor.rs index 93624a61892..1c431ef9960 100644 --- a/src/meta-srv/src/region/supervisor.rs +++ b/src/meta-srv/src/region/supervisor.rs @@ -18,7 +18,9 @@ use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; use async_trait::async_trait; +use common_error::ext::BoxedError; use common_meta::DatanodeId; +use common_meta::affinity::{AffinityFallback, AffinityReference, AffinityResolverRef}; use common_meta::datanode::Stat; use common_meta::ddl::{DetectingRegion, RegionFailureDetectorController}; use common_meta::key::runtime_switch::RuntimeSwitchManagerRef; @@ -290,6 +292,8 @@ pub struct RegionSupervisor { selector_context: SelectorContext, /// Candidate node selector. selector: RegionSupervisorSelector, + /// Resolves the failed node attributes before selection. + affinity_resolver: Option, /// Region migration manager. region_migration_manager: RegionMigrationManagerRef, /// The maintenance mode manager. @@ -408,6 +412,7 @@ impl RegionSupervisor { receiver: event_receiver, selector_context, selector, + affinity_resolver: None, region_migration_manager, runtime_switch_manager, peer_resolver, @@ -416,6 +421,12 @@ impl RegionSupervisor { } } + /// Sets the resolver used to prepare failover affinity requirements. + pub(crate) fn with_affinity_resolver(mut self, resolver: AffinityResolverRef) -> Self { + self.affinity_resolver = Some(resolver); + self + } + /// Sets the meta state. pub(crate) fn with_state(mut self, state: StateRef) -> Self { self.state = Some(state); @@ -659,17 +670,41 @@ impl RegionSupervisor { regions: &[RegionId], failure_datanodes: &[DatanodeId], ) -> Result> { - let exclude_peer_ids = HashSet::from_iter(failure_datanodes.iter().cloned()); + let affinity = if let Some(resolver) = &self.affinity_resolver { + let mut resolved = resolver + .resolve(&[AffinityReference::Datanode(from_peer_id)]) + .await + .map_err(BoxedError::new) + .context(error::OtherSnafu)?; + ensure!( + resolved.len() == 1, + error::UnexpectedSnafu { + violated: "affinity resolution must preserve the number and order of references", + } + ); + let affinity = resolved.pop().flatten(); + ensure!( + affinity.is_some() + || self.selector_context.affinity_fallback == AffinityFallback::Unconstrained, + error::UnexpectedSnafu { + violated: "source affinity requirements are unavailable" + } + ); + affinity + } else { + None + }; + let exclude_peer_ids = HashSet::from_iter(failure_datanodes.iter().copied()); match &self.selector { RegionSupervisorSelector::NaiveSelector(selector) => { - let opt = SelectorOptions { + let opts = SelectorOptions { + affinity: affinity.as_ref(), min_required_items: regions.len(), allow_duplication: true, exclude_peer_ids, workload_filter: Some(accept_ingest_workload), - affinity: None, }; - let peers = selector.select(&self.selector_context, opt).await?; + let peers = selector.select(&self.selector_context, opts).await?; ensure!( peers.len() == regions.len(), error::NoEnoughAvailableNodeSnafu { @@ -693,6 +728,7 @@ impl RegionSupervisor { from_peer_id, regions, exclude_peer_ids, + affinity.as_ref(), ) .await?; ensure!( @@ -962,6 +998,280 @@ pub(crate) mod tests { ) } + #[tokio::test] + async fn test_failover_passes_affinity_and_propagates_errors() { + use std::collections::HashSet; + use std::sync::atomic::{AtomicUsize, Ordering}; + + use common_meta::affinity::{ + AffinityFallback, AffinityReference, AffinityRequirements, AffinityResolver, + }; + use common_meta::error::{self as meta_error, Result as MetaResult}; + + use crate::error::{self, Result}; + use crate::metasrv::{SelectTarget, SelectorContext}; + use crate::selector::{RegionStatAwareSelector, Selector, SelectorOptions}; + + #[derive(Clone, Copy, PartialEq)] + enum Resolution { + Disabled, + Available, + Missing, + Error, + } + + struct Resolver { + result: Resolution, + calls: AtomicUsize, + } + + #[async_trait::async_trait] + impl AffinityResolver for Resolver { + async fn resolve( + &self, + references: &[AffinityReference], + ) -> MetaResult>> { + assert_eq!(references, &[AffinityReference::Datanode(7)]); + self.calls.fetch_add(1, Ordering::Relaxed); + match self.result { + Resolution::Available => { + Ok(vec![AffinityRequirements::new(HashMap::from([( + "pool".into(), + "a".into(), + )]))]) + } + Resolution::Disabled => unreachable!(), + Resolution::Missing => Ok(vec![None]), + Resolution::Error => meta_error::UnexpectedSnafu { + err_msg: "source lookup failed", + } + .fail(), + } + } + } + + #[derive(Clone, Copy)] + enum Outcome { + Success, + Short, + Exhausted, + Error, + } + struct FailoverSelector { + outcome: Outcome, + expected_fallback: AffinityFallback, + calls: Mutex>>, + } + impl FailoverSelector { + fn select_peers(&self, affinity: Option<&AffinityRequirements>) -> Result> { + self.calls.lock().unwrap().push(affinity.cloned()); + if affinity.is_some() { + match self.outcome { + Outcome::Success => {} + Outcome::Short => return Ok(vec![Peer::empty(9)]), + Outcome::Exhausted => { + return error::NoEnoughAvailableNodeSnafu { + required: 2usize, + available: 0usize, + select_target: SelectTarget::Datanode, + } + .fail(); + } + Outcome::Error => { + return error::UnexpectedSnafu { + violated: "selection backend failed", + } + .fail(); + } + } + } + Ok(vec![Peer::empty(9); 2]) + } + } + #[async_trait::async_trait] + impl Selector for FailoverSelector { + type Context = SelectorContext; + type Output = Vec; + async fn select( + &self, + ctx: &SelectorContext, + opts: SelectorOptions<'_>, + ) -> Result { + assert_eq!(ctx.affinity_fallback, self.expected_fallback); + assert_eq!(opts.exclude_peer_ids, HashSet::from([8])); + assert!(opts.workload_filter.is_some()); + assert!(opts.allow_duplication); + assert_eq!(opts.min_required_items, 2); + self.select_peers(opts.affinity) + } + } + #[async_trait::async_trait] + impl RegionStatAwareSelector for FailoverSelector { + type Context = SelectorContext; + type Output = Vec<(RegionId, Peer)>; + async fn select( + &self, + ctx: &SelectorContext, + source: u64, + regions: &[RegionId], + excluded: HashSet, + affinity: Option<&AffinityRequirements>, + ) -> Result { + assert_eq!(ctx.affinity_fallback, self.expected_fallback); + assert_eq!(source, 7); + assert_eq!(excluded, HashSet::from([8])); + Ok(regions + .iter() + .copied() + .zip(self.select_peers(affinity)?) + .collect()) + } + } + for stat_aware in [false, true] { + for (resolution, fallback, outcome, expected_calls, succeeds) in [ + ( + Resolution::Available, + AffinityFallback::Error, + Outcome::Success, + 1, + true, + ), + ( + Resolution::Available, + AffinityFallback::Error, + Outcome::Short, + 1, + false, + ), + ( + Resolution::Available, + AffinityFallback::Error, + Outcome::Exhausted, + 1, + false, + ), + ( + Resolution::Available, + AffinityFallback::Error, + Outcome::Error, + 1, + false, + ), + ( + Resolution::Missing, + AffinityFallback::Error, + Outcome::Success, + 0, + false, + ), + ( + Resolution::Error, + AffinityFallback::Error, + Outcome::Success, + 0, + false, + ), + ( + Resolution::Available, + AffinityFallback::Unconstrained, + Outcome::Short, + 1, + false, + ), + ( + Resolution::Available, + AffinityFallback::Unconstrained, + Outcome::Exhausted, + 1, + false, + ), + ( + Resolution::Available, + AffinityFallback::Unconstrained, + Outcome::Error, + 1, + false, + ), + ( + Resolution::Missing, + AffinityFallback::Unconstrained, + Outcome::Success, + 1, + true, + ), + ( + Resolution::Error, + AffinityFallback::Unconstrained, + Outcome::Success, + 0, + false, + ), + ( + Resolution::Disabled, + AffinityFallback::Unconstrained, + Outcome::Success, + 1, + true, + ), + ] { + let resolver = Arc::new(Resolver { + result: resolution, + calls: AtomicUsize::new(0), + }); + let selector = Arc::new(FailoverSelector { + outcome, + expected_fallback: fallback, + calls: Mutex::new(vec![]), + }); + let (mut supervisor, _sender) = new_test_supervisor(); + supervisor.selector_context.affinity_fallback = fallback; + if resolution != Resolution::Disabled { + supervisor = supervisor.with_affinity_resolver(resolver.clone()); + } + supervisor.selector = if stat_aware { + RegionSupervisorSelector::RegionStatAwareSelector(selector.clone()) + } else { + RegionSupervisorSelector::NaiveSelector(selector.clone()) + }; + let regions = [RegionId::new(1, 1), RegionId::new(1, 2)]; + // Preserve the caller's exclusion set without adding the source. + let result = supervisor.select_peers(7, ®ions, &[8]).await; + if succeeds { + assert_eq!( + result.unwrap(), + vec![(regions[0], Peer::empty(9)), (regions[1], Peer::empty(9))] + ); + } else { + let error = result.unwrap_err(); + if resolution == Resolution::Error { + assert!(matches!(error, error::Error::Other { .. })); + } else if resolution == Resolution::Missing { + assert!(matches!(error, error::Error::Unexpected { .. })); + } else { + match outcome { + Outcome::Short | Outcome::Exhausted => { + assert!(matches!(error, error::Error::NoEnoughAvailableNode { .. })) + } + Outcome::Error => { + assert!(matches!(error, error::Error::Unexpected { .. })) + } + Outcome::Success => unreachable!(), + } + } + } + assert_eq!( + resolver.calls.load(Ordering::Relaxed), + usize::from(resolution != Resolution::Disabled) + ); + let calls = selector.calls.lock().unwrap(); + assert_eq!(calls.len(), expected_calls); + if let Some(affinity) = calls.first() { + assert_eq!(affinity.is_some(), resolution == Resolution::Available); + } + } + } + } + fn filler_heartbeat() -> DatanodeHeartbeat { DatanodeHeartbeat { datanode_id: 1, diff --git a/src/meta-srv/src/selector.rs b/src/meta-srv/src/selector.rs index e5c643b7b7c..c4d29b8ecea 100644 --- a/src/meta-srv/src/selector.rs +++ b/src/meta-srv/src/selector.rs @@ -54,6 +54,7 @@ pub trait RegionStatAwareSelector: Send + Sync { from_peer_id: u64, region_ids: &[RegionId], exclude_peer_ids: HashSet, + affinity: Option<&AffinityRequirements>, ) -> Result; } diff --git a/tests-integration/src/standalone.rs b/tests-integration/src/standalone.rs index e0113f90425..5bd0d970de5 100644 --- a/tests-integration/src/standalone.rs +++ b/tests-integration/src/standalone.rs @@ -291,6 +291,7 @@ impl GreptimeDbStandaloneBuilder { let ddl_manager = Arc::new(DdlManager::new( DdlContext { affinity_resolver: None, + affinity_fallback: Default::default(), peer_alloc_context_provider: None, node_manager: node_manager.clone(), cache_invalidator: cache_registry.clone(), From ce428bb74dd54fbed374cbbe137298c4366fc94d Mon Sep 17 00:00:00 2001 From: YingNing <59612604+dxqlovestudy@users.noreply.github.com> Date: Fri, 16 Oct 2026 02:36:47 +0000 Subject: [PATCH 018/239] feat(servers): accept application/x-json-stream for log ingestion (#9581) * feat(servers): accept application/x-json-stream for log ingestion Support fluentbit's json_stream output format on the HTTP log ingestion endpoints by accepting `application/x-json-stream` as a payload content type. The payload is parsed as concatenated JSON -- multiple JSON values placed back-to-back without a separator, e.g. `{"a":1}{"b":2}` -- reusing the streaming value iterator already used for `application/json`, so the behavior for concatenated or mixed payloads is identical. Aimed at #5404 (msgpack support is left as a possible follow-up). Signed-off-by: dxqlovestudy <59612604+dxqlovestudy@users.noreply.github.com> * chore(servers): remove redundant comments per review feedback Signed-off-by: dxqlovestudy <59612604+dxqlovestudy@users.noreply.github.com> --------- Signed-off-by: dxqlovestudy <59612604+dxqlovestudy@users.noreply.github.com> --- src/servers/src/http/event.rs | 66 ++++++++++++++++++++++++++++++--- src/servers/src/http/header.rs | 3 ++ tests-integration/tests/http.rs | 56 +++++++++++++++++++++++++++- 3 files changed, 119 insertions(+), 6 deletions(-) diff --git a/src/servers/src/http/event.rs b/src/servers/src/http/event.rs index 0c96becaf59..651621dcdea 100644 --- a/src/servers/src/http/event.rs +++ b/src/servers/src/http/event.rs @@ -59,7 +59,8 @@ use crate::http::header::constants::{ GREPTIME_PIPELINE_PARAMS_HEADER, }; use crate::http::header::{ - CONTENT_TYPE_NDJSON_STR, CONTENT_TYPE_NDJSON_SUBTYPE_STR, CONTENT_TYPE_PROTOBUF_STR, + CONTENT_TYPE_JSON_STREAM_STR, CONTENT_TYPE_JSON_STREAM_SUBTYPE_STR, CONTENT_TYPE_NDJSON_STR, + CONTENT_TYPE_NDJSON_SUBTYPE_STR, CONTENT_TYPE_PROTOBUF_STR, }; use crate::http::result::greptime_manage_resp::{GreptimedbManageResponse, SqlOutput}; use crate::http::result::greptime_result_v1::GreptimedbV1Response; @@ -87,6 +88,8 @@ lazy_static! { ContentType::from_str(CONTENT_TYPE_PROTOBUF_STR).unwrap(); pub static ref NDJSON_CONTENT_TYPE: ContentType = ContentType::from_str(CONTENT_TYPE_NDJSON_STR).unwrap(); + pub static ref JSON_STREAM_CONTENT_TYPE: ContentType = + ContentType::from_str(CONTENT_TYPE_JSON_STREAM_STR).unwrap(); } /// LogIngesterQueryParams is used for query params of log ingester API. @@ -797,6 +800,10 @@ pub async fn log_ingester( #[derive(Debug, EnumIter)] enum EventPayloadResolverInner { Json, + /// `application/x-json-stream`: concatenated JSON values without any + /// separator between them, e.g. `{"a":1}{"b":2}` (the `json_stream` + /// format of fluentbit). + JsonStream, Ndjson, Text, } @@ -805,6 +812,7 @@ impl Display for EventPayloadResolverInner { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { EventPayloadResolverInner::Json => write!(f, "{}", *JSON_CONTENT_TYPE), + EventPayloadResolverInner::JsonStream => write!(f, "{}", *JSON_STREAM_CONTENT_TYPE), EventPayloadResolverInner::Ndjson => write!(f, "{}", *NDJSON_CONTENT_TYPE), EventPayloadResolverInner::Text => write!(f, "{}", *TEXT_CONTENT_TYPE), } @@ -818,6 +826,9 @@ impl TryFrom<&ContentType> for EventPayloadResolverInner { let mime: mime_guess::Mime = content_type.clone().into(); match (mime.type_(), mime.subtype()) { (mime::APPLICATION, mime::JSON) => Ok(EventPayloadResolverInner::Json), + (mime::APPLICATION, subtype) if subtype == CONTENT_TYPE_JSON_STREAM_SUBTYPE_STR => { + Ok(EventPayloadResolverInner::JsonStream) + } (mime::APPLICATION, subtype) if subtype == CONTENT_TYPE_NDJSON_SUBTYPE_STR => { Ok(EventPayloadResolverInner::Ndjson) } @@ -866,10 +877,12 @@ impl<'a> TryFrom<&'a ContentType> for EventPayloadResolver<'a> { impl EventPayloadResolver<'_> { fn parse_payload(&self, payload: Bytes, ignore_errors: bool) -> Result> { match self.inner { - EventPayloadResolverInner::Json => transform_ndjson_array_factory( - Deserializer::from_slice(&payload).into_iter(), - ignore_errors, - ), + EventPayloadResolverInner::Json | EventPayloadResolverInner::JsonStream => { + transform_ndjson_array_factory( + Deserializer::from_slice(&payload).into_iter(), + ignore_errors, + ) + } EventPayloadResolverInner::Ndjson => { let mut result = Vec::with_capacity(1000); let mut buffer = Buffers::new(1000); @@ -1077,6 +1090,49 @@ mod tests { assert_eq!(a, "[{\"a\":1},{\"b\":2}]"); } + #[test] + fn test_extract_by_json_stream_content_type() { + // `application/x-json-stream` accepts concatenated JSON values, with + // or without line breaks between them. + let json_stream = ContentType::from_str("application/x-json-stream").unwrap(); + let payload = r#"{"a": 1}{"b": 2} +{"c": 1}"# + .as_bytes(); + let payload = Bytes::from_static(payload); + + let values = + extract_pipeline_value_by_content_type(json_stream.clone(), payload, false).unwrap(); + assert_eq!( + values, + vec![ + json!({"a": 1}).into(), + json!({"b": 2}).into(), + json!({"c": 1}).into() + ] + ); + + assert!( + EventPayloadResolver::support_content_type_list().contains(&json_stream.to_string()) + ); + } + + #[test] + fn test_json_stream_matches_json_for_concatenated_values() { + // a concatenated payload is parsed exactly like `application/json` + let payload = Bytes::from_static(r#"{"a": 1}{"b": 2}"#.as_bytes()); + let via_json = + extract_pipeline_value_by_content_type(ContentType::json(), payload.clone(), false) + .unwrap(); + let via_stream = extract_pipeline_value_by_content_type( + ContentType::from_str("application/x-json-stream").unwrap(), + payload, + false, + ) + .unwrap(); + assert_eq!(via_json, via_stream); + assert_eq!(via_stream.len(), 2); + } + #[test] fn test_extract_by_content() { let payload = r#" diff --git a/src/servers/src/http/header.rs b/src/servers/src/http/header.rs index 625bee41ee7..d6f69fa72d5 100644 --- a/src/servers/src/http/header.rs +++ b/src/servers/src/http/header.rs @@ -106,6 +106,9 @@ pub static CONTENT_ENCODING_SNAPPY: HeaderValue = HeaderValue::from_static("snap pub static CONTENT_TYPE_NDJSON_STR: &str = "application/x-ndjson"; pub static CONTENT_TYPE_NDJSON_SUBTYPE_STR: &str = "x-ndjson"; +pub static CONTENT_TYPE_JSON_STREAM_STR: &str = "application/x-json-stream"; +pub static CONTENT_TYPE_JSON_STREAM_SUBTYPE_STR: &str = "x-json-stream"; + pub struct GreptimeDbName(Option); impl Header for GreptimeDbName { diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 20d6f6dde7f..8ac98bf22f6 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -189,6 +189,7 @@ macro_rules! http_tests { test_pipeline_api, test_pipeline_name_in_header, + test_json_stream_ingestion, test_plain_text_ingestion, test_pipeline_auto_transform, test_pipeline_auto_transform_with_select, @@ -8023,7 +8024,7 @@ transform: assert_eq!( body["error"], json!( - "Invalid request parameter: invalid content type: application/yaml, expected: one of application/json, application/x-ndjson, text/plain" + "Invalid request parameter: invalid content type: application/yaml, expected: one of application/json, application/x-json-stream, application/x-ndjson, text/plain" ) ); @@ -8180,6 +8181,59 @@ transform: guard.remove_all().await; } +pub async fn test_json_stream_ingestion(store_type: StorageType) { + common_telemetry::init_default_ut_logging(); + let (app, mut guard) = + setup_test_http_app_with_frontend(store_type, "test_json_stream_ingestion").await; + + // handshake + let client = TestClient::new(app).await; + + let pipeline_body = r#" +processors: + - date: + field: time + formats: + - "%Y-%m-%d %H:%M:%S%.3f" +transform: + - field: message + type: string + - field: time + type: time + index: timestamp +"#; + + // 1. create pipeline + let res = client + .post("/v1/pipelines/test") + .header("Content-Type", "application/x-yaml") + .body(pipeline_body) + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + + // 2. write data as concatenated JSON (the `json_stream` format of fluentbit) + let data_body = r#"{"time": "2024-05-25 20:16:37.217", "message": "hello"}{"time": "2024-05-25 20:16:37.218", "message": "hello world"}"#; + let res = client + .post("/v1/ingest?db=public&table=logs1&pipeline_name=test") + .header("Content-Type", "application/x-json-stream") + .body(data_body) + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + + // 3. check data + validate_data( + "json_stream_ingestion", + &client, + "select message, time from logs1", + r#"[["hello",1716668197217000000],["hello world",1716668197218000000]]"#, + ) + .await; + + guard.remove_all().await; +} + pub async fn test_pipeline_auto_transform(store_type: StorageType) { common_telemetry::init_default_ut_logging(); let (app, mut guard) = From f663569cf03108962ca762b1c65af426e0af1945 Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Fri, 16 Oct 2026 03:25:58 +0000 Subject: [PATCH 019/239] feat(auth): expose public authentication config (#9626) Signed-off-by: Lei, HUANG --- Cargo.lock | 1 + src/auth/Cargo.toml | 1 + src/auth/src/lib.rs | 2 +- src/auth/src/user_provider.rs | 27 ++++++++++++++++++++ src/servers/src/http.rs | 41 ++++++++++++++++++++++++++++++- src/servers/src/http/authorize.rs | 24 +++++++++++++++++- 6 files changed, 93 insertions(+), 3 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 658180353e3..47aa178a57f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -974,6 +974,7 @@ dependencies = [ "hmac 0.12.1", "pbkdf2", "rand 0.9.4", + "serde", "sha1 0.10.6", "sha2 0.10.9", "snafu 0.8.6", diff --git a/src/auth/Cargo.toml b/src/auth/Cargo.toml index 959eb25078e..1c2161adf22 100644 --- a/src/auth/Cargo.toml +++ b/src/auth/Cargo.toml @@ -24,6 +24,7 @@ hex.workspace = true hmac = "0.12" pbkdf2 = "0.12" rand.workspace = true +serde.workspace = true sha1 = "0.10" sha2 = "0.10" snafu.workspace = true diff --git a/src/auth/src/lib.rs b/src/auth/src/lib.rs index aacd6282d1e..c668fa9916c 100644 --- a/src/auth/src/lib.rs +++ b/src/auth/src/lib.rs @@ -39,7 +39,7 @@ pub use permission::{ }; pub use user_info::UserInfo; pub use user_provider::static_user_provider::StaticUserProvider; -pub use user_provider::{BEARER_TOKEN_USER, MysqlAuthMethod, PgAuthInfo, UserProvider}; +pub use user_provider::{AuthMethod, BEARER_TOKEN_USER, MysqlAuthMethod, PgAuthInfo, UserProvider}; /// pub type alias pub type UserInfoRef = std::sync::Arc; diff --git a/src/auth/src/user_provider.rs b/src/auth/src/user_provider.rs index 8b8c63581ac..4d3eb509b2f 100644 --- a/src/auth/src/user_provider.rs +++ b/src/auth/src/user_provider.rs @@ -24,6 +24,7 @@ use std::{fmt, io}; use common_base::secrets::ExposeSecret; use common_telemetry::warn; use pbkdf2::pbkdf2_hmac; +use serde::Serialize; use sha2::Sha256; use snafu::{OptionExt, ResultExt, ensure}; use subtle::ConstantTimeEq; @@ -49,6 +50,27 @@ use crate::{UserInfoRef, auth_mysql}; /// deployment choice. pub const BEARER_TOKEN_USER: &str = "*"; +/// A login method advertised by the HTTP authentication configuration endpoint. +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum AuthMethod { + /// Username and password authentication. + Password, + /// OpenID Connect authentication for browser clients. + Oidc { + /// Stable identifier for this identity provider. + id: String, + /// Human-readable identity provider name. + display_name: String, + /// OpenID Connect issuer URL. + issuer: String, + /// Public client identifier. + client_id: String, + /// OAuth scopes requested during login. + scopes: Vec, + }, +} + #[async_trait::async_trait] pub trait UserProvider: Send + Sync { fn name(&self) -> &str; @@ -107,6 +129,11 @@ pub trait UserProvider: Send + Sync { Ok(user_info) } + /// Returns the login methods that browser clients can use. + fn auth_methods(&self) -> Vec { + vec![AuthMethod::Password] + } + fn mysql_auth_method(&self) -> MysqlAuthMethod { if self.external() { MysqlAuthMethod::ClearPassword diff --git a/src/servers/src/http.rs b/src/servers/src/http.rs index 210fe0cee00..31ed00d443c 100644 --- a/src/servers/src/http.rs +++ b/src/servers/src/http.rs @@ -169,7 +169,8 @@ pub const AUTHORIZATION_HEADER: &str = "x-greptime-auth"; // TODO(fys): This is a temporary workaround, it will be improved later // these APIs will bypass authentication. -pub static PUBLIC_API_PREFIX: [&str; 4] = [ +pub static PUBLIC_API_PREFIX: [&str; 5] = [ + "/v1/auth/config", "/v1/influxdb/ping", "/v1/influxdb/health", "/v1/health", @@ -978,8 +979,12 @@ impl HttpServer { /// on both listeners automatically, and the API guard decides exposure. pub fn make_app(&self) -> Router { let mut router = self.router.lock().unwrap().clone(); + let auth_router = Router::new() + .route("/v1/auth/config", routing::get(authorize::auth_config)) + .with_state(AuthState::new(self.user_provider.clone())); router = router + .merge(auth_router) .route( &format!("/{HTTP_API_VERSION}/health"), routing::get(handler::health).post(handler::health), @@ -1849,6 +1854,40 @@ mod test { assert_eq!(opts.api_server_addr, "127.0.0.1:4006"); } + #[tokio::test] + async fn test_auth_config_without_provider_disables_auth() { + let server = HttpServerBuilder::new(HttpOptions::default()).build(); + let client = TestClient::new(server.build(server.make_app()).unwrap()).await; + let response = client.get("/v1/auth/config").send().await; + + assert_eq!(response.status(), StatusCode::OK); + assert_eq!( + response.json::().await, + serde_json::json!({"auth_required": false, "methods": []}) + ); + } + + #[tokio::test] + async fn test_auth_config_with_provider_is_public() { + let provider = + auth::static_user_provider_from_option("static_user_provider:cmd:user=password") + .unwrap(); + let server = HttpServerBuilder::new(HttpOptions::default()) + .with_user_provider(Arc::new(provider)) + .build(); + let client = TestClient::new(server.build(server.make_app()).unwrap()).await; + let response = client.get("/v1/auth/config").send().await; + + assert_eq!(response.status(), StatusCode::OK); + assert_eq!( + response.json::().await, + serde_json::json!({ + "auth_required": true, + "methods": [{"type": "password"}] + }) + ); + } + #[tokio::test] pub async fn test_build_serves_all_routes() { // The single (internal/full) server built via `build()` keeps serving both diff --git a/src/servers/src/http/authorize.rs b/src/servers/src/http/authorize.rs index 755bd4d4c24..5c5211ce79f 100644 --- a/src/servers/src/http/authorize.rs +++ b/src/servers/src/http/authorize.rs @@ -12,8 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. -use ::auth::UserProviderRef; +use ::auth::{AuthMethod, UserProviderRef}; use api::v1::Basic; +use axum::Json; use axum::extract::{Request, State}; use axum::http::{self, StatusCode}; use axum::middleware::Next; @@ -28,6 +29,7 @@ use common_telemetry::warn; use common_time::Timezone; use common_time::timezone::parse_timezone; use headers::Header; +use serde::Serialize; use session::context::QueryContextBuilder; use snafu::{OptionExt, ResultExt, ensure}; @@ -55,6 +57,26 @@ impl AuthState { } } +/// Public authentication configuration used by browser clients. +#[derive(Debug, Serialize)] +pub(crate) struct AuthConfig { + auth_required: bool, + methods: Vec, +} + +/// Returns the configured login methods without requiring authentication. +pub(crate) async fn auth_config(State(auth_state): State) -> Json { + let (auth_required, methods) = match auth_state.user_provider { + Some(user_provider) => (true, user_provider.auth_methods()), + None => (false, Vec::new()), + }; + + Json(AuthConfig { + auth_required, + methods, + }) +} + pub async fn inner_auth( user_provider: Option, mut req: Request, From 778f4556a8db472544888040de77391f5474a55e Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 16 Oct 2026 04:48:32 +0000 Subject: [PATCH 020/239] feat: unify query options across SQL, HTTP and gRPC (#9494) Isolated rebase verification of PR #9494. Preserve the reviewed net change; original PR branch remains untouched. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- docs/how-to/query-options.md | 18 +++ src/datanode/src/region_server.rs | 18 ++- src/datanode/src/tests.rs | 5 +- src/flow/src/batching_mode/frontend_client.rs | 53 ++++--- src/frontend/src/instance.rs | 15 ++ src/frontend/src/instance/grpc.rs | 16 +- src/operator/src/statement.rs | 12 +- src/operator/src/statement/set.rs | 66 ++++++--- src/operator/src/statement/show.rs | 9 ++ src/query/src/datafusion.rs | 56 +------ src/query/src/dist_plan/analyzer.rs | 16 +- src/query/src/optimizer/constant_term.rs | 2 +- src/query/src/planner.rs | 74 ++++++---- src/query/src/query_engine.rs | 2 +- .../src/query_engine/default_serializer.rs | 17 ++- src/query/src/query_engine/state.rs | 60 +++++++- src/query/src/sql.rs | 87 ++++++++++- src/servers/src/grpc/flight.rs | 1 + src/servers/src/grpc/greptime_handler.rs | 1 + src/servers/src/hint_headers.rs | 55 +++++++ src/servers/src/http.rs | 17 +-- src/servers/src/http/hints.rs | 41 +++++- src/servers/src/http/read_preference.rs | 7 +- src/servers/src/mysql/federated.rs | 65 +++++++- src/servers/src/postgres/handler.rs | 22 ++- src/session/src/context.rs | 139 +++++++++++++++++- src/session/src/lib.rs | 1 + src/session/src/query_options.rs | 116 +++++++++++++++ src/sql/src/ast.rs | 34 +++++ tests/cases/standalone/common/basic.result | 10 +- 30 files changed, 855 insertions(+), 180 deletions(-) create mode 100644 docs/how-to/query-options.md create mode 100644 src/session/src/query_options.rs diff --git a/docs/how-to/query-options.md b/docs/how-to/query-options.md new file mode 100644 index 00000000000..3fdaa55979d --- /dev/null +++ b/docs/how-to/query-options.md @@ -0,0 +1,18 @@ +# Query options + +Query options adjust selected behavior for a connection or request; they are not server-wide settings or resource-pool controls. + +| Canonical option | SQL aliases | Values | +| --- | --- | --- | +| `query.parallelism` | `query_parallelism`, `query.parallelism` | Integer `1`–`1024`; server default `0` means automatic CPU-based parallelism. This is not a resource-pool limit. | +| `query.allow_query_fallback` | `query_fallback`, `allow_query_fallback`, `query.allow_query_fallback` (legacy `ALLOW_QUERY_FALLBACK` syntax / `query_fallback` hint also work) | `true` or `false` | +| `query.enable_remote_dynamic_filter_pushdown` | canonical name | `true` or `false` | +| `datafusion.optimizer.*` | canonical name | Keys, types, defaults, and validation come from the linked DataFusion version (for example, numeric `max_passes` and Boolean `filter_null_join_keys`). | + +Use `SET query_parallelism = 4` or `SET query.parallelism TO 4`; `SHOW VARIABLES query.parallelism` works in MySQL and PostgreSQL, which also accepts `SHOW query.parallelism`. `SHOW VARIABLES

, ctx: QueryContextRef) -> Result<()> { - let allow_fallback_expr = exprs.first().context(NotSupportedSnafu { - feat: "No allow query fallback value find in set variable statement", - })?; - match allow_fallback_expr { - Expr::Value(ValueWithSpan { - value: Value::Boolean(allow), - span: _, - }) => { - ctx.configuration_parameter() - .set_allow_query_fallback(*allow); - Ok(()) - } - expr => NotSupportedSnafu { - feat: format!( - "Unsupported allow query fallback expr {} in set variable statement", - expr - ), - } - .fail(), +/// Applies supported query options and rejects unsupported owned namespaces. +pub fn set_query_option(variable: &str, exprs: Vec, ctx: QueryContextRef) -> Result { + let lower = variable.to_ascii_lowercase(); + if !lower.starts_with("query.") + && !lower.starts_with("datafusion.") + && session::query_options::canonical_query_option_name(variable).is_none() + { + return Ok(false); } + let [Expr::Value(value)] = exprs.as_slice() else { + return Err(InvalidConfigValueSnafu.into_error( + session::session_config::Error::InvalidConfigValue { + name: variable.to_string(), + value: exprs + .iter() + .map(ToString::to_string) + .collect::>() + .join(", "), + hint: "Expected exactly one literal value".to_string(), + location: snafu::Location::new(file!(), line!(), column!()), + }, + )); + }; + let value = match &value.value { + Value::Boolean(value) => value.to_string(), + Value::Number(value, _) + | Value::SingleQuotedString(value) + | Value::DoubleQuotedString(value) => value.clone(), + _ => { + return Err(InvalidConfigValueSnafu.into_error( + session::session_config::Error::InvalidConfigValue { + name: variable.to_string(), + value: value.value.to_string(), + hint: "Expected a boolean, number, or quoted string literal".to_string(), + location: snafu::Location::new(file!(), line!(), column!()), + }, + )); + } + }; + ctx.configuration_parameter() + .set_query_option(variable, &value) + .context(InvalidConfigValueSnafu)?; + Ok(true) } pub fn set_intervalstyle(exprs: Vec, ctx: QueryContextRef) -> Result<()> { diff --git a/src/operator/src/statement/show.rs b/src/operator/src/statement/show.rs index 4a3fa5ae9b5..b721abc4999 100644 --- a/src/operator/src/statement/show.rs +++ b/src/operator/src/statement/show.rs @@ -350,6 +350,15 @@ impl StatementExecutor { #[tracing::instrument(skip_all)] pub fn show_variable(&self, stmt: ShowVariables, query_ctx: QueryContextRef) -> Result { + if let Some(output) = query::sql::show_query_option( + stmt.clone(), + query_ctx.clone(), + self.query_engine.as_ref(), + ) + .context(error::ExecuteStatementSnafu)? + { + return Ok(output); + } query::sql::show_variable(stmt, query_ctx).context(error::ExecuteStatementSnafu) } diff --git a/src/query/src/datafusion.rs b/src/query/src/datafusion.rs index 72788461836..4b23f81eb16 100644 --- a/src/query/src/datafusion.rs +++ b/src/query/src/datafusion.rs @@ -55,9 +55,7 @@ use tracing::Span; use crate::analyze::DistAnalyzeExec; pub use crate::datafusion::planner::DfContextProviderAdapter; -use crate::dist_plan::{ - DistPlannerOptions, MergeScanLogicalPlan, RemoteDynFilterReceiverInjectorRef, -}; +use crate::dist_plan::{MergeScanLogicalPlan, RemoteDynFilterReceiverInjectorRef}; use crate::error::{ CatalogSnafu, CreateRecordBatchSnafu, MissingTableMutationHandlerSnafu, MissingTimestampColumnSnafu, QueryExecutionSnafu, Result, TableMutationSnafu, @@ -156,7 +154,8 @@ impl DatafusionQueryEngine { plan: LogicalPlan, query_ctx: QueryContextRef, ) -> Result { - let mut ctx = self.engine_context(query_ctx.clone()); + let mut ctx = self.engine_context(query_ctx)?; + let query_ctx = ctx.query_ctx(); let plan = if let Some(receiver_injector) = self.plugins.get::() { @@ -552,54 +551,15 @@ impl QueryEngine for DatafusionQueryEngine { self.state.read_table(table).map_err(Into::into) } - fn engine_context(&self, query_ctx: QueryContextRef) -> QueryEngineContext { + fn engine_context(&self, query_ctx: QueryContextRef) -> Result { let mut state = self.state.session_state(); + let query_ctx = self.state.apply_query_options(&mut state, query_ctx)?; state.config_mut().set_extension(query_ctx.clone()); state.config_mut().set_extension(self.state.clone()); - // note that hints in "x-greptime-hints" is automatically parsed - // and set to query context's extension, so we can get it from query context. - if let Some(parallelism) = query_ctx.extension(QUERY_PARALLELISM_HINT) { - if let Ok(n) = parallelism.parse::() { - if n > 0 { - let new_cfg = state.config().clone().with_target_partitions(n as usize); - *state.config_mut() = new_cfg; - } - } else { - common_telemetry::warn!( - "Failed to parse query_parallelism: {}, using default value", - parallelism - ); - } - } - // configure execution options state.config_mut().options_mut().execution.time_zone = Some(query_ctx.timezone().to_string()); - // usually it's impossible to have both `set variable` set by sql client and - // hint in header by grpc client, so only need to deal with them separately - if query_ctx.configuration_parameter().allow_query_fallback() { - state - .config_mut() - .options_mut() - .extensions - .insert(DistPlannerOptions { - allow_query_fallback: true, - }); - } else if let Some(fallback) = query_ctx.extension(QUERY_FALLBACK_HINT) { - // also check the query context for fallback hint - // if it is set, we will enable the fallback - if fallback.to_lowercase().parse::().unwrap_or(false) { - state - .config_mut() - .options_mut() - .extensions - .insert(DistPlannerOptions { - allow_query_fallback: true, - }); - } - } - state .config_mut() .options_mut() @@ -639,7 +599,7 @@ impl QueryEngine for DatafusionQueryEngine { } } - QueryEngineContext::new(state, query_ctx) + Ok(QueryEngineContext::new(state, query_ctx)) } fn engine_state(&self) -> &QueryEngineState { @@ -1194,7 +1154,7 @@ mod tests { .as_any() .downcast_ref::() .unwrap(); - let engine_ctx = engine.engine_context(QueryContext::arc()); + let engine_ctx = engine.engine_context(QueryContext::arc()).unwrap(); let state = engine_ctx.state(); let schema = Arc::new(datatypes::schema::Schema::new(vec![ColumnSchema::new( @@ -1254,7 +1214,7 @@ mod tests { .as_any() .downcast_ref::() .unwrap(); - let engine_ctx = engine.engine_context(QueryContext::arc()); + let engine_ctx = engine.engine_context(QueryContext::arc()).unwrap(); let state = engine_ctx.state(); assert!( diff --git a/src/query/src/dist_plan/analyzer.rs b/src/query/src/dist_plan/analyzer.rs index d90e5d31f9f..0c656b7b56d 100644 --- a/src/query/src/dist_plan/analyzer.rs +++ b/src/query/src/dist_plan/analyzer.rs @@ -15,6 +15,7 @@ use std::collections::{BTreeMap, BTreeSet, HashSet}; use std::sync::Arc; +use common_function::function::FunctionContext; use common_telemetry::debug; use datafusion::config::{ConfigExtension, ExtensionOptions}; use datafusion::datasource::DefaultTableSource; @@ -122,8 +123,19 @@ impl AnalyzerRule for DistPlannerAnalyzer { ) -> datafusion_common::Result { expr_depth::validate_plan(&plan)?; let mut config = config.clone(); - // Aligned with the behavior in `datafusion_optimizer::OptimizerContext::new()`. - config.optimizer.filter_null_join_keys = true; + // Preserve the historical pre-MergeScan default unless explicitly set. + let explicit = config + .extensions + .get::() + .is_some_and(|context| { + context + .query_ctx + .extension("datafusion.optimizer.filter_null_join_keys") + .is_some() + }); + if !explicit { + config.optimizer.filter_null_join_keys = true; + } let config = Arc::new(config); let opt = config.extensions.get::(); let allow_fallback = opt.map(|o| o.allow_query_fallback).unwrap_or(false); diff --git a/src/query/src/optimizer/constant_term.rs b/src/query/src/optimizer/constant_term.rs index 915691604f9..4c487680c16 100644 --- a/src/query/src/optimizer/constant_term.rs +++ b/src/query/src/optimizer/constant_term.rs @@ -470,7 +470,7 @@ mod tests { .await .unwrap(); - let engine_ctx = engine.engine_context(query_ctx); + let engine_ctx = engine.engine_context(query_ctx).unwrap(); let state = engine_ctx.state(); let analyzed_plan = state diff --git a/src/query/src/planner.rs b/src/query/src/planner.rs index 44cd1653cee..b10790634b1 100644 --- a/src/query/src/planner.rs +++ b/src/query/src/planner.rs @@ -56,6 +56,7 @@ use crate::error::{ }; use crate::log_query::planner::LogQueryPlanner; use crate::optimizer::expr_depth; +use crate::options::ScheduledTimeExtension; use crate::parser::{DEFAULT_LOOKBACK_STRING, PromQuery, QueryLanguageParser, QueryStatement}; use crate::promql::planner::PromPlanner; use crate::query_engine::{DefaultPlanDecoder, QueryEngineState}; @@ -96,24 +97,40 @@ impl DfLogicalPlanner { /// in the query context. fn derive_session_state_with_scheduled_time( &self, - query_ctx: &QueryContextRef, - ) -> Result { + query_ctx: QueryContextRef, + ) -> Result { let extensions = query_ctx.extensions(); - match crate::options::parse_scheduled_time_datetime(&extensions)? { + let mut state = match crate::options::parse_scheduled_time_datetime(&extensions)? { Some(dt) => { let execution_props = self .session_state .execution_props() .clone() .with_query_execution_start_time(dt); - Ok( - SessionStateBuilder::new_from_existing(self.session_state.clone()) - .with_execution_props(execution_props) - .build(), - ) + SessionStateBuilder::new_from_existing(self.session_state.clone()) + .with_execution_props(execution_props) + .build() } - None => Ok(self.session_state.clone()), - } + None => self.session_state.clone(), + }; + let query_ctx = self + .engine_state + .apply_query_options(&mut state, query_ctx)?; + state.config_mut().set_extension(query_ctx.clone()); + state.config_mut().set_extension(self.engine_state.clone()); + state + .config_mut() + .options_mut() + .extensions + .insert(ScheduledTimeExtension { + scheduled_time: crate::options::scheduled_time_from_ctx(&query_ctx), + }); + let config_options = state.config_options().clone(); + let _ = state + .execution_props_mut() + .config_options + .insert(config_options); + Ok(QueryEngineContext::new(state, query_ctx)) } /// Basically the same with `explain_to_plan` in DataFusion, but adapted to Greptime's @@ -221,7 +238,9 @@ impl DfLogicalPlanner { .fail()?; } - let scheduled_state = self.derive_session_state_with_scheduled_time(&query_ctx)?; + let engine_context = self.derive_session_state_with_scheduled_time(query_ctx)?; + let scheduled_state = engine_context.state().clone(); + let query_ctx = engine_context.query_ctx(); let table_provider = DfTableSourceProvider::new( self.engine_state.catalog_manager().clone(), self.engine_state.disallow_cross_catalog_query(), @@ -244,7 +263,7 @@ impl DfLogicalPlanner { ) .await?; - let config_options = self.session_state.config().options(); + let config_options = scheduled_state.config().options(); let parser_options = &config_options.sql_parser; let parser_options = ParserOptions { map_string_types_to_utf8view: false, @@ -275,10 +294,9 @@ impl DfLogicalPlanner { .await?; // Optimize logical plan by extension rules - let context = QueryEngineContext::new(scheduled_state, query_ctx); let plan = self .engine_state - .optimize_by_extension_rules(plan, &context)?; + .optimize_by_extension_rules(plan, &engine_context)?; common_telemetry::debug!("Logical planner, optimize result: {plan}"); Ok(plan) @@ -293,16 +311,17 @@ impl DfLogicalPlanner { normalize_ident: bool, query_ctx: QueryContextRef, ) -> Result { - let scheduled_state = self.derive_session_state_with_scheduled_time(&query_ctx)?; + let engine_context = self.derive_session_state_with_scheduled_time(query_ctx)?; + let query_ctx = engine_context.query_ctx(); let context_provider = DfContextProviderAdapter::try_new( self.engine_state.clone(), - scheduled_state, + engine_context.state().clone(), None, query_ctx, ) .await?; - let config_options = self.session_state.config().options(); + let config_options = engine_context.state().config().options(); let parser_options = &config_options.sql_parser; let parser_options: ParserOptions = ParserOptions { map_string_types_to_utf8view: false, @@ -317,7 +336,9 @@ impl DfLogicalPlanner { #[tracing::instrument(skip_all)] async fn plan_pql(&self, stmt: &EvalStmt, query_ctx: QueryContextRef) -> Result { - let mut scheduled_state = self.derive_session_state_with_scheduled_time(&query_ctx)?; + let engine_context = self.derive_session_state_with_scheduled_time(query_ctx)?; + let query_ctx = engine_context.query_ctx(); + let mut scheduled_state = engine_context.state().clone(); let promql_annotations = query_ctx.remote_query_id().map(promql_annotation_collector); if let Some(collector) = &promql_annotations { scheduled_state @@ -708,22 +729,19 @@ impl LogicalPlanner for DfLogicalPlanner { query: LogQuery, query_ctx: QueryContextRef, ) -> Result { - let plan_decoder = Arc::new(DefaultPlanDecoder::new( - self.session_state.clone(), - &query_ctx, - )?); + let engine_context = self.derive_session_state_with_scheduled_time(query_ctx)?; + let query_ctx = engine_context.query_ctx(); + let state = engine_context.state().clone(); + let plan_decoder = Arc::new(DefaultPlanDecoder::new(state.clone(), &query_ctx)?); let table_provider = DfTableSourceProvider::new( self.engine_state.catalog_manager().clone(), self.engine_state.disallow_cross_catalog_query(), query_ctx, plan_decoder, - self.session_state - .config_options() - .sql_parser - .enable_ident_normalization, + state.config_options().sql_parser.enable_ident_normalization, ); - let mut planner = LogQueryPlanner::new(table_provider, self.session_state.clone()); + let mut planner = LogQueryPlanner::new(table_provider, state); planner .query_to_plan(query) .await @@ -964,7 +982,7 @@ mod tests { .plan(&stmt, query_ctx.clone()) .await .unwrap(); - let context = engine.engine_context(query_ctx.clone()); + let context = engine.engine_context(query_ctx.clone()).unwrap(); let state = context.state(); state .analyzer() diff --git a/src/query/src/query_engine.rs b/src/query/src/query_engine.rs index c512b9c1e5e..0355dc2294d 100644 --- a/src/query/src/query_engine.rs +++ b/src/query/src/query_engine.rs @@ -97,7 +97,7 @@ pub trait QueryEngine: Send + Sync { fn read_table(&self, table: TableRef) -> Result; /// Create a [`QueryEngineContext`]. - fn engine_context(&self, query_ctx: QueryContextRef) -> QueryEngineContext; + fn engine_context(&self, query_ctx: QueryContextRef) -> Result; /// Retrieve the query engine state [`QueryEngineState`] fn engine_state(&self) -> &QueryEngineState; diff --git a/src/query/src/query_engine/default_serializer.rs b/src/query/src/query_engine/default_serializer.rs index 76f2bab125d..b14c5648bc3 100644 --- a/src/query/src/query_engine/default_serializer.rs +++ b/src/query/src/query_engine/default_serializer.rs @@ -433,6 +433,7 @@ mod tests { let plan_decoder = engine .engine_context(QueryContext::arc()) + .unwrap() .new_plan_decoder() .unwrap(); let catalog_list = Arc::new(DummyCatalogList::with_table_provider(table_provider)); @@ -525,6 +526,7 @@ mod tests { let table_provider = Arc::new(MemTable::try_new(schema, vec![vec![]]).unwrap()); let plan_decoder = engine .engine_context(QueryContext::arc()) + .unwrap() .new_plan_decoder() .unwrap(); @@ -831,6 +833,7 @@ mod tests { .unwrap(); let plan_decoder = engine .engine_context(query_ctx.clone()) + .unwrap() .new_plan_decoder() .unwrap(); let decoded = plan_decoder @@ -898,6 +901,7 @@ mod tests { .unwrap(); let plan_decoder = engine .engine_context(query_ctx.clone()) + .unwrap() .new_plan_decoder() .unwrap(); let catalog_list = Arc::new(DummyCatalogList::with_table_provider(Arc::new( @@ -964,6 +968,7 @@ mod tests { fn plan_decoder(engine: &dyn crate::QueryEngine) -> SubstraitPlanDecoderRef { engine .engine_context(QueryContext::arc()) + .unwrap() .new_plan_decoder() .unwrap() } @@ -997,7 +1002,11 @@ mod tests { QueryOptions::default(), ) .query_engine(); - let mut state = engine.engine_context(resolver_ctx.clone()).state().clone(); + let mut state = engine + .engine_context(resolver_ctx.clone()) + .unwrap() + .state() + .clone(); state.config_mut().set_extension(resolver_ctx); let extensions = Extensions::default(); let consumer = MergeScanSubstraitConsumer { @@ -1470,7 +1479,11 @@ mod tests { ) .query_engine(); let mut state = SessionStateBuilder::new_from_existing( - engine.engine_context(QueryContext::arc()).state().clone(), + engine + .engine_context(QueryContext::arc()) + .unwrap() + .state() + .clone(), ) .with_serializer_registry(Arc::new(DefaultSerializer)) .build(); diff --git a/src/query/src/query_engine/state.rs b/src/query/src/query_engine/state.rs index 361a0d2991d..d6c5ee6e5fc 100644 --- a/src/query/src/query_engine/state.rs +++ b/src/query/src/query_engine/state.rs @@ -59,6 +59,7 @@ use crate::dist_plan::{ MergeSortExtensionPlanner, RemoteDynFilterReceiverExtensionPlanner, RemoteDynFilterRegistryLease, }; +use crate::error::InvalidQueryContextExtensionSnafu; use crate::metrics::{QUERY_MEMORY_POOL_REJECTED_TOTAL, QUERY_MEMORY_POOL_USAGE_BYTES}; use crate::optimizer::ExtensionAnalyzerRule; use crate::optimizer::const_normalization::ConstNormalizationRule; @@ -321,7 +322,61 @@ impl QueryEngineState { rules.retain(|rule| rule.name() != name); } - /// Optimize the logical plan by the extension analyzer rules. + /// Applies validated query options to one per-query DataFusion configuration. + pub fn apply_query_options( + &self, + state: &mut SessionState, + query_ctx: QueryContextRef, + ) -> crate::error::Result { + let query_ctx = Arc::new(query_ctx.query_option_snapshot().map_err(|error| { + InvalidQueryContextExtensionSnafu { + reason: error.to_string(), + } + .build() + })?); + let options = query_ctx.extensions(); + const MASTER: &str = "datafusion.optimizer.enable_dynamic_filter_pushdown"; + // The master switch cascades to children, so apply explicit child overrides afterward. + for (key, value) in options + .get_key_value(MASTER) + .into_iter() + .chain(options.iter().filter(|(key, _)| key.as_str() != MASTER)) + { + if key.starts_with("datafusion.") { + state + .config_mut() + .options_mut() + .set(key, value) + .map_err(|error| { + InvalidQueryContextExtensionSnafu { + reason: error.to_string(), + } + .build() + })?; + } + } + if let Some(partitions) = options.get("query.parallelism") { + let partitions = partitions.parse::().map_err(|error| { + InvalidQueryContextExtensionSnafu { + reason: error.to_string(), + } + .build() + })?; + *state.config_mut() = state.config().clone().with_target_partitions(partitions); + } + if let Some(fallback) = options.get("query.allow_query_fallback") { + state + .config_mut() + .options_mut() + .extensions + .insert(DistPlannerOptions { + allow_query_fallback: fallback == "true", + }); + } + Ok(query_ctx) + } + + /// Optimize the logical plan using registered extension analyzer rules. pub fn optimize_by_extension_rules( &self, plan: DfLogicalPlan, @@ -331,8 +386,7 @@ impl QueryEngineState { self.extension_rules .iter() .try_fold(plan, |acc_plan, rule| { - let plan = - rule.analyze(acc_plan, context, self.session_state().config_options())?; + let plan = rule.analyze(acc_plan, context, context.state().config_options())?; expr_depth::validate_plan(&plan)?; Ok(plan) }) diff --git a/src/query/src/sql.rs b/src/query/src/sql.rs index 7527d44d821..260b13c7541 100644 --- a/src/query/src/sql.rs +++ b/src/query/src/sql.rs @@ -376,7 +376,7 @@ async fn query_from_information_schema_dataframe( let view = dataframe.into_view(); let dataframe = SessionContext::new_with_state( query_engine - .engine_context(query_ctx.clone()) + .engine_context(query_ctx.clone())? .state() .clone(), ) @@ -844,6 +844,91 @@ pub async fn show_charsets_dataframe( .await } +/// Shows one supported query optimizer option, returning `None` for legacy SHOW handling. +pub fn show_query_option( + stmt: ShowVariables, + query_ctx: QueryContextRef, + engine: &dyn crate::QueryEngine, +) -> Result> { + let name = sql::ast::object_name_to_string_unquoted(&stmt.variable); + let Some(canonical) = session::query_options::canonical_query_option_name(&name) else { + if name.to_ascii_lowercase().starts_with("query.") + || name.to_ascii_lowercase().starts_with("datafusion.") + { + session::query_options::parse_query_option( + &name, + "", + session::query_options::MAX_PUBLIC_PARALLELISM, + ) + .map_err(|error| { + error::InvalidQueryContextExtensionSnafu { + reason: error.to_string(), + } + .build() + })?; + } + return Ok(None); + }; + let context = engine.engine_context(query_ctx.clone())?; + let value = match canonical.as_str() { + "query.allow_query_fallback" => context + .state() + .config_options() + .extensions + .get::() + .map(|options| options.allow_query_fallback) + .unwrap_or(false) + .to_string(), + "query.enable_remote_dynamic_filter_pushdown" => { + crate::options::remote_dyn_filter_pushdown_enabled_from_extensions( + &context.query_ctx().extensions(), + )? + .to_string() + } + "query.parallelism" => context.state().config().target_partitions().to_string(), + "datafusion.optimizer.filter_null_join_keys" + if context + .state() + .analyzer() + .rules + .iter() + .any(|rule| rule.name() == "DistPlannerAnalyzer") + && context + .query_ctx() + .extension("datafusion.optimizer.filter_null_join_keys") + .is_none() => + { + "true".to_string() + } + _ => context + .state() + .config_options() + .entries() + .into_iter() + .find(|entry| entry.key == canonical) + .and_then(|entry| entry.value) + .ok_or_else(|| { + error::InvalidQueryContextExtensionSnafu { + reason: format!( + "No DataFusion configuration entry for supported option {canonical}" + ), + } + .build() + })?, + }; + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + name.clone(), + ConcreteDataType::string_datatype(), + false, + )])); + let records = RecordBatches::try_from_columns( + schema, + vec![Arc::new(StringVector::from(vec![value])) as _], + ) + .context(error::CreateRecordBatchSnafu)?; + Ok(Some(Output::new_with_record_batches(records))) +} + pub fn show_variable(stmt: ShowVariables, query_ctx: QueryContextRef) -> Result { let variable = stmt.variable.to_string().to_uppercase(); let value = match variable.as_str() { diff --git a/src/servers/src/grpc/flight.rs b/src/servers/src/grpc/flight.rs index ca05ea91844..a7514e889d6 100644 --- a/src/servers/src/grpc/flight.rs +++ b/src/servers/src/grpc/flight.rs @@ -198,6 +198,7 @@ impl FlightCraft for GreptimeRequestHandler { request: Request, ) -> TonicResult>> { let mut hints = hint_headers::extract_hints(request.metadata()); + hints = hint_headers::validate_public_hints(hints)?; hints.extend(extract_flow_extensions(request.metadata())?); let snapshot_seqs = extract_snapshot_seqs(request.metadata())?; let channel = request diff --git a/src/servers/src/grpc/greptime_handler.rs b/src/servers/src/grpc/greptime_handler.rs index c29ecc6493a..81dd4d9adf2 100644 --- a/src/servers/src/grpc/greptime_handler.rs +++ b/src/servers/src/grpc/greptime_handler.rs @@ -81,6 +81,7 @@ impl GreptimeRequestHandler { channel: Channel, ) -> Result { let header = request.header.as_ref(); + let hints = crate::hint_headers::validate_public_hints(hints)?; let query_ctx = create_query_context(channel, header, hints, HashMap::new())?; let query = request.request.context(InvalidQuerySnafu { reason: "Expecting non-empty GreptimeRequest.", diff --git a/src/servers/src/hint_headers.rs b/src/servers/src/hint_headers.rs index 4d37447e01d..7f22e25ec62 100644 --- a/src/servers/src/hint_headers.rs +++ b/src/servers/src/hint_headers.rs @@ -16,6 +16,8 @@ use http::HeaderMap; use session::hints::{HINT_KEYS, HINTS_KEY, HINTS_KEY_PREFIX}; use tonic::metadata::MetadataMap; +use crate::error::{InvalidParameterSnafu, Result}; + pub(crate) fn extract_hints(headers: &T) -> Vec<(String, String)> { let mut hints = Vec::new(); if let Some(value_str) = headers.get(HINTS_KEY) { @@ -23,6 +25,16 @@ pub(crate) fn extract_hints(headers: &T) -> Vec<(String, String) let mut parts = hint.splitn(2, '='); if let (Some(key), Some(value)) = (parts.next(), parts.next()) { hints.push((key.trim().to_string(), value.trim().to_string())); + } else { + let key = hint.trim(); + let key_lower = key.to_ascii_lowercase(); + let owned_prefix = + key_lower.starts_with("query.") || key_lower.starts_with("datafusion."); + if owned_prefix + || session::query_options::canonical_query_option_name(key).is_some() + { + hints.push((key.to_string(), String::new())); + } } }); // If hints are provided in the `x-greptime-hints` header, ignore the rest of the headers @@ -37,6 +49,49 @@ pub(crate) fn extract_hints(headers: &T) -> Vec<(String, String) hints } +/// Validate query options supplied through public protocol hints. +pub(crate) fn validate_public_hints(hints: Vec<(String, String)>) -> Result> { + use std::collections::HashMap; + + let mut normalized = Vec::with_capacity(hints.len()); + let mut seen = HashMap::new(); + for (key, value) in hints { + if session::hints::is_reserved_extension_key(&key) { + continue; + } + // Older Flownode clients send `query_parallelism=0` to request the + // server default; treat it as unset instead of rejecting the request. + if key.eq_ignore_ascii_case("query_parallelism") && value == "0" { + continue; + } + match session::query_options::parse_query_option( + &key, + &value, + session::query_options::MAX_PUBLIC_PARALLELISM, + ) { + Ok(Some((canonical, normalized_value))) => { + if let Some(existing) = seen.insert(canonical.clone(), normalized_value.clone()) + && existing != normalized_value + { + return InvalidParameterSnafu { + reason: format!("Conflicting values for hint `{canonical}`"), + } + .fail(); + } + normalized.push((canonical, normalized_value)); + } + Ok(None) => normalized.push((key, value)), + Err(error) => { + return InvalidParameterSnafu { + reason: error.to_string(), + } + .fail(); + } + } + } + Ok(normalized) +} + pub(crate) trait ToHeaderMap { fn get(&self, key: &str) -> Option<&str>; } diff --git a/src/servers/src/http.rs b/src/servers/src/http.rs index 31ed00d443c..01b13e72087 100644 --- a/src/servers/src/http.rs +++ b/src/servers/src/http.rs @@ -1042,17 +1042,10 @@ impl HttpServer { /// Attaches middlewares and debug routes to the router. /// Callers should call this method after [HttpServer::make_app()]. pub fn build(&self, router: Router) -> Result { - let timeout_layer = if self.options.timeout != Duration::default() { - Some( - ServiceBuilder::new().layer( - DynamicTimeoutLayer::new(self.options.timeout) - .with_status_code_fn(Self::request_timeout_status_code), - ), - ) - } else { - info!("HTTP server timeout is disabled"); - None - }; + let timeout_layer = ServiceBuilder::new().layer( + DynamicTimeoutLayer::new(self.options.timeout) + .with_status_code_fn(Self::request_timeout_status_code), + ); let body_limit_layer = if self.options.body_limit != ReadableSize(0) { Some( ServiceBuilder::new() @@ -1088,7 +1081,7 @@ impl HttpServer { // and we have impl IntoResponse for Error. It will print out more detailed error messages .layer(TraceLayer::new_for_http().on_failure(())) .option_layer(cors_layer) - .option_layer(timeout_layer) + .layer(timeout_layer) .option_layer(body_limit_layer) // memory limit layer - must be before body is consumed .layer(middleware::from_fn_with_state( diff --git a/src/servers/src/http/hints.rs b/src/servers/src/http/hints.rs index 7ef9401cf18..722ac66534c 100644 --- a/src/servers/src/http/hints.rs +++ b/src/servers/src/http/hints.rs @@ -15,22 +15,36 @@ use axum::body::Body; use axum::http::Request; use axum::middleware::Next; -use axum::response::Response; +use axum::response::{IntoResponse, Response}; use common_telemetry::debug; use session::context::QueryContext; use session::hints::is_reserved_extension_key; +use crate::error::{InvalidParameterSnafu, Result}; use crate::hint_headers; pub async fn extract_hints(mut request: Request, next: Next) -> Response { let hints = hint_headers::extract_hints(request.headers()); - if let Some(query_ctx) = request.extensions_mut().get_mut::() { - apply_hints(query_ctx, hints); + let dedicated_read_preference = request + .headers() + .contains_key(&crate::http::header::GREPTIME_DB_HEADER_READ_PREFERENCE); + let hints = match hint_headers::validate_public_hints(hints) { + Ok(hints) => hints, + Err(error) => return error.into_response(), + }; + if let Some(query_ctx) = request.extensions_mut().get_mut::() + && let Err(error) = apply_hints(query_ctx, hints, dedicated_read_preference) + { + return error.into_response(); } next.run(request).await } -fn apply_hints(query_ctx: &mut QueryContext, hints: Vec<(String, String)>) { +fn apply_hints( + query_ctx: &mut QueryContext, + hints: Vec<(String, String)>, + dedicated_read_preference: bool, +) -> Result<()> { for (key, value) in hints { if is_reserved_extension_key(&key) { debug!( @@ -39,8 +53,21 @@ fn apply_hints(query_ctx: &mut QueryContext, hints: Vec<(String, String)>) { ); continue; } - query_ctx.set_extension(key, value); + if key.eq_ignore_ascii_case(session::hints::READ_PREFERENCE_HINT) { + if !dedicated_read_preference { + let preference = value.parse::().map_err(|_| { + InvalidParameterSnafu { + reason: format!("Invalid read_preference hint: {value}"), + } + .build() + })?; + query_ctx.set_read_preference(preference); + } + } else { + query_ctx.set_extension(key, value); + } } + Ok(()) } #[cfg(test)] @@ -86,7 +113,9 @@ mod tests { ), ("ttl".to_string(), "7d".to_string()), ], - ); + false, + ) + .unwrap(); assert_eq!( query_ctx.remote_query_id(), diff --git a/src/servers/src/http/read_preference.rs b/src/servers/src/http/read_preference.rs index 04e9771b890..aeeddbace36 100644 --- a/src/servers/src/http/read_preference.rs +++ b/src/servers/src/http/read_preference.rs @@ -25,14 +25,13 @@ use crate::http::header::GREPTIME_DB_HEADER_READ_PREFERENCE; /// Extract read preference from the request headers. pub async fn extract_read_preference(mut request: Request, next: Next) -> Response { - let read_preference = request + if let Some(read_preference) = request .headers() .get(&GREPTIME_DB_HEADER_READ_PREFERENCE) .and_then(|header| header.to_str().ok()) .and_then(|s| ReadPreference::from_str(s).ok()) - .unwrap_or_default(); - - if let Some(query_ctx) = request.extensions_mut().get_mut::() { + && let Some(query_ctx) = request.extensions_mut().get_mut::() + { common_telemetry::debug!("Setting read preference to {}", read_preference); query_ctx.set_read_preference(read_preference); } diff --git a/src/servers/src/mysql/federated.rs b/src/servers/src/mysql/federated.rs index e6c4890df93..b594c172959 100644 --- a/src/servers/src/mysql/federated.rs +++ b/src/servers/src/mysql/federated.rs @@ -280,8 +280,31 @@ fn check_show_variables(query: &str) -> Option { )) } else if SHOW_LOWER_CASE_PATTERN.is_match(query) { Some(show_variables("lower_case_table_names", "0")) - } else if SHOW_VARIABLES_LIKE_PATTERN.is_match(query) { - Some(show_variables("", "")) + } else if let Some(matched) = SHOW_VARIABLES_LIKE_PATTERN.find(query) { + let mut remainder = strip_leading_comments(&query[matched.end()..]); + while let Some(rest) = remainder.strip_prefix(';') { + remainder = strip_leading_comments(rest); + } + if remainder.is_empty() { + Some(show_variables("", "")) + } else { + let remainder = remainder.trim_start(); + // Filter keywords need an identifier boundary; named variables do not. + let looks_like_filter = ["WHERE", "LIKE"].iter().any(|keyword| { + remainder + .get(..keyword.len()) + .is_some_and(|prefix| prefix.eq_ignore_ascii_case(keyword)) + && remainder[keyword.len()..] + .chars() + .next() + .is_none_or(|c| !(c == '_' || c == '$' || c.is_alphanumeric())) + }); + if looks_like_filter { + Some(show_variables("", "")) + } else { + None + } + } } else { None }; @@ -554,6 +577,44 @@ mod test { use super::*; + #[test] + fn test_show_variables_where_filter_is_absorbed() { + // Old Connector/J `loadServerVariables` issues `SHOW VARIABLES WHERE + // Variable_name = '...'`; the federated path answers with an empty table + // instead of falling through to our parser, which would reject it. + let session = Arc::new(Session::new(None, Channel::Mysql, Default::default(), 0)); + let output = check( + "show variables where Variable_name = 'sql_mode'", + QueryContext::arc(), + session, + ) + .expect("WHERE-filtered SHOW VARIABLES should be absorbed"); + assert!(matches!(output.data, OutputData::RecordBatches(_))); + } + + #[test] + fn test_show_variables_named_variable_falls_back() { + // A named variable (e.g. `SHOW VARIABLES MAX_EXECUTION_TIME`, + // `SHOW VARIABLES query.parallelism`) reaches the real handler. + let session = Arc::new(Session::new(None, Channel::Mysql, Default::default(), 0)); + assert!( + check( + "show variables max_execution_time", + QueryContext::arc(), + session.clone() + ) + .is_none() + ); + assert!( + check( + "show variables query.parallelism", + QueryContext::arc(), + session + ) + .is_none() + ); + } + #[test] fn test_check_abnormal() { let session = Arc::new(Session::new(None, Channel::Mysql, Default::default(), 0)); diff --git a/src/servers/src/postgres/handler.rs b/src/servers/src/postgres/handler.rs index 7bf1edca705..9d169811984 100644 --- a/src/servers/src/postgres/handler.rs +++ b/src/servers/src/postgres/handler.rs @@ -96,7 +96,7 @@ impl SimpleQueryHandler for PostgresServerHandlerInner { } statements .iter() - .map(|s| s.to_string()) + .map(statement_to_query) .collect::>() .join(";") } else { @@ -131,6 +131,24 @@ impl SimpleQueryHandler for PostgresServerHandlerInner { } } +fn statement_to_query(statement: &SqlParserStatement) -> String { + if let SqlParserStatement::ShowVariable { variable } = statement + && variable.len() > 1 + && (variable[0].value.eq_ignore_ascii_case("query") + || variable[0].value.eq_ignore_ascii_case("datafusion")) + { + return format!( + "SHOW {}", + variable + .iter() + .map(ToString::to_string) + .collect::>() + .join(".") + ); + } + statement.to_string() +} + async fn send_warning_opt(client: &mut C, query_context: QueryContextRef) -> PgWireResult<()> where C: Sink + Unpin + Send + Sync, @@ -352,7 +370,7 @@ impl QueryParser for DefaultQueryParser { } let first_stmt = statements.remove(0); let format = check_copy_to_stdout(&first_stmt); - (first_stmt.to_string(), format) + (statement_to_query(&first_stmt), format) } else { // bypass the error: it can run into error because of different // versions of sqlparser diff --git a/src/session/src/context.rs b/src/session/src/context.rs index 040d33f2405..43cd5434ca3 100644 --- a/src/session/src/context.rs +++ b/src/session/src/context.rs @@ -78,6 +78,8 @@ pub struct QueryContext { sql_dialect: Arc, #[builder(default)] extensions: HashMap, + #[builder(setter(skip))] + query_options_snapshotted: bool, /// The configuration parameter are used to store the parameters that are set by the user #[builder(default)] configuration_parameter: Arc, @@ -481,6 +483,71 @@ impl QueryContext { self.extensions.clone() } + /// Returns normalized query options from the request over session defaults. + pub fn effective_query_options( + &self, + ) -> Result, crate::session_config::Error> { + let mut request_options = HashMap::new(); + for (key, value) in &self.extensions { + if crate::hints::is_reserved_extension_key(key) { + continue; + } + // Older Frontends and Flownodes forward `query_parallelism=0` (their + // "use the server default" encoding) untouched inside the wire + // query context, so accept it as unset instead of failing the + // query on the Datanode side. The canonical + // `query.parallelism = 0` stays rejected. + if key.eq_ignore_ascii_case("query_parallelism") && value == "0" { + continue; + } + // Wire decoding preserves legacy frontend nonzero values; public entry points enforce the limit. + if let Some((canonical, normalized)) = + crate::query_options::parse_query_option(key, value, u16::MAX)? + { + if request_options + .get(&canonical) + .is_some_and(|previous| previous != &normalized) + { + return crate::session_config::InvalidConfigValueSnafu { + name: key, + value, + hint: "Conflicting aliases for the same query option", + } + .fail(); + } + request_options.insert(canonical, normalized); + } + } + let mut effective = if self.query_options_snapshotted { + HashMap::new() + } else { + self.configuration_parameter.query_options() + }; + effective.extend(request_options); + Ok(effective) + } + + /// Creates a per-query immutable option snapshot while retaining shared mutable session data. + pub fn query_option_snapshot(&self) -> Result { + let mut snapshot = self.clone(); + let options = snapshot.effective_query_options()?; + snapshot + .extensions + .retain(|key, _| crate::query_options::canonical_query_option_name(key).is_none()); + // Older Datanodes only understand the legacy `query_parallelism` hint key, so + // mirror the canonical parallelism override for the wire format while keeping + // the canonical key for in-process consumers. + let legacy_parallelism = options + .get("query.parallelism") + .map(|parallelism| ("query_parallelism".to_string(), parallelism.clone())); + snapshot.extensions.extend(options); + if let Some((key, value)) = legacy_parallelism { + snapshot.extensions.entry(key).or_insert(value); + } + snapshot.query_options_snapshotted = true; + Ok(snapshot) + } + /// Default to double quote and fallback to back quote pub fn quote_style(&self) -> char { if self.sql_dialect().is_delimited_identifier_start('"') { @@ -676,6 +743,7 @@ impl QueryContextBuilder { .sql_dialect .unwrap_or_else(|| Arc::new(GreptimeDbDialect {})), extensions, + query_options_snapshotted: false, configuration_parameter: self .configuration_parameter .unwrap_or_else(|| Arc::new(ConfigurationVariables::default())), @@ -740,7 +808,7 @@ pub struct ConfigurationVariables { postgres_bytea_output: ArcSwap, pg_datestyle_format: ArcSwap<(PGDateTimeStyle, PGDateOrder)>, pg_intervalstyle_format: ArcSwap, - allow_query_fallback: ArcSwap, + query_options: RwLock>, } impl Clone for ConfigurationVariables { @@ -749,7 +817,7 @@ impl Clone for ConfigurationVariables { postgres_bytea_output: ArcSwap::new(self.postgres_bytea_output.load().clone()), pg_datestyle_format: ArcSwap::new(self.pg_datestyle_format.load().clone()), pg_intervalstyle_format: ArcSwap::new(self.pg_intervalstyle_format.load().clone()), - allow_query_fallback: ArcSwap::new(self.allow_query_fallback.load().clone()), + query_options: RwLock::new(self.query_options.read().unwrap().clone()), } } } @@ -784,11 +852,43 @@ impl ConfigurationVariables { } pub fn allow_query_fallback(&self) -> bool { - **self.allow_query_fallback.load() + self.query_options().get("query.allow_query_fallback") == Some(&"true".to_string()) } pub fn set_allow_query_fallback(&self, allow: bool) { - self.allow_query_fallback.swap(Arc::new(allow)); + self.query_options + .write() + .unwrap() + .insert("query.allow_query_fallback".to_string(), allow.to_string()); + } + + /// Validates and stores one canonicalized per-session query option. + pub fn set_query_option( + &self, + key: &str, + value: &str, + ) -> Result<(), crate::session_config::Error> { + let (canonical, normalized) = crate::query_options::parse_query_option( + key, + value, + crate::query_options::MAX_PUBLIC_PARALLELISM, + )? + .ok_or_else(|| crate::session_config::Error::InvalidConfigValue { + name: key.to_string(), + value: value.to_string(), + hint: "Not a supported query option".to_string(), + location: snafu::Location::new(file!(), line!(), column!()), + })?; + self.query_options + .write() + .unwrap() + .insert(canonical, normalized); + Ok(()) + } + + /// Returns the session's canonical query option overrides. + pub fn query_options(&self) -> HashMap { + self.query_options.read().unwrap().clone() } } @@ -801,6 +901,37 @@ mod test { use crate::Session; use crate::context::{Channel, *}; + #[test] + fn test_query_parallelism_wire_accepts_legacy_value() { + let mut context = QueryContext::with(DEFAULT_CATALOG_NAME, "public"); + context.set_extension("query_parallelism", "2048"); + + assert_eq!( + context + .effective_query_options() + .unwrap() + .get("query.parallelism"), + Some(&"2048".to_string()) + ); + } + + #[test] + fn test_set_query_parallelism_public_limit() { + let context = QueryContext::with(DEFAULT_CATALOG_NAME, "public"); + assert!( + context + .configuration_parameter() + .set_query_option("query.parallelism", "1025") + .is_err() + ); + assert!( + context + .configuration_parameter() + .set_query_option("query.parallelism", "1024") + .is_ok() + ); + } + #[test] fn test_session() { let session = Session::new( diff --git a/src/session/src/lib.rs b/src/session/src/lib.rs index 40d2e3e67fe..478458180ef 100644 --- a/src/session/src/lib.rs +++ b/src/session/src/lib.rs @@ -16,6 +16,7 @@ pub mod context; pub mod hints; pub mod protocol_ctx; pub mod query_id; +pub mod query_options; pub mod session_config; pub mod table_name; diff --git a/src/session/src/query_options.rs b/src/session/src/query_options.rs new file mode 100644 index 00000000000..dca6456b0be --- /dev/null +++ b/src/session/src/query_options.rs @@ -0,0 +1,116 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use crate::session_config::{Error, InvalidConfigValueSnafu}; + +/// Returns a canonical query option alias or a DataFusion optimizer option name. +pub fn canonical_query_option_name(key: &str) -> Option { + let key = key.to_ascii_lowercase(); + match key.as_str() { + "query_parallelism" | "query.parallelism" => Some("query.parallelism".to_string()), + "query_fallback" | "allow_query_fallback" | "query.allow_query_fallback" => { + Some("query.allow_query_fallback".to_string()) + } + "query.enable_remote_dynamic_filter_pushdown" => { + Some("query.enable_remote_dynamic_filter_pushdown".to_string()) + } + _ if key.starts_with("datafusion.optimizer.") => Some(key), + _ => None, + } +} + +/// Maximum query parallelism accepted through public query-option inputs. +pub const MAX_PUBLIC_PARALLELISM: u16 = 1024; + +/// Validate a query option, returning its canonical name and normalized value. +pub fn parse_query_option( + key: &str, + value: &str, + max_parallelism: u16, +) -> Result, Error> { + let key_lower = key.to_ascii_lowercase(); + let canonical = match canonical_query_option_name(key) { + Some(name) => name, + None if key_lower.starts_with("query.") || key_lower.starts_with("datafusion.") => { + return InvalidConfigValueSnafu { + name: key, + value, + hint: "Unknown or unsupported query option", + } + .fail(); + } + None => return Ok(None), + }; + + let normalized = if canonical == "query.parallelism" { + match value.parse::() { + Ok(n) if (1..=max_parallelism).contains(&n) => n.to_string(), + _ => { + return InvalidConfigValueSnafu { + name: key, + value, + hint: format!("query.parallelism must be between 1 and {max_parallelism}"), + } + .fail(); + } + } + } else if canonical.starts_with("query.") { + match value.to_ascii_lowercase().as_str() { + "true" => "true".to_string(), + "false" => "false".to_string(), + _ => { + return InvalidConfigValueSnafu { + name: key, + value, + hint: "Expected true or false", + } + .fail(); + } + } + } else { + let mut options = datafusion_common::config::ConfigOptions::default(); + let set_result = options.set(&canonical, value).or_else(|error| { + if value.eq_ignore_ascii_case("true") || value.eq_ignore_ascii_case("false") { + options.set(&canonical, &value.to_ascii_lowercase()) + } else { + Err(error) + } + }); + if let Err(error) = set_result { + return InvalidConfigValueSnafu { + name: key, + value, + hint: error.to_string(), + } + .fail(); + } + match options + .entries() + .into_iter() + .find(|entry| entry.key == canonical) + .and_then(|entry| entry.value) + { + Some(value) => value, + None => { + return InvalidConfigValueSnafu { + name: key, + value, + hint: "No matching DataFusion configuration entry", + } + .fail(); + } + } + }; + Ok(Some((canonical, normalized))) +} diff --git a/src/sql/src/ast.rs b/src/sql/src/ast.rs index 122740987a0..f9e879c7569 100644 --- a/src/sql/src/ast.rs +++ b/src/sql/src/ast.rs @@ -32,3 +32,37 @@ impl ObjectNamePartExt for ObjectNamePart { ident.value.clone() } } + +/// Formats an object name by joining its identifier values without quotes. +pub fn object_name_to_string_unquoted(name: &ObjectName) -> String { + name.0 + .iter() + .map(ObjectNamePartExt::to_string_unquoted) + .collect::>() + .join(".") +} + +#[cfg(test)] +mod tests { + use sqlparser::ast::{Ident, ObjectName, ObjectNamePart}; + + use super::object_name_to_string_unquoted; + + #[test] + fn test_object_name_to_string_unquoted() { + let quoted = ObjectName(vec![ObjectNamePart::Identifier(Ident::with_quote( + '"', + "query.parallelism", + ))]); + let unquoted = ObjectName(vec![ + ObjectNamePart::Identifier(Ident::new("query")), + ObjectNamePart::Identifier(Ident::new("parallelism")), + ]); + + assert_eq!("query.parallelism", object_name_to_string_unquoted("ed)); + assert_eq!( + object_name_to_string_unquoted(&unquoted), + object_name_to_string_unquoted("ed) + ); + } +} diff --git a/tests/cases/standalone/common/basic.result b/tests/cases/standalone/common/basic.result index 37faa575fce..4c9f028eb7c 100644 --- a/tests/cases/standalone/common/basic.result +++ b/tests/cases/standalone/common/basic.result @@ -196,11 +196,11 @@ affected_rows: 0 -- SQLNESS PROTOCOL MYSQL SHOW VARIABLES MAX_EXECUTION_TIME; -+---------------+-------+ -| Variable_name | Value | -+---------------+-------+ -| | | -+---------------+-------+ ++--------------------+ +| MAX_EXECUTION_TIME | ++--------------------+ +| 2000 | ++--------------------+ -- SQLNESS PROTOCOL MYSQL SET MAX_EXECUTION_TIME = 0; From 3de8610e2e88516c30825a2c156e5b2d4f321eaf Mon Sep 17 00:00:00 2001 From: localhost Date: Fri, 16 Oct 2026 07:01:29 +0000 Subject: [PATCH 021/239] fix(ci): compare release validation against the previous stable version (#9636) Ancestry-based tag selection skips stable tags that live only on older release branches. Signed-off-by: paomian --- .github/scripts/previous-release-tag-test.py | 133 +++++++++++++++++++ .github/scripts/previous-release-tag.py | 91 +++++++++++++ .github/workflows/checks.yml | 2 + .github/workflows/release.yml | 14 +- 4 files changed, 235 insertions(+), 5 deletions(-) create mode 100644 .github/scripts/previous-release-tag-test.py create mode 100644 .github/scripts/previous-release-tag.py diff --git a/.github/scripts/previous-release-tag-test.py b/.github/scripts/previous-release-tag-test.py new file mode 100644 index 00000000000..5b1465eef93 --- /dev/null +++ b/.github/scripts/previous-release-tag-test.py @@ -0,0 +1,133 @@ +#!/usr/bin/env python3 +# 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. + +import importlib.util +import subprocess +import sys +import unittest +from pathlib import Path + + +def load_module(): + path = Path(__file__).resolve().parent / "previous-release-tag.py" + spec = importlib.util.spec_from_file_location("previous_release_tag", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +TAGS = [ + "v1.1.0", + "v1.2.0", + "v1.2.1", + "v1.2.1^{}", + "v1.3.0-alpha.1", + "v1.3.0-beta.1", + "v1.3.0-beta.2", + "v1.2.0-nightly-20260706", +] + + +class PreviousReleaseTagTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + cls.mod = load_module() + + def test_beta_uses_latest_older_stable_not_an_ancestor(self) -> None: + self.assertEqual(self.mod.previous_stable_tag("v1.3.0-beta.2", TAGS), "v1.2.1") + + def test_stable_release_skips_itself(self) -> None: + self.assertEqual(self.mod.previous_stable_tag("v1.2.1", TAGS), "v1.2.0") + + def test_same_line_stable_is_not_baseline_for_its_prerelease(self) -> None: + tags = TAGS + ["v1.3.0"] + self.assertEqual(self.mod.previous_stable_tag("v1.3.0-beta.2", tags), "v1.2.1") + self.assertEqual(self.mod.previous_stable_tag("v1.3.0", tags), "v1.2.1") + + def test_nightly_of_unreleased_version_uses_previous_stable(self) -> None: + self.assertEqual( + self.mod.previous_stable_tag("v1.3.0-nightly-20261013", TAGS), + "v1.2.1", + ) + + def test_no_older_stable_tag(self) -> None: + self.assertIsNone(self.mod.previous_stable_tag("v0.1.0", ["v0.1.0", "v1.2.1"])) + + def test_invalid_version(self) -> None: + with self.assertRaises(ValueError): + self.mod.previous_stable_tag("nightly-20261013", TAGS) + + def test_unordered_tags_use_numeric_order(self) -> None: + forward = ["v1.2.1", "v1.9.0", "v1.10.0", "v0.10.0", "v1.9.9"] + backward = list(reversed(forward)) + for tags in (forward, backward): + self.assertEqual(self.mod.previous_stable_tag("v1.10.0", tags), "v1.9.9") + self.assertEqual(self.mod.previous_stable_tag("v1.10.0-rc.1", tags), "v1.9.9") + + def test_multi_digit_components_compare_as_integers(self) -> None: + tags = ["v1.9.0", "v1.10.0", "v0.9.10", "v0.10.0", "v2.9.15", "v2.10.0", "v10.2.0"] + self.assertEqual(self.mod.previous_stable_tag("v1.10.0", tags), "v1.9.0") + self.assertEqual(self.mod.previous_stable_tag("v0.10.0", tags), "v0.9.10") + self.assertEqual(self.mod.previous_stable_tag("v2.10.0", tags), "v2.9.15") + self.assertEqual(self.mod.previous_stable_tag("v10.2.0", ["v10.0.9", "v9.20.0"]), "v10.0.9") + + +class PreviousReleaseTagCliTest(unittest.TestCase): + SCRIPT = Path(__file__).resolve().parent / "previous-release-tag.py" + + def run_cli( + self, version: str, tags: list[str] | None = None, stdin: str | None = None + ) -> subprocess.CompletedProcess[str]: + command = [sys.executable, str(self.SCRIPT), "--version", version] + if tags is not None: + command.extend(tags) + return subprocess.run( + command, + input=stdin, + text=True, + capture_output=True, + check=False, + ) + + def test_stdin_selects_baseline_and_exits_zero(self) -> None: + completed = self.run_cli( + "v1.10.0", + stdin="v1.2.1\nv1.9.0\n\nv1.10.0\nv1.9.0^{}\n", + ) + self.assertEqual(completed.returncode, 0) + self.assertEqual(completed.stdout, "v1.9.0\n") + self.assertEqual(completed.stderr, "") + + def test_positional_tags_ignore_stdin(self) -> None: + completed = self.run_cli("v1.10.0", tags=["v1.9.0"], stdin="v1.9.1\n") + self.assertEqual(completed.returncode, 0) + self.assertEqual(completed.stdout, "v1.9.0\n") + + def test_no_older_tag_exits_zero_with_empty_stdout(self) -> None: + completed = self.run_cli("v0.1.0", stdin="v0.1.0\nv1.2.1\n") + self.assertEqual(completed.returncode, 0) + self.assertEqual(completed.stdout, "") + self.assertEqual(completed.stderr, "") + + def test_invalid_version_exits_one(self) -> None: + completed = self.run_cli("nightly", stdin="v1.2.1\n") + self.assertEqual(completed.returncode, 1) + self.assertEqual(completed.stdout, "") + self.assertIn("not a release version: nightly", completed.stderr) + + +if __name__ == "__main__": + unittest.main() diff --git a/.github/scripts/previous-release-tag.py b/.github/scripts/previous-release-tag.py new file mode 100644 index 00000000000..b8a8f42ef72 --- /dev/null +++ b/.github/scripts/previous-release-tag.py @@ -0,0 +1,91 @@ +#!/usr/bin/env python3 +# 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. + +"""Pick the query-regression baseline tag for a release. + +Compatibility tests do not use this tag. They read ``tests/compatibility/ci.toml``. + +Release branches diverge, so a tag on ``release/v1.3`` is not a descendant of +``v1.2.1``. Ancestry (``git tag --merged``) therefore skips the real previous +release. Compare version numbers instead: the baseline is the greatest +``vX.Y.Z`` tag whose version is strictly older than the release being built. +Prerelease and nightly tags are not baselines. +""" + +from __future__ import annotations + +import argparse +import re +import sys + +STABLE_TAG = re.compile(r"^v(\d+)\.(\d+)\.(\d+)$") +RELEASE_VERSION = re.compile(r"^v(\d+)\.(\d+)\.(\d+)(?:-.+)?$") + + +def _triple(match: re.Match[str]) -> tuple[int, int, int]: + return int(match.group(1)), int(match.group(2)), int(match.group(3)) + + +def release_version(version: str) -> tuple[int, int, int]: + match = RELEASE_VERSION.fullmatch(version.strip()) + if match is None: + raise ValueError(f"not a release version: {version}") + return _triple(match) + + +def previous_stable_tag(version: str, tags: list[str]) -> str | None: + """Return the newest stable tag older than ``version``, or None.""" + current = release_version(version) + best: tuple[int, int, int] | None = None + best_tag: str | None = None + for raw in tags: + tag = raw.strip() + if tag.endswith("^{}"): + tag = tag[: -len("^{}")] + match = STABLE_TAG.fullmatch(tag) + if match is None: + continue + triple = _triple(match) + if triple < current and (best is None or triple > best): + best = triple + best_tag = tag + return best_tag + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--version", required=True, help="Release version, e.g. v1.3.0-beta.2") + parser.add_argument( + "tags", + nargs="*", + help="Candidate tags. Read stdin when omitted.", + ) + args = parser.parse_args(argv) + tags = list(args.tags) + if not tags and not sys.stdin.isatty(): + tags = sys.stdin.read().splitlines() + try: + tag = previous_stable_tag(args.version, tags) + except ValueError as exc: + print(exc, file=sys.stderr) + return 1 + if tag is None: + return 0 + print(tag) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml index 2aa460ec0ca..030a8e2fcbb 100644 --- a/.github/workflows/checks.yml +++ b/.github/workflows/checks.yml @@ -82,6 +82,8 @@ jobs: version: "latest" - name: Run check-version script tests run: uv run --no-project python .github/scripts/check-version-test.py + - name: Run previous-release-tag script tests + run: python3 .github/scripts/previous-release-tag-test.py - name: Run enterprise license check tests run: python3 scripts/check-enterprise-license-test.py - name: Run fuzz orchestration script tests diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 0fb22f23e61..a309c224568 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -415,11 +415,15 @@ jobs: REQUESTED_VALIDATION_POLICY: ${{ inputs.release_validation }} run: | set -euo pipefail - base_commit="${GITHUB_SHA}" - if [[ "${GITHUB_REF_TYPE}" == "tag" ]]; then - base_commit="${GITHUB_SHA}^" - fi - previous_release_tag="$(git tag --merged "${base_commit}" --sort=-v:refname | grep -E '^v[0-9]+\.[0-9]+\.[0-9]+$' | head -n 1 || true)" + # Stable tags on older release branches are not ancestors of this + # commit. Pick the newest vX.Y.Z older than the version being released. + release_version="${{ needs.allocate-runners.outputs.version }}" + previous_release_tag="$( + git ls-remote --tags origin 'refs/tags/v*' \ + | awk '{print $2}' \ + | sed -e 's#^refs/tags/##' \ + | python3 .github/scripts/previous-release-tag.py --version "${release_version}" + )" if [[ -z "${previous_release_tag}" ]]; then previous_release_tag="main" fi From a0188c8c410dec8660f7bcb6c4b2bf98f90f66b7 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Fri, 16 Oct 2026 07:36:35 +0000 Subject: [PATCH 022/239] docs: cover pre-releases and upgrade notes in the release-note skill (#9635) * docs: cover pre-releases and upgrade notes in the release-note skill Signed-off-by: Dennis Zhuang * docs: add style and length rules to the release-note skill Signed-off-by: Dennis Zhuang * docs: drop machine-specific notes from the release-note skill Signed-off-by: Dennis Zhuang * docs: base the first pre-release on the previous release branch point Signed-off-by: Dennis Zhuang * docs: subtract the whole previous line when building a minor release note Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- .../skills/greptimedb-release-note/SKILL.md | 235 +++++++++++++++--- 1 file changed, 199 insertions(+), 36 deletions(-) diff --git a/.agents/skills/greptimedb-release-note/SKILL.md b/.agents/skills/greptimedb-release-note/SKILL.md index 893b99445ce..026ad90cd8d 100644 --- a/.agents/skills/greptimedb-release-note/SKILL.md +++ b/.agents/skills/greptimedb-release-note/SKILL.md @@ -1,6 +1,6 @@ --- name: greptimedb-release-note -description: Generate a GreptimeDB release changelog with git cliff (correct range, subtract already-released patch PRs, rebuild contributors, add human-curated highlights), output to a file, and prepare the docs-repo blog PR. Use when asked to write/generate a GreptimeDB release note or changelog. +description: Generate a GreptimeDB release changelog (minor, patch, or alpha/beta/rc pre-release) with git cliff (correct range, subtract already-released patch PRs, rebuild contributors, add human-curated highlights and upgrade notes), output to a file, and prepare the docs-repo blog PR. Use when asked to write/generate a GreptimeDB release note or changelog. --- # GreptimeDB Release Note / Changelog @@ -27,41 +27,54 @@ a token it still runs, but may be rate-limited / incomplete. ## 2. Determine the previous version (skip nightlies) -From `gh release list --repo GreptimeTeam/greptimedb`, ignore `*-nightly-*`, `-rc.*`, `-beta.*`, -and build-suffixed tags; focus on formal `vX.Y.Z`. +From `gh release list --repo GreptimeTeam/greptimedb`, ignore `*-nightly-*` and build-suffixed +tags (e.g. `v1.3.0-alpha.1-3c0e2a8d5-20260928-1790639248`). - Patch `vX.Y.Z` (Z>0): previous = `vX.Y.(Z-1)`. - New minor `vX.Y.0`: previous = the latest `vX.(Y-1).*` (for `v1.0.0`, the biggest `0.x`). +- Pre-release `vX.Y.Z-{alpha,beta,rc}.N`: previous = the latest pre-release tag of the same + `vX.Y.Z` (e.g. `v1.3.0-beta.1` for `v1.3.0-beta.2`). If there is none (the first alpha), + see "First pre-release" in §3: the previous minor's `.0` may not be published yet, and since + v1.2 a `.0` tag sits on `release/vX.Y`, not on `main`. Double-check with the user. ## 3. Pick the range and generate -Key topology fact: a **minor tag** (e.g. `v1.0.0`) is an ancestor of `main`; **patch tags** -(`v1.0.1`, `v1.0.2`) live on the `release/v1.0` branch and are **NOT** ancestors of `main` -(they are cherry-picks with different SHAs). Verify with +Key topology fact: **patch tags** (`v1.0.1`, `v1.0.2`) live on `release/vX.Y` and are **NOT** +ancestors of `main` (they are cherry-picks with different SHAs). Up to v1.1 the `.0` tag was +the point where `release/vX.Y` left `main`; from v1.2 the branch is cut before the `.0`, and +pre-releases and the `.0` itself are tagged on the branch with backports. Verify with `git merge-base --is-ancestor /main`. -### New minor (`X.Y.0`, cut from main) -Note the two different tags here: the `git cliff` **base is the previous minor `.0` tag** -(e.g. `v1.0.0`, an ancestor of `main`) — **not** the "previous version" from §2 (the latest -patch, e.g. `v1.0.2`), which is only used below to decide which patch PRs to subtract. +### New minor (`X.Y.0`) +The note compares against the whole previous line: everything new since `release/vX.(Y-1)` +left `main`, minus what any published `vX.(Y-1).*` release (pre-releases, `.0`, patches) +already shipped. The line's own pre-releases (`vX.Y.0-beta.*`) are not subtracted; the `.0` +note covers the whole minor. -Base = the **previous minor tag** (e.g. `v1.0.0`); tip = the release commit (= `release/vX.Y` -tip, usually `/main`): +Base = the **previous minor tag** (e.g. `v1.1.0`); tip = the tagged `vX.Y.0` commit, resolved as +in "Pre-release" below. Split the range at the point where `release/vX.Y` left `main`: ``` -GITHUB_TOKEN=$(gh auth token) git cliff .. --tag vX.Y.0 \ - -o /path/CHANGELOG-vX.Y.0.md +bp=$(git merge-base /release/vX.Y /main) ``` -`cliff.toml`'s `ignore_tags` folds in-range nightly tags into the single section. +- `..$bp` is on `main`; generate it with `git cliff`: + ``` + GITHUB_TOKEN=$(gh auth token) git cliff ..$bp --tag vX.Y.0 \ + -o /path/CHANGELOG-vX.Y.0.md + ``` + `cliff.toml`'s `ignore_tags` folds in-range nightly tags into the single section. +- `$bp..` is the release branch: backports, squashed picks, and commits without a PR. + `git cliff` fails on such commits (`Filter call 'sort' failed`), so build this part as a + pre-release: collect the original PRs with the loop in "Pre-release" and add their bullets. + List the commits that carry no PR with + `git log --format='%h %s' $bp.. | grep -vE '\(#[0-9]+\)$'`; drop mechanical ones, and + for a manual port of a `main` PR (the body says `cherry picked from commit` a SHA that is not + on `main`), find that PR by title and cite it. -**Then subtract PRs already shipped in the intermediate patch releases** (`vX.(Y-1).1`, -`.2`, …) — their main-branch commits are inside the range and would duplicate, and the -audience cares about what's new vs the latest patch. Collect the PR set from the patch release -bodies and remove matching lines: -``` -gh release view vX.(Y-1).Z --repo GreptimeTeam/greptimedb | grep -oE 'pull/[0-9]+' -``` -Remove every changelog bullet whose `#NNNN` is in that combined set (a small Python script is -the reliable way — match `pull/)` on lines starting with `*`). +**Then subtract PRs already shipped in published `vX.(Y-1).*` releases**: their `main` commits +are inside the range. Build the set as for the first pre-release (below): the PR links in +every published `vX.(Y-1).*` release body, plus the `(#N)` line endings in the history of the +newest published `vX.(Y-1).*` tag. Remove every changelog bullet whose `#NNNN` is in that set (a +small Python script is the reliable way — match `pull/)` on lines starting with `*`). ### Patch (`X.Y.Z`, Z>0, cherry-picks on the release branch) The previous patch tag is an ancestor of the release branch, so a plain range works **when the @@ -99,11 +112,78 @@ user which they prefer**): If the branch history is otherwise messy, the same "identify the picked PRs and keep only those" rule applies. +### First pre-release (`X.Y.0-alpha.1`, cut from main) +The base is the commit where the previous line's release branch left `main`. It exists when the +alpha is cut, even if `vX.(Y-1).0` does not: +``` +base=$(git merge-base /release/vX.(Y-1) ) +``` +If there is no `release/vX.(Y-1)`, use the newest published release tag that is an ancestor of +the tag (`git merge-base --is-ancestor `). Run `git cliff ..`, then +subtract every PR that already shipped in a published `vX.(Y-1).*` release, pre-releases +included. Take the union of the PR links in those release bodies and the `(#N)` line endings in +the history of the newest published `vX.(Y-1).*` tag: +``` +gh release view --repo GreptimeTeam/greptimedb | grep -oE 'pull/[0-9]+' +git log --format=%B | grep -oE '\(#[0-9]+\)$' +``` +Either source alone misses PRs: a release body can omit a backport that its tag contains, and a +backport commit can name only the backport PR. + +### Pre-release (`X.Y.Z-{alpha,beta,rc}.N` after an earlier pre-release) +**The tip is the tagged commit, not the release branch tip.** Resolve it from the remote and +write this SHA in the docs PR body (§8) as what the note was checked against: +``` +git ls-remote --tags 'vX.Y.Z-beta.N*' +``` +Use the `refs/tags/vX.Y.Z-beta.N` line (its `^{}` line if the tag is annotated), not the +build-suffixed tags it also lists. PRs that land on `release/vX.Y` after the tag +(`git log --oneline ../release/vX.Y`) belong to the next release, even if they are +already merged when you write the note. + +The range `..` usually mixes two kinds of commits: individual backports +(`... [Backport release/vX.Y] (#NNNN)`) and squashed pick commits (`chore: backport selected +changes for ... (#NNNN)`). In both, the original PRs are the `(#NN)` at the end of body lines; +a backport commit can also carry prerequisite PRs. Collect the PR set, falling back to the +subject for commits whose body lists none: +``` +for c in $(git rev-list ..); do + prs=$(git log -1 --format=%b "$c" | grep -oE '\(#[0-9]+\)$') + [ -n "$prs" ] || prs=$(git log -1 --format=%s "$c" | grep -oE '\(#[0-9]+\)$') + echo "$prs" +done | tr -dc '0-9\n' | sort -un +``` +List the bullets under the **original** PRs (title and author of the PR merged to `main`), not +the backport PRs. Build them with patch (a) or (b) above; for (b), fetch +`gh api repos/GreptimeTeam/greptimedb/pulls/ --jq '{number,title,author:.user.login,merged_at}'`, +group by the `cliff.toml` parsers (first match wins), order by `merged_at`, and put `!` titles +under Breaking changes (§6). + +A squashed pick PR can also contain release-only commits that have no PR. List them: +``` +gh api repos/GreptimeTeam/greptimedb/pulls//commits --paginate \ + --jq '.[].commit.message | split("\n")[0]' | grep -vE '\(#[0-9]+\)$' +``` +Mechanical ones (version bump, toolchain or lint adaptations; example: +`fix: adapt release OptionMap iteration to Rust 1.98`) get no bullet. A release-only +commit that changes behavior gets a bullet in its group. + > Flaky-network note for (a): `git cliff` fetches the repo's full commit history from the GitHub > API for author enrichment (paginates back thousands of commits, even for a tiny range). On a > flaky connection it can time out mid-pagination; it caches pages, so just **re-run until it > completes** — each retry resumes from the cache. If it keeps failing, fall back to (b). +### Groups (all paths) +`git cliff` drops PRs whose title isn't a conventional commit (`filter_unconventional`). List +them and add the user-facing ones by hand: +``` +git log --format=%s | grep -E '\(#[0-9]+\)$' | grep -vE '^[a-z]+(\([^)]*\))?!?: ' +``` +Use only the groups in `cliff.toml`. `git cliff` prints a `### Build` section for `build:` commits +because no parser matches them; move a user-facing `build(deps)` fix to Bug Fixes and the rest to +Miscellaneous Tasks. Keep each bullet list free of blank lines; a blank line turns it into a +loose list that renders with different spacing. + ## 4. Rebuild contributor sections after subtracting `cliff` computes `New Contributors` / `All Contributors` over the full range, so after removing @@ -114,6 +194,21 @@ lines, recompute from the **remaining** commit bullets: in-range PR shipped in a patch). Do this in the same script that subtracts the PRs. +When the bullets are built by hand (§3 (b), pre-releases), `git cliff` computes no New +Contributors, and GitHub's generated notes don't either (they only see the backport PRs). An +author is new if none of their merged PRs shipped before, i.e. none appears in the history of the +base tag: +``` +git log --format=%B | grep -oE '\(#[0-9]+\)$' | tr -dc '0-9\n' | sort -u > /tmp/shipped.txt +gh search prs --repo GreptimeTeam/greptimedb --author --merged --limit 1000 \ + --json number --jq '.[].number' | sort -u > /tmp/author.txt +comm -12 /tmp/author.txt /tmp/shipped.txt # empty => new contributor +``` +Run it for every non-bot author in the release, and cite their earliest-merged PR in this +release. "Earliest merged PR is in range" is not the same test: a contributor whose first PR +merged on `main` but was not picked into any release is still new when a later PR of theirs +ships. + ## 5. Verify Cross-check the result against `git log ..`: no subtracted PR remains, and no @@ -122,17 +217,71 @@ line (ask the user; some prefer to keep it). ## 6. Human-curated sections (insert after the `Release date:` line) +**Style.** Follow the Google developer documentation style guide: sentence-case headings, active voice, present +tense for behavior, American spelling, serial comma, code font for identifiers, bold only for UI +names. Keep it short: +- Intro: two sentences at most, what the release covers and its main theme. +- Highlight: one paragraph of up to four sentences, plus an example if it earns one (see below). +- Upgrade note or breaking change: one bullet, one or two sentences: what changes, who is + affected, what to do. +- Every sentence states a fact the reader acts on or checks. Delete sentences that restate the + heading or the previous sentence, transitions, and summaries. +- No marketing words (`significantly`, `powerful`, `seamless`, `robust`, `enhanced`), no vague + verbs (`supports`, `handles`, `improves`; say what happens), no `several`/`many` when the number + is known. Keep the source's hedges; don't add new ones. + +Example: +> Bad: This release significantly enhances write reliability by introducing a robust retry +> mechanism that seamlessly handles transient failures. +> +> Good: In distributed clusters, the frontend retries a failed write up to three times with +> exponential backoff. Set `max_retries = 0` to disable retries. + Mirror past release bodies (`gh release view v1.0.0 --repo GreptimeTeam/greptimedb`): -- **Short intro**, terse, engineer-written tone (no marketing adjectives). -- **`### 👍 Highlights`** — *few, deep* highlights, each with a **working example** - (SQL / TOML config). To get examples right, **read the highlight PR and the related docs - PR/page** in the docs repo (`GreptimeTeam/docs`, often checked out locally) — verify syntax - in `docs/reference/sql/*.md`, `config/*.example.toml`, `config/config.md`. Do **not** mention - implementation details, tiny features, or unfinished/experimental-but-not-ready features. - Always have the user review and edit the highlights; iterate. +- **Short intro** (see Style). +- **`### 👍 Highlights`** — *few, deep* highlights, ranked by user impact. Fixes for silently + wrong query results, lost or invisible data, and security issues rank with new features and + can be highlights. + - Write a fix as its symptom and the affected deployments or usage, not the PR title. Example: + "with no traffic, `errors / requests > 0.05` was true, so the alert fired". + - State whether a feature is on by default. If an example config uses the default values, say + so. + - Give an example (SQL / TOML config) only where it helps readers act. Don't pick highlights, + or the example, by what is easy to demonstrate. + - One highlight covers one change or one user-visible problem. Don't merge unrelated PRs into + one entry because they touch the same component. + - Sources: the PR body's sections on impact, compatibility, and behavior changes, plus the + code and config it touches, not the title alone. Verify examples against the related docs + PR/page in the docs repo (`GreptimeTeam/docs`, often checked out locally), + `docs/reference/sql/*.md`, `config/*.example.toml`, `config/config.md`. + Do **not** mention implementation details, tiny features, or unfinished/experimental-but-not-ready + features. Always have the user review and edit the highlights; iterate. - **Dashboard** subsection — don't just bump the version. Read the dashboard PRs in the bundled release (`gh release view --repo GreptimeTeam/dashboard`; then the linked PRs) and - describe the user-facing change. + describe the user-facing change. Use the names the UI shows, not the PR titles; check them in + the dashboard's locale strings at the bundled tag: + ``` + gh api 'repos/GreptimeTeam/dashboard/contents/src/locale/en-US/menu.ts?ref=' \ + --jq .content | base64 -d + ``` + Screenshots need absolute `https://` URLs: the body is copied to the GitHub release, where + relative and site-root (`/img/...`) paths don't resolve. +- **`### Upgrade notes`** and **`### Breaking changes`** — check every PR in the release for: + `!` titles; `Compatibility`, `Behavior change`, `Breaking` paragraphs in the body (often a + plain `Compatibility:` line, not a heading); unchecked "API changes are backward compatible" / + "Schema or data changes are backward compatible" boxes; new config options and their defaults; + changed defaults; persisted format changes; rolling-upgrade order; new hard limits. These + signals set reading order, not a filter: a limit can sit in plain prose with every box + checked, so read the body of every user-facing `feat`/`fix` PR. Examples of what is easy to + miss: a 512 MiB cap on decompressed request bodies (#9264); a default step for stepless + subqueries (#9481). + - **Upgrade notes**: steps operators take — upgrade order, index rebuilds, settings to review. + - **Breaking changes**: behavior that used to work and now errors or returns different + results. Put the `!` PR bullets here too, after these entries (see `v1.3.0-beta.1`). +- **Benchmark results** (optional) — state the baseline version (often the previous formal + release, not the previous beta), data size, cache state, and hardware. Don't attribute results + to single PRs. Leave out results the benchmark marks as not for publication, and smoke-level + runs. ## 7. Output @@ -144,7 +293,8 @@ it**. The release runbook deletes it after the release + docs PR are done. The release note is also published as a blog post in **`GreptimeTeam/docs`**. - **Ask the user for their local `GreptimeTeam/docs` checkout path**; if they have none, offer to clone (`git clone git@github.com:GreptimeTeam/docs.git `). -- File: `blog/release-X-Y-Z.md` (version with dashes; see `blog/release-1-0-0.md`). +- File: `blog/release-X-Y-Z.md` (version with dashes; see `blog/release-1-0-0.md`). Pre-releases + use dashes throughout: `blog/release-1-3-0-beta-2.md`. - Content = docs frontmatter (below) followed by the GitHub release body, **keeping the `# vX.Y.Z` H1 directly under the frontmatter**. The blog frontmatter has no `title:`, so Docusaurus uses that H1 as the page title and sidebar label — omitting it makes the page @@ -156,11 +306,13 @@ The release note is also published as a blog post in **`GreptimeTeam/docs`**. description: GreptimeDB vX.Y.Z Changelog date: YYYY-MM-DD --- - # vX.Y.Z Release date: ... ``` + The docs CI check `check-release-note-sync` requires the post minus its frontmatter to equal + the GitHub release body byte for byte, so the H1 goes on the line right after the closing + `---`, with no blank line. - The docs working tree may have unrelated WIP — **don't disturb it**: use a `git worktree` off `origin/main`: ``` @@ -169,8 +321,19 @@ The release note is also published as a blog post in **`GreptimeTeam/docs`**. ``` - **Read and follow the current docs repo PR template** (`.github/pull_request_template.md`) from the checkout used for the release. - Fill in every required section and leave reviewer-owned checklist boxes unchecked; - do not rely on hard-coded section names from an older template. + Fill in every required section; do not rely on hard-coded section names from an older + template. The checklist items are the author's to confirm: check the ones you did, and for one + that doesn't apply, check it and say why. For pre-releases, state the tagged SHA (§3) under + Verification. +- **Build the site before opening the PR.** A new worktree has no `node_modules`. Install with + the pnpm version pinned in `package.json` (`packageManager`), as CI does: + ``` + cd /tmp/docs-release-note + pm=$(jq -r .packageManager package.json) # e.g. pnpm@8.6.0 + npx -y "$pm" install --frozen-lockfile + npx -y "$pm" build + ``` + If the build fails, fix it or stop; don't mark the PR ready for review. - Commit with sign-off, push, open a **draft** PR, then remove the worktree: ``` git -C /tmp/docs-release-note add blog/release-X-Y-Z.md From 9089ea7a4b00c7b945caced0192ce91a3f58ff44 Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Fri, 16 Oct 2026 07:50:40 +0000 Subject: [PATCH 023/239] fix(mito): avoid pinning region version when rebuilding indexes (#9597) refactor(mito): avoid pinning region version when rebuilding indexes handle_rebuild_index kept the whole region version alive across the manifest read and schedule_build awaits, which pins memtables and the SST version even though index build tasks only need the region metadata, index options and the candidate SST file handles. Extract those inputs into a small IndexBuildSnapshot and drop the version before awaiting. Signed-off-by: Lei, HUANG --- src/mito2/src/worker/handle_rebuild_index.rs | 60 +++++++++++++------- 1 file changed, 39 insertions(+), 21 deletions(-) diff --git a/src/mito2/src/worker/handle_rebuild_index.rs b/src/mito2/src/worker/handle_rebuild_index.rs index 59f42ef4232..17ba1ece1c8 100644 --- a/src/mito2/src/worker/handle_rebuild_index.rs +++ b/src/mito2/src/worker/handle_rebuild_index.rs @@ -19,6 +19,7 @@ use std::sync::Arc; use api::v1::region::build_index_request; use common_telemetry::{debug, error, warn}; +use store_api::metadata::RegionMetadataRef; use store_api::region_request::RegionBuildIndexRequest; use store_api::storage::{FileId, RegionId}; use tokio::sync::mpsc; @@ -27,7 +28,7 @@ use crate::cache::CacheStrategy; use crate::error::{InvalidRequestSnafu, Result, WorkerStoppedSnafu}; use crate::manifest::action::RegionEdit; use crate::metrics::INDEX_PUBLICATION_STALE_TOTAL; -use crate::region::version::VersionRef; +use crate::region::options::IndexOptions; use crate::region::{IndexBuildSource, MitoRegionRef}; use crate::request::{ BuildIndexRequest, IndexBuildFailed, IndexBuildFinished, IndexBuildStopped, OptionOutputTx, @@ -39,11 +40,20 @@ use crate::sst::index::{ }; use crate::worker::RegionWorkerLoop; +/// Parts of a region version that index build tasks need. +/// +/// Holding this instead of the whole version avoids pinning memtables and the +/// SST version while the worker awaits. +pub(crate) struct IndexBuildSnapshot { + metadata: RegionMetadataRef, + index_options: IndexOptions, +} + impl RegionWorkerLoop { pub(crate) fn new_index_build_task( &self, region: &MitoRegionRef, - version: &VersionRef, + snapshot: &IndexBuildSnapshot, file: FileHandle, file_meta: FileMeta, build_type: IndexBuildType, @@ -65,11 +75,11 @@ impl RegionWorkerLoop { let indexer_builder_ref = Arc::new(IndexerBuilderImpl { build_type: build_type.clone(), - metadata: version.metadata.clone(), + metadata: snapshot.metadata.clone(), inverted_index_config: self.config.inverted_index.clone(), fulltext_index_config: self.config.fulltext_index.clone(), bloom_filter_index_config: self.config.bloom_filter_index.clone(), - index_options: version.options.index_options.clone(), + index_options: snapshot.index_options.clone(), intermediate_manager, puffin_manager, write_cache_enabled: self.cache_manager.write_cache().is_some(), @@ -78,8 +88,8 @@ impl RegionWorkerLoop { IndexBuildTask { region_id: region.region_id, file: file.clone(), - target_region_metadata: version.metadata.clone(), - source: IndexBuildSource::new(file_meta, version.metadata.schema_version), + target_region_metadata: snapshot.metadata.clone(), + source: IndexBuildSource::new(file_meta, snapshot.metadata.schema_version), reason: build_type, access_layer: access_layer.clone(), listener: self.listener.clone(), @@ -159,11 +169,27 @@ impl RegionWorkerLoop { return; }; - let version_control = region.version_control.clone(); - let version = version_control.current().version; + // Only keep the parts of the version that index builds need, so the + // memtables and SST version are not pinned across the awaits below. + let (snapshot, all_files) = { + let version = region.version_control.current().version; + let all_files: HashMap = version + .ssts + .levels() + .iter() + .flat_map(|level| level.files.iter()) + .filter(|(_, handle)| !handle.is_deleted() && !handle.compacting()) + .map(|(id, handle)| (*id, handle.clone())) + .collect(); + let snapshot = IndexBuildSnapshot { + metadata: version.metadata.clone(), + index_options: version.options.index_options.clone(), + }; + (snapshot, all_files) + }; // A committed index publication may not have reached version control // yet. Use the manifest's FileMeta as the conditional publication - // source while keeping the builder's schema generation from `version`. + // source while keeping the builder's schema generation from the version. let manifest = region.manifest_ctx.manifest().await; let current_file_meta = |file_id| { let file_meta = manifest.files.get(&file_id); @@ -176,23 +202,15 @@ impl RegionWorkerLoop { file_meta.cloned() }; - let all_files: HashMap = version - .ssts - .levels() - .iter() - .flat_map(|level| level.files.iter()) - .filter(|(_, handle)| !handle.is_deleted() && !handle.compacting()) - .map(|(id, handle)| (*id, handle.clone())) - .collect(); - let build_tasks = if request.file_metas.is_empty() { // If no specific files are provided, find files whose index is inconsistent with the region metadata. all_files .values() .filter_map(|file| { let file_meta = current_file_meta(file.meta_ref().file_id)?; - (!file_meta.is_index_consistent_with_region(&version.metadata.column_metadatas)) - .then(|| (file.clone(), file_meta)) + (!file_meta + .is_index_consistent_with_region(&snapshot.metadata.column_metadatas)) + .then(|| (file.clone(), file_meta)) }) .collect::>() } else { @@ -238,7 +256,7 @@ impl RegionWorkerLoop { let task = self.new_index_build_task( ®ion, - &version, + &snapshot, file_handle.clone(), file_meta, request.build_type.clone(), From 88bdab0f6616005b23909f5bbf3bdca9f26bf8a9 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Fri, 16 Oct 2026 09:13:57 +0000 Subject: [PATCH 024/239] fix(mito2): keep processing bulk requests after one fails to encode (#9638) * fix(mito2): keep processing bulk requests after one fails to encode When `push_bulk` fails to build the WAL entry of a bulk request, it answers that request with the error. `process_bulk_requests` then returned from the loop, so the remaining requests of the drained batch were dropped without an answer and their callers only saw a receive error. Skip the failed request and continue with the rest of the batch instead. Signed-off-by: jeremyhi * test(mito2): keep bulk install barriers independent per target The bulk install test barrier was a single process-wide slot, so tests that arm it while running in the same process replaced each other's barrier. Keep one barrier per armed target instead. Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- src/mito2/src/engine/basic_test.rs | 180 ++++++++++++++++++++++++++- src/mito2/src/region_write_ctx.rs | 39 +++--- src/mito2/src/worker/handle_write.rs | 2 +- 3 files changed, 192 insertions(+), 29 deletions(-) diff --git a/src/mito2/src/engine/basic_test.rs b/src/mito2/src/engine/basic_test.rs index bd0e6827880..c54691a5a8f 100644 --- a/src/mito2/src/engine/basic_test.rs +++ b/src/mito2/src/engine/basic_test.rs @@ -16,16 +16,22 @@ use std::collections::HashMap; use std::sync::atomic::Ordering; +use std::time::Duration; use api::v1::helper::row; use api::v1::value::ValueData; -use api::v1::{Rows, SemanticType}; +use api::v1::{ArrowIpc, Rows, SemanticType}; use common_base::readable_size::ReadableSize; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; -use common_recordbatch::RecordBatches; +use common_recordbatch::{DfRecordBatch, RecordBatches}; +use common_test_util::flight::encode_to_flight_data; use common_time::Timestamp; use common_wal::options::WAL_OPTIONS_KEY; +use datatypes::arrow::array::{ + ArrayRef, DictionaryArray, Float64Array, StringArray, TimestampMillisecondArray, +}; +use datatypes::arrow::datatypes::{DataType, Field, Schema, TimeUnit, UInt32Type}; use datatypes::prelude::ConcreteDataType; use datatypes::schema::ColumnSchema; use futures::TryStreamExt; @@ -35,13 +41,14 @@ use rstest::rstest; use rstest_reuse::{self, apply}; use store_api::metadata::ColumnMetadata; use store_api::region_request::{ - PathType, RegionCreateRequest, RegionFlushRequest, RegionOpenRequest, RegionPutRequest, - RegionRequirements, + PathType, RegionBulkInsertsRequest, RegionCreateRequest, RegionFlushRequest, RegionOpenRequest, + RegionPutRequest, RegionRequirements, }; use store_api::storage::RegionId; use super::*; use crate::region::version::VersionControlData; +use crate::region_write_ctx::test_hooks::arm_bulk_install_barrier; use crate::test_util::{ CreateRequestBuilder, LogStoreFactory, TestEnv, build_delete_rows_for_key, build_rows, build_rows_for_key, delete_rows, delete_rows_schema, flush_region, kafka_log_store_factory, @@ -1872,3 +1879,168 @@ async fn restart_without_region_close( .unwrap(); engine } + +/// Builds a bulk insert request for a region created by [CreateRequestBuilder] +/// with its default columns `(tag_0, field_0, ts)`. +fn build_bulk_insert_request(region_id: RegionId, num_rows: usize) -> RegionBulkInsertsRequest { + let schema = Arc::new(Schema::new(vec![ + Field::new("tag_0", DataType::Utf8, true), + Field::new("field_0", DataType::Float64, true), + Field::new( + "ts", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + ])); + let tag = Arc::new(StringArray::from_iter_values( + (0..num_rows).map(|v| format!("a{v}")), + )) as ArrayRef; + let field = Arc::new(Float64Array::from_iter_values( + (0..num_rows).map(|v| v as f64), + )) as ArrayRef; + let ts = Arc::new(TimestampMillisecondArray::from_iter_values( + (0..num_rows).map(|v| v as i64 * 1000), + )) as ArrayRef; + let payload = DfRecordBatch::try_new(schema, vec![tag, field, ts]).unwrap(); + let (schema, record_batch) = encode_to_flight_data(payload.clone()); + + RegionBulkInsertsRequest { + skip_wal: false, + region_id, + payload, + raw_data: ArrowIpc { + schema: schema.data_header, + data_header: record_batch.data_header, + payload: record_batch.data_body, + }, + partition_expr_version: None, + aligned_schema_version: None, + } +} + +/// Builds a bulk insert request whose WAL entry can't be encoded. The batch +/// misses `field_0`, so the worker fills it and clears `raw_data`, and its +/// all-null dictionary `tag_0` column can't be encoded into a single IPC +/// message. +fn build_unencodable_bulk_insert_request(region_id: RegionId) -> RegionBulkInsertsRequest { + let schema = Arc::new(Schema::new(vec![ + Field::new( + "tag_0", + DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)), + true, + ), + Field::new( + "ts", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + ])); + let tag = Arc::new(DictionaryArray::::from_iter([None::<&str>])) as ArrayRef; + let ts = Arc::new(TimestampMillisecondArray::from_iter_values([0])) as ArrayRef; + let payload = DfRecordBatch::try_new(schema, vec![tag, ts]).unwrap(); + + RegionBulkInsertsRequest { + skip_wal: false, + region_id, + payload, + // The worker clears it after filling the missing column. + raw_data: ArrowIpc::default(), + partition_expr_version: None, + aligned_schema_version: None, + } +} + +async fn scan_num_rows(engine: &MitoEngine, region_id: RegionId) -> usize { + let stream = engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + batches.iter().map(|batch| batch.num_rows()).sum() +} + +#[tokio::test] +async fn test_bulk_insert_encode_failure_keeps_worker_batch() { + common_telemetry::init_default_ut_logging(); + + let mut env = TestEnv::new().await; + // A single worker so requests of all regions share one batch. + let engine = env + .create_engine(MitoConfig { + num_workers: 1, + ..Default::default() + }) + .await; + + let blocker_region_id = RegionId::new(1, 1); + let region_id = RegionId::new(1, 2); + let other_region_id = RegionId::new(1, 3); + for id in [blocker_region_id, region_id, other_region_id] { + engine + .handle_request( + id, + RegionRequest::Create(CreateRequestBuilder::new().build()), + ) + .await + .unwrap(); + } + + // Pauses the worker in a bulk write of the blocker region so the requests + // below queue up and reach the worker as one batch. + let mut barrier = arm_bulk_install_barrier( + blocker_region_id, + engine + .get_region(blocker_region_id) + .unwrap() + .version_control + .clone(), + ); + let blocker_engine = engine.clone(); + let blocker = tokio::spawn(async move { + blocker_engine + .handle_request( + blocker_region_id, + RegionRequest::BulkInserts(build_bulk_insert_request(blocker_region_id, 1)), + ) + .await + }); + tokio::time::timeout(Duration::from_secs(10), barrier.wait_until_reached()) + .await + .expect("bulk write never reached the install barrier"); + + let mut receivers = Vec::new(); + for (id, request) in [ + (region_id, build_unencodable_bulk_insert_request(region_id)), + (region_id, build_bulk_insert_request(region_id, 3)), + ( + other_region_id, + build_bulk_insert_request(other_region_id, 2), + ), + ] { + let metadata = engine.inner.get_metadata(id).ok(); + let (request, receiver) = WorkerRequest::try_from_region_request( + id, + RegionRequest::BulkInserts(request), + metadata, + ) + .unwrap(); + engine + .inner + .workers + .submit_to_worker(id, request) + .await + .unwrap(); + receivers.push(receiver); + } + barrier.release(); + assert_eq!(1, blocker.await.unwrap().unwrap().affected_rows); + + let [failed, same_region, other_region] = receivers.try_into().unwrap(); + let err = failed.await.unwrap().unwrap_err(); + assert_eq!(StatusCode::Unsupported, err.status_code(), "{err:?}"); + // Requests after the failed one in the same batch are still written. + assert_eq!(3, same_region.await.unwrap().unwrap()); + assert_eq!(2, other_region.await.unwrap().unwrap()); + assert_eq!(3, scan_num_rows(&engine, region_id).await); + assert_eq!(2, scan_num_rows(&engine, other_region_id).await); +} diff --git a/src/mito2/src/region_write_ctx.rs b/src/mito2/src/region_write_ctx.rs index f8d7eb9c840..8c306effd1d 100644 --- a/src/mito2/src/region_write_ctx.rs +++ b/src/mito2/src/region_write_ctx.rs @@ -470,13 +470,13 @@ pub(crate) mod test_hooks { release: watch::Sender, } - static ACTIVE_BARRIER: Mutex> = Mutex::new(None); + static ACTIVE_BARRIERS: Mutex> = Mutex::new(Vec::new()); static NEXT_BARRIER_ID: AtomicU64 = AtomicU64::new(1); - fn lock_active_barrier() -> std::sync::MutexGuard<'static, Option> { + fn lock_active_barriers() -> std::sync::MutexGuard<'static, Vec> { // Never let a poisoned mutex (e.g. a panic in another test while // holding the lock) hang or break unrelated tests. - ACTIVE_BARRIER + ACTIVE_BARRIERS .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()) } @@ -513,8 +513,8 @@ pub(crate) mod test_hooks { } } - /// Arms the bulk-install barrier and returns the owning guard; any - /// previously armed barrier is replaced. + /// Arms a bulk-install barrier for the target and returns the owning + /// guard; barriers of concurrently running tests stay independent. pub(crate) fn arm_bulk_install_barrier( target_region_id: RegionId, target_version_control: VersionControlRef, @@ -522,8 +522,7 @@ pub(crate) mod test_hooks { let (reached_tx, reached_rx) = watch::channel(false); let (release_tx, _release_rx) = watch::channel(false); let id = NEXT_BARRIER_ID.fetch_add(1, Ordering::Relaxed); - let mut active = lock_active_barrier(); - *active = Some(ActiveBarrier { + lock_active_barriers().push(ActiveBarrier { id, target_region_id, target_version_control, @@ -539,10 +538,7 @@ pub(crate) mod test_hooks { } fn disarm_barrier(id: u64) { - let mut active = lock_active_barrier(); - if active.as_ref().is_some_and(|barrier| barrier.id == id) { - *active = None; - } + lock_active_barriers().retain(|barrier| barrier.id != id); } /// Pauses a bulk write before installing its parts until the barrier is @@ -552,19 +548,14 @@ pub(crate) mod test_hooks { version_control: &VersionControlRef, ) { let (reached_tx, release_rx) = { - let active = lock_active_barrier(); - match active.as_ref() { - Some(barrier) - if barrier.target_region_id == region_id - && std::sync::Arc::ptr_eq( - &barrier.target_version_control, - version_control, - ) => - { - (barrier.reached.clone(), barrier.release.subscribe()) - } - _ => return, - } + let active = lock_active_barriers(); + let Some(barrier) = active.iter().find(|barrier| { + barrier.target_region_id == region_id + && std::sync::Arc::ptr_eq(&barrier.target_version_control, version_control) + }) else { + return; + }; + (barrier.reached.clone(), barrier.release.subscribe()) }; let _ = reached_tx.send(true); let mut release_rx = release_rx; diff --git a/src/mito2/src/worker/handle_write.rs b/src/mito2/src/worker/handle_write.rs index 721235e71b5..2b9551308f0 100644 --- a/src/mito2/src/worker/handle_write.rs +++ b/src/mito2/src/worker/handle_write.rs @@ -564,7 +564,7 @@ impl RegionWorkerLoop { // Collect requests by region. if !region_ctx.push_bulk(bulk_req.sender, bulk_req.request, None, bulk_req.skip_wal) { - return; + continue; } } } From 672e94f61d44089d4369865e51d5f6a9122405dc Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Fri, 16 Oct 2026 09:40:54 +0000 Subject: [PATCH 025/239] feat: add column encoding option and use it for trace span ids (#9629) * feat: add column encoding option and set it for trace span ids COMPRESSION WITH accepts encoding = 'plain' | 'delta_binary_packed'. Either disables the dictionary for the column; delta_binary_packed is limited to integer and timestamp columns. OTLP trace v1/v2 tables now create span_id and parent_span_id with plain encoding and 64KiB pages, and timestamp_end with delta_binary_packed. Signed-off-by: Dennis Zhuang * test: add compatibility cases for column encoding options Signed-off-by: Dennis Zhuang * test: move column encoding downgrade checks into downgrade_compatibility Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/api/src/v1/column_def.rs | 12 ++- src/common/catalog/src/consts.rs | 1 + src/datatypes/src/schema.rs | 17 ++-- src/datatypes/src/schema/column_schema.rs | 95 +++++++++++++++++++ src/mito2/src/sst/parquet.rs | 42 +++++--- src/operator/src/insert.rs | 43 ++++++++- src/servers/src/otlp/trace.rs | 4 +- src/sql/src/parsers/utils.rs | 13 +-- tests-integration/src/otlp.rs | 6 +- tests-integration/tests/http.rs | 10 +- .../create/create_with_compression.result | 20 +++- .../common/create/create_with_compression.sql | 10 +- .../cases/column_encoding/case.toml | 8 ++ .../cases/column_encoding/setup.sql | 3 + .../cases/column_encoding/verify.result | 67 +++++++++++++ .../cases/column_encoding/verify.sql | 20 ++++ .../cases/downgrade_compatibility/case.toml | 4 +- .../cases/downgrade_compatibility/setup.sql | 10 ++ .../downgrade_compatibility/verify.result | 10 ++ .../cases/downgrade_compatibility/verify.sql | 2 + 20 files changed, 347 insertions(+), 50 deletions(-) create mode 100644 tests/compatibility/cases/column_encoding/case.toml create mode 100644 tests/compatibility/cases/column_encoding/setup.sql create mode 100644 tests/compatibility/cases/column_encoding/verify.result create mode 100644 tests/compatibility/cases/column_encoding/verify.sql diff --git a/src/api/src/v1/column_def.rs b/src/api/src/v1/column_def.rs index af62ed716e5..73c56339474 100644 --- a/src/api/src/v1/column_def.rs +++ b/src/api/src/v1/column_def.rs @@ -16,8 +16,8 @@ use std::collections::HashMap; use arrow_schema::extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY}; use datatypes::schema::{ - COMMENT_KEY, COMPRESSION_KEY, ColumnDefaultConstraint, ColumnSchema, FULLTEXT_KEY, - FulltextAnalyzer, FulltextBackend, FulltextOptions, INVERTED_INDEX_KEY, Metadata, + COMMENT_KEY, COMPRESSION_KEY, ColumnDefaultConstraint, ColumnSchema, CompressionOptions, + FULLTEXT_KEY, FulltextAnalyzer, FulltextBackend, FulltextOptions, INVERTED_INDEX_KEY, Metadata, SKIPPING_INDEX_KEY, SkippingIndexOptions, SkippingIndexType, }; use greptime_proto::v1::{ @@ -240,6 +240,14 @@ pub fn options_from_skipping(skipping: &SkippingIndexOptions) -> Result Result { + let mut options = ColumnOptions::default(); + let v = serde_json::to_string(compression).context(error::SerializeJsonSnafu)?; + options.options.insert(COMPRESSION_GRPC_KEY.to_string(), v); + Ok(options) +} + /// Tries to construct a `ColumnOptions` for inverted index. pub fn options_from_inverted() -> ColumnOptions { let mut options = ColumnOptions::default(); diff --git a/src/common/catalog/src/consts.rs b/src/common/catalog/src/consts.rs index 7c71de2e1fa..d999f871487 100644 --- a/src/common/catalog/src/consts.rs +++ b/src/common/catalog/src/consts.rs @@ -217,6 +217,7 @@ pub const TRACE_TIMESTAMP_COLUMN: &str = "timestamp"; pub const SPAN_KIND_COLUMN: &str = "span_kind"; pub const SPAN_STATUS_CODE_COLUMN: &str = "span_status_code"; pub const DURATION_NANO_COLUMN: &str = "duration_nano"; +pub const TIMESTAMP_END_COLUMN: &str = "timestamp_end"; pub const SPAN_KIND_CLIENT: &str = "SPAN_KIND_CLIENT"; pub const SPAN_KIND_SERVER: &str = "SPAN_KIND_SERVER"; pub const SPAN_STATUS_ERROR: &str = "STATUS_CODE_ERROR"; diff --git a/src/datatypes/src/schema.rs b/src/datatypes/src/schema.rs index 66867f6141d..4652b0104d1 100644 --- a/src/datatypes/src/schema.rs +++ b/src/datatypes/src/schema.rs @@ -28,14 +28,15 @@ use snafu::{ResultExt, ensure}; use crate::error::{self, DuplicateColumnSnafu, Error, ProjectArrowSchemaSnafu, Result}; use crate::prelude::ConcreteDataType; pub use crate::schema::column_schema::{ - COLUMN_COMPRESSION_OPT_KEY_LEVEL, COLUMN_COMPRESSION_OPT_KEY_PAGE_SIZE, - COLUMN_COMPRESSION_OPT_KEY_TYPE, COLUMN_FULLTEXT_CHANGE_OPT_KEY_ENABLE, - COLUMN_FULLTEXT_OPT_KEY_ANALYZER, COLUMN_FULLTEXT_OPT_KEY_BACKEND, - COLUMN_FULLTEXT_OPT_KEY_CASE_SENSITIVE, COLUMN_FULLTEXT_OPT_KEY_FALSE_POSITIVE_RATE, - COLUMN_FULLTEXT_OPT_KEY_GRANULARITY, COLUMN_SKIPPING_INDEX_OPT_KEY_FALSE_POSITIVE_RATE, - COLUMN_SKIPPING_INDEX_OPT_KEY_GRANULARITY, COLUMN_SKIPPING_INDEX_OPT_KEY_TYPE, COMMENT_KEY, - COMPRESSION_KEY, ColumnExtType, ColumnSchema, CompressionOptions, CompressionType, - FULLTEXT_KEY, FulltextAnalyzer, FulltextBackend, FulltextOptions, INVERTED_INDEX_KEY, Metadata, + COLUMN_COMPRESSION_OPT_KEY_ENCODING, COLUMN_COMPRESSION_OPT_KEY_LEVEL, + COLUMN_COMPRESSION_OPT_KEY_PAGE_SIZE, COLUMN_COMPRESSION_OPT_KEY_TYPE, + COLUMN_FULLTEXT_CHANGE_OPT_KEY_ENABLE, COLUMN_FULLTEXT_OPT_KEY_ANALYZER, + COLUMN_FULLTEXT_OPT_KEY_BACKEND, COLUMN_FULLTEXT_OPT_KEY_CASE_SENSITIVE, + COLUMN_FULLTEXT_OPT_KEY_FALSE_POSITIVE_RATE, COLUMN_FULLTEXT_OPT_KEY_GRANULARITY, + COLUMN_SKIPPING_INDEX_OPT_KEY_FALSE_POSITIVE_RATE, COLUMN_SKIPPING_INDEX_OPT_KEY_GRANULARITY, + COLUMN_SKIPPING_INDEX_OPT_KEY_TYPE, COMMENT_KEY, COMPRESSION_KEY, ColumnEncoding, + ColumnExtType, ColumnSchema, CompressionOptions, CompressionType, FULLTEXT_KEY, + FulltextAnalyzer, FulltextBackend, FulltextOptions, INVERTED_INDEX_KEY, Metadata, SKIPPING_INDEX_KEY, SkippingIndexOptions, SkippingIndexType, TIME_INDEX_KEY, }; pub use crate::schema::constraint::ColumnDefaultConstraint; diff --git a/src/datatypes/src/schema/column_schema.rs b/src/datatypes/src/schema/column_schema.rs index 2d48ca72ad5..9dac9258116 100644 --- a/src/datatypes/src/schema/column_schema.rs +++ b/src/datatypes/src/schema/column_schema.rs @@ -69,6 +69,7 @@ pub const COLUMN_SKIPPING_INDEX_OPT_KEY_TYPE: &str = "type"; pub const COLUMN_COMPRESSION_OPT_KEY_TYPE: &str = "type"; pub const COLUMN_COMPRESSION_OPT_KEY_LEVEL: &str = "level"; pub const COLUMN_COMPRESSION_OPT_KEY_PAGE_SIZE: &str = "page_size"; +pub const COLUMN_COMPRESSION_OPT_KEY_ENCODING: &str = "encoding"; pub const DEFAULT_GRANULARITY: u32 = 10240; @@ -474,6 +475,17 @@ impl ColumnSchema { ), } ); + if options.encoding == Some(ColumnEncoding::DeltaBinaryPacked) { + ensure!( + self.data_type.is_timestamp() || self.data_type.as_arrow_type().is_integer(), + error::InvalidCompressionOptionSnafu { + msg: format!( + "encoding delta_binary_packed requires an integer or timestamp column, but column {} is {}", + self.name, self.data_type + ), + } + ); + } Ok(()) } @@ -1018,6 +1030,24 @@ impl fmt::Display for CompressionType { } } +/// Parquet value encodings for SST columns. Both replace the default dictionary encoding. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ColumnEncoding { + Plain, + /// Only valid for integer and timestamp columns. + DeltaBinaryPacked, +} + +impl fmt::Display for ColumnEncoding { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + ColumnEncoding::Plain => write!(f, "plain"), + ColumnEncoding::DeltaBinaryPacked => write!(f, "delta_binary_packed"), + } + } +} + /// Parquet compression settings for a column in SSTs. Unset fields keep the SST defaults. #[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] pub struct CompressionOptions { @@ -1027,6 +1057,8 @@ pub struct CompressionOptions { pub level: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub page_size: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub encoding: Option, } impl CompressionOptions { @@ -1070,6 +1102,12 @@ impl CompressionOptions { }; map.insert(COLUMN_COMPRESSION_OPT_KEY_PAGE_SIZE.to_string(), page_size); } + if let Some(encoding) = self.encoding { + map.insert( + COLUMN_COMPRESSION_OPT_KEY_ENCODING.to_string(), + encoding.to_string(), + ); + } map } } @@ -1129,10 +1167,25 @@ impl TryFrom> for CompressionOptions { }) .transpose()?; + let encoding = options + .get(COLUMN_COMPRESSION_OPT_KEY_ENCODING) + .map(|value| match value.to_ascii_lowercase().as_str() { + "plain" => Ok(ColumnEncoding::Plain), + "delta_binary_packed" => Ok(ColumnEncoding::DeltaBinaryPacked), + _ => error::InvalidCompressionOptionSnafu { + msg: format!( + "invalid encoding: {value}, expected: 'plain' or 'delta_binary_packed'" + ), + } + .fail(), + }) + .transpose()?; + Ok(CompressionOptions { compression_type, level, page_size, + encoding, }) } } @@ -1525,6 +1578,7 @@ mod tests { compression_type: CompressionType::Zstd, level: Some(3), page_size: Some(ReadableSize::mb(8)), + encoding: None, }, options ); @@ -1558,6 +1612,17 @@ mod tests { serde_json::to_string(&page_only).unwrap() ); + let delta = parse(&[("encoding", "DELTA_BINARY_PACKED")]).unwrap(); + assert_eq!(Some(ColumnEncoding::DeltaBinaryPacked), delta.encoding); + assert_eq!( + r#"{"type":"zstd","encoding":"delta_binary_packed"}"#, + serde_json::to_string(&delta).unwrap() + ); + assert_eq!( + delta, + parse(&[("encoding", delta.to_options_map()["encoding"].as_str())]).unwrap() + ); + for invalid in [ &[("type", "lz4")][..], &[("level", "0")], @@ -1565,6 +1630,7 @@ mod tests { &[("level", "high")], &[("page_size", "0")], &[("page_size", "large")], + &[("encoding", "dictionary")], ] { assert!(parse(invalid).is_err(), "{invalid:?}"); } @@ -1598,6 +1664,35 @@ mod tests { } } + #[test] + fn test_delta_encoding_requires_integer_or_timestamp() { + let options = CompressionOptions { + encoding: Some(ColumnEncoding::DeltaBinaryPacked), + ..Default::default() + }; + for data_type in [ + ConcreteDataType::int32_datatype(), + ConcreteDataType::uint64_datatype(), + ConcreteDataType::timestamp_nanosecond_datatype(), + ] { + assert!( + ColumnSchema::new("c", data_type, true) + .with_compression_options(options) + .is_ok() + ); + } + for data_type in [ + ConcreteDataType::string_datatype(), + ConcreteDataType::float64_datatype(), + ] { + assert!( + ColumnSchema::new("c", data_type, true) + .with_compression_options(options) + .is_err() + ); + } + } + #[test] fn test_skipping_index_options_deserialization() { let original_options = "{\"granularity\":1024,\"false-positive-rate-in-10000\":10,\"index-type\":\"BloomFilter\"}"; diff --git a/src/mito2/src/sst/parquet.rs b/src/mito2/src/sst/parquet.rs index e6968447dce..79ad957b75a 100644 --- a/src/mito2/src/sst/parquet.rs +++ b/src/mito2/src/sst/parquet.rs @@ -20,7 +20,8 @@ use std::sync::Arc; use api::v1::SemanticType; use common_base::readable_size::ReadableSize; use datatypes::json::JsonSettings; -use parquet::basic::{Compression, ZstdLevel}; +use datatypes::schema::ColumnEncoding; +use parquet::basic::{Compression, Encoding, ZstdLevel}; use parquet::file::metadata::ParquetMetaData; use parquet::file::properties::{DEFAULT_DICTIONARY_PAGE_SIZE_LIMIT, WriterPropertiesBuilder}; use parquet::schema::types::ColumnPath; @@ -109,6 +110,15 @@ pub(crate) fn apply_column_compression( continue; }; let path = ColumnPath::new(vec![column.column_schema.name.clone()]); + if let Some(encoding) = options.encoding { + let encoding = match encoding { + ColumnEncoding::Plain => Encoding::PLAIN, + ColumnEncoding::DeltaBinaryPacked => Encoding::DELTA_BINARY_PACKED, + }; + builder = builder + .set_column_dictionary_enabled(path.clone(), false) + .set_column_encoding(path.clone(), encoding); + } if let Some(level) = options.level { let level = ZstdLevel::try_new(level).context(WriteParquetSnafu)?; builder = builder.set_column_compression(path.clone(), Compression::ZSTD(level)); @@ -883,7 +893,7 @@ mod tests { async fn write( values: &[Vec], compression: Option, - ) -> (usize, usize) { + ) -> (usize, usize, bool) { let mut metadata = (*build_test_binary_test_region_metadata()).clone(); if let Some(options) = compression { let column = &mut metadata.column_metadatas[1]; @@ -961,16 +971,16 @@ mod tests { let pages = builder.metadata().offset_index().unwrap()[0][field] .page_locations() .len(); - let size = builder - .metadata() - .row_group(0) - .column(field) - .compressed_size() as usize; - (size, pages) + let chunk = builder.metadata().row_group(0).column(field); + ( + chunk.compressed_size() as usize, + pages, + chunk.dictionary_page_offset().is_some(), + ) } - let (default_size, default_pages) = write(&values, None).await; - let (level_size, level_pages) = write( + let (default_size, default_pages, default_dict) = write(&values, None).await; + let (level_size, level_pages, _) = write( &values, Some(CompressionOptions { level: Some(19), @@ -978,7 +988,7 @@ mod tests { }), ) .await; - let (_, pages) = write( + let (_, pages, _) = write( &values, Some(CompressionOptions { page_size: Some(ReadableSize::kb(64)), @@ -986,11 +996,21 @@ mod tests { }), ) .await; + let (_, _, plain_dict) = write( + &values, + Some(CompressionOptions { + encoding: Some(ColumnEncoding::Plain), + ..Default::default() + }), + ) + .await; assert!(default_pages > 1, "{default_pages}"); assert_eq!(default_pages, level_pages); assert!(level_size < default_size, "{level_size} >= {default_size}"); assert!(pages > default_pages, "{pages} <= {default_pages}"); + assert!(default_dict); + assert!(!plain_dict); } #[tokio::test] diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index e8f183bf7cc..fbf65986c56 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -22,7 +22,7 @@ use std::time::{Duration, Instant}; use ahash::{HashMap, HashMapExt, HashSet, HashSetExt}; use api::v1::alter_table_expr::Kind; -use api::v1::column_def::{options_from_skipping, try_as_column_def}; +use api::v1::column_def::{options_from_compression, options_from_skipping, try_as_column_def}; use api::v1::region::{ InsertRequest as RegionInsertRequest, InsertRequests as RegionInsertRequests, RegionRequestHeader, @@ -35,10 +35,11 @@ use api::v1::{ use arrow::datatypes::{DataType as ArrowDataType, Schema as ArrowSchema}; use catalog::CatalogManagerRef; use client::{OutputData, OutputMeta}; +use common_base::readable_size::ReadableSize; use common_catalog::consts::{ - DEFAULT_PRIVATE_SCHEMA_NAME, PARENT_SPAN_ID_COLUMN, SERVICE_NAME_COLUMN, TRACE_ID_COLUMN, - TRACE_TABLE_NAME, TRACE_TABLE_NAME_SESSION_KEY, default_engine, is_ddl_reserved_table, - trace_operations_table_name, trace_services_table_name, + DEFAULT_PRIVATE_SCHEMA_NAME, PARENT_SPAN_ID_COLUMN, SERVICE_NAME_COLUMN, SPAN_ID_COLUMN, + TIMESTAMP_END_COLUMN, TRACE_ID_COLUMN, TRACE_TABLE_NAME, TRACE_TABLE_NAME_SESSION_KEY, + default_engine, is_ddl_reserved_table, trace_operations_table_name, trace_services_table_name, }; use common_event_recorder::DEFAULT_EVENTS_TABLE_NAME; use common_frontend::slow_query_event::SLOW_QUERY_TABLE_NAME; @@ -55,7 +56,7 @@ use common_telemetry::tracing_context::TracingContext; use common_telemetry::{debug, error, warn}; use common_time::Timestamp; use common_time::timestamp::TimeUnit; -use datatypes::schema::SkippingIndexOptions; +use datatypes::schema::{ColumnEncoding, CompressionOptions, SkippingIndexOptions}; use futures_util::future; use meter_core::data::MeterRecord; use meter_macros::write_meter; @@ -1365,6 +1366,38 @@ impl Inserter { } } + // Span ids and end timestamps rarely repeat within an SST, so a + // dictionary costs more than it saves. Small span id pages let + // lookups by trace id decode only the pages holding matched rows. + let id_compression = CompressionOptions { + encoding: Some(ColumnEncoding::Plain), + page_size: Some(ReadableSize::kb(64)), + ..Default::default() + }; + let compressed_columns = [ + (SPAN_ID_COLUMN, id_compression), + (PARENT_SPAN_ID_COLUMN, id_compression), + ( + TIMESTAMP_END_COLUMN, + CompressionOptions { + encoding: Some(ColumnEncoding::DeltaBinaryPacked), + ..Default::default() + }, + ), + ]; + for (name, compression) in compressed_columns { + if let Some(col) = + create_table.column_defs.iter_mut().find(|c| c.name == name) + { + let compression = options_from_compression(&compression) + .context(ColumnOptionsSnafu)?; + col.options + .get_or_insert_with(Default::default) + .options + .extend(compression.options); + } + } + // use table_options to mark table model version create_table.table_options.insert( TABLE_DATA_MODEL.to_string(), diff --git a/src/servers/src/otlp/trace.rs b/src/servers/src/otlp/trace.rs index 73ce493388b..4b6c4d8ce62 100644 --- a/src/servers/src/otlp/trace.rs +++ b/src/servers/src/otlp/trace.rs @@ -25,7 +25,8 @@ use api::v1::RowInsertRequests; // in `common_catalog::consts` (shared with the read-time graph derivation). pub use common_catalog::consts::{ DURATION_NANO_COLUMN, SERVICE_NAME_COLUMN, SPAN_KIND_COLUMN, - SPAN_STATUS_CODE_COLUMN as SPAN_STATUS_CODE, TRACE_TIMESTAMP_COLUMN as TIMESTAMP_COLUMN, + SPAN_STATUS_CODE_COLUMN as SPAN_STATUS_CODE, TIMESTAMP_END_COLUMN, + TRACE_TIMESTAMP_COLUMN as TIMESTAMP_COLUMN, }; pub use common_catalog::consts::{ PARENT_SPAN_ID_COLUMN, RESOURCE_ATTRIBUTES_COLUMN, SCOPE_ATTRIBUTES_COLUMN, @@ -36,7 +37,6 @@ use pipeline::PipelineWay; use crate::error::{NotSupportedSnafu, Result}; use crate::otlp::trace::span::TraceSpan; -pub const TIMESTAMP_END_COLUMN: &str = "timestamp_end"; pub const SPAN_STATUS_MESSAGE_COLUMN: &str = "span_status_message"; pub const SPAN_EVENTS_COLUMN: &str = "span_events"; pub const SPAN_LINKS_COLUMN: &str = "span_links"; diff --git a/src/sql/src/parsers/utils.rs b/src/sql/src/parsers/utils.rs index 65744708232..01b136a997a 100644 --- a/src/sql/src/parsers/utils.rs +++ b/src/sql/src/parsers/utils.rs @@ -28,12 +28,12 @@ use datafusion_expr::{AggregateUDF, Expr, HigherOrderUDF, ScalarUDF, TableSource use datafusion_sql::planner::{ContextProvider, SqlToRel}; use datatypes::arrow::datatypes::DataType; use datatypes::schema::{ - COLUMN_COMPRESSION_OPT_KEY_LEVEL, COLUMN_COMPRESSION_OPT_KEY_PAGE_SIZE, - COLUMN_COMPRESSION_OPT_KEY_TYPE, COLUMN_FULLTEXT_OPT_KEY_ANALYZER, - COLUMN_FULLTEXT_OPT_KEY_BACKEND, COLUMN_FULLTEXT_OPT_KEY_CASE_SENSITIVE, - COLUMN_FULLTEXT_OPT_KEY_FALSE_POSITIVE_RATE, COLUMN_FULLTEXT_OPT_KEY_GRANULARITY, - COLUMN_SKIPPING_INDEX_OPT_KEY_FALSE_POSITIVE_RATE, COLUMN_SKIPPING_INDEX_OPT_KEY_GRANULARITY, - COLUMN_SKIPPING_INDEX_OPT_KEY_TYPE, + COLUMN_COMPRESSION_OPT_KEY_ENCODING, COLUMN_COMPRESSION_OPT_KEY_LEVEL, + COLUMN_COMPRESSION_OPT_KEY_PAGE_SIZE, COLUMN_COMPRESSION_OPT_KEY_TYPE, + COLUMN_FULLTEXT_OPT_KEY_ANALYZER, COLUMN_FULLTEXT_OPT_KEY_BACKEND, + COLUMN_FULLTEXT_OPT_KEY_CASE_SENSITIVE, COLUMN_FULLTEXT_OPT_KEY_FALSE_POSITIVE_RATE, + COLUMN_FULLTEXT_OPT_KEY_GRANULARITY, COLUMN_SKIPPING_INDEX_OPT_KEY_FALSE_POSITIVE_RATE, + COLUMN_SKIPPING_INDEX_OPT_KEY_GRANULARITY, COLUMN_SKIPPING_INDEX_OPT_KEY_TYPE, }; use snafu::{ResultExt, ensure}; use sqlparser::dialect::Dialect; @@ -382,6 +382,7 @@ pub fn validate_column_compression_create_option(key: &str) -> bool { COLUMN_COMPRESSION_OPT_KEY_TYPE, COLUMN_COMPRESSION_OPT_KEY_LEVEL, COLUMN_COMPRESSION_OPT_KEY_PAGE_SIZE, + COLUMN_COMPRESSION_OPT_KEY_ENCODING, ] .contains(&key) } diff --git a/tests-integration/src/otlp.rs b/tests-integration/src/otlp.rs index 2349d580c5f..94fac669010 100644 --- a/tests-integration/src/otlp.rs +++ b/tests-integration/src/otlp.rs @@ -713,11 +713,11 @@ mod test { }; let expected = r#"CREATE TABLE IF NOT EXISTS "trace_v2" ( "timestamp" TIMESTAMP(9) NOT NULL, - "timestamp_end" TIMESTAMP(9) NULL, + "timestamp_end" TIMESTAMP(9) NULL COMPRESSION WITH(encoding = 'delta_binary_packed', type = 'zstd'), "duration_nano" BIGINT NULL, - "parent_span_id" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'), + "parent_span_id" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM') COMPRESSION WITH(encoding = 'plain', page_size = '64KiB', type = 'zstd'), "trace_id" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'), - "span_id" STRING NULL, + "span_id" STRING NULL COMPRESSION WITH(encoding = 'plain', page_size = '64KiB', type = 'zstd'), "span_kind" STRING NULL, "span_name" STRING NULL, "span_status_code" STRING NULL, diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 8ac98bf22f6..5b357f6904f 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -9770,7 +9770,7 @@ pub async fn test_otlp_traces_v1(store_type: StorageType) { ) .await; - let expected_ddl = r#"[["mytable","CREATE TABLE IF NOT EXISTS \"mytable\" (\n \"timestamp\" TIMESTAMP(9) NOT NULL,\n \"timestamp_end\" TIMESTAMP(9) NULL,\n \"duration_nano\" BIGINT NULL,\n \"parent_span_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"trace_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_id\" STRING NULL,\n \"span_kind\" STRING NULL,\n \"span_name\" STRING NULL,\n \"span_status_code\" STRING NULL,\n \"span_status_message\" STRING NULL,\n \"trace_state\" STRING NULL,\n \"scope_name\" STRING NULL,\n \"scope_version\" STRING NULL,\n \"service_name\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_attributes.net.peer.ip\" STRING NULL,\n \"span_attributes.peer.service\" STRING NULL,\n \"span_events\" JSON NULL,\n \"span_links\" JSON NULL,\n TIME INDEX (\"timestamp\"),\n PRIMARY KEY (\"service_name\")\n)\nPARTITION ON COLUMNS (\"trace_id\") (\n trace_id < '1',\n trace_id >= '1' AND trace_id < '2',\n trace_id >= '2' AND trace_id < '3',\n trace_id >= '3' AND trace_id < '4',\n trace_id >= '4' AND trace_id < '5',\n trace_id >= '5' AND trace_id < '6',\n trace_id >= '6' AND trace_id < '7',\n trace_id >= '7' AND trace_id < '8',\n trace_id >= '8' AND trace_id < '9',\n trace_id >= '9' AND trace_id < 'a',\n trace_id >= 'a' AND trace_id < 'b',\n trace_id >= 'b' AND trace_id < 'c',\n trace_id >= 'c' AND trace_id < 'd',\n trace_id >= 'd' AND trace_id < 'e',\n trace_id >= 'e' AND trace_id < 'f',\n trace_id >= 'f'\n)\nENGINE=mito\nWITH(\n 'comment' = 'Created on insertion',\n append_mode = 'true',\n 'greptime.semantic.entity.service.id' = 'service_name',\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\n 'greptime.semantic.signal_type' = 'trace',\n 'greptime.semantic.source' = 'opentelemetry',\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\n table_data_model = 'greptime_trace_v1'\n)"]]"#; + let expected_ddl = r#"[["mytable","CREATE TABLE IF NOT EXISTS \"mytable\" (\n \"timestamp\" TIMESTAMP(9) NOT NULL,\n \"timestamp_end\" TIMESTAMP(9) NULL COMPRESSION WITH(encoding = 'delta_binary_packed', type = 'zstd'),\n \"duration_nano\" BIGINT NULL,\n \"parent_span_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM') COMPRESSION WITH(encoding = 'plain', page_size = '64KiB', type = 'zstd'),\n \"trace_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_id\" STRING NULL COMPRESSION WITH(encoding = 'plain', page_size = '64KiB', type = 'zstd'),\n \"span_kind\" STRING NULL,\n \"span_name\" STRING NULL,\n \"span_status_code\" STRING NULL,\n \"span_status_message\" STRING NULL,\n \"trace_state\" STRING NULL,\n \"scope_name\" STRING NULL,\n \"scope_version\" STRING NULL,\n \"service_name\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_attributes.net.peer.ip\" STRING NULL,\n \"span_attributes.peer.service\" STRING NULL,\n \"span_events\" JSON NULL,\n \"span_links\" JSON NULL,\n TIME INDEX (\"timestamp\"),\n PRIMARY KEY (\"service_name\")\n)\nPARTITION ON COLUMNS (\"trace_id\") (\n trace_id < '1',\n trace_id >= '1' AND trace_id < '2',\n trace_id >= '2' AND trace_id < '3',\n trace_id >= '3' AND trace_id < '4',\n trace_id >= '4' AND trace_id < '5',\n trace_id >= '5' AND trace_id < '6',\n trace_id >= '6' AND trace_id < '7',\n trace_id >= '7' AND trace_id < '8',\n trace_id >= '8' AND trace_id < '9',\n trace_id >= '9' AND trace_id < 'a',\n trace_id >= 'a' AND trace_id < 'b',\n trace_id >= 'b' AND trace_id < 'c',\n trace_id >= 'c' AND trace_id < 'd',\n trace_id >= 'd' AND trace_id < 'e',\n trace_id >= 'e' AND trace_id < 'f',\n trace_id >= 'f'\n)\nENGINE=mito\nWITH(\n 'comment' = 'Created on insertion',\n append_mode = 'true',\n 'greptime.semantic.entity.service.id' = 'service_name',\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\n 'greptime.semantic.signal_type' = 'trace',\n 'greptime.semantic.source' = 'opentelemetry',\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\n table_data_model = 'greptime_trace_v1'\n)"]]"#; validate_data( "otlp_traces", &client, @@ -9823,7 +9823,7 @@ pub async fn test_otlp_traces_v1(store_type: StorageType) { ) .await; assert_eq!(StatusCode::OK, res.status()); - let expected_ddl = r#"[["trace_table_part1","CREATE TABLE IF NOT EXISTS \"trace_table_part1\" (\n \"timestamp\" TIMESTAMP(9) NOT NULL,\n \"timestamp_end\" TIMESTAMP(9) NULL,\n \"duration_nano\" BIGINT NULL,\n \"parent_span_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"trace_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_id\" STRING NULL,\n \"span_kind\" STRING NULL,\n \"span_name\" STRING NULL,\n \"span_status_code\" STRING NULL,\n \"span_status_message\" STRING NULL,\n \"trace_state\" STRING NULL,\n \"scope_name\" STRING NULL,\n \"scope_version\" STRING NULL,\n \"service_name\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_attributes.net.peer.ip\" STRING NULL,\n \"span_attributes.peer.service\" STRING NULL,\n \"span_events\" JSON NULL,\n \"span_links\" JSON NULL,\n TIME INDEX (\"timestamp\"),\n PRIMARY KEY (\"service_name\")\n)\n\nENGINE=mito\nWITH(\n 'comment' = 'Created on insertion',\n append_mode = 'true',\n 'greptime.semantic.entity.service.id' = 'service_name',\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\n 'greptime.semantic.signal_type' = 'trace',\n 'greptime.semantic.source' = 'opentelemetry',\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\n table_data_model = 'greptime_trace_v1'\n)"]]"#; + let expected_ddl = r#"[["trace_table_part1","CREATE TABLE IF NOT EXISTS \"trace_table_part1\" (\n \"timestamp\" TIMESTAMP(9) NOT NULL,\n \"timestamp_end\" TIMESTAMP(9) NULL COMPRESSION WITH(encoding = 'delta_binary_packed', type = 'zstd'),\n \"duration_nano\" BIGINT NULL,\n \"parent_span_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM') COMPRESSION WITH(encoding = 'plain', page_size = '64KiB', type = 'zstd'),\n \"trace_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_id\" STRING NULL COMPRESSION WITH(encoding = 'plain', page_size = '64KiB', type = 'zstd'),\n \"span_kind\" STRING NULL,\n \"span_name\" STRING NULL,\n \"span_status_code\" STRING NULL,\n \"span_status_message\" STRING NULL,\n \"trace_state\" STRING NULL,\n \"scope_name\" STRING NULL,\n \"scope_version\" STRING NULL,\n \"service_name\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_attributes.net.peer.ip\" STRING NULL,\n \"span_attributes.peer.service\" STRING NULL,\n \"span_events\" JSON NULL,\n \"span_links\" JSON NULL,\n TIME INDEX (\"timestamp\"),\n PRIMARY KEY (\"service_name\")\n)\n\nENGINE=mito\nWITH(\n 'comment' = 'Created on insertion',\n append_mode = 'true',\n 'greptime.semantic.entity.service.id' = 'service_name',\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\n 'greptime.semantic.signal_type' = 'trace',\n 'greptime.semantic.source' = 'opentelemetry',\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\n table_data_model = 'greptime_trace_v1'\n)"]]"#; validate_data( "otlp_traces", &client, @@ -9860,7 +9860,7 @@ pub async fn test_otlp_traces_v1(store_type: StorageType) { ) .await; assert_eq!(StatusCode::OK, res.status()); - let expected_ddl = r#"[["trace_table_part4","CREATE TABLE IF NOT EXISTS \"trace_table_part4\" (\n \"timestamp\" TIMESTAMP(9) NOT NULL,\n \"timestamp_end\" TIMESTAMP(9) NULL,\n \"duration_nano\" BIGINT NULL,\n \"parent_span_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"trace_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_id\" STRING NULL,\n \"span_kind\" STRING NULL,\n \"span_name\" STRING NULL,\n \"span_status_code\" STRING NULL,\n \"span_status_message\" STRING NULL,\n \"trace_state\" STRING NULL,\n \"scope_name\" STRING NULL,\n \"scope_version\" STRING NULL,\n \"service_name\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_attributes.net.peer.ip\" STRING NULL,\n \"span_attributes.peer.service\" STRING NULL,\n \"span_events\" JSON NULL,\n \"span_links\" JSON NULL,\n TIME INDEX (\"timestamp\"),\n PRIMARY KEY (\"service_name\")\n)\nPARTITION ON COLUMNS (\"trace_id\") (\n trace_id < '4',\n trace_id >= '4' AND trace_id < '8',\n trace_id >= '8' AND trace_id < 'c',\n trace_id >= 'c'\n)\nENGINE=mito\nWITH(\n 'comment' = 'Created on insertion',\n append_mode = 'true',\n 'greptime.semantic.entity.service.id' = 'service_name',\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\n 'greptime.semantic.signal_type' = 'trace',\n 'greptime.semantic.source' = 'opentelemetry',\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\n table_data_model = 'greptime_trace_v1'\n)"]]"#; + let expected_ddl = r#"[["trace_table_part4","CREATE TABLE IF NOT EXISTS \"trace_table_part4\" (\n \"timestamp\" TIMESTAMP(9) NOT NULL,\n \"timestamp_end\" TIMESTAMP(9) NULL COMPRESSION WITH(encoding = 'delta_binary_packed', type = 'zstd'),\n \"duration_nano\" BIGINT NULL,\n \"parent_span_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM') COMPRESSION WITH(encoding = 'plain', page_size = '64KiB', type = 'zstd'),\n \"trace_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_id\" STRING NULL COMPRESSION WITH(encoding = 'plain', page_size = '64KiB', type = 'zstd'),\n \"span_kind\" STRING NULL,\n \"span_name\" STRING NULL,\n \"span_status_code\" STRING NULL,\n \"span_status_message\" STRING NULL,\n \"trace_state\" STRING NULL,\n \"scope_name\" STRING NULL,\n \"scope_version\" STRING NULL,\n \"service_name\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_attributes.net.peer.ip\" STRING NULL,\n \"span_attributes.peer.service\" STRING NULL,\n \"span_events\" JSON NULL,\n \"span_links\" JSON NULL,\n TIME INDEX (\"timestamp\"),\n PRIMARY KEY (\"service_name\")\n)\nPARTITION ON COLUMNS (\"trace_id\") (\n trace_id < '4',\n trace_id >= '4' AND trace_id < '8',\n trace_id >= '8' AND trace_id < 'c',\n trace_id >= 'c'\n)\nENGINE=mito\nWITH(\n 'comment' = 'Created on insertion',\n append_mode = 'true',\n 'greptime.semantic.entity.service.id' = 'service_name',\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\n 'greptime.semantic.signal_type' = 'trace',\n 'greptime.semantic.source' = 'opentelemetry',\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\n table_data_model = 'greptime_trace_v1'\n)"]]"#; validate_data( "otlp_traces", &client, @@ -9897,7 +9897,7 @@ pub async fn test_otlp_traces_v1(store_type: StorageType) { ) .await; assert_eq!(StatusCode::OK, res.status()); - let expected_ddl = r#"[["trace_table_part32","CREATE TABLE IF NOT EXISTS \"trace_table_part32\" (\n \"timestamp\" TIMESTAMP(9) NOT NULL,\n \"timestamp_end\" TIMESTAMP(9) NULL,\n \"duration_nano\" BIGINT NULL,\n \"parent_span_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"trace_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_id\" STRING NULL,\n \"span_kind\" STRING NULL,\n \"span_name\" STRING NULL,\n \"span_status_code\" STRING NULL,\n \"span_status_message\" STRING NULL,\n \"trace_state\" STRING NULL,\n \"scope_name\" STRING NULL,\n \"scope_version\" STRING NULL,\n \"service_name\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_attributes.net.peer.ip\" STRING NULL,\n \"span_attributes.peer.service\" STRING NULL,\n \"span_events\" JSON NULL,\n \"span_links\" JSON NULL,\n TIME INDEX (\"timestamp\"),\n PRIMARY KEY (\"service_name\")\n)\nPARTITION ON COLUMNS (\"trace_id\") (\n trace_id < '08',\n trace_id >= '08' AND trace_id < '10',\n trace_id >= '10' AND trace_id < '18',\n trace_id >= '18' AND trace_id < '20',\n trace_id >= '20' AND trace_id < '28',\n trace_id >= '28' AND trace_id < '30',\n trace_id >= '30' AND trace_id < '38',\n trace_id >= '38' AND trace_id < '40',\n trace_id >= '40' AND trace_id < '48',\n trace_id >= '48' AND trace_id < '50',\n trace_id >= '50' AND trace_id < '58',\n trace_id >= '58' AND trace_id < '60',\n trace_id >= '60' AND trace_id < '68',\n trace_id >= '68' AND trace_id < '70',\n trace_id >= '70' AND trace_id < '78',\n trace_id >= '78' AND trace_id < '80',\n trace_id >= '80' AND trace_id < '88',\n trace_id >= '88' AND trace_id < '90',\n trace_id >= '90' AND trace_id < '98',\n trace_id >= '98' AND trace_id < 'a0',\n trace_id >= 'a0' AND trace_id < 'a8',\n trace_id >= 'a8' AND trace_id < 'b0',\n trace_id >= 'b0' AND trace_id < 'b8',\n trace_id >= 'b8' AND trace_id < 'c0',\n trace_id >= 'c0' AND trace_id < 'c8',\n trace_id >= 'c8' AND trace_id < 'd0',\n trace_id >= 'd0' AND trace_id < 'd8',\n trace_id >= 'd8' AND trace_id < 'e0',\n trace_id >= 'e0' AND trace_id < 'e8',\n trace_id >= 'e8' AND trace_id < 'f0',\n trace_id >= 'f0' AND trace_id < 'f8',\n trace_id >= 'f8'\n)\nENGINE=mito\nWITH(\n 'comment' = 'Created on insertion',\n append_mode = 'true',\n 'greptime.semantic.entity.service.id' = 'service_name',\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\n 'greptime.semantic.signal_type' = 'trace',\n 'greptime.semantic.source' = 'opentelemetry',\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\n table_data_model = 'greptime_trace_v1'\n)"]]"#; + let expected_ddl = r#"[["trace_table_part32","CREATE TABLE IF NOT EXISTS \"trace_table_part32\" (\n \"timestamp\" TIMESTAMP(9) NOT NULL,\n \"timestamp_end\" TIMESTAMP(9) NULL COMPRESSION WITH(encoding = 'delta_binary_packed', type = 'zstd'),\n \"duration_nano\" BIGINT NULL,\n \"parent_span_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM') COMPRESSION WITH(encoding = 'plain', page_size = '64KiB', type = 'zstd'),\n \"trace_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_id\" STRING NULL COMPRESSION WITH(encoding = 'plain', page_size = '64KiB', type = 'zstd'),\n \"span_kind\" STRING NULL,\n \"span_name\" STRING NULL,\n \"span_status_code\" STRING NULL,\n \"span_status_message\" STRING NULL,\n \"trace_state\" STRING NULL,\n \"scope_name\" STRING NULL,\n \"scope_version\" STRING NULL,\n \"service_name\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_attributes.net.peer.ip\" STRING NULL,\n \"span_attributes.peer.service\" STRING NULL,\n \"span_events\" JSON NULL,\n \"span_links\" JSON NULL,\n TIME INDEX (\"timestamp\"),\n PRIMARY KEY (\"service_name\")\n)\nPARTITION ON COLUMNS (\"trace_id\") (\n trace_id < '08',\n trace_id >= '08' AND trace_id < '10',\n trace_id >= '10' AND trace_id < '18',\n trace_id >= '18' AND trace_id < '20',\n trace_id >= '20' AND trace_id < '28',\n trace_id >= '28' AND trace_id < '30',\n trace_id >= '30' AND trace_id < '38',\n trace_id >= '38' AND trace_id < '40',\n trace_id >= '40' AND trace_id < '48',\n trace_id >= '48' AND trace_id < '50',\n trace_id >= '50' AND trace_id < '58',\n trace_id >= '58' AND trace_id < '60',\n trace_id >= '60' AND trace_id < '68',\n trace_id >= '68' AND trace_id < '70',\n trace_id >= '70' AND trace_id < '78',\n trace_id >= '78' AND trace_id < '80',\n trace_id >= '80' AND trace_id < '88',\n trace_id >= '88' AND trace_id < '90',\n trace_id >= '90' AND trace_id < '98',\n trace_id >= '98' AND trace_id < 'a0',\n trace_id >= 'a0' AND trace_id < 'a8',\n trace_id >= 'a8' AND trace_id < 'b0',\n trace_id >= 'b0' AND trace_id < 'b8',\n trace_id >= 'b8' AND trace_id < 'c0',\n trace_id >= 'c0' AND trace_id < 'c8',\n trace_id >= 'c8' AND trace_id < 'd0',\n trace_id >= 'd0' AND trace_id < 'd8',\n trace_id >= 'd8' AND trace_id < 'e0',\n trace_id >= 'e0' AND trace_id < 'e8',\n trace_id >= 'e8' AND trace_id < 'f0',\n trace_id >= 'f0' AND trace_id < 'f8',\n trace_id >= 'f8'\n)\nENGINE=mito\nWITH(\n 'comment' = 'Created on insertion',\n append_mode = 'true',\n 'greptime.semantic.entity.service.id' = 'service_name',\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\n 'greptime.semantic.signal_type' = 'trace',\n 'greptime.semantic.source' = 'opentelemetry',\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\n table_data_model = 'greptime_trace_v1'\n)"]]"#; validate_data( "otlp_traces", &client, @@ -13012,7 +13012,7 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { .await; assert_eq!(StatusCode::OK, res.status()); - let trace_table_sql = "[[\"mytable\",\"CREATE TABLE IF NOT EXISTS \\\"mytable\\\" (\\n \\\"timestamp\\\" TIMESTAMP(9) NOT NULL,\\n \\\"timestamp_end\\\" TIMESTAMP(9) NULL,\\n \\\"duration_nano\\\" BIGINT NULL,\\n \\\"parent_span_id\\\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\\n \\\"trace_id\\\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\\n \\\"span_id\\\" STRING NULL,\\n \\\"span_kind\\\" STRING NULL,\\n \\\"span_name\\\" STRING NULL,\\n \\\"span_status_code\\\" STRING NULL,\\n \\\"span_status_message\\\" STRING NULL,\\n \\\"trace_state\\\" STRING NULL,\\n \\\"scope_name\\\" STRING NULL,\\n \\\"scope_version\\\" STRING NULL,\\n \\\"service_name\\\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\\n \\\"span_attributes.operation.type\\\" STRING NULL,\\n \\\"span_attributes.net.peer.ip\\\" STRING NULL,\\n \\\"span_attributes.peer.service\\\" STRING NULL,\\n \\\"span_events\\\" JSON NULL,\\n \\\"span_links\\\" JSON NULL,\\n TIME INDEX (\\\"timestamp\\\"),\\n PRIMARY KEY (\\\"service_name\\\")\\n)\\nPARTITION ON COLUMNS (\\\"trace_id\\\") (\\n trace_id < '1',\\n trace_id >= '1' AND trace_id < '2',\\n trace_id >= '2' AND trace_id < '3',\\n trace_id >= '3' AND trace_id < '4',\\n trace_id >= '4' AND trace_id < '5',\\n trace_id >= '5' AND trace_id < '6',\\n trace_id >= '6' AND trace_id < '7',\\n trace_id >= '7' AND trace_id < '8',\\n trace_id >= '8' AND trace_id < '9',\\n trace_id >= '9' AND trace_id < 'a',\\n trace_id >= 'a' AND trace_id < 'b',\\n trace_id >= 'b' AND trace_id < 'c',\\n trace_id >= 'c' AND trace_id < 'd',\\n trace_id >= 'd' AND trace_id < 'e',\\n trace_id >= 'e' AND trace_id < 'f',\\n trace_id >= 'f'\\n)\\nENGINE=mito\\nWITH(\\n 'comment' = 'Created on insertion',\\n append_mode = 'true',\\n 'greptime.semantic.entity.service.id' = 'service_name',\\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\\n 'greptime.semantic.signal_type' = 'trace',\\n 'greptime.semantic.source' = 'opentelemetry',\\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\\n table_data_model = 'greptime_trace_v1',\\n ttl = '7days'\\n)\"]]"; + let trace_table_sql = "[[\"mytable\",\"CREATE TABLE IF NOT EXISTS \\\"mytable\\\" (\\n \\\"timestamp\\\" TIMESTAMP(9) NOT NULL,\\n \\\"timestamp_end\\\" TIMESTAMP(9) NULL COMPRESSION WITH(encoding = 'delta_binary_packed', type = 'zstd'),\\n \\\"duration_nano\\\" BIGINT NULL,\\n \\\"parent_span_id\\\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM') COMPRESSION WITH(encoding = 'plain', page_size = '64KiB', type = 'zstd'),\\n \\\"trace_id\\\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\\n \\\"span_id\\\" STRING NULL COMPRESSION WITH(encoding = 'plain', page_size = '64KiB', type = 'zstd'),\\n \\\"span_kind\\\" STRING NULL,\\n \\\"span_name\\\" STRING NULL,\\n \\\"span_status_code\\\" STRING NULL,\\n \\\"span_status_message\\\" STRING NULL,\\n \\\"trace_state\\\" STRING NULL,\\n \\\"scope_name\\\" STRING NULL,\\n \\\"scope_version\\\" STRING NULL,\\n \\\"service_name\\\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\\n \\\"span_attributes.operation.type\\\" STRING NULL,\\n \\\"span_attributes.net.peer.ip\\\" STRING NULL,\\n \\\"span_attributes.peer.service\\\" STRING NULL,\\n \\\"span_events\\\" JSON NULL,\\n \\\"span_links\\\" JSON NULL,\\n TIME INDEX (\\\"timestamp\\\"),\\n PRIMARY KEY (\\\"service_name\\\")\\n)\\nPARTITION ON COLUMNS (\\\"trace_id\\\") (\\n trace_id < '1',\\n trace_id >= '1' AND trace_id < '2',\\n trace_id >= '2' AND trace_id < '3',\\n trace_id >= '3' AND trace_id < '4',\\n trace_id >= '4' AND trace_id < '5',\\n trace_id >= '5' AND trace_id < '6',\\n trace_id >= '6' AND trace_id < '7',\\n trace_id >= '7' AND trace_id < '8',\\n trace_id >= '8' AND trace_id < '9',\\n trace_id >= '9' AND trace_id < 'a',\\n trace_id >= 'a' AND trace_id < 'b',\\n trace_id >= 'b' AND trace_id < 'c',\\n trace_id >= 'c' AND trace_id < 'd',\\n trace_id >= 'd' AND trace_id < 'e',\\n trace_id >= 'e' AND trace_id < 'f',\\n trace_id >= 'f'\\n)\\nENGINE=mito\\nWITH(\\n 'comment' = 'Created on insertion',\\n append_mode = 'true',\\n 'greptime.semantic.entity.service.id' = 'service_name',\\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\\n 'greptime.semantic.signal_type' = 'trace',\\n 'greptime.semantic.source' = 'opentelemetry',\\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\\n table_data_model = 'greptime_trace_v1',\\n ttl = '7days'\\n)\"]]"; validate_data( "trace_v1_create_table", &client, diff --git a/tests/cases/standalone/common/create/create_with_compression.result b/tests/cases/standalone/common/create/create_with_compression.result index 8b0e25a7c0f..77a21288c94 100644 --- a/tests/cases/standalone/common/create/create_with_compression.result +++ b/tests/cases/standalone/common/create/create_with_compression.result @@ -2,8 +2,8 @@ CREATE TABLE spans ( ts TIMESTAMP TIME INDEX, trace_id STRING SKIPPING INDEX, prompt STRING COMPRESSION WITH (type = 'zstd', level = 3, page_size = '8MiB'), - reply STRING COMPRESSION WITH (level = 9), - tokens BIGINT COMPRESSION WITH (page_size = '64KiB'), + reply STRING COMPRESSION WITH (level = 9, encoding = 'plain'), + tokens BIGINT COMPRESSION WITH (page_size = '64KiB', encoding = 'delta_binary_packed'), ) WITH (append_mode = 'true'); Affected Rows: 0 @@ -17,8 +17,8 @@ SHOW CREATE TABLE spans; | | "ts" TIMESTAMP(3) NOT NULL, | | | "trace_id" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'), | | | "prompt" STRING NULL COMPRESSION WITH(level = '3', page_size = '8MiB', type = 'zstd'), | -| | "reply" STRING NULL COMPRESSION WITH(level = '9', type = 'zstd'), | -| | "tokens" BIGINT NULL COMPRESSION WITH(page_size = '64KiB', type = 'zstd'), | +| | "reply" STRING NULL COMPRESSION WITH(encoding = 'plain', level = '9', type = 'zstd'), | +| | "tokens" BIGINT NULL COMPRESSION WITH(encoding = 'delta_binary_packed', page_size = '64KiB', type = 'zstd'), | | | TIME INDEX ("ts") | | | ) | | | | @@ -53,6 +53,10 @@ SELECT ts, trace_id, prompt, reply, tokens FROM spans ORDER BY ts; | 1970-01-01T00:00:00.003 | t2 | another prompt | | | +-------------------------+----------+--------------------------------+------------+--------+ +ALTER TABLE spans MODIFY COLUMN tokens STRING; + +Error: 1004(InvalidArguments), Invalid compression option: encoding delta_binary_packed requires an integer or timestamp column, but column tokens is String + DROP TABLE spans; Affected Rows: 0 @@ -73,6 +77,14 @@ CREATE TABLE bad (ts TIMESTAMP TIME INDEX, s STRING COMPRESSION WITH (codec = 'z Error: 1004(InvalidArguments), Invalid column option, column name: s, error: invalid COMPRESSION option: codec +CREATE TABLE bad (ts TIMESTAMP TIME INDEX, s STRING COMPRESSION WITH (encoding = 'dictionary')); + +Error: 1004(InvalidArguments), Invalid compression option: invalid encoding: dictionary, expected: 'plain' or 'delta_binary_packed' + +CREATE TABLE bad (ts TIMESTAMP TIME INDEX, s STRING COMPRESSION WITH (encoding = 'delta_binary_packed')); + +Error: 1004(InvalidArguments), Invalid compression option: encoding delta_binary_packed requires an integer or timestamp column, but column s is String + CREATE TABLE bad (ts TIMESTAMP TIME INDEX, j JSON2 COMPRESSION WITH (level = 3)); Error: 1004(InvalidArguments), Invalid compression option: compression options are not supported on column j of type Json2"" diff --git a/tests/cases/standalone/common/create/create_with_compression.sql b/tests/cases/standalone/common/create/create_with_compression.sql index 166c92be952..af817e250c0 100644 --- a/tests/cases/standalone/common/create/create_with_compression.sql +++ b/tests/cases/standalone/common/create/create_with_compression.sql @@ -2,8 +2,8 @@ CREATE TABLE spans ( ts TIMESTAMP TIME INDEX, trace_id STRING SKIPPING INDEX, prompt STRING COMPRESSION WITH (type = 'zstd', level = 3, page_size = '8MiB'), - reply STRING COMPRESSION WITH (level = 9), - tokens BIGINT COMPRESSION WITH (page_size = '64KiB'), + reply STRING COMPRESSION WITH (level = 9, encoding = 'plain'), + tokens BIGINT COMPRESSION WITH (page_size = '64KiB', encoding = 'delta_binary_packed'), ) WITH (append_mode = 'true'); SHOW CREATE TABLE spans; @@ -17,6 +17,8 @@ ADMIN flush_table('spans'); SELECT ts, trace_id, prompt, reply, tokens FROM spans ORDER BY ts; +ALTER TABLE spans MODIFY COLUMN tokens STRING; + DROP TABLE spans; CREATE TABLE bad (ts TIMESTAMP TIME INDEX, s STRING COMPRESSION WITH (level = 0)); @@ -27,6 +29,10 @@ CREATE TABLE bad (ts TIMESTAMP TIME INDEX, s STRING COMPRESSION WITH (page_size CREATE TABLE bad (ts TIMESTAMP TIME INDEX, s STRING COMPRESSION WITH (codec = 'zstd')); +CREATE TABLE bad (ts TIMESTAMP TIME INDEX, s STRING COMPRESSION WITH (encoding = 'dictionary')); + +CREATE TABLE bad (ts TIMESTAMP TIME INDEX, s STRING COMPRESSION WITH (encoding = 'delta_binary_packed')); + CREATE TABLE bad (ts TIMESTAMP TIME INDEX, j JSON2 COMPRESSION WITH (level = 3)); CREATE TABLE bad (ts TIMESTAMP TIME INDEX, s STRING COMPRESSION WITH (level = [3])); diff --git a/tests/compatibility/cases/column_encoding/case.toml b/tests/compatibility/cases/column_encoding/case.toml new file mode 100644 index 00000000000..567515eb10f --- /dev/null +++ b/tests/compatibility/cases/column_encoding/case.toml @@ -0,0 +1,8 @@ +name = "column_encoding" +reason = "Persist COMPRESSION WITH (encoding = ...) column options and read plain and delta encoded SSTs across restart, after upgrading from a version without the option." +introduced_by = "PR #9629" +topologies = ["distributed", "standalone"] +from_range = ["*"] +to_range = [">=v1.4.0-alpha.0"] +features = ["table", "compression", "encoding", "write", "read"] +owner = "storage" diff --git a/tests/compatibility/cases/column_encoding/setup.sql b/tests/compatibility/cases/column_encoding/setup.sql new file mode 100644 index 00000000000..3d0e8ce78e0 --- /dev/null +++ b/tests/compatibility/cases/column_encoding/setup.sql @@ -0,0 +1,3 @@ +CREATE TABLE legacy_spans (ts TIMESTAMP TIME INDEX, span_id STRING, ts_end TIMESTAMP(9)); +INSERT INTO legacy_spans VALUES (1, 'a1', 1000), (2, 'a2', 2000); +ADMIN FLUSH_TABLE('legacy_spans'); diff --git a/tests/compatibility/cases/column_encoding/verify.result b/tests/compatibility/cases/column_encoding/verify.result new file mode 100644 index 00000000000..4ff40a6d31a --- /dev/null +++ b/tests/compatibility/cases/column_encoding/verify.result @@ -0,0 +1,67 @@ +SELECT * FROM legacy_spans ORDER BY ts; + ++-------------------------+---------+----------------------------+ +| ts | span_id | ts_end | ++-------------------------+---------+----------------------------+ +| 1970-01-01T00:00:00.001 | a1 | 1970-01-01T00:00:00.000001 | +| 1970-01-01T00:00:00.002 | a2 | 1970-01-01T00:00:00.000002 | ++-------------------------+---------+----------------------------+ + +CREATE TABLE encoded_spans ( + ts TIMESTAMP TIME INDEX, + span_id STRING COMPRESSION WITH (encoding = 'plain', page_size = '64KiB'), + ts_end TIMESTAMP(9) COMPRESSION WITH (encoding = 'delta_binary_packed') +); + +Affected Rows: 0 + +INSERT INTO encoded_spans SELECT * FROM legacy_spans; + +Affected Rows: 2 + +ADMIN FLUSH_TABLE('encoded_spans'); + ++------------------------------------+ +| ADMIN FLUSH_TABLE('encoded_spans') | ++------------------------------------+ +| 0 | ++------------------------------------+ + +-- SQLNESS ARG restart=true +SHOW CREATE TABLE encoded_spans; + ++---------------+---------------------------------------------------------------------------------------------------+ +| Table | Create Table | ++---------------+---------------------------------------------------------------------------------------------------+ +| encoded_spans | CREATE TABLE IF NOT EXISTS "encoded_spans" ( | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | "span_id" STRING NULL COMPRESSION WITH(encoding = 'plain', page_size = '64KiB', type = 'zstd'), | +| | "ts_end" TIMESTAMP(9) NULL COMPRESSION WITH(encoding = 'delta_binary_packed', type = 'zstd'), | +| | TIME INDEX ("ts") | +| | ) | +| | | +| | ENGINE=mito | +| | | ++---------------+---------------------------------------------------------------------------------------------------+ + +INSERT INTO encoded_spans VALUES (3, 'b3', 3000); + +Affected Rows: 1 + +ADMIN FLUSH_TABLE('encoded_spans'); + ++------------------------------------+ +| ADMIN FLUSH_TABLE('encoded_spans') | ++------------------------------------+ +| 0 | ++------------------------------------+ + +SELECT * FROM encoded_spans ORDER BY ts; + ++-------------------------+---------+----------------------------+ +| ts | span_id | ts_end | ++-------------------------+---------+----------------------------+ +| 1970-01-01T00:00:00.001 | a1 | 1970-01-01T00:00:00.000001 | +| 1970-01-01T00:00:00.002 | a2 | 1970-01-01T00:00:00.000002 | +| 1970-01-01T00:00:00.003 | b3 | 1970-01-01T00:00:00.000003 | ++-------------------------+---------+----------------------------+ diff --git a/tests/compatibility/cases/column_encoding/verify.sql b/tests/compatibility/cases/column_encoding/verify.sql new file mode 100644 index 00000000000..5ae6cd787b9 --- /dev/null +++ b/tests/compatibility/cases/column_encoding/verify.sql @@ -0,0 +1,20 @@ +SELECT * FROM legacy_spans ORDER BY ts; + +CREATE TABLE encoded_spans ( + ts TIMESTAMP TIME INDEX, + span_id STRING COMPRESSION WITH (encoding = 'plain', page_size = '64KiB'), + ts_end TIMESTAMP(9) COMPRESSION WITH (encoding = 'delta_binary_packed') +); + +INSERT INTO encoded_spans SELECT * FROM legacy_spans; + +ADMIN FLUSH_TABLE('encoded_spans'); + +-- SQLNESS ARG restart=true +SHOW CREATE TABLE encoded_spans; + +INSERT INTO encoded_spans VALUES (3, 'b3', 3000); + +ADMIN FLUSH_TABLE('encoded_spans'); + +SELECT * FROM encoded_spans ORDER BY ts; diff --git a/tests/compatibility/cases/downgrade_compatibility/case.toml b/tests/compatibility/cases/downgrade_compatibility/case.toml index ab0496e523d..560efad0b32 100644 --- a/tests/compatibility/cases/downgrade_compatibility/case.toml +++ b/tests/compatibility/cases/downgrade_compatibility/case.toml @@ -1,8 +1,8 @@ name = "downgrade_compatibility" -reason = "Verify v1.1.4 can reopen tables whose region WAL options or byte-stream-split (BSS) float SSTs were written by the current binary, and can read all rows of an append-only table flushed and compacted with preserve_row_sequence enabled." +reason = "Verify v1.1.4 can reopen tables whose region WAL options, byte-stream-split (BSS) float SSTs, or plain and delta encoded column SSTs were written by the current binary, and can read all rows of an append-only table flushed and compacted with preserve_row_sequence enabled." introduced_by = "fix: preserve legacy region WAL options format; preserve_row_sequence" topologies = ["distributed", "standalone"] from_range = [">=v1.2.0"] to_range = ["=v1.1.4"] -features = ["table", "wal", "downgrade", "append", "preserve_row_sequence", "sst", "float", "byte_stream_split"] +features = ["table", "wal", "downgrade", "append", "preserve_row_sequence", "sst", "float", "byte_stream_split", "encoding"] owner = "metasrv" diff --git a/tests/compatibility/cases/downgrade_compatibility/setup.sql b/tests/compatibility/cases/downgrade_compatibility/setup.sql index d86ea21dfa5..bd0610693e3 100644 --- a/tests/compatibility/cases/downgrade_compatibility/setup.sql +++ b/tests/compatibility/cases/downgrade_compatibility/setup.sql @@ -58,3 +58,13 @@ INSERT INTO t_fulltext_bloom_v2_downgrade VALUES (2, '错误error日志'); ADMIN FLUSH_TABLE('t_fulltext_bloom_v2_downgrade'); + +CREATE TABLE t_column_encoding_downgrade ( + ts TIMESTAMP TIME INDEX, + span_id STRING COMPRESSION WITH (encoding = 'plain', page_size = '64KiB'), + ts_end TIMESTAMP(9) COMPRESSION WITH (encoding = 'delta_binary_packed') +); + +INSERT INTO t_column_encoding_downgrade VALUES (1, 'a1', 1000), (2, 'a2', 2000), (3, NULL, NULL); + +ADMIN FLUSH_TABLE('t_column_encoding_downgrade'); diff --git a/tests/compatibility/cases/downgrade_compatibility/verify.result b/tests/compatibility/cases/downgrade_compatibility/verify.result index 698403244a9..d73e23f58fc 100644 --- a/tests/compatibility/cases/downgrade_compatibility/verify.result +++ b/tests/compatibility/cases/downgrade_compatibility/verify.result @@ -47,3 +47,13 @@ SELECT ts, msg FROM t_fulltext_bloom_v2_downgrade WHERE msg @@ '错误error日 +-------------------------+---------------+ | 1970-01-01T00:00:00.002 | 错误error日志 | +-------------------------+---------------+ + +SELECT ts, span_id, ts_end FROM t_column_encoding_downgrade ORDER BY ts; + ++-------------------------+---------+----------------------------+ +| ts | span_id | ts_end | ++-------------------------+---------+----------------------------+ +| 1970-01-01T00:00:00.001 | a1 | 1970-01-01T00:00:00.000001 | +| 1970-01-01T00:00:00.002 | a2 | 1970-01-01T00:00:00.000002 | +| 1970-01-01T00:00:00.003 | | | ++-------------------------+---------+----------------------------+ diff --git a/tests/compatibility/cases/downgrade_compatibility/verify.sql b/tests/compatibility/cases/downgrade_compatibility/verify.sql index 56d476c1557..26664ff3ac9 100644 --- a/tests/compatibility/cases/downgrade_compatibility/verify.sql +++ b/tests/compatibility/cases/downgrade_compatibility/verify.sql @@ -11,3 +11,5 @@ ORDER BY ts, host; SELECT ts, msg FROM t_fulltext_bloom_v2_downgrade WHERE msg @@ 'hello_world' ORDER BY ts; SELECT ts, msg FROM t_fulltext_bloom_v2_downgrade WHERE msg @@ '错误error日志' ORDER BY ts; + +SELECT ts, span_id, ts_end FROM t_column_encoding_downgrade ORDER BY ts; From 8b0a1557025f5c27b5e66135b30dd18307cb64f8 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Fri, 16 Oct 2026 09:53:01 +0000 Subject: [PATCH 026/239] perf(otlp): chunk trace spans across resource groups (#9621) * perf(otlp): chunk trace spans across resource groups Signed-off-by: Dennis Zhuang * fix(otlp): plan spans of a mixed binary/JSON chunk individually Signed-off-by: Dennis Zhuang * docs(otlp): describe cross-group trace chunking Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/frontend/src/instance/otlp/README.md | 20 ++--- .../src/instance/otlp/trace_ingest.rs | 78 ++++++++++--------- .../src/instance/otlp/trace_ingest/tests.rs | 47 ++++++++++- tests-integration/tests/http.rs | 50 ++++++++++++ 4 files changed, 147 insertions(+), 48 deletions(-) diff --git a/src/frontend/src/instance/otlp/README.md b/src/frontend/src/instance/otlp/README.md index e308ebec184..db633dcd52c 100644 --- a/src/frontend/src/instance/otlp/README.md +++ b/src/frontend/src/instance/otlp/README.md @@ -19,7 +19,7 @@ flowchart TD A["POST /v1/otlp/v1/traces
OTLP protobuf"] --> B["HTTP handler
decode, select table and pipeline"] B --> C["Instance::traces
permission and interceptor"] C --> D["Derive trace conventions
parse resource/scope groups"] - D --> E["Chunk spans within each group"] + D --> E["Chunk spans across groups"] E --> F{"Trace pipeline"} F -->|greptime_trace_v0| G["v0 conversion
fixed columns and JSON attributes"] @@ -72,8 +72,10 @@ After permission checks, admission counts all spans in the request once per database. Chunks, schema retries, and derived lookup-table writes retain that admission; their actual writes still contribute WCU accounting. -`ingest_trace_spans` in [`trace_ingest.rs`](trace_ingest.rs) splits the spans in -each group into owned chunks. `trace_ingest_chunk_size` defaults to 512; setting +`ingest_trace_spans` in [`trace_ingest.rs`](trace_ingest.rs) splits the spans of +all resource/scope groups in the request into owned chunks; a chunk can mix +groups because every span carries its own resource and scope attributes. +`trace_ingest_chunk_size` defaults to 512; setting it to 0 disables splitting. The option is defined in [`service_config/otlp.rs`](../../service_config/otlp.rs). @@ -113,11 +115,11 @@ The v1 path performs these steps before the normal chunk writes: `TraceBatchSchema` that records observed scalar types, sparse row presence, and binary logical types. Drop the dense request after collecting its lightweight schema observations. -2. Mark a chunk for span-only fallback if one column contains both raw OTLP - bytes and arrays or key-value lists encoded as JSONB. Both use protobuf - `BinaryValue`, but a Greptime column cannot represent both logical schemas. - The chunk bypasses request-wide schema planning and is reconciled span by - span. +2. Split a chunk into single-span planning units if one column contains both + raw OTLP bytes and arrays or key-value lists encoded as JSONB. Both use + protobuf `BinaryValue`, but a Greptime column cannot represent both logical + schemas. Only the spans whose logical types conflict are removed from the + request-wide plan; the other spans keep contributing to it. 3. Aggregate the compatible observations into `TraceRequestSchema` and compare them with the existing table schema and the fixed types in [`trace_semconv.rs`](trace_semconv.rs). @@ -150,7 +152,7 @@ determines the next action when that preparation fails: | Reaction | Behavior | | --- | --- | -| `RetryPerSpan` | Prepare and write each span independently. Successful spans are accepted; deterministic per-span preparation failures are rejected. A v1 chunk with incompatible raw-binary and JSONB values starts here without attempting a chunk write. | +| `RetryPerSpan` | Prepare and write each span independently. Successful spans are accepted; deterministic per-span preparation failures are rejected. | | `DiscardChunk` | Do not retry the chunk at span granularity. Count every span in the chunk as rejected. | | `Propagate` | Return retryable or unclassified failures instead of reporting a partial success. Retry hints take precedence over status-code classification. | diff --git a/src/frontend/src/instance/otlp/trace_ingest.rs b/src/frontend/src/instance/otlp/trace_ingest.rs index 8b90a19ba49..99fc999398a 100644 --- a/src/frontend/src/instance/otlp/trace_ingest.rs +++ b/src/frontend/src/instance/otlp/trace_ingest.rs @@ -310,7 +310,6 @@ impl TraceFailureMessages { enum TraceChunkSchemaState { Prepared, ReconcilePerChunk, - SpanFallbackOnly, } impl TraceChunkSchemaState { @@ -1132,30 +1131,36 @@ impl Instance { let mut chunk_states = Vec::new(); // Convert each chunk once, then retain only its sparse values while // request-wide schema planning runs. - for chunk in groups - .into_iter() - .flat_map(|group| chunk_owned(group.spans, self.trace_ingest_chunk_size)) - { + for chunk in chunk_spans_across_groups(groups, self.trace_ingest_chunk_size) { let batch_index = chunks.len(); let span_metadata = chunk.iter().map(TraceSpanMetadata::from).collect(); let (table_data, batch_schema) = otlp::trace::v1::v1_to_main_table_data_with_schema(chunk)?; - let schema_state = if batch_schema.has_incompatible_logical_types() { - TraceChunkSchemaState::SpanFallbackOnly - } else { - TraceChunkSchemaState::Prepared - }; let (schema, rows) = table_data.into_schema_and_rows(); - if schema_state == TraceChunkSchemaState::Prepared { + let has_incompatible_logical_types = + batch_schema.has_incompatible_logical_types(); + if !has_incompatible_logical_types { request_schema.observe_batch_schema(batch_index, &schema, &batch_schema); } - chunks.push(TraceChunkRetry::try_new( + let retry = TraceChunkRetry::try_new( ingest_ctx.table_name.to_string(), Rows { schema, rows }, span_metadata, batch_schema.into_retry_columns(), - )?); - chunk_states.push(schema_state); + )?; + if has_incompatible_logical_types { + // A chunk schema cannot hold raw bytes and JSONB in one column. + // Plan each span on its own so the conflict only affects the + // spans that carry it, not the other columns of the chunk. + for retry in retry.into_single_span_chunks()? { + request_schema.observe_retry_chunk(chunks.len(), &retry)?; + chunks.push(retry); + chunk_states.push(TraceChunkSchemaState::Prepared); + } + } else { + chunks.push(retry); + chunk_states.push(TraceChunkSchemaState::Prepared); + } } if !chunks.is_empty() { @@ -1170,17 +1175,14 @@ impl Instance { .await?; } } else { - for group in groups { - let chunks = chunk_owned(group.spans, self.trace_ingest_chunk_size); - for chunk in chunks { - self.ingest_trace_chunk( - &ingest_ctx, - chunk, - main_ctx.clone(), - &mut ingest_state, - ) - .await?; - } + for chunk in chunk_spans_across_groups(groups, self.trace_ingest_chunk_size) { + self.ingest_trace_chunk( + &ingest_ctx, + chunk, + main_ctx.clone(), + &mut ingest_state, + ) + .await?; } } @@ -1564,18 +1566,6 @@ impl Instance { ctx: QueryContextRef, ingest_state: &mut TraceIngestState, ) -> ServerResult<()> { - if schema_state == TraceChunkSchemaState::SpanFallbackOnly { - Self::push_trace_failure_message( - &mut ingest_state.failure_messages, - ChunkFailureReaction::RetryPerSpan.as_metric_label(), - "incompatible_binary_and_json", - "Chunk fallback triggered by incompatible binary and JSON values".to_string(), - ); - return self - .ingest_trace_v1_rows_span_by_span(retry, ctx, ingest_state) - .await; - } - let mut request = retry.to_request()?; let _ = Self::prepare_trace_v1_chunk_rewrite( request_schema, @@ -2450,6 +2440,20 @@ fn trace_logical_types_incompatible( || !is_trace_reconcile_candidate_type(right_datatype) } +/// Chunks the spans of all resource/scope groups together. Every span carries +/// its own resource and scope attributes, so a chunk may mix groups; chunking per +/// group instead would turn a collector batch of many small groups into many +/// tiny inserts, each fanned out to every trace partition. +fn chunk_spans_across_groups( + groups: Vec, + chunk_size: usize, +) -> Vec> { + chunk_owned( + groups.into_iter().flat_map(|group| group.spans).collect(), + chunk_size, + ) +} + fn chunk_owned(items: Vec, chunk_size: usize) -> Vec> { if items.is_empty() { return Vec::new(); diff --git a/src/frontend/src/instance/otlp/trace_ingest/tests.rs b/src/frontend/src/instance/otlp/trace_ingest/tests.rs index d497db83e13..690621b6a3f 100644 --- a/src/frontend/src/instance/otlp/trace_ingest/tests.rs +++ b/src/frontend/src/instance/otlp/trace_ingest/tests.rs @@ -33,7 +33,7 @@ use opentelemetry_proto::tonic::common::v1::any_value::Value as OtlpValue; use opentelemetry_proto::tonic::common::v1::{AnyValue, ArrayValue, KeyValue}; use prometheus::IntCounter; use servers::otlp::trace::attributes::Attributes; -use servers::otlp::trace::span::{SpanEvents, SpanLinks, TraceSpan}; +use servers::otlp::trace::span::{SpanEvents, SpanLinks, TraceSpan, TraceSpanGroup}; use servers::otlp::trace::v1::{TraceBinaryType, TraceRetryColumn}; use servers::otlp::trace::{SERVICE_NAME_COLUMN, TraceAuxData}; use table::table_name::TableName; @@ -42,7 +42,7 @@ use super::{ ChunkFailureReaction, Instance, TraceAuxCache, TraceAuxCacheKey, TraceAuxEntry, TraceChunkRetry, TraceChunkSchemaState, TraceFailureMessages, TraceRequestSchema, TraceRequestSchemaPlan, TraceSpanMetadata, TraceTablePreAlter, chunk_owned, - wrap_trace_alter_failure, + chunk_spans_across_groups, wrap_trace_alter_failure, }; use crate::metrics::OTLP_TRACES_FAILURE_COUNT; @@ -296,6 +296,49 @@ fn test_chunk_owned() { assert!(chunk_owned::(Vec::new(), 0).is_empty()); } +#[test] +fn test_chunks_span_resource_groups() { + // Chunks fill up across resource groups instead of restarting per group. + let groups = ["a", "b", "c"] + .into_iter() + .map(|service| { + let mut spans = (0..3) + .map(|i| trace_span(&format!("{service}{i}"), vec![])) + .collect::>(); + for span in &mut spans { + span.service_name = Some(service.to_string()); + } + TraceSpanGroup { + service_name: Some(service.to_string()), + resource_attributes: Attributes::from(vec![]).into(), + scope_name: String::new(), + scope_version: String::new(), + scope_attributes: Attributes::from(vec![]).into(), + spans, + } + }) + .collect::>(); + + let chunks = chunk_spans_across_groups(groups, 4); + let services = chunks + .iter() + .map(|chunk| { + chunk + .iter() + .map(|span| span.service_name.as_deref().unwrap()) + .collect::>() + }) + .collect::>(); + assert_eq!( + services, + vec![ + vec!["a", "a", "a", "b"], + vec!["b", "b", "c", "c"], + vec!["c"] + ] + ); +} + #[test] fn test_classify_trace_prewrite_failure() { for status in [ diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 5b357f6904f..f522ad39ee5 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -10539,6 +10539,56 @@ pub async fn test_otlp_traces_v1(store_type: StorageType) { "unexpected error body: {body:?}" ); + // Raw bytes and JSONB values of `x` from different resources share one chunk. + // That conflict must not keep the unrelated resource `c` out of schema + // planning: `n` is planned as String, so `c`'s value is not coerced to the + // Int64 that per-span writes would pick from resource `a` first. + let mixed_table_name = "trace_type_mixed_binary_json"; + let mixed_resources = [ + ( + "mixed-a", + vec![ + json!({"key": "x", "value": {"bytesValue": "AQI="}}), + make_int_attr("n", 1), + ], + ), + ( + "mixed-b", + vec![json!({"key": "x", "value": {"arrayValue": {"values": [{"intValue": "1"}]}}})], + ), + ("mixed-c", vec![make_string_attr("n", "abc")]), + ]; + let mut mixed_req = ExportTraceServiceRequest::default(); + for (i, (service_name, attributes)) in mixed_resources.into_iter().enumerate() { + let id = 0x71 + i; + mixed_req.resource_spans.extend( + make_trace_v1_request( + service_name, + vec![make_trace_v1_span( + &format!("{id:032x}"), + &format!("{id:016x}"), + service_name, + 1_736_480_942_446_000_000, + 1_736_480_942_446_100_000, + attributes, + )], + ) + .resource_spans, + ); + } + let res = send_trace_v1_req(&client, mixed_table_name, mixed_req, false).await; + assert_eq!(StatusCode::OK, res.status()); + validate_data( + "otlp_traces_v1_mixed_binary_json_rows", + &client, + &format!( + "select service_name, \"span_attributes.n\" from {} order by service_name", + mixed_table_name + ), + r#"[["mixed-a","1"],["mixed-c","abc"]]"#, + ) + .await; + guard.remove_all().await; } From 0e82b29b23c0c4349954b11c7fe9df8de86650e9 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Fri, 16 Oct 2026 09:57:50 +0000 Subject: [PATCH 027/239] ci: add crash, restart, multi-region and drop cells to the S3 WAL benchmark (#9632) * ci: add crash, restart, multi-region and drop cells to the S3 WAL benchmark Extend the object store WAL S3 benchmark beyond one load per configuration: - Inputs for a full image reference, the configurations to run (R0 adds Raft Engine with the standalone default sync_write = false), extra environment for every server container, a restart after each load, and a JSON list of named cell groups that override those inputs. - Restart cells SIGKILL the loaded server, restart it on the same prefix and data home, and time /health, the first correct count and the log store recovery; min_objects fails a cell that left too few objects indexed. - Crash groups write striped 1000-row batches from several writers, kill the server at a random point, restart it and check every acknowledged batch against the table (durable: present exactly once; enqueued: no duplicate, no partial batch, losses counted). A pinned variant keeps one row of a second table unflushed and samples the indexed objects while it runs. - Idle windows scrape the server before and after a quiet period. - A load group may split the loaded table into equal event_time ranges, so every batch writes to every region, and checks that each region took rows. - Drop groups write a table, drop it, kill and restart the server, then write and flush another table and wait for the dropped table's objects to be collected. - The summary adds WAL traffic, objects per second, mean object size, entries and appends per object, S3 requests per operation, server and host CPU, the loader's own request latency, rows per region, crash cycles, and each restart's replay per region from the server log. Signed-off-by: jeremyhi * ci: report the object sequences a replay spans instead of objects replayed The width of a region's replayed entry range counts every object sequence in it, including objects that hold no entry of the region and are skipped, so the summary labels it as a span. Signed-off-by: jeremyhi * ci: fail a crash cycle that loses a batch an earlier restart recovered The oracle accepted every missing acknowledged batch in the enqueued mode, so a batch that one restart recovered and a later restart no longer found passed. Every cycle keeps what its restart found, and a cycle now fails in either mode unless every batch an earlier restart recovered is still present, so only the current cycle's backlog that was not durable at its kill may be missing. Signed-off-by: jeremyhi * ci: keep the pinned WAL pinned and label OpenDAL operations as such The pin table inherited the engine's 10-minute auto flush, so a longer pinned cycle could flush its row and release the WAL before the kill. It now sets its own auto_flush_interval of one day, longer than any crash group can run, and a pinned cycle fails unless its indexed objects never fall while it runs. The summary's request counts come from opendal_operation_duration_seconds, which counts OpenDAL calls rather than HTTP requests (a paginated list or a multipart write is one operation), so the fields and columns now say S3 operations. Signed-off-by: jeremyhi * ci: check a pinned cycle's retention by the collected-object counter Net growth of the indexed objects can hide collection, and short cycles take no periodic sample, so a pinned cycle now fails unless deleted_objects_total is the same in the scrape before its writers start and in the one before the kill, both of the same server process. The operation columns name OpenDAL. Signed-off-by: jeremyhi * ci: let a pinned cell collect only the objects below its pin The first write after a start lets the store collect the empty start object asynchronously, so a deletion after a cycle's baseline does not mean the pin let go. A pinned cell now scrapes the server before it writes the pin row and fails only once its servers together have collected more objects than were indexed then, counting the first server from that scrape and every restarted one from its start. Signed-off-by: jeremyhi * ci: test that a pinned cell's allowance spans its server restarts Two servers that each collect fewer objects than lie below the pin but more together now fail the cell in the second cycle. Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- .github/scripts/oswal-s3-bench-crash.py | 229 +++++++ .github/scripts/oswal-s3-bench-load.sh | 27 + .github/scripts/oswal-s3-bench-run.sh | 341 ++++++++-- .github/scripts/oswal-s3-bench-summary.py | 368 +++++++++- .github/scripts/oswal-s3-bench.py | 236 ++++++- .github/workflows/oswal-s3-bench.yml | 35 +- tests/perf/test_oswal_s3_bench_scripts.py | 790 +++++++++++++++++++++- 7 files changed, 1907 insertions(+), 119 deletions(-) create mode 100644 .github/scripts/oswal-s3-bench-crash.py create mode 100644 .github/scripts/oswal-s3-bench-load.sh diff --git a/.github/scripts/oswal-s3-bench-crash.py b/.github/scripts/oswal-s3-bench-crash.py new file mode 100644 index 00000000000..35b382aa24b --- /dev/null +++ b/.github/scripts/oswal-s3-bench-crash.py @@ -0,0 +1,229 @@ +#!/usr/bin/env python3 +# 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. + +"""Writer and oracle of the benchmark's crash cells. + +Every writer owns the stripe `first + index, first + index + writers, ...` of +batch sequence numbers, inserts `rows` rows per batch and stops at the first +statement that is not acknowledged, so a stripe has at most one attempted but +unacknowledged batch: the one in flight at the kill. Every request is logged +with its status and the time it ended, on the same clock as the kill stamps. +With `--seconds` the writers also stop once that long has passed, and the +command fails unless every writer wrote until then without a failure. +""" + +from __future__ import annotations + +import argparse +import collections +import json +import os +import sys +import threading +import time +import urllib.error +import urllib.parse +import urllib.request +from pathlib import Path + + +def insert(url: str, table: str, seq: int, rows: int) -> str: + values = ", ".join(f"({seq}, {seq * rows + i})" for i in range(rows)) + body = urllib.parse.urlencode({"sql": f"INSERT INTO {table} (seq, ts) VALUES {values}"}).encode() + try: + with urllib.request.urlopen(urllib.request.Request(url, data=body), timeout=60) as response: + return str(response.status) + except urllib.error.HTTPError as error: + return str(error.code) + except (OSError, urllib.error.URLError): + # The connection dropped: the request may or may not have been applied. + return "000" + + +def writer(args, index: int, deadline: float, failures: list) -> None: + directory, seq = Path(args.dir), args.first_seq + index + with open(directory / f"attempted-{index}.log", "a") as attempted, \ + open(directory / f"requests-{index}.log", "a") as requests, \ + open(directory / f"acked-{index}.log", "a") as acked: + while time.monotonic() < deadline: + attempted.write(f"{seq}\n") + attempted.flush() + code = insert(args.url, args.table, seq, args.rows) + requests.write(f"{seq} {code} {time.time_ns()}\n") + requests.flush() + if code != "200": + failures.append(f"writer {index}: batch {seq} ended with {code}") + break + acked.write(f"{seq}\n") + acked.flush() + seq += args.writers + (directory / f"exit-{index}.log").write_text(f"{time.time_ns()}\n") + + +def write(args) -> int: + os.makedirs(args.dir, exist_ok=True) + deadline = time.monotonic() + args.seconds if args.seconds else float("inf") + failures = [] + + def run(index): + try: + writer(args, index, deadline, failures) + except Exception as error: + failures.append(f"writer {index}: {error!r}") + raise + + threads = [threading.Thread(target=run, args=(i,)) for i in range(args.writers)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + # Crash cells end their writers with a kill; timed writers must have run to the deadline. + if args.seconds and (failures or not any(lines(Path(args.dir) / f"acked-{i}.log") for i in range(args.writers))): + print("; ".join(failures) or "no batch was acknowledged", file=sys.stderr) + return 1 + return 0 + + +def lines(path: Path) -> list: + return [line.split() for line in path.read_text().splitlines() if line.strip()] if path.exists() else [] + + +def found_batches(path: Path) -> dict: + """GreptimeDB's JSON answer to SELECT seq, count(*) ... GROUP BY seq, as batch -> row count.""" + return {int(seq): int(count) for seq, count in json.loads(path.read_text())["output"][0]["records"]["rows"]} + + +def check_kill(directory: Path, writers: int, before_ns: int, after_ns: int, window_max_ms: int) -> dict: + """A dropped connection is in flight when it ended at or after the stamp taken before the kill.""" + errors, in_flight, running = [], 0, 0 + if after_ns - before_ns > window_max_ms * 1_000_000: + errors.append(f"kill window {after_ns - before_ns}ns exceeds {window_max_ms}ms") + for index in range(writers): + exit_path = directory / f"exit-{index}.log" + if not exit_path.exists(): + errors.append(f"writer {index} left no exit time") + continue + running += int(exit_path.read_text()) >= before_ns + for seq, code, end_ns in lines(directory / f"requests-{index}.log"): + if code == "200": + continue + if code != "000": + errors.append(f"writer {index}: batch {seq} rejected with HTTP {code}") + elif int(end_ns) < before_ns: + errors.append(f"writer {index}: batch {seq} lost its connection before the kill") + else: + in_flight += 1 + if running == 0: + errors.append("no writer was running at the kill") + return {"kill_window_ns": after_ns - before_ns, "in_flight": in_flight, "writers_at_kill": running, + "kill_errors": errors} + + +def verify(run_dir: Path, found: dict, rows: int, enqueued: bool, cycle: int) -> dict: + """Checks the table, read as batch -> row count, against every cycle's writers so far and + against what every earlier restart recovered, and counts what the given cycle's writers + acknowledged, left and lost on their own.""" + acked, stripes, cycle_attempted, cycle_acked = [], collections.OrderedDict(), set(), set() + recovered_before = set() + for cycle_dir in sorted(run_dir.glob("cycle-*"), key=lambda p: int(p.name.split("-")[1])): + current = cycle_dir.name == f"cycle-{cycle}" + if int(cycle_dir.name.split("-")[1]) < cycle and (cycle_dir / "found.json").exists(): + recovered_before.update(found_batches(cycle_dir / "found.json")) + for path in sorted(cycle_dir.glob("attempted-*.log")): + stripes[(cycle_dir.name, path.name)] = [int(seq) for (seq,) in lines(path)] + if current: + cycle_attempted.update(stripes[(cycle_dir.name, path.name)]) + for path in sorted(cycle_dir.glob("acked-*.log")): + acked += [int(seq) for (seq,) in lines(path)] + if current: + cycle_acked.update(int(seq) for (seq,) in lines(path)) + acked_set = set(acked) + attempted = {seq for stripe in stripes.values() for seq in stripe} + errors = [] + duplicated = sorted(seq for seq, count in found.items() if count > rows) + partial = sorted(seq for seq, count in found.items() if count < rows) + if duplicated: + errors.append(f"batches present more than once: {duplicated[:20]}") + if partial: + errors.append(f"batches present in part: {partial[:20]}") + missing = sorted(acked_set - set(found)) + if missing and not enqueued: + errors.append(f"acknowledged but absent: {missing[:20]}") + # A batch an earlier restart recovered was durable then, so in either mode only the + # current cycle's backlog that was not durable at its kill may be missing. + vanished = sorted(recovered_before - set(found)) + if vanished: + errors.append(f"recovered by an earlier restart but absent: {vanished[:20]}") + unattempted = sorted(set(found) - attempted) + if unattempted: + errors.append(f"present but never attempted: {unattempted[:20]}") + # The unacknowledged tail of each stripe was in flight at a kill. + in_flight = set() + for key, stripe in stripes.items(): + prefix = 0 + while prefix < len(stripe) and stripe[prefix] in acked_set: + prefix += 1 + if any(seq in acked_set for seq in stripe[prefix:]): + errors.append(f"{key}: acknowledged after a failure") + in_flight.update(stripe[prefix:]) + extra = sorted(set(found) - acked_set) + stray = [seq for seq in extra if seq not in in_flight] + if stray and not enqueued: + errors.append(f"present, unacknowledged and not in flight at a kill: {stray[:20]}") + if len(acked) != len(acked_set): + errors.append("a batch was acknowledged twice") + present = set(found) & cycle_attempted + return {"cycle_acked_batches": len(cycle_acked), "cycle_present_batches": len(present), + "cycle_lost_batches": len(cycle_acked - present), "cycle_lost_rows": len(cycle_acked - present) * rows, + "cycle_extra_batches": len(present - cycle_acked), + "total_found_batches": len(found), "total_acked_batches": len(acked_set), + "total_lost_batches": len(missing), "total_extra_batches": len(extra), "verify_errors": errors} + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + commands = parser.add_subparsers(dest="command", required=True) + w = commands.add_parser("write") + w.add_argument("--dir", required=True) + w.add_argument("--url", required=True) + w.add_argument("--table", default="crash") + for name in ("writers", "first-seq", "rows"): + w.add_argument(f"--{name}", type=int, required=True) + w.add_argument("--seconds", type=float, default=0, help="stop after this long; 0 writes until a failure") + c = commands.add_parser("check") + c.add_argument("--run-dir", type=Path, required=True) + c.add_argument("--cycle", type=int, required=True) + c.add_argument("--found", type=Path, required=True, help="GreptimeDB JSON output of SELECT seq, count(*)") + c.add_argument("--mode", choices=("durable", "enqueued"), required=True) + for name in ("writers", "rows", "before-ns", "after-ns", "window-max-ms"): + c.add_argument(f"--{name}", type=int, required=True) + c.add_argument("--pin-rows", type=int, help="rows of the pin table, whose one row was acknowledged first") + args = parser.parse_args() + if args.command == "write": + return write(args) + result = {"cycle": args.cycle} + result |= check_kill(args.run_dir / f"cycle-{args.cycle}", args.writers, args.before_ns, args.after_ns, + args.window_max_ms) + result |= verify(args.run_dir, found_batches(args.found), args.rows, + args.mode == "enqueued", args.cycle) + if args.pin_rows is not None and (args.pin_rows > 1 or args.mode == "durable" and args.pin_rows != 1): + result["verify_errors"].append(f"the pin table holds {args.pin_rows} rows, not 1") + result["passed"] = not result["kill_errors"] and not result["verify_errors"] + print(json.dumps(result)) + return 0 if result["passed"] else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/.github/scripts/oswal-s3-bench-load.sh b/.github/scripts/oswal-s3-bench-load.sh new file mode 100644 index 00000000000..1748fb40d66 --- /dev/null +++ b/.github/scripts/oswal-s3-bench-load.sh @@ -0,0 +1,27 @@ +#!/usr/bin/env bash +# Runs inside the o11ybench runtime image: the load of scripts/agent-observability/greptimedb.sh +# with the same arguments plus the loader's --skip-create-table, for a table the cell created. +set -euo pipefail + +mkdir "/results/$RESULT_NAME" +mkdir -p /results/corpus +cp /work/corpus/summary.json /work/corpus/evidence.json /results/corpus/ +w="$PWD/workloads/logbench-agent-observability" +r="/results/$RESULT_NAME" +prepare_workers=${PREPARE_WORKERS:-$(python3 -c 'import sys; sys.path.insert(0, "scripts"); from benchmark_workers import effective_cpus; print(effective_cpus())')} +o11ybench-greptimedb-loader --workload logbench-agent-observability --prepare-workers "$prepare_workers" \ + --skip-create-table \ + --grpc-endpoint "$GRPC_ENDPOINT" --sql-endpoint "$ENDPOINT" \ + --database public --table agent_observations --batch-size "${GREPTIMEDB_BATCH_SIZE:-10000}" --write-mode stream --max-in-flight "${MAX_IN_FLIGHT:-1}" \ + --flush-table --request-timeout-seconds 900 \ + --schema "$w/schema.json" --physical-index-profiles "$w/physical-index-profiles.json" \ + --data /work/corpus/agent_observations.jsonl \ + --profile-summary /work/corpus/summary.json --evidence /work/corpus/evidence.json \ + --transcript-output "$r/load-transcript.jsonl" \ + --target-metadata-output "$r/target-metadata.json" \ + --greptimedb-image "$TARGET_IMAGE" --greptimedb-version "$TARGET_VERSION" \ + --greptimedb-digest "$TARGET_DIGEST" \ + --startup-input=standalone --startup-input=start \ + --startup-input=--http-addr=0.0.0.0:4000 \ + --startup-input=--grpc-bind-addr=0.0.0.0:4001 \ + > "$r/load-result.json" diff --git a/.github/scripts/oswal-s3-bench-run.sh b/.github/scripts/oswal-s3-bench-run.sh index a9835c1e5f1..f822a4a9b6f 100644 --- a/.github/scripts/oswal-s3-bench-run.sh +++ b/.github/scripts/oswal-s3-bench-run.sh @@ -22,12 +22,12 @@ check_runtime() { } pull_images() { - : "${RESULTS:?}" "${RUNTIME_IMAGE:?}" "${GREPTIMEDB_TAG:?}" + : "${RESULTS:?}" "${RUNTIME_IMAGE:?}" "${GREPTIMEDB_REF:?}" mkdir -p "$RESULTS/images" bash "$here/pull-benchmark-image.sh" "$RUNTIME_IMAGE" - bash "$here/pull-benchmark-image.sh" "greptime/greptimedb:$GREPTIMEDB_TAG" + bash "$here/pull-benchmark-image.sh" "$GREPTIMEDB_REF" docker image inspect "$RUNTIME_IMAGE" > "$RESULTS/images/runtime.json" - docker image inspect "greptime/greptimedb:$GREPTIMEDB_TAG" > "$RESULTS/images/greptimedb.json" + docker image inspect "$GREPTIMEDB_REF" > "$RESULTS/images/greptimedb.json" } wal_disk() { @@ -46,25 +46,240 @@ wal_disk() { findmnt -no SOURCE,FSTYPE -T "$root" | tee "$RESULTS/wal-disk.txt" } +# One corpus per row count the load cells use; crash cells write their own rows. generate() { : "${PLAN:?}" "${RESULTS:?}" "${CORPUS_DIR:?}" "${OWNER:?}" "${RUNTIME_IMAGE:?}" - local args=() arg - while IFS= read -r arg; do args+=("$arg"); done < <(jq -r '.generator_args[]' "$PLAN") - mkdir -p "$CORPUS_DIR" - runtime --name "$OWNER-generator" --mount "type=bind,src=$CORPUS_DIR,dst=/work/corpus" \ - "$RUNTIME_IMAGE" scripts/generate_agent_observations.py "${args[@]}" \ - --output /work/corpus/agent_observations.jsonl \ - --summary-output /work/corpus/summary.json --evidence-output /work/corpus/evidence.json \ - > "$RESULTS/generator.log" 2>&1 - jq -e --argjson rows "$(jq .rows "$PLAN")" '.line_count == $rows' "$CORPUS_DIR/summary.json" >/dev/null + local corpus args arg dir rows + while IFS= read -r corpus; do + args=() dir="$CORPUS_DIR/$(jq -r .dir <<< "$corpus")" rows=$(jq .rows <<< "$corpus") + while IFS= read -r arg; do args+=("$arg"); done < <(jq -r '.generator_args[]' <<< "$corpus") + mkdir -p "$dir" + runtime --name "$OWNER-generator-$rows" --mount "type=bind,src=$dir,dst=/work/corpus" \ + "$RUNTIME_IMAGE" scripts/generate_agent_observations.py "${args[@]}" \ + --output /work/corpus/agent_observations.jsonl \ + --summary-output /work/corpus/summary.json --evidence-output /work/corpus/evidence.json \ + >> "$RESULTS/generator.log" 2>&1 + jq -e --argjson rows "$rows" '.line_count == $rows' "$dir/summary.json" >/dev/null + done < <(jq -c '.corpora[]' "$PLAN") } -# Scrape once, after the load succeeded or failed, while the server still runs. +# Scrape once per server, after the load succeeded or failed, while it still runs: +# the loaded server into metrics.prom, a restarted one into metrics-restart.prom. scrape_metrics() { curl --noproxy '*' -fsS --max-time 30 http://127.0.0.1:4000/metrics > "$dir/metrics.all" || { rm -f "$dir/metrics.all"; return 1; } - grep -E '^(# (HELP|TYPE) )?greptime_logstore_object_store_wal_' "$dir/metrics.all" > "$dir/metrics.prom" || true + grep -E '^(# (HELP|TYPE) )?(greptime_logstore_object_store_wal_|opendal_operation_(duration_seconds_count|errors_total))' \ + "$dir/metrics.all" > "$dir/metrics$phase.prom" || true rm -f "$dir/metrics.all" - metrics_scraped=true + case "$phase" in "") metrics_scraped=true ;; -restart) restart_scraped=true ;; esac +} + +# A crash cycle cannot repeat a scrape once its server is killed, so it retries. +# The named WAL sample must be in the scrape, or it does not count. +scrape_retry() { + local required=$1 attempt + for attempt in 1 2 3; do + scrape_metrics && grep -q "^greptime_logstore_object_store_wal_$required " "$dir/metrics$phase.prom" && return 0 + sleep 1 + done + return 1 +} + +sql() { curl --noproxy '*' -fsS --max-time 30 --data-urlencode "sql=$1" 'http://127.0.0.1:4000/v1/sql?db=public'; } +now_ns() { python3 -c 'import time; print(time.time_ns())'; } + +now() { date +%s.%N; } +elapsed() { awk -v a="$1" -v b="$(now)" 'BEGIN {printf "%.3f", b - a}'; } +# CPU seconds the server's cgroup used, and busy and total jiffies of the host. +server_cpu() { docker exec "$cid" cat /sys/fs/cgroup/cpu.stat 2>/dev/null | awk '/^usage_usec / {printf "%.3f", $2 / 1e6}' || true; } +host_cpu() { awk '/^cpu / {print $2 + $3 + $4 + $7 + $8 + $9, $2 + $3 + $4 + $5 + $6 + $7 + $8 + $9}' /proc/stat 2>/dev/null || true; } + +# A server's log is named after the phase it started in, whatever phase stops it. +start_server() { + server_phase=$phase + cid=$(docker run -d --init --name "$name$phase" --label "o11ybench.owner=$OWNER" \ + --add-host host.docker.internal:host-gateway \ + -p 127.0.0.1:4000:4000 -p 127.0.0.1:4001:4001 "${docker_env[@]}" ${mounts[@]+"${mounts[@]}"} \ + "$GREPTIMEDB_IMAGE" standalone start --http-addr=0.0.0.0:4000 --grpc-bind-addr=0.0.0.0:4001) + local deadline=$((SECONDS + 300)) + until curl --noproxy '*' -fsS --max-time 5 http://127.0.0.1:4000/health >/dev/null 2>&1; do + [[ $(docker inspect --format '{{.State.Running}}' "$cid") == true ]] || { echo "$id: server$phase exited" >&2; exit 1; } + ((SECONDS < deadline)) || { echo "$id: server$phase not healthy in 300s" >&2; exit 1; } + sleep 0.2 + done +} + +# kill sends SIGKILL, killed collects a server the caller already killed, and +# no argument stops it gracefully. +stop_server() { + case "${1:-}" in + kill) docker kill --signal KILL "$cid" >/dev/null ;; + killed) docker wait "$cid" >/dev/null ;; + *) docker stop -t 120 "$cid" >/dev/null ;; + esac + docker logs "$cid" > "$dir/server$server_phase.log" 2>&1 + docker rm -fv "$cid" >/dev/null + cid= +} + +# Kill the loaded server, start the same configuration on the same prefix and +# data home, and time it until /health and until the loaded table answers. +restart() { + local table count started deadline + table=$(jq -er .table "$dir/greptimedb-1/load-result.json") + scrape_metrics || { echo "$id: metrics scrape failed" >&2; exit 4; } + # The objects recovery has to read are those still indexed at the kill. + indexed_before_kill=$(awk '/^greptime_logstore_object_store_wal_indexed_objects / {print $2}' "$dir/metrics.prom") + stop_server kill + phase=-restart started=$(now) + start_server + restart_seconds=$(elapsed "$started") + deadline=$((SECONDS + 300)) + # Regions may still be opening after /health. + until count=$(curl --noproxy '*' -fsS --max-time 30 --data-urlencode "sql=SELECT count(*) FROM $table" \ + 'http://127.0.0.1:4000/v1/sql?db=public' | jq -er '.output[0].records.rows[0][0]'); do + ((SECONDS < deadline)) || { echo "$id: $table not readable in 300s after restart" >&2; exit 5; } + sleep 0.2 + done + restart_query_seconds=$(elapsed "$started") restart_count=$count + [[ "$count" == "$(jq .rows <<< "$spec")" ]] || { echo "$id: $count rows after restart" >&2; exit 5; } + # The restart was measured, but on fewer objects than the cell asked for. + awk -v n="${indexed_before_kill:-0}" -v min="$(jq -r '.min_objects // 0' <<< "$spec")" 'BEGIN {exit !(n >= min)}' \ + || { echo "$id: ${indexed_before_kill:-no} objects indexed at the kill, fewer than min_objects" >&2; exit 7; } +} + +# Pins the WAL with one row in a second table before the writers start, if asked, +# and samples the object counts every 30 s while they run. +sample_objects() { + set +e + local body + while sleep "${OSWAL_SAMPLE_SECONDS:-30}"; do + # A failed scrape is a missing sample, never zero objects. + if ! body=$(curl --noproxy '*' -fsS --max-time 10 http://127.0.0.1:4000/metrics 2>/dev/null) \ + || ! grep -q '^greptime_logstore_object_store_wal_indexed_objects ' <<< "$body"; then + echo "$(date +%s) missing" >> "$dir/objects.tsv" + continue + fi + awk -v t="$(date +%s)" ' + /^greptime_logstore_object_store_wal_indexed_objects / {i = $2} + /^greptime_logstore_object_store_wal_deleted_objects_total / {d = $2} + /^greptime_logstore_object_store_wal_created_objects_total / {c = $2} + END {print t, i + 0, d + 0, c + 0}' <<< "$body" >> "$dir/objects.tsv" + done +} + +# Writers insert striped batches until a SIGKILL at a random point stops the +# server; it restarts on the same prefix and data home and the table is checked +# against what was acknowledged. Nothing is cleaned between cycles. +crash_cycles() { + local writers rows cycles kill_min kill_max cycle cdir wpid spid= delay before after started restart_s query_s next=0 pin_count pin_args + local load_start complete pin_left pin_base deleted sql_url=${OSWAL_SQL_URL:-http://127.0.0.1:4000/v1/sql?db=public} + writers=$(jq -r .concurrency <<< "$spec") rows=$(jq -r .batch_rows "$PLAN") cycles=$(jq -r .cycles <<< "$spec") + kill_min=$(jq -r '.kill_after[0]' <<< "$spec") kill_max=$(jq -r '.kill_after[1]' <<< "$spec") + if jq -e .pin <<< "$spec" >/dev/null; then + # Only the objects indexed before the pin row lie below it, such as the empty start object that + # the row's write lets the store collect at any later point; the servers may collect no more. + phase=-pin; scrape_retry indexed_objects || { echo "$id: metrics scrape failed" >&2; exit 4; } + read -r pin_left pin_base < <(awk '/^greptime_logstore_object_store_wal_indexed_objects / {i = $2} + /^greptime_logstore_object_store_wal_deleted_objects_total / {d = $2} END {printf "%d %d\n", i, d}' "$dir/metrics-pin.prom") + # The pin table's own auto flush outlasts every cycle a crash group allows, so its row holds the WAL. + sql "CREATE TABLE pin (v BIGINT, ts TIMESTAMP TIME INDEX) WITH ('auto_flush_interval' = '1d')" >/dev/null + sql "INSERT INTO pin VALUES (1, 0)" >/dev/null + fi + sql "CREATE TABLE crash (seq BIGINT, ts TIMESTAMP TIME INDEX) WITH (append_mode = 'true')" >/dev/null + for ((cycle = 1; cycle <= cycles; cycle++)); do + cdir="$dir/cycle-$cycle" complete=true; mkdir -p "$cdir" + # The counters of this server process before its writers start are the cycle's baseline. + phase=-start-$cycle; scrape_retry indexed_objects || complete=false + load_start=$(now_ns) + python3 "$here/oswal-s3-bench-crash.py" write --dir "$cdir" --url "$sql_url" \ + --writers "$writers" --first-seq "$next" --rows "$rows" & + wpid=$! + if jq -e .pin <<< "$spec" >/dev/null; then sample_objects & spid=$!; fi + delay=$(awk -v a="$kill_min" -v b="$kill_max" 'BEGIN {srand(); printf "%.3f", a + rand() * (b - a)}') + sleep "$delay" + phase=-cycle-$cycle; scrape_retry indexed_objects || complete=false + before=$(now_ns); docker kill --signal KILL "$cid" >/dev/null; after=$(now_ns) + wait "$wpid" + if [[ -n "$spid" ]]; then kill "$spid" 2>/dev/null || true; wait "$spid" 2>/dev/null || true; spid=; fi + stop_server killed + phase=-restart-$cycle started=$(now) + start_server + restart_s=$(elapsed "$started") + until sql "SELECT seq, count(*) FROM crash GROUP BY seq" > "$cdir/found.json" 2>/dev/null; do + (($(elapsed "$started" | cut -d. -f1) < 300)) || { echo "$id: crash not readable after restart $cycle" >&2; exit 5; } + sleep 0.2 + done + query_s=$(elapsed "$started") + scrape_retry recovery_seconds_count || complete=false + pin_count=null pin_args=() + if jq -e .pin <<< "$spec" >/dev/null; then + pin_count=$(sql "SELECT count(*) FROM pin" | jq -er '.output[0].records.rows[0][0]') pin_args=(--pin-rows "$pin_count") + fi + local verdict=0 + python3 "$here/oswal-s3-bench-crash.py" check --run-dir "$dir" --cycle "$cycle" --found "$cdir/found.json" \ + --mode "$(jq -r .wal.ACK_MODE <<< "$spec")" --writers "$writers" --rows "$rows" \ + --before-ns "$before" --after-ns "$after" --window-max-ms 2000 ${pin_args[@]+"${pin_args[@]}"} \ + > "$cdir/check.json" || verdict=$? + jq -c --arg delay "$delay" --arg restart "$restart_s" --arg query "$query_s" --argjson pin "$pin_count" \ + --argjson load "$(awk -v a="$load_start" -v b="$before" 'BEGIN {printf "%.3f", (b - a) / 1e9}')" \ + --argjson complete "$complete" \ + '. + {kill_after_seconds: ($delay | tonumber), load_seconds: $load, restart_seconds: ($restart | tonumber), + restart_query_seconds: ($query | tonumber), pin_rows: $pin, metrics_complete: $complete}' \ + "$cdir/check.json" >> "$dir/cycles.jsonl" + ((verdict == 0)) || { echo "$id: cycle $cycle failed its check" >&2; exit 6; } + # A cycle without its scrapes is not a measurement, whatever the table held. + [[ "$complete" == true ]] || { echo "$id: cycle $cycle lost a metrics scrape" >&2; exit 4; } + # Each server counts its collections from its own start: the first one from its pin scrape, + # every restarted one from zero, and an absent counter has not moved. + if jq -e .pin <<< "$spec" >/dev/null; then + deleted=$(awk '/^greptime_logstore_object_store_wal_deleted_objects_total / {d = $2} END {printf "%d", d}' \ + "$dir/metrics-cycle-$cycle.prom") + pin_left=$((pin_left - (deleted - pin_base))) pin_base=0 + ((pin_left >= 0)) || { echo "$id: cycle $cycle collected WAL objects the pin held" >&2; exit 11; } + fi + next=$(cat "$cdir"/attempted-*.log | sort -n | tail -1); next=$((next + 1)) + done +} + +# Writes the table `dropped` for write_seconds, drops it and kills the server right +# away, while most of its objects are still indexed; restarts on the same prefix and +# data home, then writes and flushes the table `kept` and samples the indexed objects +# until the dropped table's objects are collected. +drop_restart() { + local writers rows seconds started deadline indexed recovered sql_url=${OSWAL_SQL_URL:-http://127.0.0.1:4000/v1/sql?db=public} + writers=$(jq -r .concurrency <<< "$spec") rows=$(jq -r .batch_rows "$PLAN") seconds=$(jq -r .write_seconds <<< "$spec") + sql "CREATE TABLE dropped (seq BIGINT, ts TIMESTAMP TIME INDEX) WITH (append_mode = 'true')" >/dev/null + sql "CREATE TABLE kept (seq BIGINT, ts TIMESTAMP TIME INDEX) WITH (append_mode = 'true')" >/dev/null + python3 "$here/oswal-s3-bench-crash.py" write --dir "$dir/dropped" --url "$sql_url" --table dropped \ + --writers "$writers" --first-seq 0 --rows "$rows" --seconds "$seconds" \ + || { echo "$id: the writers of dropped stopped before $seconds s" >&2; exit 10; } + phase=-written; scrape_retry indexed_objects || { echo "$id: metrics scrape failed" >&2; exit 4; } + sql "DROP TABLE dropped" >/dev/null + stop_server kill + phase=-restart started=$(now) + start_server + restart_seconds=$(elapsed "$started") + scrape_retry recovery_seconds_count || { echo "$id: metrics scrape failed" >&2; exit 4; } + recovered=$(awk '/^greptime_logstore_object_store_wal_recovered_objects_total / {print $2}' "$dir/metrics-restart.prom") + python3 "$here/oswal-s3-bench-crash.py" write --dir "$dir/kept" --url "$sql_url" --table kept \ + --writers 1 --first-seq 0 --rows "$rows" --seconds 2 \ + || { echo "$id: the writer of kept stopped before 2 s" >&2; exit 10; } + sql "ADMIN FLUSH_TABLE('kept')" >/dev/null + # The flushed watermark becomes durable with the next object, so every sample writes one row. + deadline=$((SECONDS + ${OSWAL_COLLECT_SECONDS:-120})) + while :; do + sql "INSERT INTO kept VALUES (-1, $SECONDS)" >/dev/null + # The samples have their own file; the final scrape keeps the restart's name. + phase=-collected + scrape_retry indexed_objects || { phase=-restart; echo "$id: metrics scrape failed" >&2; exit 4; } + phase=-restart + indexed=$(awk '/^greptime_logstore_object_store_wal_indexed_objects / {print $2}' "$dir/metrics-collected.prom") + printf '%s %s\n' "$(date +%s)" "$indexed" >> "$dir/collected.tsv" + # A start object, the newest object and kept's own may stay; the dropped table's go. + awk -v i="$indexed" -v r="${recovered:-0}" 'BEGIN {exit !(i <= 3 || 2 * i < r)}' && break + ((SECONDS < deadline)) || { echo "$id: $indexed of ${recovered:-?} recovered objects still indexed" >&2; exit 9; } + sleep 1 + done } finish() { @@ -74,15 +289,24 @@ finish() { if [[ -n "$cid" ]]; then # A load that succeeded without its metrics is not a measurement. scrape_metrics || { echo "$id: metrics scrape failed" >&2; [[ $rc != 0 ]] || rc=4; } - docker stop -t 120 "$cid" >/dev/null - docker logs "$cid" > "$dir/server.log" 2>&1 - docker rm -fv "$cid" >/dev/null + stop_server fi [[ -z "$wal_dir" ]] || sudo -n rm -rf "$wal_dir" + [[ -z "$home_dir" ]] || sudo -n rm -rf "$home_dir" jq -n --argjson spec "$spec" --arg device "$wal_device" --arg version "$version" \ --arg load_seconds "$load_seconds" --argjson exit_code "$rc" --argjson scraped "$metrics_scraped" \ + --argjson restart_scraped "$restart_scraped" --arg restart_seconds "$restart_seconds" \ + --arg restart_query_seconds "$restart_query_seconds" --arg restart_count "$restart_count" \ + --arg indexed_before_kill "$indexed_before_kill" \ + --arg server_cpu_seconds "$server_cpu_seconds" --arg host_cpu_busy "$host_cpu_busy" --arg host_cpus "$(nproc 2>/dev/null)" \ '$spec + {wal_device: $device, version: $version, exit_code: $exit_code, metrics_scraped: $scraped, - load_wall_seconds: ($load_seconds | tonumber? // null)}' > "$dir/cell.json" + load_wall_seconds: ($load_seconds | tonumber? // null), + server_cpu_seconds: ($server_cpu_seconds | tonumber? // null), + host_cpu_busy: ($host_cpu_busy | tonumber? // null), host_cpus: ($host_cpus | tonumber? // null), + restart_scraped: $restart_scraped, restart_seconds: ($restart_seconds | tonumber? // null), + restart_query_seconds: ($restart_query_seconds | tonumber? // null), + restart_count: ($restart_count | tonumber? // null), + indexed_before_kill: ($indexed_before_kill | tonumber? // null)}' > "$dir/cell.json" exit "$rc" } @@ -91,22 +315,30 @@ finish() { cell() { : "${PLAN:?}" "${RESULTS:?}" "${CORPUS_DIR:?}" "${OWNER:?}" "${RUNTIME_IMAGE:?}" "${GREPTIMEDB_IMAGE:?}" : "${RUN_PREFIX:?}" "${WAL_ROOT:?}" "${S3_BUCKET:?}" "${S3_REGION:?}" "${AWS_ACCESS_KEY_ID:?}" "${AWS_SECRET_ACCESS_KEY:?}" - spec=$1 cid= wal_dir= version= load_seconds= metrics_scraped=false id= - local concurrency name setting started ended + spec=$1 cid= wal_dir= home_dir= version= load_seconds= metrics_scraped=false id= phase= server_phase= + restart_scraped=false restart_seconds= restart_query_seconds= restart_count= server_cpu_seconds= host_cpu_busy= + indexed_before_kill= + local concurrency setting started cpu_before host_before host_after corpus partitions load script id=$(jq -r .id <<< "$spec"); concurrency=$(jq -r .concurrency <<< "$spec") - dir="$RESULTS/cells/$id"; name="$OWNER-$id" + dir="$RESULTS/cells/$id" name="$OWNER-$id" mkdir -p "$dir" # The root holds data/ and the WAL's default prefix wal/ side by side. local settings=(STORAGE__TYPE=S3 "STORAGE__BUCKET=$S3_BUCKET" "STORAGE__REGION=$S3_REGION" "STORAGE__ROOT=$RUN_PREFIX/$id") [[ -z "${S3_CONTAINER_ENDPOINT:-}" ]] || settings+=("STORAGE__ENDPOINT=$S3_CONTAINER_ENDPOINT") + # A restarted server needs the metadata the loaded one kept in its data home. + jq -e '.restart or .mode != "load"' <<< "$spec" >/dev/null && settings+=(STORAGE__DATA_HOME=/oswal_home) while IFS= read -r setting; do settings+=("$setting"); done \ < <(jq -r '.wal | to_entries[] | "WAL__\(.key)=\(.value)"' <<< "$spec") - local docker_env=() mounts=() + docker_env=() mounts=() for setting in "${settings[@]}"; do docker_env+=(-e "GREPTIMEDB_STANDALONE__$setting") printf 'GREPTIMEDB_STANDALONE__%s\n' "$setting" >> "$dir/config.env" done + while IFS= read -r setting; do + docker_env+=(-e "$setting") + printf '%s\n' "$setting" >> "$dir/config.env" + done < <(jq -r '.env // {} | to_entries[] | "\(.key)=\(.value)"' <<< "$spec") # Pass the two keys by name: their values never reach argv, a file or the artifact. export GREPTIMEDB_STANDALONE__STORAGE__ACCESS_KEY_ID=$AWS_ACCESS_KEY_ID export GREPTIMEDB_STANDALONE__STORAGE__SECRET_ACCESS_KEY=$AWS_SECRET_ACCESS_KEY @@ -119,33 +351,60 @@ cell() { else wal_device="s3://$S3_BUCKET/$RUN_PREFIX/$id/wal" fi + if jq -e '.restart or .mode != "load"' <<< "$spec" >/dev/null; then + home_dir="$WAL_ROOT/$id-home"; mkdir -p "$home_dir" + mounts+=(--mount "type=bind,src=$home_dir,dst=/oswal_home") + fi trap finish EXIT - cid=$(docker run -d --init --name "$name" --label "o11ybench.owner=$OWNER" \ - --add-host host.docker.internal:host-gateway \ - -p 127.0.0.1:4000:4000 -p 127.0.0.1:4001:4001 "${docker_env[@]}" ${mounts[@]+"${mounts[@]}"} \ - "$GREPTIMEDB_IMAGE" standalone start --http-addr=0.0.0.0:4000 --grpc-bind-addr=0.0.0.0:4001) - local deadline=$((SECONDS + 300)) - until curl --noproxy '*' -fsS --max-time 5 http://127.0.0.1:4000/health >/dev/null 2>&1; do - [[ $(docker inspect --format '{{.State.Running}}' "$cid") == true ]] || { echo "$id: server exited" >&2; exit 1; } - ((SECONDS < deadline)) || { echo "$id: server not healthy in 300s" >&2; exit 1; } - sleep 2 - done + start_server version=$(docker exec "$cid" greptime --version | sed -n 's/^version: //p') [[ -n "$version" ]] - started=$(date +%s.%N) - # Only the load phase, run exactly as run-agent-observability.sh runs it. + if jq -e '.mode == "crash"' <<< "$spec" >/dev/null; then crash_cycles; return; fi + if jq -e '.mode == "drop"' <<< "$spec" >/dev/null; then drop_restart; return; fi + corpus="$CORPUS_DIR/corpus-$(jq -r .rows <<< "$spec")" partitions=$(jq -r '.partitions // 1' <<< "$spec") + # Only the load phase, run exactly as run-agent-observability.sh runs it, or with the + # loader writing into a table split into regions that the cell created first. + load=(scripts/agent-observability/greptimedb.sh load) script=() + if ((partitions > 1)); then + sql "$(python3 "$here/oswal-s3-bench.py" table --summary "$corpus/summary.json" --partitions "$partitions")" >/dev/null + load=(/work/oswal-s3-bench-load.sh) + script=(--mount "type=bind,src=$here/oswal-s3-bench-load.sh,dst=/work/oswal-s3-bench-load.sh,readonly") + fi + cpu_before=$(server_cpu) host_before=$(host_cpu) + started=$(now) runtime --name "$name-load" --network host \ - --mount "type=bind,src=$CORPUS_DIR,dst=/work/corpus,readonly" --mount "type=bind,src=$dir,dst=/results" \ - -e GRPC_ENDPOINT=127.0.0.1:4001 -e "DATA_PROFILE=$(jq -r .data_profile "$PLAN")" -e RESULT_NAME=greptimedb-1 \ + --mount "type=bind,src=$corpus,dst=/work/corpus,readonly" ${script[@]+"${script[@]}"} \ + --mount "type=bind,src=$dir,dst=/results" \ + -e GRPC_ENDPOINT=127.0.0.1:4001 -e "DATA_PROFILE=$(jq -r .data_profile <<< "$spec")" -e RESULT_NAME=greptimedb-1 \ -e "ENDPOINT=http://127.0.0.1:4000/v1/sql?db=public&format=greptimedb_v1" \ -e "TARGET_IMAGE=$GREPTIMEDB_IMAGE" -e "TARGET_VERSION=$version" -e "TARGET_DIGEST=${GREPTIMEDB_IMAGE##*@}" \ -e "TARGET_ID=$name" -e TARGET_HOST=127.0.0.1 -e PREPARE_WORKERS= -e "MAX_IN_FLIGHT=$concurrency" \ -e "GREPTIMEDB_BATCH_SIZE=$(jq -r .batch_rows "$PLAN")" --entrypoint /bin/bash \ - "$RUNTIME_IMAGE" scripts/agent-observability/greptimedb.sh load > "$dir/load.log" 2>&1 - ended=$(date +%s.%N) - load_seconds=$(awk -v a="$started" -v b="$ended" 'BEGIN {printf "%.3f", b - a}') - jq -e --argjson rows "$(jq .rows "$PLAN")" '.count_before == 0 and .count_after == $rows' \ + "$RUNTIME_IMAGE" "${load[@]}" > "$dir/load.log" 2>&1 + load_seconds=$(elapsed "$started") + host_after=$(host_cpu) + server_cpu_seconds=$(awk -v a="$cpu_before" -v b="$(server_cpu)" 'BEGIN {if (a != "" && b != "") printf "%.3f", b - a}') + # Busy share of all host CPUs during the load, the loader included. + host_cpu_busy=$(awk -v a="$host_before" -v b="$host_after" \ + 'BEGIN {split(a, x, " "); split(b, y, " "); if (y[2] > x[2]) printf "%.4f", (y[1] - x[1]) / (y[2] - x[2])}') + jq -e --argjson rows "$(jq .rows <<< "$spec")" '.count_before == 0 and .count_after == $rows' \ "$dir/greptimedb-1/load-result.json" >/dev/null + if ((partitions > 1)); then + # The server's own account of the rows each region took. + sql "SELECT p.partition_name, s.region_id, s.region_rows, s.written_bytes_since_open + FROM information_schema.partitions p JOIN information_schema.region_statistics s + ON p.greptime_partition_id = s.region_id + WHERE p.table_schema = 'public' AND p.table_name = 'agent_observations' + ORDER BY s.region_id" > "$dir/regions.json" + jq -e --argjson n "$partitions" '[.output[0].records.rows[] | select(.[2] > 0)] | length == $n' \ + "$dir/regions.json" >/dev/null || { echo "$id: fewer than $partitions regions took rows" >&2; exit 8; } + fi + # An idle window: what the server sends with no writes, between two scrapes. + if jq -e '.idle_seconds > 0' <<< "$spec" >/dev/null; then + phase=-idle; scrape_metrics || { echo "$id: metrics scrape failed" >&2; exit 4; }; phase= + sleep "$(jq -r .idle_seconds <<< "$spec")" + fi + if jq -e '.restart' <<< "$spec" >/dev/null; then restart; fi } scrub() { diff --git a/.github/scripts/oswal-s3-bench-summary.py b/.github/scripts/oswal-s3-bench-summary.py index 67186d422da..305a647a9a4 100644 --- a/.github/scripts/oswal-s3-bench-summary.py +++ b/.github/scripts/oswal-s3-bench-summary.py @@ -26,14 +26,30 @@ from pathlib import Path PREFIX = "greptime_logstore_object_store_wal_" COUNTERS = ("created_objects_total", "stalled_appends_total", "poisoned_total", "create_failures_total", "create_conflicts_total") +# The Raft Engine controls every object store WAL cell is compared with: fsync per write and the default. +BASELINES = ("R", "R0") SAMPLE = re.compile(r'([a-zA-Z_:][a-zA-Z0-9_:]*)(?:\{([^}]*)\})?\s+(\S+)') +# Mito's line for every region it opened, with the WAL range it replayed. +REPLAY = re.compile(r"Replay WAL for region: (\d+)\(.*?provider: (\w+).*?rows recovered: (\d+), " + r"replay from entry id: (\d+), last entry id: (\d+),.*?elapsed: ([0-9.]+)(ns|µs|us|ms|s)\b") +UNITS = {"ns": 1e-9, "µs": 1e-6, "us": 1e-6, "ms": 1e-3, "s": 1.0} +# An object store WAL entry id is the object sequence above a 20-bit position. +POSITION_BITS = 20 def parse_metrics(text: str) -> dict: - """Map each unlabeled sample to its value and each histogram to its buckets.""" - samples, buckets = {}, {} + """Map each unlabeled sample to its value, each histogram to its buckets, and the S3 + operations (data files and WAL together, as OpenDAL counts its calls: a paginated list or a + multipart write is one operation over several HTTP requests) to their count per kind; the local cache + stores report through the same layer under another scheme and are left out.""" + samples, buckets, operations = {}, {}, {} for line in text.splitlines(): match = SAMPLE.fullmatch(line.strip()) + if match and match.group(1) == "opendal_operation_duration_seconds_count": + operation = re.search(r'operation="([^"]+)"', match.group(2) or "") + if operation and re.search(r'scheme="s3"', match.group(2) or ""): + operations[operation.group(1)] = operations.get(operation.group(1), 0) + float(match.group(3)) + continue if not match or not match.group(1).startswith(PREFIX): continue name, labels, value = match.group(1)[len(PREFIX):], match.group(2) or "", float(match.group(3)) @@ -42,7 +58,8 @@ def parse_metrics(text: str) -> dict: buckets.setdefault(name[:-len("_bucket")], []).append((float(le.group(1)), value)) else: samples[name] = value - return {"samples": samples, "buckets": {name: sorted(items) for name, items in buckets.items()}} + return {"samples": samples, "buckets": {name: sorted(items) for name, items in buckets.items()}, + "operations": operations} def quantile(q: float, buckets: list) -> float | None: @@ -64,25 +81,149 @@ def geomean(values: list) -> float | None: return math.exp(sum(map(math.log, values)) / len(values)) if values and all(v > 0 for v in values) else None +def read_metrics(path: Path) -> dict: + return parse_metrics(path.read_text() if path.exists() else "") + + +def ratio(a, b): + return a / b if a is not None and b else None + + +def replay(log: Path, metrics: dict) -> dict | None: + """What a restarted server replayed: every region it opened, the span of object sequences + its WAL range covers, and the log store's time serving region reads (`read_seconds`). The span + bounds the objects a region read from above: objects without its entries are skipped.""" + if not log.exists(): + return None + regions = [] + for match in REPLAY.finditer(log.read_text(errors="replace")): + region, provider, rows, first, last, elapsed, unit = match.groups() + # The range runs from the first entry to replay to the last one; it is empty when first > last. + span = None if provider != "ObjectStore" else 0 if int(first) > int(last) else \ + (int(last) >> POSITION_BITS) - (int(first) >> POSITION_BITS) + 1 + regions.append({"region": int(region), "rows": int(rows), "first_entry_id": int(first), + "last_entry_id": int(last), "seconds": float(elapsed) * UNITS[unit], "object_span": span}) + samples = metrics["samples"] + return {"regions": regions, "rows": sum(r["rows"] for r in regions), + "seconds": max((r["seconds"] for r in regions), default=None), + "object_span": sum(r["object_span"] for r in regions) + if regions and regions[0]["object_span"] is not None else None, + "read_seconds": samples.get("read_seconds_sum"), "reads": samples.get("read_seconds_count")} + + +def replayed(regions: list) -> str: + """Each region that replayed rows: its rows, the object sequences its range spans and its time.""" + return "; ".join(f"{r['region']}: {r['rows']:,} rows" + ( + f", object sequences {r['first_entry_id'] >> POSITION_BITS:,}–{r['last_entry_id'] >> POSITION_BITS:,}" + if r["object_span"] else "") + f", {r['seconds']:.2f} s" for r in regions if r["rows"]) or "—" + + +def client_latency(path: Path) -> tuple: + """The 50th and 99th percentile of the loader's own time per write request.""" + durations = [] + if path.exists(): + for line in path.read_text().splitlines(): + record = json.loads(line) if line.strip() else {} + if record.get("step") == "grpc_write_batch" and record.get("status") == "success": + durations.append(record["duration_seconds"]) + durations.sort() + return tuple(durations[max(0, math.ceil(q * len(durations)) - 1)] if durations else None for q in (0.5, 0.99)) + + +def region_rows(path: Path) -> list | None: + """Rows per region of a partitioned load cell, as the server reported them after the load.""" + if not path.exists(): + return None + return [int(row[2] or 0) for row in json.loads(path.read_text())["output"][0]["records"]["rows"]] + + +def interval_seconds(wal: dict) -> float: + value = wal.get("FLUSH_INTERVAL", "100ms") + return float(value[:-2]) / 1000 if value.endswith("ms") else float(value.rstrip("s")) + + +def traffic(end: dict, start: dict, seconds, wal: dict) -> dict: + """WAL traffic between two scrapes of one process over `seconds`; more objects than seal + ticks, or objects near the 8 MiB threshold, mean batches sealed on size rather than on the timer.""" + def delta(name): + value = end["samples"].get(name) + return None if value is None else value - start["samples"].get(name, 0) + created, written, objects = delta("created_objects_total"), delta("object_bytes_sum"), delta("object_bytes_count") + per_second = created / seconds if created and seconds else None + mean = written / objects if written and objects else None + return {"wal_mb_per_second": written / seconds / 1e6 if written and seconds else None, + "objects_per_second": per_second, "mean_object_bytes": mean, + "sealed_on_size": None if per_second is None else + per_second > 1.05 / interval_seconds(wal) or (mean or 0) >= 0.9 * 8 * 2**20, + "operations": {op: n - start["operations"].get(op, 0) for op, n in end["operations"].items() + if n != start["operations"].get(op, 0)}} + + +def crash_cycles(path: Path, wal: dict) -> list: + """Each cycle's check plus the loaded server's metrics before the kill and the restarted one's after.""" + cycles = [json.loads(line) for line in (path / "cycles.jsonl").read_text().splitlines() + if line.strip()] if (path / "cycles.jsonl").exists() else [] + for cycle in cycles: + start = read_metrics(path / f"metrics-start-{cycle['cycle']}.prom") + killed = read_metrics(path / f"metrics-cycle-{cycle['cycle']}.prom") + before = killed["samples"] + after = read_metrics(path / f"metrics-restart-{cycle['cycle']}.prom") + once = after["samples"].get("recovery_seconds_count") == 1 + cycle |= traffic(killed, start, cycle.get("load_seconds"), wal) + cycle |= {"created_before_kill": before.get("created_objects_total"), + "indexed_before_kill": before.get("indexed_objects"), + "recovery_seconds": after["samples"].get("recovery_seconds_sum") if once else None, + "recovered_objects": after["samples"].get("recovered_objects_total", 0 if once else None), + "restart_operations": after["operations"], + "replay": replay(path / f"server-restart-{cycle['cycle']}.log", after)} + return cycles + + def load_cell(path: Path, spec: dict | None = None) -> dict: """One cell's row; a planned cell that never recorded a result has no values, so no ratio uses its round.""" recorded = path / "cell.json" cell = json.loads(recorded.read_text()) if recorded.exists() else spec | {"exit_code": None, "missing": True} result_path = path / "greptimedb-1" / "load-result.json" result = json.loads(result_path.read_text()) if result_path.exists() else {} - metrics_path = path / "metrics.prom" - metrics = parse_metrics(metrics_path.read_text() if metrics_path.exists() else "") - ack = metrics["buckets"].get("append_ack_seconds", []) + final, restarted = read_metrics(path / "metrics.prom"), read_metrics(path / "metrics-restart.prom") + # With an idle window the scrape taken when the load ended describes the load. + metrics = read_metrics(path / "metrics-idle.prom") if cell.get("idle_seconds") else final load_seconds = result.get("load_seconds") + wall, cpu = cell.get("load_wall_seconds"), cell.get("server_cpu_seconds") row = { - "id": cell["id"], "round": cell["round"], "concurrency": cell["concurrency"], + "id": cell["id"], "group": cell.get("group", ""), "mode": cell.get("mode", "load"), + "round": cell["round"], "concurrency": cell["concurrency"], "config": cell["config"], "exit_code": cell["exit_code"], "wal_device": cell.get("wal_device"), - "missing": cell.get("missing", False), - "version": cell.get("version"), "load_wall_seconds": cell.get("load_wall_seconds"), + "missing": cell.get("missing", False), "env": cell.get("env", {}), + "version": cell.get("version"), "load_wall_seconds": wall, "rows": result.get("count_after"), "load_seconds": load_seconds, "rows_per_second": result.get("rows_per_second") if cell["exit_code"] == 0 else None, - "append_ack_p50_seconds": quantile(0.5, ack), "append_ack_p99_seconds": quantile(0.99, ack), + "server_cpu_cores": cpu / wall if cpu is not None and wall else None, + "host_cpu_busy": cell.get("host_cpu_busy"), "host_cpus": cell.get("host_cpus"), } + for name, key in (("append_ack_seconds", "append_ack"), ("seal_to_durable_seconds", "seal_to_durable"), + ("object_bytes", "object_bytes")): + buckets = metrics["buckets"].get(name, []) + row[f"{key}_p50"], row[f"{key}_p99"] = quantile(0.5, buckets), quantile(0.99, buckets) + row["client_p50"], row["client_p99"] = client_latency(path / "greptimedb-1" / "load-transcript.jsonl") + # An entry is one region's write; an append, counted only in the durable mode, carries one or more. + samples = metrics["samples"] + row["entries_per_object"] = ratio(samples.get("object_entries_sum"), samples.get("object_entries_count")) + row["appends_per_object"] = ratio(samples.get("append_ack_seconds_count"), samples.get("object_bytes_count")) + row["partitions"] = cell.get("partitions", 1) + row["region_rows"] = region_rows(path / "regions.json") + if cell.get("restart"): + # A restarted process observed exactly one recovery, so the sum is its duration. + count = restarted["samples"].get("recovery_seconds_count") + row["restart"] = { + "wall_seconds": cell.get("restart_seconds"), "query_seconds": cell.get("restart_query_seconds"), + "rows": cell.get("restart_count"), + "recovery_seconds": restarted["samples"].get("recovery_seconds_sum") if count == 1 else None, + "recovered_objects": restarted["samples"].get("recovered_objects_total") + if cell.get("restart_scraped") else None, + "indexed_before_kill": cell.get("indexed_before_kill"), "min_objects": cell.get("min_objects", 0), + "replay": replay(path / "server-restart.log", restarted), + } # The counters register on their first event, so a scrape of an object store WAL # that created objects omits the ones that never moved: those are zero, not unknown. active = (cell.get("metrics_scraped") and cell.get("wal", {}).get("PROVIDER") == "experimental_object_store" @@ -91,26 +232,74 @@ def load_cell(path: Path, spec: dict | None = None) -> dict: row[counter] = metrics["samples"].get(counter, 0 if active else None) created = row["created_objects_total"] row["creates_per_minute"] = created * 60 / load_seconds if created is not None and load_seconds else None + row |= traffic(metrics, {"samples": {}, "operations": {}}, load_seconds, cell.get("wal", {})) + row["restart_operations"] = restarted["operations"] + if cell.get("idle_seconds"): + idle = traffic(final, metrics, cell["idle_seconds"], cell.get("wal", {})) + end = final["samples"].get("created_objects_total") + row["idle"] = {"seconds": cell["idle_seconds"], "operations": idle["operations"], + "creates": None if end is None else end - metrics["samples"].get("created_objects_total", 0)} + if row["mode"] == "drop": + written = read_metrics(path / "metrics-written.prom")["samples"] + restarted_samples = restarted["samples"] + samples = [line.split() for line in (path / "collected.tsv").read_text().splitlines() + if line.strip()] if (path / "collected.tsv").exists() else [] + row["drop"] = { + "acked_batches": sum(len((p.read_text()).split()) for p in sorted((path / "dropped").glob("acked-*.log"))), + "indexed_before_drop": written.get("indexed_objects"), + "restart_seconds": cell.get("restart_seconds"), + "recovered_objects": restarted_samples.get("recovered_objects_total"), + "indexed_at_end": int(float(samples[-1][1])) if samples else None, + "deleted_after_restart": restarted_samples.get("deleted_objects_total"), + "seconds_to_collect": int(samples[-1][0]) - int(samples[0][0]) if samples else None, + } + if row["mode"] == "crash": + row["cycles"] = crash_cycles(path, cell.get("wal", {})) + if (path / "objects.tsv").exists(): + # A sample is time, indexed, deleted, created; a failed scrape is time and "missing". + row["objects_samples"] = [[int(float(line.split()[0]))] + ( + [None] * 3 if line.split()[1:] == ["missing"] else [int(float(v)) for v in line.split()[1:]]) + for line in (path / "objects.tsv").read_text().splitlines() if line.strip()] return row def ratios(rows: list) -> list: - """D/R and E/R per round and their geometric mean over the rounds, per concurrency.""" - speed = {(r["round"], r["concurrency"], r["config"]): r["rows_per_second"] for r in rows} + """Each object store WAL cell over each Raft Engine control per round, and their geometric mean over the rounds, per concurrency.""" + rows = [r for r in rows if r["mode"] == "load"] + speed = {(r["group"], r["round"], r["concurrency"], r["config"]): r["rows_per_second"] + for r in rows} result = [] - for concurrency in sorted({r["concurrency"] for r in rows}): - rounds = sorted({r["round"] for r in rows if r["concurrency"] == concurrency}) - for config in ("D", "E"): - per_round = [] - for round_number in rounds: - base, value = speed.get((round_number, concurrency, "R")), speed.get((round_number, concurrency, config)) - per_round.append(value / base if base and value else None) - complete = [value for value in per_round if value is not None] - result.append({"concurrency": concurrency, "ratio": f"{config}/R", "per_round": per_round, - "geomean": geomean(complete) if len(complete) == len(per_round) else None}) + for group in dict.fromkeys(r["group"] for r in rows): + members = [r for r in rows if r["group"] == group] + series = sorted({r["config"] for r in members if r["config"] not in BASELINES}) + baselines = [b for b in BASELINES if any(r["config"] == b for r in members)] + for concurrency in sorted({r["concurrency"] for r in members}): + rounds = sorted({r["round"] for r in members if r["concurrency"] == concurrency}) + for config, baseline in ((s, b) for s in series for b in baselines): + per_round = [] + for round_number in rounds: + base = speed.get((group, round_number, concurrency, baseline)) + measured = speed.get((group, round_number, concurrency, config)) + per_round.append(measured / base if base and measured else None) + complete = [ratio for ratio in per_round if ratio is not None] + result.append({"group": group, "concurrency": concurrency, "ratio": f"{config}/{baseline}", + "per_round": per_round, + "geomean": geomean(complete) if len(complete) == len(per_round) else None}) return result +def operations(counts: dict) -> str: + return ", ".join(f"{op} {n:,.0f}" for op, n in sorted(counts.items())) or "—" + + +def overrides(env: dict) -> str: + return " ".join(f"{k}={v}" for k, v in sorted(env.items())) or "—" + + +def pair(a, b, digits=0, scale=1.0) -> str: + return f"{fmt(a, digits, scale)} / {fmt(b, digits, scale)}" + + def fmt(value, digits=0, scale=1.0) -> str: return "—" if value is None else f"{value * scale:,.{digits}f}" @@ -119,21 +308,134 @@ def render(run: dict, rows: list, pairs: list) -> str: lines = [f"## Object store WAL against S3 ({'smoke, MinIO' if run.get('smoke') else run.get('instance_type')})", ""] lines += [f"GreptimeDB `{run.get('greptimedb_image')}` ({next((r['version'] for r in rows if r['version']), '—')}); " f"{fmt(run.get('rows'))} rows in {run.get('batch_rows')}-row batches; bucket region {run.get('region')}.", ""] - lines += ["| Cell | Rows/s | Load s | Wall s | Rows | Ack p50 ms | Ack p99 ms | Creates | Creates/min " - "| Stalled | Poisoned | Create failures | Conflicts | WAL on |", - "| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | --- |"] - for r in rows: - cell = f"{r['id']} (not run)" if r["missing"] else r["id"] if r["exit_code"] == 0 \ + lines += ["| Cell | Overrides | Rows/s | Load s | Wall s | Ack p50 / p99 ms | Seal to durable p50 / p99 ms " + "| Object p50 / p99 MiB | Creates | Creates/min | Stalled | Poisoned | Create failures | Conflicts " + "| Server CPUs | Host busy | WAL on |", + "| --- | --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: " + "| ---: | ---: | --- |"] + def label(r): + return f"{r['id']} (not run)" if r["missing"] else r["id"] if r["exit_code"] == 0 \ else f"{r['id']} (failed, exit {r['exit_code']})" + loads = [r for r in rows if r["mode"] == "load"] + for r in loads: + cell = label(r) lines.append(" | ".join([ - f"| {cell}", fmt(r["rows_per_second"]), fmt(r["load_seconds"], 1), fmt(r["load_wall_seconds"], 1), - fmt(r["rows"]), fmt(r["append_ack_p50_seconds"], 1, 1000), fmt(r["append_ack_p99_seconds"], 1, 1000), + f"| {cell}", overrides(r["env"]), fmt(r["rows_per_second"]), fmt(r["load_seconds"], 1), + fmt(r["load_wall_seconds"], 1), pair(r["append_ack_p50"], r["append_ack_p99"], 1, 1000), + pair(r["seal_to_durable_p50"], r["seal_to_durable_p99"], 1, 1000), + pair(r["object_bytes_p50"], r["object_bytes_p99"], 1, 1 / 2**20), *(fmt(r[c]) for c in COUNTERS[:1]), fmt(r["creates_per_minute"], 1), - *(fmt(r[c]) for c in COUNTERS[1:]), f"`{r['wal_device'] or '—'}` |"])) - lines += ["", "| Concurrency | Ratio | Per round | Geometric mean |", "| ---: | --- | --- | ---: |"] + *(fmt(r[c]) for c in COUNTERS[1:]), fmt(r["server_cpu_cores"], 2), + "—" if r["host_cpu_busy"] is None else f"{r['host_cpu_busy']:.0%} of {r['host_cpus'] or '?'}", + f"`{r['wal_device'] or '—'}` |"])) + lines += ["", "| Cell | WAL MB/s | Objects/s | Mean object MiB | Sealed on size | OpenDAL S3 operations during the load " + "| OpenDAL S3 operations of the restart | Idle s | Idle creates | Idle OpenDAL S3 operations |", + "| --- | ---: | ---: | ---: | --- | --- | --- | ---: | ---: | --- |"] + for r in loads: + idle = r.get("idle") or {} + lines.append(" | ".join([ + f"| {label(r)}", fmt(r["wal_mb_per_second"], 1), fmt(r["objects_per_second"], 1), + fmt(r["mean_object_bytes"], 2, 1 / 2**20), + "—" if r["sealed_on_size"] is None else ("yes" if r["sealed_on_size"] else "no"), + operations(r["operations"]), operations(r["restart_operations"]) if "restart" in r else "—", + fmt(idle.get("seconds")), fmt(idle.get("creates")), operations(idle["operations"]) if idle else "—"]) + " |") + crashes = [r for r in rows if r["mode"] == "crash"] + if crashes: + lines += ["", "| Cell | Cycle | Kill after s | Acked | Present | Lost | Lost rows | Extra | In flight " + "| All found | All acked | All lost | Indexed at kill | Restart to /health s | Recovery s " + "| Recovered objects | Pin rows | Verdict |", + "| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: " + "| ---: | ---: | ---: | --- |"] + for r in crashes: + for c in r["cycles"] or [{}]: + verdict = "—" if not c else "FAIL: " + "; ".join(c.get("kill_errors", []) + c.get("verify_errors", [])) \ + if not c.get("passed") else "pass" if c.get("metrics_complete", True) else "pass, scrape missing" + lines.append(" | ".join([ + f"| {label(r)}", fmt(c.get("cycle")), fmt(c.get("kill_after_seconds"), 1), + fmt(c.get("cycle_acked_batches")), fmt(c.get("cycle_present_batches")), + fmt(c.get("cycle_lost_batches")), fmt(c.get("cycle_lost_rows")), fmt(c.get("cycle_extra_batches")), + fmt(c.get("in_flight")), fmt(c.get("total_found_batches")), fmt(c.get("total_acked_batches")), + fmt(c.get("total_lost_batches")), fmt(c.get("indexed_before_kill")), + fmt(c.get("restart_seconds"), 2), fmt(c.get("recovery_seconds"), 3), + fmt(c.get("recovered_objects")), fmt(c.get("pin_rows")), verdict]) + " |") + lines += ["", "| Cell | Cycle | Load s | WAL MB/s | Objects/s | Mean object MiB | Sealed on size " + "| OpenDAL S3 operations until the kill | OpenDAL S3 operations of the restart |", + "| --- | ---: | ---: | ---: | ---: | ---: | --- | --- | --- |"] + for r in crashes: + for c in r["cycles"]: + lines.append(" | ".join([ + f"| {r['id']}", fmt(c["cycle"]), fmt(c.get("load_seconds"), 1), fmt(c["wal_mb_per_second"], 1), + fmt(c["objects_per_second"], 1), fmt(c["mean_object_bytes"], 2, 1 / 2**20), + "—" if c["sealed_on_size"] is None else ("yes" if c["sealed_on_size"] else "no"), + operations(c["operations"]), operations(c["restart_operations"])]) + " |") + pinned = [r for r in crashes if r.get("objects_samples")] + if pinned: + lines += ["", "| Cell | Indexed objects every 30 s | Missing samples | Deleted objects at the end " + "| Created objects at the end |", "| --- | --- | ---: | ---: | ---: |"] + for r in pinned: + samples = r["objects_samples"] + last = next((s for s in reversed(samples) if s[1] is not None), [None] * 4) + lines.append(f"| {r['id']} | {', '.join('missing' if s[1] is None else str(s[1]) for s in samples)} " + f"| {sum(s[1] is None for s in samples)} | {fmt(last[2])} | {fmt(last[3])} |") + details = [r for r in loads if any(r[k] is not None for k in ("client_p50", "entries_per_object", "region_rows"))] + if details: + lines += ["", "| Cell | Client p50 / p99 ms | Entries per object | Appends per object | Regions " + "| Rows per region min / max |", "| --- | ---: | ---: | ---: | ---: | ---: |"] + for r in details: + spread = r["region_rows"] or [] + lines.append(" | ".join([ + f"| {label(r)}", pair(r["client_p50"], r["client_p99"], 1, 1000), fmt(r["entries_per_object"], 1), + fmt(r["appends_per_object"], 2), fmt(len(spread) if spread else r["partitions"]), + pair(min(spread), max(spread)) if spread else "—"]) + " |") + replays = [(r["id"], "—", r["restart"]) for r in loads if r.get("restart", {}).get("replay")] + replays += [(r["id"], fmt(c["cycle"]), c | {"wall_seconds": c.get("restart_seconds")}) + for r in crashes for c in r["cycles"] if c.get("replay")] + if replays: + lines += ["", "| Cell | Cycle | Restart to /health s | Recovery s | Recovered objects | Regions opened " + "| Rows replayed | Object span | Slowest region replay s | Read seconds (reads) " + "| Regions with rows |", + "| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | --- | --- |"] + for cell, cycle, x in replays: + y = x["replay"] + lines.append(" | ".join([ + f"| {cell}", cycle, fmt(x.get("wall_seconds"), 2), fmt(x.get("recovery_seconds"), 3), + fmt(x.get("recovered_objects")), fmt(len(y["regions"])), fmt(y["rows"]), fmt(y["object_span"]), + fmt(y["seconds"], 2), f"{fmt(y['read_seconds'], 2)} ({fmt(y['reads'])})", + replayed(y["regions"])]) + " |") + drops = [r for r in rows if r.get("drop")] + if drops: + lines += ["", "| Cell | Acked batches before the drop | Indexed before the drop | Restart to /health s " + "| Recovered objects | Indexed at the end | Deleted after the restart | Seconds to collect | Verdict |", + "| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | --- |"] + for r in drops: + x = r["drop"] + # Objects the release already collected before the kill leave nothing to show. + verdict = "—" if x["indexed_at_end"] is None else \ + "nothing left at the kill" if (x["recovered_objects"] or 0) <= 3 else \ + "collected" if r["exit_code"] == 0 else "still indexed" + lines.append(" | ".join([ + f"| {label(r)}", fmt(x["acked_batches"]), fmt(x["indexed_before_drop"]), fmt(x["restart_seconds"], 2), + fmt(x["recovered_objects"]), fmt(x["indexed_at_end"]), + fmt(x["deleted_after_restart"]), fmt(x["seconds_to_collect"]), verdict]) + " |") + restarts = [r for r in rows if "restart" in r] + if restarts: + lines += ["", "| Cell | Overrides | Objects created | Indexed at kill | Restart to /health s | Restart to count s " + "| Recovery s | Recovered objects | Rows after restart |", + "| --- | --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: |"] + for r in restarts: + x = r["restart"] + short = x["min_objects"] and (x["indexed_before_kill"] or 0) < x["min_objects"] + lines.append(" | ".join([ + f"| {label(r)}", overrides(r["env"]), fmt(r["created_objects_total"]), + fmt(x["indexed_before_kill"]) + (f" (under {x['min_objects']:,})" if short else ""), + fmt(x["wall_seconds"], 2), + fmt(x["query_seconds"], 2), fmt(x["recovery_seconds"], 3), fmt(x["recovered_objects"]), + f"{fmt(x['rows'])} |"])) + lines += ["", "| Group | Concurrency | Ratio | Per round | Geometric mean |", + "| --- | ---: | --- | --- | ---: |"] for p in pairs: - lines.append(f"| {p['concurrency']} | {p['ratio']} | {', '.join(fmt(v, 3) for v in p['per_round'])} " - f"| {fmt(p['geomean'], 3)} |") + lines.append(f"| {p['group'] or '—'} | {p['concurrency']} | {p['ratio']} " + f"| {', '.join(fmt(v, 3) for v in p['per_round'])} | {fmt(p['geomean'], 3)} |") return "\n".join(lines) + "\n" diff --git a/.github/scripts/oswal-s3-bench.py b/.github/scripts/oswal-s3-bench.py index faad8ace674..9d674400fe2 100644 --- a/.github/scripts/oswal-s3-bench.py +++ b/.github/scripts/oswal-s3-bench.py @@ -29,18 +29,62 @@ import re import sys import time import urllib.request +from datetime import datetime from pathlib import Path SMOKE_ROWS = 50_000 # The o11ybench data profiles; a matching row count generates the canonical dataset. DATA_PROFILES = {5_000: "S", 2_000_000: "D", 10_000_000: "P", 100_000_000: "M"} -# Identical [storage] for every configuration; only [wal] differs. +# Identical [storage] for every configuration; only [wal] differs. R0 is the standalone default. WAL_CONFIGS = { "R": {"PROVIDER": "raft_engine", "SYNC_WRITE": "true", "DIR": "/wal"}, + "R0": {"PROVIDER": "raft_engine", "SYNC_WRITE": "false", "DIR": "/wal"}, "D": {"PROVIDER": "experimental_object_store", "ACK_MODE": "durable"}, "E": {"PROVIDER": "experimental_object_store", "ACK_MODE": "enqueued"}, } -ORDERS = (("R", "D", "E"), ("E", "D", "R")) +BASELINES = ("R", "R0") +# Restart cells seal ten times as often as the default so recovery has thousands of objects. +RESTART_FLUSH_INTERVAL = "10ms" +# The loader's CREATE TABLE for the agent observability workload at the pinned runtime image. +AGENT_TABLE = """CREATE TABLE `agent_observations` ( + `event_time` TIMESTAMP(3) NOT NULL, + `biz_date` DATE NOT NULL, + `trace_id` STRING NOT NULL SKIPPING INDEX, + `session_id` STRING NOT NULL, + `observation_id` STRING NOT NULL SKIPPING INDEX, + `parent_observation_id` STRING NULL SKIPPING INDEX, + `seq_no` BIGINT NOT NULL, + `type` STRING NOT NULL INVERTED INDEX, + `status` STRING NOT NULL INVERTED INDEX, + `tenant` STRING NOT NULL, + `app` STRING NOT NULL, + `environment` STRING NOT NULL, + `task_category` STRING NULL, + `trace_archetype` STRING NULL, + `model` STRING NULL, + `tool_name` STRING NULL, + `input` STRING NULL FULLTEXT INDEX WITH(analyzer = 'English', backend = 'bloom', case_sensitive = 'false', granularity = '10240', false_positive_rate = '0.01'), + `output` STRING NULL FULLTEXT INDEX WITH(analyzer = 'English', backend = 'bloom', case_sensitive = 'false', granularity = '10240', false_positive_rate = '0.01'), + `input_tokens` BIGINT NULL, + `output_tokens` BIGINT NULL, + `total_cost` DOUBLE NULL, + `latency_ms` BIGINT NULL, + `payload` JSON2 NOT NULL, + TIME INDEX (`event_time`), + PRIMARY KEY (`tenant`, `app`, `environment`) +) +{partitions}ENGINE=mito +WITH( + append_mode = 'true', + 'compaction.type' = 'twcs', + 'compaction.twcs.time_window' = '2h' +)""" +GROUP_NAME = re.compile(r"[a-z][a-z0-9]{0,11}") +RESERVED_ENV = ("GREPTIMEDB_STANDALONE__STORAGE__", "GREPTIMEDB_STANDALONE__WAL__") +ENV_NAME = re.compile(r"[A-Za-z_][A-Za-z0-9_]*") +# Docker's reference grammar: path components separated by ".", "_", "__" or runs of "-". +IMAGE = re.compile(r"[a-z0-9]+(?:(?:[._]|__|-+)[a-z0-9]+)*(?::[0-9]+)?(?:/[a-z0-9]+(?:(?:[._]|__|-+)[a-z0-9]+)*)+" + r"(?::[A-Za-z0-9_][A-Za-z0-9_.-]{0,127})?(?:@sha256:[0-9a-f]{64})?") NIGHTLY = re.compile(r"v\d+\.\d+\.\d+(?:-[0-9A-Za-z.]+)?-nightly-(\d{8})") TAG = re.compile(r"[A-Za-z0-9_][A-Za-z0-9_.-]{0,127}") DIGEST = re.compile(r".+@sha256:[0-9a-f]{64}") @@ -78,32 +122,155 @@ def build_manifest(plan: dict, greptimedb_inspect: list, runtime_image: str, raise ValueError("the GreptimeDB image has no sha256 repository digest") return {"smoke": plan["smoke"], "rows": plan["rows"], "rounds": plan["rounds"], "concurrency": plan["concurrency"], "batch_rows": plan["batch_rows"], + "configs": plan["configs"], "extra_env": plan["extra_env"], + "restart": plan["restart"], "groups": plan["groups"], "greptimedb_image": image, "runtime_image": runtime_image, "instance_type": instance_type, "region": region, "prefix": prefix} -def build_plan(concurrency: str, rows: str, batch_rows: str, rounds: str, smoke: bool) -> dict: +def parse_extra_env(text: str) -> dict: + """A JSON object or K=V lines; the storage and WAL settings that define a cell stay the benchmark's own.""" + if text.strip().startswith("{"): + parsed = json.loads(text) + if not all(isinstance(v, (str, int)) and not isinstance(v, bool) for v in parsed.values()): + raise ValueError("extra_env values must be strings or integers") + env = {k: str(v) for k, v in parsed.items()} + else: + env = {} + for line in filter(None, (line.strip() for line in text.splitlines())): + name, sep, value = line.partition("=") + if not sep or name.strip() in env: + raise ValueError(f"extra_env line is not a new K=V: {line!r}") + env[name.strip()] = value.strip() + for name, value in env.items(): + # GreptimeDB matches its environment keys without regard to case. + if not ENV_NAME.fullmatch(name) or name.upper().startswith(RESERVED_ENV) or "\n" in value: + raise ValueError(f"extra_env may not set {name!r}") + return env + + +def partitioned_table(summary: dict, partitions: int) -> str: + """The loader's table split into equal event_time ranges over the corpus window. The + generator draws event times uniformly from that window, so every range gets an equal + share of the rows and every batch writes to every region.""" + def millis(text): + return int(datetime.fromisoformat(text.replace("Z", "+00:00")).timestamp() * 1000) + start, end = millis(summary["start"]), millis(summary["end"]) + bounds = [start + (end - start) * i // partitions for i in range(1, partitions)] + if len(set(bounds)) != len(bounds) or not start < bounds[0]: + raise ValueError(f"the corpus window {summary['start']}..{summary['end']} is too short for {partitions} ranges") + ranges = [f"`event_time` < {bounds[0]}"] + ranges += [f"`event_time` >= {low} AND `event_time` < {high}" for low, high in zip(bounds, bounds[1:])] + ranges.append(f"`event_time` >= {bounds[-1]}") + clause = "PARTITION ON COLUMNS (`event_time`) (\n " + ",\n ".join(ranges) + "\n)\n" + return AGENT_TABLE.format(partitions=clause) + + +def group_cells(group: dict, smoke: bool) -> list: + """The cells of one group, in run order; a group without a name is the top-level inputs.""" + name, mode = group.get("name", ""), group.get("mode", "load") + if mode not in ("load", "crash", "drop"): + raise ValueError(f"group {name!r}: mode must be load, crash or drop, got {mode!r}") + concurrency = str(group["concurrency"]) levels = [positive("concurrency", level.strip(), 256) for level in concurrency.split(",")] if len(set(levels)) != len(levels): raise ValueError(f"concurrency repeats a level: {concurrency!r}") - row_count = SMOKE_ROWS if smoke else positive("rows", rows, 100_000_000) - round_count = 1 if smoke else positive("rounds", rounds, 10) - profile = DATA_PROFILES.get(row_count) + chosen = [config.strip() for config in str(group["configs"]).split(",")] + allowed = set(WAL_CONFIGS) - (set(BASELINES) if mode != "load" else set()) + if not chosen or len(set(chosen)) != len(chosen) or not set(chosen) <= allowed: + raise ValueError(f"group {name!r}: configs must be distinct names among {', '.join(sorted(allowed))}, " + f"got {group['configs']!r}") + extra = group.get("extra_env", "") + env = parse_extra_env(json.dumps(extra) if isinstance(extra, dict) else str(extra)) + restart = group.get("restart", False) is True + rows = SMOKE_ROWS if smoke else positive("rows", str(group["rows"]), 100_000_000) + rounds = 1 if smoke else positive("rounds", str(group["rounds"]), 10) + idle = int(group.get("idle_seconds", 0)) + common = {"group": name, "mode": mode, "restart": restart, "idle_seconds": idle} + partitions = int(group.get("partitions", 1)) + if partitions != 1 and (mode != "load" or not 2 <= partitions <= 64): + raise ValueError(f"group {name!r}: partitions must be in 2..64 and needs mode load") + if mode == "load": + if not 0 <= idle <= 600: + raise ValueError(f"group {name!r}: idle_seconds must be in 0..600") + # The objects a restart cell must leave indexed at its kill; a smoke run is too small for any. + min_objects = int(group.get("min_objects", 0)) + if min_objects < 0 or (min_objects and not restart): + raise ValueError(f"group {name!r}: min_objects needs restart and must not be negative") + common |= {"rows": rows, "data_profile": DATA_PROFILES.get(rows, "custom"), + "min_objects": 0 if smoke else min_objects, "partitions": partitions} + elif mode == "drop": + # Writes one table for write_seconds, drops it, kills and restarts the server, then + # writes and flushes another table and watches the dropped table's objects go. + write_seconds = int(group.get("write_seconds", 30)) + if not 1 <= write_seconds <= 600 or restart or idle: + raise ValueError(f"group {name!r}: write_seconds must be in 1..600, without restart or idle") + common |= {"write_seconds": min(write_seconds, 10) if smoke else write_seconds} + else: + cycles = int(group.get("cycles", 5)) + kill_after = [int(v) for v in group.get("kill_after", [3, 10])] + if not 1 <= cycles <= 20 or len(kill_after) != 2 or not 1 <= kill_after[0] <= kill_after[1] <= 3600: + raise ValueError(f"group {name!r}: cycles must be in 1..20 and kill_after a [min, max] in 1..3600") + if restart or idle: + raise ValueError(f"group {name!r}: a crash group restarts after every kill and has no idle window") + if smoke: + cycles, kill_after = min(cycles, 2), [min(v, 60) for v in kill_after] + common |= {"cycles": cycles, "kill_after": kill_after, "pin": group.get("pin", False) is True} + series = [config for config in WAL_CONFIGS if config in chosen] cells = [] - for round_number in range(1, round_count + 1): - # Alternate the order between rounds so no configuration always runs first. + for round_number in range(1, rounds + 1): + # Reverse the order every other round so no configuration always runs first. for level in levels: - for config in ORDERS[(round_number - 1) % 2]: - cells.append({ - "id": f"{round_number}-{level}-{config}", "round": round_number, - "concurrency": level, "config": config, "wal": WAL_CONFIGS[config], + for config in series if round_number % 2 else reversed(series): + wal = WAL_CONFIGS[config] | ({"FLUSH_INTERVAL": RESTART_FLUSH_INTERVAL} if restart and config not in BASELINES else {}) + cells.append(common | { + "id": (f"{name}-" if name else "") + f"{round_number}-{level}-{config}", + "round": round_number, "concurrency": level, "config": config, "wal": wal, "env": env, }) + return cells + + +def parse_groups(text: str, defaults: dict) -> list: + """A JSON list of groups; every key a group leaves out takes the top-level input's value.""" + if not text.strip(): + return [defaults] + groups = json.loads(text) + if not isinstance(groups, list) or not groups or not all(isinstance(g, dict) for g in groups): + raise ValueError("groups must be a non-empty JSON list of objects") + known = set(defaults) | {"name", "mode", "idle_seconds", "cycles", "kill_after", "pin", "min_objects", + "partitions", "write_seconds"} + names = [g.get("name") for g in groups] + for group, name in zip(groups, names): + if not isinstance(name, str) or not GROUP_NAME.fullmatch(name): + raise ValueError(f"every group needs a name of up to 12 lowercase letters and digits, got {name!r}") + if set(group) - known: + raise ValueError(f"group {name!r} has unknown keys {sorted(set(group) - known)}") + if len(set(names)) != len(names): + raise ValueError(f"group names repeat: {names}") + return [defaults | group for group in groups] + + +def build_plan(concurrency: str, rows: str, batch_rows: str, rounds: str, smoke: bool, + configs: str = "R,D,E", extra_env: str = "", restart: bool = False, groups: str = "") -> dict: + defaults = {"concurrency": concurrency, "rows": rows, "rounds": rounds, "configs": configs, + "extra_env": extra_env, "restart": restart} + resolved = parse_groups(groups, defaults) + cells = [cell for group in resolved for cell in group_cells(group, smoke)] + if len({cell["id"] for cell in cells}) != len(cells): + raise ValueError("two cells share an id") + corpora = [] + for count in dict.fromkeys(cell["rows"] for cell in cells if cell["mode"] == "load"): + profile = DATA_PROFILES.get(count) + corpora.append({"rows": count, "dir": f"corpus-{count}", "data_profile": profile or "custom", + "generator_args": ["--data-profile", profile] if profile else ["--line-count", str(count)]}) + first = cells[0] return { - "smoke": smoke, "rows": row_count, "rounds": round_count, - "batch_rows": positive("batch_rows", batch_rows, 100_000), "concurrency": levels, - "data_profile": profile or "custom", - "generator_args": ["--data-profile", profile] if profile else ["--line-count", str(row_count)], - "cells": cells, + "smoke": smoke, "rows": first.get("rows"), "rounds": max(cell["round"] for cell in cells), + "batch_rows": positive("batch_rows", batch_rows, 100_000), + "concurrency": sorted({cell["concurrency"] for cell in cells}), + "configs": list(dict.fromkeys(cell["config"] for cell in cells)), + "extra_env": parse_extra_env(extra_env), "groups": [g.get("name", "") for g in resolved], + "restart": restart, "corpora": corpora, "cells": cells, } @@ -114,6 +281,15 @@ def newest_nightly(tags: list[str]) -> str: return max(dated)[1] +def resolve_image(image: str, tag: str) -> str: + """The full image reference when given, else greptime/greptimedb at the resolved tag.""" + if image: + if not IMAGE.fullmatch(image) or not re.search(r"(:[^/]+|@sha256:[0-9a-f]{64})$", image): + raise ValueError(f"greptimedb_image must be a full reference with a tag or digest: {image!r}") + return image + return f"greptime/greptimedb:{resolve_tag(tag)}" + + def resolve_tag(requested: str) -> str: if requested != "nightly": if not TAG.fullmatch(requested): @@ -169,12 +345,21 @@ def main() -> int: for name in ("concurrency", "rows", "batch-rows", "rounds"): plan.add_argument(f"--{name}", required=True) plan.add_argument("--smoke", choices=("true", "false"), required=True) + plan.add_argument("--restart", choices=("true", "false"), default="false") + plan.add_argument("--configs", default="R,D,E") + for name in ("extra-env", "groups"): + plan.add_argument(f"--{name}", default="") for name in ("timeout-minutes", "ttl-hours", "output"): plan.add_argument(f"--{name}") - commands.add_parser("tag").add_argument("requested") + image = commands.add_parser("image") + image.add_argument("--tag", required=True) + image.add_argument("--image", default="") regions = commands.add_parser("regions") for name in ("bucket", "runner"): regions.add_argument(f"--{name}", required=True) + table = commands.add_parser("table") + table.add_argument("--summary", type=Path, required=True) + table.add_argument("--partitions", type=int, required=True) manifest = commands.add_parser("manifest") for name in ("results", "runtime-image", "instance-type", "region", "prefix"): manifest.add_argument(f"--{name}", required=True) @@ -187,18 +372,21 @@ def main() -> int: if args.command == "plan": if args.timeout_minutes or args.ttl_hours: validate_budgets(args.timeout_minutes or "", args.ttl_hours or "") - result = build_plan(args.concurrency, args.rows, args.batch_rows, args.rounds, args.smoke == "true") + result = build_plan(args.concurrency, args.rows, args.batch_rows, args.rounds, args.smoke == "true", + args.configs, args.extra_env, args.restart == "true", args.groups) if args.output: with open(args.output, "w") as output: json.dump(result, output, indent=2) - print(f"{len(result['cells'])} cells, {result['rows']} rows, profile {result['data_profile']}") - elif args.command == "tag": - tag = resolve_tag(args.requested) - append_github("GITHUB_OUTPUT", "greptimedb_tag", tag) - print(f"GreptimeDB image tag: {tag}") + print(f"{len(result['cells'])} cells; corpora {[c['rows'] for c in result['corpora']]}") + elif args.command == "image": + reference = resolve_image(args.image, args.tag) + append_github("GITHUB_OUTPUT", "greptimedb_ref", reference) + print(f"GreptimeDB image: {reference}") elif args.command == "regions": check_regions(args.bucket, args.runner) print(f"bucket and runner both in {args.runner}") + elif args.command == "table": + print(partitioned_table(json.loads(args.summary.read_text()), args.partitions)) elif args.command == "manifest": results = Path(args.results) manifest = build_manifest(json.loads((results / "plan.json").read_text()), diff --git a/.github/workflows/oswal-s3-bench.yml b/.github/workflows/oswal-s3-bench.yml index c31dfe8f1c3..f5e92d2cad5 100644 --- a/.github/workflows/oswal-s3-bench.yml +++ b/.github/workflows/oswal-s3-bench.yml @@ -10,6 +10,26 @@ on: description: 'Docker Hub greptime/greptimedb tag; nightly = the newest v*-nightly-* tag (latest is a release without the object store WAL)' type: string default: nightly + greptimedb_image: + description: 'Full image reference with a tag or digest; used instead of greptimedb_tag when set' + type: string + default: '' + configs: + description: 'WAL configurations: R (Raft Engine, sync_write), R0 (Raft Engine, standalone default), D (durable), E (enqueued)' + type: string + default: R,D,E + extra_env: + description: 'Environment for every GreptimeDB container: a JSON object or K=V lines' + type: string + default: '' + restart: + description: 'After each load, SIGKILL the server, restart it on the same prefix and data home, and time recovery; object store cells flush every 10ms' + type: boolean + default: false + groups: + description: 'JSON list of cell groups, each overriding configs, extra_env, concurrency, rounds, rows, restart, plus idle_seconds and partitions (regions of the loaded table), or mode crash with cycles, kill_after [min, max] and pin, or mode drop with write_seconds; empty = one group from the inputs above' + type: string + default: '' instance_type: description: EC2 instance type in the bucket region; the d of c6id gives a local NVMe instance store for the Raft Engine WAL type: string @@ -47,6 +67,11 @@ permissions: env: SMOKE: ${{ inputs.smoke }} GREPTIMEDB_TAG: ${{ inputs.greptimedb_tag }} + GREPTIMEDB_IMAGE_INPUT: ${{ inputs.greptimedb_image }} + CONFIGS: ${{ inputs.configs }} + EXTRA_ENV: ${{ inputs.extra_env }} + RESTART: ${{ inputs.restart }} + OSWAL_GROUPS: ${{ inputs.groups }} REQUESTED_INSTANCE_TYPE: ${{ inputs.instance_type }} CONCURRENCY: ${{ inputs.concurrency }} ROWS: ${{ inputs.rows }} @@ -61,7 +86,7 @@ jobs: timeout-minutes: 15 outputs: instance_type: ${{ steps.runner.outputs.instance_type }} - greptimedb_tag: ${{ steps.plan.outputs.greptimedb_tag }} + greptimedb_ref: ${{ steps.plan.outputs.greptimedb_ref }} runtime_image: ${{ steps.runtime.outputs.runtime_image }} steps: - uses: actions/checkout@v4 @@ -77,8 +102,9 @@ jobs: run: | python3 greptimedb/.github/scripts/oswal-s3-bench.py plan --smoke "$SMOKE" \ --concurrency "$CONCURRENCY" --rows "$ROWS" --batch-rows "$BATCH_ROWS" --rounds "$ROUNDS" \ + --configs "$CONFIGS" --extra-env "$EXTRA_ENV" --restart "$RESTART" --groups "$OSWAL_GROUPS" \ --timeout-minutes "$BENCHMARK_TIMEOUT_MINUTES" --ttl-hours "$JANITOR_TTL_HOURS" - python3 greptimedb/.github/scripts/oswal-s3-bench.py tag "$GREPTIMEDB_TAG" + python3 greptimedb/.github/scripts/oswal-s3-bench.py image --tag "$GREPTIMEDB_TAG" --image "$GREPTIMEDB_IMAGE_INPUT" - name: Resolve CI runner id: runner run: python3 greptimedb/.github/scripts/ci-runner-config.py --provider AWS --instance-type "$REQUESTED_INSTANCE_TYPE" @@ -132,7 +158,7 @@ jobs: timeout-minutes: ${{ fromJSON(format('{0}', inputs.benchmark_timeout_minutes)) }} env: INSTANCE_TYPE: ${{ inputs.smoke && 'ubuntu-latest' || needs.validate.outputs.instance_type }} - GREPTIMEDB_TAG: ${{ needs.validate.outputs.greptimedb_tag }} + GREPTIMEDB_REF: ${{ needs.validate.outputs.greptimedb_ref }} RUNTIME_IMAGE: ${{ needs.validate.outputs.runtime_image }} RUN_PREFIX: oswal-bench/${{ github.run_id }}-${{ github.run_attempt }} OWNER: o11ybench-oswal-${{ github.run_id }}-${{ github.run_attempt }} @@ -161,7 +187,8 @@ jobs: run: | bash greptimedb/.github/scripts/oswal-s3-bench-run.sh pull python3 greptimedb/.github/scripts/oswal-s3-bench.py plan --smoke "$SMOKE" --output "$PLAN" \ - --concurrency "$CONCURRENCY" --rows "$ROWS" --batch-rows "$BATCH_ROWS" --rounds "$ROUNDS" + --concurrency "$CONCURRENCY" --rows "$ROWS" --batch-rows "$BATCH_ROWS" --rounds "$ROUNDS" \ + --configs "$CONFIGS" --extra-env "$EXTRA_ENV" --restart "$RESTART" --groups "$OSWAL_GROUPS" python3 greptimedb/.github/scripts/oswal-s3-bench.py manifest --results "$RESULTS" \ --runtime-image "$RUNTIME_IMAGE" --instance-type "$INSTANCE_TYPE" --region "$S3_REGION" --prefix "$RUN_PREFIX" - name: Probe the bucket diff --git a/tests/perf/test_oswal_s3_bench_scripts.py b/tests/perf/test_oswal_s3_bench_scripts.py index dc81240f1ae..6934ff0db13 100644 --- a/tests/perf/test_oswal_s3_bench_scripts.py +++ b/tests/perf/test_oswal_s3_bench_scripts.py @@ -20,9 +20,12 @@ import json import os import subprocess import sys +import re import tempfile import unittest +import urllib.parse from pathlib import Path +from types import SimpleNamespace SCRIPTS = Path(__file__).parents[2] / ".github/scripts" @@ -35,6 +38,7 @@ def load(name, filename): bench = load("oswal_s3_bench", "oswal-s3-bench.py") +crash = load("oswal_s3_bench_crash", "oswal-s3-bench-crash.py") summary = load("oswal_s3_bench_summary", "oswal-s3-bench-summary.py") ACK_BUCKETS = """# TYPE greptime_logstore_object_store_wal_append_ack_seconds histogram @@ -57,8 +61,9 @@ class PlanTest(unittest.TestCase): self.assertEqual([c["id"] for c in plan["cells"]], [ "1-1-R", "1-1-D", "1-1-E", "1-8-R", "1-8-D", "1-8-E", "2-1-E", "2-1-D", "2-1-R", "2-8-E", "2-8-D", "2-8-R"]) - self.assertEqual(plan["generator_args"], ["--data-profile", "D"]) - self.assertEqual((plan["rows"], plan["batch_rows"], plan["data_profile"]), (2_000_000, 1000, "D")) + self.assertEqual(plan["corpora"], [{"rows": 2_000_000, "dir": "corpus-2000000", "data_profile": "D", + "generator_args": ["--data-profile", "D"]}]) + self.assertEqual((plan["rows"], plan["batch_rows"]), (2_000_000, 1000)) self.assertEqual({c["config"]: c["wal"] for c in plan["cells"]}, { "R": {"PROVIDER": "raft_engine", "SYNC_WRITE": "true", "DIR": "/wal"}, "D": {"PROVIDER": "experimental_object_store", "ACK_MODE": "durable"}, @@ -66,8 +71,8 @@ class PlanTest(unittest.TestCase): def test_smoke_overrides_rows_and_rounds(self): plan = bench.build_plan("8", "2000000", "1000", "2", smoke=True) - self.assertEqual((plan["rows"], plan["rounds"], plan["data_profile"]), (50_000, 1, "custom")) - self.assertEqual(plan["generator_args"], ["--line-count", "50000"]) + self.assertEqual((plan["rows"], plan["rounds"], plan["corpora"][0]["data_profile"]), (50_000, 1, "custom")) + self.assertEqual(plan["corpora"][0]["generator_args"], ["--line-count", "50000"]) self.assertEqual([c["id"] for c in plan["cells"]], ["1-8-R", "1-8-D", "1-8-E"]) def test_invalid_inputs_are_rejected(self): @@ -76,6 +81,129 @@ class PlanTest(unittest.TestCase): with self.subTest(args=args), self.assertRaises(ValueError): bench.build_plan(*args, smoke=False) + def test_the_baselines_and_the_order_reverse_every_other_round(self): + plan = bench.build_plan("8", "2000000", "1000", "2", smoke=False, configs="R,R0,D,E") + self.assertEqual([c["id"] for c in plan["cells"]], [ + "1-8-R", "1-8-R0", "1-8-D", "1-8-E", "2-8-E", "2-8-D", "2-8-R0", "2-8-R"]) + cells = {c["id"]: c for c in plan["cells"]} + self.assertEqual((cells["1-8-R0"]["env"], cells["1-8-R0"]["wal"]), + ({}, {"PROVIDER": "raft_engine", "SYNC_WRITE": "false", "DIR": "/wal"})) + + def test_extra_env_and_restart_reach_every_cell(self): + for extra in ('{"RUST_BACKTRACE": 1, "RUST_LOG": "info"}', "RUST_BACKTRACE=1\n\nRUST_LOG = info\n"): + plan = bench.build_plan("8", "2000000", "1000", "1", smoke=False, configs="R,D", extra_env=extra, + restart=True) + with self.subTest(extra=extra): + self.assertEqual([c["id"] for c in plan["cells"]], ["1-8-R", "1-8-D"]) + cells = {c["id"]: c for c in plan["cells"]} + self.assertEqual(cells["1-8-D"]["env"], {"RUST_BACKTRACE": "1", "RUST_LOG": "info"}) + self.assertEqual(cells["1-8-D"]["wal"]["FLUSH_INTERVAL"], "10ms") + self.assertNotIn("FLUSH_INTERVAL", cells["1-8-R"]["wal"]) + self.assertTrue(all(c["restart"] for c in plan["cells"])) + + def test_invalid_extra_env_and_configs_are_rejected(self): + for kwargs in ({"extra_env": "GREPTIMEDB_STANDALONE__STORAGE__BUCKET=other"}, {"extra_env": "NO_VALUE"}, + {"extra_env": "GREPTIMEDB_STANDALONE__WAL__FLUSH_INTERVAL=100ms", "restart": True}, + {"extra_env": "GREPTIMEDB_STANDALONE__WAL__SYNC_WRITE=false"}, + {"extra_env": "GREPTIMEDB_STANDALONE__wal__FLUSH_INTERVAL=100ms", "restart": True}, + {"extra_env": "greptimedb_standalone__Storage__bucket=other"}, + {"extra_env": "A=1\nA=2"}, {"extra_env": '{"A": true}'}, {"extra_env": "1A=x"}, + {"configs": "R,X"}, {"configs": "D,D"}): + with self.subTest(kwargs=kwargs), self.assertRaises(ValueError): + bench.build_plan("8", "6000", "1000", "1", smoke=False, **kwargs) + + def test_the_image_is_the_full_reference_or_a_greptimedb_tag(self): + image = "docker.io/greptime/greptimedb-dev:dev-20261013-1791848986-a5a1a9f4" + for valid in (image, "ghcr.io/acme/greptimedb--bench:dev", "ghcr.io/acme/greptimedb__bench:dev", + "localhost:5000/a/b@sha256:" + "c" * 64): + self.assertEqual(bench.resolve_image(valid, "nightly"), valid) + self.assertEqual(bench.resolve_image("", "v1.4.0-nightly-20261012"), "greptime/greptimedb:v1.4.0-nightly-20261012") + for bad in ("greptimedb", "docker.io/greptime/greptimedb", "Docker.io/x/y:1", "x/y:bad tag", + "x/y___z:1", "x/-y:1"): + with self.subTest(bad=bad), self.assertRaises(ValueError): + bench.resolve_image(bad, "nightly") + + def test_groups_inherit_the_inputs_and_keep_their_own_settings(self): + groups = [{"name": "recovery", "configs": "D,E", "rows": 4000000, "rounds": 1, "restart": True}, + {"name": "crashe", "mode": "crash", "configs": "E", "rounds": 1, "cycles": 3, "kill_after": [5, 9]}, + {"name": "pin", "mode": "crash", "configs": "D", "rounds": 1, "cycles": 1, + "kill_after": [540, 540], "pin": True}, + {"name": "idle", "configs": "D", "rounds": 1, "idle_seconds": 180}] + plan = bench.build_plan("8", "2000000", "1000", "2", smoke=False, groups=json.dumps(groups)) + self.assertEqual([c["id"] for c in plan["cells"]], [ + "recovery-1-8-D", "recovery-1-8-E", "crashe-1-8-E", "pin-1-8-D", "idle-1-8-D"]) + self.assertEqual([c["rows"] for c in plan["corpora"]], [4_000_000, 2_000_000]) + cells = {c["id"]: c for c in plan["cells"]} + self.assertEqual(cells["recovery-1-8-D"]["wal"]["FLUSH_INTERVAL"], "10ms") + self.assertEqual((cells["crashe-1-8-E"]["cycles"], cells["crashe-1-8-E"]["kill_after"]), (3, [5, 9])) + self.assertTrue(cells["pin-1-8-D"]["pin"]) + self.assertNotIn("FLUSH_INTERVAL", cells["pin-1-8-D"]["wal"]) + self.assertEqual(cells["idle-1-8-D"]["idle_seconds"], 180) + smoke = bench.build_plan("8", "2000000", "1000", "2", smoke=True, groups=json.dumps(groups)) + pin = next(c for c in smoke["cells"] if c["group"] == "pin") + self.assertEqual((pin["cycles"], pin["kill_after"]), (1, [60, 60])) + self.assertEqual([c["rows"] for c in smoke["corpora"]], [50_000]) + recovery = [{"name": "rec", "configs": "D", "restart": True, "min_objects": 6000}] + self.assertEqual(bench.build_plan("8", "4000000", "1000", "1", smoke=False, + groups=json.dumps(recovery))["cells"][0]["min_objects"], 6000) + self.assertEqual(bench.build_plan("8", "4000000", "1000", "1", smoke=True, + groups=json.dumps(recovery))["cells"][0]["min_objects"], 0) + + def test_the_workflow_passes_groups_through_bash_unchanged(self): + workflow = (SCRIPTS.parent / "workflows" / "oswal-s3-bench.yml").read_text() + self.assertIn("OSWAL_GROUPS: ${{ inputs.groups }}", workflow) + self.assertEqual(workflow.count('--groups "$OSWAL_GROUPS"'), 2) + command = (f'python3 "{SCRIPTS}/oswal-s3-bench.py" plan --smoke false --concurrency 8 --rows 6000 ' + '--batch-rows 1000 --rounds 1 --groups "$OSWAL_GROUPS"') + for groups, expected in (("", "3 cells"), ('[{"name": "a", "configs": "D,E"}]', "2 cells")): + done = subprocess.run(["bash", "-c", command], env=os.environ | {"OSWAL_GROUPS": groups}, + capture_output=True, text=True, timeout=60) + self.assertEqual(done.returncode, 0, done.stderr) + self.assertTrue(done.stdout.startswith(expected), done.stdout) + + def test_a_partitioned_group_and_a_drop_group(self): + groups = [{"name": "multi", "configs": "R,D", "concurrency": "8,32", "partitions": 8}, + {"name": "drop", "mode": "drop", "configs": "D", "rounds": 1, "write_seconds": 45}] + plan = bench.build_plan("8", "2000000", "1000", "1", smoke=False, groups=json.dumps(groups)) + cells = {c["id"]: c for c in plan["cells"]} + self.assertEqual(list(cells), ["multi-1-8-R", "multi-1-8-D", "multi-1-32-R", "multi-1-32-D", "drop-1-8-D"]) + self.assertEqual({c["partitions"] for c in plan["cells"] if c["mode"] == "load"}, {8}) + self.assertEqual((cells["drop-1-8-D"]["write_seconds"], cells["drop-1-8-D"]["mode"]), (45, "drop")) + self.assertNotIn("rows", cells["drop-1-8-D"]) + self.assertEqual([c["rows"] for c in plan["corpora"]], [2_000_000]) + smoke = bench.build_plan("8", "2000000", "1000", "1", smoke=True, groups=json.dumps(groups)) + self.assertEqual(next(c for c in smoke["cells"] if c["mode"] == "drop")["write_seconds"], 10) + self.assertEqual(bench.build_plan("8", "6000", "1000", "1", smoke=False)["cells"][0]["partitions"], 1) + + def test_the_partitioned_table_splits_the_corpus_window_into_equal_ranges(self): + ddl = bench.partitioned_table({"start": "2026-07-10T00:00:00Z", "end": "2026-07-10T01:00:00Z"}, 8) + start = 1783641600000 + bounds = [int(b) for b in re.findall(r"`event_time` < (\d+)", ddl)] + self.assertEqual(bounds, [start + 450_000 * i for i in range(1, 8)]) + self.assertIn(f"`event_time` >= {start + 450_000 * 7}\n)\nENGINE=mito", ddl) + self.assertEqual(ddl.count("`event_time` >="), 7) + self.assertTrue(ddl.startswith("CREATE TABLE `agent_observations` (")) + self.assertIn("PRIMARY KEY (`tenant`, `app`, `environment`)\n)\nPARTITION ON COLUMNS (`event_time`) (", ddl) + self.assertNotIn("{partitions}", ddl) + with self.assertRaises(ValueError): + bench.partitioned_table({"start": "2026-07-10T00:00:00Z", "end": "2026-07-10T00:00:00.004Z"}, 8) + + def test_invalid_groups_are_rejected(self): + for groups in ([], [{"configs": "D"}], [{"name": "Bad"}], [{"name": "a"}, {"name": "a"}], + [{"name": "a", "colour": "red"}], [{"name": "a", "mode": "crash", "configs": "R"}], + [{"name": "a", "mode": "crash", "configs": "D", "restart": True}], + [{"name": "a", "mode": "crash", "configs": "D", "kill_after": [9, 5]}], + [{"name": "a", "mode": "soak"}], [{"name": "a", "idle_seconds": 601}], + [{"name": "a", "configs": "D", "min_objects": 6000}], + [{"name": "a", "configs": "D", "restart": True, "min_objects": -1}], + [{"name": "a", "configs": "D", "partitions": 65}], [{"name": "a", "configs": "D", "partitions": 0}], + [{"name": "a", "mode": "crash", "configs": "D", "partitions": 8}], + [{"name": "a", "mode": "drop", "configs": "R"}], + [{"name": "a", "mode": "drop", "configs": "D", "write_seconds": 0}], + [{"name": "a", "mode": "drop", "configs": "D", "restart": True}]): + with self.subTest(groups=groups), self.assertRaises(ValueError): + bench.build_plan("8", "6000", "1000", "1", smoke=False, groups=json.dumps(groups)) + def test_nightly_tag_is_the_newest_by_date(self): tags = ["latest", "v1.3.0-alpha.1-nightly-20260907", "v1.4.0-nightly-20261012", "v1.4.0-nightly-20261005", "nightly-arm64"] @@ -107,6 +235,185 @@ class PlanTest(unittest.TestCase): bench.build_manifest(plan, [{"RepoDigests": []}], "", "", "", "") +def write_lines(path, lines): + path.write_text("".join(f"{line}\n" for line in lines)) + + +def crash_cycle(root, cycle, stripes, acked, requests, exits, found=None): + """One cycle of two writers: attempted stripes, acked batches, requests and exit times per + writer, and what the restart after its kill found, if given.""" + directory = root / f"cycle-{cycle}" + directory.mkdir(parents=True) + for index in range(2): + write_lines(directory / f"attempted-{index}.log", stripes[index]) + write_lines(directory / f"acked-{index}.log", acked[index]) + write_lines(directory / f"requests-{index}.log", requests[index]) + (directory / f"exit-{index}.log").write_text(f"{exits[index]}\n") + if found is not None: + (directory / "found.json").write_text(json.dumps( + {"output": [{"records": {"rows": [[seq, count] for seq, count in found.items()]}}]})) + return directory + + +class CrashOracleTest(unittest.TestCase): + # Writer 0 got batches 0 and 2 acknowledged and lost 4 in flight; writer 1 got 1 and lost 3. + STRIPES = ([0, 2, 4], [1, 3]) + ACKED = ([0, 2], [1]) + REQUESTS = (["0 200 10", "2 200 20", "4 000 105"], ["1 200 15", "3 000 101"]) + + def check(self, root, found, mode="durable", requests=REQUESTS, exits=(106, 102)): + crash_cycle(root, 1, self.STRIPES, self.ACKED, requests, exits) + result = crash.check_kill(root / "cycle-1", 2, 100, 104, 2000) + return result | crash.verify(root, found, 10, mode == "enqueued", 1) + + def test_acked_batches_present_once_and_an_in_flight_extra_pass(self): + with tempfile.TemporaryDirectory() as directory: + result = self.check(Path(directory), {0: 10, 1: 10, 2: 10, 4: 10}) + self.assertEqual((result["kill_errors"], result["verify_errors"]), ([], [])) + self.assertEqual((result["in_flight"], result["writers_at_kill"], result["cycle_extra_batches"], + result["cycle_lost_batches"], result["cycle_present_batches"]), (2, 2, 1, 0, 4)) + + def test_durable_violations_are_reported(self): + cases = {"acknowledged but absent": {0: 10, 1: 10}, "present more than once": {0: 20, 1: 10, 2: 10}, + "present in part": {0: 10, 1: 10, 2: 7}, "never attempted": {0: 10, 1: 10, 2: 10, 9: 10}} + for error, found in cases.items(): + with self.subTest(error=error), tempfile.TemporaryDirectory() as directory: + result = self.check(Path(directory), found) + self.assertTrue(any(error in e for e in result["verify_errors"]), result["verify_errors"]) + + def test_enqueued_loses_acked_batches_without_failing(self): + with tempfile.TemporaryDirectory() as directory: + result = self.check(Path(directory), {0: 10}, mode="enqueued") + self.assertEqual(result["verify_errors"], []) + self.assertEqual((result["cycle_lost_batches"], result["cycle_lost_rows"], result["total_lost_batches"]), + (2, 20, 2)) + + def test_a_later_cycle_reports_only_its_own_losses(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + crash_cycle(root, 1, self.STRIPES, self.ACKED, self.REQUESTS, (106, 102)) + crash_cycle(root, 2, ([10, 12], [11]), ([10], [11]), (["10 200 210", "12 000 305"], ["11 200 215"]), + (306, 302)) + # Cycle 1 lost batches 1 and 2; cycle 2 lost nothing. + result = crash.verify(root, {0: 10, 10: 10, 11: 10}, 10, True, 2) + self.assertEqual((result["cycle_acked_batches"], result["cycle_lost_batches"], + result["total_lost_batches"], result["total_acked_batches"]), (2, 0, 2, 5)) + + def later_cycle(self, root): + """Cycle 1's restart recovered batches 0, 1 and 2 and the in-flight 4; cycle 2 acknowledged + 10 and 11 and had 12 in flight.""" + crash_cycle(root, 1, self.STRIPES, self.ACKED, self.REQUESTS, (106, 102), found={0: 10, 1: 10, 2: 10, 4: 10}) + crash_cycle(root, 2, ([10, 12], [11]), ([10], [11]), (["10 200 210", "12 000 305"], ["11 200 215"]), + (306, 302)) + + def test_a_batch_an_earlier_restart_recovered_must_stay_present(self): + for mode in ("enqueued", "durable"): + with self.subTest(mode=mode), tempfile.TemporaryDirectory() as directory: + root = Path(directory) + self.later_cycle(root) + # Acknowledged batch 1 and the in-flight 4 vanished after cycle 1's restart found them. + result = crash.verify(root, {0: 10, 2: 10, 10: 10, 11: 10}, 10, mode == "enqueued", 2) + self.assertIn("recovered by an earlier restart but absent: [1, 4]", result["verify_errors"]) + self.assertEqual(result["cycle_lost_batches"], 0) + + def test_only_the_current_cycle_may_lose_acknowledged_batches_in_enqueued_mode(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + self.later_cycle(root) + found = {0: 10, 1: 10, 2: 10, 4: 10, 10: 10} + result = crash.verify(root, found, 10, True, 2) + self.assertEqual(result["verify_errors"], []) + self.assertEqual((result["cycle_lost_batches"], result["total_lost_batches"]), (1, 1)) + durable = crash.verify(root, found, 10, False, 2) + self.assertEqual(durable["verify_errors"], ["acknowledged but absent: [11]"]) + duplicated = crash.verify(root, found | {10: 20}, 10, True, 2) + self.assertEqual(duplicated["verify_errors"], ["batches present more than once: [10]"]) + + def test_a_connection_lost_before_the_kill_fails(self): + with tempfile.TemporaryDirectory() as directory: + requests = (["0 200 10", "2 200 20", "4 000 50"], ["1 200 15", "3 500 101"]) + result = self.check(Path(directory), {0: 10, 1: 10, 2: 10}, requests=requests, exits=(51, 99)) + self.assertEqual(len(result["kill_errors"]), 3, result["kill_errors"]) + + def test_the_writer_stops_each_stripe_at_its_first_unacknowledged_batch(self): + import http.server + import threading + inserted, lock = [], threading.Lock() + + class Handler(http.server.BaseHTTPRequestHandler): + def do_POST(self): + body = self.rfile.read(int(self.headers["Content-Length"])).decode() + with lock: + if len(inserted) >= 6: + self.close_connection = True + self.connection.close() + return + inserted.append(body) + self.send_response(200) + self.end_headers() + + def log_message(self, *args): + pass + + server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), Handler) + threading.Thread(target=server.serve_forever, daemon=True).start() + try: + with tempfile.TemporaryDirectory() as directory: + run = Path(directory) + args = SimpleNamespace(dir=str(run / "cycle-1"), url=f"http://127.0.0.1:{server.server_port}/v1/sql", + table="crash", writers=2, first_seq=0, rows=3, seconds=0) + crash.write(args) + found = {} + for body in inserted: + for seq in set(re.findall(r"%28(\d+)%2C", body)): + found[int(seq)] = 3 + result = crash.verify(run, found, 3, False, 1) + self.assertEqual((result["verify_errors"], result["total_acked_batches"]), ([], 6)) + for index in range(2): + attempted = (run / "cycle-1" / f"attempted-{index}.log").read_text().split() + self.assertEqual([int(seq) % 2 for seq in attempted], [index] * len(attempted)) + finally: + server.shutdown() + server.server_close() + + +class WriterDeadlineTest(unittest.TestCase): + def test_writers_stop_at_the_deadline_with_every_batch_acknowledged(self): + import http.server + import threading + status = [200] + + class Handler(http.server.BaseHTTPRequestHandler): + def do_POST(self): + self.rfile.read(int(self.headers["Content-Length"])) + self.send_response(status[0]) + self.end_headers() + + def log_message(self, *args): + pass + + server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), Handler) + threading.Thread(target=server.serve_forever, daemon=True).start() + try: + with tempfile.TemporaryDirectory() as directory: + run = Path(directory) + args = SimpleNamespace(dir=str(run / "cycle-1"), url=f"http://127.0.0.1:{server.server_port}/", + table="dropped", writers=2, first_seq=0, rows=2, seconds=0.3) + self.assertEqual(crash.write(args), 0) + result = crash.verify(run, {}, 2, True, 1) + self.assertGreater(result["cycle_acked_batches"], 0) + self.assertEqual(result["verify_errors"], []) + for index in range(2): + self.assertTrue((run / "cycle-1" / f"exit-{index}.log").exists()) + # A rejected batch ends a timed writer early, which fails the command. + status[0] = 503 + args.dir, args.seconds = str(run / "cycle-2"), 30 + self.assertEqual(crash.write(args), 1) + finally: + server.shutdown() + server.server_close() + + class SummaryTest(unittest.TestCase): def test_quantile_interpolates_like_prometheus(self): buckets = summary.parse_metrics(ACK_BUCKETS)["buckets"]["append_ack_seconds"] @@ -117,14 +424,14 @@ class SummaryTest(unittest.TestCase): self.assertEqual(summary.quantile(0.99, [(0.1, 1.0), (float("inf"), 4.0)]), 0.1) self.assertNotIn("kafka_client_produce_elapsed_count", summary.parse_metrics(ACK_BUCKETS)["samples"]) - def write_cell(self, root, cell_id, rows_per_second, exit_code=0, metrics=""): + def write_cell(self, root, cell_id, rows_per_second, exit_code=0, metrics="", **extra): round_number, concurrency, config = cell_id.split("-") path = root / "cells" / cell_id (path / "greptimedb-1").mkdir(parents=True) (path / "cell.json").write_text(json.dumps({ "id": cell_id, "round": int(round_number), "concurrency": int(concurrency), "config": config, - "wal": bench.WAL_CONFIGS[config], "metrics_scraped": True, - "exit_code": exit_code, "wal_device": "/dev/nvme1n1", "version": "1.4.0", "load_wall_seconds": 61.0})) + "wal": bench.WAL_CONFIGS[config], "metrics_scraped": True, "env": {}, + "exit_code": exit_code, "wal_device": "/dev/nvme1n1", "version": "1.4.0", "load_wall_seconds": 61.0} | extra)) if exit_code == 0: (path / "greptimedb-1" / "load-result.json").write_text(json.dumps({ "count_after": 6000, "load_seconds": 60.0, "rows_per_second": rows_per_second})) @@ -149,10 +456,172 @@ class SummaryTest(unittest.TestCase): d_cell = next(c for c in result["cells"] if c["id"] == "1-8-D") self.assertEqual(d_cell["creates_per_minute"], 120.0) self.assertEqual(d_cell["stalled_appends_total"], 3) - self.assertIn("| 1-8-D | 80 | 60.0 | 61.0 | 6,000 | 128.0 | 253.4 | 120 | 120.0 | 3 | 0 | 0 | 0 |", output) + self.assertIn("| 1-8-D | — | 80 | 60.0 | 61.0 | 128.0 / 253.4 | — / — | — / — | 120 | 120.0 | 3 | 0 | 0 | 0 | — | — |", output) r_cell = next(c for c in result["cells"] if c["id"] == "1-8-R") self.assertIsNone(r_cell["stalled_appends_total"]) + def test_ratios_cover_both_controls_and_restart_cells_report_recovery(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + restart = {"restart": True, "restart_scraped": True, "restart_seconds": 3.5, "indexed_before_kill": 117, + "min_objects": 6000, + "restart_query_seconds": 4.25, "restart_count": 6000, + "server_cpu_seconds": 122.0, "host_cpu_busy": 0.5, "host_cpus": 8} + self.write_cell(root, "1-8-R", 100.0) + self.write_cell(root, "1-8-R0", 125.0) + self.write_cell(root, "1-8-D", 50.0, metrics=ACK_BUCKETS, env={"RUST_LOG": "info"}, **restart) + self.write_cell(root, "1-8-E", 75.0) + (root / "cells" / "1-8-D" / "metrics-restart.prom").write_text( + "greptime_logstore_object_store_wal_recovery_seconds_sum 1.5\n" + "greptime_logstore_object_store_wal_recovery_seconds_count 1\n" + "greptime_logstore_object_store_wal_recovered_objects_total 118\n") + rows = [summary.load_cell(path) for path in sorted((root / "cells").iterdir())] + pairs = {p["ratio"]: p["geomean"] for p in summary.ratios(rows)} + self.assertEqual(pairs, {"D/R": 0.5, "D/R0": 0.4, "E/R": 0.75, "E/R0": 0.6}) + cell = next(r for r in rows if r["id"] == "1-8-D") + self.assertEqual(cell["restart"], {"wall_seconds": 3.5, "query_seconds": 4.25, "rows": 6000, + "recovery_seconds": 1.5, "recovered_objects": 118, + "indexed_before_kill": 117, "min_objects": 6000, "replay": None}) + self.assertEqual(cell["server_cpu_cores"], 2.0) + output = summary.render({}, rows, summary.ratios(rows)) + self.assertIn("| 1-8-D | RUST_LOG=info | 120 | 117 (under 6,000) | 3.50 | 4.25 | 1.500 | 118 " + "| 6,000 |", output) + self.assertIn("| 2.00 | 50% of 8 |", output) + self.assertIn("| 8 | E/R0 | 0.600 | 0.600 |", output) + + def test_traffic_idle_crash_and_pinned_cells_are_summarized(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + traffic = ("greptime_logstore_object_store_wal_created_objects_total 600\n" + "greptime_logstore_object_store_wal_object_bytes_sum 3000000000\n" + "greptime_logstore_object_store_wal_object_bytes_count 600\n" + 'opendal_operation_duration_seconds_count{scheme="s3",operation="write"} 650\n' + 'opendal_operation_duration_seconds_count{scheme="s3",operation="read"} 40\n') + # The final scrape follows a 180 s idle window that created 2 objects of 1 MB. + self.write_cell(root, "1-32-E", 90.0, idle_seconds=180, group="", metrics=traffic.replace( + " 600\n", " 602\n", 1).replace("3000000000", "3002000000").replace("count 600", "count 602") + .replace("} 650", "} 652")) + (root / "cells" / "1-32-E" / "metrics-idle.prom").write_text(traffic.replace("} 650", "} 648")) + cell = summary.load_cell(root / "cells" / "1-32-E") + self.assertEqual((cell["objects_per_second"], cell["mean_object_bytes"], cell["sealed_on_size"]), + (10.0, 5e6, False)) + self.assertAlmostEqual(cell["wal_mb_per_second"], 50.0) + self.assertEqual(cell["idle"], {"seconds": 180, "creates": 2, "operations": {"write": 4}}) + path = root / "cells" / "pin-1-8-D" + path.mkdir(parents=True) + (path / "cell.json").write_text(json.dumps({ + "id": "pin-1-8-D", "group": "pin", "mode": "crash", "round": 1, "concurrency": 8, "config": "D", + "wal": bench.WAL_CONFIGS["D"], "exit_code": 0, "metrics_scraped": True})) + (path / "cycles.jsonl").write_text(json.dumps({ + "cycle": 1, "kill_after_seconds": 540.0, "load_seconds": 540.0, "cycle_acked_batches": 4000, + "cycle_present_batches": 4003, "cycle_lost_batches": 0, "cycle_lost_rows": 0, "cycle_extra_batches": 3, + "in_flight": 8, "total_found_batches": 4003, "total_acked_batches": 4000, "total_lost_batches": 0, + "restart_seconds": 4.5, "pin_rows": 1, "passed": True, "metrics_complete": True, + "kill_errors": [], "verify_errors": []}) + "\n") + (path / "metrics-start-1.prom").write_text("greptime_logstore_object_store_wal_created_objects_total 1\n") + (path / "metrics-cycle-1.prom").write_text( + "greptime_logstore_object_store_wal_indexed_objects 5400\n" + "greptime_logstore_object_store_wal_created_objects_total 5401\n" + "greptime_logstore_object_store_wal_object_bytes_sum 10800000000\n" + "greptime_logstore_object_store_wal_object_bytes_count 5400\n" + 'opendal_operation_duration_seconds_count{scheme="s3",operation="write"} 9000\n' + 'opendal_operation_duration_seconds_count{scheme="fs",operation="write"} 70\n') + (path / "metrics-restart-1.prom").write_text( + "greptime_logstore_object_store_wal_recovery_seconds_sum 2.5\n" + "greptime_logstore_object_store_wal_recovery_seconds_count 1\n" + "greptime_logstore_object_store_wal_recovered_objects_total 5399\n") + (path / "objects.tsv").write_text("100 300 0 301\n130 600 0 601\n160 missing\n") + rows = [summary.load_cell(p) for p in sorted((root / "cells").iterdir())] + output = summary.render({}, rows, summary.ratios(rows)) + self.assertIn("| 1-32-E | 50.0 | 10.0 | 4.77 | no | read 40, write 648 | — | 180 | 2 | write 4 |", output) + self.assertIn("| pin-1-8-D | 1 | 540.0 | 4,000 | 4,003 | 0 | 0 | 3 | 8 | 4,003 | 4,000 | 0 | 5,400 " + "| 4.50 | 2.500 | 5,399 | 1 | pass |", output) + self.assertIn("| pin-1-8-D | 1 | 540.0 | 20.0 | 10.0 | 1.91 | no | write 9,000 | — |", output) + self.assertIn("| pin-1-8-D | 300, 600, missing | 1 | 0 | 601 |", output) + + def test_partitioned_load_cells_report_client_latency_entries_and_regions(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + metrics = (ACK_BUCKETS + "greptime_logstore_object_store_wal_object_entries_sum 4000\n" + "greptime_logstore_object_store_wal_object_entries_count 500\n" + "greptime_logstore_object_store_wal_object_bytes_count 500\n") + self.write_cell(root, "1-8-D", 80.0, metrics=metrics, partitions=8) + path = root / "cells" / "1-8-D" + transcript = [{"step": "create_table", "duration_seconds": 9.0}] + transcript += [{"step": "grpc_write_batch", "status": "success", "duration_seconds": d / 1000} + for d in range(1, 101)] + (path / "greptimedb-1" / "load-transcript.jsonl").write_text( + "".join(json.dumps(r) + "\n" for r in transcript)) + (path / "regions.json").write_text(json.dumps({"output": [{"records": {"rows": [ + [f"p{i}", 4398046511104 + i, 750 + i, 1000] for i in range(8)]}}]})) + cell = summary.load_cell(path) + self.assertEqual((cell["client_p50"], cell["client_p99"]), (0.05, 0.099)) + self.assertEqual((cell["entries_per_object"], cell["appends_per_object"]), (8.0, 0.2)) + self.assertEqual(cell["region_rows"], [750 + i for i in range(8)]) + output = summary.render({}, [cell], []) + self.assertIn("| 1-8-D | 50.0 / 99.0 | 8.0 | 0.20 | 8 | 750 / 757 |", output) + + def test_restarts_report_what_each_region_replayed(self): + log = ("INFO mito2::region::opener: Replay WAL for region: 4398046511104(1024, 0), provider: ObjectStore(" + 'ObjectStoreProvider { region_id: 4398046511104(1024, 0), prefix: "wal" }), rows recovered: 0, ' + "replay from entry id: 1048578, last entry id: 1048577, total timeseries replayed: 0, elapsed: 60.581µs\n" + "INFO mito2::region::opener: Replay WAL for region: 4402341478400(1025, 0), provider: ObjectStore(" + 'ObjectStoreProvider { region_id: 4402341478400(1025, 0), prefix: "wal" }), rows recovered: 12625000, ' + "replay from entry id: 4853858306, last entry id: 8164212737, total timeseries replayed: 82, " + "elapsed: 95.150155247s\n" + "INFO mito2::region::opener: Replay WAL for region: 4402341478401(1025, 1), provider: ObjectStore(" + 'ObjectStoreProvider { region_id: 4402341478401(1025, 1), prefix: "wal" }), rows recovered: 5, ' + "replay from entry id: 8164212736, last entry id: 8164212737, total timeseries replayed: 1, " + "elapsed: 12.5ms\n") + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) + (path / "server-restart.log").write_text(log) + metrics = summary.parse_metrics("greptime_logstore_object_store_wal_read_seconds_sum 95.15\n" + "greptime_logstore_object_store_wal_read_seconds_count 3\n") + result = summary.replay(path / "server-restart.log", metrics) + self.assertEqual((result["rows"], result["object_span"], result["read_seconds"], result["reads"]), + (12_625_005, 3158 + 1, 95.15, 3)) + self.assertAlmostEqual(result["seconds"], 95.150155247) + self.assertEqual([(r["region"], r["rows"], r["object_span"]) for r in result["regions"]], + [(4398046511104, 0, 0), (4402341478400, 12_625_000, 3158), (4402341478401, 5, 1)]) + self.assertEqual((result["regions"][1]["first_entry_id"], result["regions"][1]["last_entry_id"]), + (4853858306, 8164212737)) + self.assertAlmostEqual(result["regions"][0]["seconds"], 60.581e-6) + self.assertEqual(summary.replayed(result["regions"]), + "4402341478400: 12,625,000 rows, object sequences 4,629–7,786, 95.15 s; " + "4402341478401: 5 rows, object sequences 7,786–7,786, 0.01 s") + self.assertIsNone(summary.replay(path / "absent.log", metrics)) + raft = log.replace("ObjectStore(", "RaftEngine(") + (path / "server-restart.log").write_text(raft) + self.assertIsNone(summary.replay(path / "server-restart.log", metrics)["object_span"]) + # A region flushed in object 100 that next wrote in object 200 reads one object but spans 101. + gap = log.replace("replay from entry id: 4853858306", f"replay from entry id: {(100 << 20) + 2}") \ + .replace("last entry id: 8164212737", f"last entry id: {(200 << 20) + 1}") + (path / "server-restart.log").write_text(gap) + self.assertEqual(summary.replay(path / "server-restart.log", metrics)["regions"][1]["object_span"], 101) + + def test_a_drop_cell_reports_the_objects_collected_after_its_restart(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + path = root / "cells" / "drop-1-8-D" + (path / "dropped").mkdir(parents=True) + (path / "cell.json").write_text(json.dumps({ + "id": "drop-1-8-D", "group": "drop", "mode": "drop", "round": 1, "concurrency": 8, "config": "D", + "wal": bench.WAL_CONFIGS["D"], "exit_code": 0, "restart_seconds": 2.5, "restart_scraped": True})) + write_lines(path / "dropped" / "acked-0.log", [0, 8, 16]) + write_lines(path / "dropped" / "acked-1.log", [1, 9]) + (path / "metrics-written.prom").write_text("greptime_logstore_object_store_wal_indexed_objects 300\n") + (path / "metrics-restart.prom").write_text( + "greptime_logstore_object_store_wal_recovered_objects_total 280\n" + "greptime_logstore_object_store_wal_deleted_objects_total 279\n") + (path / "collected.tsv").write_text("100 281\n101 120\n104 2\n") + cell = summary.load_cell(path) + self.assertEqual(cell["drop"], {"acked_batches": 5, "indexed_before_drop": 300, "restart_seconds": 2.5, + "recovered_objects": 280, "indexed_at_end": 2, + "deleted_after_restart": 279, "seconds_to_collect": 4}) + output = summary.render({}, [cell], []) + self.assertIn("| drop-1-8-D | 5 | 300 | 2.50 | 280 | 2 | 279 | 4 | collected |", output) + def test_untouched_counters_are_zero_only_for_an_active_object_store_wal(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) @@ -199,25 +668,54 @@ FAKE_DOCKER = r"""#!/usr/bin/env bash printf '%s\n' "docker $*" >> "$FAKE_LOG" case "$1" in run) - for arg in "$@"; do [[ "$arg" != -d ]] || { echo "keys=${GREPTIMEDB_STANDALONE__STORAGE__ACCESS_KEY_ID:+set}${GREPTIMEDB_STANDALONE__STORAGE__SECRET_ACCESS_KEY:+set}" >> "$FAKE_LOG"; echo fakecid; exit 0; }; done + for arg in "$@"; do [[ "$arg" != -d ]] || { echo $(($(cat "$FAKE_LOG.runs" 2>/dev/null || echo 0) + 1)) > "$FAKE_LOG.runs"; rm -f "${FAKE_KILL_FLAG:-/nonexistent}"; echo "keys=${GREPTIMEDB_STANDALONE__STORAGE__ACCESS_KEY_ID:+set}${GREPTIMEDB_STANDALONE__STORAGE__SECRET_ACCESS_KEY:+set}" >> "$FAKE_LOG"; echo fakecid; exit 0; }; done for arg in "$@"; do case "$arg" in type=bind,src=*,dst=/results) results=${arg#type=bind,src=}; results=${results%,dst=/results} ;; esac done [[ "$FAKE_LOAD_RC" == 0 ]] || exit "$FAKE_LOAD_RC" mkdir -p "$results/greptimedb-1" - echo "{\"count_before\": 0, \"count_after\": $FAKE_ROWS, \"load_seconds\": 2, \"rows_per_second\": 3000}" \ + echo "{\"count_before\": 0, \"count_after\": $FAKE_ROWS, \"load_seconds\": 2, \"rows_per_second\": 3000, \"table\": \"agent_observations\"}" \ > "$results/greptimedb-1/load-result.json" ;; + kill) touch "${FAKE_KILL_FLAG:-/dev/null}" ;; inspect) echo true ;; - exec) echo 'version: 1.4.0-test' ;; - logs) echo 'server log' ;; + exec) [[ "$*" == *cpu.stat* ]] && echo 'usage_usec 5000000' || echo 'version: 1.4.0-test' ;; + logs) echo 'server log' + # The rows a server replayed are its start number, so each log can be told apart. + [[ -z "${FAKE_REPLAY:-}" ]] || echo "INFO Replay WAL for region: 4398046511104(1024, 0), provider: ObjectStore(x), rows recovered: $(cat "$FAKE_LOG.runs"), replay from entry id: 1048577, last entry id: 2097153, total timeseries replayed: 1, elapsed: 1.5s" ;; esac """ FAKE_CURL = r"""#!/usr/bin/env bash printf '%s\n' "curl $*" >> "$FAKE_LOG" +[[ "$*" != *"--max-time 10"* || "${FAKE_SAMPLE_RC:-0}" == 0 ]] || exit "$FAKE_SAMPLE_RC" for arg in "$@"; do + [[ "$arg" != *"GROUP BY seq"* ]] || { python3 -c 'import collections, json, sys +counts = collections.Counter(int(l) for l in open(sys.argv[1]) if l.strip()) +print(json.dumps({"output": [{"records": {"rows": [[k, v * int(sys.argv[2])] for k, v in counts.items()]}}]}))' \ + "$FAKE_FOUND" "$FAKE_BATCH_ROWS"; exit 0; } +done +for arg in "$@"; do + [[ "$arg" != *region_statistics* ]] || { python3 -c 'import json, sys +print(json.dumps({"output": [{"records": {"rows": [[f"p{i}", i, int(n), 1] for i, n in enumerate(sys.argv[1].split(","))]}}]}))' \ + "$FAKE_REGION_ROWS"; exit 0; } +done +for arg in "$@"; do + [[ "$arg" != */v1/sql* ]] || { echo "{\"output\": [{\"records\": {\"rows\": [[${FAKE_COUNT:-6000}]]}}]}"; exit 0; } [[ "$arg" != */metrics ]] || { [[ "$FAKE_METRICS_RC" == 0 ]] || exit "$FAKE_METRICS_RC" - printf '%s\n' 'greptime_logstore_object_store_wal_poisoned_total 1' 'other_total 2'; exit 0; + printf '%s\n' 'greptime_logstore_object_store_wal_poisoned_total 1' 'other_total 2' + [[ -z "${FAKE_FULL_METRICS:-}" ]] || printf '%s\n' \ + "greptime_logstore_object_store_wal_indexed_objects $FAKE_FULL_METRICS" \ + 'greptime_logstore_object_store_wal_recovery_seconds_count 1' \ + "greptime_logstore_object_store_wal_recovered_objects_total ${FAKE_RECOVERED:-0}" + # Every server counts the objects it collects from its start: with FAKE_COLLECTING one more on + # every scrape; with FAKE_FIRST_COLLECTS_AFTER the first server collects one object once it has + # been scraped that many times. + run=$(cat "$FAKE_LOG.runs") n=$(($(cat "$FAKE_LOG.scrapes.$run" 2>/dev/null || echo 0) + 1)) + echo "$n" > "$FAKE_LOG.scrapes.$run" + [[ -z "${FAKE_COLLECTING:-}" ]] || echo "greptime_logstore_object_store_wal_deleted_objects_total $((n - 1))" + [[ -z "${FAKE_FIRST_COLLECTS_AFTER:-}" || "$run" != 1 || "$n" -le "$FAKE_FIRST_COLLECTS_AFTER" ]] \ + || echo 'greptime_logstore_object_store_wal_deleted_objects_total 1' + exit 0; } done """ @@ -225,19 +723,20 @@ done class CellLifecycleTest(unittest.TestCase): # the run script against fake docker and curl - def run_cells(self, root, load_rc, metrics_rc=0): + def run_cells(self, root, load_rc, metrics_rc=0, restart=False, count=6000, groups=""): bin_dir = root / "bin" bin_dir.mkdir() for name, body in (("docker", FAKE_DOCKER), ("curl", FAKE_CURL)): (bin_dir / name).write_text(body) (bin_dir / name).chmod(0o755) - plan = bench.build_plan("8", "6000", "1000", "1", smoke=False) - plan["cells"] = [cell for cell in plan["cells"] if cell["config"] != "R"] + plan = bench.build_plan("8", "6000", "1000", "1", smoke=False, configs="D,E", restart=restart, + extra_env="RUST_LOG=info", groups=groups) (root / "results").mkdir() (root / "results" / "plan.json").write_text(json.dumps(plan)) env = os.environ | { "PATH": f"{bin_dir}:{os.environ['PATH']}", "FAKE_LOG": str(root / "fake.log"), - "FAKE_LOAD_RC": str(load_rc), "FAKE_METRICS_RC": str(metrics_rc), "FAKE_ROWS": "6000", "PLAN": str(root / "results" / "plan.json"), + "FAKE_LOAD_RC": str(load_rc), "FAKE_METRICS_RC": str(metrics_rc), "FAKE_ROWS": "6000", "FAKE_COUNT": str(count), + "PLAN": str(root / "results" / "plan.json"), "RESULTS": str(root / "results"), "CORPUS_DIR": str(root / "corpus"), "OWNER": "o11ybench-oswal-1-1", "RUNTIME_IMAGE": "runtime@sha256:" + "b" * 64, "GREPTIMEDB_IMAGE": "greptimedb@sha256:" + "a" * 64, "RUN_PREFIX": "oswal-bench/1-1", "WAL_ROOT": str(root / "wal"), "S3_BUCKET": "bucket", @@ -302,6 +801,9 @@ class CellLifecycleTest(unittest.TestCase): # the run script against fake docke self.assertIn("GREPTIMEDB_STANDALONE__STORAGE__ROOT=oswal-bench/1-1/1-8-D", config) self.assertIn("GREPTIMEDB_STANDALONE__WAL__ACK_MODE=durable", config) self.assertIn("GREPTIMEDB_STANDALONE__STORAGE__SECRET_ACCESS_KEY=", config) + self.assertIn("RUST_LOG=info", config) + self.assertIn("-e RUST_LOG=info ", log) + self.assertNotIn("docker kill", log) recorded = json.loads((cell / "cell.json").read_text()) self.assertEqual((recorded["exit_code"], recorded["version"], recorded["metrics_scraped"]), (0, "1.4.0-test", True)) @@ -313,6 +815,260 @@ class CellLifecycleTest(unittest.TestCase): # the run script against fake docke self.assertNotEqual(leaked.returncode, 0) self.assertFalse((root / "results").exists()) + def test_a_restart_cell_kills_restarts_and_counts_the_rows(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + cells, _, log, _ = self.run_cells(root, load_rc=0, restart=True) + self.assertEqual(cells.returncode, 0, cells.stdout + cells.stderr) + self.assertEqual(log.count("docker kill --signal KILL fakecid"), 2) + self.assertEqual(log.count("docker stop -t 120 fakecid"), 2) + self.assertIn("--name o11ybench-oswal-1-1-1-8-D-restart", log) + self.assertIn("sql=SELECT count(*) FROM agent_observations", log) + cell = root / "results" / "cells" / "1-8-D" + recorded = json.loads((cell / "cell.json").read_text()) + self.assertEqual((recorded["exit_code"], recorded["restart_count"], recorded["restart_scraped"]), (0, 6000, True)) + self.assertIsInstance(recorded["restart_seconds"], float) + config = (cell / "config.env").read_text().splitlines() + self.assertIn("GREPTIMEDB_STANDALONE__STORAGE__DATA_HOME=/oswal_home", config) + self.assertIn("GREPTIMEDB_STANDALONE__WAL__FLUSH_INTERVAL=10ms", config) + self.assertIn(f"type=bind,src={root}/wal/1-8-D-home,dst=/oswal_home", log) + for name in ("metrics.prom", "metrics-restart.prom", "server.log", "server-restart.log"): + self.assertTrue((cell / name).exists(), name) + + def test_an_idle_window_scrapes_before_and_after_it(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + groups = json.dumps([{"name": "idle", "configs": "D", "idle_seconds": 1}]) + cells, _, log, _ = self.run_cells(root, load_rc=0, groups=groups) + self.assertEqual(cells.returncode, 0, cells.stdout + cells.stderr) + cell = root / "results" / "cells" / "idle-1-8-D" + self.assertTrue((cell / "metrics-idle.prom").exists()) + self.assertTrue((cell / "metrics.prom").exists()) + self.assertEqual(log.count("/metrics"), 2) + self.assertIn(f"type=bind,src={root}/corpus/corpus-6000,dst=/work/corpus,readonly", log) + + def run_crash(self, root, metrics_rc=0, pin=False, sample_rc=0, full_metrics="5", groups=None, env=None, + status=200): + """A crash cell, or another group that writes through the crash writer, against fake docker + and a stub SQL endpoint that the fake kill takes down.""" + import http.server + import threading + flag, found = root / "killed", root / "found.log" + found.write_text("") + lock = threading.Lock() + + class Handler(http.server.BaseHTTPRequestHandler): + def do_POST(self): + body = urllib.parse.unquote_plus(self.rfile.read(int(self.headers["Content-Length"])).decode()) + if flag.exists(): + self.close_connection = True + self.connection.close() + return + seq = re.search(r"VALUES [(](\d+),", body).group(1) + with lock, open(found, "a") as out: + out.write(seq + "\n") + self.send_response(status) + self.end_headers() + + def log_message(self, *args): + pass + + class Server(http.server.ThreadingHTTPServer): + request_queue_size = 64 # all eight writers connect at once + + server = Server(("127.0.0.1", 0), Handler) + threading.Thread(target=server.serve_forever, daemon=True).start() + groups = groups or json.dumps([{"name": "crash", "configs": "D", "mode": "crash", "cycles": 2, + "kill_after": [1, 1], "pin": pin}]) + old = os.environ.copy() + os.environ.update({"OSWAL_SQL_URL": f"http://127.0.0.1:{server.server_port}/v1/sql", + "FAKE_KILL_FLAG": str(flag), "FAKE_FOUND": str(found), "FAKE_BATCH_ROWS": "1000", + "OSWAL_SAMPLE_SECONDS": "0.2", "FAKE_SAMPLE_RC": str(sample_rc), + "FAKE_FULL_METRICS": full_metrics} | (env or {})) + try: + return self.run_cells(root, load_rc=0, metrics_rc=metrics_rc, groups=groups, + count=(env or {}).get("FAKE_COUNT", 6000)) + finally: + os.environ.clear() + os.environ.update(old) + server.shutdown() + server.server_close() + + def test_a_crash_cell_kills_restarts_and_checks_every_cycle(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + cells, _, log, _ = self.run_crash(root, env={"FAKE_REPLAY": "1"}) + self.assertEqual(cells.returncode, 0, cells.stdout + cells.stderr) + cell = root / "results" / "cells" / "crash-1-8-D" + # Server 1 is loaded and killed; servers 2 and 3 are the restarts of cycles 1 and 2. + replays = [c["replay"]["rows"] for c in summary.load_cell(cell)["cycles"]] + self.assertEqual(replays, [2, 3]) + cycles = [json.loads(line) for line in (cell / "cycles.jsonl").read_text().splitlines()] + self.assertEqual([c["cycle"] for c in cycles], [1, 2]) + for c in cycles: + self.assertTrue(c["passed"] and c["metrics_complete"], c) + self.assertEqual((c["cycle_lost_batches"], c["writers_at_kill"]), (0, 8)) + self.assertGreater(c["cycle_acked_batches"], 0) + self.assertEqual(log.count("docker kill --signal KILL fakecid"), 2) + self.assertEqual(log.count("docker wait fakecid"), 2) + self.assertIn("GREPTIMEDB_STANDALONE__STORAGE__DATA_HOME=/oswal_home", (cell / "config.env").read_text()) + for name in ("metrics-start-1.prom", "metrics-cycle-2.prom", "metrics-restart-2.prom", "server.log", + "server-restart-1.log", "server-restart-2.log"): + self.assertTrue((cell / name).exists(), name) + + def test_failed_pinned_samples_are_missing_not_zero(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + cells, _, log, _ = self.run_crash(root, pin=True, sample_rc=28, env={"FAKE_COUNT": "1"}) + self.assertEqual(cells.returncode, 0, cells.stdout + cells.stderr) + cell = root / "results" / "cells" / "crash-1-8-D" + samples = (cell / "objects.tsv").read_text().splitlines() + self.assertTrue(samples and all(line.split()[1:] == ["missing"] for line in samples), samples) + self.assertIn("sql=INSERT INTO pin VALUES (1, 0)", log) + self.assertIn("WITH ('auto_flush_interval' = '1d')", urllib.parse.unquote_plus(log)) + self.assertEqual(len((cell / "cycles.jsonl").read_text().splitlines()), 2) + + def test_a_pinned_cell_fails_when_objects_the_pin_held_are_collected(self): + # Each server collects two objects before its kill and no periodic sample is taken. With one + # object below the pin the first server already exceeds it; with two, only both together do. + for below_pin, failed_cycle in (("1", 1), ("2", 2)): + with self.subTest(below_pin=below_pin), tempfile.TemporaryDirectory() as directory: + root = Path(directory) + cells, _, _, _ = self.run_crash(root, pin=True, full_metrics=below_pin, env={ + "FAKE_COUNT": "1", "FAKE_COLLECTING": "1", "OSWAL_SAMPLE_SECONDS": "30"}) + cell = root / "results" / "cells" / "crash-1-8-D" + self.assertEqual(json.loads((cell / "cell.json").read_text())["exit_code"], 11, + cells.stdout + cells.stderr) + self.assertFalse((cell / "objects.tsv").exists()) + self.assertEqual(len((cell / "cycles.jsonl").read_text().splitlines()), failed_cycle) + self.assertIn(f"cycle {failed_cycle} collected WAL objects the pin held", cells.stdout + cells.stderr) + + def test_a_pinned_cell_may_collect_the_objects_below_its_pin_late(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + # The start object, the one object indexed before the pin, goes only after the first cycle's + # baseline: the pin scrape and that baseline are the first server's first two scrapes. + cells, _, _, _ = self.run_crash(root, pin=True, full_metrics="1", env={ + "FAKE_COUNT": "1", "FAKE_FIRST_COLLECTS_AFTER": "2", "OSWAL_SAMPLE_SECONDS": "30"}) + cell = root / "results" / "cells" / "crash-1-8-D" + self.assertEqual(json.loads((cell / "cell.json").read_text())["exit_code"], 0, cells.stdout + cells.stderr) + self.assertNotIn("deleted_objects_total", (cell / "metrics-start-1.prom").read_text()) + self.assertIn("deleted_objects_total 1", (cell / "metrics-cycle-1.prom").read_text()) + self.assertEqual(len((cell / "cycles.jsonl").read_text().splitlines()), 2) + + def test_a_durable_pinned_cycle_fails_unless_its_pin_row_is_there_once(self): + for count in ("0", "2"): + with self.subTest(count=count), tempfile.TemporaryDirectory() as directory: + root = Path(directory) + self.run_crash(root, pin=True, env={"FAKE_COUNT": count}) + cell = root / "results" / "cells" / "crash-1-8-D" + cycle = json.loads((cell / "cycles.jsonl").read_text().splitlines()[0]) + self.assertEqual((cycle["passed"], cycle["pin_rows"]), (False, int(count))) + self.assertIn(f"the pin table holds {count} rows, not 1", cycle["verify_errors"]) + self.assertEqual(json.loads((cell / "cell.json").read_text())["exit_code"], 6) + + def test_a_scrape_without_the_wal_samples_is_incomplete(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + cells, _, _, _ = self.run_crash(root, full_metrics="") + self.assertNotEqual(cells.returncode, 0) + cell = root / "results" / "cells" / "crash-1-8-D" + cycle = json.loads((cell / "cycles.jsonl").read_text().splitlines()[0]) + self.assertEqual((cycle["passed"], cycle["metrics_complete"]), (True, False)) + self.assertEqual(json.loads((cell / "cell.json").read_text())["exit_code"], 4) + + def test_a_restart_cell_must_leave_its_minimum_objects_indexed(self): + for minimum, exit_code in ((5, 0), (6, 7)): + with self.subTest(minimum=minimum), tempfile.TemporaryDirectory() as directory: + root = Path(directory) + groups = json.dumps([{"name": "rec", "configs": "D", "restart": True, "min_objects": minimum}]) + old = os.environ.copy() + os.environ["FAKE_FULL_METRICS"] = "5" + try: + cells, _, _, _ = self.run_cells(root, load_rc=0, groups=groups) + finally: + os.environ.clear() + os.environ.update(old) + recorded = json.loads((root / "results" / "cells" / "rec-1-8-D" / "cell.json").read_text()) + self.assertEqual((recorded["exit_code"], recorded["indexed_before_kill"], recorded["restart_count"]), + (exit_code, 5, 6000)) + + def test_a_partitioned_cell_creates_its_table_and_checks_every_region_took_rows(self): + for region_rows, exit_code in (("750,750,750,750,750,750,750,750", 0), ("1000,1000,1000,0,1000,1000,1000,0", 8)): + with self.subTest(region_rows=region_rows), tempfile.TemporaryDirectory() as directory: + root = Path(directory) + (root / "corpus" / "corpus-6000").mkdir(parents=True) + (root / "corpus" / "corpus-6000" / "summary.json").write_text(json.dumps( + {"start": "2026-07-10T00:00:00Z", "end": "2026-07-10T01:00:00Z"})) + groups = json.dumps([{"name": "multi", "configs": "D", "partitions": 8}]) + old = os.environ.copy() + os.environ["FAKE_REGION_ROWS"] = region_rows + try: + cells, _, log, _ = self.run_cells(root, load_rc=0, groups=groups) + finally: + os.environ.clear() + os.environ.update(old) + cell = root / "results" / "cells" / "multi-1-8-D" + self.assertEqual(json.loads((cell / "cell.json").read_text())["exit_code"], exit_code, + cells.stdout + cells.stderr) + self.assertIn("PARTITION ON COLUMNS (`event_time`)", urllib.parse.unquote_plus(log)) + self.assertIn(f"type=bind,src={SCRIPTS}/oswal-s3-bench-load.sh,dst=/work/oswal-s3-bench-load.sh,readonly", + log) + self.assertIn(" /work/oswal-s3-bench-load.sh", log) + self.assertNotIn("greptimedb.sh load", log) + self.assertEqual(len(json.loads((cell / "regions.json").read_text())["output"][0]["records"]["rows"]), 8) + + def test_a_drop_cell_restarts_and_waits_for_the_dropped_objects_to_go(self): + for indexed, recovered, exit_code in (("2", "0", 0), ("50", "200", 0), ("50", "60", 9)): + with self.subTest(indexed=indexed, recovered=recovered), tempfile.TemporaryDirectory() as directory: + root = Path(directory) + groups = json.dumps([{"name": "drop", "mode": "drop", "configs": "D", "write_seconds": 1}]) + cells, _, log, _ = self.run_crash(root, full_metrics=indexed, groups=groups, + env={"FAKE_RECOVERED": recovered, "OSWAL_COLLECT_SECONDS": "1"}) + cell = root / "results" / "cells" / "drop-1-8-D" + recorded = json.loads((cell / "cell.json").read_text()) + self.assertEqual(recorded["exit_code"], exit_code, cells.stdout + cells.stderr) + self.assertIsInstance(recorded["restart_seconds"], float) + sql = urllib.parse.unquote_plus(log) + for statement in ("CREATE TABLE dropped", "CREATE TABLE kept", "DROP TABLE dropped", + "ADMIN FLUSH_TABLE('kept')", "INSERT INTO kept VALUES (-1, "): + self.assertIn(statement, sql) + self.assertLess(sql.index("DROP TABLE dropped"), log.index("docker kill --signal KILL fakecid")) + self.assertEqual(log.count("docker kill --signal KILL fakecid"), 1) + self.assertTrue(list((cell / "dropped").glob("acked-*.log"))) + self.assertTrue(list((cell / "kept").glob("acked-*.log"))) + for name in ("metrics-written.prom", "metrics-restart.prom", "server.log", + "server-restart.log", "collected.tsv"): + self.assertTrue((cell / name).exists(), name) + + def test_a_drop_cell_whose_writers_fail_stops_before_the_drop(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + groups = json.dumps([{"name": "drop", "mode": "drop", "configs": "D", "write_seconds": 1}]) + cells, _, log, _ = self.run_crash(root, full_metrics="2", groups=groups, status=503) + recorded = json.loads((root / "results" / "cells" / "drop-1-8-D" / "cell.json").read_text()) + self.assertEqual(recorded["exit_code"], 10, cells.stdout + cells.stderr) + self.assertNotIn("DROP TABLE", urllib.parse.unquote_plus(log)) + self.assertNotIn("docker kill", log) + + def test_a_crash_cycle_without_its_scrapes_fails_the_cell(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + cells, _, _, _ = self.run_crash(root, metrics_rc=22) + self.assertNotEqual(cells.returncode, 0) + cell = root / "results" / "cells" / "crash-1-8-D" + cycle = json.loads((cell / "cycles.jsonl").read_text().splitlines()[0]) + self.assertEqual((cycle["passed"], cycle["metrics_complete"]), (True, False)) + self.assertEqual(json.loads((cell / "cell.json").read_text())["exit_code"], 4) + + def test_a_restart_that_loses_rows_fails_the_cell(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + cells, _, _, _ = self.run_cells(root, load_rc=0, restart=True, count=5999) + self.assertNotEqual(cells.returncode, 0) + recorded = json.loads((root / "results" / "cells" / "1-8-D" / "cell.json").read_text()) + self.assertEqual((recorded["exit_code"], recorded["restart_count"]), (5, 5999)) + if __name__ == "__main__": unittest.main() From 98f79d0cf9a61a27025e3ae0cc97abc158bdf59f Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Fri, 16 Oct 2026 09:57:54 +0000 Subject: [PATCH 028/239] fix(meta): allow affinity fallback during failover (#9641) Signed-off-by: WenyXu --- src/meta-srv/AGENTS.md | 6 +- src/meta-srv/src/region/supervisor.rs | 164 +++++++------------------- 2 files changed, 48 insertions(+), 122 deletions(-) diff --git a/src/meta-srv/AGENTS.md b/src/meta-srv/AGENTS.md index f1153b4cc35..c1aff378496 100644 --- a/src/meta-srv/AGENTS.md +++ b/src/meta-srv/AGENTS.md @@ -70,8 +70,10 @@ API under `src/meta-srv/src/service/admin/`. supervisor; without one, allocation skips source lookup. DDL configurators can install a query provider before procedure loaders are registered. Repartition and failover resolve source requirements before selection. Each selector allows - one unconstrained retry on candidate shortage only when explicitly configured; - metadata errors propagate. Procedure state is unchanged. + one unconstrained retry on candidate shortage when its context permits it. + Failover always permits this retry and continues without affinity when source + resolution fails or returns no requirements; DDL/repartition retain their own + policy. Candidate discovery errors still propagate. Procedure state is unchanged. - **Selector stats fields**: keep `selector/` implementations in sync with the stats they consume. diff --git a/src/meta-srv/src/region/supervisor.rs b/src/meta-srv/src/region/supervisor.rs index 1c431ef9960..5cf3f8e7fa8 100644 --- a/src/meta-srv/src/region/supervisor.rs +++ b/src/meta-srv/src/region/supervisor.rs @@ -18,7 +18,6 @@ use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; use async_trait::async_trait; -use common_error::ext::BoxedError; use common_meta::DatanodeId; use common_meta::affinity::{AffinityFallback, AffinityReference, AffinityResolverRef}; use common_meta::datanode::Stat; @@ -410,7 +409,11 @@ impl RegionSupervisor { failure_detector: RegionFailureDetector::new(options), failover_counts: HashMap::new(), receiver: event_receiver, - selector_context, + // Failover prioritizes recovery over affinity, independently of DDL policy. + selector_context: SelectorContext { + affinity_fallback: AffinityFallback::Unconstrained, + ..selector_context + }, selector, affinity_resolver: None, region_migration_manager, @@ -671,26 +674,31 @@ impl RegionSupervisor { failure_datanodes: &[DatanodeId], ) -> Result> { let affinity = if let Some(resolver) = &self.affinity_resolver { - let mut resolved = resolver + match resolver .resolve(&[AffinityReference::Datanode(from_peer_id)]) .await - .map_err(BoxedError::new) - .context(error::OtherSnafu)?; - ensure!( - resolved.len() == 1, - error::UnexpectedSnafu { - violated: "affinity resolution must preserve the number and order of references", + { + Ok(mut resolved) => { + ensure!( + resolved.len() == 1, + error::UnexpectedSnafu { + violated: "affinity resolution must preserve the number and order of references", + } + ); + let affinity = resolved.pop().flatten(); + if affinity.is_none() { + warn!( + "Source affinity unavailable for datanode {}; selecting failover peers without affinity", + from_peer_id + ); + } + affinity } - ); - let affinity = resolved.pop().flatten(); - ensure!( - affinity.is_some() - || self.selector_context.affinity_fallback == AffinityFallback::Unconstrained, - error::UnexpectedSnafu { - violated: "source affinity requirements are unavailable" + Err(error) => { + warn!(error; "Failed to resolve failover affinity for datanode {}; selecting without affinity", from_peer_id); + None } - ); - affinity + } } else { None }; @@ -999,7 +1007,7 @@ pub(crate) mod tests { } #[tokio::test] - async fn test_failover_passes_affinity_and_propagates_errors() { + async fn test_failover_affinity_fallback_and_selector_errors() { use std::collections::HashSet; use std::sync::atomic::{AtomicUsize, Ordering}; @@ -1128,91 +1136,14 @@ pub(crate) mod tests { } } for stat_aware in [false, true] { - for (resolution, fallback, outcome, expected_calls, succeeds) in [ - ( - Resolution::Available, - AffinityFallback::Error, - Outcome::Success, - 1, - true, - ), - ( - Resolution::Available, - AffinityFallback::Error, - Outcome::Short, - 1, - false, - ), - ( - Resolution::Available, - AffinityFallback::Error, - Outcome::Exhausted, - 1, - false, - ), - ( - Resolution::Available, - AffinityFallback::Error, - Outcome::Error, - 1, - false, - ), - ( - Resolution::Missing, - AffinityFallback::Error, - Outcome::Success, - 0, - false, - ), - ( - Resolution::Error, - AffinityFallback::Error, - Outcome::Success, - 0, - false, - ), - ( - Resolution::Available, - AffinityFallback::Unconstrained, - Outcome::Short, - 1, - false, - ), - ( - Resolution::Available, - AffinityFallback::Unconstrained, - Outcome::Exhausted, - 1, - false, - ), - ( - Resolution::Available, - AffinityFallback::Unconstrained, - Outcome::Error, - 1, - false, - ), - ( - Resolution::Missing, - AffinityFallback::Unconstrained, - Outcome::Success, - 1, - true, - ), - ( - Resolution::Error, - AffinityFallback::Unconstrained, - Outcome::Success, - 0, - false, - ), - ( - Resolution::Disabled, - AffinityFallback::Unconstrained, - Outcome::Success, - 1, - true, - ), + for (resolution, outcome, succeeds) in [ + (Resolution::Available, Outcome::Success, true), + (Resolution::Available, Outcome::Short, false), + (Resolution::Available, Outcome::Exhausted, false), + (Resolution::Available, Outcome::Error, false), + (Resolution::Missing, Outcome::Success, true), + (Resolution::Error, Outcome::Success, true), + (Resolution::Disabled, Outcome::Success, true), ] { let resolver = Arc::new(Resolver { result: resolution, @@ -1220,11 +1151,10 @@ pub(crate) mod tests { }); let selector = Arc::new(FailoverSelector { outcome, - expected_fallback: fallback, + expected_fallback: AffinityFallback::Unconstrained, calls: Mutex::new(vec![]), }); let (mut supervisor, _sender) = new_test_supervisor(); - supervisor.selector_context.affinity_fallback = fallback; if resolution != Resolution::Disabled { supervisor = supervisor.with_affinity_resolver(resolver.clone()); } @@ -1243,20 +1173,14 @@ pub(crate) mod tests { ); } else { let error = result.unwrap_err(); - if resolution == Resolution::Error { - assert!(matches!(error, error::Error::Other { .. })); - } else if resolution == Resolution::Missing { - assert!(matches!(error, error::Error::Unexpected { .. })); - } else { - match outcome { - Outcome::Short | Outcome::Exhausted => { - assert!(matches!(error, error::Error::NoEnoughAvailableNode { .. })) - } - Outcome::Error => { - assert!(matches!(error, error::Error::Unexpected { .. })) - } - Outcome::Success => unreachable!(), + match outcome { + Outcome::Short | Outcome::Exhausted => { + assert!(matches!(error, error::Error::NoEnoughAvailableNode { .. })) } + Outcome::Error => { + assert!(matches!(error, error::Error::Unexpected { .. })) + } + Outcome::Success => unreachable!(), } } assert_eq!( @@ -1264,7 +1188,7 @@ pub(crate) mod tests { usize::from(resolution != Resolution::Disabled) ); let calls = selector.calls.lock().unwrap(); - assert_eq!(calls.len(), expected_calls); + assert_eq!(calls.len(), 1); if let Some(affinity) = calls.first() { assert_eq!(affinity.is_some(), resolution == Resolution::Available); } From 822c90c15f9837f5e49603cabae3a2067b244a76 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Fri, 16 Oct 2026 12:10:27 +0000 Subject: [PATCH 029/239] fix(test): download darwin release assets for compat runs on macOS (#9630) Signed-off-by: Dennis Zhuang --- tests/runner/src/util.rs | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/runner/src/util.rs b/tests/runner/src/util.rs index 1be4183b49a..e448f660590 100644 --- a/tests/runner/src/util.rs +++ b/tests/runner/src/util.rs @@ -76,7 +76,9 @@ async fn download_files(url: &str, path: &str) { .get(url) .send() .await - .expect("Failed to send download request"); + .expect("Failed to send download request") + .error_for_status() + .unwrap_or_else(|e| panic!("Failed to download {url}: {e}")); let len = resp.content_length(); let mut stream = resp.bytes_stream(); let mut size_downloaded = 0; @@ -111,7 +113,11 @@ fn decompress(archive: &str, dest: &str) { /// /// * `version` - The version of the binary to download. i.e. "v0.9.5" pub async fn pull_binary(version: &str) { - let os = std::env::consts::OS; + // Release assets name macOS builds `darwin`. + let os = match std::env::consts::OS { + "macos" => "darwin", + os => os, + }; let arch = match std::env::consts::ARCH { "x86_64" => "amd64", "aarch64" => "arm64", From e68e8c46d8bd3a3f3494e3bd416ee1ab4abbe0dd Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Fri, 16 Oct 2026 12:43:27 +0000 Subject: [PATCH 030/239] perf(log-store): reuse the release records the object store WAL recovery already read (#9645) * perf(log-store): reuse the release records the object store WAL recovery already read Recovery reads every object it lists to verify its header and footer: an object of at most 64 KiB whole, a larger one as its header and the last 64 KiB. It then read the segment of release records of every indexed object that holds one again with a range read, although the segment, the last one before the footer, is almost always among the bytes already read. A prefix of many dropped regions, each released in an object of its own, paid two reads per such object. The footer read now keeps a copy of that segment when the bytes it read hold it, and recovery decodes the copy, or reads the segment as before when the copy is missing, only for the objects on the selected chain. The segment is copied so that the read window is not held until the releases are decoded. Chain selection, the merge of the largest released entry id per region, and the failure of recovery on a release segment that cannot be fetched or decoded, whatever on_corrupted_segment says, are unchanged. Tests cover the reads of a recovery whose release segments are all in the read bytes, a release segment that starts before the tail window, and that a damaged release segment fails recovery on the chain and is never decoded off it. An ignored measurement test counts the reads and times recovery over a prefix of dropped regions. Signed-off-by: jeremyhi * perf(log-store): bound the release segments object store WAL recovery keeps Recovery keeps the segment of release records of every object it fetches before it selects the chain, so an object off the chain could hold up to the tail window until recovery ends. Keep a segment only up to 4 KiB, far above the records of the few regions a release seals at once, and read a larger one again only when its object is on the chain, as before. Tests cover a release segment inside the tail window of an object read as its header and its tail, which exercises the offset into the window, and a segment above the kept size in an object read whole. Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- src/log-store/src/object_store_wal/store.rs | 390 +++++++++++++++++++- 1 file changed, 373 insertions(+), 17 deletions(-) diff --git a/src/log-store/src/object_store_wal/store.rs b/src/log-store/src/object_store_wal/store.rs index 0a8c68c83e7..df037b43771 100644 --- a/src/log-store/src/object_store_wal/store.rs +++ b/src/log-store/src/object_store_wal/store.rs @@ -123,8 +123,16 @@ const READ_PREFETCH_DEPTH: usize = 8; const READ_PREFETCH_BYTES: usize = 64 * 1024 * 1024; /// Bytes recovery reads from the end of an object in one request. The window /// holds the trailer and the footer of an object with up to 1364 regions of -/// 48 bytes each, so a second request for the footer is rare. +/// 48 bytes each, so a second request for the footer is rare, and usually the +/// segment of release records, which is the last segment. const RECOVERY_TAIL_WINDOW: usize = 64 * 1024; +/// Largest segment of release records that recovery keeps from the bytes it +/// read for a footer. It keeps them before it knows which objects are on the +/// chain, so the bound caps what objects off the chain hold until recovery +/// ends; a larger segment is read again if its object is on the chain. A +/// release seals its record at once, so a segment usually holds the records of +/// one or a few regions, 16 bytes each. +const MAX_KEPT_RELEASE_BYTES: u64 = 4 * 1024; /// A log store over the immutable WAL objects under one prefix. /// @@ -3031,11 +3039,13 @@ struct FetchedObjects { vanished: Vec<(u64, Error)>, } -/// An object as recovery fetched it: its key, header and footer. +/// An object as recovery fetched it: its key, header and footer, and its +/// segment of release records when the reads of the footer returned it. struct FetchedObject { object: ListedObject, header: Header, footer: Vec, + releases: Option, } /// Rebuilds the catalog from object headers and footers, @@ -3048,16 +3058,25 @@ struct FetchedObject { /// Only the objects on the chain [`select_chain`] picks are indexed. async fn recover(io: &dyn WalObjectIo) -> Result { let objects = io.list().await?; - let mut recovered = finish_recovery(fetch_footers(io, objects, RECOVERY_CONCURRENCY).await?)?; - recovered.released_entry_ids = read_releases(io, &recovered.catalog).await?; + let mut fetched = fetch_footers(io, objects, RECOVERY_CONCURRENCY).await?; + let releases = fetched + .objects + .iter_mut() + .filter_map(|fetched| Some((fetched.object.object_seq, fetched.releases.take()?))) + .collect(); + let mut recovered = finish_recovery(fetched)?; + recovered.released_entry_ids = read_releases(io, &recovered.catalog, releases).await?; Ok(recovered) } /// Reads the release records of the indexed objects, which seed the -/// watermarks of their regions: one range read of the segment of -/// [`RELEASE_REGION_ID`] per object that holds one. A segment that cannot be -/// fetched or does not decode fails recovery, whatever `on_corrupted_segment` -/// says: skipping it would lose the watermarks it seeds. +/// watermarks of their regions, from the segment of [`RELEASE_REGION_ID`] of +/// every indexed object that holds one: the bytes in `fetched` by object +/// sequence, which the reads of the footers returned, or else one range read. +/// The segments of objects that are not indexed are never decoded. A segment +/// that cannot be fetched or does not decode fails recovery, whatever +/// `on_corrupted_segment` says: skipping it would lose the watermarks it +/// seeds. /// /// A release record needs no bookkeeping of its own. It is sealed once every /// entry it covers was admitted, so it lies in the object holding the last @@ -3073,6 +3092,7 @@ async fn recover(io: &dyn WalObjectIo) -> Result { async fn read_releases( io: &dyn WalObjectIo, catalog: &ObjectCatalog, + mut fetched: HashMap, ) -> Result> { let segments = catalog .objects_in_order() @@ -3080,14 +3100,18 @@ async fn read_releases( footer .iter() .find(|entry| entry.region_id == RELEASE_REGION_ID) - .map(|entry| (object_seq, entry.clone())) + .map(|entry| (object_seq, entry.clone(), fetched.remove(&object_seq))) }) .collect::>(); let releases = futures::stream::iter(segments) - .map(|(object_seq, entry)| async move { - let bytes = io - .get_range(object_seq, entry.segment_offset, entry.segment_len) - .await?; + .map(|(object_seq, entry, fetched)| async move { + let bytes = match fetched { + Some(bytes) => bytes, + None => { + io.get_range(object_seq, entry.segment_offset, entry.segment_len) + .await? + } + }; decode_segment(&bytes, &entry) .and_then(|records| decode_releases(&records[0].payload)) .with_context(|_| InvalidWalObjectSnafu { @@ -3345,10 +3369,11 @@ async fn fetch_footers( let outcomes = futures::stream::iter(objects) .map(|object| async move { match fetch_footer(io, &object).await { - Ok((header, footer)) => Ok(Ok(FetchedObject { + Ok((header, footer, releases)) => Ok(Ok(FetchedObject { object, header, footer, + releases, })), Err(error) if is_not_found(&error) => Ok(Err((object.object_seq, error))), Err(error) => Err(error), @@ -3376,7 +3401,9 @@ fn is_not_found(error: &Error) -> bool { /// Reads the header, trailer and footer of `object` and verifies them: the /// header must carry the sequence of the key, the trailer must be well formed, /// the footer must match the checksum the trailer holds and its segments must -/// tile the object body. +/// tile the object body. Also returns a copy of the segment of release records +/// when the bytes read hold it and it is at most [`MAX_KEPT_RELEASE_BYTES`], +/// which is not decoded here. /// /// A short object is read whole. Otherwise the header and a window at the end /// of the object are read concurrently, and the footer is read separately only @@ -3384,7 +3411,7 @@ fn is_not_found(error: &Error) -> bool { async fn fetch_footer( io: &dyn WalObjectIo, object: &ListedObject, -) -> Result<(Header, Vec)> { +) -> Result<(Header, Vec, Option)> { let ListedObject { object_seq, size, .. } = *object; @@ -3442,7 +3469,18 @@ async fn fetch_footer( }; let footer = decode_footer(&footer, trailer).map_err(invalid)?; verify_segment_ranges(&footer, footer_range.start).map_err(invalid)?; - Ok((header, footer)) + // A copy, so the window is not kept until the releases are decoded. + let releases = footer + .iter() + .find(|entry| entry.region_id == RELEASE_REGION_ID) + .filter(|entry| { + entry.segment_offset >= tail_start as u64 && entry.segment_len <= MAX_KEPT_RELEASE_BYTES + }) + .map(|entry| { + let start = entry.segment_offset as usize - tail_start; + Bytes::copy_from_slice(&tail[start..start + entry.segment_len as usize]) + }); + Ok((header, footer, releases)) } /// Verifies the header and trailer of the object `object_seq` of `object_len` @@ -3968,6 +4006,7 @@ mod tests { object, header: decoded.header, footer: decoded.footer, + releases: None, }); } finish_recovery(FetchedObjects { @@ -9167,6 +9206,206 @@ mod tests { store.stop().await.unwrap(); } + #[tokio::test] + async fn test_store_recovery_takes_release_records_from_the_footer_reads() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + let (one, two, three) = (region(1), region(2), region(3)); + let store = open(object_store.clone(), &eager()).await; + // Object 1 holds entries of regions one and two, object 2 the release + // of region one and object 3 the drop of region three; region two + // holds the boundary at object 1. + store + .append_batch(vec![entry(&store, one, "a1"), entry(&store, two, "b1")]) + .await + .unwrap(); + store.release_region(&provider(one), one).await.unwrap(); + store.obsolete_all(&provider(three), three).await.unwrap(); + store.stop().await.unwrap(); + + // Every object is shorter than the tail window, so recovery reads + // each one whole, once, and decodes the release records from those + // bytes without a range read. + let (io, reads) = RecordingIo::over(object_store); + let store = open_over(io, &eager()).await; + assert_eq!(Some(id(1, 1)), watermark(&store, one)); + assert_eq!(None, watermark(&store, two)); + assert_eq!(Some(EntryId::MAX), watermark(&store, three)); + assert!(reads.lock().unwrap().is_empty()); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_recovery_reads_a_release_segment_outside_the_tail_window() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + // So many release records that their segment starts before the + // window at the end of the object. + let releases = (1..=5_000) + .map(|number| (region(number), id(0, 1))) + .collect::>(); + let record = Record { + region_id: RELEASE_REGION_ID, + entry_id: id(0, 1), + payload: encode_releases(&releases), + }; + put_records(&object_store, 0, &[record]).await; + let (io, reads) = RecordingIo::over(object_store); + let object = io.list().await.unwrap().remove(0); + let (_, footer) = footer_of(&io.get(0).await.unwrap()); + let segment = &footer[0]; + assert!(segment.segment_offset < object.size - RECOVERY_TAIL_WINDOW as u64); + + let store = open_over(io, &eager()).await; + for number in [1, 5_000] { + assert_eq!(Some(id(0, 1)), watermark(&store, region(number))); + } + let mut reads = reads.lock().unwrap().clone(); + reads.sort_unstable(); + assert_eq!( + vec![ + (0, 0, HEADER_LEN as u64), + (0, segment.segment_offset, segment.segment_len), + ( + 0, + object.size - RECOVERY_TAIL_WINDOW as u64, + RECOVERY_TAIL_WINDOW as u64 + ), + ], + reads + ); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_recovery_takes_a_release_segment_from_the_tail_window() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + // A segment longer than the window comes first, so the object is read + // as its header and its tail, and the segment of release records, the + // last one, lies inside the tail but not at its start. + let records = [ + Record { + region_id: region(1), + entry_id: id(0, 1), + payload: Bytes::from(vec![0; RECOVERY_TAIL_WINDOW * 2]), + }, + Record { + region_id: RELEASE_REGION_ID, + entry_id: id(0, 1), + payload: encode_releases(&BTreeMap::from([(region(2), 7)])), + }, + ]; + put_records(&object_store, 0, &records).await; + let (io, reads) = RecordingIo::over(object_store); + let object = io.list().await.unwrap().remove(0); + let tail_start = object.size - RECOVERY_TAIL_WINDOW as u64; + let (_, footer) = footer_of(&io.get(0).await.unwrap()); + let segment = footer + .iter() + .find(|entry| entry.region_id == RELEASE_REGION_ID) + .unwrap(); + assert!(segment.segment_offset > tail_start); + + let store = open_over(io, &eager()).await; + assert_eq!(Some(7), watermark(&store, region(2))); + let mut reads = reads.lock().unwrap().clone(); + reads.sort_unstable(); + assert_eq!( + vec![ + (0, 0, HEADER_LEN as u64), + (0, tail_start, RECOVERY_TAIL_WINDOW as u64) + ], + reads + ); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_recovery_reads_a_release_segment_above_the_kept_size_again() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + // The object is read whole, but its segment of release records is + // larger than recovery keeps. + let releases = (1..=300) + .map(|number| (region(number), id(0, 1))) + .collect::>(); + let record = Record { + region_id: RELEASE_REGION_ID, + entry_id: id(0, 1), + payload: encode_releases(&releases), + }; + put_records(&object_store, 0, &[record]).await; + let (io, reads) = RecordingIo::over(object_store); + let object = io.list().await.unwrap().remove(0); + let (_, footer) = footer_of(&io.get(0).await.unwrap()); + let segment = &footer[0]; + assert!(object.size <= RECOVERY_TAIL_WINDOW as u64); + assert!(segment.segment_len > MAX_KEPT_RELEASE_BYTES); + + let store = open_over(io, &eager()).await; + for number in [1, 300] { + assert_eq!(Some(id(0, 1)), watermark(&store, region(number))); + } + assert_eq!( + vec![(0, segment.segment_offset, segment.segment_len)], + *reads.lock().unwrap() + ); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_recovery_decodes_release_records_only_on_the_chain() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + let io = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); + let release_record = |object_seq, number| Record { + region_id: RELEASE_REGION_ID, + entry_id: id(object_seq, 1), + payload: encode_releases(&BTreeMap::from([(region(number), 1)])), + }; + let damage_release_segment = |bytes: &mut Vec| { + let (_, footer) = footer_of(bytes); + let segment = footer + .iter() + .find(|entry| entry.region_id == RELEASE_REGION_ID) + .unwrap(); + let last = segment.segment_offset + segment.segment_len - 1; + bytes[last as usize] ^= 1; + }; + // Object 0 holds an entry of region one and the release of region + // two. Object 1 releases region three but is off the chain: it links + // to an object 0 of another epoch. + let entry_record = Record { + region_id: region(1), + entry_id: id(0, 1), + payload: Bytes::from_static(b"a1"), + }; + put_records(&object_store, 0, &[entry_record, release_record(0, 2)]).await; + let orphan = Header { + object_seq: 1, + epoch: 1, + prev: Some(ChainLink { + object_seq: 0, + epoch: 2, + }), + reclaim_boundary: 0, + }; + put_object_with_header(&io, orphan, &[release_record(1, 3)]).await; + + // A damaged release segment off the chain is never decoded. + corrupt_object(&object_store, &io.object_path(1), damage_release_segment).await; + let store = open(object_store.clone(), &eager()).await; + assert_eq!(Some(1), watermark(&store, region(2))); + assert_eq!(None, watermark(&store, region(3))); + store.stop().await.unwrap(); + + // A damaged release segment on the chain fails recovery. + let path = io.object_path(0); + corrupt_object(&object_store, &path, damage_release_segment).await; + assert_recovery_rejects(&object_store, &path, "checksum mismatch").await; + } + #[tokio::test] async fn test_store_carries_the_record_of_a_failed_release() { let _serialized = SERIALIZED.lock().await; @@ -9687,6 +9926,123 @@ mod tests { } } + /// Measurements of costs that grow with the prefix, taken through the + /// store. They are ignored by default and print what they measure; run + /// them with `--run-ignored only` and `--no-capture`, preferably with + /// `--cargo-profile release`. + #[allow(clippy::print_stdout)] + mod measurements { + use super::*; + + /// Object access that counts the reads of whole objects and of byte + /// ranges. + struct CountingIo { + inner: ObjectStoreIo, + gets: AtomicUsize, + range_gets: AtomicUsize, + } + + #[async_trait::async_trait] + impl WalObjectIo for CountingIo { + async fn put_if_absent(&self, object_seq: u64, content: Bytes) -> Result { + self.inner.put_if_absent(object_seq, content).await + } + + async fn get(&self, object_seq: u64) -> Result { + self.gets.fetch_add(1, Ordering::Relaxed); + self.inner.get(object_seq).await + } + + async fn get_range(&self, object_seq: u64, offset: u64, len: u64) -> Result { + self.range_gets.fetch_add(1, Ordering::Relaxed); + self.inner.get_range(object_seq, offset, len).await + } + + async fn delete(&self, object_seq: u64) -> Result<()> { + self.inner.delete(object_seq).await + } + + async fn list(&self) -> Result> { + self.inner.list().await + } + + fn object_path(&self, object_seq: u64) -> String { + self.inner.object_path(object_seq) + } + + fn storage(&self) -> String { + self.inner.storage() + } + } + + /// The object store a measurement of recovery runs on: a fresh root + /// in the bucket the `GT_S3_*` variables name when they are set, an + /// in-memory store otherwise. + fn measured_object_store() -> (ObjectStore, &'static str) { + let Ok(bucket) = std::env::var("GT_S3_BUCKET") else { + return (memory_store(), "memory"); + }; + let variable = |name: &str| std::env::var(name).unwrap(); + let mut builder = object_store::services::S3::default() + .root(&format!("measurements/{}", uuid::Uuid::new_v4())) + .bucket(&bucket) + .region(&variable("GT_S3_REGION")) + .access_key_id(&variable("GT_S3_ACCESS_KEY_ID")) + .secret_access_key(&variable("GT_S3_ACCESS_KEY")); + if let Ok(endpoint) = std::env::var("GT_S3_ENDPOINT_URL") { + builder = builder.endpoint(&endpoint); + } + (ObjectStore::new(builder).unwrap(), "s3") + } + + /// Counts the reads of a recovery over a prefix where every dropped + /// region left its release record in an object of its own, while one + /// region still holds the boundary at the object of the first + /// entries, and times the recovery. + #[tokio::test] + #[ignore] + async fn measure_recovery_of_dropped_regions() { + let _serialized = SERIALIZED.lock().await; + for regions in [1_000, 4_000] { + let (object_store, storage) = measured_object_store(); + let store = open(object_store.clone(), &eager()).await; + let entries = (1..=regions + 1) + .map(|number| entry(&store, region(number), "a")) + .collect(); + store.append_batch(entries).await.unwrap(); + for number in 1..=regions { + let region_id = region(number); + store + .obsolete_all(&provider(region_id), region_id) + .await + .unwrap(); + } + store.stop().await.unwrap(); + + let io = Arc::new(CountingIo { + inner: ObjectStoreIo::new(object_store, PREFIX).unwrap(), + gets: AtomicUsize::new(0), + range_gets: AtomicUsize::new(0), + }); + let objects = io.list().await.unwrap().len(); + let started_at = Instant::now(); + let store = open_over(io.clone(), &eager()).await; + let elapsed = started_at.elapsed(); + for number in 1..=regions { + assert_eq!(Some(EntryId::MAX), watermark(&store, region(number))); + } + assert_eq!(None, watermark(&store, region(regions + 1))); + println!( + "recovery of dropped regions: storage={storage} regions={regions} objects={objects} gets={} range_gets={} recovery_ms={:.3}", + io.gets.load(Ordering::Relaxed), + io.range_gets.load(Ordering::Relaxed), + elapsed.as_secs_f64() * 1e3, + ); + store.stop().await.unwrap(); + } + } + } + /// Drives every outcome of a conditional create and every failure of the /// object store access through the store, one case per fault, and checks /// that every row of the *Failure matrix* has one. From 095c288f16f4ee359a515c9a30cc355d6f6f8310 Mon Sep 17 00:00:00 2001 From: DeviousCardi <115358213+DeviousCardi@users.noreply.github.com> Date: Fri, 16 Oct 2026 13:17:55 +0000 Subject: [PATCH 031/239] fix(promql): honor @ modifier on subqueries (#9446) * fix(promql): honor @ modifier on subqueries `prom_subquery_expr_to_plan` never read the `at` field of a subquery, so `sum_over_time(m[20s:10s] @ 60)` was evaluated like `m[20s:10s]` and followed the outer evaluation grid. Handle `@` the way Prometheus does: fold the anchor into the subquery offset (`offset + (eval_start - anchor)`, as `setOffsetForAtModifier` does), evaluate the subquery once at the start of the evaluation, and replay the folded window over the outer grid, reusing the anchored range selector machinery. A call whose range argument is an anchored subquery is promoted as step-invariant, like one over an anchored range selector. `@ start()`/`@ end()` and `offset` combine as for selectors, and an unrepresentable anchor is rejected with the existing error. Closes #9411 Signed-off-by: Aarav * test(promql): add sqlness cases for @ on subqueries Cover the #9411 reproducer end to end: `@ `, `@ start()`/`@ end()`, `@` combined with `offset` in either order, an instant query away from the anchor, `*_over_time` functions and `rate` over an anchored subquery, a computed inner expression, and an aggregation above the anchored call. Two series per table so a dropped or merged series shows up as a wrong row. Refs #9411 Signed-off-by: Aarav --------- Signed-off-by: Aarav --- src/query/src/promql/planner.rs | 36 ++- src/query/src/promql/planner/at_modifier.rs | 13 +- src/query/src/promql/planner/test.rs | 137 +++++++++++ .../common/promql/subquery_at_modifier.result | 229 ++++++++++++++++++ .../common/promql/subquery_at_modifier.sql | 92 +++++++ 5 files changed, 494 insertions(+), 13 deletions(-) create mode 100644 tests/cases/standalone/common/promql/subquery_at_modifier.result create mode 100644 tests/cases/standalone/common/promql/subquery_at_modifier.sql diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index afe52c7b560..6cb0d671196 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -577,7 +577,7 @@ impl PromPlanner { range, step, offset, - .. + at, } = subquery_expr; // Shift the child window back; `RangeManipulate` restores the outer timeline. @@ -603,10 +603,18 @@ impl PromPlanner { end: current_end, } ); + // `@` makes the subquery step-invariant: like Prometheus (`setOffsetForAtModifier`), fold + // the anchor into the offset and evaluate once at `current_start`, then replay the result + // over the outer grid. See [`Self::at_modifier_offset`]. + let at_offset = self.at_modifier_offset(at, offset)?; + let (eval_end, offset_ms) = match at_offset { + Some(at_offset) => (current_start, at_offset), + None => (current_end, offset_ms), + }; // An empty child needs only its schema; plan it over the valid caller window. let child_window = Self::subquery_child_window( current_start, - current_end, + eval_end, current_interval, offset_ms, range.as_millis() as Millisecond, @@ -707,23 +715,31 @@ impl PromPlanner { }; let manipulate = RangeManipulate::new( - self.ctx.start, - self.ctx.end, + current_start, + eval_end, self.ctx.interval, offset_ms, range_ms, - time_index_column, + time_index_column.clone(), self.ctx.field_columns.clone(), divide_plan, ) .context(DataFusionPlanningSnafu)?; - // The payload timestamps are shifted by the subquery offset; see - // [`Self::create_range_eval_ts_expr`]. + // The payload timestamps are shifted by the subquery offset (the rewritten one under `@`); + // see [`Self::create_range_eval_ts_expr`]. self.ctx.range_fold_offset = Some(offset_ms); - - Ok(LogicalPlan::Extension(Extension { + let manipulate = LogicalPlan::Extension(Extension { node: Arc::new(manipulate), - })) + }); + + // An anchored subquery was folded once, at `current_start`; report that window at every + // step of the outer grid, like an anchored range selector. + Ok(match at_offset { + Some(_) => { + self.replay_over_grid(manipulate, current_start, current_end, time_index_column) + } + None => manipulate, + }) } async fn prom_aggr_expr_to_plan( diff --git a/src/query/src/promql/planner/at_modifier.rs b/src/query/src/promql/planner/at_modifier.rs index 545f1714a22..81b03d29ed3 100644 --- a/src/query/src/promql/planner/at_modifier.rs +++ b/src/query/src/promql/planner/at_modifier.rs @@ -21,7 +21,8 @@ use datafusion::logical_expr::{Extension, LogicalPlan, LogicalPlanBuilder}; use datafusion::prelude::{Column, Expr as DfExpr}; use promql::extension_plan::{InstantManipulate, Millisecond, SeriesDivide}; use promql_parser::parser::{ - AtModifier, Call, Expr as PromExpr, MatrixSelector, Offset, ParenExpr, + AtModifier, Call, Expr as PromExpr, MatrixSelector, Offset, ParenExpr, SubqueryExpr, + VectorSelector, }; use snafu::{OptionExt, ResultExt, ensure}; @@ -138,8 +139,14 @@ impl PromPlanner { arg = expr; } match arg { - // The window is pinned by `@`, so every step folds the same samples. - PromExpr::MatrixSelector(MatrixSelector { vs, .. }) if vs.at.is_some() => { + // The window is pinned by `@`, so every step folds the same samples. That holds + // for a subquery with `@` too, which Prometheus treats as step-invariant whatever + // its inner expression is. + PromExpr::MatrixSelector(MatrixSelector { + vs: VectorSelector { at: Some(_), .. }, + .. + }) + | PromExpr::Subquery(SubqueryExpr { at: Some(_), .. }) => { if anchored_range { return false; } diff --git a/src/query/src/promql/planner/test.rs b/src/query/src/promql/planner/test.rs index ae40640b50b..be3140a6a52 100644 --- a/src/query/src/promql/planner/test.rs +++ b/src/query/src/promql/planner/test.rs @@ -5754,6 +5754,143 @@ async fn subquery_rejects_a_child_window_out_of_the_representable_range() { assert_eq!(err.status_code(), StatusCode::InvalidArguments, "{err}"); } +/// `@` on a subquery pins its evaluation instant, like `@` on a range selector: Prometheus folds +/// the anchor into the subquery offset (`setOffsetForAtModifier`) and treats the subquery as +/// step-invariant, so every step reports the window `(anchor - offset - range, anchor - offset]`. +/// See . +/// +/// The probe samples are at 0s..60s every 10s with values 1..7, on the 10s subquery grid. +#[tokio::test] +async fn subquery_at_modifier_pins_the_evaluation_instant() { + // The issue reproducer: `@ 60` folds `(40s, 60s]` -- the points 50s and 60s, 6 + 7 -- at every + // step, the same as the range selector form `m[20s] @ 60`. + for query in [ + "sum_over_time(subquery_offset_probe[20s] @ 60)", + "sum_over_time(subquery_offset_probe[20s:10s] @ 60)", + "sum_over_time(subquery_offset_probe[20s:10s] @ end())", + ] { + assert_eq!( + run_subquery_offset_probe(query, 30, 60, 30).await, + vec![(30_000, 13.0), (60_000, 13.0)], + "{query}" + ); + } + // `@ start()` folds `(10s, 30s]` -- 3 + 4 -- at every step. + assert_eq!( + run_subquery_offset_probe( + "sum_over_time(subquery_offset_probe[20s:10s] @ start())", + 30, + 60, + 30 + ) + .await, + vec![(30_000, 7.0), (60_000, 7.0)], + ); + // `offset` is applied relative to the anchor: `@ 60 offset 30s` folds `(10s, 30s]`, in either + // order of the modifiers. + for query in [ + "sum_over_time(subquery_offset_probe[20s:10s] @ 60 offset 30s)", + "sum_over_time(subquery_offset_probe[20s:10s] offset 30s @ 60)", + ] { + assert_eq!( + run_subquery_offset_probe(query, 30, 60, 30).await, + vec![(30_000, 7.0), (60_000, 7.0)], + "{query}" + ); + } + // An instant query evaluated away from the anchor still folds the anchored window. + assert_eq!( + run_subquery_offset_probe( + "sum_over_time(subquery_offset_probe[20s:10s] @ 60)", + 0, + 0, + 1 + ) + .await, + vec![(0, 13.0)], + ); + // An anchor whose window starts before the first sample: `@ 0` folds `(-20s, 0s]`, where only + // the point 0s has a sample. + assert_eq!( + run_subquery_offset_probe( + "sum_over_time(subquery_offset_probe[20s:10s] @ 0)", + 30, + 60, + 30 + ) + .await, + vec![(30_000, 1.0), (60_000, 1.0)], + ); + // `predict_linear` is not promoted as a whole, so it consumes the window replayed by the + // subquery itself; that window holds the same points as the range selector form. + let anchored = run_subquery_offset_probe( + "predict_linear(subquery_offset_probe[20s:10s] @ 60, 0)", + 30, + 60, + 30, + ) + .await; + assert_eq!(anchored.len(), 2, "{anchored:?}"); + assert_eq!( + anchored, + run_subquery_offset_probe( + "predict_linear(subquery_offset_probe[20s] @ 60, 0)", + 30, + 60, + 30 + ) + .await, + ); +} + +/// An anchored subquery evaluates its child over the anchored window only, not over the whole +/// outer evaluation range, and an anchor that cannot be represented in milliseconds is rejected. +#[tokio::test] +async fn subquery_at_modifier_scans_only_the_anchored_window() { + let plan = plan_subquery_offset_probe( + "sum_over_time(subquery_offset_probe[20s:10s] @ 60)", + 0, + 100_000, + 30, + ) + .await + .unwrap(); + let plan_str = plan.display_indent_schema().to_string(); + // The child grid is the points 50s and 60s of `(40s, 60s]`. + assert!( + plan_str.contains("PromInstantManipulate: range=[50000..60000]"), + "{plan_str}" + ); + // The subquery is folded once, at the start of the evaluation, and only the outer replay spans + // the evaluation range. + assert!( + plan_str.contains("PromRangeManipulate: req range=[0..0]"), + "{plan_str}" + ); + assert!( + plan_str.contains("subquery_offset_probe.timestamp <= TimestampMillisecond(60000, None)"), + "the child must not be evaluated over the outer range:\n{plan_str}" + ); + + #[cfg(not(windows))] + { + let err = plan_subquery_offset_probe( + "sum_over_time(subquery_offset_probe[20s:10s] @ 1e16)", + 0, + 60, + 30, + ) + .await + .unwrap_err(); + assert!( + err.to_string() + .contains("Timestamp out of range for the `@` modifier"), + "{err}" + ); + assert_eq!(err.status_code(), StatusCode::InvalidArguments); + } +} + #[tokio::test] async fn test_hash_join() { let mut eval_stmt = EvalStmt { diff --git a/tests/cases/standalone/common/promql/subquery_at_modifier.result b/tests/cases/standalone/common/promql/subquery_at_modifier.result new file mode 100644 index 00000000000..02dc4111c8d --- /dev/null +++ b/tests/cases/standalone/common/promql/subquery_at_modifier.result @@ -0,0 +1,229 @@ +-- Tests for the PromQL `@` modifier on subqueries. +-- +-- Prometheus folds the `@` of a subquery into its offset (`setOffsetForAtModifier`) and treats +-- the subquery as step-invariant: every evaluation step reports the window +-- `(anchor - offset - range, anchor - offset]`, sampled on the absolute multiples of the subquery +-- step, exactly like `@` on a range selector. `@ start()` / `@ end()` anchor at the start/end of +-- the statement's evaluation range. The output timestamps still follow the evaluation grid. +-- +-- Two series with different values, so that a dropped or merged series shows up as a wrong row. +-- Sample timestamps are in milliseconds, while `@` and `TQL EVAL` timestamps are in seconds. +-- The samples lie on the 10s subquery step grid. +CREATE TABLE subquery_at_total ( + ts TIMESTAMP TIME INDEX, + host STRING PRIMARY KEY, + val DOUBLE, +); + +Affected Rows: 0 + +INSERT INTO subquery_at_total VALUES + (0, 'a', 1), + (10000, 'a', 2), + (20000, 'a', 3), + (30000, 'a', 4), + (40000, 'a', 5), + (50000, 'a', 6), + (60000, 'a', 7), + (0, 'b', 10), + (10000, 'b', 20), + (20000, 'b', 30), + (30000, 'b', 40), + (40000, 'b', 50), + (50000, 'b', 60), + (60000, 'b', 70); + +Affected Rows: 14 + +-- The range selector form folds (40s, 60s] at every step: 6 + 7 = 13 (and 130 for 'b'). +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (30, 60, '30s') sum_over_time(subquery_at_total[20s] @ 60); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:00:30 | 13.0 | a | +| 1970-01-01T00:00:30 | 130.0 | b | +| 1970-01-01T00:01:00 | 13.0 | a | +| 1970-01-01T00:01:00 | 130.0 | b | ++---------------------+----------------------------------+------+ + +-- The subquery with `@ 60` must report the same window at every step. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (30, 60, '30s') sum_over_time(subquery_at_total[20s:10s] @ 60); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:00:30 | 13.0 | a | +| 1970-01-01T00:00:30 | 130.0 | b | +| 1970-01-01T00:01:00 | 13.0 | a | +| 1970-01-01T00:01:00 | 130.0 | b | ++---------------------+----------------------------------+------+ + +-- Baseline without `@`: the window follows the grid, 3 + 4 at 30s and 6 + 7 at 60s. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (30, 60, '30s') sum_over_time(subquery_at_total[20s:10s]); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:00:30 | 7.0 | a | +| 1970-01-01T00:00:30 | 70.0 | b | +| 1970-01-01T00:01:00 | 13.0 | a | +| 1970-01-01T00:01:00 | 130.0 | b | ++---------------------+----------------------------------+------+ + +-- `@ end()` anchors at 60s: 13 at both steps. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (30, 60, '30s') sum_over_time(subquery_at_total[20s:10s] @ end()); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:00:30 | 13.0 | a | +| 1970-01-01T00:00:30 | 130.0 | b | +| 1970-01-01T00:01:00 | 13.0 | a | +| 1970-01-01T00:01:00 | 130.0 | b | ++---------------------+----------------------------------+------+ + +-- `@ start()` anchors at 30s: (10s, 30s] holds 3 + 4 = 7 at both steps. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (30, 60, '30s') sum_over_time(subquery_at_total[20s:10s] @ start()); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:00:30 | 7.0 | a | +| 1970-01-01T00:00:30 | 70.0 | b | +| 1970-01-01T00:01:00 | 7.0 | a | +| 1970-01-01T00:01:00 | 70.0 | b | ++---------------------+----------------------------------+------+ + +-- `offset` applies relative to the anchor, in either modifier order: `@ 60 offset 30s` folds +-- (10s, 30s] = 7 at both steps. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (30, 60, '30s') sum_over_time(subquery_at_total[20s:10s] @ 60 offset 30s); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:00:30 | 7.0 | a | +| 1970-01-01T00:00:30 | 70.0 | b | +| 1970-01-01T00:01:00 | 7.0 | a | +| 1970-01-01T00:01:00 | 70.0 | b | ++---------------------+----------------------------------+------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (30, 60, '30s') sum_over_time(subquery_at_total[20s:10s] offset 30s @ 60); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:00:30 | 7.0 | a | +| 1970-01-01T00:00:30 | 70.0 | b | +| 1970-01-01T00:01:00 | 7.0 | a | +| 1970-01-01T00:01:00 | 70.0 | b | ++---------------------+----------------------------------+------+ + +-- An instant query away from the anchor still folds the anchored window: 13. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') sum_over_time(subquery_at_total[20s:10s] @ 60); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:00:00 | 13.0 | a | +| 1970-01-01T00:00:00 | 130.0 | b | ++---------------------+----------------------------------+------+ + +-- Other `*_over_time` functions over the anchored window (40s, 60s]: points 50s and 60s. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (30, 60, '30s') max_over_time(subquery_at_total[20s:10s] @ 60); + ++---------------------+----------------------------------+------+ +| ts | prom_max_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:00:30 | 7.0 | a | +| 1970-01-01T00:00:30 | 70.0 | b | +| 1970-01-01T00:01:00 | 7.0 | a | +| 1970-01-01T00:01:00 | 70.0 | b | ++---------------------+----------------------------------+------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (30, 60, '30s') count_over_time(subquery_at_total[20s:10s] @ 60); + ++---------------------+------------------------------------+------+ +| ts | prom_count_over_time(ts_range,val) | host | ++---------------------+------------------------------------+------+ +| 1970-01-01T00:00:30 | 2.0 | a | +| 1970-01-01T00:00:30 | 2.0 | b | +| 1970-01-01T00:01:00 | 2.0 | a | +| 1970-01-01T00:01:00 | 2.0 | b | ++---------------------+------------------------------------+------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (30, 60, '30s') avg_over_time(subquery_at_total[20s:10s] @ 60); + ++---------------------+----------------------------------+------+ +| ts | prom_avg_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:00:30 | 6.5 | a | +| 1970-01-01T00:00:30 | 65.0 | b | +| 1970-01-01T00:01:00 | 6.5 | a | +| 1970-01-01T00:01:00 | 65.0 | b | ++---------------------+----------------------------------+------+ + +-- `rate` over (30s, 60s] -- points 40s, 50s, 60s -- extrapolates from the anchored window, not +-- from the evaluation step, so the subquery form matches the range selector form at every step. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (30, 60, '30s') rate(subquery_at_total[30s] @ 60); + ++---------------------+-----------------------------------------+------+ +| ts | prom_rate(ts_range,val,ts,Int64(30000)) | host | ++---------------------+-----------------------------------------+------+ +| 1970-01-01T00:00:30 | 0.1 | a | +| 1970-01-01T00:00:30 | 1.0 | b | +| 1970-01-01T00:01:00 | 0.1 | a | +| 1970-01-01T00:01:00 | 1.0 | b | ++---------------------+-----------------------------------------+------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (30, 60, '30s') rate(subquery_at_total[30s:10s] @ 60); + ++---------------------+-----------------------------------------+------+ +| ts | prom_rate(ts_range,val,ts,Int64(30000)) | host | ++---------------------+-----------------------------------------+------+ +| 1970-01-01T00:00:30 | 0.1 | a | +| 1970-01-01T00:00:30 | 1.0 | b | +| 1970-01-01T00:01:00 | 0.1 | a | +| 1970-01-01T00:01:00 | 1.0 | b | ++---------------------+-----------------------------------------+------+ + +-- A computed inner expression and an aggregation above the anchored call. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (30, 60, '30s') sum_over_time((subquery_at_total * 2)[20s:10s] @ 60); + ++---------------------+-----------------------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val * Float64(2)) | host | ++---------------------+-----------------------------------------------+------+ +| 1970-01-01T00:00:30 | 26.0 | a | +| 1970-01-01T00:00:30 | 260.0 | b | +| 1970-01-01T00:01:00 | 26.0 | a | +| 1970-01-01T00:01:00 | 260.0 | b | ++---------------------+-----------------------------------------------+------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (30, 60, '30s') sum(sum_over_time(subquery_at_total[20s:10s] @ 60)); + ++---------------------+---------------------------------------+ +| ts | sum(prom_sum_over_time(ts_range,val)) | ++---------------------+---------------------------------------+ +| 1970-01-01T00:00:30 | 143.0 | +| 1970-01-01T00:01:00 | 143.0 | ++---------------------+---------------------------------------+ + +DROP TABLE subquery_at_total; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/promql/subquery_at_modifier.sql b/tests/cases/standalone/common/promql/subquery_at_modifier.sql new file mode 100644 index 00000000000..5ae70098664 --- /dev/null +++ b/tests/cases/standalone/common/promql/subquery_at_modifier.sql @@ -0,0 +1,92 @@ +-- Tests for the PromQL `@` modifier on subqueries. +-- +-- Prometheus folds the `@` of a subquery into its offset (`setOffsetForAtModifier`) and treats +-- the subquery as step-invariant: every evaluation step reports the window +-- `(anchor - offset - range, anchor - offset]`, sampled on the absolute multiples of the subquery +-- step, exactly like `@` on a range selector. `@ start()` / `@ end()` anchor at the start/end of +-- the statement's evaluation range. The output timestamps still follow the evaluation grid. +-- +-- Two series with different values, so that a dropped or merged series shows up as a wrong row. +-- Sample timestamps are in milliseconds, while `@` and `TQL EVAL` timestamps are in seconds. +-- The samples lie on the 10s subquery step grid. + +CREATE TABLE subquery_at_total ( + ts TIMESTAMP TIME INDEX, + host STRING PRIMARY KEY, + val DOUBLE, +); + +INSERT INTO subquery_at_total VALUES + (0, 'a', 1), + (10000, 'a', 2), + (20000, 'a', 3), + (30000, 'a', 4), + (40000, 'a', 5), + (50000, 'a', 6), + (60000, 'a', 7), + (0, 'b', 10), + (10000, 'b', 20), + (20000, 'b', 30), + (30000, 'b', 40), + (40000, 'b', 50), + (50000, 'b', 60), + (60000, 'b', 70); + +-- The range selector form folds (40s, 60s] at every step: 6 + 7 = 13 (and 130 for 'b'). +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (30, 60, '30s') sum_over_time(subquery_at_total[20s] @ 60); + +-- The subquery with `@ 60` must report the same window at every step. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (30, 60, '30s') sum_over_time(subquery_at_total[20s:10s] @ 60); + +-- Baseline without `@`: the window follows the grid, 3 + 4 at 30s and 6 + 7 at 60s. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (30, 60, '30s') sum_over_time(subquery_at_total[20s:10s]); + +-- `@ end()` anchors at 60s: 13 at both steps. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (30, 60, '30s') sum_over_time(subquery_at_total[20s:10s] @ end()); + +-- `@ start()` anchors at 30s: (10s, 30s] holds 3 + 4 = 7 at both steps. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (30, 60, '30s') sum_over_time(subquery_at_total[20s:10s] @ start()); + +-- `offset` applies relative to the anchor, in either modifier order: `@ 60 offset 30s` folds +-- (10s, 30s] = 7 at both steps. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (30, 60, '30s') sum_over_time(subquery_at_total[20s:10s] @ 60 offset 30s); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (30, 60, '30s') sum_over_time(subquery_at_total[20s:10s] offset 30s @ 60); + +-- An instant query away from the anchor still folds the anchored window: 13. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') sum_over_time(subquery_at_total[20s:10s] @ 60); + +-- Other `*_over_time` functions over the anchored window (40s, 60s]: points 50s and 60s. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (30, 60, '30s') max_over_time(subquery_at_total[20s:10s] @ 60); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (30, 60, '30s') count_over_time(subquery_at_total[20s:10s] @ 60); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (30, 60, '30s') avg_over_time(subquery_at_total[20s:10s] @ 60); + +-- `rate` over (30s, 60s] -- points 40s, 50s, 60s -- extrapolates from the anchored window, not +-- from the evaluation step, so the subquery form matches the range selector form at every step. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (30, 60, '30s') rate(subquery_at_total[30s] @ 60); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (30, 60, '30s') rate(subquery_at_total[30s:10s] @ 60); + +-- A computed inner expression and an aggregation above the anchored call. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (30, 60, '30s') sum_over_time((subquery_at_total * 2)[20s:10s] @ 60); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (30, 60, '30s') sum(sum_over_time(subquery_at_total[20s:10s] @ 60)); + +DROP TABLE subquery_at_total; From baa5630119341a8ffc173dd49571917c09df5123 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Fri, 16 Oct 2026 14:27:04 +0000 Subject: [PATCH 032/239] perf(log-store): keep the object store WAL reclaim boundary independent of the region count (#9644) * perf(log-store): keep the object store WAL reclaim boundary independent of the region count The WAL actor recomputed the reclaim boundary after every watermark update, release and indexed object by scanning the footers of the objects from the previous boundary, checking each segment against the obsolete watermarks. Each scan started at the first segment of the boundary object, so when one object holds segments of many regions whose watermarks advance one region at a time, every update inspected the segments already known to be obsolete again, about R^2/2 footer entries in total for R regions. The catalog now keeps the obsolete watermarks itself and counts the live segments of every object. A watermark move visits only the segments of its region that it newly reaches, through the per-region index, and the boundary is the first object at or above the floor that still holds a live segment. The boundary is unchanged: segments of release records never hold it, a region without a watermark does, a watermark of EntryId::MAX releases every segment of its region, and it never moves below the floor. The store records each watermark in the catalog together with the one reads use. The per-region index is now keyed by the largest entry id of each segment, which is also object sequence order because the entry ranges of a region increase strictly with the sequence. Indexing an object already enforced that; it now checks the neighbours by entry id instead of by sequence, which accepts and rejects the same objects. A property test checks the boundary against a scan of the footers over random appends, watermark moves, releases, drops, collections and recoveries, and ignored measurement tests time watermark updates and tail reads through the store. Signed-off-by: jeremyhi * perf(log-store): seek the object store WAL objects of a region by entry id A read found the objects of its region by filtering every object of the region that the catalog retains, so a read of the last few objects cost time in proportion to the whole retained history of the region, under the catalog read lock. The per-region index is keyed by the largest entry id of each segment, and the entry ranges of a region increase strictly with the object sequence, which the catalog enforces when it indexes an object, whatever entry ids the objects carry. So the lookup seeks to the first segment that ends at or after the start of the range and stops at the first that starts after its end, and the result is still in sequence order. The property test of the catalog now also checks range lookups against a scan of the footers. Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- src/log-store/src/object_store_wal/catalog.rs | 454 ++++++++++++++---- src/log-store/src/object_store_wal/store.rs | 203 +++++++- 2 files changed, 532 insertions(+), 125 deletions(-) diff --git a/src/log-store/src/object_store_wal/catalog.rs b/src/log-store/src/object_store_wal/catalog.rs index 61b4d2c09cc..49567db3060 100644 --- a/src/log-store/src/object_store_wal/catalog.rs +++ b/src/log-store/src/object_store_wal/catalog.rs @@ -15,7 +15,7 @@ //! In-memory index over the footers of the objects of one WAL prefix. use std::collections::{BTreeMap, HashMap}; -use std::ops::Bound::{Excluded, Unbounded}; +use std::ops::Bound::{Excluded, Included, Unbounded}; use snafu::{OptionExt, ensure}; use store_api::logstore::EntryId; @@ -28,14 +28,24 @@ use crate::object_store_wal::batch::{OBJECT_SEQ_LIMIT, sequence_floor}; use crate::object_store_wal::format::{FooterEntry, RELEASE_REGION_ID}; /// Indexes objects by sequence and, per region, the objects that hold entries -/// of that region. +/// of that region, and counts the live segments of every object under the +/// obsolete watermarks. #[derive(Debug, Default)] pub(crate) struct ObjectCatalog { objects: BTreeMap>, - regions: BTreeMap>, + /// The objects that hold a segment of the region, as the sequence and + /// the footer entry of the segment, keyed by the largest entry id of the + /// segment. The entry ranges of a region increase strictly with the + /// object sequence, so this is also sequence order. + regions: BTreeMap>, /// Largest entry id ever indexed per region, kept after the object that /// held it was removed. max_entry_ids: HashMap, + /// The obsolete watermark per region, see [`obsolete`](Self::obsolete). + watermarks: HashMap, + /// The number of live segments of every indexed object that holds one, + /// see [`reclaim_boundary`](Self::reclaim_boundary). + live_segments: BTreeMap, } impl ObjectCatalog { @@ -84,49 +94,46 @@ impl ObjectCatalog { ); // Validate every region before mutating either index so insertion is atomic. + // While the indexed ranges of a region increase with the sequence, a + // segment can only conflict with the segments right below and right + // above its largest entry id: each must lie on the side of it that + // its sequence is on. for entry in &footer { let Some(region_objects) = self.regions.get(&entry.region_id) else { continue; }; - if let Some((&previous_seq, previous)) = region_objects.range(..object_seq).next_back() - { + let below = region_objects.range(..entry.max_entry_id).next_back(); + let above = region_objects.range(entry.max_entry_id..).next(); + for (_, (neighbor_seq, neighbor)) in below.into_iter().chain(above) { + let ((lower_seq, lower), (upper_seq, upper)) = if *neighbor_seq < object_seq { + ((*neighbor_seq, neighbor), (object_seq, entry)) + } else { + ((object_seq, entry), (*neighbor_seq, neighbor)) + }; ensure!( - previous.max_entry_id < entry.min_entry_id, + lower.max_entry_id < upper.min_entry_id, CorruptedWalObjectSnafu { reason: out_of_order_reason( entry.region_id, - previous_seq, - previous.max_entry_id, - object_seq, - entry.min_entry_id - ), - } - ); - } - if let Some((&next_seq, next)) = region_objects - .range((Excluded(object_seq), Unbounded)) - .next() - { - ensure!( - entry.max_entry_id < next.min_entry_id, - CorruptedWalObjectSnafu { - reason: out_of_order_reason( - entry.region_id, - object_seq, - entry.max_entry_id, - next_seq, - next.min_entry_id + lower_seq, + lower.max_entry_id, + upper_seq, + upper.min_entry_id ), } ); } } + let live = footer.iter().filter(|entry| self.is_live(entry)).count(); + if live > 0 { + self.live_segments.insert(object_seq, live); + } for entry in &footer { self.regions .entry(entry.region_id) .or_default() - .insert(object_seq, entry.clone()); + .insert(entry.max_entry_id, (object_seq, entry.clone())); self.max_entry_ids .entry(entry.region_id) .and_modify(|max| *max = (*max).max(entry.max_entry_id)) @@ -141,9 +148,10 @@ impl ObjectCatalog { /// largest entry id of every region it held is kept. pub(crate) fn remove_object(&mut self, object_seq: u64) -> Option> { let footer = self.objects.remove(&object_seq)?; + self.live_segments.remove(&object_seq); for entry in &footer { if let Some(objects) = self.regions.get_mut(&entry.region_id) { - objects.remove(&object_seq); + objects.remove(&entry.max_entry_id); if objects.is_empty() { self.regions.remove(&entry.region_id); } @@ -157,34 +165,54 @@ impl ObjectCatalog { self.objects.contains_key(&object_seq) } - /// Returns the reclaim boundary the indexed objects and the obsolete - /// watermarks allow, given that every object below `from` holds no live - /// segment: the first indexed object at or above `from` that holds a - /// live segment, or the sequence after the last indexed object when none - /// does, and never less than `from`. - /// - /// A segment is live when its maximum entry id is above the watermark of - /// its region, or its region has no watermark; a region whose watermark - /// is [`EntryId::MAX`] holds no live segment, and neither does + /// Moves the obsolete watermark of `region_id` up to `entry_id`; it never + /// moves down. The segments of the region whose largest entry id the + /// watermark reaches stop being live, so a move costs a lookup per + /// segment it reaches, however many segments their objects hold. + pub(crate) fn obsolete(&mut self, region_id: RegionId, entry_id: EntryId) { + let previous = self.watermarks.get(®ion_id).copied(); + if previous.is_some_and(|previous| previous >= entry_id) { + return; + } + self.watermarks.insert(region_id, entry_id); + if region_id == RELEASE_REGION_ID { + return; + } + let Some(objects) = self.regions.get(®ion_id) else { + return; + }; + let reached = (previous.map_or(Unbounded, Excluded), Included(entry_id)); + for (object_seq, _) in objects.range(reached).map(|(_, object)| object) { + if let Some(live) = self.live_segments.get_mut(object_seq) { + *live -= 1; + if *live == 0 { + self.live_segments.remove(object_seq); + } + } + } + } + + /// Returns whether a segment is live: its region is not /// [`RELEASE_REGION_ID`], whose records recovery reads only from the - /// objects it indexes. Watermarks never move down - /// and a region first appears in an object above the boundary, so the - /// scan starts at the previous boundary and never reads an object twice - /// once the boundary has passed it. - pub(crate) fn reclaim_boundary( - &self, - from: u64, - obsolete_entry_ids: &HashMap, - ) -> u64 { - let live = self.objects.range(from..).find(|(_, footer)| { - footer.iter().any(|entry| { - entry.region_id != RELEASE_REGION_ID - && obsolete_entry_ids - .get(&entry.region_id) - .is_none_or(|obsolete| entry.max_entry_id > *obsolete) - }) - }); - match live { + /// objects it indexes, and either has no watermark or a watermark below + /// the largest entry id of the segment. A region whose watermark is + /// [`EntryId::MAX`] holds no live segment. + fn is_live(&self, entry: &FooterEntry) -> bool { + entry.region_id != RELEASE_REGION_ID + && self + .watermarks + .get(&entry.region_id) + .is_none_or(|watermark| entry.max_entry_id > *watermark) + } + + /// Returns the reclaim boundary the indexed objects and the obsolete + /// watermarks allow: the first indexed object at or above `from` that + /// holds a live segment, see [`is_live`](Self::is_live), or the sequence + /// after the last indexed object when none does, and never less than + /// `from`. The objects holding a live segment are tracked as objects are + /// indexed and watermarks move, so the boundary is a lookup. + pub(crate) fn reclaim_boundary(&self, from: u64) -> u64 { + match self.live_segments.range(from..).next() { Some((&object_seq, _)) => object_seq, None => self .objects @@ -196,7 +224,11 @@ impl ObjectCatalog { } /// Returns the objects that hold entries of `region_id` overlapping - /// `start_entry_id..=end_entry_id`, ordered by object sequence. + /// `start_entry_id..=end_entry_id`, ordered by object sequence. The + /// segments of a region are ordered by entry id, so the lookup seeks to + /// the first segment that ends at or after `start_entry_id` and stops + /// before the first that starts after `end_entry_id`, whatever the number + /// of objects the region holds. pub(crate) fn objects_for_entry_range( &self, region_id: RegionId, @@ -216,11 +248,9 @@ impl ObjectCatalog { return Ok(Vec::new()); }; Ok(objects - .iter() - .filter(|(_, entry)| { - entry.max_entry_id >= start_entry_id && entry.min_entry_id <= end_entry_id - }) - .map(|(&object_seq, entry)| (object_seq, entry)) + .range(start_entry_id..) + .map(|(_, (object_seq, entry))| (*object_seq, entry)) + .take_while(|(_, entry)| entry.min_entry_id <= end_entry_id) .collect()) } @@ -286,6 +316,10 @@ fn out_of_order_reason( #[cfg(test)] mod tests { + use rand::rngs::StdRng; + use rand::seq::SliceRandom; + use rand::{Rng, SeedableRng}; + use super::*; use crate::error::Error; use crate::object_store_wal::batch::entry_id; @@ -467,6 +501,35 @@ mod tests { "are not strictly increasing", ); assert_eq!(1, catalog.objects_in_order().count()); + + // A segment that fits between two others by its entry ids must also + // fit between them by its sequence. + catalog + .insert_object(8, vec![footer_entry(region_id, 50, 60)]) + .unwrap(); + assert_corrupted( + catalog.insert_object(1, vec![footer_entry(region_id, 30, 40)]), + "object 1 ends at 40, object 2 starts at 10", + ); + assert_corrupted( + catalog.insert_object(9, vec![footer_entry(region_id, 30, 40)]), + "object 8 ends at 60, object 9 starts at 30", + ); + assert_corrupted( + catalog.insert_object(5, vec![footer_entry(region_id, 61, 70)]), + "object 5 ends at 70, object 8 starts at 50", + ); + catalog + .insert_object(5, vec![footer_entry(region_id, 30, 40)]) + .unwrap(); + assert_eq!( + vec![2, 5, 8], + object_seqs( + &catalog + .objects_for_entry_range(region_id, 0, EntryId::MAX) + .unwrap() + ) + ); } #[test] @@ -500,6 +563,7 @@ mod tests { fn test_catalog_reclaim_boundary_follows_the_watermarks() { let region_one = RegionId::new(1, 1); let region_two = RegionId::new(2, 1); + let region_three = RegionId::new(3, 1); let mut catalog = ObjectCatalog::default(); catalog.insert_object(0, Vec::new()).unwrap(); catalog @@ -526,49 +590,47 @@ mod tests { vec![footer_entry(region_two, entry_id(3, 1), entry_id(3, 1))], ) .unwrap(); - let boundary = |catalog: &ObjectCatalog, obsolete: &[(RegionId, EntryId)]| { - catalog.reclaim_boundary(0, &obsolete.iter().copied().collect()) - }; // A region without a watermark holds the boundary at its first object. - assert_eq!(1, boundary(&catalog, &[])); - assert_eq!(1, boundary(&catalog, &[(region_one, entry_id(2, 1))])); - // A watermark inside an object keeps that object. - assert_eq!( - 1, - boundary( - &catalog, - &[(region_one, entry_id(1, 1)), (region_two, EntryId::MAX)] - ) - ); - assert_eq!( - 2, - boundary( - &catalog, - &[(region_one, entry_id(1, 2)), (region_two, entry_id(1, 1))] - ) - ); + assert_eq!(1, catalog.reclaim_boundary(0)); + catalog.obsolete(region_two, entry_id(1, 1)); + assert_eq!(1, catalog.reclaim_boundary(0)); + // A watermark inside a segment keeps its object. + catalog.obsolete(region_one, entry_id(1, 1)); + assert_eq!(1, catalog.reclaim_boundary(0)); + catalog.obsolete(region_one, entry_id(1, 2)); + assert_eq!(2, catalog.reclaim_boundary(0)); // The segment of release records never holds the boundary. - assert_eq!( - 3, - boundary( - &catalog, - &[(region_one, entry_id(2, 1)), (region_two, entry_id(1, 1))] - ) - ); - // A region of `obsolete_all` contributes nothing, and when no region - // does, the boundary is past the last object. - assert_eq!( - 4, - boundary( - &catalog, - &[(region_one, entry_id(2, 1)), (region_two, EntryId::MAX)] - ) - ); + catalog.obsolete(region_one, entry_id(2, 1)); + assert_eq!(3, catalog.reclaim_boundary(0)); + // A region whose watermark is `EntryId::MAX` holds nothing, and when + // no region does, the boundary is past the last object. + catalog.obsolete(region_two, EntryId::MAX); + assert_eq!(4, catalog.reclaim_boundary(0)); + // A watermark never moves down, and the boundary is never below + // `from`. + catalog.obsolete(region_one, entry_id(1, 1)); + assert_eq!(4, catalog.reclaim_boundary(0)); + assert_eq!(5, catalog.reclaim_boundary(5)); - // Objects below the previous boundary are not scanned again. - assert_eq!(3, catalog.reclaim_boundary(3, &HashMap::new())); - assert_eq!(5, catalog.reclaim_boundary(5, &HashMap::new())); + // A segment indexed at or below the watermark of its region is not + // live, one of a region without a watermark is. + catalog + .insert_object( + 4, + vec![footer_entry(region_two, entry_id(4, 1), entry_id(4, 1))], + ) + .unwrap(); + assert_eq!(5, catalog.reclaim_boundary(0)); + catalog + .insert_object( + 5, + vec![footer_entry(region_three, entry_id(5, 1), entry_id(5, 1))], + ) + .unwrap(); + assert_eq!(5, catalog.reclaim_boundary(0)); + catalog.obsolete(region_three, entry_id(5, 1)); + assert_eq!(6, catalog.reclaim_boundary(0)); // Removing objects keeps the largest entry id of every region, and a // region without indexed objects contributes nothing. @@ -588,8 +650,194 @@ mod tests { .unwrap() .is_empty() ); - assert_eq!(3, boundary(&catalog, &[])); - assert_eq!(4, catalog.next_object_seq().unwrap()); + assert_eq!(6, catalog.reclaim_boundary(0)); + assert_eq!(6, catalog.next_object_seq().unwrap()); + } + + /// Random appends, watermark moves, releases, drops, collections and + /// recoveries give the reclaim boundary that a scan of the footers from + /// `from` gives, whatever `from` is, and the objects of a region that a + /// scan finds for a range of entry ids. + #[test] + fn test_catalog_matches_a_scan_of_the_footers() { + let regions = (1..=6) + .map(|number| RegionId::new(1, number)) + .collect::>(); + for seed in 0..256 { + let mut rng = StdRng::seed_from_u64(seed); + let mut catalog = ObjectCatalog::default(); + let mut watermarks = HashMap::new(); + // The largest entry id written per region, which a later segment + // of the region is above. + let mut written = HashMap::new(); + let mut next_object_seq = 0; + let mut boundary = 0; + for step in 0..128 { + let mut releases = rng.random_bool(0.2); + match rng.random_range(0..8) { + 0..=2 => {} + // A watermark move, which may be below the watermark or + // above every id written, of a region or, rarely, of the + // reserved region, which holds nothing live either way. + 3 | 4 => { + let region_id = if rng.random_bool(0.1) { + RELEASE_REGION_ID + } else { + regions[rng.random_range(0..regions.len())] + }; + let ceiling = written.get(®ion_id).copied().unwrap_or(0) + 2; + let entry_id = rng.random_range(0..=ceiling); + obsolete(&mut catalog, &mut watermarks, region_id, entry_id); + } + // A release through the largest id written, or a drop, + // seals an object with a release record. + 5 => { + let region_id = regions[rng.random_range(0..regions.len())]; + let entry_id = if rng.random_bool(0.5) { + written.get(®ion_id).copied().unwrap_or(0) + } else { + EntryId::MAX + }; + obsolete(&mut catalog, &mut watermarks, region_id, entry_id); + releases = true; + } + // A collection deletes objects below the boundary, in any + // order and not always all of them. Removing an object + // above it, which may hold live segments, is followed too. + 6 => { + let collected = catalog + .objects_in_order() + .map(|(object_seq, _)| object_seq) + .filter(|object_seq| { + (*object_seq < boundary || rng.random_bool(0.1)) + && rng.random_bool(0.7) + }) + .collect::>(); + for object_seq in collected { + assert!(catalog.remove_object(object_seq).is_some()); + } + } + // A recovery indexes the objects in any order, then the + // watermarks. + _ => { + let mut objects = catalog + .objects_in_order() + .map(|(object_seq, footer)| (object_seq, footer.to_vec())) + .collect::>(); + objects.shuffle(&mut rng); + catalog = ObjectCatalog::default(); + for (object_seq, footer) in objects { + catalog.insert_object(object_seq, footer).unwrap(); + } + for (region_id, entry_id) in &watermarks { + catalog.obsolete(*region_id, *entry_id); + } + } + } + if rng.random_bool(0.5) || releases { + let segments = regions + .iter() + .copied() + .filter(|_| rng.random_bool(0.5)) + .chain(releases.then_some(RELEASE_REGION_ID)) + .collect::>(); + let mut footer = Vec::new(); + for region_id in segments { + let min_entry_id = written.get(®ion_id).copied().unwrap_or(0) + + 1 + + rng.random_range(0..3); + let max_entry_id = min_entry_id + rng.random_range(0..3); + written.insert(region_id, max_entry_id); + footer.push(footer_entry(region_id, min_entry_id, max_entry_id)); + } + catalog.insert_object(next_object_seq, footer).unwrap(); + next_object_seq += rng.random_range(1..=2); + } + + let expected = scanned_reclaim_boundary(&catalog, boundary, &watermarks); + assert_eq!( + expected, + catalog.reclaim_boundary(boundary), + "seed {seed}, step {step}" + ); + let from = rng.random_range(0..=next_object_seq + 1); + assert_eq!( + scanned_reclaim_boundary(&catalog, from, &watermarks), + catalog.reclaim_boundary(from), + "seed {seed}, step {step}, from {from}" + ); + boundary = expected; + + for region_id in regions.iter().chain([&RELEASE_REGION_ID]) { + let ceiling = written.get(region_id).copied().unwrap_or(0) + 2; + let start = rng.random_range(0..=ceiling); + let end = rng.random_range(start..=ceiling); + let scanned = catalog + .objects_in_order() + .filter(|(_, footer)| { + footer.iter().any(|entry| { + entry.region_id == *region_id + && entry.max_entry_id >= start + && entry.min_entry_id <= end + }) + }) + .map(|(object_seq, _)| object_seq) + .collect::>(); + assert_eq!( + scanned, + object_seqs( + &catalog + .objects_for_entry_range(*region_id, start, end) + .unwrap() + ), + "seed {seed}, step {step}, region {region_id}, {start}..={end}" + ); + } + } + } + } + + fn obsolete( + catalog: &mut ObjectCatalog, + watermarks: &mut HashMap, + region_id: RegionId, + entry_id: EntryId, + ) { + catalog.obsolete(region_id, entry_id); + watermarks + .entry(region_id) + .and_modify(|watermark: &mut EntryId| *watermark = (*watermark).max(entry_id)) + .or_insert(entry_id); + } + + /// The first object at or above `from` holding a segment that is live + /// under `watermarks`, found by reading every footer, or the sequence + /// after the last object when none does, and never less than `from`. + fn scanned_reclaim_boundary( + catalog: &ObjectCatalog, + from: u64, + watermarks: &HashMap, + ) -> u64 { + let live = catalog + .objects_in_order() + .filter(|(object_seq, _)| *object_seq >= from) + .find(|(_, footer)| { + footer.iter().any(|entry| { + entry.region_id != RELEASE_REGION_ID + && watermarks + .get(&entry.region_id) + .is_none_or(|watermark| entry.max_entry_id > *watermark) + }) + }); + match live { + Some((object_seq, _)) => object_seq, + None => catalog + .objects_in_order() + .last() + .map_or(from, |(object_seq, _)| { + object_seq.saturating_add(1).max(from) + }), + } } fn footer_entry(region_id: RegionId, min_entry_id: u64, max_entry_id: u64) -> FooterEntry { diff --git a/src/log-store/src/object_store_wal/store.rs b/src/log-store/src/object_store_wal/store.rs index df037b43771..d5096c073eb 100644 --- a/src/log-store/src/object_store_wal/store.rs +++ b/src/log-store/src/object_store_wal/store.rs @@ -394,8 +394,10 @@ impl ObjectStoreLogStore { METRIC_OBJECT_STORE_WAL_INDEXED_OBJECTS.inc(); METRIC_OBJECT_STORE_WAL_INDEXED_BYTES.add(object_len(&[]) as i64); // The release records are the only watermarks recorded yet. - let writer_reclaim_boundary = - catalog.reclaim_boundary(reclaim_boundary, &released_entry_ids); + for (®ion_id, &entry_id) in &released_entry_ids { + catalog.obsolete(region_id, entry_id); + } + let writer_reclaim_boundary = catalog.reclaim_boundary(reclaim_boundary); let catalog = Arc::new(RwLock::new(catalog)); let obsolete_entry_ids = Arc::new(Mutex::new(released_entry_ids)); let wal_holes = WalHoles::default(); @@ -710,14 +712,24 @@ fn injected_create_failure(io: &dyn WalObjectIo, object_seq: u64) -> Result, + region_id: RegionId, + entry_id: EntryId, +) { obsolete_entry_ids .lock() .unwrap_or_else(PoisonError::into_inner) .entry(region_id) .and_modify(|current| *current = (*current).max(entry_id)) .or_insert(entry_id); + catalog + .write() + .unwrap_or_else(PoisonError::into_inner) + .obsolete(region_id, entry_id); } fn positive_bytes(bytes: u64, name: &str) -> Result { @@ -1130,7 +1142,12 @@ impl LogStore for ObjectStoreLogStore { // still queued behind the stop: nothing is assigned an id any // more, so the watermark alone is consistent. None => { - record_obsolete(&self.obsolete_entry_ids, region_id, watermark); + record_obsolete( + &self.obsolete_entry_ids, + &self.catalog, + region_id, + watermark, + ); Ok(()) } } @@ -2412,7 +2429,12 @@ impl Actor { response: oneshot::Sender>, ) { let result = self.raise_sequence_floor(region_id, entry_id).map(|_| { - record_obsolete(&self.obsolete_entry_ids, region_id, watermark); + record_obsolete( + &self.obsolete_entry_ids, + &self.catalog, + region_id, + watermark, + ); }); if result.is_ok() { self.update_reclaim_boundary(); @@ -2451,7 +2473,7 @@ impl Actor { } else { self.issued_entry_ids.get(®ion_id).copied().unwrap_or(0) }; - record_obsolete(&self.obsolete_entry_ids, region_id, entry_id); + record_obsolete(&self.obsolete_entry_ids, &self.catalog, region_id, entry_id); self.update_reclaim_boundary(); self.collect(); let release = Release { @@ -2481,17 +2503,11 @@ impl Actor { /// Raises the reclaim boundary to what the catalog and the obsolete /// watermarks allow. fn update_reclaim_boundary(&mut self) { - let boundary = { - let obsolete_entry_ids = self - .obsolete_entry_ids - .lock() - .unwrap_or_else(PoisonError::into_inner); - self.catalog - .read() - .unwrap_or_else(PoisonError::into_inner) - .reclaim_boundary(self.reclaim_boundary, &obsolete_entry_ids) - }; - self.reclaim_boundary = boundary; + self.reclaim_boundary = self + .catalog + .read() + .unwrap_or_else(PoisonError::into_inner) + .reclaim_boundary(self.reclaim_boundary); } /// Starts a collection, which also attempts the deletes that failed since @@ -9926,14 +9942,157 @@ mod tests { } } - /// Measurements of costs that grow with the prefix, taken through the - /// store. They are ignored by default and print what they measure; run - /// them with `--run-ignored only` and `--no-capture`, preferably with - /// `--cargo-profile release`. + /// Measurements of costs that grow with the catalog or the prefix, + /// taken through the store. They are ignored by default and print what + /// they measure; run them with `--run-ignored only` and `--no-capture`, + /// preferably with `--cargo-profile release`. #[allow(clippy::print_stdout)] mod measurements { use super::*; + /// Times every watermark update while one object holds a segment of + /// each region and the regions are flushed in region id order, so + /// each update leaves the next region as the first live segment of + /// that object. An update is timed as the caller of `obsolete` sees + /// it, which includes the round trip of its command to the actor. + #[tokio::test] + #[ignore] + async fn measure_watermark_updates_over_one_shared_object() { + let _serialized = SERIALIZED.lock().await; + for regions in [1_000, 5_000, 10_000, 20_000] { + let store = open_with_one_shared_object(regions).await; + let mut pauses = Vec::with_capacity(regions as usize); + for number in 1..=regions { + let region_id = region(number); + let started_at = Instant::now(); + store + .obsolete(&provider(region_id), region_id, id(1, 1)) + .await + .unwrap(); + pauses.push(started_at.elapsed()); + } + // The next object carries the boundary past every flushed + // object. + append(&store, region(regions + 1), "b").await.unwrap(); + assert_eq!(vec![2], collected_object_seqs(&store).await); + print_pauses("watermark updates", regions, pauses); + store.stop().await.unwrap(); + } + } + + /// Times the steps the actor takes for each of the same watermark + /// updates that depend on the catalog, recording the watermark and + /// recomputing the reclaim boundary, without the round trip: the + /// store is stopped first and the steps run on its catalog. + #[tokio::test] + #[ignore] + async fn measure_actor_watermark_steps_over_one_shared_object() { + let _serialized = SERIALIZED.lock().await; + for regions in [1_000, 5_000, 10_000, 20_000] { + let store = open_with_one_shared_object(regions).await; + store.stop().await.unwrap(); + let mut boundary = 0; + let mut pauses = Vec::with_capacity(regions as usize); + for number in 1..=regions { + let started_at = Instant::now(); + record_obsolete( + &store.obsolete_entry_ids, + &store.catalog, + region(number), + id(1, 1), + ); + boundary = store.catalog.read().unwrap().reclaim_boundary(boundary); + pauses.push(started_at.elapsed()); + } + assert_eq!(2, boundary); + print_pauses("actor watermark steps", regions, pauses); + } + } + + /// Opens a store whose object 1 holds one segment of each of + /// `regions` regions. + async fn open_with_one_shared_object(regions: u32) -> Arc { + let store = open(memory_store(), &eager()).await; + let entries = (1..=regions) + .map(|number| entry(&store, region(number), "a")) + .collect(); + store.append_batch(entries).await.unwrap(); + store + } + + fn print_pauses(label: &str, regions: u32, mut pauses: Vec) { + let total = pauses.iter().sum::(); + pauses.sort_unstable(); + let micros = |pause: Duration| pause.as_secs_f64() * 1e6; + println!( + "{label}: regions={regions} total_ms={:.3} mean_us={:.3} p50_us={:.3} p99_us={:.3} max_us={:.3}", + total.as_secs_f64() * 1e3, + micros(total) / pauses.len() as f64, + micros(pauses[pauses.len() / 2]), + micros(pauses[pauses.len() * 99 / 100]), + micros(pauses[pauses.len() - 1]), + ); + } + + /// Times 100 reads of the last ten objects of a region that holds a + /// segment in every retained object, and 100 lookups of those objects + /// in the catalog alone. + #[tokio::test] + #[ignore] + async fn measure_tail_reads_over_retained_objects() { + let _serialized = SERIALIZED.lock().await; + let region_id = region(1); + for objects in [10_000, 100_000, 1_000_000] { + let object_store = memory_store(); + let io = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); + futures::stream::iter(0..objects) + .for_each_concurrent(RECOVERY_CONCURRENCY, |object_seq| { + let io = &io; + async move { + let header = Header { + object_seq, + epoch: 1, + prev: object_seq.checked_sub(1).map(|prev| ChainLink { + object_seq: prev, + epoch: 1, + }), + reclaim_boundary: 0, + }; + let record = Record { + region_id, + entry_id: id(object_seq, 1), + payload: Bytes::from_static(b"a"), + }; + put_object_with_header(io, header, &[record]).await; + } + }) + .await; + let store = open(object_store, &eager()).await; + let start = id(objects - 10, 1); + + let started_at = Instant::now(); + for _ in 0..100 { + assert_eq!(10, read_entries(&store, region_id, start).await.len()); + } + let reads = started_at.elapsed(); + let started_at = Instant::now(); + for _ in 0..100 { + let catalog = store.catalog.read().unwrap(); + let found = catalog + .objects_for_entry_range(region_id, start, EntryId::MAX) + .unwrap(); + assert_eq!(10, found.len()); + } + let lookups = started_at.elapsed(); + println!( + "tail reads: objects={objects} read_100_ms={:.3} lookup_100_ms={:.3}", + reads.as_secs_f64() * 1e3, + lookups.as_secs_f64() * 1e3, + ); + store.stop().await.unwrap(); + } + } + /// Object access that counts the reads of whole objects and of byte /// ranges. struct CountingIo { From cca95b2637a797a9a0b8e6824718500d0fe692e7 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Mon, 19 Oct 2026 04:10:34 +0000 Subject: [PATCH 033/239] fix(datanode): stop the object store WAL before the region server (#9639) * fix(mito2): cancel flushes waiting for the WAL when the engine stops A flush waits for the WAL to be durable through the entry id it records before it writes its manifest edit. With the object store WAL in the `enqueued` mode that wait lasts until an upload succeeds, while the log store is stopped only after the engines. A graceful shutdown during an outage of the WAL's object store therefore never finished: closing a region does not cancel its running flush, stopping the worker group waits for running flushes and runs the queued ones, and a region holding writes that skipped the WAL is flushed on close. The worker group now owns a cancellation token, raised before its schedulers stop and kept raised. `RegionEngine::begin_stop`, which the region server calls before it closes regions, raises it earlier, so that a close-time flush cannot wait either. A flush polls its durability barrier first and gives up at the barrier once the token is raised: it writes no manifest edit, deletes its SSTs and fails as cancelled. A flush whose entries are already durable commits as before. The region server reports a failed close, which loses the region's writes that skipped the WAL, in its stop result. Signed-off-by: jeremyhi * test(mito2): wait for the outage on the store under test Signed-off-by: jeremyhi * fix(mito2): give a healthy WAL a grace period before a stopping engine cancels flushes A flush of a stopping engine no longer gives up its wait for the WAL as soon as the stop begins. It gives up at once while the log store reports that it fails to make entries durable, through the new `LogStore::is_failing`, and otherwise 5 s after the engine began to stop, so a healthy object store WAL uploads its last entries first and a close-time flush keeps the writes that skipped the WAL. The object store WAL reports that it fails while the last create of its backlog failed, which covers a refused and an unreachable object store, and once it is poisoned or stopped. Signed-off-by: jeremyhi * fix(mito2): drain the WAL for a stopping flush and cancel its failed waits A flush of a stopping engine now asks the log store to drain before it waits out the grace period, through a new `LogStore::begin_drain` that does nothing by default. The object store WAL then seals its open batch at once and, until it stops, as soon as a durability wait needs a region it holds, so a healthy store makes the entries durable after about one upload whatever its `flush_interval`, which has no upper bound. A barrier wait that fails while the engine stops, for instance on a poisoned or stopped log store, now fails the flush as cancelled like a wait that gives up. The healthy-store test no longer depends on the seal timer: its store never seals on its own, and the close must finish well within the grace period. Signed-off-by: jeremyhi * fix(mito2): drain the WAL when the engine begins to stop The drain started only when a stopping flush reached its durability barrier, so a close behind slow flushes that reached it after the grace period had no time left to upload its entries. The worker group now asks the log store to drain as soon as it begins to stop, and the object store WAL seals its open batch then and every later admission right away, so its entries become durable while the flushes still write their SSTs. `LogStore::begin_drain` becomes synchronous for that. The healthy-store test now writes after the stop began and closes the region only after the grace period. Signed-off-by: jeremyhi * fix(log-store): refuse appends once the object store WAL drains A WAL write admitted late in a stop could still be uploading when the close-time flush of its region reached an expired grace period, so the writes that skipped the WAL in that region were lost although the store was healthy. A draining object store WAL now refuses every later append, also one held back, with the stopped error, so the entries a stopping engine waits for are fixed when the stop begins and become durable while the regions still close; writes that skip the WAL are still taken. Tests cover a close-time flush whose wait completes within the grace period once a parked create is released, a close that reaches its wait after the grace period, and a drain refusing a held-back append. Signed-off-by: jeremyhi * fix(datanode): word a failed region close neutrally in the shutdown log A close can fail for reasons other than its flush, such as a stopped worker, and for regions without writes that skipped the WAL, so the error line no longer states that such writes were lost: it says they may be. Signed-off-by: jeremyhi * fix(datanode): stop the object store WAL before the region server The datanode stopped the object store WAL only after the region server, whose stop waits for flushes that wait for the WAL to become durable, so a shutdown never finished while the WAL's uploads kept failing. The WAL now stops first: its stop uploads the backlog or drops it and answers every wait for its entries, so a running, queued or close-time flush either commits or fails with the WAL-stopped error without writing a manifest edit. This replaces the stop signal and grace period the engine used to give up those waits, and the log store and engine hooks they needed. Signed-off-by: jeremyhi * test(datanode): check rows after a healthy shutdown and wait for the WAL stop Signed-off-by: jeremyhi * test: keep only the shutdown tests that prove the new stop order The engine-level tests pinned the log store's stop contract, which the log store's own tests already cover, and the stalled-create test asserted a limitation rather than a guarantee. Drop them with the testing hooks only they used; the datanode and region server tests cover the change end to end. Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- src/datanode/Cargo.toml | 1 + src/datanode/src/datanode.rs | 277 +++++++++++++++++++++++++++++- src/datanode/src/region_server.rs | 53 +++++- 3 files changed, 325 insertions(+), 6 deletions(-) diff --git a/src/datanode/Cargo.toml b/src/datanode/Cargo.toml index 7fedd2fc263..bef2641415e 100644 --- a/src/datanode/Cargo.toml +++ b/src/datanode/Cargo.toml @@ -81,5 +81,6 @@ common-meta = { workspace = true, features = ["testing"] } common-query.workspace = true common-test-util.workspace = true datafusion-common.workspace = true +log-store = { workspace = true, features = ["testing"] } mito2 = { workspace = true, features = ["test"] } session.workspace = true diff --git a/src/datanode/src/datanode.rs b/src/datanode/src/datanode.rs index 8edcd2ce4a1..df12cf4f2a8 100644 --- a/src/datanode/src/datanode.rs +++ b/src/datanode/src/datanode.rs @@ -167,7 +167,10 @@ impl Datanode { .context(ShutdownInstanceSnafu), ); } - record_shutdown_error(&mut first_error, self.region_server.stop().await); + // The object store WAL stops before the regions close: its stop + // uploads the backlog or drops it, and answers every wait for its + // entries, so a flush that waits for the WAL, queued or run by the + // close of a region, cannot hold up the shutdown. if let Some(log_store) = &self.object_store_log_store { record_shutdown_error( &mut first_error, @@ -178,6 +181,7 @@ impl Datanode { .context(ShutdownInstanceSnafu), ); } + record_shutdown_error(&mut first_error, self.region_server.stop().await); first_error.map_or(Ok(()), Err) } @@ -1055,11 +1059,13 @@ mod tests { use std::assert_matches; use std::collections::{BTreeMap, HashMap}; use std::io::Write; + use std::ops::Range; use std::path::Path; use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Duration; + use api::v1::{ColumnSchema, Rows}; use cache::build_datanode_layered_cache_registry; use common_base::Plugins; use common_base::readable_size::ReadableSize; @@ -1073,28 +1079,34 @@ mod tests { use common_meta::node_manager::NodeManagerRef; use common_test_util::temp_dir::{create_named_temp_file, create_temp_dir}; use common_wal::config::DatanodeWalConfig; - use common_wal::config::object_store::{ObjectStoreWalConfig, STANDALONE_GENERATION}; + use common_wal::config::object_store::{AckMode, ObjectStoreWalConfig, STANDALONE_GENERATION}; + use common_wal::options::{ObjectStoreWalOptions, WAL_OPTIONS_KEY, WalOptions}; use log_store::ObjectStoreLogStore; use log_store::error::Error as LogStoreError; use meta_client::client::MetaClientBuilder; use mito2::engine::MITO_ENGINE_NAME; + use mito2::test_util::{CreateRequestBuilder, build_rows, rows_schema}; use object_store::ObjectStore; use object_store::manager::ObjectStoreManager; use partition::manager::PartitionRuleManagerRef; use query::region_query::{RegionQueryHandlerFactory, RegionQueryHandlerFactoryRef}; use store_api::logstore::LogStore; - use store_api::region_request::RegionRequest; + use store_api::logstore::provider::Provider; + use store_api::region_request::{PathType, RegionOpenRequest, RegionPutRequest, RegionRequest}; use store_api::storage::RegionId; use crate::config::{DatanodeOptions, RegionEngineConfig, StorageConfig}; use crate::datanode::{ - DatanodeBuilder, validate_object_store_wal_config, validate_region_engine_config, + Datanode, DatanodeBuilder, validate_object_store_wal_config, validate_region_engine_config, wal_object_store, }; use crate::error::{self, Error}; use crate::region_query::DatanodeRegionQueryHandler; + use crate::region_server::RegionServer; use crate::tests::{MockRegionEngine, mock_region_server}; + const WAIT: Duration = Duration::from_secs(30); + async fn setup_table_datanode(kv: &KvBackendRef) { let mgr = DatanodeTableManager::new(kv.clone()); let txn = mgr @@ -1579,6 +1591,263 @@ mod tests { assert!(is_stopped(log_store).await); } + const REGION: RegionId = RegionId::new(1024, 1); + + fn enqueued_wal_config(flush_interval: Duration) -> ObjectStoreWalConfig { + ObjectStoreWalConfig { + ack_mode: AckMode::Enqueued, + flush_interval, + ..Default::default() + } + } + + /// The options that put a region on the WAL of a datanode with `config`. + fn region_wal_options(config: &ObjectStoreWalConfig) -> HashMap { + let prefix = config.node_prefix(0, STANDALONE_GENERATION); + let wal_options = + serde_json::to_string(&WalOptions::ObjectStore(ObjectStoreWalOptions::new(prefix))) + .unwrap(); + HashMap::from([(WAL_OPTIONS_KEY.to_string(), wal_options)]) + } + + async fn put_rows( + region_server: &RegionServer, + schema: &[ColumnSchema], + rows: Range, + skip_wal: bool, + ) { + if rows.is_empty() { + return; + } + region_server + .handle_request( + REGION, + RegionRequest::Put(RegionPutRequest { + rows: Rows { + schema: schema.to_vec(), + rows: build_rows(rows.start, rows.end), + }, + skip_wal, + hint: None, + partition_expr_version: None, + }), + ) + .await + .unwrap(); + } + + /// Builds a datanode on an `enqueued` object store WAL with + /// `flush_interval` and creates [`REGION`] on it with `durable_rows` + /// written through the WAL and uploaded. Then makes every later create + /// of the WAL fail with a transient error if `uploads_fail`, and writes + /// `rows` through the WAL and `skipped_rows` past it. Returns the datanode + /// and its log store. + async fn datanode_with_object_store_wal( + data_home: &str, + flush_interval: Duration, + durable_rows: usize, + uploads_fail: bool, + rows: usize, + skipped_rows: usize, + ) -> (Datanode, Arc) { + let config = enqueued_wal_config(flush_interval); + let prefix = config.node_prefix(0, STANDALONE_GENERATION); + let mut request = CreateRequestBuilder::new().build(); + request.options.extend(region_wal_options(&config)); + let datanode = object_store_wal_builder(data_home, config) + .build() + .await + .unwrap(); + let log_store = datanode.object_store_log_store.clone().unwrap(); + let region_server = datanode.region_server(); + let schema = rows_schema(&request); + region_server + .handle_request(REGION, RegionRequest::Create(request)) + .await + .unwrap(); + + put_rows(®ion_server, &schema, 0..durable_rows, false).await; + if durable_rows > 0 { + let provider = Provider::object_store_provider(REGION, prefix); + tokio::time::timeout(WAIT, async { + while log_store.latest_entry_id(&provider).unwrap() == 0 { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("the rows must become durable"); + } + if uploads_fail { + log_store.fail_creates(); + } + let skipped = durable_rows + rows; + put_rows(®ion_server, &schema, durable_rows..skipped, false).await; + put_rows( + ®ion_server, + &schema, + skipped..skipped + skipped_rows, + true, + ) + .await; + (datanode, log_store) + } + + /// Restarts a datanode on the files of `data_home`, opens [`REGION`] and + /// returns the number of rows it holds after replaying the WAL. + async fn rows_after_restart(data_home: &str) -> u64 { + let config = enqueued_wal_config(ObjectStoreWalConfig::default().flush_interval); + let options = region_wal_options(&config); + let mut datanode = object_store_wal_builder(data_home, config) + .build() + .await + .unwrap(); + let region_server = datanode.region_server(); + region_server + .handle_request( + REGION, + RegionRequest::Open(RegionOpenRequest { + engine: MITO_ENGINE_NAME.to_string(), + table_dir: CreateRequestBuilder::new().build().table_dir, + path_type: PathType::Bare, + options, + skip_wal_replay: false, + checkpoint: None, + requirements: Default::default(), + }), + ) + .await + .unwrap(); + let rows = region_server.region_statistic(REGION).unwrap().num_rows; + datanode.shutdown().await.unwrap(); + rows + } + + /// Returns the Parquet files under `dir` and its subdirectories. + fn sst_files(dir: &Path) -> Vec { + let mut files = Vec::new(); + for entry in std::fs::read_dir(dir).unwrap() { + let path = entry.unwrap().path(); + if path.is_dir() { + files.extend(sst_files(&path)); + } else if path.extension().is_some_and(|ext| ext == "parquet") { + files.push(path.display().to_string()); + } + } + files + } + + /// A flush that waits for WAL uploads that keep failing does not hold up + /// the shutdown: the WAL stops first and drops its backlog, which fails + /// the flush with the stopped error before it writes its manifest edit, + /// and the shutdown reports the dropped backlog. After a restart the rows + /// that were durable before the outage replay once. + #[tokio::test] + async fn test_shutdown_while_a_flush_waits_for_failing_wal_uploads() { + common_telemetry::init_default_ut_logging(); + let data_home = create_temp_dir("object-store-wal-shutdown-flush"); + let data_home = data_home.path().to_str().unwrap(); + let (mut datanode, log_store) = datanode_with_object_store_wal( + data_home, + ObjectStoreWalConfig::default().flush_interval, + 3, + true, + 2, + 0, + ) + .await; + let flush = { + let region_server = datanode.region_server(); + tokio::spawn(async move { + region_server + .handle_request(REGION, RegionRequest::Flush(Default::default())) + .await + }) + }; + tokio::time::timeout(WAIT, log_store.wait_for_durability_waits(1)) + .await + .unwrap() + .unwrap(); + + let err = tokio::time::timeout(WAIT, datanode.shutdown()) + .await + .expect("the shutdown must finish") + .unwrap_err(); + + assert_matches!(err, Error::ShutdownInstance { .. }); + assert!( + format!("{err:?}").contains("injected create failure"), + "{err:?}" + ); + let flush_err = flush.await.unwrap().unwrap_err(); + assert!( + format!("{flush_err:?}").contains("Object store WAL log store is stopped"), + "{flush_err:?}" + ); + assert_eq!(Vec::::new(), sst_files(Path::new(data_home))); + assert!(is_stopped(&log_store).await); + drop(datanode); + + assert_eq!(3, rows_after_restart(data_home).await); + } + + /// A region that holds writes that skipped the WAL is flushed when it + /// closes. When the WAL's uploads keep failing, the WAL stops first and + /// drops its backlog, so that flush fails at once instead of waiting: the + /// close fails, the writes that skipped the WAL are lost, and the shutdown + /// returns the dropped backlog. + #[tokio::test] + async fn test_shutdown_closes_a_region_after_the_wal_dropped_its_backlog() { + common_telemetry::init_default_ut_logging(); + let data_home = create_temp_dir("object-store-wal-shutdown-closing-flush"); + let (mut datanode, log_store) = datanode_with_object_store_wal( + data_home.path().to_str().unwrap(), + ObjectStoreWalConfig::default().flush_interval, + 0, + true, + 3, + 2, + ) + .await; + + let err = tokio::time::timeout(WAIT, datanode.shutdown()) + .await + .expect("the shutdown must finish") + .unwrap_err(); + + assert_matches!(err, Error::ShutdownInstance { .. }); + assert!( + format!("{err:?}").contains("injected create failure"), + "{err:?}" + ); + assert_eq!(Vec::::new(), sst_files(data_home.path())); + assert!(is_stopped(&log_store).await); + } + + /// With a healthy WAL whose seal interval is longer than any test, the + /// close-time flush of a region that holds writes that skipped the WAL + /// commits: the WAL's stop seals and uploads its open batch before the + /// region closes, so the flush finds its entries durable, and every row + /// survives a restart. + #[tokio::test] + async fn test_shutdown_flushes_a_region_while_the_wal_is_healthy() { + common_telemetry::init_default_ut_logging(); + let data_home = create_temp_dir("object-store-wal-shutdown-healthy-flush"); + let data_home = data_home.path().to_str().unwrap(); + let (mut datanode, log_store) = + datanode_with_object_store_wal(data_home, Duration::from_secs(3600), 0, false, 3, 2) + .await; + + tokio::time::timeout(WAIT, datanode.shutdown()) + .await + .expect("the shutdown must finish") + .unwrap(); + assert_eq!(1, sst_files(Path::new(data_home)).len()); + assert!(is_stopped(&log_store).await); + drop(datanode); + + assert_eq!(5, rows_after_restart(data_home).await); + } + #[test] fn test_node_prefix_passes_validation() { let config = ObjectStoreWalConfig::default(); diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index 23a8fb91c48..0b7c7645e00 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -1986,6 +1986,7 @@ impl RegionServerInner { .collect::>(); let num_regions = regions.len(); + let mut first_error = None; for (region_id, engine) in regions { let closed = engine .handle_request( @@ -1995,7 +1996,15 @@ impl RegionServerInner { .await; match closed { Ok(_) => debug!("Region {region_id} is closed"), - Err(e) => warn!("Failed to close region {region_id}, err: {e}"), + Err(e) => { + // A close that failed may not have flushed writes that + // skipped the WAL, which nothing replays. + error!(e; "Failed to close region {region_id}; writes that skipped the WAL may be lost"); + record_shutdown_error( + &mut first_error, + Err(e).context(HandleRegionRequestSnafu { region_id }), + ); + } } } self.region_map.clear(); @@ -2003,7 +2012,6 @@ impl RegionServerInner { drop(self.mito_engine.write().unwrap().take()); let engines = self.engines.write().unwrap().drain().collect::>(); - let mut first_error = None; for (engine_name, engine) in engines { let result = engine .stop() @@ -2539,6 +2547,47 @@ mod tests { assert!(status.is_some()); } + /// A region that fails to close is reported in the result of the stop, + /// which still closes the other regions. + #[tokio::test] + async fn test_stop_reports_a_region_that_failed_to_close() { + common_telemetry::init_default_ut_logging(); + + let mut mock_region_server = mock_region_server(); + let failing = RegionId::new(1, 1); + let closing = RegionId::new(1, 2); + let closed = Arc::new(std::sync::Mutex::new(Vec::new())); + let (engine, _receiver) = MockRegionEngine::with_mock_fn(MITO_ENGINE_NAME, { + let closed = closed.clone(); + Box::new(move |region_id, request| { + if matches!(request, RegionRequest::Close(_)) { + closed.lock().unwrap().push(region_id); + if region_id == failing { + return error::UnexpectedSnafu { + violated: "close failed".to_string(), + } + .fail(); + } + } + Ok(0) + }) + }); + mock_region_server.register_engine(engine.clone()); + for region_id in [failing, closing] { + mock_region_server + .inner + .region_map + .insert(region_id, RegionEngineWithStatus::Ready(engine.clone())); + } + + let err = mock_region_server.stop().await.unwrap_err(); + + assert_matches!(err, error::Error::HandleRegionRequest { region_id, .. } if region_id == failing); + let mut closed = closed.lock().unwrap().clone(); + closed.sort(); + assert_eq!(vec![failing, closing], closed); + } + #[tokio::test] async fn test_batch_open_region_ignore_nonexistent_regions() { common_telemetry::init_default_ut_logging(); From 4e8bc6df86bbb1aeb6636bec0bebfcf5d5b49cad Mon Sep 17 00:00:00 2001 From: fys <40801205+fengys1996@users.noreply.github.com> Date: Mon, 19 Oct 2026 06:58:40 +0000 Subject: [PATCH 034/239] fix(prom): align remote read label matcher semantics (#9633) * fix(prom): align remote read label matcher semantics Signed-off-by: fys * chore: optimize comment of validate_matchers * fix(prom): preserve existing binary label regex handling Signed-off-by: fys * refactor: condition * test(prom): simplify remote read matcher tests and clarify comments * fix(prom): return binary remote read labels as UTF-8 * fix(prom): reject binary labels in remote read * refactor(prom): align remote read null label handling with PromQL Signed-off-by: fys * refactor(prom): simplify null handling in label matcher branches Signed-off-by: fys * fix(prom): restrict remote read labels to string types Signed-off-by: fys * test(prom): use string label in default value column test Signed-off-by: fys * fix: code review --------- Signed-off-by: fys --- src/frontend/src/error.rs | 8 + src/frontend/src/instance.rs | 4 +- src/frontend/src/instance/prom_store.rs | 84 ++++++- src/servers/src/prom_store.rs | 315 ++++++++++++++++++++++-- tests-integration/src/prom_store.rs | 6 + 5 files changed, 389 insertions(+), 28 deletions(-) diff --git a/src/frontend/src/error.rs b/src/frontend/src/error.rs index 451040c884a..923e6879f35 100644 --- a/src/frontend/src/error.rs +++ b/src/frontend/src/error.rs @@ -231,6 +231,13 @@ pub enum Error { source: servers::error::Error, }, + #[snafu(display("Invalid Prometheus remote read schema: {detail}"))] + InvalidPromRemoteReadSchema { + detail: String, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Failed to create logical plan for prometheus query"))] PromStoreRemoteQueryPlan { #[snafu(implicit)] @@ -394,6 +401,7 @@ impl ErrorExt for Error { | Error::SchemaExists { .. } | Error::ColumnNotFound { .. } | Error::AmbiguousValueColumn { .. } + | Error::InvalidPromRemoteReadSchema { .. } | Error::UnsupportedFormat { .. } | Error::IllegalAuthConfig { .. } | Error::ColumnNoneDefaultValue { .. } diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index e7d2e828b28..d72d0677552 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -3336,7 +3336,7 @@ mod tests { false, ) .with_time_index(true), - ColumnSchema::new("extra_field", ConcreteDataType::float64_datatype(), false), + ColumnSchema::new("extra_field", ConcreteDataType::string_datatype(), false), ColumnSchema::new( greptime_value(), ConcreteDataType::float64_datatype(), @@ -3347,7 +3347,7 @@ mod tests { schema, vec![ Arc::new(TimestampMillisecondVector::from_vec(vec![1000, 2000, 3000])) as VectorRef, - Arc::new(Float64Vector::from_vec(vec![99.0, 99.0, 99.0])) as VectorRef, + Arc::new(StringVector::from(vec!["99", "99", "99"])) as VectorRef, Arc::new(Float64Vector::from_vec(vec![1.0, 2.0, 3.0])) as VectorRef, ], ) diff --git a/src/frontend/src/instance/prom_store.rs b/src/frontend/src/instance/prom_store.rs index 875cbaa873d..7238c38f83d 100644 --- a/src/frontend/src/instance/prom_store.rs +++ b/src/frontend/src/instance/prom_store.rs @@ -35,6 +35,8 @@ use common_query::Output; use common_query::prelude::{GREPTIME_PHYSICAL_TABLE, greptime_value}; use common_recordbatch::RecordBatches; use common_telemetry::{debug, tracing}; +use datatypes::data_type::ConcreteDataType; +use datatypes::schema::Schema; use operator::insert::{ AutoCreateTableType, InserterRef, PerTableSemanticIndex, apply_per_table_semantic_options, build_create_table_expr, fill_table_options_for_create, parse_per_table_semantic_index, @@ -52,7 +54,7 @@ use servers::query_handler::{ PromStoreProtocolHandler, PromStoreProtocolHandlerRef, PromStoreResponse, }; use session::context::QueryContextRef; -use snafu::{OptionExt, ResultExt}; +use snafu::{OptionExt, ResultExt, ensure}; use store_api::metric_engine_consts::{METRIC_ENGINE_NAME, PHYSICAL_TABLE_METADATA_KEY}; use store_api::mito_engine_options::SST_FORMAT_KEY; use table::TableRef; @@ -61,7 +63,7 @@ use tracing::instrument; use crate::error::{ AmbiguousValueColumnSnafu, CatalogSnafu, ColumnNotFoundSnafu, ExecLogicalPlanSnafu, - PromStoreRemoteQueryPlanSnafu, ReadTableSnafu, Result, + InvalidPromRemoteReadSchemaSnafu, PromStoreRemoteQueryPlanSnafu, ReadTableSnafu, Result, }; use crate::instance::Instance; @@ -172,7 +174,7 @@ async fn to_query_result( }) } -fn resolve_column_names(table_name: &str, table: &TableRef) -> Result { +fn resolve_value_column_name(table_name: &str, table: &TableRef) -> Result { let columns = table .field_columns() .map(|column| column.name) @@ -198,6 +200,29 @@ fn resolve_column_names(table_name: &str, table: &TableRef) -> Result { } } +fn ensure_string_label_columns( + schema: &Schema, + timestamp_column_name: &str, + value_column_name: &str, +) -> Result<()> { + for column in schema.column_schemas() { + if column.name == timestamp_column_name || column.name == value_column_name { + continue; + } + // All remaining columns are emitted as labels, including non-key columns. + ensure!( + matches!(column.data_type, ConcreteDataType::String(_)), + InvalidPromRemoteReadSchemaSnafu { + detail: format!( + "Prometheus remote read only supports string labels: column '{}' has type {:?}", + column.name, column.data_type + ), + } + ); + } + Ok(()) +} + impl Instance { #[tracing::instrument(skip_all)] async fn handle_remote_query( @@ -232,7 +257,9 @@ impl Instance { .name .clone(); - let value_column_name = resolve_column_names(table_name, &table)?; + let value_column_name = resolve_value_column_name(table_name, &table)?; + + ensure_string_label_columns(&table.schema(), ×tamp_column_name, &value_column_name)?; let dataframe = self .query_engine @@ -817,10 +844,59 @@ mod tests { use std::sync::Arc; use api::prom_store::remote::LabelMatcher; + use common_error::ext::ErrorExt; + use common_error::status_code::StatusCode; use session::context::QueryContext; use super::*; + #[test] + fn test_remote_read_label_types() { + use datatypes::schema::ColumnSchema; + + for (label_type, supported) in [ + (ConcreteDataType::string_datatype(), true), + (ConcreteDataType::large_string_datatype(), true), + (ConcreteDataType::utf8_view_datatype(), true), + (ConcreteDataType::binary_datatype(), false), + (ConcreteDataType::binary_view_datatype(), false), + (ConcreteDataType::json_datatype(), false), + (ConcreteDataType::vector_datatype(3), false), + (ConcreteDataType::int64_datatype(), false), + (ConcreteDataType::boolean_datatype(), false), + ] { + let schema = Schema::new(vec![ + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + ColumnSchema::new("value", label_type.clone(), true), + ColumnSchema::new("label", ConcreteDataType::string_datatype(), true), + ]); + // Value columns are validated separately, regardless of their type. + ensure_string_label_columns(&schema, "ts", "value").unwrap(); + + let result = ensure_string_label_columns(&schema, "ts", "label"); + assert_eq!(supported, result.is_ok(), "{label_type:?}"); + if let Err(err) = result { + assert!(matches!( + err, + crate::error::Error::InvalidPromRemoteReadSchema { .. } + )); + assert_eq!(err.status_code(), StatusCode::InvalidArguments); + assert!( + err.to_string() + .starts_with("Invalid Prometheus remote read schema: ") + ); + assert!( + err.to_string() + .contains("only supports string labels: column 'value'") + ); + } + } + } + #[test] fn test_auto_create_table_type_for_prom_remote_write_metric_engine() { let mut query_ctx = QueryContext::with( diff --git a/src/servers/src/prom_store.rs b/src/servers/src/prom_store.rs index 0568d967010..b991cc9c524 100644 --- a/src/servers/src/prom_store.rs +++ b/src/servers/src/prom_store.rs @@ -148,15 +148,17 @@ pub fn validate_matchers(matchers: &[LabelMatcher]) -> Result<()> { .build() })?; if matches!(matcher_type, MatcherType::Re | MatcherType::Nre) { - // Validate using the same regex engine as query execution. - // Its syntax differs from Prometheus's; e.g. `(?x) a` is accepted here - // but rejected by Prometheus. - regex::Regex::new(&matcher.value).map_err(|e| { - error::InvalidPromRemoteRequestSnafu { - msg: format!("invalid regex for label '{}': {e}", matcher.name), - } - .build() - })?; + // Like Prometheus, validate the raw regex first, then the anchored pattern. + // We use the same regex engine as query execution, whose syntax differs + // from Prometheus's: e.g. `(?x) a` is accepted here but rejected by Prometheus. + regex::Regex::new(&matcher.value) + .and_then(|_| regex::Regex::new(&anchored_label_regex(&matcher.value))) + .map_err(|e| { + error::InvalidPromRemoteRequestSnafu { + msg: format!("invalid regex for label '{}': {e}", matcher.name), + } + .build() + })?; } } Ok(()) @@ -196,22 +198,61 @@ pub fn query_to_plan( .build() })?; - match m_type { + // Prometheus treats an absent label as "". + let label = if dataframe.schema().has_column_with_unqualified_name(name) { + ident(name) + } else { + lit("") + }; + + // Like PromQL, include NULL when the matcher accepts the empty string. + // Keep EQ/NEQ on the original column for index recognition. + let condition = match m_type { MatcherType::Eq => { - conditions.push(ident(name).eq(lit(value))); + if value.is_empty() { + label.clone().eq(lit(value)).or(label.is_null()) + } else { + label.eq(lit(value)) + } } MatcherType::Neq => { - conditions.push(ident(name).not_eq(lit(value))); + if value.is_empty() { + label.not_eq(lit(value)) + } else { + label.clone().not_eq(lit(value)).or(label.is_null()) + } } - // Case sensitive regexp match - MatcherType::Re => { - conditions.push(regexp_match(ident(name), lit(value), None).is_not_null()); + MatcherType::Re | MatcherType::Nre => { + let pattern = anchored_label_regex(value); + let matches_empty = regex::Regex::new(&pattern) + .map_err(|e| { + error::InvalidPromRemoteRequestSnafu { + msg: format!("invalid regex for label '{name}': {e}"), + } + .build() + })? + .is_match(""); + let matched = regexp_match(label.clone(), lit(pattern), None); + if m_type == MatcherType::Re { + let condition = matched.is_not_null(); + if matches_empty { + condition.or(label.is_null()) + } else { + condition + } + } else { + let condition = matched.is_null(); + // A NULL input also produces NULL. Keep it only when the + // negative matcher accepts the empty string. + if matches_empty { + condition.and(label.is_not_null()) + } else { + condition + } + } } - // Case sensitive regexp not match - MatcherType::Nre => { - conditions.push(regexp_match(ident(name), lit(value), None).is_null()); - } - } + }; + conditions.push(condition); } // Safety: conditions MUST not be empty, reduce always return Some(expr). @@ -224,6 +265,12 @@ pub fn query_to_plan( Ok(dataframe.into_parts().1) } +/// Prometheus's NewFastRegexMatcher also wraps patterns in ^(?s:...)$: +/// https://github.com/prometheus/prometheus/blob/main/model/labels/regexp.go +fn anchored_label_regex(value: &str) -> String { + format!("^(?s:{value})$") +} + #[inline] fn new_label(name: String, value: String) -> Label { Label { name, value } @@ -941,6 +988,32 @@ mod tests { const NEQ_TYPE: i32 = MatcherType::Neq as i32; const RE_TYPE: i32 = MatcherType::Re as i32; + #[test] + fn test_validate_matchers_regex_syntax() { + for op in [MatcherType::Re, MatcherType::Nre] { + for (value, valid) in [ + ("api|web", true), + ("(?s)api.web", true), + ("(?x) api", true), + ("(?x) api # comment", false), + ("[", false), + // Wrapping must not make an invalid original pattern valid. + ("a)|(?:b", false), + ] { + let matcher = LabelMatcher { + name: "job".into(), + value: value.into(), + r#type: op as i32, + }; + assert_eq!( + valid, + validate_matchers(&[matcher]).is_ok(), + "{op:?}: {value}" + ); + } + } + } + #[test] fn test_table_name() { let q = Query { @@ -1095,7 +1168,7 @@ mod tests { }, LabelMatcher { name: "job".to_string(), - value: "*prom*".to_string(), + value: ".*prom.*".to_string(), r#type: RE_TYPE, }, LabelMatcher { @@ -1123,12 +1196,210 @@ mod tests { let ts_col = greptime_timestamp(); let expected = format!( - "Filter: ?table?.{} >= TimestampMillisecond(1000, None) AND ?table?.{} <= TimestampMillisecond(2000, None) AND regexp_match(?table?.job, Utf8(\"*prom*\")) IS NOT NULL AND ?table?.instance != Utf8(\"localhost\") AND ?table?.service.name = Utf8(\"api\") AND ?table?.Region = Utf8(\"us\")\n TableScan: ?table?", + "Filter: ?table?.{} >= TimestampMillisecond(1000, None) AND ?table?.{} <= TimestampMillisecond(2000, None) AND regexp_match(?table?.job, Utf8(\"^(?s:.*prom.*)$\")) IS NOT NULL AND (?table?.instance != Utf8(\"localhost\") OR ?table?.instance IS NULL) AND ?table?.service.name = Utf8(\"api\") AND ?table?.Region = Utf8(\"us\")\n TableScan: ?table?", ts_col, ts_col ); assert_eq!(expected, display_string); } + #[tokio::test] + async fn test_remote_read_label_matcher_semantics() { + let ctx = SessionContext::new(); + let batch = DfRecordBatch::try_from_iter(vec![ + ( + greptime_timestamp(), + Arc::new(TimestampMillisecondArray::from(vec![1000; 6])) as ArrayRef, + ), + ( + "job", + Arc::new(StringArray::from(vec![ + Some("api"), + Some("my-api-server"), + Some(""), + None, + Some("web"), + Some("api\nweb"), + ])) as ArrayRef, + ), + ( + "id", + Arc::new(UInt32Array::from(vec![0, 1, 2, 3, 4, 5])) as ArrayRef, + ), + ]) + .unwrap(); + ctx.register_batch("metrics", batch).unwrap(); + for (name, op, value, expected) in [ + ("job", MatcherType::Eq, "api", vec![0]), + ("job", MatcherType::Neq, "api", vec![1, 2, 3, 4, 5]), + ("job", MatcherType::Eq, "", vec![2, 3]), + ("job", MatcherType::Neq, "", vec![0, 1, 4, 5]), + ("job", MatcherType::Re, "api", vec![0]), + ("job", MatcherType::Nre, "api", vec![1, 2, 3, 4, 5]), + ("job", MatcherType::Re, "api|web", vec![0, 4]), + ("job", MatcherType::Nre, "api|web", vec![1, 2, 3, 5]), + ("job", MatcherType::Re, "", vec![2, 3]), + ("job", MatcherType::Nre, "", vec![0, 1, 4, 5]), + ("job", MatcherType::Re, ".*", vec![0, 1, 2, 3, 4, 5]), + ("job", MatcherType::Nre, ".*", vec![]), + ("job", MatcherType::Re, ".+", vec![0, 1, 4, 5]), + ("job", MatcherType::Nre, ".+", vec![2, 3]), + ("job", MatcherType::Re, "(?m)api", vec![0]), + ("job", MatcherType::Re, "api.web", vec![5]), + ("job", MatcherType::Nre, "api.web", vec![0, 1, 2, 3, 4]), + ("job", MatcherType::Re, "(?-s:api.web)", vec![]), + ("job", MatcherType::Re, "(api)?", vec![0, 2, 3]), + ("missing", MatcherType::Eq, "", vec![0, 1, 2, 3, 4, 5]), + ("missing", MatcherType::Eq, "us", vec![]), + ("missing", MatcherType::Neq, "us", vec![0, 1, 2, 3, 4, 5]), + ("missing", MatcherType::Neq, "", vec![]), + ("missing", MatcherType::Re, ".*", vec![0, 1, 2, 3, 4, 5]), + ("missing", MatcherType::Re, ".+", vec![]), + ("missing", MatcherType::Nre, ".*", vec![]), + ("missing", MatcherType::Nre, ".+", vec![0, 1, 2, 3, 4, 5]), + ("missing", MatcherType::Re, "", vec![0, 1, 2, 3, 4, 5]), + ("missing", MatcherType::Nre, "", vec![]), + ] { + let query = Query { + start_timestamp_ms: 1000, + end_timestamp_ms: 1000, + matchers: vec![LabelMatcher { + name: name.into(), + value: value.into(), + r#type: op as i32, + }], + ..Default::default() + }; + let plan = query_to_plan( + ctx.table("metrics").await.unwrap(), + &query, + greptime_timestamp(), + ) + .unwrap_or_else(|err| panic!("{name} {op:?} {value:?}: {err}")); + let batches = ctx + .execute_logical_plan(plan) + .await + .unwrap() + .collect() + .await + .unwrap(); + let actual: Vec = batches + .iter() + .flat_map(|batch| { + batch + .column(2) + .as_primitive::() + .values() + .to_vec() + }) + .collect(); + assert_eq!(expected, actual, "{name} {op:?} {value:?}"); + } + } + + #[tokio::test] + async fn test_remote_read_matcher_combinations_and_label_types() { + let labels = StringArray::from(vec![Some("api"), Some("api"), None, Some("api")]); + for label_type in [ + DataType::Utf8, + DataType::LargeUtf8, + DataType::Utf8View, + DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)), + ] { + let ctx = SessionContext::new(); + let batch = DfRecordBatch::try_from_iter(vec![ + ( + greptime_timestamp(), + Arc::new(TimestampMillisecondArray::from(vec![999, 1000, 2000, 2001])) + as ArrayRef, + ), + ("Job.Name", casts::cast(&labels, &label_type).unwrap()), + ]) + .unwrap(); + ctx.register_batch("metrics", batch).unwrap(); + let mut cases = vec![ + (vec![("Job.Name", MatcherType::Re, ".*")], vec![1000, 2000]), + (vec![("Job.Name", MatcherType::Nre, ".+")], vec![2000]), + ]; + // Exercise combinations once; other representations only need regex/NULL coverage. + if label_type == DataType::Utf8 { + cases.extend([ + ( + vec![ + ("Job.Name", MatcherType::Re, "api|web"), + ("Job.Name", MatcherType::Neq, "web"), + ], + vec![1000], + ), + ( + vec![ + ("Job.Name", MatcherType::Eq, ""), + ("Job.Name", MatcherType::Neq, "api"), + ], + vec![2000], + ), + ( + vec![ + ("Job.Name", MatcherType::Eq, "api"), + ("Job.Name", MatcherType::Eq, "web"), + ], + vec![], + ), + ( + vec![ + ("job.name", MatcherType::Eq, ""), + ("Job.Name", MatcherType::Re, ".*"), + ], + vec![1000, 2000], + ), + ]); + } + for (matchers, expected) in cases { + let mut query = Query { + start_timestamp_ms: 1000, + end_timestamp_ms: 2000, + matchers: matchers + .iter() + .map(|(name, op, value)| LabelMatcher { + name: (*name).into(), + value: (*value).into(), + r#type: *op as i32, + }) + .collect(), + ..Default::default() + }; + query.matchers.push(LabelMatcher { + name: METRIC_NAME_LABEL.into(), + value: "metrics".into(), + r#type: EQ_TYPE, + }); + let plan = query_to_plan( + ctx.table("metrics").await.unwrap(), + &query, + greptime_timestamp(), + ) + .unwrap(); + let batches = ctx + .execute_logical_plan(plan) + .await + .unwrap() + .collect() + .await + .unwrap(); + let actual: Vec = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_primitive::() + .values() + .to_vec() + }) + .collect(); + assert_eq!(expected, actual, "{label_type:?}: {matchers:?}"); + } + } + } + fn column_schemas_with( mut kts_iter: Vec<(&str, ColumnDataType, SemanticType)>, ) -> Vec { diff --git a/tests-integration/src/prom_store.rs b/tests-integration/src/prom_store.rs index 80da82c0e89..87798c282c1 100644 --- a/tests-integration/src/prom_store.rs +++ b/tests-integration/src/prom_store.rs @@ -708,6 +708,12 @@ mod tests { (MatcherType::Neq as i32, "[", 200), (MatcherType::Re as i32, "spark", 200), (MatcherType::Nre as i32, "api.*", 200), + (MatcherType::Re as i32, "(?x) spark", 200), + (MatcherType::Nre as i32, "(?x) spark", 200), + (MatcherType::Re as i32, "(?x) spark # comment", 400), + (MatcherType::Nre as i32, "(?x) spark # comment", 400), + (MatcherType::Re as i32, "a)|(?:b", 400), + (MatcherType::Nre as i32, "a)|(?:b", 400), ] { let mut query = remote_query(metric); query.matchers.push(LabelMatcher { From 4684b95f72a86b563f9d43cd6e7952acc9392366 Mon Sep 17 00:00:00 2001 From: yihong Date: Mon, 19 Oct 2026 07:17:58 +0000 Subject: [PATCH 035/239] ci: do not run these actions on fork (#9721) Signed-off-by: yihong0618 --- .github/workflows/aws-ec2-runner-janitor.yml | 1 + .github/workflows/query-regression-janitor.yml | 1 + 2 files changed, 2 insertions(+) diff --git a/.github/workflows/aws-ec2-runner-janitor.yml b/.github/workflows/aws-ec2-runner-janitor.yml index 1c4435022b4..0355a1159d9 100644 --- a/.github/workflows/aws-ec2-runner-janitor.yml +++ b/.github/workflows/aws-ec2-runner-janitor.yml @@ -10,6 +10,7 @@ permissions: jobs: sweep: + if: ${{ github.repository == 'GreptimeTeam/greptimedb' }} runs-on: ubuntu-latest timeout-minutes: 30 steps: diff --git a/.github/workflows/query-regression-janitor.yml b/.github/workflows/query-regression-janitor.yml index ce277ce2f02..c615729206c 100644 --- a/.github/workflows/query-regression-janitor.yml +++ b/.github/workflows/query-regression-janitor.yml @@ -14,6 +14,7 @@ permissions: jobs: sweep: + if: ${{ github.repository == 'GreptimeTeam/greptimedb' }} runs-on: ubuntu-latest timeout-minutes: 15 steps: From 791f42924626d4ac16051316b39799d36aba94db Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Mon, 19 Oct 2026 07:34:50 +0000 Subject: [PATCH 036/239] fix(mito2): validate bulk insert batches against the region schema (#9698) * fix(mito2): validate bulk insert batches against the region schema Bulk inserts were only checked for the time index column, unlike row writes, which WriteRequest::check_schema checks against the region schema. A batch with columns of another type, nulls in the time index or in non-nullable columns, duplicate or unknown columns reached the region worker: a time index that is not a timestamp or holds only nulls made the worker fail, and other mismatches failed after the WAL append or were stored as written, depending on the SST format. Move the column checks of check_schema into a core shared by row writes and bulk inserts, and check bulk insert batches at the same points as rows: when the engine builds the worker request, and again in the worker if the region schema changed since. check_schema keeps its behavior. Bulk insert columns must have the column's Arrow type, except that tags may use any string encoding, as the memtables already convert them. Sparse batches must carry one encoded primary key column of binary without nulls. A batch that misses columns gets them filled by the worker even if it claims to match the current schema. The metric engine checks logical bulk inserts against the logical region before rewriting them into the physical region, as the rewrite copies columns without checking them and would encode a duplicated tag into the primary key twice. Signed-off-by: jeremyhi * fix(mito2): report an error when a write is interrupted by a panic WriteNotify answers its waiter when it is dropped and reported success unless an error had been set, also when it was dropped while unwinding from a panic in the middle of the write. Report an error in that case. Signed-off-by: jeremyhi * fix(mito2): tolerate a poisoned series lock when dropping a series map SeriesMap::drop reads every series to update the memtable metrics and unwrapped the lock, so dropping the map while unwinding from a panic that poisoned a series lock panicked again. Read through the poisoned lock instead. Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- src/metric-engine/src/engine/bulk_insert.rs | 95 ++++- src/mito2/src/engine.rs | 2 + .../src/engine/bulk_insert_validation_test.rs | 329 ++++++++++++++++ src/mito2/src/memtable/time_series.rs | 24 +- src/mito2/src/region_write_ctx.rs | 23 +- src/mito2/src/request.rs | 351 +++++++++++++----- src/mito2/src/worker/handle_bulk_insert.rs | 6 +- src/mito2/src/worker/handle_write.rs | 22 +- 8 files changed, 751 insertions(+), 101 deletions(-) create mode 100644 src/mito2/src/engine/bulk_insert_validation_test.rs diff --git a/src/metric-engine/src/engine/bulk_insert.rs b/src/metric-engine/src/engine/bulk_insert.rs index b7ad7646e3c..a9080745e2e 100644 --- a/src/metric-engine/src/engine/bulk_insert.rs +++ b/src/metric-engine/src/engine/bulk_insert.rs @@ -16,9 +16,11 @@ use std::collections::HashSet; use api::v1::{ArrowIpc, SemanticType}; use bytes::Bytes; +use common_error::ext::BoxedError; use common_grpc::flight::{FlightEncoder, FlightMessage}; use datatypes::arrow::record_batch::RecordBatch; -use snafu::{OptionExt, ensure}; +use mito2::request::check_bulk_insert_batch; +use snafu::{OptionExt, ResultExt, ensure}; use store_api::codec::PrimaryKeyEncoding; use store_api::metadata::RegionMetadataRef; use store_api::region_request::{AffectedRows, RegionBulkInsertsRequest, RegionRequest}; @@ -111,6 +113,11 @@ impl MetricEngineInner { &batch, &logical_metadata, )?; + // The rewrite below copies columns without checking them, and would encode a + // duplicated tag into the primary key twice. + check_bulk_insert_batch(&logical_metadata, &batch) + .map_err(BoxedError::new) + .context(error::MitoWriteOperationSnafu)?; let modified_batch = modify_batch_sparse( batch.clone(), region_id.table_id(), @@ -307,6 +314,92 @@ mod tests { }) } + #[tokio::test] + async fn test_bulk_insert_rejects_duplicate_tag() { + use common_error::status_code::StatusCode; + + let env = TestEnv::new().await; + env.init_metric_region().await; + let logical_region_id = env.default_logical_region_id(); + + let batch = build_logical_batch(0, 2); + let mut fields = batch.schema().fields().to_vec(); + let mut columns = batch.columns().to_vec(); + fields.push(fields[2].clone()); + columns.push(columns[2].clone()); + let batch = RecordBatch::try_new(Arc::new(ArrowSchema::new(fields)), columns).unwrap(); + let err = env + .metric() + .handle_request( + logical_region_id, + build_bulk_request(logical_region_id, batch, false), + ) + .await + .unwrap_err(); + assert_eq!(err.status_code(), StatusCode::InvalidArguments, "{err:?}"); + } + + #[tokio::test] + async fn test_bulk_insert_physical_region_rejects_malformed_primary_key() { + use common_error::status_code::StatusCode; + use datatypes::arrow::array::{ArrayRef, BinaryArray, Int64Array}; + use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME; + + let env = TestEnv::new().await; + env.init_metric_region().await; + let physical_region_id = env.default_physical_region_id(); + let logical_region_id = env.default_logical_region_id(); + let tag_columns = vec![TagColumnInfo { + name: "job".to_string(), + index: 2, + column_id: 2, + }]; + let batch = modify_batch_sparse( + build_logical_batch(0, 2), + logical_region_id.table_id(), + &tag_columns, + &[0, 1], + ) + .unwrap(); + let primary_key_index = batch.schema().index_of(PRIMARY_KEY_COLUMN_NAME).unwrap(); + let with_primary_key = |primary_key: Option, extra: Option| { + let mut fields = batch.schema().fields().to_vec(); + let mut columns = batch.columns().to_vec(); + fields.remove(primary_key_index); + columns.remove(primary_key_index); + for array in [primary_key, extra].into_iter().flatten() { + fields.push(Arc::new(Field::new( + PRIMARY_KEY_COLUMN_NAME, + array.data_type().clone(), + true, + ))); + columns.push(array); + } + RecordBatch::try_new(Arc::new(ArrowSchema::new(fields)), columns).unwrap() + }; + let encoded = batch.column(primary_key_index).clone(); + let rejected = [ + with_primary_key(None, None), + with_primary_key(Some(Arc::new(Int64Array::from(vec![1, 2]))), None), + with_primary_key( + Some(Arc::new(BinaryArray::from_opt_vec(vec![Some(b"a"), None]))), + None, + ), + with_primary_key(Some(encoded.clone()), Some(encoded)), + ]; + for batch in rejected { + let err = env + .metric() + .handle_request( + physical_region_id, + build_bulk_request(physical_region_id, batch, false), + ) + .await + .unwrap_err(); + assert_eq!(err.status_code(), StatusCode::InvalidArguments, "{err:?}"); + } + } + async fn init_dense_metric_region(env: &TestEnv) -> RegionId { let physical_region_id = env.default_physical_region_id(); env.create_physical_region( diff --git a/src/mito2/src/engine.rs b/src/mito2/src/engine.rs index 47752adbe75..3036c61e7ff 100644 --- a/src/mito2/src/engine.rs +++ b/src/mito2/src/engine.rs @@ -25,6 +25,8 @@ mod batch_catchup_test; #[cfg(test)] mod batch_open_test; #[cfg(test)] +mod bulk_insert_validation_test; +#[cfg(test)] mod bump_committed_sequence_test; #[cfg(test)] mod catchup_test; diff --git a/src/mito2/src/engine/bulk_insert_validation_test.rs b/src/mito2/src/engine/bulk_insert_validation_test.rs new file mode 100644 index 00000000000..9652152031f --- /dev/null +++ b/src/mito2/src/engine/bulk_insert_validation_test.rs @@ -0,0 +1,329 @@ +// 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. + +//! Bulk inserts checked against the region schema. + +use std::sync::Arc; + +use api::v1::{ArrowIpc, Rows}; +use common_error::ext::ErrorExt; +use common_error::status_code::StatusCode; +use common_recordbatch::{DfRecordBatch, RecordBatches}; +use common_test_util::flight::encode_to_flight_data; +use datatypes::arrow::array::{ + ArrayRef, Float64Array, Int64Array, LargeStringArray, ListArray, StringArray, StringViewArray, + TimestampMillisecondArray, TimestampSecondArray, +}; +use datatypes::arrow::datatypes::{Field, Int64Type, Schema}; +use datatypes::prelude::ConcreteDataType; +use store_api::region_engine::RegionEngine; +use store_api::region_request::{ + AlterKind, ModifyColumnType, RegionAlterRequest, RegionBulkInsertsRequest, RegionPutRequest, + RegionRequest, +}; +use store_api::storage::{RegionId, ScanRequest}; + +use crate::config::MitoConfig; +use crate::engine::MitoEngine; +use crate::request::WorkerRequest; +use crate::test_util::{CreateRequestBuilder, TestEnv, build_rows, reopen_region, rows_schema}; + +const SST_FORMATS: [&str; 2] = ["flat", "primary_key"]; + +fn column(name: &str, array: ArrayRef) -> (Field, ArrayRef) { + (Field::new(name, array.data_type().clone(), true), array) +} + +fn tag() -> (Field, ArrayRef) { + column("tag_0", Arc::new(StringArray::from(vec!["a", "b"]))) +} + +fn field() -> (Field, ArrayRef) { + column("field_0", Arc::new(Float64Array::from(vec![1.0, 2.0]))) +} + +fn ts() -> (Field, ArrayRef) { + column( + "ts", + Arc::new(TimestampMillisecondArray::from(vec![0, 1000])), + ) +} + +fn bulk_request(region_id: RegionId, columns: Vec<(Field, ArrayRef)>) -> RegionRequest { + RegionRequest::BulkInserts(bulk_insert(region_id, columns)) +} + +fn bulk_insert(region_id: RegionId, columns: Vec<(Field, ArrayRef)>) -> RegionBulkInsertsRequest { + let (fields, arrays): (Vec<_>, Vec<_>) = columns.into_iter().unzip(); + let payload = DfRecordBatch::try_new(Arc::new(Schema::new(fields)), arrays).unwrap(); + let (schema, record_batch) = encode_to_flight_data(payload.clone()); + RegionBulkInsertsRequest { + skip_wal: false, + region_id, + payload, + raw_data: ArrowIpc { + schema: schema.data_header, + data_header: record_batch.data_header, + payload: record_batch.data_body, + }, + partition_expr_version: None, + aligned_schema_version: None, + } +} + +async fn scan_num_rows(engine: &MitoEngine, region_id: RegionId) -> usize { + let stream = engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + batches.iter().map(|batch| batch.num_rows()).sum() +} + +/// Batches that do not match the region schema are rejected before they are +/// appended to the WAL, and the region keeps accepting writes. +#[tokio::test] +async fn test_bulk_insert_rejects_batches_not_matching_schema() { + common_telemetry::init_default_ut_logging(); + let list = ListArray::from_iter_primitive::(vec![ + Some(vec![Some(1)]), + Some(vec![Some(2)]), + ]); + let rejected = vec![ + vec![ + tag(), + field(), + column("ts", Arc::new(Int64Array::from(vec![0, 1000]))), + ], + vec![ + tag(), + field(), + column( + "ts", + Arc::new(TimestampMillisecondArray::from(vec![Some(0), None])), + ), + ], + vec![ + tag(), + field(), + column( + "ts", + Arc::new(TimestampMillisecondArray::from(vec![None, None])), + ), + ], + vec![ + tag(), + field(), + column("ts", Arc::new(TimestampSecondArray::from(vec![0, 1]))), + ], + vec![tag(), field()], + vec![ + tag(), + column("field_0", Arc::new(StringArray::from(vec!["1", "2"]))), + ts(), + ], + vec![tag(), column("field_0", Arc::new(list)), ts()], + vec![ + tag(), + column( + "field_0", + Arc::new(Float64Array::from(vec![Some(1.0), None])), + ), + ts(), + ], + vec![tag(), field(), field(), ts()], + vec![ + tag(), + field(), + ts(), + column("unknown", Arc::new(Int64Array::from(vec![1, 2]))), + ], + ]; + for (i, sst_format) in SST_FORMATS.into_iter().enumerate() { + let mut env = TestEnv::new().await; + let engine = env.create_engine(MitoConfig::default()).await; + let region_id = RegionId::new(1, i as u32 + 1); + let request = CreateRequestBuilder::new() + .all_not_null(true) + .insert_option("sst_format", sst_format) + .build(); + let column_schemas = rows_schema(&request); + let table_dir = request.table_dir.clone(); + let options = request.options.clone(); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + for columns in &rejected { + let err = engine + .handle_request(region_id, bulk_request(region_id, columns.clone())) + .await + .unwrap_err(); + assert_eq!( + StatusCode::InvalidArguments, + err.status_code(), + "{sst_format}: {err:?}" + ); + } + + let put = RegionRequest::Put(RegionPutRequest { + rows: Rows { + schema: column_schemas, + rows: build_rows(0, 3), + }, + hint: None, + skip_wal: false, + partition_expr_version: None, + }); + engine.handle_request(region_id, put).await.unwrap(); + // Reopening replays the WAL, which holds none of the rejected batches. + reopen_region(&engine, region_id, table_dir, true, options).await; + assert_eq!(3, scan_num_rows(&engine, region_id).await, "{sst_format}"); + } +} + +/// Tags may use any string encoding, which the memtables convert. +#[tokio::test] +async fn test_bulk_insert_accepts_tag_encodings() { + common_telemetry::init_default_ut_logging(); + let tags: [ArrayRef; 3] = [ + Arc::new(StringArray::from(vec!["a", "b"])), + Arc::new(LargeStringArray::from(vec!["a", "b"])), + Arc::new(StringViewArray::from(vec!["a", "b"])), + ]; + for (i, sst_format) in SST_FORMATS.into_iter().enumerate() { + let mut env = TestEnv::new().await; + let engine = env.create_engine(MitoConfig::default()).await; + for (j, tag) in tags.iter().enumerate() { + let region_id = RegionId::new(i as u32 + 1, j as u32 + 1); + let request = CreateRequestBuilder::new() + .insert_option("sst_format", sst_format) + .build(); + let table_dir = request.table_dir.clone(); + let options = request.options.clone(); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + let columns = vec![column("tag_0", tag.clone()), field(), ts()]; + let response = engine + .handle_request(region_id, bulk_request(region_id, columns)) + .await + .unwrap(); + assert_eq!(2, response.affected_rows, "{sst_format}: {tag:?}"); + reopen_region(&engine, region_id, table_dir, true, options).await; + assert_eq!( + 2, + scan_num_rows(&engine, region_id).await, + "{sst_format}: {tag:?}" + ); + } + } +} + +/// The worker checks a batch again if the region schema changed after the +/// request was checked. +#[tokio::test] +async fn test_bulk_insert_rechecks_batch_after_schema_change() { + common_telemetry::init_default_ut_logging(); + for (i, sst_format) in SST_FORMATS.into_iter().enumerate() { + let mut env = TestEnv::new().await; + let engine = env.create_engine(MitoConfig::default()).await; + let region_id = RegionId::new(1, i as u32 + 1); + let request = CreateRequestBuilder::new() + .insert_option("sst_format", sst_format) + .build(); + let table_dir = request.table_dir.clone(); + let options = request.options.clone(); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + let metadata = engine.get_region(region_id).unwrap().metadata(); + + // The batch matches the schema before `field_0` becomes a string. + let alter = RegionAlterRequest { + kind: AlterKind::ModifyColumnTypes { + columns: vec![ModifyColumnType { + column_name: "field_0".to_string(), + target_type: ConcreteDataType::string_datatype(), + }], + }, + }; + engine + .handle_request(region_id, RegionRequest::Alter(alter)) + .await + .unwrap(); + let (request, receiver) = WorkerRequest::try_from_region_request( + region_id, + bulk_request(region_id, vec![tag(), field(), ts()]), + Some(metadata), + ) + .unwrap(); + engine + .inner + .workers + .submit_to_worker(region_id, request) + .await + .unwrap(); + let err = receiver.await.unwrap().unwrap_err(); + assert_eq!( + StatusCode::InvalidArguments, + err.status_code(), + "{sst_format}: {err:?}" + ); + + reopen_region(&engine, region_id, table_dir, true, options).await; + assert_eq!(0, scan_num_rows(&engine, region_id).await, "{sst_format}"); + } +} + +/// A batch that claims to match the current schema but misses a column still +/// gets the column filled. +#[tokio::test] +async fn test_bulk_insert_fills_columns_missing_from_aligned_batch() { + common_telemetry::init_default_ut_logging(); + for (i, sst_format) in SST_FORMATS.into_iter().enumerate() { + let mut env = TestEnv::new().await; + let engine = env.create_engine(MitoConfig::default()).await; + let region_id = RegionId::new(1, i as u32 + 1); + let request = CreateRequestBuilder::new() + .insert_option("sst_format", sst_format) + .build(); + let table_dir = request.table_dir.clone(); + let options = request.options.clone(); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + let mut request = bulk_insert(region_id, vec![tag(), ts()]); + request.aligned_schema_version = Some( + engine + .get_region(region_id) + .unwrap() + .metadata() + .schema_version, + ); + let response = engine + .handle_request(region_id, RegionRequest::BulkInserts(request)) + .await + .unwrap(); + assert_eq!(2, response.affected_rows, "{sst_format}"); + reopen_region(&engine, region_id, table_dir, true, options).await; + assert_eq!(2, scan_num_rows(&engine, region_id).await, "{sst_format}"); + } +} diff --git a/src/mito2/src/memtable/time_series.rs b/src/mito2/src/memtable/time_series.rs index fad746927d4..1fa6cb24b15 100644 --- a/src/mito2/src/memtable/time_series.rs +++ b/src/mito2/src/memtable/time_series.rs @@ -17,7 +17,7 @@ use std::collections::{BTreeMap, Bound, HashSet}; use std::fmt::{Debug, Formatter}; use std::iter; use std::sync::atomic::{AtomicI64, AtomicU64, AtomicUsize, Ordering}; -use std::sync::{Arc, RwLock}; +use std::sync::{Arc, PoisonError, RwLock}; use std::time::{Duration, Instant}; use api::v1::OpType; @@ -474,7 +474,13 @@ impl Drop for SeriesMap { let num_field_builders = self .0 .values() - .map(|v| v.read().unwrap().active.num_field_builders()) + // The map may be dropped while unwinding from a panic that poisoned a lock. + .map(|v| { + v.read() + .unwrap_or_else(PoisonError::into_inner) + .active + .num_field_builders() + }) .sum::(); MEMTABLE_ACTIVE_SERIES_COUNT.sub(num_series as i64); MEMTABLE_ACTIVE_FIELD_BUILDER_COUNT.sub(num_field_builders as i64); @@ -1493,6 +1499,20 @@ mod tests { assert_eq!(expect, &read); } + #[test] + fn test_series_map_drop_with_poisoned_lock() { + let region_metadata = schema_for_test(); + let series = Arc::new(RwLock::new(Series::new(®ion_metadata))); + let map = SeriesMap(BTreeMap::from([(vec![1], series.clone())])); + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + let _guard = series.write().unwrap(); + panic!("poison the series lock"); + })) + .unwrap_err(); + assert!(series.is_poisoned()); + drop(map); + } + #[test] fn test_series() { let region_metadata = schema_for_test(); diff --git a/src/mito2/src/region_write_ctx.rs b/src/mito2/src/region_write_ctx.rs index 8c306effd1d..b0b781eb43c 100644 --- a/src/mito2/src/region_write_ctx.rs +++ b/src/mito2/src/region_write_ctx.rs @@ -23,7 +23,7 @@ use store_api::logstore::LogStore; use store_api::logstore::provider::Provider; use store_api::storage::{RegionId, SequenceNumber}; -use crate::error::{Error, Result, WriteGroupSnafu}; +use crate::error::{Error, Result, UnexpectedSnafu, WriteGroupSnafu}; use crate::memtable::KeyValues; use crate::memtable::bulk::part::BulkPart; use crate::metrics; @@ -66,6 +66,15 @@ impl WriteNotify { impl Drop for WriteNotify { fn drop(&mut self) { + // A panic may have interrupted the write, so it must not be reported as written. + if self.err.is_none() && std::thread::panicking() { + self.err = Some(Arc::new( + UnexpectedSnafu { + reason: "the write was interrupted by a panic", + } + .build(), + )); + } self.notify_result(); } } @@ -583,6 +592,18 @@ mod tests { use crate::memtable::bulk::part::BulkPart; use crate::test_util::version_util::VersionControlBuilder; + #[test] + fn test_write_notify_dropped_by_panic_reports_error() { + let (tx, rx) = oneshot::channel(); + let notify = WriteNotify::new(OptionOutputTx::from(tx), 3); + std::panic::catch_unwind(std::panic::AssertUnwindSafe(move || { + let _notify = notify; + panic!("interrupt the write"); + })) + .unwrap_err(); + assert!(rx.blocking_recv().unwrap().is_err()); + } + #[test] fn test_request_skip_wal_preserves_sequences_and_other_writes() { // Ablate only the request flag: the workload and sequence allocation stay identical. diff --git a/src/mito2/src/request.rs b/src/mito2/src/request.rs index c7ec359d6b0..45e9041006d 100644 --- a/src/mito2/src/request.rs +++ b/src/mito2/src/request.rs @@ -25,6 +25,9 @@ use api::helper::{ use api::v1::column_def::options_from_column_schema; use api::v1::{ColumnDataType, ColumnSchema, OpType, Rows, SemanticType, Value, WriteHint}; use common_telemetry::info; +use datatypes::arrow::array::{Array, ArrayRef}; +use datatypes::arrow::datatypes::DataType as ArrowDataType; +use datatypes::arrow::record_batch::RecordBatch; use datatypes::prelude::DataType; use partition::expr::PartitionExpr; use prometheus::HistogramTimer; @@ -44,6 +47,7 @@ use store_api::region_request::{ RegionFlushRequest, RegionOpenRequest, RegionRequest, RegionTruncateRequest, StagingPartitionDirective, }; +use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME; use store_api::storage::{FileId, RegionId, SequenceNumber}; use tokio::sync::oneshot::{self, Receiver, Sender}; @@ -194,85 +198,22 @@ impl WriteRequest { let region_id = self.region_id; // Index all columns in rows. - let mut rows_columns: HashMap<_, _> = self + let rows_columns = self .rows .schema .iter() - .map(|column| (&column.column_name, column)) + .map(|schema| { + let column = RowColumn { + request: self, + schema, + }; + (schema.column_name.as_str(), column) + }) .collect(); - - let mut need_fill_default = false; - // Checks all columns in this region. - for column in &metadata.column_metadatas { - if let Some(input_col) = rows_columns.remove(&column.column_schema.name) { - // Check data type. - ensure!( - is_column_type_value_eq( - input_col.datatype, - input_col.datatype_extension.clone(), - &column.column_schema.data_type - ), - InvalidRequestSnafu { - region_id, - reason: format!( - "column {} expect type {:?}, given: {}({})", - column.column_schema.name, - column.column_schema.data_type, - ColumnDataType::try_from(input_col.datatype) - .map(|v| v.as_str_name()) - .unwrap_or("Unknown"), - input_col.datatype, - ) - } - ); - - // Check semantic type. - ensure!( - is_semantic_type_eq(input_col.semantic_type, column.semantic_type), - InvalidRequestSnafu { - region_id, - reason: format!( - "column {} has semantic type {:?}, given: {}({})", - column.column_schema.name, - column.semantic_type, - api::v1::SemanticType::try_from(input_col.semantic_type) - .map(|v| v.as_str_name()) - .unwrap_or("Unknown"), - input_col.semantic_type - ), - } - ); - - // Check nullable. - // Safety: `rows_columns` ensures this column exists. - let has_null = self.has_null[self.name_to_index[&column.column_schema.name]]; - ensure!( - !has_null || column.column_schema.is_nullable(), - InvalidRequestSnafu { - region_id, - reason: format!( - "column {} is not null but input has null", - column.column_schema.name - ), - } - ); - } else { - // Rows don't have this column. - self.check_missing_column(column)?; - - need_fill_default = true; - } - } - - // Checks all columns in rows exist in the region. - if !rows_columns.is_empty() { - let names: Vec<_> = rows_columns.into_keys().collect(); - return InvalidRequestSnafu { - region_id, - reason: format!("unknown columns: {:?}", names), - } - .fail(); - } + let need_fill_default = + check_request_columns(region_id, metadata, rows_columns, |column| { + self.check_missing_column(column) + })?; // If we need to fill default values, return a special error. ensure!(!need_fill_default, FillDefaultSnafu { region_id }); @@ -369,16 +310,7 @@ impl WriteRequest { } // Not a delete request. Checks whether they have default value. - ensure!( - column.column_schema.is_nullable() - || column.column_schema.default_constraint().is_some(), - InvalidRequestSnafu { - region_id: self.region_id, - reason: format!("missing column {}", column.column_schema.name), - } - ); - - Ok(()) + check_column_can_be_filled(self.region_id, column) } /// Returns the default value for specific column. @@ -442,6 +374,242 @@ impl WriteRequest { } } +/// A column of a write request, as [check_request_columns] sees it. +trait RequestColumn { + /// Returns whether the column has the data type of `column`. + fn has_type_of(&self, column: &ColumnMetadata) -> bool; + + /// Describes the data type of the column in errors. + fn type_name(&self) -> String; + + /// Returns the semantic type of the column if the request carries one. + fn semantic_type(&self) -> Option; + + /// Returns whether the column contains nulls. + fn has_null(&self) -> bool; +} + +/// Checks the columns of a write request against the region schema. Each column +/// must be a column of the region, have its data type and, if the request carries +/// one, its semantic type, and contain nulls only if the region column is +/// nullable. `check_missing` checks the region columns that the request doesn't +/// have. Returns whether there are such columns. +fn check_request_columns( + region_id: RegionId, + metadata: &RegionMetadata, + mut columns: HashMap<&str, C>, + mut check_missing: impl FnMut(&ColumnMetadata) -> Result<()>, +) -> Result { + let mut has_missing = false; + for column in &metadata.column_metadatas { + let Some(input_col) = columns.remove(column.column_schema.name.as_str()) else { + check_missing(column)?; + has_missing = true; + continue; + }; + // Check data type. + ensure!( + input_col.has_type_of(column), + InvalidRequestSnafu { + region_id, + reason: format!( + "column {} expect type {:?}, given: {}", + column.column_schema.name, + column.column_schema.data_type, + input_col.type_name(), + ) + } + ); + + // Check semantic type. + if let Some(semantic_type) = input_col.semantic_type() { + ensure!( + is_semantic_type_eq(semantic_type, column.semantic_type), + InvalidRequestSnafu { + region_id, + reason: format!( + "column {} has semantic type {:?}, given: {}({})", + column.column_schema.name, + column.semantic_type, + api::v1::SemanticType::try_from(semantic_type) + .map(|v| v.as_str_name()) + .unwrap_or("Unknown"), + semantic_type + ), + } + ); + } + + // Check nullable. + ensure!( + !input_col.has_null() || column.column_schema.is_nullable(), + InvalidRequestSnafu { + region_id, + reason: format!( + "column {} is not null but input has null", + column.column_schema.name + ), + } + ); + } + + // Checks all columns in the request exist in the region. + if !columns.is_empty() { + let names: Vec<_> = columns.into_keys().collect(); + return InvalidRequestSnafu { + region_id, + reason: format!("unknown columns: {:?}", names), + } + .fail(); + } + + Ok(has_missing) +} + +/// Checks that a write may leave out `column`, which is then filled with its +/// default value or null. +fn check_column_can_be_filled(region_id: RegionId, column: &ColumnMetadata) -> Result<()> { + ensure!( + column.column_schema.is_nullable() || column.column_schema.default_constraint().is_some(), + InvalidRequestSnafu { + region_id, + reason: format!("missing column {}", column.column_schema.name), + } + ); + Ok(()) +} + +/// A column of the rows of a [WriteRequest]. +struct RowColumn<'a> { + request: &'a WriteRequest, + schema: &'a ColumnSchema, +} + +impl RequestColumn for RowColumn<'_> { + fn has_type_of(&self, column: &ColumnMetadata) -> bool { + is_column_type_value_eq( + self.schema.datatype, + self.schema.datatype_extension.clone(), + &column.column_schema.data_type, + ) + } + + fn type_name(&self) -> String { + format!( + "{}({})", + ColumnDataType::try_from(self.schema.datatype) + .map(|v| v.as_str_name()) + .unwrap_or("Unknown"), + self.schema.datatype, + ) + } + + fn semantic_type(&self) -> Option { + Some(self.schema.semantic_type) + } + + fn has_null(&self) -> bool { + // Safety: `name_to_index` has every column of the request. + self.request.has_null[self.request.name_to_index[&self.schema.column_name]] + } +} + +/// A column of the batch of a bulk insert. +struct BulkColumn<'a> { + array: &'a ArrayRef, +} + +impl RequestColumn for BulkColumn<'_> { + fn has_type_of(&self, column: &ColumnMetadata) -> bool { + let actual = self.array.data_type(); + let data_type = &column.column_schema.data_type; + // The memtables convert tags of any string encoding. + if column.semantic_type == SemanticType::Tag && data_type.is_string() { + return match actual { + ArrowDataType::Dictionary(key, value) => key.is_integer() && value.is_string(), + actual => actual.is_string(), + }; + } + *actual == data_type.as_arrow_type() + } + + fn type_name(&self) -> String { + self.array.data_type().to_string() + } + + fn semantic_type(&self) -> Option { + None + } + + fn has_null(&self) -> bool { + self.array.logical_null_count() > 0 + } +} + +/// Checks the batch of a bulk insert against the region schema with the rules +/// [WriteRequest::check_schema] applies to rows: column names are unique, every +/// column is a column of the region with its Arrow type and contains nulls only +/// if the region column is nullable, and a missing column is nullable or has a +/// default value. Tags may use any string encoding. Sparse batches carry the +/// tags in one encoded primary key column instead. Returns whether the batch +/// misses columns that the writer fills. +pub fn check_bulk_insert_batch(metadata: &RegionMetadata, batch: &RecordBatch) -> Result { + let region_id = metadata.region_id; + let sparse = metadata.primary_key_encoding == PrimaryKeyEncoding::Sparse; + let mut primary_key = None; + let mut columns = HashMap::with_capacity(batch.num_columns()); + for (field, array) in batch.schema_ref().fields().iter().zip(batch.columns()) { + let duplicate = if sparse && field.name() == PRIMARY_KEY_COLUMN_NAME { + primary_key.replace(array).is_some() + } else { + columns + .insert(field.name().as_str(), BulkColumn { array }) + .is_some() + }; + ensure!( + !duplicate, + InvalidRequestSnafu { + region_id, + reason: format!("duplicate column {}", field.name()), + } + ); + } + if sparse { + let primary_key = primary_key.with_context(|| InvalidRequestSnafu { + region_id, + reason: format!("missing column {PRIMARY_KEY_COLUMN_NAME}"), + })?; + // The writer casts the encoded primary key to a dictionary of binary. + let type_matches = match primary_key.data_type() { + ArrowDataType::Binary => true, + ArrowDataType::Dictionary(key, value) => { + key.is_integer() && **value == ArrowDataType::Binary + } + _ => false, + }; + ensure!( + type_matches && primary_key.logical_null_count() == 0, + InvalidRequestSnafu { + region_id, + reason: format!( + "column {PRIMARY_KEY_COLUMN_NAME} must be binary without nulls, given: {}", + primary_key.data_type() + ), + } + ); + } + let mut needs_fill = false; + check_request_columns(region_id, metadata, columns, |column| { + if sparse && column.semantic_type == SemanticType::Tag { + Ok(()) + } else { + needs_fill = true; + check_column_can_be_filled(region_id, column) + } + })?; + Ok(needs_fill) +} + /// Validate proto value schema. pub(crate) fn validate_proto_value( region_id: RegionId, @@ -762,7 +930,18 @@ impl WorkerRequest { sender: sender.into(), request: DdlRequest::EnterStaging(v), }), - RegionRequest::BulkInserts(region_bulk_inserts_request) => { + RegionRequest::BulkInserts(mut region_bulk_inserts_request) => { + // The worker answers empty batches without writing them. + if let Some(region_metadata) = ®ion_metadata + && region_bulk_inserts_request.payload.num_rows() > 0 + && check_bulk_insert_batch( + region_metadata, + ®ion_bulk_inserts_request.payload, + )? + { + // Lets the worker fill the missing columns. + region_bulk_inserts_request.aligned_schema_version = None; + } WorkerRequest::BulkInserts(BulkInsertRequest { metadata: region_metadata, sender: sender.into(), diff --git a/src/mito2/src/worker/handle_bulk_insert.rs b/src/mito2/src/worker/handle_bulk_insert.rs index 9ecbf537661..d48f06178c6 100644 --- a/src/mito2/src/worker/handle_bulk_insert.rs +++ b/src/mito2/src/worker/handle_bulk_insert.rs @@ -71,10 +71,12 @@ impl RegionWorkerLoop { return; } - // safety: ts data type must be a timestamp type. + // safety: `WorkerRequest::try_from_region_request` checked the batch against + // `region_metadata`, so the time index is a timestamp array without nulls. let (ts_primitive, _) = datatypes::timestamp::timestamp_array_to_primitive(ts).unwrap(); - // safety: we've checked ts.len() == batch.num_rows() and batch is not empty + // safety: we've checked ts.len() == batch.num_rows(), batch is not empty and + // the time index has no nulls. let min_ts = arrow::compute::min(&ts_primitive).unwrap(); let max_ts = arrow::compute::max(&ts_primitive).unwrap(); diff --git a/src/mito2/src/worker/handle_write.rs b/src/mito2/src/worker/handle_write.rs index 2b9551308f0..5de0439e26e 100644 --- a/src/mito2/src/worker/handle_write.rs +++ b/src/mito2/src/worker/handle_write.rs @@ -34,7 +34,9 @@ use crate::metrics::{ }; use crate::region::{RegionLeaderState, RegionRoleState}; use crate::region_write_ctx::{RegionWriteCtx, WriteSource}; -use crate::request::{SenderBulkRequest, SenderWriteRequest, WriteRequest}; +use crate::request::{ + SenderBulkRequest, SenderWriteRequest, WriteRequest, check_bulk_insert_batch, +}; use crate::wal::Wal; use crate::worker::RegionWorkerLoop; @@ -552,14 +554,16 @@ impl RegionWorkerLoop { aligned_schema.schema_version == region_ctx.version().metadata.schema_version }); - // Fill missing columns if needed - if need_fill_missing_columns - && let Err(e) = bulk_req - .request - .fill_missing_columns(®ion_ctx.version().metadata) - { - bulk_req.sender.send(Err(e)); - continue; + // Fill missing columns if needed. The batch is checked again as the region + // schema may have changed since the request was checked. + if need_fill_missing_columns { + let metadata = ®ion_ctx.version().metadata; + if let Err(e) = check_bulk_insert_batch(metadata, &bulk_req.request.batch) + .and_then(|_| bulk_req.request.fill_missing_columns(metadata)) + { + bulk_req.sender.send(Err(e)); + continue; + } } // Collect requests by region. From 52554db52fe36666240b9db1839057f6dfdc776d Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Mon, 19 Oct 2026 08:44:15 +0000 Subject: [PATCH 037/239] perf(mito2): keep short series unsplit in compaction merges (#9706) * perf(mito2): keep short series unsplit in compaction merges Compaction writes every merged batch to the output SST. Splitting short series before the merge made it emit one small batch per series and input file, which the SST writer and indexer then handled one by one. Compaction now splits only when series are long enough to fill most of a batch. With 16 L0 files of 10k series x 30 rows, one output went from 160,560 merged batches to 592, and merging plus writing took 1.66s instead of 2.67s. Signed-off-by: Lei, HUANG * test(mito2): benchmark compacting L0 files into one output The benchmark writes 16 L0 files through the engine, each holding every series over its own time span as consecutive flushes do, and times a strict window compaction that merges them into one output. Every iteration compacts a fresh copy of the same files. With 10k series of 30 rows per file, keeping short series unsplit in compaction merges cut the time from 830ms to 520ms. With 60 series of 5,000 rows per file, which compaction still splits, both took 266ms. Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- src/mito2/Cargo.toml | 5 + src/mito2/benches/bench_compaction_merge.rs | 236 ++++++++++++++++++++ src/mito2/src/compaction/reader.rs | 78 +++++++ src/mito2/src/read/scan_util.rs | 3 + src/mito2/src/read/seq_scan.rs | 11 +- 5 files changed, 330 insertions(+), 3 deletions(-) create mode 100644 src/mito2/benches/bench_compaction_merge.rs diff --git a/src/mito2/Cargo.toml b/src/mito2/Cargo.toml index 02c134e964c..ba6cbe5d8cb 100644 --- a/src/mito2/Cargo.toml +++ b/src/mito2/Cargo.toml @@ -154,3 +154,8 @@ required-features = ["testing"] name = "bench_dense_index_update" harness = false required-features = ["testing"] + +[[bench]] +name = "bench_compaction_merge" +harness = false +required-features = ["test"] diff --git a/src/mito2/benches/bench_compaction_merge.rs b/src/mito2/benches/bench_compaction_merge.rs new file mode 100644 index 00000000000..a4b9b6bc6fe --- /dev/null +++ b/src/mito2/benches/bench_compaction_merge.rs @@ -0,0 +1,236 @@ +// 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. + +//! Benchmarks compacting L0 files into one output through the engine. +//! +//! Every input file holds all series over its own time span, as consecutive +//! flushes do. Short series are where splitting the inputs by series hands the +//! SST writer one small batch per series and file. +//! +//! Run with: +//! ```sh +//! cargo bench -p mito2 --features test --bench bench_compaction_merge +//! ``` + +use std::collections::HashMap; +use std::path::Path; +use std::time::{Duration, Instant}; + +use api::v1::Rows; +use api::v1::region::{StrictWindow, compact_request}; +use common_test_util::temp_dir::create_temp_dir; +use criterion::{Criterion, SamplingMode, Throughput, criterion_group, criterion_main}; +use mito2::config::MitoConfig; +use mito2::engine::MitoEngine; +use mito2::test_util::{ + CreateRequestBuilder, TestEnv, build_rows_for_key, flush_region, rows_schema, +}; +use store_api::region_engine::{RegionEngine, RegionRole}; +use store_api::region_request::{ + PathType, RegionCompactRequest, RegionOpenRequest, RegionPutRequest, RegionRequest, +}; +use store_api::storage::RegionId; + +const REGION_ID: RegionId = RegionId::new(1, 1); +/// Series written by each put request. +const SERIES_PER_PUT: usize = 1000; +/// The inputs of every shape fit in one day, so each flush writes one file and +/// one compaction window covers all of them. +const WINDOW_SECONDS: i64 = 24 * 3600; + +struct Shape { + name: &'static str, + num_files: usize, + num_series: usize, + rows_per_series: usize, +} + +impl Shape { + fn num_rows(&self) -> usize { + self.num_files * self.num_series * self.rows_per_series + } +} + +const SHAPES: [Shape; 2] = [ + // Few enough series for scans to split the inputs, each far shorter than a batch. + Shape { + name: "10k_series_x_30_rows", + num_files: 16, + num_series: 10_000, + rows_per_series: 30, + }, + // Series long enough to fill most of a batch. + Shape { + name: "60_series_x_5000_rows", + num_files: 16, + num_series: 60, + rows_per_series: 5_000, + }, +]; + +/// L0 files of a shape and the options to open the region holding them. +struct Inputs { + env: TestEnv, + table_dir: String, + path_type: PathType, + options: HashMap, +} + +impl Inputs { + /// Opens a copy of the inputs in a new engine, so every compaction starts from the same files. + async fn open_copy(&self) -> (TestEnv, MitoEngine) { + let data_home = create_temp_dir("bench_compaction_merge"); + copy_dir(self.env.data_home(), data_home.path()); + let mut env = TestEnv::with_data_home(data_home).await; + let engine = env.create_engine(engine_config()).await; + let request = RegionOpenRequest { + engine: String::new(), + table_dir: self.table_dir.clone(), + path_type: self.path_type, + options: self.options.clone(), + skip_wal_replay: false, + checkpoint: None, + requirements: Default::default(), + }; + engine + .handle_request(REGION_ID, RegionRequest::Open(request)) + .await + .unwrap(); + engine + .set_region_role(REGION_ID, RegionRole::Leader) + .unwrap(); + (env, engine) + } +} + +fn engine_config() -> MitoConfig { + MitoConfig { + // Keeps automatic compaction from merging the inputs. + min_compaction_interval: Duration::from_secs(3600), + ..Default::default() + } +} + +/// Writes `shape.num_files` L0 files, each with every series over its own time span. +async fn write_inputs(shape: &Shape) -> Inputs { + let mut env = TestEnv::new().await; + let engine = env.create_engine(engine_config()).await; + let request = CreateRequestBuilder::new().build(); + let schema = rows_schema(&request); + let table_dir = request.table_dir.clone(); + let path_type = request.path_type; + let options = request.options.clone(); + engine + .handle_request(REGION_ID, RegionRequest::Create(request)) + .await + .unwrap(); + + for file in 0..shape.num_files { + let start = file * shape.rows_per_series; + let end = start + shape.rows_per_series; + for first in (0..shape.num_series).step_by(SERIES_PER_PUT) { + let last = (first + SERIES_PER_PUT).min(shape.num_series); + let rows = (first..last) + .flat_map(|series| build_rows_for_key(&format!("host-{series:05}"), start, end, 0)) + .collect(); + let request = RegionPutRequest { + rows: Rows { + schema: schema.clone(), + rows, + }, + // The rows only need to reach the SSTs. + skip_wal: true, + hint: None, + partition_expr_version: None, + }; + engine + .handle_request(REGION_ID, RegionRequest::Put(request)) + .await + .unwrap(); + } + flush_region(&engine, REGION_ID, None).await; + } + let statistic = engine.get_region_statistic(REGION_ID).unwrap(); + assert_eq!( + shape.num_files as u64, statistic.sst_num, + "every flush must leave its own input file" + ); + engine.stop().await.unwrap(); + + Inputs { + env, + table_dir, + path_type, + options, + } +} + +fn copy_dir(from: &Path, to: &Path) { + std::fs::create_dir_all(to).unwrap(); + for entry in std::fs::read_dir(from).unwrap() { + let entry = entry.unwrap(); + let target = to.join(entry.file_name()); + if entry.file_type().unwrap().is_dir() { + copy_dir(&entry.path(), &target); + } else { + std::fs::copy(entry.path(), target).unwrap(); + } + } +} + +fn compact_request() -> RegionRequest { + RegionRequest::Compact(RegionCompactRequest { + options: compact_request::Options::StrictWindow(StrictWindow { + window_seconds: WINDOW_SECONDS, + }), + ..Default::default() + }) +} + +fn bench_compaction_merge(c: &mut Criterion) { + let runtime = tokio::runtime::Runtime::new().unwrap(); + let mut group = c.benchmark_group("compaction_merge"); + group.sample_size(10); + group.sampling_mode(SamplingMode::Flat); + group.warm_up_time(Duration::from_secs(1)); + group.measurement_time(Duration::from_secs(20)); + for shape in &SHAPES { + let inputs = runtime.block_on(write_inputs(shape)); + group.throughput(Throughput::Elements(shape.num_rows() as u64)); + group.bench_function(shape.name, |b| { + b.iter_custom(|iterations| { + runtime.block_on(async { + let mut elapsed = Duration::ZERO; + for _ in 0..iterations { + let (_env, engine) = inputs.open_copy().await; + let start = Instant::now(); + engine + .handle_request(REGION_ID, compact_request()) + .await + .unwrap(); + elapsed += start.elapsed(); + let statistic = engine.get_region_statistic(REGION_ID).unwrap(); + assert_eq!(1, statistic.sst_num, "inputs must merge into one output"); + engine.stop().await.unwrap(); + } + elapsed + }) + }); + }); + } + group.finish(); +} + +criterion_group!(benches, bench_compaction_merge); +criterion_main!(benches); diff --git a/src/mito2/src/compaction/reader.rs b/src/mito2/src/compaction/reader.rs index 7fde251cb80..159ebe2f22a 100644 --- a/src/mito2/src/compaction/reader.rs +++ b/src/mito2/src/compaction/reader.rs @@ -242,3 +242,81 @@ fn ts_to_lit(ts: Timestamp, ts_col_unit: TimeUnit) -> Result { }; Ok(datafusion_expr::lit(scalar_value)) } + +#[cfg(test)] +mod tests { + use api::v1::Rows; + use store_api::region_engine::RegionEngine; + use store_api::region_request::{RegionFlushRequest, RegionRequest}; + use store_api::storage::RegionId; + + use super::*; + use crate::test_util::{ + CreateRequestBuilder, TestEnv, build_rows_for_key, put_rows, rows_schema, + }; + + #[tokio::test] + async fn test_compaction_reader_keeps_short_series_batches_whole() { + let mut env = TestEnv::new().await; + let engine = env.create_engine(Default::default()).await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let schema = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + // Each file has 100 short series, which queries split for merging. + for file in 0..2 { + let rows = (0..100) + .flat_map(|series| { + build_rows_for_key(&format!("key-{series:03}"), file * 10, file * 10 + 10, 0) + }) + .collect(); + put_rows( + &engine, + region_id, + Rows { + schema: schema.clone(), + rows, + }, + ) + .await; + engine + .handle_request( + region_id, + RegionRequest::Flush(RegionFlushRequest::default()), + ) + .await + .unwrap(); + } + let region = engine.get_region(region_id).unwrap(); + let version = region.version(); + let inputs: Vec<_> = version.ssts.levels()[0].files().cloned().collect(); + assert_eq!(2, inputs.len()); + + let mut source = CompactionSstReaderBuilder { + metadata: version.metadata.clone(), + sst_layer: region.access_layer.clone(), + cache: engine.cache_manager(), + inputs: &inputs, + append_mode: false, + filter_deleted: true, + time_range: None, + merge_mode: MergeMode::LastRow, + } + .build_flat_sst_reader() + .await + .unwrap(); + let mut num_batches = 0; + let mut num_rows = 0; + while let Some(batch) = source.next_batch().await.unwrap() { + num_batches += 1; + num_rows += batch.num_rows(); + } + + // Splitting by series would yield a batch per series and file, 200 in total. + assert_eq!(2000, num_rows); + assert!(num_batches < 20, "got {num_batches} batches"); + } +} diff --git a/src/mito2/src/read/scan_util.rs b/src/mito2/src/read/scan_util.rs index 7224143f2c4..53967949387 100644 --- a/src/mito2/src/read/scan_util.rs +++ b/src/mito2/src/read/scan_util.rs @@ -1204,6 +1204,9 @@ const NUM_SERIES_THRESHOLD: u64 = 10240; /// Minimum batch size after splitting. The batch size is less than 60 because a series may only have /// 60 samples per hour. const BATCH_SIZE_THRESHOLD: u64 = 50; +/// Minimum estimated rows per series for compaction to split batches. Shorter series +/// would make compaction write one small batch per series. +pub(crate) const MIN_COMPACTION_SPLIT_ROWS: usize = DEFAULT_READ_BATCH_SIZE / 2; /// Returns the estimated rows per batch after splitting if splitting flat record batches /// may improve merge performance. Returns `None` if splitting is not beneficial. diff --git a/src/mito2/src/read/seq_scan.rs b/src/mito2/src/read/seq_scan.rs index bbf6ff1bf92..24c35481669 100644 --- a/src/mito2/src/read/seq_scan.rs +++ b/src/mito2/src/read/seq_scan.rs @@ -46,8 +46,9 @@ use crate::read::range_cache::{ }; use crate::read::scan_region::{ScanInput, StreamContext}; use crate::read::scan_util::{ - PartitionMetrics, PartitionMetricsList, SplitRecordBatchStream, compute_parallel_channel_size, - scan_flat_file_ranges, scan_flat_mem_ranges, should_split_flat_batches_for_merge, + MIN_COMPACTION_SPLIT_ROWS, PartitionMetrics, PartitionMetricsList, SplitRecordBatchStream, + compute_parallel_channel_size, scan_flat_file_ranges, scan_flat_mem_ranges, + should_split_flat_batches_for_merge, }; use crate::read::stream::{ConvertBatchStream, ScanBatch, ScanBatchStream}; use crate::read::{BoxedRecordBatchStream, ScannerMetrics, scan_util}; @@ -739,7 +740,11 @@ pub(crate) async fn build_flat_sources( let needs_merge = ordered_sources.iter().flatten().count() > 1 || (compaction && stream_ctx.ranges.len() > 1); let split_batch_size = if needs_merge { - should_split_flat_batches_for_merge(stream_ctx, range_meta) + should_split_flat_batches_for_merge(stream_ctx, range_meta).filter(|rows_per_series| { + // Compaction writes every merged batch out, so splitting short series + // would hand the SST writer and indexer one tiny batch per series. + !compaction || *rows_per_series >= MIN_COMPACTION_SPLIT_ROWS + }) } else { None }; From b3dc5674e9245414b5470c8ca2d5eb9505c0b65d Mon Sep 17 00:00:00 2001 From: Yingwen Date: Mon, 19 Oct 2026 10:09:30 +0000 Subject: [PATCH 038/239] feat(mito2)!: account SST decoding against shared scan memory limits (#9101) * feat(mito2): add shared estimated SST scan memory budget Signed-off-by: evenyag * feat: unify scan memory limits with configurable accounting modes Signed-off-by: evenyag * refactor(mito2): reuse query runtime memory pool for series merges Signed-off-by: evenyag * test(mito2): wait for cancelled scan reservations before exhaustion check Signed-off-by: evenyag * docs(mito2): explain parallel scan cancellation cleanup Signed-off-by: evenyag * fix(mito2): deprecate waiting for scan memory quota Signed-off-by: evenyag * refactor(recordbatch): remove scan memory wait handling Signed-off-by: evenyag * fix(table): initialize memory pool in scan test fixtures Signed-off-by: evenyag --------- Signed-off-by: evenyag --- Cargo.lock | 1 + Cargo.toml | 2 + config/config.md | 20 +- config/datanode.example.toml | 36 ++- config/frontend.example.toml | 3 +- config/standalone.example.toml | 36 ++- src/cmd/src/datanode/scanbench.rs | 1 + src/cmd/tests/load_config_test.rs | 6 +- src/common/recordbatch/src/lib.rs | 366 ++++++---------------- src/common/recordbatch/src/memory.rs | 31 ++ src/mito2/AGENTS.md | 4 + src/mito2/src/config.rs | 77 ++++- src/mito2/src/engine.rs | 83 ++--- src/mito2/src/engine/scan_test.rs | 181 ++++++++++- src/mito2/src/read/scan_region.rs | 56 ++-- src/mito2/src/read/seq_scan.rs | 8 +- src/mito2/src/read/series_scan.rs | 38 +-- src/mito2/src/read/unordered_scan.rs | 3 +- src/mito2/src/sst/parquet/push_decoder.rs | 172 +++++++++- src/mito2/src/sst/parquet/reader.rs | 17 +- src/store-api/Cargo.toml | 1 + src/store-api/src/region_engine.rs | 16 +- src/table/src/table/scan.rs | 35 ++- tests-integration/tests/http.rs | 3 +- 24 files changed, 766 insertions(+), 430 deletions(-) create mode 100644 src/common/recordbatch/src/memory.rs diff --git a/Cargo.lock b/Cargo.lock index 47aa178a57f..9b895626f1b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -14412,6 +14412,7 @@ dependencies = [ "common-time", "common-wal", "datafusion-common", + "datafusion-execution", "datafusion-expr", "datafusion-physical-plan", "datatypes", diff --git a/Cargo.toml b/Cargo.toml index e2f91a54336..3cb4e77fb77 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -133,6 +133,7 @@ dashmap = "6.1" datafusion = "=55.1.0" datafusion-common = "=55.1.0" datafusion-datasource = "=55.1.0" +datafusion-execution = "=55.1.0" datafusion-expr = "=55.1.0" datafusion-expr-common = "=55.1.0" datafusion-functions = "=55.1.0" @@ -368,6 +369,7 @@ substrait = { git = "https://github.com/GreptimeTeam/substrait-rs.git", rev = "9 datafusion = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } datafusion-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } datafusion-datasource = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } +datafusion-execution = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } datafusion-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } datafusion-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } datafusion-functions = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } diff --git a/config/config.md b/config/config.md index 82e98fe96ec..8663ff1cc28 100644 --- a/config/config.md +++ b/config/config.md @@ -142,7 +142,7 @@ | `procedure.max_running_procedures` | Integer | `128` | Max running procedures.
The maximum number of procedures that can be running at the same time.
If the number of running procedures exceeds this limit, the procedure will be rejected. | | `query` | -- | -- | The query engine options. | | `query.parallelism` | Integer | `0` | Parallelism of the query engine.
Default to 0, which means the number of CPU cores. | -| `query.memory_pool_size` | String | 0 | Memory pool size for query execution operators (aggregation, sorting, join).
Supports absolute size (e.g., "2GB", "4GB") or percentage of system memory (e.g., "20%").
Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`.
When this limit is reached, queries will fail with ResourceExhausted error.
NOTE: This does NOT limit memory used by table scans. | +| `query.memory_pool_size` | String | 0 | Memory pool size for query execution operators (aggregation, sorting, join,
and internal series merges during table scans).
Supports absolute size (e.g., "2GB", "4GB") or percentage of system memory (e.g., "20%").
Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`.
When this limit is reached, queries will fail with ResourceExhausted error.
Projected Parquet row groups and streamed scan batches use scan_memory_limit separately. | | `query.experimental_memory_pool_policy` | String | `greedy` | Experimental memory pool allocation policy:
- "greedy" (default): first-come-first-served allocation; preserves current behavior.
- "fair": divides available memory among spillable operators and may spill earlier.
Only effective when `memory_pool_size` is bounded (>0). | | `query.experimental_spill_mode` | String | `default` | Spill mode:
- "default": preserve DataFusion built-in OS temp directory (default).
- "custom": explicitly configure spill path, quota, and compression.
- "disabled": explicitly disable disk spilling.
Set this to "custom" before using the path/quota/compression keys below. | | `query.experimental_spill_path` | String | Unset | Spill directory path. Ignored unless mode is "custom". | @@ -210,10 +210,11 @@ | `region_engine.mito.enable_refill_cache_on_read` | Bool | `true` | Enable refilling cache on read operations (default: true).
When disabled, cache refilling on read won't happen. | | `region_engine.mito.manifest_cache_size` | String | `256MB` | Capacity for manifest cache (default: 256MB). | | `region_engine.mito.sst_write_buffer_size` | String | `8MB` | Buffer size for SST writing. | -| `region_engine.mito.max_concurrent_scan_files` | Integer | `384` | Maximum number of SST files to scan concurrently. | +| `region_engine.mito.max_concurrent_scan_files` | Integer | `384` | Maximum number of SST files to scan concurrently when the memory limit is disabled
or a scan includes extension sources not covered by row-group accounting. | | `region_engine.mito.allow_stale_entries` | Bool | `false` | Whether to allow stale WAL entries read during replay. | -| `region_engine.mito.scan_memory_limit` | String | `unlimited` | Memory limit for table scans across all queries.
Supports absolute size (e.g., "2GB") or percentage of system memory (e.g., "20%").
Setting it to 0 or "unlimited" disables the limit. | -| `region_engine.mito.scan_memory_on_exhausted` | String | `fail` | Controls what happens when a scan cannot get memory immediately.
"fail" (default) fails fast and is the recommended option for most users.
"wait" / "wait()" waits for memory to become available. This is mainly
for advanced tuning in bursty workloads where temporary contention is common and
higher latency is acceptable.
"wait" means "wait(10s)", not unlimited waiting. | +| `region_engine.mito.scan_memory_limit` | String | `25%` | Memory limit for table scans across all queries and regions.
Defaults to 25% of cgroup memory, or host memory when no cgroup limit is available.
Supports absolute size (e.g., "2GB") or percentage of system memory (e.g., "20%").
Setting it to 0 or "unlimited" disables the limit and restores file-count checks. | +| `region_engine.mito.scan_memory_mode` | String | `basic` | Memory accounting mode: "basic" (default) counts projected Parquet columns'
uncompressed row-group sizes while decoding, including files from older versions.
"extended" also counts cumulative streamed batches until stream drop.
Limited reservations round up to KB; unlimited accounting is byte-exact.
Internal series merges use query.memory_pool_size in both modes, sharing capacity
with other query operators on the executing node. Neither budget measures total process memory. | +| `region_engine.mito.scan_memory_on_exhausted` | String | `fail` | Scan memory exhaustion always fails immediately.
Legacy "wait" / "wait()" settings are deprecated. They are accepted
for upgrade compatibility, but startup logs a warning and treats them as "fail".
Change existing wait settings to "fail"; scans no longer wait for quota to free up.
This policy applies only to scan reservations. Series merges use the query pool
and fail immediately when that pool is exhausted. | | `region_engine.mito.min_compaction_interval` | String | `0m` | Minimum time interval between two compactions.
To align with the old behavior, the default value is 0 (no restrictions). | | `region_engine.mito.schedule_compaction_after_edit` | Bool | `true` | Whether to allow to schedule a compaction after a successful region edit.

Setting this to "true" is a necessary but not sufficient condition for scheduling compaction after a region edit.
Other constraints, such as "min_compaction_interval", may still prevent compaction from being scheduled.
Setting this to "false", however, guarantees that compaction will not be scheduled after a region edit. | | `region_engine.mito.default_flat_format` | Bool | `true` | Whether to enable flat format as the default SST format. | @@ -387,7 +388,7 @@ | `query` | -- | -- | The query engine options. | | `query.parallelism` | Integer | `0` | Parallelism of the query engine.
Default to 0, which means the number of CPU cores. | | `query.allow_query_fallback` | Bool | `false` | Whether to allow query fallback when push down optimize fails.
Default to false, meaning when push down optimize failed, return error msg | -| `query.memory_pool_size` | String | 0 | Memory pool size for query execution operators (aggregation, sorting, join).
Supports absolute size (e.g., "4GB", "8GB") or percentage of system memory (e.g., "30%").
Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`.
When this limit is reached, queries will fail with ResourceExhausted error.
NOTE: This does NOT limit memory used by table scans (only applies to datanodes). | +| `query.memory_pool_size` | String | 0 | Memory pool size for query execution operators (aggregation, sorting, join).
Supports absolute size (e.g., "4GB", "8GB") or percentage of system memory (e.g., "30%").
Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`.
When this limit is reached, queries will fail with ResourceExhausted error.
This pool applies to operators on this frontend. Datanodes use their own query
pools for local operators and internal series merges during table scans. | | `query.experimental_memory_pool_policy` | String | `greedy` | Experimental memory pool allocation policy:
- "greedy" (default): first-come-first-served allocation; preserves current behavior.
- "fair": divides available memory among spillable operators and may spill earlier.
Only effective when `memory_pool_size` is bounded (>0). | | `query.experimental_spill_mode` | String | `default` | Spill mode:
- "default": preserve DataFusion built-in OS temp directory (default).
- "custom": explicitly configure spill path, quota, and compression.
- "disabled": explicitly disable disk spilling.
Set this to "custom" before using the path/quota/compression keys below. | | `query.experimental_spill_path` | String | Unset | Spill directory path. Ignored unless mode is "custom". | @@ -618,7 +619,7 @@ | `wal.accept_missing_history` | Bool | `false` | Whether a region opens although the WAL objects under the prefix do not hold its history.
An existing region fails to open when the prefix held no object when the WAL opened, or when the flushed entry id of the region names an object above every object of the prefix that holds entries. Both happen after `storage_provider`, or the bucket or root of the provider it names, changed without the WAL objects being copied, or after the objects were deleted.
To move the WAL, stop the instance, copy every object under the prefix to the new location and then change the configuration; the regions open with all their writes.
When the old WAL is lost, set this to `true` for one start: the regions open without their writes that were not flushed, which are lost, a warning names each region, and `greptime_logstore_object_store_wal_missing_history_accepted_total` counts them. Turn it off again once the instance has written to the WAL after that start; until then a later start refuses a region whose flushed entry id names an object above every object of the prefix that holds entries.
**It's only used when the provider is `experimental_object_store`**. | | `query` | -- | -- | The query engine options. | | `query.parallelism` | Integer | `0` | Parallelism of the query engine.
Default to 0, which means the number of CPU cores. | -| `query.memory_pool_size` | String | 0 | Memory pool size for query execution operators (aggregation, sorting, join).
Supports absolute size (e.g., "2GB", "4GB") or percentage of system memory (e.g., "20%").
Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`.
When this limit is reached, queries will fail with ResourceExhausted error.
NOTE: This does NOT limit memory used by table scans. | +| `query.memory_pool_size` | String | 0 | Memory pool size for query execution operators (aggregation, sorting, join,
and internal series merges during table scans).
Supports absolute size (e.g., "2GB", "4GB") or percentage of system memory (e.g., "20%").
Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`.
When this limit is reached, queries will fail with ResourceExhausted error.
Projected Parquet row groups and streamed scan batches use scan_memory_limit separately. | | `query.experimental_memory_pool_policy` | String | `greedy` | Experimental memory pool allocation policy:
- "greedy" (default): first-come-first-served allocation; preserves current behavior.
- "fair": divides available memory among spillable operators and may spill earlier.
Only effective when `memory_pool_size` is bounded (>0). | | `query.experimental_spill_mode` | String | `default` | Spill mode:
- "default": preserve DataFusion built-in OS temp directory (default).
- "custom": explicitly configure spill path, quota, and compression.
- "disabled": explicitly disable disk spilling.
Set this to "custom" before using the path/quota/compression keys below. | | `query.experimental_spill_path` | String | Unset | Spill directory path. Ignored unless mode is "custom". | @@ -687,10 +688,11 @@ | `region_engine.mito.enable_refill_cache_on_read` | Bool | `true` | Enable refilling cache on read operations (default: true).
When disabled, cache refilling on read won't happen. | | `region_engine.mito.manifest_cache_size` | String | `256MB` | Capacity for manifest cache (default: 256MB). | | `region_engine.mito.sst_write_buffer_size` | String | `8MB` | Buffer size for SST writing. | -| `region_engine.mito.max_concurrent_scan_files` | Integer | `384` | Maximum number of SST files to scan concurrently. | +| `region_engine.mito.max_concurrent_scan_files` | Integer | `384` | Maximum number of SST files to scan concurrently when the memory limit is disabled
or a scan includes extension sources not covered by row-group accounting. | | `region_engine.mito.allow_stale_entries` | Bool | `false` | Whether to allow stale WAL entries read during replay. | -| `region_engine.mito.scan_memory_limit` | String | `unlimited` | Memory limit for table scans across all queries.
Supports absolute size (e.g., "2GB") or percentage of system memory (e.g., "20%").
Setting it to 0 or "unlimited" disables the limit. | -| `region_engine.mito.scan_memory_on_exhausted` | String | `fail` | Controls what happens when a scan cannot get memory immediately.
"fail" (default) fails fast and is the recommended option for most users.
"wait" / "wait()" waits for memory to become available. This is mainly
for advanced tuning in bursty workloads where temporary contention is common and
higher latency is acceptable.
"wait" means "wait(10s)", not unlimited waiting. | +| `region_engine.mito.scan_memory_limit` | String | `25%` | Memory limit for table scans across all queries and regions.
Defaults to 25% of cgroup memory, or host memory when no cgroup limit is available.
Supports absolute size (e.g., "2GB") or percentage of system memory (e.g., "20%").
Setting it to 0 or "unlimited" disables the limit and restores file-count checks. | +| `region_engine.mito.scan_memory_mode` | String | `basic` | Memory accounting mode: "basic" (default) counts projected Parquet columns'
uncompressed row-group sizes while decoding, including files from older versions.
"extended" also counts cumulative streamed batches until stream drop.
Limited reservations round up to KB; unlimited accounting is byte-exact.
Internal series merges use query.memory_pool_size in both modes, sharing capacity
with other query operators on the executing node. Neither budget measures total process memory. | +| `region_engine.mito.scan_memory_on_exhausted` | String | `fail` | Scan memory exhaustion always fails immediately.
Legacy "wait" / "wait()" settings are deprecated. They are accepted
for upgrade compatibility, but startup logs a warning and treats them as "fail".
Change existing wait settings to "fail"; scans no longer wait for quota to free up.
This policy applies only to scan reservations. Series merges use the query pool
and fail immediately when that pool is exhausted. | | `region_engine.mito.min_compaction_interval` | String | `0m` | Minimum time interval between two compactions.
To align with the old behavior, the default value is 0 (no restrictions). | | `region_engine.mito.schedule_compaction_after_edit` | Bool | `true` | Whether to allow to schedule a compaction after a successful region edit.

Setting this to "true" is a necessary but not sufficient condition for scheduling compaction after a region edit.
Other constraints, such as "min_compaction_interval", may still prevent compaction from being scheduled.
Setting this to "false", however, guarantees that compaction will not be scheduled after a region edit. | | `region_engine.mito.default_flat_format` | Bool | `true` | Whether to enable flat format as the default SST format. | diff --git a/config/datanode.example.toml b/config/datanode.example.toml index d5f8e1e9be0..4b8eb41242e 100644 --- a/config/datanode.example.toml +++ b/config/datanode.example.toml @@ -337,11 +337,12 @@ overwrite_entry_start_id = false ## Default to 0, which means the number of CPU cores. parallelism = 0 -## Memory pool size for query execution operators (aggregation, sorting, join). +## Memory pool size for query execution operators (aggregation, sorting, join, +## and internal series merges during table scans). ## Supports absolute size (e.g., "2GB", "4GB") or percentage of system memory (e.g., "20%"). ## Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`. ## When this limit is reached, queries will fail with ResourceExhausted error. -## NOTE: This does NOT limit memory used by table scans. +## Projected Parquet row groups and streamed scan batches use scan_memory_limit separately. ## @toml2docs:none-default="0" memory_pool_size = "50%" @@ -665,22 +666,33 @@ manifest_cache_size = "256MB" ## Buffer size for SST writing. sst_write_buffer_size = "8MB" -## Maximum number of SST files to scan concurrently. +## Maximum number of SST files to scan concurrently when the memory limit is disabled +## or a scan includes extension sources not covered by row-group accounting. max_concurrent_scan_files = 384 ## Whether to allow stale WAL entries read during replay. allow_stale_entries = false -## Memory limit for table scans across all queries. +## Memory limit for table scans across all queries and regions. +## Defaults to 25% of cgroup memory, or host memory when no cgroup limit is available. ## Supports absolute size (e.g., "2GB") or percentage of system memory (e.g., "20%"). -## Setting it to 0 or "unlimited" disables the limit. -scan_memory_limit = "unlimited" -## Controls what happens when a scan cannot get memory immediately. -## "fail" (default) fails fast and is the recommended option for most users. -## "wait" / "wait()" waits for memory to become available. This is mainly -## for advanced tuning in bursty workloads where temporary contention is common and -## higher latency is acceptable. -## "wait" means "wait(10s)", not unlimited waiting. +## Setting it to 0 or "unlimited" disables the limit and restores file-count checks. +scan_memory_limit = "25%" + +## Memory accounting mode: "basic" (default) counts projected Parquet columns' +## uncompressed row-group sizes while decoding, including files from older versions. +## "extended" also counts cumulative streamed batches until stream drop. +## Limited reservations round up to KB; unlimited accounting is byte-exact. +## Internal series merges use query.memory_pool_size in both modes, sharing capacity +## with other query operators on the executing node. Neither budget measures total process memory. +scan_memory_mode = "basic" + +## Scan memory exhaustion always fails immediately. +## Legacy "wait" / "wait()" settings are deprecated. They are accepted +## for upgrade compatibility, but startup logs a warning and treats them as "fail". +## Change existing wait settings to "fail"; scans no longer wait for quota to free up. +## This policy applies only to scan reservations. Series merges use the query pool +## and fail immediately when that pool is exhausted. scan_memory_on_exhausted = "fail" ## Minimum time interval between two compactions. diff --git a/config/frontend.example.toml b/config/frontend.example.toml index 24aef8bd796..97a939c9829 100644 --- a/config/frontend.example.toml +++ b/config/frontend.example.toml @@ -370,7 +370,8 @@ allow_query_fallback = false ## Supports absolute size (e.g., "4GB", "8GB") or percentage of system memory (e.g., "30%"). ## Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`. ## When this limit is reached, queries will fail with ResourceExhausted error. -## NOTE: This does NOT limit memory used by table scans (only applies to datanodes). +## This pool applies to operators on this frontend. Datanodes use their own query +## pools for local operators and internal series merges during table scans. ## @toml2docs:none-default="0" memory_pool_size = "50%" diff --git a/config/standalone.example.toml b/config/standalone.example.toml index 6fba1d91398..e16210d3644 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -548,11 +548,12 @@ max_running_procedures = 128 ## Default to 0, which means the number of CPU cores. parallelism = 0 -## Memory pool size for query execution operators (aggregation, sorting, join). +## Memory pool size for query execution operators (aggregation, sorting, join, +## and internal series merges during table scans). ## Supports absolute size (e.g., "2GB", "4GB") or percentage of system memory (e.g., "20%"). ## Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`. ## When this limit is reached, queries will fail with ResourceExhausted error. -## NOTE: This does NOT limit memory used by table scans. +## Projected Parquet row groups and streamed scan batches use scan_memory_limit separately. ## @toml2docs:none-default="0" memory_pool_size = "50%" @@ -872,22 +873,33 @@ manifest_cache_size = "256MB" ## Buffer size for SST writing. sst_write_buffer_size = "8MB" -## Maximum number of SST files to scan concurrently. +## Maximum number of SST files to scan concurrently when the memory limit is disabled +## or a scan includes extension sources not covered by row-group accounting. max_concurrent_scan_files = 384 ## Whether to allow stale WAL entries read during replay. allow_stale_entries = false -## Memory limit for table scans across all queries. +## Memory limit for table scans across all queries and regions. +## Defaults to 25% of cgroup memory, or host memory when no cgroup limit is available. ## Supports absolute size (e.g., "2GB") or percentage of system memory (e.g., "20%"). -## Setting it to 0 or "unlimited" disables the limit. -scan_memory_limit = "unlimited" -## Controls what happens when a scan cannot get memory immediately. -## "fail" (default) fails fast and is the recommended option for most users. -## "wait" / "wait()" waits for memory to become available. This is mainly -## for advanced tuning in bursty workloads where temporary contention is common and -## higher latency is acceptable. -## "wait" means "wait(10s)", not unlimited waiting. +## Setting it to 0 or "unlimited" disables the limit and restores file-count checks. +scan_memory_limit = "25%" + +## Memory accounting mode: "basic" (default) counts projected Parquet columns' +## uncompressed row-group sizes while decoding, including files from older versions. +## "extended" also counts cumulative streamed batches until stream drop. +## Limited reservations round up to KB; unlimited accounting is byte-exact. +## Internal series merges use query.memory_pool_size in both modes, sharing capacity +## with other query operators on the executing node. Neither budget measures total process memory. +scan_memory_mode = "basic" + +## Scan memory exhaustion always fails immediately. +## Legacy "wait" / "wait()" settings are deprecated. They are accepted +## for upgrade compatibility, but startup logs a warning and treats them as "fail". +## Change existing wait settings to "fail"; scans no longer wait for quota to free up. +## This policy applies only to scan reservations. Series merges use the query pool +## and fail immediately when that pool is exhausted. scan_memory_on_exhausted = "fail" ## Minimum time interval between two compactions. diff --git a/src/cmd/src/datanode/scanbench.rs b/src/cmd/src/datanode/scanbench.rs index 9fbde16fb7b..667a5f54be1 100644 --- a/src/cmd/src/datanode/scanbench.rs +++ b/src/cmd/src/datanode/scanbench.rs @@ -1017,6 +1017,7 @@ impl ScanbenchCommand { let num_partitions = scanner.properties().partitions.len(); let ctx = QueryScanContext { explain_verbose: collect_scanner_explain, + ..Default::default() }; let metrics_set = ExecutionPlanMetricsSet::new(); diff --git a/src/cmd/tests/load_config_test.rs b/src/cmd/tests/load_config_test.rs index 00cdafd9b8a..2781183c0c8 100644 --- a/src/cmd/tests/load_config_test.rs +++ b/src/cmd/tests/load_config_test.rs @@ -201,7 +201,8 @@ fn test_load_datanode_example_config() { auto_flush_interval: Duration::from_secs(10 * 60), default_region_write_buffer_size: ReadableSize::mb(0), write_cache_ttl: Some(Duration::from_secs(60 * 60 * 8)), - scan_memory_limit: MemoryLimit::Unlimited, + scan_memory_limit: MemoryLimit::Percentage(25), + scan_memory_mode: mito2::config::ScanMemoryMode::Basic, ..Default::default() }), RegionEngineConfig::File(FileEngineConfig {}), @@ -440,7 +441,8 @@ fn test_load_standalone_example_config() { auto_flush_interval: Duration::from_secs(10 * 60), default_region_write_buffer_size: ReadableSize::mb(0), write_cache_ttl: Some(Duration::from_secs(60 * 60 * 8)), - scan_memory_limit: MemoryLimit::Unlimited, + scan_memory_limit: MemoryLimit::Percentage(25), + scan_memory_mode: mito2::config::ScanMemoryMode::Basic, ..Default::default() }), RegionEngineConfig::File(FileEngineConfig {}), diff --git a/src/common/recordbatch/src/lib.rs b/src/common/recordbatch/src/lib.rs index 0ccb8a4c28e..a3455b366dd 100644 --- a/src/common/recordbatch/src/lib.rs +++ b/src/common/recordbatch/src/lib.rs @@ -17,21 +17,18 @@ pub mod cursor; pub mod error; pub mod ext; pub mod filter; +mod memory; pub mod recordbatch; pub mod util; use std::fmt::{self, Write}; -use std::future::Future; use std::pin::Pin; use std::sync::Arc; use adapter::RecordBatchMetrics; use arc_swap::ArcSwapOption; use common_base::readable_size::ReadableSize; -use common_error::ext::BoxedError; -use common_memory_manager::{ - MemoryGuard, MemoryManager, MemoryMetrics, OnExhaustedPolicy, PermitGranularity, -}; +use common_memory_manager::{MemoryManager, MemoryMetrics, PermitGranularity}; use common_telemetry::tracing::Span; pub use datafusion::physical_plan::SendableRecordBatchStream as DfSendableRecordBatchStream; use datatypes::arrow::array::{Array, ArrayRef, AsArray, StringBuilder}; @@ -53,9 +50,10 @@ use error::Result; use futures::task::{Context, Poll}; use futures::{Stream, TryStreamExt}; pub use recordbatch::RecordBatch; -use snafu::{IntoError, ResultExt, ensure}; +use snafu::{ResultExt, ensure}; use crate::error::{ArrowComputeSnafu, NewDfRecordBatchSnafu}; +pub use crate::memory::QueryMemoryReservation; pub trait RecordBatchStream: Stream> { fn name(&self) -> &str { @@ -595,14 +593,14 @@ impl> + Unpin> Stream for RecordBatchStream } } -/// Memory tracker for RecordBatch streams. Clone to share the same limit across queries. -/// -/// Each stream acquires quota independently from this tracker. +/// Memory tracker for RecordBatch streams and reader reservations. Clone to share +/// the same limit across queries. Limited accounting rounds up to KB; unlimited +/// accounting is byte-exact. Exhaustion fails immediately. DataFusion operators use +/// the separate query runtime pool. #[derive(Clone)] pub struct QueryMemoryTracker { manager: MemoryManager, metrics: CallbackMemoryMetrics, - on_exhausted_policy: OnExhaustedPolicy, } impl fmt::Debug for QueryMemoryTracker { @@ -610,7 +608,6 @@ impl fmt::Debug for QueryMemoryTracker { f.debug_struct("QueryMemoryTracker") .field("current", &self.current()) .field("limit", &self.limit()) - .field("on_exhausted_policy", &self.on_exhausted_policy) .field("on_update", &self.metrics.has_on_update()) .field("on_exhausted", &self.metrics.has_on_exhausted()) .field("on_rejected", &self.metrics.has_on_rejected()) @@ -620,13 +617,9 @@ impl fmt::Debug for QueryMemoryTracker { impl QueryMemoryTracker { /// Create a builder for a query memory tracker. - pub fn builder( - limit: usize, - on_exhausted_policy: OnExhaustedPolicy, - ) -> QueryMemoryTrackerBuilder { + pub fn builder(limit: usize) -> QueryMemoryTrackerBuilder { QueryMemoryTrackerBuilder { limit, - on_exhausted_policy, on_update: None, on_exhausted: None, on_reject: None, @@ -636,7 +629,10 @@ impl QueryMemoryTracker { fn new_stream_tracker(&self) -> StreamMemoryTracker { StreamMemoryTracker { tracker: self.clone(), - guard: self.manager.try_acquire(0).unwrap(), + guard: QueryMemoryReservation { + // A zero-byte acquisition always succeeds; this manager is never closed. + inner: self.manager.try_acquire(0).unwrap(), + }, tracked_bytes: 0, } } @@ -645,10 +641,22 @@ impl QueryMemoryTracker { self.manager.used_bytes() as usize } - fn limit(&self) -> usize { + /// Returns the memory limit in bytes, or zero when unlimited. + pub fn limit(&self) -> usize { self.manager.limit_bytes() as usize } + /// Reserves bytes until the returned guard is dropped, failing immediately if + /// the shared budget cannot accommodate the reservation. + pub fn reserve(&self, bytes: usize) -> Result { + let mut tracker = self.new_stream_tracker(); + if let Err(error) = tracker.try_track(bytes) { + tracker.inc_rejected(); + return Err(error); + } + Ok(tracker.guard) + } + fn reject_error( &self, current: usize, @@ -660,7 +668,7 @@ impl QueryMemoryTracker { "{} requested, {} used globally ({}%), {} used by this stream, hard limit: {}", ReadableSize(additional as u64), ReadableSize(current as u64), - (current * 100).checked_div(limit).unwrap_or(0), + current.saturating_mul(100).checked_div(limit).unwrap_or(0), ReadableSize(stream_tracked as u64), ReadableSize(limit as u64) ); @@ -675,7 +683,6 @@ impl QueryMemoryTracker { /// Builder for constructing a [`QueryMemoryTracker`] with optional callbacks. pub struct QueryMemoryTrackerBuilder { limit: usize, - on_exhausted_policy: OnExhaustedPolicy, on_update: Option, on_exhausted: Option, on_reject: Option, @@ -687,7 +694,7 @@ impl QueryMemoryTrackerBuilder { /// /// # Note /// The callback is called after both successful `track()` and stream drop. - /// Usage is exact in unlimited mode and 1KB-aligned in limited mode. + /// Usage is byte-exact in unlimited mode; otherwise it is 1KB-aligned. pub fn on_update(mut self, on_update: F) -> Self where F: Fn(usize) + Send + Sync + 'static, @@ -700,7 +707,6 @@ impl QueryMemoryTrackerBuilder { /// /// # Note /// This is called when the non-blocking allocation fast path fails. - /// Requests using `OnExhaustedPolicy::Wait` may still succeed after waiting. /// It is never called when `limit == 0` (unlimited mode). pub fn on_exhausted(mut self, on_exhausted: F) -> Self where @@ -728,22 +734,16 @@ impl QueryMemoryTrackerBuilder { metrics.clone(), ); - QueryMemoryTracker { - manager, - metrics, - on_exhausted_policy: self.on_exhausted_policy, - } + QueryMemoryTracker { manager, metrics } } } struct StreamMemoryTracker { tracker: QueryMemoryTracker, - guard: MemoryGuard, + guard: QueryMemoryReservation, tracked_bytes: usize, } -type MemoryAcquireResult = std::result::Result<(), common_memory_manager::Error>; - impl StreamMemoryTracker { fn inc_rejected(&self) { self.tracker.inc_rejected(); @@ -758,51 +758,13 @@ impl StreamMemoryTracker { } } - async fn track_with_policy(mut self, additional: usize) -> (Self, MemoryAcquireResult) { - let result = self - .guard - .acquire_additional_with_policy(additional as u64, self.tracker.on_exhausted_policy) - .await; - if result.is_ok() { - self.tracked_bytes = self.tracked_bytes.saturating_add(additional); - } - (self, result) - } - fn reject_error(&self, additional: usize) -> error::Error { let current = self.tracker.current(); self.tracker .reject_error(current, additional, self.tracked_bytes) } - - fn wait_error(&self, additional: usize, source: common_memory_manager::Error) -> error::Error { - match source { - common_memory_manager::Error::MemoryLimitExceeded { .. } => { - self.reject_error(additional) - } - common_memory_manager::Error::MemoryAcquireTimeout { waited, .. } => { - let current = self.tracker.current(); - let limit = self.tracker.limit(); - let msg = format!( - "timed out waiting {:?} for {}, {} used globally ({}%), {} used by this stream, hard limit: {}", - waited, - ReadableSize(additional as u64), - ReadableSize(current as u64), - (current * 100).checked_div(limit).unwrap_or(0), - ReadableSize(self.tracked_bytes as u64), - ReadableSize(limit as u64) - ); - error::ExceedMemoryLimitSnafu { msg }.build() - } - error => error::ExternalSnafu.into_error(BoxedError::new(error)), - } - } } -type PendingTrackFuture = Pin< - Box + Send>, ->; - #[derive(Clone)] struct CallbackMemoryMetrics { inner: Arc, @@ -871,99 +833,24 @@ impl MemoryMetrics for CallbackMemoryMetrics { /// A wrapper stream that tracks memory usage of RecordBatches. pub struct MemoryTrackedStream { inner: SendableRecordBatchStream, - tracker: Option, - // Waiting stores a batch that has already been pulled from the inner stream but has not yet - // acquired additional quota. This keeps `poll_next()` non-blocking and allows bounded waits, - // at the cost of temporarily holding one untracked batch per blocked stream in memory. - waiting: Option, + tracker: StreamMemoryTracker, } impl MemoryTrackedStream { pub fn new(inner: SendableRecordBatchStream, tracker: QueryMemoryTracker) -> Self { Self { inner, - tracker: Some(tracker.new_stream_tracker()), - waiting: None, + tracker: tracker.new_stream_tracker(), } } - fn ready_tracker_mut(&mut self) -> &mut StreamMemoryTracker { - debug_assert!( - self.waiting.is_none(), - "a ready tracker must not coexist with a waiting future" - ); - self.tracker.as_mut().unwrap() - } - - fn enter_waiting(&mut self, batch: RecordBatch, additional: usize) { - debug_assert!( - self.waiting.is_none(), - "enter_waiting should only be called from the ready state" - ); - debug_assert!( - self.tracker.is_some(), - "enter_waiting requires a tracker in the ready state" - ); - let tracker = self.tracker.take().unwrap(); - self.waiting = Some(Self::start_waiting(tracker, batch, additional)); - } - - fn start_waiting( - tracker: StreamMemoryTracker, - batch: RecordBatch, - additional: usize, - ) -> PendingTrackFuture { - Box::pin(async move { - let (tracker, result) = tracker.track_with_policy(additional).await; - (tracker, batch, additional, result) - }) - } - - fn poll_waiting(&mut self, cx: &mut Context<'_>) -> Poll>> { - let future = self.waiting.as_mut().unwrap(); - match future.as_mut().poll(cx) { - Poll::Ready((tracker, batch, additional, result)) => { - let output = match result { - Ok(()) => Ok(batch), - Err(error) => { - tracker.inc_rejected(); - Err(tracker.wait_error(additional, error)) - } - }; - self.waiting = None; - self.tracker = Some(tracker); - Poll::Ready(Some(output)) - } - Poll::Pending => Poll::Pending, - } - } - - fn poll_batch( - &mut self, - batch: RecordBatch, - cx: &mut Context<'_>, - ) -> Poll>> { + fn track_batch(&mut self, batch: RecordBatch) -> Result { let additional = batch.logical_slice_memory_size(); - let tracker = self.ready_tracker_mut(); - - if let Err(error) = tracker.try_track(additional) { - match tracker.tracker.on_exhausted_policy { - OnExhaustedPolicy::Fail => { - tracker.inc_rejected(); - return Poll::Ready(Some(Err(error))); - } - // `Wait` is a deliberate tradeoff: the batch has already been materialized, so we - // keep it in memory while waiting for quota instead of failing immediately. Under - // contention, real memory usage can therefore exceed `scan_memory_limit` by up to - // one buffered batch per blocked stream. - OnExhaustedPolicy::Wait { .. } => { - self.enter_waiting(batch, additional); - return self.poll_waiting(cx); - } - } + if let Err(error) = self.tracker.try_track(additional) { + self.tracker.inc_rejected(); + return Err(error); } - - Poll::Ready(Some(Ok(batch))) + Ok(batch) } } @@ -971,12 +858,8 @@ impl Stream for MemoryTrackedStream { type Item = Result; fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - if self.waiting.is_some() { - return self.poll_waiting(cx); - } - match Pin::new(&mut self.inner).poll_next(cx) { - Poll::Ready(Some(Ok(batch))) => self.poll_batch(batch, cx), + Poll::Ready(Some(Ok(batch))) => Poll::Ready(Some(self.track_batch(batch))), Poll::Ready(Some(Err(error))) => Poll::Ready(Some(Err(error))), Poll::Ready(None) => Poll::Ready(None), Poll::Pending => Poll::Pending, @@ -1006,9 +889,8 @@ impl RecordBatchStream for MemoryTrackedStream { mod tests { use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; - use std::time::Duration; - use common_memory_manager::{OnExhaustedPolicy, PermitGranularity}; + use common_memory_manager::PermitGranularity; use datatypes::arrow::array::{ DictionaryArray, Int32Array, ListArray, StringArray, UInt32Array, }; @@ -1020,7 +902,6 @@ mod tests { use datatypes::schema::{ColumnSchema, Schema}; use datatypes::vectors::{BooleanVector, Int32Vector, StringVector}; use futures::StreamExt; - use tokio::time::{sleep, timeout}; use super::*; @@ -1041,6 +922,51 @@ mod tests { as usize } + #[test] + fn test_reader_reservation_rejects_and_releases_capacity() { + let tracker = QueryMemoryTracker::builder(2048).build(); + let first = tracker.reserve(1025).unwrap(); + assert_eq!(2048, tracker.current()); + assert!(tracker.reserve(1).is_err()); + assert_eq!(2048, tracker.current()); + drop(first); + assert_eq!(0, tracker.current()); + + assert!(tracker.reserve(2049).is_err()); + assert_eq!(0, tracker.current()); + let guard = tracker.reserve(2048).unwrap(); + assert_eq!(2048, tracker.current()); + drop(guard); + assert_eq!(0, tracker.current()); + } + + #[tokio::test] + async fn test_reader_and_batch_reservations_share_scan_budget() { + let batch = large_string_batch(512); + let batch_bytes = aligned_tracked_bytes(batch.logical_slice_memory_size()); + let tracker = QueryMemoryTracker::builder(batch_bytes + 1024).build(); + let reader = tracker.reserve(1).unwrap(); + let mut stream = MemoryTrackedStream::new( + RecordBatches::try_new(batch.schema.clone(), vec![batch.clone(), batch]) + .unwrap() + .as_stream(), + tracker.clone(), + ); + stream.next().await.unwrap().unwrap(); + assert_eq!(tracker.limit(), tracker.current()); + assert!(tracker.reserve(1).is_err()); + assert!(stream.next().await.unwrap().is_err()); + drop(reader); + assert_eq!(batch_bytes, tracker.current()); + drop(stream); + assert_eq!(0, tracker.current()); + let unlimited = QueryMemoryTracker::builder(0).build(); + let guard = unlimited.reserve(17).unwrap(); + assert_eq!(17, unlimited.current()); + drop(guard); + assert_eq!(0, unlimited.current()); + } + #[tokio::test] async fn test_memory_tracked_stream_charges_logical_slice_size() { let schema = Arc::new(Schema::new(vec![ColumnSchema::new( @@ -1057,7 +983,7 @@ mod tests { .unwrap(); let expected_bytes = aligned_tracked_bytes(batch.logical_slice_memory_size()); assert!(expected_bytes < aligned_tracked_bytes(batch.buffer_memory_size())); - let tracker = QueryMemoryTracker::builder(MB, OnExhaustedPolicy::Fail).build(); + let tracker = QueryMemoryTracker::builder(MB).build(); let mut stream = MemoryTrackedStream::new( RecordBatches::try_new(batch.schema.clone(), vec![batch]) .unwrap() @@ -1282,8 +1208,7 @@ mod tests { #[test] fn test_query_memory_tracker_basic() { - let tracker = - Arc::new(QueryMemoryTracker::builder(10 * MB, OnExhaustedPolicy::Fail).build()); + let tracker = Arc::new(QueryMemoryTracker::builder(10 * MB).build()); let mut stream1 = tracker.new_stream_tracker(); assert!(stream1.try_track(5 * MB).is_ok()); @@ -1300,8 +1225,7 @@ mod tests { #[test] fn test_query_memory_tracker_shared_global_limit() { - let tracker = - Arc::new(QueryMemoryTracker::builder(10 * MB, OnExhaustedPolicy::Fail).build()); + let tracker = Arc::new(QueryMemoryTracker::builder(10 * MB).build()); let mut stream1 = tracker.new_stream_tracker(); let mut stream2 = tracker.new_stream_tracker(); @@ -1325,8 +1249,7 @@ mod tests { #[test] fn test_query_memory_tracker_hard_limit() { - let tracker = - Arc::new(QueryMemoryTracker::builder(10 * MB, OnExhaustedPolicy::Fail).build()); + let tracker = Arc::new(QueryMemoryTracker::builder(10 * MB).build()); let mut stream = tracker.new_stream_tracker(); assert!(stream.try_track(9 * MB).is_ok()); @@ -1347,7 +1270,7 @@ mod tests { #[test] fn test_query_memory_tracker_unlimited() { - let tracker = Arc::new(QueryMemoryTracker::builder(0, OnExhaustedPolicy::Fail).build()); + let tracker = Arc::new(QueryMemoryTracker::builder(0).build()); let mut stream = tracker.new_stream_tracker(); assert!(stream.try_track(10 * MB).is_ok()); @@ -1358,8 +1281,7 @@ mod tests { #[test] fn test_query_memory_tracker_rounds_to_kilobytes() { - let tracker = - Arc::new(QueryMemoryTracker::builder(10 * MB, OnExhaustedPolicy::Fail).build()); + let tracker = Arc::new(QueryMemoryTracker::builder(10 * MB).build()); let mut stream = tracker.new_stream_tracker(); assert!(stream.try_track(1_537).is_ok()); @@ -1370,110 +1292,12 @@ mod tests { } #[tokio::test] - async fn test_memory_tracked_stream_waits_for_capacity() { + async fn test_memory_tracked_stream_rejects_immediately() { let exhausted = Arc::new(AtomicUsize::new(0)); let rejected = Arc::new(AtomicUsize::new(0)); let exhausted_counter = exhausted.clone(); let rejected_counter = rejected.clone(); - let tracker = QueryMemoryTracker::builder( - MB, - OnExhaustedPolicy::Wait { - timeout: Duration::from_millis(200), - }, - ) - .on_exhausted(move || { - exhausted_counter.fetch_add(1, Ordering::Relaxed); - }) - .on_reject(move || { - rejected_counter.fetch_add(1, Ordering::Relaxed); - }) - .build(); - let batch = large_string_batch(700 * 1024); - let expected_bytes = aligned_tracked_bytes(batch.logical_slice_memory_size()); - - let mut stream1 = MemoryTrackedStream::new( - RecordBatches::try_new(batch.schema.clone(), vec![batch.clone()]) - .unwrap() - .as_stream(), - tracker.clone(), - ); - let first = stream1.next().await.unwrap().unwrap(); - assert_eq!(first.num_rows(), 1); - assert_eq!(tracker.current(), expected_bytes); - - let stream2 = MemoryTrackedStream::new( - RecordBatches::try_new(batch.schema.clone(), vec![batch]) - .unwrap() - .as_stream(), - tracker.clone(), - ); - let waiter = tokio::spawn(async move { - let mut stream2 = stream2; - stream2.next().await.unwrap() - }); - - sleep(Duration::from_millis(50)).await; - assert!(!waiter.is_finished()); - - drop(stream1); - let second = waiter.await.unwrap().unwrap(); - assert_eq!(second.num_rows(), 1); - assert_eq!(exhausted.load(Ordering::Relaxed), 1); - assert_eq!(rejected.load(Ordering::Relaxed), 0); - } - - #[tokio::test] - async fn test_memory_tracked_stream_wait_times_out() { - let exhausted = Arc::new(AtomicUsize::new(0)); - let rejected = Arc::new(AtomicUsize::new(0)); - let exhausted_counter = exhausted.clone(); - let rejected_counter = rejected.clone(); - let tracker = QueryMemoryTracker::builder( - MB, - OnExhaustedPolicy::Wait { - timeout: Duration::from_millis(50), - }, - ) - .on_exhausted(move || { - exhausted_counter.fetch_add(1, Ordering::Relaxed); - }) - .on_reject(move || { - rejected_counter.fetch_add(1, Ordering::Relaxed); - }) - .build(); - let batch = large_string_batch(700 * 1024); - - let mut stream1 = MemoryTrackedStream::new( - RecordBatches::try_new(batch.schema.clone(), vec![batch.clone()]) - .unwrap() - .as_stream(), - tracker.clone(), - ); - let first = stream1.next().await.unwrap().unwrap(); - assert_eq!(first.num_rows(), 1); - - let mut stream2 = MemoryTrackedStream::new( - RecordBatches::try_new(batch.schema.clone(), vec![batch]) - .unwrap() - .as_stream(), - tracker, - ); - let result = timeout(Duration::from_secs(1), stream2.next()) - .await - .unwrap(); - let error = result.unwrap().unwrap_err(); - assert!(error.to_string().contains("timed out waiting")); - assert_eq!(exhausted.load(Ordering::Relaxed), 1); - assert_eq!(rejected.load(Ordering::Relaxed), 1); - } - - #[tokio::test] - async fn test_memory_tracked_stream_fail_policy_rejects_immediately() { - let exhausted = Arc::new(AtomicUsize::new(0)); - let rejected = Arc::new(AtomicUsize::new(0)); - let exhausted_counter = exhausted.clone(); - let rejected_counter = rejected.clone(); - let tracker = QueryMemoryTracker::builder(MB, OnExhaustedPolicy::Fail) + let tracker = QueryMemoryTracker::builder(MB) .on_exhausted(move || { exhausted_counter.fetch_add(1, Ordering::Relaxed); }) @@ -1498,8 +1322,10 @@ mod tests { .as_stream(), tracker, ); - let result = stream2.next().await.unwrap(); - assert!(result.is_err()); + assert!(matches!( + futures::poll!(stream2.next()), + Poll::Ready(Some(Err(_))) + )); assert_eq!(exhausted.load(Ordering::Relaxed), 1); assert_eq!(rejected.load(Ordering::Relaxed), 1); } diff --git a/src/common/recordbatch/src/memory.rs b/src/common/recordbatch/src/memory.rs new file mode 100644 index 00000000000..402599f6de4 --- /dev/null +++ b/src/common/recordbatch/src/memory.rs @@ -0,0 +1,31 @@ +// 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. + +//! Scan memory reservations held for reader lifetimes. + +use common_memory_manager::MemoryGuard; + +use crate::CallbackMemoryMetrics; + +/// Reserved scan memory, released automatically on drop. +#[derive(Debug)] +pub struct QueryMemoryReservation { + pub(crate) inner: MemoryGuard, +} + +impl QueryMemoryReservation { + pub(crate) fn try_acquire_additional(&mut self, bytes: u64) -> bool { + self.inner.try_acquire_additional(bytes) + } +} diff --git a/src/mito2/AGENTS.md b/src/mito2/AGENTS.md index f70aa622a27..08600cd00e4 100644 --- a/src/mito2/AGENTS.md +++ b/src/mito2/AGENTS.md @@ -51,6 +51,10 @@ immutable `Version` (`region/version.rs`) → scans memtables and Parquet SSTs (`sst/parquet.rs`) → merges (`read/`) and dedups by sequence → projected, filtered `RecordBatch` stream. +Internal series merges use the query runtime pool forwarded through +`QueryScanContext`; reader and output-batch accounting uses the separate scan +memory tracker. Direct scanner callers default to an unbounded operator pool. + ## Public surface - Entry: `MitoEngine` in `src/mito2/src/engine.rs`, built via `MitoEngineBuilder`. diff --git a/src/mito2/src/config.rs b/src/mito2/src/config.rs index c0fa74281ca..98b4ba4be42 100644 --- a/src/mito2/src/config.rs +++ b/src/mito2/src/config.rs @@ -51,11 +51,22 @@ const PAGE_CACHE_SIZE_FACTOR: u64 = 8; /// Use `1/INDEX_CREATE_MEM_THRESHOLD_FACTOR` of OS memory size as mem threshold for creating index const INDEX_CREATE_MEM_THRESHOLD_FACTOR: u64 = 16; +/// Memory accounted against the shared scan limit. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ScanMemoryMode { + /// Count projected Parquet columns' uncompressed row-group sizes. + #[default] + Basic, + /// Also count cumulative streamed batches. + Extended, +} + /// Fetch option timeout pub(crate) const FETCH_OPTION_TIMEOUT: Duration = Duration::from_secs(3); /// Configuration for [MitoEngine](crate::engine::MitoEngine). -/// Before using the config, make sure to call `MitoConfig::validate()` to check if the config is valid. +/// Before using the config, call [`MitoConfig::sanitize()`] to validate and normalize it. #[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq)] #[serde(default)] pub struct MitoConfig { @@ -170,18 +181,22 @@ pub struct MitoConfig { /// Buffer size for SST writing. // TODO(#9240): Support per-object-store write buffer sizes for backend-specific tuning. pub sst_write_buffer_size: ReadableSize, - /// Maximum number of SST files to scan concurrently (default 384). + /// Maximum number of SST files to scan concurrently when memory accounting + /// does not cover the scan or the limit is disabled (default 384). pub max_concurrent_scan_files: usize, /// Whether to allow stale entries read during replay. pub allow_stale_entries: bool, - /// Memory limit for table scans across all queries. + /// Memory limit for table scans across all queries (default: 25% of system memory). /// Setting it to 0 or "unlimited" disables the limit. /// Supports absolute size (e.g., "2GB") or percentage of system memory (e.g., "50%"). pub scan_memory_limit: MemoryLimit, - /// Behavior when scan memory tracking cannot acquire memory from the budget. - /// `wait` means `wait(10s)`, not unlimited waiting. - /// Defaults to [`OnExhaustedPolicy::Fail`], which intentionally differs from - /// [`OnExhaustedPolicy::default()`]. + /// Scan memory accounting: basic counts projected Parquet row groups; + /// extended also counts cumulative streamed batches. Limited reservations round up to KB. + /// Internal series merges use the query runtime pool in both modes. + pub scan_memory_mode: ScanMemoryMode, + /// Scan memory exhaustion always fails immediately (default: `fail`). + /// Legacy `wait` and `wait()` values are accepted but deprecated: + /// sanitization warns and replaces them with [`OnExhaustedPolicy::Fail`]. pub scan_memory_on_exhausted: OnExhaustedPolicy, /// Index configs. @@ -253,7 +268,8 @@ impl Default for MitoConfig { sst_write_buffer_size: DEFAULT_WRITE_BUFFER_SIZE, max_concurrent_scan_files: DEFAULT_MAX_CONCURRENT_SCAN_FILES, allow_stale_entries: false, - scan_memory_limit: MemoryLimit::default(), + scan_memory_limit: MemoryLimit::Percentage(25), + scan_memory_mode: ScanMemoryMode::Basic, scan_memory_on_exhausted: OnExhaustedPolicy::Fail, index: IndexConfig::default(), inverted_index: InvertedIndexConfig::default(), @@ -281,6 +297,18 @@ impl MitoConfig { /// /// Returns an error if there is a configuration that unable to sanitize. pub fn sanitize(&mut self, data_home: &str) -> Result<()> { + if matches!( + self.scan_memory_on_exhausted, + OnExhaustedPolicy::Wait { .. } + ) { + warn!( + "scan_memory_on_exhausted = {:?} is deprecated and is treated as fail; \ + scan memory exhaustion now fails immediately. Set scan_memory_on_exhausted = \"fail\"", + self.scan_memory_on_exhausted + ); + self.scan_memory_on_exhausted = OnExhaustedPolicy::Fail; + } + if self.experimental_enable_series_index { snafu::ensure!( self.experimental_series_index_max_size.as_bytes() >= 1024, @@ -419,6 +447,39 @@ impl MitoConfig { mod tests { use super::*; + #[test] + fn test_scan_memory_policy_upgrade() { + for policy in [ + None, + Some("fail"), + Some("skip"), + Some("wait"), + Some("wait(30s)"), + ] { + let input = policy + .map(|policy| format!("scan_memory_on_exhausted = {policy:?}")) + .unwrap_or_default(); + let mut config: MitoConfig = toml::from_str(&input).unwrap(); + let compaction_policy = config.experimental_compaction_on_exhausted; + config.sanitize("/data").unwrap(); + assert_eq!(OnExhaustedPolicy::Fail, config.scan_memory_on_exhausted); + assert_eq!( + compaction_policy, + config.experimental_compaction_on_exhausted + ); + let serialized = toml::to_string(&config).unwrap(); + assert!(serialized.contains("scan_memory_on_exhausted = \"fail\"")); + config.sanitize("/data").unwrap(); + assert_eq!(serialized, toml::to_string(&config).unwrap()); + } + for policy in ["unknown", "wait(invalid)", "wait()"] { + assert!( + toml::from_str::(&format!("scan_memory_on_exhausted = {policy:?}")) + .is_err() + ); + } + } + #[test] fn test_default_auto_flush_interval() { assert_eq!( diff --git a/src/mito2/src/engine.rs b/src/mito2/src/engine.rs index 3036c61e7ff..cdfc78ff0c4 100644 --- a/src/mito2/src/engine.rs +++ b/src/mito2/src/engine.rs @@ -108,7 +108,6 @@ use common_recordbatch::{QueryMemoryTracker, SendableRecordBatchStream}; use common_stat::get_total_memory_bytes; use common_telemetry::{debug, info, tracing, warn}; use common_wal::options::WalOptions; -use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool, UnboundedMemoryPool}; use futures::future::{join_all, try_join_all}; use futures::stream::{self, Stream, StreamExt}; use object_store::manager::ObjectStoreManagerRef; @@ -138,7 +137,7 @@ use tokio::sync::{Semaphore, oneshot}; use crate::access_layer::RegionFilePathFactory; use crate::cache::{CacheManagerRef, CacheStrategy}; -use crate::config::MitoConfig; +use crate::config::{MitoConfig, ScanMemoryMode}; use crate::engine::puffin_index::{IndexEntryContext, collect_index_entries_from_puffin}; use crate::error::{ IncrementalQueryStaleSnafu, InvalidRequestSnafu, JoinSnafu, MitoManifestInfoSnafu, RecvSnafu, @@ -244,26 +243,23 @@ impl<'a, S: LogStore> MitoEngineBuilder<'a, S> { let wal_raw_entry_reader = Arc::new(LogStoreRawEntryReader::new(self.log_store)); let total_memory = get_total_memory_bytes().max(0) as u64; let scan_memory_limit = config.scan_memory_limit.resolve(total_memory) as usize; - let scan_memory_pool = new_scan_memory_pool(scan_memory_limit); - let scan_memory_tracker = - QueryMemoryTracker::builder(scan_memory_limit, config.scan_memory_on_exhausted) - .on_update(|usage| { - SCAN_MEMORY_USAGE_BYTES.set(usage as i64); - }) - .on_exhausted(|| { - SCAN_MEMORY_EXHAUSTED_TOTAL.inc(); - }) - .on_reject(|| { - SCAN_REQUESTS_REJECTED_TOTAL.inc(); - }) - .build(); + let scan_memory_tracker = QueryMemoryTracker::builder(scan_memory_limit) + .on_update(|usage| { + SCAN_MEMORY_USAGE_BYTES.set(usage as i64); + }) + .on_exhausted(|| { + SCAN_MEMORY_EXHAUSTED_TOTAL.inc(); + }) + .on_reject(|| { + SCAN_REQUESTS_REJECTED_TOTAL.inc(); + }) + .build(); let inner = EngineInner { workers, config, wal_raw_entry_reader, scan_memory_tracker, - scan_memory_pool, region_hook, #[cfg(feature = "enterprise")] extension_range_provider_factory: None, @@ -749,8 +745,6 @@ struct EngineInner { wal_raw_entry_reader: Arc, /// Memory tracker for table scans. scan_memory_tracker: QueryMemoryTracker, - /// Memory pool shared by internal scan operators across all queries. - scan_memory_pool: Arc, /// The region hook (if any) registered via plugins; exposed for the GC worker /// to fire [`RegionHook::on_region_gc`]. region_hook: Option, @@ -1165,7 +1159,7 @@ impl EngineInner { .with_ignore_range_index(!self.config.experimental_enable_range_index) .with_query_stat_counters(region.region_stats.query_stat_counters()) .with_max_concurrent_scan_files(self.config.max_concurrent_scan_files) - .with_scan_memory_pool(self.scan_memory_pool.clone()) + .with_scan_memory_tracker(self.scan_memory_tracker.clone()) .with_experimental_series_scan_v2(self.config.experimental_series_scan_v2) .with_ignore_inverted_index(self.config.inverted_index.apply_on_query.disabled()) .with_ignore_fulltext_index(self.config.fulltext_index.apply_on_query.disabled()) @@ -1457,7 +1451,8 @@ impl RegionEngine for MitoEngine { } fn query_memory_tracker(&self) -> Option { - Some(self.inner.scan_memory_tracker.clone()) + (self.inner.config.scan_memory_mode == ScanMemoryMode::Extended) + .then(|| self.inner.scan_memory_tracker.clone()) } async fn get_committed_sequence( @@ -1622,19 +1617,17 @@ impl MitoEngine { let wal_raw_entry_reader = Arc::new(LogStoreRawEntryReader::new(log_store.clone())); let total_memory = get_total_memory_bytes().max(0) as u64; let scan_memory_limit = config.scan_memory_limit.resolve(total_memory) as usize; - let scan_memory_pool = new_scan_memory_pool(scan_memory_limit); - let scan_memory_tracker = - QueryMemoryTracker::builder(scan_memory_limit, config.scan_memory_on_exhausted) - .on_update(|usage| { - SCAN_MEMORY_USAGE_BYTES.set(usage as i64); - }) - .on_exhausted(|| { - SCAN_MEMORY_EXHAUSTED_TOTAL.inc(); - }) - .on_reject(|| { - SCAN_REQUESTS_REJECTED_TOTAL.inc(); - }) - .build(); + let scan_memory_tracker = QueryMemoryTracker::builder(scan_memory_limit) + .on_update(|usage| { + SCAN_MEMORY_USAGE_BYTES.set(usage as i64); + }) + .on_exhausted(|| { + SCAN_MEMORY_EXHAUSTED_TOTAL.inc(); + }) + .on_reject(|| { + SCAN_REQUESTS_REJECTED_TOTAL.inc(); + }) + .build(); Ok(MitoEngine { inner: Arc::new(EngineInner { workers: WorkerGroup::start_for_test( @@ -1653,7 +1646,6 @@ impl MitoEngine { config, wal_raw_entry_reader, scan_memory_tracker, - scan_memory_pool, region_hook: None, #[cfg(feature = "enterprise")] extension_range_provider_factory: None, @@ -1667,20 +1659,10 @@ impl MitoEngine { } } -fn new_scan_memory_pool(scan_memory_limit: usize) -> Arc { - if scan_memory_limit == 0 { - Arc::new(UnboundedMemoryPool::default()) - } else { - Arc::new(GreedyMemoryPool::new(scan_memory_limit)) - } -} - #[cfg(test)] mod tests { use std::time::Duration; - use datafusion::execution::memory_pool::MemoryConsumer; - use super::*; use crate::sst::file::FileMeta; @@ -1766,17 +1748,4 @@ mod tests { }; assert!(!is_valid_region_edit(&edit)); } - - #[test] - fn test_scan_memory_pool_is_shared_across_consumers() { - let pool = new_scan_memory_pool(100); - let cloned_pool = pool.clone(); - let first = MemoryConsumer::new("first-scan").register(&pool); - let second = MemoryConsumer::new("second-scan").register(&cloned_pool); - - first.try_grow(60).unwrap(); - assert!(second.try_grow(50).is_err()); - second.try_grow(40).unwrap(); - assert_eq!(100, pool.reserved()); - } } diff --git a/src/mito2/src/engine/scan_test.rs b/src/mito2/src/engine/scan_test.rs index f55f084ebd7..39613efb383 100644 --- a/src/mito2/src/engine/scan_test.rs +++ b/src/mito2/src/engine/scan_test.rs @@ -1275,15 +1275,26 @@ async fn test_scan_with_min_sst_sequence_with_format(flat_format: bool) { #[tokio::test] async fn test_max_concurrent_scan_files() { - test_max_concurrent_scan_files_with_format(false).await; - test_max_concurrent_scan_files_with_format(true).await; + test_max_concurrent_scan_files_with_format(false, false).await; + test_max_concurrent_scan_files_with_format(true, false).await; } -async fn test_max_concurrent_scan_files_with_format(flat_format: bool) { +#[tokio::test] +async fn test_scan_memory_limit_small_legacy_files() { + test_max_concurrent_scan_files_with_format(false, true).await; + test_max_concurrent_scan_files_with_format(true, true).await; +} + +async fn test_max_concurrent_scan_files_with_format(flat_format: bool, estimated: bool) { let mut env = TestEnv::with_prefix("test_max_concurrent_scan_files").await; let config = MitoConfig { default_flat_format: flat_format, max_concurrent_scan_files: 2, + scan_memory_limit: if estimated { + common_base::memory_limit::MemoryLimit::Size(ReadableSize::mb(1)) + } else { + common_base::memory_limit::MemoryLimit::Unlimited + }, ..Default::default() }; let engine = env.create_engine(config).await; @@ -1311,13 +1322,45 @@ async fn test_max_concurrent_scan_files_with_format(flat_format: bool) { put_and_flush(3, 7).await; put_and_flush(6, 9).await; + if estimated { + // Simulate SSTs created before the manifest recorded row-group sizes. + let region = engine.get_region(region_id).unwrap(); + let version = region.version_control.current().version; + let legacy_files = version + .ssts + .levels() + .iter() + .flat_map(|level| level.files()) + .map(|file| { + let mut meta = file.meta_ref().clone(); + meta.max_row_group_uncompressed_size = 0; + meta + }) + .collect::>(); + let version = crate::region::version::VersionBuilder::from_version(version) + .clear_files() + .add_files( + Arc::new(crate::sst::file_purger::NoopFilePurger), + legacy_files.into_iter(), + ) + .build(); + region.version_control.overwrite_current(Arc::new(version)); + } + let request = ScanRequest::default(); let scanner = engine.scanner(region_id, request).await.unwrap(); let Scanner::Seq(scanner) = scanner else { panic!("Scanner should be seq scan"); }; - let error = scanner.check_scan_limit().unwrap_err(); - assert_eq!(StatusCode::RateLimited, error.status_code()); + if estimated { + scanner.check_scan_limit().unwrap(); + RecordBatches::try_collect(scanner.build_stream().unwrap()) + .await + .unwrap(); + } else { + let error = scanner.check_scan_limit().unwrap_err(); + assert_eq!(StatusCode::RateLimited, error.status_code()); + } let request = ScanRequest { distribution: Some(TimeSeriesDistribution::PerSeries), @@ -1327,8 +1370,15 @@ async fn test_max_concurrent_scan_files_with_format(flat_format: bool) { let Scanner::Series(scanner) = scanner else { panic!("Scanner should be series scan"); }; - let error = scanner.check_scan_limit().unwrap_err(); - assert_eq!(StatusCode::RateLimited, error.status_code()); + if estimated { + scanner.check_scan_limit().unwrap(); + RecordBatches::try_collect(scanner.build_stream().await.unwrap()) + .await + .unwrap(); + } else { + let error = scanner.check_scan_limit().unwrap_err(); + assert_eq!(StatusCode::RateLimited, error.status_code()); + } } #[tokio::test] @@ -1443,9 +1493,15 @@ async fn test_series_scan_with_format(flat_format: bool) { async fn test_two_phase_series_scan() { for use_index in [false, true] { for use_range_index in [false, true] { - check_two_phase_series_scan(use_index, use_range_index, true).await; + for mode in [ + crate::config::ScanMemoryMode::Basic, + crate::config::ScanMemoryMode::Extended, + ] { + check_two_phase_series_scan(use_index, use_range_index, true, mode).await; + } } - check_two_phase_series_scan(use_index, true, false).await; + check_two_phase_series_scan(use_index, true, false, crate::config::ScanMemoryMode::Basic) + .await; } } @@ -1453,10 +1509,15 @@ async fn check_two_phase_series_scan( use_index: bool, use_range_index: bool, enable_range_index: bool, + scan_memory_mode: crate::config::ScanMemoryMode, ) { let mut env = TestEnv::with_prefix("test_two_phase_series_scan").await; let engine = env .create_engine(MitoConfig { + scan_memory_mode, + scan_memory_on_exhausted: common_memory_manager::OnExhaustedPolicy::Wait { + timeout: std::time::Duration::from_secs(30), + }, experimental_series_scan_v2: true, experimental_enable_series_index: true, experimental_enable_range_index: enable_range_index, @@ -1685,7 +1746,12 @@ async fn check_two_phase_series_scan( .unwrap(); let metrics_set = ExecutionPlanMetricsSet::default(); - let context = store_api::region_engine::QueryScanContext::default(); + use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryConsumer, MemoryPool}; + let pool: Arc = Arc::new(GreedyMemoryPool::new(1024 * 1024)); + let context = store_api::region_engine::QueryScanContext { + memory_pool: pool.clone(), + ..Default::default() + }; let partition_batches = try_join_all((0..3).map(|partition| { let stream = scanner .scan_partition(&context, &metrics_set, partition) @@ -1744,6 +1810,65 @@ async fn check_two_phase_series_scan( assert_eq!(Some(&2), series_to_partition.get("b")); assert_eq!(Some(&2), series_to_partition.get("d")); + assert_eq!(0, pool.reserved()); + // Existing query operators consume the same pool in either scan mode. + scanner.reset_state(); + let other = MemoryConsumer::new("other-query-operator").register(&pool); + other.try_grow(1024 * 1024).unwrap(); + let failed = try_join_all((0..3).map(|partition| { + let stream = scanner + .scan_partition(&context, &metrics_set, partition) + .unwrap(); + async move { stream.try_collect::>().await } + })) + .await; + assert!(failed.is_err()); + other.free(); + assert_eq!(0, pool.reserved()); + + // Exhausting the scan budget must fail even with an unbounded query pool. + scanner.reset_state(); + let scan_tracker = &engine.inner.scan_memory_tracker; + // Cancelled background readers release their reservations asynchronously. + tokio::time::timeout(std::time::Duration::from_secs(5), async { + while scan_tracker.current() > 0 { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + let scan_guard = scan_tracker.reserve(scan_tracker.limit()).unwrap(); + let fresh_context = store_api::region_engine::QueryScanContext::default(); + let fresh_pool = Arc::downgrade(&fresh_context.memory_pool); + let failed = tokio::time::timeout( + std::time::Duration::from_secs(5), + try_join_all((0..3).map(|partition| { + let stream = scanner + .scan_partition(&fresh_context, &metrics_set, partition) + .unwrap(); + async move { stream.try_collect::>().await } + })), + ) + .await + .expect("legacy wait policy must fail immediately on scan memory exhaustion"); + assert!(failed.is_err()); + drop(fresh_context); + // Keep the external quota occupied while failure tears down the distributor + // and readers. Cleanup must not depend on quota becoming available. + tokio::time::timeout(std::time::Duration::from_secs(5), async { + while fresh_pool.upgrade().is_some() { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + assert_eq!(scan_tracker.limit(), scan_tracker.current()); + drop(scan_guard); + assert_eq!(0, scan_tracker.current()); + + // Keep the previous execution's pool exhausted while replay uses a fresh pool. + other.try_grow(1024 * 1024).unwrap(); + let context = store_api::region_engine::QueryScanContext::default(); scanner.reset_state(); assert_eq!("two_phase", scanner.mode()); assert_eq!( @@ -1789,6 +1914,42 @@ async fn check_two_phase_series_scan( replay_rows.sort(); assert_eq!(actual_rows, replay_rows); + // Dropping every partition must let the shared distributor release its pool. + scanner.reset_state(); + let cancelled_context = store_api::region_engine::QueryScanContext::default(); + let cancelled_pool = Arc::downgrade(&cancelled_context.memory_pool); + let streams = (0..3) + .map(|partition| { + scanner + .scan_partition(&cancelled_context, &metrics_set, partition) + .unwrap() + }) + .collect::>(); + drop(streams); + drop(cancelled_context); + tokio::time::timeout(std::time::Duration::from_secs(5), async { + while cancelled_pool.upgrade().is_some() { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + scanner.reset_state(); + let direct_batches = scanner + .build_stream() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!( + actual_rows.len(), + direct_batches + .iter() + .map(|batch| batch.num_rows()) + .sum::() + ); + // Exercise precise field/time filtering and candidate tag filtering on both paths. let filtered = engine .scanner( diff --git a/src/mito2/src/read/scan_region.rs b/src/mito2/src/read/scan_region.rs index 0f9cea941b3..d2bc53167be 100644 --- a/src/mito2/src/read/scan_region.rs +++ b/src/mito2/src/read/scan_region.rs @@ -22,13 +22,12 @@ use std::time::Instant; use api::v1::SemanticType; use common_error::ext::BoxedError; -use common_recordbatch::SendableRecordBatchStream; use common_recordbatch::adapter::RegionQueryStatCounters; use common_recordbatch::filter::SimpleFilterEvaluator; +use common_recordbatch::{QueryMemoryTracker, SendableRecordBatchStream}; use common_telemetry::tracing::Instrument; use common_telemetry::{debug, error, tracing, warn}; use common_time::range::TimestampRange; -use datafusion::execution::memory_pool::{MemoryPool, UnboundedMemoryPool}; use datafusion::optimizer::simplify_expressions::ExprSimplifier; use datafusion::physical_plan::expressions::DynamicFilterPhysicalExpr; use datafusion_common::pruning::PruningStatistics; @@ -244,8 +243,7 @@ pub(crate) struct ScanRegion { cache_strategy: CacheStrategy, /// Maximum number of SST files to scan concurrently. max_concurrent_scan_files: usize, - /// Memory pool shared by internal scan operators across all queries. - scan_memory_pool: Arc, + scan_memory_tracker: Option, /// Whether to enable the experimental two-phase metric series scan. experimental_series_scan_v2: bool, /// Whether to ignore range indexes during scans. @@ -284,7 +282,7 @@ impl ScanRegion { request, cache_strategy, max_concurrent_scan_files: DEFAULT_MAX_CONCURRENT_SCAN_FILES, - scan_memory_pool: Arc::new(UnboundedMemoryPool::default()), + scan_memory_tracker: None, experimental_series_scan_v2: false, ignore_range_index: false, ignore_inverted_index: false, @@ -326,10 +324,8 @@ impl ScanRegion { self } - /// Sets the memory pool shared by internal scan operators. - #[must_use] - pub(crate) fn with_scan_memory_pool(mut self, scan_memory_pool: Arc) -> Self { - self.scan_memory_pool = scan_memory_pool; + pub(crate) fn with_scan_memory_tracker(mut self, tracker: QueryMemoryTracker) -> Self { + self.scan_memory_tracker = Some(tracker); self } @@ -583,7 +579,7 @@ impl ScanRegion { .with_bloom_filter_index_appliers(bloom_filter_appliers) .with_fulltext_index_appliers(fulltext_index_appliers) .with_max_concurrent_scan_files(self.max_concurrent_scan_files) - .with_scan_memory_pool(self.scan_memory_pool) + .with_scan_memory_tracker(self.scan_memory_tracker) .with_start_time(self.start_time) .with_append_mode(self.version.options.append_mode) .with_filter_deleted(self.filter_deleted) @@ -946,8 +942,8 @@ pub struct ScanInput { ignore_file_not_found: bool, /// Maximum number of SST files to scan concurrently. pub(crate) max_concurrent_scan_files: usize, - /// Memory pool shared by internal scan operators across all queries. - pub(crate) scan_memory_pool: Arc, + /// Shared scan quota for projected SST row groups; absent during compaction. + scan_memory_tracker: Option, /// Index appliers. inverted_index_appliers: [Option; 2], bloom_filter_index_appliers: [Option; 2], @@ -1026,7 +1022,7 @@ impl ScanInput { cache_strategy: CacheStrategy::Disabled, ignore_file_not_found: false, max_concurrent_scan_files: DEFAULT_MAX_CONCURRENT_SCAN_FILES, - scan_memory_pool: Arc::new(UnboundedMemoryPool::default()), + scan_memory_tracker: None, inverted_index_appliers: [None, None], bloom_filter_index_appliers: [None, None], fulltext_index_appliers: [None, None], @@ -1119,6 +1115,9 @@ impl ScanInputBuilder { /// Builds a finalized [ScanInput] and computes its scan analysis. #[must_use] pub(crate) fn build(mut self) -> ScanInput { + if self.input.compaction { + self.input.scan_memory_tracker = None; + } let input = &self.input; let eligible = !input.compaction && !input.files.is_empty() @@ -1275,10 +1274,8 @@ impl ScanInputBuilder { self } - /// Sets the memory pool shared by internal scan operators. - #[must_use] - pub(crate) fn with_scan_memory_pool(mut self, scan_memory_pool: Arc) -> Self { - self.input.scan_memory_pool = scan_memory_pool; + pub(crate) fn with_scan_memory_tracker(mut self, tracker: Option) -> Self { + self.input.scan_memory_tracker = tracker; self } @@ -1415,6 +1412,16 @@ impl ScanInputBuilder { } impl ScanInput { + pub(crate) fn uses_scan_memory_limit(&self) -> bool { + let covered = self + .scan_memory_tracker + .as_ref() + .is_some_and(|tracker| tracker.limit() > 0); + #[cfg(feature = "enterprise")] + let covered = covered && self.extension_ranges.is_empty(); + covered + } + /// Builds memtable ranges to scan by `index`. pub(crate) fn build_mem_ranges(&self, index: RowGroupIndex) -> SmallVec<[MemtableRange; 2]> { let memtable = &self.memtables[index.index]; @@ -1547,6 +1554,7 @@ impl ScanInput { .clone() .filter(|_| !self.ignore_range_index), ) + .memory_tracker(self.scan_memory_tracker.clone()) .predicate(predicate) .projection(Some(self.read_cols.clone())) .json2_rewrite_targets(self.json2_rewrite_targets.clone()) @@ -1659,9 +1667,17 @@ impl ScanInput { // We release the permit before sending result to avoid the task waiting on // the channel with the permit held. let maybe_batch = { - // Safety: We never close the semaphore. - let _permit = semaphore.acquire().await.unwrap(); - input.next().await + tokio::select! { + // Cancellation or early completion can drop the receiver while + // we wait for a permit or I/O. Exit and drop the + // input to release its reservations without waiting for a batch. + _ = sender.closed() => break, + batch = async { + // Safety: We never close the semaphore. + let _permit = semaphore.acquire().await.unwrap(); + input.next().await + } => batch, + } }; match maybe_batch { Some(Ok(batch)) => { diff --git a/src/mito2/src/read/seq_scan.rs b/src/mito2/src/read/seq_scan.rs index 24c35481669..ecfa41bcf2e 100644 --- a/src/mito2/src/read/seq_scan.rs +++ b/src/mito2/src/read/seq_scan.rs @@ -107,10 +107,9 @@ impl SeqScan { )] pub fn build_stream(&self) -> Result { let metrics_set = ExecutionPlanMetricsSet::new(); + let ctx = QueryScanContext::default(); let streams = (0..self.properties.partitions.len()) - .map(|partition: usize| { - self.scan_partition(&QueryScanContext::default(), &metrics_set, partition) - }) + .map(|partition: usize| self.scan_partition(&ctx, &metrics_set, partition)) .collect::, _>>()?; let aggr_stream = ChainedRecordBatchStream::new(streams).map_err(BoxedError::new)?; @@ -506,6 +505,9 @@ impl SeqScan { /// Checks resource limit for the scanner. pub(crate) fn check_scan_limit(&self) -> Result<()> { + if self.stream_ctx.input.uses_scan_memory_limit() { + return Ok(()); + } // Check max file count limit for all partitions since we scan them in parallel. let total_max_files: usize = self .properties diff --git a/src/mito2/src/read/series_scan.rs b/src/mito2/src/read/series_scan.rs index 4328a20f59c..a77cbebb557 100644 --- a/src/mito2/src/read/series_scan.rs +++ b/src/mito2/src/read/series_scan.rs @@ -25,6 +25,7 @@ use common_recordbatch::util::ChainedRecordBatchStream; use common_recordbatch::{RecordBatchStreamWrapper, SendableRecordBatchStream}; use common_telemetry::tracing::{self, Instrument}; use common_telemetry::warn; +use datafusion::execution::memory_pool::MemoryPool; use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; use datafusion::physical_plan::{DisplayAs, DisplayFormatType}; use datatypes::arrow::array::BinaryArray; @@ -266,12 +267,9 @@ impl SeriesScan { metrics_set, ctx.explain_verbose, ), - SeriesScanMode::TwoPhase => self.scan_two_phase_batch_in_partition( - partition, - part_metrics, - metrics_set, - ctx.explain_verbose, - ), + SeriesScanMode::TwoPhase => { + self.scan_two_phase_batch_in_partition(ctx, partition, part_metrics, metrics_set) + } } } @@ -311,12 +309,12 @@ impl SeriesScan { fn scan_two_phase_batch_in_partition( &self, + ctx: &QueryScanContext, partition: usize, part_metrics: PartitionMetrics, metrics_set: &ExecutionPlanMetricsSet, - explain_verbose: bool, ) -> Result { - self.maybe_start_candidate_distributor(metrics_set, explain_verbose); + self.maybe_start_candidate_distributor(ctx, metrics_set); // Safety: `scan_batch_in_partition` validates the index, and the receiver // list always contains one entry for each output partition. let mut receiver = self.candidate_receivers.lock().unwrap()[partition] @@ -406,8 +404,8 @@ impl SeriesScan { fn maybe_start_candidate_distributor( &self, + ctx: &QueryScanContext, metrics_set: &ExecutionPlanMetricsSet, - explain_verbose: bool, ) { let mut rx_list = self.candidate_receivers.lock().unwrap(); if !rx_list.is_empty() { @@ -416,7 +414,10 @@ impl SeriesScan { let (senders, receivers, active_receivers) = new_candidate_channel_list(self.properties.num_partitions()); + // Capture the first partition's runtime pool once for the shared distributor. + // reset_state clears the receivers so the next execution captures a fresh pool. let mut distributor = SeriesCandidateDistributor { + memory_pool: ctx.memory_pool.clone(), stream_ctx: self.stream_ctx.clone(), range_semaphore: Arc::new(Semaphore::new(self.properties.num_partitions())), partitions: self.properties.partitions.clone(), @@ -425,7 +426,7 @@ impl SeriesScan { active_receivers, metrics_set: metrics_set.clone(), metrics_list: self.metrics_list.clone(), - explain_verbose, + explain_verbose: ctx.explain_verbose, }; let region_id = distributor.stream_ctx.input.mapper.metadata().region_id; let span = tracing::info_span!("SeriesScan::candidate_distributor", region_id = %region_id); @@ -447,8 +448,9 @@ impl SeriesScan { pub(crate) async fn build_stream(&self) -> Result { let part_num = self.properties.num_partitions(); let metrics_set = ExecutionPlanMetricsSet::default(); + let ctx = QueryScanContext::default(); let streams = (0..part_num) - .map(|i| self.scan_partition(&QueryScanContext::default(), &metrics_set, i)) + .map(|i| self.scan_partition(&ctx, &metrics_set, i)) .collect::, BoxedError>>()?; let chained_stream = ChainedRecordBatchStream::new(streams).map_err(BoxedError::new)?; Ok(Box::pin(chained_stream)) @@ -458,6 +460,7 @@ impl SeriesScan { pub(crate) fn scan_all_partitions(&self) -> Result { let metrics_set = ExecutionPlanMetricsSet::new(); + let ctx = QueryScanContext::default(); let streams = (0..self.properties.partitions.len()) .map(|partition| { let metrics = new_partition_metrics( @@ -468,12 +471,7 @@ impl SeriesScan { &self.metrics_list, ); - self.scan_batch_in_partition( - &QueryScanContext::default(), - partition, - metrics, - &metrics_set, - ) + self.scan_batch_in_partition(&ctx, partition, metrics, &metrics_set) }) .collect::>>()?; @@ -482,6 +480,9 @@ impl SeriesScan { /// Checks resource limit for the scanner. pub(crate) fn check_scan_limit(&self) -> Result<()> { + if self.stream_ctx.input.uses_scan_memory_limit() { + return Ok(()); + } // Sum the total number of files across all partitions let total_files: usize = self .properties @@ -842,6 +843,7 @@ impl SeriesDistributor { /// Discovers metric series and sends one reader assignment to each partition. struct SeriesCandidateDistributor { + memory_pool: Arc, stream_ctx: Arc, range_semaphore: Arc, partitions: Vec>, @@ -875,7 +877,7 @@ impl SeriesCandidateDistributor { self.partitions.clone(), self.pruner.clone(), self.range_semaphore.clone(), - self.stream_ctx.input.scan_memory_pool.clone(), + self.memory_pool.clone(), self.metrics_set.clone(), part_metrics.clone(), )?; diff --git a/src/mito2/src/read/unordered_scan.rs b/src/mito2/src/read/unordered_scan.rs index 8bc4f2e0038..9362898023f 100644 --- a/src/mito2/src/read/unordered_scan.rs +++ b/src/mito2/src/read/unordered_scan.rs @@ -89,8 +89,9 @@ impl UnorderedScan { pub(crate) async fn build_stream(&self) -> Result { let metrics_set = ExecutionPlanMetricsSet::new(); let part_num = self.properties.num_partitions(); + let ctx = QueryScanContext::default(); let streams = (0..part_num) - .map(|i| self.scan_partition(&QueryScanContext::default(), &metrics_set, i)) + .map(|i| self.scan_partition(&ctx, &metrics_set, i)) .collect::, BoxedError>>()?; let stream = stream! { for mut stream in streams { diff --git a/src/mito2/src/sst/parquet/push_decoder.rs b/src/mito2/src/sst/parquet/push_decoder.rs index 4ae9057479a..1a1446c13c2 100644 --- a/src/mito2/src/sst/parquet/push_decoder.rs +++ b/src/mito2/src/sst/parquet/push_decoder.rs @@ -17,6 +17,7 @@ use std::ops::Range; use bytes::{Bytes, BytesMut}; +use common_recordbatch::QueryMemoryTracker; use datatypes::arrow::record_batch::RecordBatch; use futures::StreamExt; use futures::stream::BoxStream; @@ -25,7 +26,7 @@ use parquet::DecodeResult; use parquet::arrow::ProjectionMask; use parquet::arrow::arrow_reader::{ArrowReaderMetadata, RowSelection}; use parquet::arrow::push_decoder::ParquetPushDecoderBuilder; -use snafu::{ResultExt, ensure}; +use snafu::{OptionExt, ResultExt, ensure}; use crate::cache::file_cache::{FileType, IndexKey}; use crate::cache::{CacheStrategy, PageRangePart}; @@ -53,6 +54,8 @@ pub struct SstParquetRangeFetcher { row_group_idx: usize, /// Optional metrics for tracking fetch operations. fetch_metrics: Option, + /// Shared quota for decoding this fetcher's row group. + memory_tracker: Option, } impl SstParquetRangeFetcher { @@ -72,9 +75,16 @@ impl SstParquetRangeFetcher { cache_strategy, row_group_idx, fetch_metrics, + memory_tracker: None, } } + /// Sets the shared quota for this decoder's projected row group. + pub(crate) fn with_memory_tracker(mut self, tracker: Option) -> Self { + self.memory_tracker = tracker; + self + } + /// Fetches byte ranges from page cache, write cache, or object store. async fn fetch_bytes_with_cache(&self, ranges: Vec>) -> Result> { let fetch_start = self @@ -321,6 +331,24 @@ pub fn build_sst_parquet_record_batch_stream( file_path: String, batch_size: usize, ) -> Result>> { + let uncompressed_bytes = if fetcher.memory_tracker.is_some() { + let row_group = arrow_metadata.metadata().row_group(row_group_idx); + row_group + .columns() + .iter() + .enumerate() + .filter(|(index, _)| projection.leaf_included(*index)) + .try_fold(0usize, |total, (_, column)| { + usize::try_from(column.uncompressed_size()) + .ok() + .and_then(|size| total.checked_add(size)) + }) + .context(UnexpectedSnafu { + reason: "Invalid projected row-group uncompressed size", + })? + } else { + 0 + }; let mut builder = ParquetPushDecoderBuilder::new_with_metadata(arrow_metadata) .with_row_groups(vec![row_group_idx]) .with_projection(projection) @@ -335,12 +363,17 @@ pub fn build_sst_parquet_record_batch_stream( .context(ReadParquetSnafu { path: &file_path })?; Ok(async_stream::try_stream! { + let _reservation = if let Some(tracker) = &fetcher.memory_tracker { + Some(tracker.reserve(uncompressed_bytes) + .context(crate::error::RecordBatchSnafu)?) + } else { + None + }; loop { match decoder.try_decode().context(ReadParquetSnafu { path: &file_path })? { DecodeResult::NeedsData(ranges) => { let data = fetcher.fetch_bytes_with_cache(ranges.clone()).await?; - decoder - .push_ranges(ranges, data) + decoder.push_ranges(ranges, data) .context(ReadParquetSnafu { path: &file_path })?; } DecodeResult::Data(batch) => yield batch, @@ -355,6 +388,139 @@ pub fn build_sst_parquet_record_batch_stream( mod tests { use super::*; + #[tokio::test] + async fn test_decoder_memory_projection_lifetime_and_error() { + use std::sync::Arc; + + use datatypes::arrow::array::{Int64Array, StringArray}; + use datatypes::arrow::datatypes::{DataType, Field, Schema}; + use object_store::services::Memory; + use parquet::arrow::ArrowWriter; + use parquet::arrow::arrow_reader::ArrowReaderOptions; + use parquet::file::properties::WriterProperties; + + use crate::sst::parquet::metadata::MetadataLoader; + use crate::sst::parquet::reader::MetadataCacheMetrics; + + let schema = Arc::new(Schema::new(vec![ + Field::new("value", DataType::Int64, false), + Field::new("payload", DataType::Utf8, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int64Array::from_iter_values(0..8)), + Arc::new(StringArray::from_iter_values( + (0..8).map(|i| format!("{i}:{}", "x".repeat(512))), + )), + ], + ) + .unwrap(); + let mut bytes = Vec::new(); + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(4)) + .build(); + let mut writer = ArrowWriter::try_new(&mut bytes, schema, Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + let store = ObjectStore::new(Memory::default()).unwrap(); + let file_size = bytes.len() as u64; + store.write("memory.parquet", bytes).await.unwrap(); + let metadata = MetadataLoader::new(store.clone(), "memory.parquet", file_size) + .load(&mut MetadataCacheMetrics::default()) + .await + .unwrap(); + let sizes: Vec<_> = metadata + .row_groups() + .iter() + .map(|rg| rg.column(0).uncompressed_size() as usize) + .collect(); + let aligned = |bytes: usize| bytes.div_ceil(1024) * 1024; + let limit = aligned(*sizes.iter().max().unwrap()); + assert!(metadata.row_group(0).total_byte_size() as usize > limit); + let projection = ProjectionMask::leaves(metadata.file_metadata().schema_descr(), [0]); + let metadata = + ArrowReaderMetadata::try_new(Arc::new(metadata), ArrowReaderOptions::new()).unwrap(); + let tracker = QueryMemoryTracker::builder(limit).build(); + let cache = CacheStrategy::EnableAll(Arc::new( + crate::cache::CacheManager::builder() + .page_cache_size(1024 * 1024) + .build(), + )); + let region_file_id = RegionFileId::new( + store_api::storage::RegionId::new(1, 1), + store_api::storage::FileId::random(), + ); + let build = |row_group_idx, path: &str, selection| { + let fetcher = SstParquetRangeFetcher::new( + region_file_id, + path.into(), + store.clone(), + if path == "memory.parquet" { + cache.clone() + } else { + CacheStrategy::Disabled + }, + row_group_idx, + None, + ) + .with_memory_tracker(Some(tracker.clone())); + build_sst_parquet_record_batch_stream( + metadata.clone(), + row_group_idx, + selection, + projection.clone(), + fetcher, + path.into(), + 1, + ) + .unwrap() + }; + // Construction (including an unpolled stream's drop) must not reserve memory. + drop(build(0, "memory.parquet", None)); + assert_eq!(0, tracker.current()); + for (row_group_idx, size) in sizes.iter().enumerate() { + let mut first = build(row_group_idx, "memory.parquet", None); + assert_eq!(1, first.next().await.unwrap().unwrap().num_rows()); + assert_eq!(aligned(*size), tracker.current()); + let mut competing = build(row_group_idx, "memory.parquet", None); + assert!(competing.next().await.unwrap().is_err()); + assert_eq!(aligned(*size), tracker.current()); + // Reserve before decoding, even when the selection would require no page reads. + let mut skipped = build( + row_group_idx, + "memory.parquet", + Some(RowSelection::from(vec![ + parquet::arrow::arrow_reader::RowSelector::skip(4), + ])), + ); + assert!(skipped.next().await.unwrap().is_err()); + assert_eq!(aligned(*size), tracker.current()); + while first.next().await.transpose().unwrap().is_some() {} + // Completed streams may remain alive without holding their quota. + assert_eq!(0, tracker.current()); + } + // The second read must acquire quota even when all pages are cached. + store.delete("memory.parquet").await.unwrap(); + let mut cancelled = build(0, "memory.parquet", None); + cancelled.next().await.unwrap().unwrap(); + drop(cancelled); + assert_eq!(0, tracker.current()); + let mut failed = build(0, "missing.parquet", None); + assert!(failed.next().await.unwrap().is_err()); + drop(failed); + assert_eq!(0, tracker.current()); + let mut skipped = build( + 0, + "memory.parquet", + Some(RowSelection::from(vec![ + parquet::arrow::arrow_reader::RowSelector::skip(4), + ])), + ); + assert!(skipped.next().await.is_none()); + assert_eq!(0, tracker.current()); + } + #[test] fn test_assemble_range_from_cached_subrange_and_fetched_tail() { let cached_parts = vec![vec![PageRangePart { diff --git a/src/mito2/src/sst/parquet/reader.rs b/src/mito2/src/sst/parquet/reader.rs index a3cbdef90f1..56488669f5a 100644 --- a/src/mito2/src/sst/parquet/reader.rs +++ b/src/mito2/src/sst/parquet/reader.rs @@ -22,6 +22,7 @@ use api::v1::SemanticType; use arrow_schema::extension::{ EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY, ExtensionType, }; +use common_recordbatch::QueryMemoryTracker; use common_recordbatch::filter::{SimpleFilterEvaluator, TimestampUnitCast}; use common_telemetry::{debug, error, tracing, warn}; use datafusion::physical_plan::PhysicalExpr; @@ -226,9 +227,17 @@ pub struct ParquetReaderBuilder { defer_optional_page_index: bool, /// Scan-wide hint for rows in a decoded batch. batch_size: usize, + /// Optional shared quota for projected row-group decoding. + memory_tracker: Option, } impl ParquetReaderBuilder { + /// Sets the quota used while decoding row groups. + pub(crate) fn memory_tracker(mut self, tracker: Option) -> Self { + self.memory_tracker = tracker; + self + } + /// Returns a new [ParquetReaderBuilder] to read specific SST. pub fn new( table_dir: String, @@ -238,6 +247,7 @@ impl ParquetReaderBuilder { ) -> ParquetReaderBuilder { ParquetReaderBuilder { series_index: None, + memory_tracker: None, table_dir, path_type, file_handle, @@ -656,6 +666,7 @@ impl ParquetReaderBuilder { .collect::>>()?; let reader_builder = RowGroupReaderBuilder { + memory_tracker: self.memory_tracker.clone(), file_handle: self.file_handle.clone(), file_path, parquet_meta, @@ -1768,6 +1779,8 @@ pub(crate) struct RowGroupReaderBuilder { prefilter_builder: Option, /// Hint for rows in a decoded batch. batch_size: usize, + /// Optional shared quota for projected row-group decoding. + memory_tracker: Option, } /// Context passed to [RowGroupReaderBuilder::build()] carrying all information @@ -1956,7 +1969,8 @@ impl RowGroupReaderBuilder { self.cache_strategy.clone(), row_group_idx, fetch_metrics.cloned(), - ); + ) + .with_memory_tracker(self.memory_tracker.clone()); build_sst_parquet_record_batch_stream( self.arrow_metadata.clone(), @@ -2575,6 +2589,7 @@ mod tests { ArrowReaderMetadata::try_new(parquet_meta.clone(), ArrowReaderOptions::new()).unwrap(); ( RowGroupReaderBuilder { + memory_tracker: None, file_handle: file_handle.clone(), file_path, parquet_meta, diff --git a/src/store-api/Cargo.toml b/src/store-api/Cargo.toml index 4a3616afb8a..fefebeacaa0 100644 --- a/src/store-api/Cargo.toml +++ b/src/store-api/Cargo.toml @@ -22,6 +22,7 @@ common-sql.workspace = true common-time.workspace = true common-wal.workspace = true datafusion-common.workspace = true +datafusion-execution.workspace = true datafusion-expr.workspace = true datafusion-physical-plan.workspace = true datatypes.workspace = true diff --git a/src/store-api/src/region_engine.rs b/src/store-api/src/region_engine.rs index 6577880b415..36b8a471394 100644 --- a/src/store-api/src/region_engine.rs +++ b/src/store-api/src/region_engine.rs @@ -26,6 +26,7 @@ use common_error::ext::BoxedError; use common_recordbatch::adapter::RegionQueryStatCounters; use common_recordbatch::{EmptyRecordBatchStream, QueryMemoryTracker, SendableRecordBatchStream}; use common_time::Timestamp; +use datafusion_execution::memory_pool::{MemoryPool, UnboundedMemoryPool}; use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; use datafusion_physical_plan::{DisplayAs, DisplayFormatType, PhysicalExpr}; use datatypes::schema::SchemaRef; @@ -461,10 +462,23 @@ impl PrepareRequest { } /// Necessary context of the query for the scanner. -#[derive(Clone, Default)] +/// Direct callers should create one context per scan and reuse it across partitions. +/// The default pool is unbounded; cloning the context shares its pool. +#[derive(Clone)] pub struct QueryScanContext { /// Whether the query is EXPLAIN ANALYZE VERBOSE. pub explain_verbose: bool, + /// Executing node's query runtime pool, shared by all partitions of this scan. + pub memory_pool: Arc, +} + +impl Default for QueryScanContext { + fn default() -> Self { + Self { + explain_verbose: false, + memory_pool: Arc::new(UnboundedMemoryPool::default()), + } + } } /// A scanner that provides a way to scan the region concurrently. diff --git a/src/table/src/table/scan.rs b/src/table/src/table/scan.rs index fc7aaa24dc6..cae9f5f042c 100644 --- a/src/table/src/table/scan.rs +++ b/src/table/src/table/scan.rs @@ -456,6 +456,7 @@ impl ExecutionPlan for RegionScanExec { let ctx = QueryScanContext { explain_verbose: self.explain_verbose, + memory_pool: context.memory_pool().clone(), }; let stream = self .scanner @@ -732,6 +733,7 @@ mod test { properties: ScannerProperties, metadata: store_api::metadata::RegionMetadataRef, dynamic_filters: Arc>>>, + observed_pool: Arc>>>, } impl fmt::Display for RepeatableScanner { @@ -770,10 +772,11 @@ mod test { fn scan_partition( &self, - _ctx: &QueryScanContext, + ctx: &QueryScanContext, _metrics_set: &ExecutionPlanMetricsSet, _partition: usize, ) -> Result { + *self.observed_pool.lock().unwrap() = Some(ctx.memory_pool.clone()); Ok(self.batches.as_stream()) } @@ -960,6 +963,7 @@ mod test { batches, properties: ScannerProperties::new(vec![vec![], vec![]], false, 0), metadata: metadata.clone(), + observed_pool: Default::default(), dynamic_filters: Arc::new(Mutex::new(Vec::new())), }); let request = ScanRequest { @@ -1069,6 +1073,7 @@ mod test { properties, metadata: metadata.clone(), dynamic_filters: Arc::new(Mutex::new(Vec::new())), + observed_pool: Default::default(), }); let plan = RegionScanExec::new( scanner, @@ -1166,6 +1171,7 @@ mod test { properties, metadata: metadata.clone(), dynamic_filters: Arc::new(Mutex::new(Vec::new())), + observed_pool: Default::default(), }); let plan = RegionScanExec::new( scanner, @@ -1219,6 +1225,31 @@ mod test { } } + #[test] + fn test_execute_forwards_runtime_memory_pool() { + let (schema, batch, metadata) = dynamic_filter_test_data(5686, false); + let observed_pool = Arc::new(Mutex::new(None)); + let scanner = RepeatableScanner { + batches: RecordBatches::try_new(schema, vec![batch]).unwrap(), + properties: ScannerProperties::default(), + metadata, + dynamic_filters: Default::default(), + observed_pool: observed_pool.clone(), + }; + let plan = + Arc::new(RegionScanExec::new(Box::new(scanner), ScanRequest::default(), None).unwrap()); + for _ in 0..2 { + let context = SessionContext::new().task_ctx(); + let pool = context.memory_pool().clone(); + drop(plan.execute(0, context).unwrap()); + assert!(Arc::ptr_eq( + observed_pool.lock().unwrap().as_ref().unwrap(), + &pool + )); + plan.scanner.lock().unwrap().reset_state(); + } + } + #[test] fn test_count_statistics_require_exact_source_rows() { for (append_mode, exact, expected) in [ @@ -1236,6 +1267,7 @@ mod test { .with_total_rows_is_exact(exact), metadata, dynamic_filters: Arc::new(Mutex::new(Vec::new())), + observed_pool: Default::default(), }); let plan = RegionScanExec::new(scanner, ScanRequest::default(), None).unwrap(); let stats = StatisticsContext::new() @@ -1271,6 +1303,7 @@ mod test { properties: ScannerProperties::default(), metadata, dynamic_filters: scanner_dynamic_filters.clone(), + observed_pool: Default::default(), }; let plan = Arc::new(RegionScanExec::new(Box::new(scanner), ScanRequest::default(), None).unwrap()); diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index f522ad39ee5..0bf6228f4c9 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -3455,6 +3455,8 @@ manifest_cache_size = "256MiB" sst_write_buffer_size = "8MiB" max_concurrent_scan_files = 384 allow_stale_entries = false +scan_memory_limit = "25%" +scan_memory_mode = "basic" scan_memory_on_exhausted = "fail" min_compaction_interval = "0s" schedule_compaction_after_edit = true @@ -3548,7 +3550,6 @@ fn drop_lines_with_inconsistent_results(input: String) -> String { "cache_capacity =", "trace_aux_cache_size =", "memory_pool_size =", - "scan_memory_limit =", "sas_token =", "scope =", "num_workers =", From 660c4efe1f6ff4c4958db180b41378fe0fa4de42 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Mon, 19 Oct 2026 10:36:35 +0000 Subject: [PATCH 039/239] fix(mito2): fence writes before a graceful downgrade (#9716) * test(mito2): cover graceful role changes racing with writes A graceful role change flips the region state in a background task while the worker keeps handling writes. Add bulk and row write tests that check a write handled after the role change request is either rejected or covered by the returned last_entry_id. Add tests for DowngradingLeader while an ALTER keeps the region in Altering, alone and in the downgrade sequence the datanode runs (DowngradingLeader, Flush, Follower) with a migration candidate on its own local WAL. Add a test that the downgrade sequence covers a write in flight when a lease renewal sets Downgrading. Signed-off-by: jeremyhi * fix(mito2): fence writes before a graceful downgrade set_region_role_state_gracefully hands the role change to a background task that waits for the manifest lock. The worker keeps handling writes until the task flips the state, so a write could be acknowledged with an entry id above the returned last_entry_id. Set Downgrading in the worker before spawning the task for Follower and DowngradingLeader requests, as lease renewal already does, so every write the worker handles afterwards is rejected. DowngradingLeader also did nothing and returned success when the region was in another leader state such as Altering. The region became writable again when the operation finished, and the datanode then flushed and turned it into a follower in one step. Return a retryable error instead, so metasrv retries the downgrade. Signed-off-by: jeremyhi * test(mito2): assert the retryable downgrade error during ALTER Check that DowngradingLeader during an ALTER returns a retryable error, and release the paused ALTER on that path. The local WAL test now retries the downgrade after the ALTER, as metasrv does. Signed-off-by: jeremyhi * test(mito2): write before the retried downgrade in the local WAL test Rows written before an ALTER were already flushed, so the candidate check did not depend on the downgrade's final flush. Write rows after the ALTER and before the downgrade is retried, which only that flush can carry to the candidate. Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- src/mito2/src/engine/basic_test.rs | 5 +- src/mito2/src/engine/object_store_wal_test.rs | 4 +- src/mito2/src/engine/set_role_state_test.rs | 540 +++++++++++++++++- src/mito2/src/region.rs | 14 +- src/mito2/src/worker.rs | 11 +- 5 files changed, 564 insertions(+), 10 deletions(-) diff --git a/src/mito2/src/engine/basic_test.rs b/src/mito2/src/engine/basic_test.rs index c54691a5a8f..2576158bab6 100644 --- a/src/mito2/src/engine/basic_test.rs +++ b/src/mito2/src/engine/basic_test.rs @@ -1882,7 +1882,10 @@ async fn restart_without_region_close( /// Builds a bulk insert request for a region created by [CreateRequestBuilder] /// with its default columns `(tag_0, field_0, ts)`. -fn build_bulk_insert_request(region_id: RegionId, num_rows: usize) -> RegionBulkInsertsRequest { +pub(crate) fn build_bulk_insert_request( + region_id: RegionId, + num_rows: usize, +) -> RegionBulkInsertsRequest { let schema = Arc::new(Schema::new(vec![ Field::new("tag_0", DataType::Utf8, true), Field::new("field_0", DataType::Float64, true), diff --git a/src/mito2/src/engine/object_store_wal_test.rs b/src/mito2/src/engine/object_store_wal_test.rs index cd27c63b48d..5a61d8b6733 100644 --- a/src/mito2/src/engine/object_store_wal_test.rs +++ b/src/mito2/src/engine/object_store_wal_test.rs @@ -79,13 +79,13 @@ fn open_request(table_dir: String) -> RegionOpenRequest { } } -async fn new_engine(env: &mut TestEnv) -> (MitoEngine, Arc) { +pub(crate) async fn new_engine(env: &mut TestEnv) -> (MitoEngine, Arc) { env.create_engine_with_object_store_wal(MitoConfig::default(), &wal_config()) .await } /// Creates a region with object store WAL options and returns its table dir and row schema. -async fn create_region( +pub(crate) async fn create_region( engine: &MitoEngine, region_id: RegionId, ) -> (String, Vec) { diff --git a/src/mito2/src/engine/set_role_state_test.rs b/src/mito2/src/engine/set_role_state_test.rs index 53aa3511459..94a52c4f3c6 100644 --- a/src/mito2/src/engine/set_role_state_test.rs +++ b/src/mito2/src/engine/set_role_state_test.rs @@ -12,23 +12,36 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::sync::Arc; use std::time::Duration; -use api::v1::Rows; -use common_error::ext::ErrorExt; +use api::v1::{Rows, SemanticType}; +use common_error::ext::{ErrorExt, RetryHint}; use common_error::status_code::StatusCode; +use common_recordbatch::RecordBatches; +use datatypes::prelude::ConcreteDataType; +use datatypes::schema::ColumnSchema; +use store_api::metadata::ColumnMetadata; use store_api::region_engine::{ RegionEngine, RegionRole, SetRegionRoleStateResponse, SetRegionRoleStateSuccess, SettableRegionRoleState, }; use store_api::region_request::{ - EnterStagingRequest, RegionFlushRequest, RegionPutRequest, RegionRequest, + AddColumn, AffectedRows, AlterKind, EnterStagingRequest, PathType, RegionAlterRequest, + RegionCatchupRequest, RegionFlushRequest, RegionOpenRequest, RegionPutRequest, RegionRequest, StagingPartitionDirective, }; -use store_api::storage::RegionId; +use store_api::storage::{RegionId, ScanRequest}; +use tokio::sync::oneshot; use crate::config::MitoConfig; -use crate::region::{RegionLeaderState, RegionRoleState}; +use crate::engine::MitoEngine; +use crate::engine::basic_test::build_bulk_insert_request; +use crate::engine::listener::NotifyRegionChangeResultListener; +use crate::engine::object_store_wal_test; +use crate::region::{MitoRegionRef, RegionLeaderState, RegionRoleState}; +use crate::region_write_ctx::test_hooks::arm_bulk_install_barrier; +use crate::request::WorkerRequest; use crate::test_util::{ CheckpointTaskBlocker, CreateRequestBuilder, TestEnv, build_rows, put_rows, rows_schema, }; @@ -1001,3 +1014,520 @@ async fn test_restricted_state_transitions_with_format(flat_format: bool) { .unwrap(); assert_success_response(&result, 0); } + +const WAIT: Duration = Duration::from_secs(10); + +type WriteReceiver = oneshot::Receiver>; + +/// Submits a graceful role change straight to the region's worker. Requests +/// submitted after this returns reach the worker after it. +async fn submit_role_change( + engine: &MitoEngine, + region_id: RegionId, + state: SettableRegionRoleState, +) -> oneshot::Receiver { + let (request, receiver) = WorkerRequest::new_set_readonly_gracefully(region_id, state); + engine + .inner + .workers + .submit_to_worker(region_id, request) + .await + .unwrap(); + receiver +} + +/// Submits a region request straight to the region's worker. +async fn submit_request( + engine: &MitoEngine, + region_id: RegionId, + request: RegionRequest, +) -> WriteReceiver { + let metadata = engine.inner.get_metadata(region_id).ok(); + let (request, receiver) = + WorkerRequest::try_from_region_request(region_id, request, metadata).unwrap(); + engine + .inner + .workers + .submit_to_worker(region_id, request) + .await + .unwrap(); + receiver +} + +fn put_request(rows: Rows) -> RegionRequest { + RegionRequest::Put(RegionPutRequest { + skip_wal: false, + rows, + hint: None, + partition_expr_version: None, + }) +} + +fn reported_last_entry_id(response: &SetRegionRoleStateResponse) -> u64 { + match response { + SetRegionRoleStateResponse::Success(success) => success.last_entry_id().unwrap(), + _ => panic!("Expected success response, got: {:?}", response), + } +} + +/// Lets the role change task waiting behind `manifest_guard` run and returns +/// the `last_entry_id` it reports. +async fn finish_role_change( + manifest_guard: G, + role_change: oneshot::Receiver, +) -> u64 { + drop(manifest_guard); + let response = tokio::time::timeout(WAIT, role_change) + .await + .expect("role change did not finish") + .unwrap(); + reported_last_entry_id(&response) +} + +/// Waits for the write's result if it is still in flight. +async fn write_result( + early: Option>, + write: WriteReceiver, +) -> crate::error::Result { + match early { + Some(result) => result, + None => tokio::time::timeout(WAIT, write) + .await + .expect("write did not finish") + .unwrap(), + } +} + +/// Checks the promise of `set_region_role_state_gracefully`: an acknowledged +/// write must be covered by the `last_entry_id` the call returned. +fn assert_write_covered( + region: &MitoRegionRef, + result: crate::error::Result, + reported: u64, +) { + let last_entry_id = region.version_control.current().last_entry_id; + assert!( + result.is_err() || last_entry_id <= reported, + "a write was acknowledged ({result:?}) after the role change returned last_entry_id {reported}, \ + but the region's last_entry_id is now {last_entry_id}" + ); +} + +/// The role change task waits for the manifest lock before it changes the +/// role, while the worker keeps handling requests. A bulk write handled after +/// the role change request must be rejected, or be covered by the returned +/// `last_entry_id` if it was accepted. +#[tokio::test(flavor = "multi_thread")] +async fn test_set_follower_gracefully_covers_inflight_bulk_write() { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::with_prefix("follower-inflight-bulk-write").await; + let engine = env.create_engine(MitoConfig::default()).await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let column_schemas = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + put_rows( + &engine, + region_id, + Rows { + schema: column_schemas, + rows: build_rows(0, 3), + }, + ) + .await; + + let region = engine.get_region(region_id).unwrap(); + let manifest_guard = region.manifest_ctx.manifest_manager.read().await; + let role_change = + submit_role_change(&engine, region_id, SettableRegionRoleState::Follower).await; + // An accepted write pauses after its WAL append, before installing its rows. + let mut barrier = arm_bulk_install_barrier(region_id, region.version_control.clone()); + let mut write = submit_request( + &engine, + region_id, + RegionRequest::BulkInserts(build_bulk_insert_request(region_id, 2)), + ) + .await; + let early = tokio::select! { + _ = barrier.wait_until_reached() => None, + result = &mut write => Some(result.unwrap()), + _ = tokio::time::sleep(WAIT) => panic!("bulk write never reached the install barrier"), + }; + let reported = finish_role_change(manifest_guard, role_change).await; + assert_eq!(RegionRoleState::Follower, region.state()); + + barrier.release(); + let result = write_result(early, write).await; + assert_write_covered(®ion, result, reported); +} + +/// Same as [test_set_follower_gracefully_covers_inflight_bulk_write] for a row +/// write. An accepted put waits in its WAL append: the object store WAL +/// acknowledges an append once its object exists, and the test parks that +/// object. +#[tokio::test(flavor = "multi_thread")] +async fn test_set_follower_gracefully_covers_inflight_row_write() { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::with_prefix("follower-inflight-row-write").await; + let (engine, store) = object_store_wal_test::new_engine(&mut env).await; + let region_id = RegionId::new(1, 1); + let (_, schema) = object_store_wal_test::create_region(&engine, region_id).await; + put_rows( + &engine, + region_id, + Rows { + schema: schema.clone(), + rows: build_rows(0, 3), + }, + ) + .await; + + let region = engine.get_region(region_id).unwrap(); + let manifest_guard = region.manifest_ctx.manifest_manager.read().await; + let role_change = + submit_role_change(&engine, region_id, SettableRegionRoleState::Follower).await; + store.hold_creates(); + let mut write = submit_request( + &engine, + region_id, + put_request(Rows { + schema, + rows: build_rows(3, 5), + }), + ) + .await; + let early = tokio::select! { + result = store.wait_for_parked_creates(1) => { + result.unwrap(); + None + } + result = &mut write => Some(result.unwrap()), + _ = tokio::time::sleep(WAIT) => panic!("put never reached the WAL"), + }; + let reported = finish_role_change(manifest_guard, role_change).await; + assert_eq!(RegionRoleState::Follower, region.state()); + + store.release_creates(); + let result = write_result(early, write).await; + assert_write_covered(®ion, result, reported); +} + +/// A lease renewal sets Downgrading outside the worker, so a write that passed +/// the state check before it can still be in flight. The datanode's downgrade +/// steps (DowngradingLeader, Flush, Follower) must still flush that write and +/// report its entry. +#[tokio::test(flavor = "multi_thread")] +async fn test_downgrade_flush_follower_covers_inflight_write() { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::with_prefix("downgrade-flush-follower").await; + let engine = env.create_engine(MitoConfig::default()).await; + let region_id = RegionId::new(1, 1); + engine + .handle_request( + region_id, + RegionRequest::Create(CreateRequestBuilder::new().build()), + ) + .await + .unwrap(); + + let region = engine.get_region(region_id).unwrap(); + let mut barrier = arm_bulk_install_barrier(region_id, region.version_control.clone()); + let write = submit_request( + &engine, + region_id, + RegionRequest::BulkInserts(build_bulk_insert_request(region_id, 2)), + ) + .await; + tokio::time::timeout(WAIT, barrier.wait_until_reached()) + .await + .expect("bulk write never reached the install barrier"); + engine + .set_region_role(region_id, RegionRole::DowngradingLeader) + .unwrap(); + + let downgrade = submit_role_change( + &engine, + region_id, + SettableRegionRoleState::DowngradingLeader, + ) + .await; + barrier.release(); + assert_eq!(2, write.await.unwrap().unwrap()); + reported_last_entry_id(&downgrade.await.unwrap()); + engine + .handle_request( + region_id, + RegionRequest::Flush(RegionFlushRequest::default()), + ) + .await + .unwrap(); + + let response = engine + .set_region_role_state_gracefully(region_id, SettableRegionRoleState::Follower) + .await + .unwrap(); + let reported = reported_last_entry_id(&response); + let current = region.version_control.current(); + assert_eq!(current.last_entry_id, reported); + assert_eq!(current.version.flushed_entry_id, reported); + assert!(current.version.memtables.is_empty()); +} + +fn add_nullable_field1() -> RegionAlterRequest { + RegionAlterRequest { + kind: AlterKind::AddColumns { + columns: vec![AddColumn { + column_metadata: ColumnMetadata { + column_schema: ColumnSchema::new( + "field_1", + ConcreteDataType::float64_datatype(), + true, + ), + semantic_type: SemanticType::Field, + column_id: 3, + }, + location: None, + }], + }, + } +} + +async fn wait_for_state(region: &MitoRegionRef, state: RegionRoleState) { + tokio::time::timeout(WAIT, async { + while region.state() != state { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .unwrap_or_else(|_| panic!("region never reached {state:?}, now {:?}", region.state())); +} + +async fn count_rows(engine: &MitoEngine, region_id: RegionId) -> usize { + let stream = engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + batches.iter().map(|batch| batch.num_rows()).sum() +} + +/// Creates a region with rows `[0, 3)` in an SST and starts an ALTER whose +/// result the worker does not handle until `listener` is woken, so the region +/// stays in Altering. Returns the ALTER task. +async fn create_region_and_start_alter( + env: &TestEnv, + engine: &MitoEngine, + region_id: RegionId, + request: store_api::region_request::RegionCreateRequest, +) -> tokio::task::JoinHandle<()> { + env.get_schema_metadata_manager() + .register_region_table_info( + region_id.table_id(), + "test_table", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + let column_schemas = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + put_rows( + engine, + region_id, + Rows { + schema: column_schemas, + rows: build_rows(0, 3), + }, + ) + .await; + engine + .handle_request( + region_id, + RegionRequest::Flush(RegionFlushRequest::default()), + ) + .await + .unwrap(); + + let cloned_engine = engine.clone(); + let alter = tokio::spawn(async move { + cloned_engine + .handle_request(region_id, RegionRequest::Alter(add_nullable_field1())) + .await + .unwrap(); + }); + let region = engine.get_region(region_id).unwrap(); + wait_for_state( + ®ion, + RegionRoleState::Leader(RegionLeaderState::Altering), + ) + .await; + alter +} + +/// DowngradingLeader while an ALTER keeps the region in Altering must fail +/// with a retryable error, because the region becomes writable again once the +/// ALTER finishes. A retry after the ALTER fences writes. +#[tokio::test(flavor = "multi_thread")] +async fn test_downgrading_leader_during_alter_stops_writes() { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::with_prefix("downgrading-during-alter").await; + let listener = Arc::new(NotifyRegionChangeResultListener::default()); + let engine = env + .create_engine_with(MitoConfig::default(), None, Some(listener.clone()), None) + .await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let column_schemas = rows_schema(&request); + let alter = create_region_and_start_alter(&env, &engine, region_id, request).await; + let region = engine.get_region(region_id).unwrap(); + + let response = engine + .set_region_role_state_gracefully(region_id, SettableRegionRoleState::DowngradingLeader) + .await + .unwrap(); + let state_after_downgrade = region.state(); + listener.wake_notify(); + tokio::time::timeout(WAIT, alter).await.unwrap().unwrap(); + match &response { + SetRegionRoleStateResponse::InvalidTransition(err) => { + assert_eq!(RetryHint::Retryable, err.retry_hint(), "{err:?}") + } + _ => panic!( + "DowngradingLeader returned {response:?} with the region in {state_after_downgrade:?}" + ), + } + + let response = engine + .set_region_role_state_gracefully(region_id, SettableRegionRoleState::DowngradingLeader) + .await + .unwrap(); + reported_last_entry_id(&response); + let err = engine + .handle_request( + region_id, + put_request(Rows { + schema: column_schemas, + rows: build_rows(3, 5), + }), + ) + .await + .unwrap_err(); + assert_eq!(StatusCode::RegionNotReady, err.status_code()); +} + +/// The datanode's downgrade steps (DowngradingLeader, Flush, Follower) with +/// the DowngradingLeader step overlapping an ALTER. The region accepts writes +/// until a downgrade takes effect. A migration candidate opens the region from +/// the object store with its own local WAL, so the final flush must carry +/// every write the old leader acknowledged. +#[tokio::test(flavor = "multi_thread")] +async fn test_downgrade_during_alter_keeps_acknowledged_writes_with_local_wal() { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::with_prefix("downgrade-during-alter-local-wal").await; + let listener = Arc::new(NotifyRegionChangeResultListener::default()); + let engine = env + .create_engine_with(MitoConfig::default(), None, Some(listener.clone()), None) + .await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let column_schemas = rows_schema(&request); + let table_dir = request.table_dir.clone(); + let options = request.options.clone(); + let alter = create_region_and_start_alter(&env, &engine, region_id, request).await; + + // Mirrors `DowngradeRegionsHandler::handle_downgrade_region` in the + // datanode. Metasrv retries the downgrade on a retryable error. + let mut response = engine + .set_region_role_state_gracefully(region_id, SettableRegionRoleState::DowngradingLeader) + .await + .unwrap(); + listener.wake_notify(); + tokio::time::timeout(WAIT, alter).await.unwrap().unwrap(); + put_rows( + &engine, + region_id, + Rows { + schema: column_schemas.clone(), + rows: build_rows(3, 5), + }, + ) + .await; + let mut acknowledged_rows = 5; + if let SetRegionRoleStateResponse::InvalidTransition(err) = &response { + assert_eq!(RetryHint::Retryable, err.retry_hint(), "{err:?}"); + response = engine + .set_region_role_state_gracefully(region_id, SettableRegionRoleState::DowngradingLeader) + .await + .unwrap(); + } + reported_last_entry_id(&response); + engine + .handle_request( + region_id, + RegionRequest::Flush(RegionFlushRequest::default()), + ) + .await + .unwrap(); + if let Ok(output) = engine + .handle_request( + region_id, + put_request(Rows { + schema: column_schemas, + rows: build_rows(5, 8), + }), + ) + .await + { + acknowledged_rows += output.affected_rows; + } + let response = engine + .set_region_role_state_gracefully(region_id, SettableRegionRoleState::Follower) + .await + .unwrap(); + let reported = reported_last_entry_id(&response); + + // The candidate shares the object store but not the local WAL. + let mut candidate_env = TestEnv::with_prefix("downgrade-during-alter-candidate").await; + candidate_env.object_store_manager = env.get_object_store_manager(); + candidate_env.log_store = Some(candidate_env.create_log_store().await); + let candidate = candidate_env.new_mito_engine(MitoConfig::default()).await; + candidate + .handle_request( + region_id, + RegionRequest::Open(RegionOpenRequest { + engine: String::new(), + table_dir, + options, + skip_wal_replay: true, + path_type: PathType::Bare, + checkpoint: None, + requirements: Default::default(), + }), + ) + .await + .unwrap(); + candidate + .handle_request( + region_id, + RegionRequest::Catchup(RegionCatchupRequest { + set_writable: true, + entry_id: Some(reported), + metadata_entry_id: None, + checkpoint: None, + }), + ) + .await + .unwrap(); + assert_eq!( + acknowledged_rows, + count_rows(&candidate, region_id).await, + "the candidate misses rows the old leader acknowledged; Follower reported last_entry_id {reported}" + ); +} diff --git a/src/mito2/src/region.rs b/src/mito2/src/region.rs index d4753fe14bf..ba32f7aff56 100644 --- a/src/mito2/src/region.rs +++ b/src/mito2/src/region.rs @@ -599,12 +599,24 @@ impl MitoRegion { info!("Region {} already in downgrading mode", self.region_id); wait_for_checkpoint = true; } - _ => { + RegionRoleState::Follower => { warn!( "Cannot start downgrade for region {} from state {:?}", self.region_id, current_state ); } + RegionRoleState::Leader(_) => { + // The region may become writable again when the + // current operation, such as an ALTER, finishes. Fail + // with a retryable error instead of reporting a + // downgrade that does not fence writes. + return Err(RegionStateSnafu { + region_id: self.region_id, + state: current_state, + expect: RegionRoleState::Leader(RegionLeaderState::Writable), + } + .build()); + } } None } diff --git a/src/mito2/src/worker.rs b/src/mito2/src/worker.rs index 4bf9d4b244a..763a068e106 100644 --- a/src/mito2/src/worker.rs +++ b/src/mito2/src/worker.rs @@ -52,7 +52,7 @@ use rand::{Rng, rng}; use snafu::{ResultExt, ensure}; use store_api::logstore::LogStore; use store_api::region_engine::{ - SetRegionRoleStateResponse, SetRegionRoleStateSuccess, SettableRegionRoleState, + RegionRole, SetRegionRoleStateResponse, SetRegionRoleStateSuccess, SettableRegionRoleState, }; use store_api::storage::{FileId, RegionId}; use tokio::sync::mpsc::{Receiver, Sender}; @@ -1417,6 +1417,15 @@ impl RegionWorkerLoop { sender: oneshot::Sender, ) { if let Some(region) = self.regions.get_region(region_id) { + if matches!( + region_role_state, + SettableRegionRoleState::Follower | SettableRegionRoleState::DowngradingLeader + ) { + // Fence writes here: writes the worker handles later see + // Downgrading, and earlier ones have finished. This only + // changes a writable or staging leader. + region.set_role(RegionRole::DowngradingLeader); + } // We need to do this in background as we need the manifest lock. common_runtime::spawn_global(async move { match region.set_role_state_gracefully(region_role_state).await { From e0032534d22c111e98880c7840538a6ca4b89462 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Mon, 19 Oct 2026 13:35:07 +0000 Subject: [PATCH 040/239] fix(mito2): build indexes for encoded bulk parts flushed in sync mode (#9722) * fix(mito2): build indexes for encoded bulk parts flushed in sync mode Signed-off-by: Dennis Zhuang * fix(mito2): skip reading the SST when an index build has nothing to create Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/mito2/src/engine/index_build_test.rs | 164 +++++++++++++++++++++++ src/mito2/src/flush.rs | 62 +++++++-- src/mito2/src/request.rs | 2 + src/mito2/src/sst/index.rs | 18 ++- src/mito2/src/worker/handle_flush.rs | 14 +- 5 files changed, 244 insertions(+), 16 deletions(-) diff --git a/src/mito2/src/engine/index_build_test.rs b/src/mito2/src/engine/index_build_test.rs index 6d2241846b4..a00b4798a03 100644 --- a/src/mito2/src/engine/index_build_test.rs +++ b/src/mito2/src/engine/index_build_test.rs @@ -1914,3 +1914,167 @@ async fn test_build_series_index_rejects_non_metric_region() { ); engine.stop().await.unwrap(); } + +/// Counts bytes read from SST files while `COUNTING_SST_READS` is set. +struct SstReadCounter { + inner: object_store::layers::mock::Reader, + is_sst: bool, + bytes: Arc, +} + +impl object_store::layers::mock::Read for SstReadCounter { + async fn open( + &self, + range: object_store::layers::mock::BytesRange, + ) -> object_store::layers::mock::Result<( + object_store::layers::mock::RpRead, + Box, + )> { + self.record(&range); + self.inner.open_dyn(range).await + } + + async fn read( + &self, + range: object_store::layers::mock::BytesRange, + ) -> object_store::layers::mock::Result<( + object_store::layers::mock::RpRead, + object_store::layers::mock::Buffer, + )> { + self.record(&range); + self.inner.read_dyn(range).await + } +} + +impl SstReadCounter { + fn record(&self, range: &object_store::layers::mock::BytesRange) { + if self.is_sst { + self.bytes + .fetch_add(range.size().unwrap_or(u64::MAX) as usize, Ordering::Relaxed); + } + } +} + +#[rstest::rstest] +#[tokio::test] +async fn test_sync_flush_indexes_encoded_bulk_parts( + #[values(false, true)] disable_inverted_on_flush: bool, +) { + let sst_bytes_read = Arc::new(AtomicUsize::new(0)); + let counter = sst_bytes_read.clone(); + let layer = object_store::layers::mock::MockLayerBuilder::default() + .reader_factory(Arc::new(move |path, _op, inner| { + Box::new(SstReadCounter { + inner, + is_sst: path.ends_with(".parquet"), + bytes: counter.clone(), + }) as object_store::layers::mock::Reader + })) + .build() + .unwrap(); + let mut env = TestEnv::with_prefix("test_sync_flush_indexes_encoded_bulk_parts_") + .await + .with_mock_layer(layer); + let listener = Arc::new(IndexBuildListener::default()); + let mut config = MitoConfig { + default_flat_format: true, + min_compaction_interval: Duration::from_secs(3600), + ..Default::default() + }; + assert_eq!(config.index.build_mode, IndexBuildMode::Sync); + if disable_inverted_on_flush { + config.inverted_index.create_on_flush = Mode::Disable; + } + let engine = env + .create_engine_with(config, None, Some(listener.clone()), None) + .await; + + let region_id = RegionId::new(1, 1); + env.get_schema_metadata_manager() + .register_region_table_info( + region_id.table_id(), + "test_table", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + let mut request = CreateRequestBuilder::new() + .insert_option("memtable.type", "bulk") + .insert_option("memtable.bulk.merge_threshold", "2") + .insert_option("memtable.bulk.encode_row_threshold", "1") + .build_with_index(); + if !disable_inverted_on_flush { + // `build_with_index` only adds inverted indexes on tags; also cover a skipping index. + let field = request + .column_metadatas + .iter_mut() + .find(|c| c.column_schema.name == "field_0") + .unwrap(); + field.column_schema = field + .column_schema + .clone() + .with_skipping_options(SkippingIndexOptions::new_unchecked( + 1024, + 0.01, + SkippingIndexType::BloomFilter, + )) + .unwrap(); + } + let column_schemas = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + // Each write is large enough to become its own bulk part, and the SST is larger than the + // footer prefetch so reading the rows is distinguishable from reading the metadata. + for (start, end) in [(0, 100_000), (100_000, 200_000)] { + let rows = Rows { + schema: column_schemas.clone(), + rows: build_rows(start, end), + }; + put_rows(&engine, region_id, rows).await; + } + // Merges the two parts into an encoded part, which flush uploads as-is. + let version = engine.get_region(region_id).unwrap().version(); + for memtable in version.memtables.list_memtables() { + memtable.compact(false).unwrap(); + } + flush_region(&engine, region_id, None).await; + + let files = current_file_metas(&engine, region_id).await; + if disable_inverted_on_flush { + // Only inverted indexes are configured and they are disabled on flush, so the builds + // have nothing to create. They must stop before reading the rows of the SSTs. + let total_size: usize = files.iter().map(|f| f.file_size as usize).sum(); + // Stopping the engine waits for scheduled index builds. + engine.stop().await.unwrap(); + assert!(listener.begin_count() > 0); + assert!(files.iter().all(|f| f.available_indexes.is_empty())); + let read = sst_bytes_read.load(Ordering::Relaxed); + assert!(read < total_size / 2, "read {read} of {total_size} bytes"); + return; + } + + let metadata = engine.get_region(region_id).unwrap().metadata(); + let indexed = async { + loop { + let files = current_file_metas(&engine, region_id).await; + if files + .iter() + .all(|f| f.is_index_consistent_with_region(&metadata.column_metadatas)) + { + return files; + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + }; + let files = tokio::time::timeout(Duration::from_secs(10), indexed) + .await + .expect("flushed encoded parts have no index"); + assert!(listener.begin_count() > 0); + assert!(files.iter().all(|f| f.inverted_index_available())); + assert!(files.iter().all(|f| f.bloom_filter_index_available())); +} diff --git a/src/mito2/src/flush.rs b/src/mito2/src/flush.rs index d5191bdcadb..52eb311de73 100644 --- a/src/mito2/src/flush.rs +++ b/src/mito2/src/flush.rs @@ -29,7 +29,7 @@ use smallvec::{SmallVec, smallvec}; use snafu::{ResultExt, ensure}; use store_api::metadata::RegionMetadataRef; use store_api::region_request::RegionFlushReason; -use store_api::storage::{RegionId, SequenceNumber}; +use store_api::storage::{FileId, RegionId, SequenceNumber}; use strum::IntoStaticStr; use tokio::sync::{Semaphore, mpsc, watch}; @@ -394,7 +394,10 @@ impl RegionFlushTask { }; let worker_request = match flush_result { - Ok(edit) => { + Ok(FlushedEdit { + edit, + unindexed_files, + }) => { let memtables_to_remove = version_data .version .memtables @@ -410,6 +413,7 @@ impl RegionFlushTask { senders: std::mem::take(&mut self.senders), _timer: timer, edit, + unindexed_files, memtables_to_remove, is_staging: self.is_staging, }; @@ -447,7 +451,7 @@ impl RegionFlushTask { version_data: &VersionControlData, state: &CancellableTaskState, uncommitted: &UncommittedSsts, - ) -> Result { + ) -> Result { // We must use the immutable memtables list and entry ids from the `version_data` // for consistency as others might already modify the version in the `version_control`. let version = &version_data.version; @@ -466,6 +470,7 @@ impl RegionFlushTask { let DoFlushMemtablesResult { file_metas, + unindexed_files, flushed_bytes, series_count, encoded_part_count, @@ -589,7 +594,10 @@ impl RegionFlushTask { self.reason.as_str() ); - Ok(edit) + Ok(FlushedEdit { + edit, + unindexed_files, + }) } async fn do_flush_memtables( @@ -604,6 +612,7 @@ impl RegionFlushTask { let mut flushed_bytes = 0; let mut series_count = 0; let mut encoded_part_count = 0; + let mut unindexed_files = Vec::new(); let mut flush_metrics = Metrics::new(WriteType::Flush); let partition_expr = parse_partition_expr(self.partition_expr.as_deref())?; let hook = self.manifest_ctx.hook(); @@ -650,7 +659,12 @@ impl RegionFlushTask { .await?; encoded_part_count += num_encoded; for (source_idx, result) in results.into_iter().enumerate() { - let (max_sequence, ssts_written, metrics) = result?; + let FlushedSsts { + max_sequence, + ssts: ssts_written, + metrics, + indexed, + } = result?; if ssts_written.is_empty() { // No data written. continue; @@ -669,6 +683,9 @@ impl RegionFlushTask { for sst_info in &ssts_written { flushed_bytes += sst_info.file_size; + if !indexed { + unindexed_files.push(sst_info.file_id); + } let pk_range = sst_info .file_metadata .as_ref() @@ -702,6 +719,7 @@ impl RegionFlushTask { Ok(DoFlushMemtablesResult { file_metas, + unindexed_files, flushed_bytes, series_count, encoded_part_count, @@ -747,7 +765,12 @@ impl RegionFlushTask { .await?; uncommitted.track(&ssts); FLUSH_FILE_TOTAL.inc_by(ssts.len() as u64); - Ok((max_sequence, ssts, metrics)) + Ok(FlushedSsts { + max_sequence, + ssts, + metrics, + indexed: true, + }) }); tasks.push(task); } @@ -770,7 +793,13 @@ impl RegionFlushTask { .await?; uncommitted.track(std::slice::from_ref(&encoded.sst_info)); FLUSH_FILE_TOTAL.inc(); - Ok((max_sequence, smallvec![encoded.sst_info], metrics)) + // Encoded parts are uploaded as-is, without passing through the indexer. + Ok(FlushedSsts { + max_sequence, + ssts: smallvec![encoded.sst_info], + metrics, + indexed: false, + }) }); tasks.push(task); } @@ -906,14 +935,30 @@ impl RegionFlushTask { } } +struct FlushedEdit { + edit: RegionEdit, + /// Files in `edit` written without building their indexes. + unindexed_files: Vec, +} + +struct FlushedSsts { + max_sequence: SequenceNumber, + ssts: SstInfoArray, + metrics: Metrics, + /// Whether the indexer ran while writing `ssts`. + indexed: bool, +} + struct FlushFlatMemResult { num_encoded: usize, num_sources: usize, - results: Vec>, + results: Vec>, } struct DoFlushMemtablesResult { file_metas: Vec, + /// Files written without building their indexes. + unindexed_files: Vec, flushed_bytes: u64, series_count: usize, encoded_part_count: usize, @@ -1925,6 +1970,7 @@ mod tests { flushed_sequence: None, committed_sequence: None, }, + unindexed_files: Vec::new(), memtables_to_remove: smallvec![], is_staging: false, }), diff --git a/src/mito2/src/request.rs b/src/mito2/src/request.rs index 45e9041006d..4cf475c7278 100644 --- a/src/mito2/src/request.rs +++ b/src/mito2/src/request.rs @@ -1138,6 +1138,8 @@ pub(crate) struct FlushFinished { pub(crate) _timer: HistogramTimer, /// Region edit to apply. pub(crate) edit: RegionEdit, + /// Files in `edit` written without building their indexes. + pub(crate) unindexed_files: Vec, /// Memtables to remove. pub(crate) memtables_to_remove: SmallVec<[MemtableId; 2]>, /// Whether the region is in staging mode. diff --git a/src/mito2/src/sst/index.rs b/src/mito2/src/sst/index.rs index 8cc332f8d29..6b90d337202 100644 --- a/src/mito2/src/sst/index.rs +++ b/src/mito2/src/sst/index.rs @@ -266,6 +266,13 @@ pub struct Indexer { } impl Indexer { + /// Returns true if no index creator is active. + pub(crate) fn is_empty(&self) -> bool { + self.inverted_indexer.is_none() + && self.fulltext_indexer.is_none() + && self.bloom_filter_indexer.is_none() + } + /// Wraps real creators for update-only benchmarks without a Puffin output. #[cfg(feature = "testing")] pub fn for_bench( @@ -421,11 +428,7 @@ impl IndexerBuilder for IndexerBuilderImpl { indexer.prepare_dense_pk_decoder(&self.metadata); indexer.intermediate_manager = Some(self.intermediate_manager.clone()); - let has_any_indexer = indexer.inverted_indexer.is_some() - || indexer.fulltext_indexer.is_some() - || indexer.bloom_filter_indexer.is_some(); - - if !has_any_indexer { + if indexer.is_empty() { indexer.abort().await; return Indexer::default(); } @@ -840,6 +843,11 @@ impl IndexBuildTask { .indexer_builder .build(region_file_id, new_index_version, row_group_size) .await; + // The build type or index options may leave nothing to build. Return before reading + // any row; an empty indexer finishes with no output, which publishes nothing. + if indexer.is_empty() { + return Ok(IndexBuildOutcome::Finished); + } if let Some(mut parquet_reader) = parquet_reader.take() { // TODO(SNC123): optimize index batch diff --git a/src/mito2/src/worker/handle_flush.rs b/src/mito2/src/worker/handle_flush.rs index 1d4442d8d03..88d911dccc3 100644 --- a/src/mito2/src/worker/handle_flush.rs +++ b/src/mito2/src/worker/handle_flush.rs @@ -439,10 +439,18 @@ impl RegionWorkerLoop { let flush_on_close = request.flush_reason == FlushReason::Closing; - let index_build_file_metas = std::mem::take(&mut request.edit.files_to_add); + let mut index_build_file_metas = std::mem::take(&mut request.edit.files_to_add); + if self.config.index.build_mode == IndexBuildMode::Sync { + // Sync flush builds indexes while writing SSTs, except for encoded bulk parts, + // which are uploaded as-is. Build the missing indexes for those files here. + let metadata = region.metadata(); + index_build_file_metas.retain(|meta| { + request.unindexed_files.contains(&meta.file_id) + && !meta.is_index_consistent_with_region(&metadata.column_metadatas) + }); + } - // In async mode, create indexes after flush. - if self.config.index.build_mode == IndexBuildMode::Async { + if !index_build_file_metas.is_empty() { self.handle_rebuild_index( BuildIndexRequest { region_id, From 1e29e8d4b688f183d446f84e62de533e67c9579b Mon Sep 17 00:00:00 2001 From: fengyue-xve <1759519137@qq.com> Date: Mon, 19 Oct 2026 14:49:14 +0000 Subject: [PATCH 041/239] fix(meta): validate persisted partition bounds before altering regions (#9487) * fix(meta): validate persisted partition bounds before altering regions A partition bound is persisted with the type of the column at the time it was written. When an ALTER TABLE changes the type of a column referenced by a partition expression (widening the time index unit is the common case), each region rewrites the bound literal to the new type while building its scan predicate. A millisecond bound like year 3000 fits milliseconds but overflows nanoseconds, so every scan of the region fails once the table schema is updated -- and some regions may already be altered when the failure surfaces. Validate every region's persisted expression on prepare, before any region is modified: parse the persisted JSON (mirrored types, since common-meta cannot depend on the partition crate which depends on it), replay the same value -> ScalarValue conversion and checked cast the region performs, and reject the ALTER naming the offending region when a bound cannot be represented with the new column type. Incompatible changes are rejected atomically; compatible widening proceeds unchanged. Adds meta-level tests (multi-region rejection naming the second region, compatible widening completing) and an end-to-end sqlness case. Fixes #9454 Signed-off-by: fengyue-xve <1759519137@qq.com> * fix: access ColumnSchema fields directly and use plain format args ColumnSchema exposes name and data_type as public fields, not methods, and the telemetry macros only accept a format string after the error. Signed-off-by: fengyue-xve <1759519137@qq.com> * test: access ColumnSchema data_type field in alter table tests Same fix as the previous commit: data_type is a public field on ColumnSchema, not a method. Signed-off-by: fengyue-xve <1759519137@qq.com> * test: drop borrows now that data_type is read as a field Signed-off-by: fengyue-xve <1759519137@qq.com> * test: drop useless format! in nested expression test Signed-off-by: fengyue-xve <1759519137@qq.com> * test: cover the partition tail in change_col_type sqlness case Signed-off-by: fengyue-xve <1759519137@qq.com> * test: drop spurious blank line in change_col_type result Signed-off-by: fengyue-xve <1759519137@qq.com> * fix(meta): compare partition bounds with the scan comparison coercion The preflight validator cast the persisted bound directly to the altered column type, which rejected ALTERs the scan itself accepts: the region coerces both the column and the bound to the common comparison type (ExprSimplifier::coerce in PredicateGroup::new), so narrowing a field column with an out-of-range bound still scans. Compare through comparison_coercion(target_type, bound_type) instead and only require the bound to be castable to the common type. Also add a round-trip test in the partition crate that serializes real PartitionExprs and runs them through the validator, so drift between the Persisted* mirror and the partition types fails CI, and fix the error message to match the comparison semantics. Signed-off-by: fengyue-xve <1759519137@qq.com> * test: update the sqlness fixture for the new incompatibility message The error message now reflects the comparison semantics: the persisted bound can no longer be compared with the new column type. Signed-off-by: fengyue-xve <1759519137@qq.com> * build(deps): refresh Cargo.lock for the common-meta datafusion-expr dependency common-meta gained a datafusion-expr dependency in ff9c3f4 but the lock file was not refreshed, so cargo check --locked refused to run on CI. Signed-off-by: fengyue-xve <1759519137@qq.com> * fix(meta): reject every non-widening type change of a partitioned column The region routes rows on write through `MultiDimPartitionRule::find_region`, which compares the persisted bound against row values directly: `Value::cmp` panics on mismatched variants and equality misroutes silently. After a column partitioned at `v < 3000000000` is altered from BIGINT to INT, every INSERT compares `Value::Int32` rows against the persisted `Value::Int64` bound and panics in the frontend, even though the scan-side coercion this check mirrors still works. Reject any type change of a column referenced by a persisted partition expression unless it is a timestamp unit widening, which keeps the bound and the rows comparable because `Timestamp` comparisons convert units first; the widening keeps the scan cast check unchanged. Signed-off-by: fengyue-xve <1759519137@qq.com> --------- Signed-off-by: fengyue-xve <1759519137@qq.com> --- Cargo.lock | 1 + src/common/meta/Cargo.toml | 1 + src/common/meta/src/ddl/alter_table.rs | 45 ++ .../src/ddl/alter_table/partition_compat.rs | 565 ++++++++++++++++++ src/common/meta/src/ddl/tests/alter_table.rs | 173 ++++++ src/common/meta/src/error.rs | 17 +- src/partition/src/expr.rs | 117 ++++ .../common/alter/change_col_type.result | 67 +++ .../common/alter/change_col_type.sql | 31 + 9 files changed, 1016 insertions(+), 1 deletion(-) create mode 100644 src/common/meta/src/ddl/alter_table/partition_compat.rs diff --git a/Cargo.lock b/Cargo.lock index 9b895626f1b..946472d1df0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2673,6 +2673,7 @@ dependencies = [ "common-wal", "common-workload", "datafusion-common", + "datafusion-expr", "datatypes", "deadpool", "deadpool-postgres", diff --git a/src/common/meta/Cargo.toml b/src/common/meta/Cargo.toml index bd974deb4cf..64303b070c8 100644 --- a/src/common/meta/Cargo.toml +++ b/src/common/meta/Cargo.toml @@ -46,6 +46,7 @@ common-time.workspace = true common-wal.workspace = true common-workload.workspace = true datafusion-common.workspace = true +datafusion-expr.workspace = true datatypes.workspace = true deadpool = { workspace = true, optional = true } deadpool-postgres = { workspace = true, optional = true } diff --git a/src/common/meta/src/ddl/alter_table.rs b/src/common/meta/src/ddl/alter_table.rs index 2b4ac1c3655..232191def88 100644 --- a/src/common/meta/src/ddl/alter_table.rs +++ b/src/common/meta/src/ddl/alter_table.rs @@ -14,8 +14,11 @@ mod executor; mod metadata; +mod partition_compat; mod region_request; +pub use partition_compat::validate_region_partition_bounds; + use std::collections::HashSet; use std::vec; @@ -147,6 +150,10 @@ impl AlterTableProcedure { )?; self.new_table_info = Some(new_table_info); + // Reject schema changes that the persisted partition expressions of some + // region cannot survive, before any region is altered. + self.validate_region_partition_exprs().await?; + // Safety: Checked in `AlterTableProcedure::new`. let alter_kind = self.data.task.alter_table.kind.as_ref().unwrap(); if sets_skip_wal(alter_kind) { @@ -241,6 +248,44 @@ impl AlterTableProcedure { Ok(Status::executing(true)) } + /// Validates the persisted partition expressions of every region against the + /// schema proposed by this ALTER TABLE, so an ALTER that would break a region's + /// partition predicate is rejected before any region is modified. + async fn validate_region_partition_exprs(&self) -> Result<()> { + // Safety: filled in `fill_table_info`. + let old_schema = &self + .data + .table_info_value + .as_ref() + .unwrap() + .table_info + .meta + .schema; + // Safety: set right after the new table info is built in `on_prepare`. + let new_schema = &self.new_table_info.as_ref().unwrap().meta.schema; + let changed_columns = partition_compat::changed_column_types(old_schema, new_schema); + if changed_columns.is_empty() { + return Ok(()); + } + + let (physical_table_id, physical_table_route) = self + .context + .table_metadata_manager + .table_route_manager() + .get_physical_table_route(self.data.table_id()) + .await?; + if physical_table_id != self.data.table_id() { + // Partition expressions are persisted on the regions of the physical + // table; the altered schema here belongs to a logical table. + return Ok(()); + } + + partition_compat::validate_region_partition_bounds( + &changed_columns, + &physical_table_route.region_routes, + ) + } + fn table_poison_key(&self) -> PoisonKey { table_poison_key(self.data.table_id()) } diff --git a/src/common/meta/src/ddl/alter_table/partition_compat.rs b/src/common/meta/src/ddl/alter_table/partition_compat.rs new file mode 100644 index 00000000000..9fc94004d5c --- /dev/null +++ b/src/common/meta/src/ddl/alter_table/partition_compat.rs @@ -0,0 +1,565 @@ +// 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. + +//! Preflight validation of persisted partition expressions against an ALTER TABLE schema change. +//! +//! Persisted partition bounds keep the types they were written with. A column whose +//! type changes while a partition expression references it breaks the region on both +//! paths that consume the bounds: +//! +//! * On scan, the region coerces the bound literal and the new column type to a +//! common comparison type while building the partition predicate. If no common type +//! exists, or the bound cannot be cast to it (a classic example is a millisecond +//! timestamp bound that overflows when the time index column is widened to +//! nanoseconds), every scan of that region fails at that point. +//! * On write, the region routes rows through `MultiDimPartitionRule::find_region`, +//! which compares the persisted bound against row values directly (`Value::cmp`): a +//! bound whose variant no longer matches the column's panics on ordered comparisons +//! and misroutes silently on equality. Only a timestamp unit widening keeps the two +//! comparable — `Timestamp` comparisons convert units before comparing — so every +//! other type change of a referenced column is rejected outright, and the widening +//! keeps the scan check below. +//! +//! Such an ALTER must be rejected before it is submitted to any region, otherwise +//! already-altered regions keep failing while the rest of the table is untouched. +//! +//! The `partition` crate depends on `common-meta`, so this crate cannot depend on +//! `partition` back and cannot reuse [`partition::expr::PartitionExpr`]. The `Persisted*` +//! types below mirror the JSON representation of `partition::partition::PartitionBound` and +//! `partition::expr::PartitionExpr` and must be kept in sync with `partition/src/partition.rs` +//! and `partition/src/expr.rs`; round-trip tests in `partition` run real serialized +//! expressions through this validator so drift fails in CI. The compatibility check itself +//! deliberately repeats the same conversion (value -> [`ScalarValue`]) and the same +//! comparison coercion the region performs when it builds the partition predicate +//! (`ExprSimplifier::coerce` / [`comparison_coercion`]): both sides are cast to the common +//! comparison type, and the persisted bound must be castable to that type with +//! [`ScalarValue::cast_to`]. A bound that would fail on scan is rejected here as well. + +use std::collections::HashMap; + +use common_telemetry::warn; +use datafusion_common::ScalarValue; +use datafusion_expr::type_coercion::binary::comparison_coercion; +use datatypes::data_type::{ConcreteDataType, DataType}; +use datatypes::schema::Schema; +use datatypes::value::{ + Value, duration_to_scalar_value, time_to_scalar_value, timestamp_to_scalar_value, +}; +use serde::Deserialize; +use store_api::storage::RegionId; + +use crate::error::{PartitionExprIncompatibleSnafu, Result}; +use crate::rpc::router::RegionRoute; + +/// Mirrors `partition::partition::PartitionBound`. +#[derive(Debug, Deserialize)] +enum PersistedPartitionBound { + /// Deprecated since 0.9.0. Only deserialized to keep unknown old formats parseable. + #[allow(dead_code)] + Value(Value), + /// Deprecated since 0.15.0. + MaxValue, + Expr(PersistedPartitionExpr), +} + +/// Mirrors `partition::expr::PartitionExpr`. +#[derive(Debug, Deserialize)] +struct PersistedPartitionExpr { + lhs: Box, + op: PersistedOp, + rhs: Box, +} + +/// Mirrors `partition::expr::Operand`. +#[derive(Debug, Deserialize)] +enum PersistedOperand { + Column(String), + Value(Value), + Expr(PersistedPartitionExpr), +} + +/// Mirrors `partition::expr::RestrictedOp`. +#[derive(Debug, Deserialize)] +enum PersistedOp { + Eq, + NotEq, + Lt, + LtEq, + Gt, + GtEq, + And, + Or, +} + +/// Returns the columns whose data type differs between `old_schema` and `new_schema`, +/// mapped to (old type, new type). Columns only present in one of the schemas (added or +/// dropped) are not included; dropping a column referenced by a partition expression is +/// already rejected by the table metadata validation. +pub(crate) fn changed_column_types( + old_schema: &Schema, + new_schema: &Schema, +) -> HashMap { + let mut changed = HashMap::new(); + for new_column in new_schema.column_schemas() { + let Some(old_column) = old_schema.column_schema_by_name(&new_column.name) else { + continue; + }; + if old_column.data_type != new_column.data_type { + changed.insert( + new_column.name.clone(), + (old_column.data_type.clone(), new_column.data_type.clone()), + ); + } + } + + changed +} + +/// Validates that the persisted partition bounds of every region are still representable +/// with the types of `changed_columns`. Returns an error naming the first offending +/// region without modifying anything. +/// +/// Public so that round-trip tests in the `partition` crate can run real serialized +/// partition expressions through the same mirror used here. +pub fn validate_region_partition_bounds( + changed_columns: &HashMap, + region_routes: &[RegionRoute], +) -> Result<()> { + if changed_columns.is_empty() { + return Ok(()); + } + + for route in region_routes { + let expr_json = route.region.partition_expr(); + if expr_json.is_empty() { + continue; + } + + let bound: PersistedPartitionBound = match serde_json::from_str(&expr_json) { + Ok(bound) => bound, + Err(err) => { + // A malformed persisted expression already fails every scan of this + // region. It is unrelated to the schema change, so don't block the ALTER. + warn!( + "Skipping partition expression validation for region {}, \ + the persisted expression is not parseable: {}, expr: {}", + route.region.id, err, expr_json + ); + continue; + } + }; + let PersistedPartitionBound::Expr(expr) = bound else { + // Deprecated bound formats don't carry expressions and are not used + // when the region builds its partition predicate. + continue; + }; + + validate_expr(&expr, changed_columns, route.region.id)?; + } + + Ok(()) +} + +fn validate_expr( + expr: &PersistedPartitionExpr, + changed_columns: &HashMap, + region_id: RegionId, +) -> Result<()> { + match expr.op { + PersistedOp::And | PersistedOp::Or => { + validate_operand(&expr.lhs, changed_columns, region_id)?; + validate_operand(&expr.rhs, changed_columns, region_id)?; + } + PersistedOp::Eq + | PersistedOp::NotEq + | PersistedOp::Lt + | PersistedOp::LtEq + | PersistedOp::Gt + | PersistedOp::GtEq => { + // Comparison between a column and a bound literal in either order. + validate_comparison(&expr.lhs, &expr.rhs, changed_columns, region_id)?; + validate_comparison(&expr.rhs, &expr.lhs, changed_columns, region_id)?; + } + } + + Ok(()) +} + +fn validate_operand( + operand: &PersistedOperand, + changed_columns: &HashMap, + region_id: RegionId, +) -> Result<()> { + match operand { + PersistedOperand::Expr(expr) => validate_expr(expr, changed_columns, region_id), + PersistedOperand::Column(_) | PersistedOperand::Value(_) => Ok(()), + } +} + +fn validate_comparison( + operand: &PersistedOperand, + other: &PersistedOperand, + changed_columns: &HashMap, + region_id: RegionId, +) -> Result<()> { + let (PersistedOperand::Column(column), PersistedOperand::Value(value)) = (operand, other) + else { + return Ok(()); + }; + let Some((old_type, target_type)) = changed_columns.get(column) else { + return Ok(()); + }; + // Values not representable as a partition bound (lists, structs, json) are + // rejected by the partition parser and cannot appear in a persisted expression. + let Some(scalar) = value_to_scalar_value(value) else { + return Ok(()); + }; + + // The write path (`MultiDimPartitionRule::find_region`) compares the persisted + // bound against row values directly: `Value::cmp` panics on mismatched variants + // and equality misroutes silently. Only a timestamp unit widening keeps the + // bound and the rows comparable (`Timestamp` comparisons convert units first), + // so every other type change of a referenced column is rejected here. + if !old_type.is_timestamp_unit_widening_to(target_type) { + warn!( + "Persisted partition bound would be compared against values of a different \ + type after the ALTER, region_id: {}, column: '{}', target_type: {}", + region_id, column, target_type + ); + return PartitionExprIncompatibleSnafu { + region_id, + column: column.clone(), + target_type: target_type.to_string(), + } + .fail(); + } + + // For the widening, mirror `ExprSimplifier::coerce` in `PredicateGroup::new`: + // the region casts both the column and the persisted bound to the common + // comparison type. The bound is not cast to the column type itself, so the + // widening is accepted as long as the bound survives the common type; a bound + // that cannot be cast to it — or a pair of types without a common comparison + // type — fails every scan. + let target_arrow_type = target_type.as_arrow_type(); + let Some(common_type) = comparison_coercion(&target_arrow_type, &scalar.data_type()) else { + warn!( + "Persisted partition bound has no common comparison type with the altered \ + column type, region_id: {}, column: '{}', target_type: {}", + region_id, column, target_type + ); + return PartitionExprIncompatibleSnafu { + region_id, + column: column.clone(), + target_type: target_type.to_string(), + } + .fail(); + }; + if scalar.data_type() != common_type + && let Err(error) = scalar.cast_to(&common_type) + { + warn!( + error; + "Persisted partition bound is not representable with the altered column type, \ + region_id: {}, column: '{}', target_type: {}", + region_id, + column, + target_type + ); + return PartitionExprIncompatibleSnafu { + region_id, + column: column.clone(), + target_type: target_type.to_string(), + } + .fail(); + } + + Ok(()) +} + +/// Converts a persisted bound value to a [`ScalarValue`], mirroring +/// `partition::expr::Operand::try_as_logical_expr`. Returns `None` for values that +/// cannot appear in a persisted partition expression. +fn value_to_scalar_value(value: &Value) -> Option { + Some(match value { + Value::Boolean(v) => ScalarValue::Boolean(Some(*v)), + Value::UInt8(v) => ScalarValue::UInt8(Some(*v)), + Value::UInt16(v) => ScalarValue::UInt16(Some(*v)), + Value::UInt32(v) => ScalarValue::UInt32(Some(*v)), + Value::UInt64(v) => ScalarValue::UInt64(Some(*v)), + Value::Int8(v) => ScalarValue::Int8(Some(*v)), + Value::Int16(v) => ScalarValue::Int16(Some(*v)), + Value::Int32(v) => ScalarValue::Int32(Some(*v)), + Value::Int64(v) => ScalarValue::Int64(Some(*v)), + Value::Float32(v) => ScalarValue::Float32(Some(v.0)), + Value::Float64(v) => ScalarValue::Float64(Some(v.0)), + Value::String(v) => ScalarValue::Utf8(Some(v.as_utf8().to_string())), + Value::Binary(v) => ScalarValue::Binary(Some(v.to_vec())), + Value::Date(v) => ScalarValue::Date32(Some(v.val())), + Value::Null => ScalarValue::Null, + Value::Timestamp(t) => timestamp_to_scalar_value(t.unit(), Some(t.value())), + Value::Time(t) => time_to_scalar_value(*t.unit(), Some(t.value())).ok()?, + Value::IntervalYearMonth(v) => ScalarValue::IntervalYearMonth(Some(v.to_i32())), + Value::IntervalDayTime(v) => ScalarValue::IntervalDayTime(Some((*v).into())), + Value::IntervalMonthDayNano(v) => ScalarValue::IntervalMonthDayNano(Some((*v).into())), + Value::Duration(d) => duration_to_scalar_value(d.unit(), Some(d.value())), + Value::Decimal128(d) => { + let (v, p, s) = d.to_scalar_value(); + ScalarValue::Decimal128(v, p, s) + } + Value::List(_) | Value::Struct(_) | Value::Json(_) => return None, + }) +} + +#[cfg(test)] +mod tests { + use datatypes::schema::ColumnSchema; + use store_api::storage::RegionId; + + use super::*; + use crate::rpc::router::Region; + + fn schema(ts_type: ConcreteDataType) -> Schema { + Schema::new(vec![ + ColumnSchema::new("ts", ts_type, false), + ColumnSchema::new("host", ConcreteDataType::string_datatype(), true), + ]) + } + + fn ts_bound(millis: i64) -> String { + format!( + r#"{{"Expr":{{"lhs":{{"Column":"ts"}},"op":"Lt","rhs":{{"Value":{{"Timestamp":{{"value":{millis},"unit":"Millisecond"}}}}}}}}}}"# + ) + } + + fn routes(exprs: &[&str]) -> Vec { + exprs + .iter() + .enumerate() + .map(|(index, expr)| RegionRoute { + region: Region { + id: RegionId::new(1024, index as u32 + 1), + partition_expr: expr.to_string(), + ..Default::default() + }, + ..Default::default() + }) + .collect() + } + + fn changed( + column: &str, + old_type: ConcreteDataType, + new_type: ConcreteDataType, + ) -> HashMap { + HashMap::from([(column.to_string(), (old_type, new_type))]) + } + + #[test] + fn test_changed_column_types() { + let old = schema(ConcreteDataType::timestamp_millisecond_datatype()); + let new = schema(ConcreteDataType::timestamp_nanosecond_datatype()); + let changed = changed_column_types(&old, &new); + assert_eq!(changed.len(), 1); + assert_eq!( + changed["ts"], + ( + ConcreteDataType::timestamp_millisecond_datatype(), + ConcreteDataType::timestamp_nanosecond_datatype() + ) + ); + + // Added and dropped columns are not type changes. + let mut new_columns = old.column_schemas().to_vec(); + new_columns.push(ColumnSchema::new( + "extra", + ConcreteDataType::int64_datatype(), + true, + )); + let new = Schema::new(new_columns); + assert!(changed_column_types(&old, &new).is_empty()); + } + + #[test] + fn test_overflowing_bound_is_rejected() { + // Year 3000 fits milliseconds but overflows nanoseconds. + let routes = routes(&[&ts_bound(32_503_680_000_000)]); + let err = validate_region_partition_bounds( + &changed( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + ConcreteDataType::timestamp_nanosecond_datatype(), + ), + &routes, + ) + .unwrap_err(); + assert!( + err.to_string().contains("region 4398046511105"), + "unexpected error: {err}" + ); + assert!(err.to_string().contains("'ts'"), "unexpected error: {err}"); + } + + #[test] + fn test_representable_bound_is_accepted() { + // 2024-01-01T00:00:00Z fits both microseconds and nanoseconds. + let routes = routes(&[&ts_bound(1_704_067_200_000)]); + validate_region_partition_bounds( + &changed( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + ConcreteDataType::timestamp_nanosecond_datatype(), + ), + &routes, + ) + .unwrap(); + validate_region_partition_bounds( + &changed( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + ConcreteDataType::timestamp_microsecond_datatype(), + ), + &routes, + ) + .unwrap(); + } + + #[test] + fn test_later_region_rejection_names_the_region() { + let routes = routes(&[&ts_bound(1_704_067_200_000), &ts_bound(32_503_680_000_000)]); + let err = validate_region_partition_bounds( + &changed( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + ConcreteDataType::timestamp_nanosecond_datatype(), + ), + &routes, + ) + .unwrap_err(); + assert!( + err.to_string().contains("region 4398046511106"), + "unexpected error: {err}" + ); + } + + #[test] + fn test_unrelated_columns_are_skipped() { + let host_bound = + r#"{"Expr":{"lhs":{"Column":"host"},"op":"Eq","rhs":{"Value":{"String":"a"}}}}"#; + // The changed column is not referenced by this bound. + validate_region_partition_bounds( + &changed( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + ConcreteDataType::timestamp_nanosecond_datatype(), + ), + &routes(&[host_bound]), + ) + .unwrap(); + // The referenced column did not change. + validate_region_partition_bounds( + &changed( + "host", + ConcreteDataType::int32_datatype(), + ConcreteDataType::string_datatype(), + ), + &routes(&[&ts_bound(32_503_680_000_000)]), + ) + .unwrap(); + } + + #[test] + fn test_nested_expression_is_validated() { + let expr = r#"{"Expr":{"lhs":{"Expr":{"lhs":{"Column":"ts"},"op":"GtEq","rhs":{"Value":{"Timestamp":{"value":0,"unit":"Millisecond"}}}}},"op":"And","rhs":{"Expr":{"lhs":{"Column":"ts"},"op":"Lt","rhs":{"Value":{"Timestamp":{"value":32503680000000,"unit":"Millisecond"}}}}}}}"#.to_string(); + let err = validate_region_partition_bounds( + &changed( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + ConcreteDataType::timestamp_nanosecond_datatype(), + ), + &routes(&[&expr]), + ) + .unwrap_err(); + assert!(err.to_string().contains("'ts'"), "unexpected error: {err}"); + } + + #[test] + fn test_deprecated_and_empty_bounds_are_skipped() { + validate_region_partition_bounds( + &changed( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + ConcreteDataType::timestamp_nanosecond_datatype(), + ), + &routes(&["", r#""MaxValue""#, r#"{"Value":{"UInt32":1}}"#, "not-json"]), + ) + .unwrap(); + } + + #[test] + fn test_field_narrowing_is_rejected_for_persisted_bounds() { + // BIGINT bound outside the INT range with the column narrowed to INT: the + // scan would coerce both sides to Int64, but the region's write path + // compares the persisted Int64 bound against Int32 rows directly and + // panics, so the ALTER must be rejected. + let bound = + r#"{"Expr":{"lhs":{"Column":"v"},"op":"Lt","rhs":{"Value":{"Int64":3000000000}}}}"#; + let err = validate_region_partition_bounds( + &changed( + "v", + ConcreteDataType::int64_datatype(), + ConcreteDataType::int32_datatype(), + ), + &routes(&[bound]), + ) + .unwrap_err(); + assert!(err.to_string().contains("'v'"), "unexpected error: {err}"); + } + + #[test] + fn test_non_timestamp_widening_is_rejected_for_persisted_bounds() { + // Widening INT to BIGINT still changes the variant rows carry: the + // region's write path compares the persisted Int32 bound against Int64 + // rows directly and panics, so the ALTER must be rejected as well. + let bound = r#"{"Expr":{"lhs":{"Column":"v"},"op":"Lt","rhs":{"Value":{"Int32":10}}}}"#; + let err = validate_region_partition_bounds( + &changed( + "v", + ConcreteDataType::int32_datatype(), + ConcreteDataType::int64_datatype(), + ), + &routes(&[bound]), + ) + .unwrap_err(); + assert!(err.to_string().contains("'v'"), "unexpected error: {err}"); + } + + #[test] + fn test_bound_without_common_comparison_type_is_rejected() { + // A boolean bound has no comparison coercion with a timestamp column, so + // even a widening whose bounds are normally representable fails every + // predicate build; the ALTER must be rejected. + let bound = + r#"{"Expr":{"lhs":{"Column":"ts"},"op":"Lt","rhs":{"Value":{"Boolean":true}}}}"#; + let err = validate_region_partition_bounds( + &changed( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + ConcreteDataType::timestamp_nanosecond_datatype(), + ), + &routes(&[bound]), + ) + .unwrap_err(); + assert!(err.to_string().contains("'ts'"), "unexpected error: {err}"); + } +} diff --git a/src/common/meta/src/ddl/tests/alter_table.rs b/src/common/meta/src/ddl/tests/alter_table.rs index 8233d63b1e5..c67034b1da4 100644 --- a/src/common/meta/src/ddl/tests/alter_table.rs +++ b/src/common/meta/src/ddl/tests/alter_table.rs @@ -1478,3 +1478,176 @@ async fn test_annotation_alter_is_metadata_only() { } } } + +/// Builds the persisted partition expression `ts < ` in the format the +/// partition crate writes to region metadata. +fn test_ts_partition_bound(millis: i64) -> String { + format!( + r#"{{"Expr":{{"lhs":{{"Column":"ts"}},"op":"Lt","rhs":{{"Value":{{"Timestamp":{{"value":{millis},"unit":"Millisecond"}}}}}}}}}}"# + ) +} + +fn test_partitioned_table_route(table_id: u32, region_bounds: &[i64]) -> TableRouteValue { + TableRouteValue::physical( + region_bounds + .iter() + .enumerate() + .map(|(index, bound)| RegionRoute { + region: Region { + id: RegionId::new(table_id, index as u32 + 1), + partition_expr: test_ts_partition_bound(*bound), + ..Default::default() + }, + leader_peer: Some(Peer::empty(index as u64 + 1)), + ..Default::default() + }) + .collect(), + ) +} + +fn test_widen_ts_unit_task(table_name: &str, target_type: ColumnDataType) -> AlterTableTask { + AlterTableTask { + alter_table: AlterTableExpr { + catalog_name: DEFAULT_CATALOG_NAME.to_string(), + schema_name: DEFAULT_SCHEMA_NAME.to_string(), + table_name: table_name.to_string(), + kind: Some(Kind::ModifyColumnTypes(api::v1::ModifyColumnTypes { + modify_column_types: vec![api::v1::ModifyColumnType { + column_name: "ts".to_string(), + target_type: target_type as i32, + target_type_extension: None, + }], + })), + }, + } +} + +#[tokio::test] +async fn test_widen_ts_unit_rejects_incompatible_partition_bounds_before_altering() { + let node_manager = Arc::new(MockDatanodeManager::new(())); + let ddl_context = new_ddl_context(node_manager); + let table_name = "foo"; + let table_id = 1024; + let task = test_create_table_task(table_name, table_id); + + // Region 1 holds 2024-01-01, which fits nanoseconds. Region 2 holds 3000-01-01, + // which fits milliseconds but overflows nanoseconds. + let routes = test_partitioned_table_route(table_id, &[1_704_067_200_000, 32_503_680_000_000]); + ddl_context + .table_metadata_manager + .create_table_metadata(task.table_info.clone(), routes.clone(), HashMap::new()) + .await + .unwrap(); + + let task = test_widen_ts_unit_task(table_name, ColumnDataType::TimestampNanosecond); + let mut procedure = AlterTableProcedure::new(table_id, task, ddl_context.clone()).unwrap(); + let err = procedure.on_prepare().await.unwrap_err(); + + assert_matches!(err, Error::PartitionExprIncompatible { .. }); + assert_eq!(err.status_code(), StatusCode::InvalidArguments); + // The rejection names the offending region and column. + assert!( + err.to_string().contains("4398046511106"), + "unexpected error: {err}" + ); + assert!(err.to_string().contains("'ts'"), "unexpected error: {err}"); + + // The rejected ALTER left the table schema and the whole route unchanged. + let table_info = ddl_context + .table_metadata_manager + .table_info_manager() + .get(table_id) + .await + .unwrap() + .unwrap() + .into_inner() + .table_info; + assert_eq!( + table_info + .meta + .schema + .column_schema_by_name("ts") + .unwrap() + .data_type, + ConcreteDataType::timestamp_millisecond_datatype() + ); + let (physical_table_id, physical_table_route) = ddl_context + .table_metadata_manager + .table_route_manager() + .get_physical_table_route(table_id) + .await + .unwrap(); + assert_eq!(physical_table_id, table_id); + assert_eq!( + routes, + TableRouteValue::physical(physical_table_route.region_routes) + ); +} + +#[tokio::test] +async fn test_widen_ts_unit_with_compatible_partition_bounds_completes() { + let node_manager = Arc::new(MockDatanodeManager::new(())); + let ddl_context = new_ddl_context(node_manager); + let table_name = "foo"; + let table_id = 1024; + let task = test_create_table_task(table_name, table_id); + + // Both persisted bounds fit nanoseconds. + let routes = test_partitioned_table_route(table_id, &[0, 1_704_067_200_000]); + ddl_context + .table_metadata_manager + .create_table_metadata(task.table_info.clone(), routes, HashMap::new()) + .await + .unwrap(); + + let task = test_widen_ts_unit_task(table_name, ColumnDataType::TimestampNanosecond); + let procedure_id = ProcedureId::random(); + let provider = Arc::new(MockContextProvider::default()); + let mut procedure = AlterTableProcedure::new(table_id, task, ddl_context.clone()).unwrap(); + procedure.on_prepare().await.unwrap(); + procedure + .submit_alter_region_requests(procedure_id, provider.as_ref()) + .await + .unwrap(); + procedure.mut_data().set_column_metadatas(vec![ + ColumnMetadata { + column_schema: ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_nanosecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: 0, + }, + ColumnMetadata { + column_schema: ColumnSchema::new("host", ConcreteDataType::string_datatype(), false), + semantic_type: SemanticType::Tag, + column_id: 1, + }, + ColumnMetadata { + column_schema: ColumnSchema::new("cpu", ConcreteDataType::float64_datatype(), false), + semantic_type: SemanticType::Field, + column_id: 2, + }, + ]); + procedure.on_update_metadata().await.unwrap(); + + let table_info = ddl_context + .table_metadata_manager + .table_info_manager() + .get(table_id) + .await + .unwrap() + .unwrap() + .into_inner() + .table_info; + assert_eq!( + table_info + .meta + .schema + .column_schema_by_name("ts") + .unwrap() + .data_type, + ConcreteDataType::timestamp_nanosecond_datatype() + ); +} diff --git a/src/common/meta/src/error.rs b/src/common/meta/src/error.rs index adfa158adf0..3de9e4a3a46 100644 --- a/src/common/meta/src/error.rs +++ b/src/common/meta/src/error.rs @@ -1149,6 +1149,20 @@ pub enum Error { actual_column_id: u32, }, + #[snafu(display( + "Partition expression of region {} is incompatible with the altered column '{}': the persisted bound can no longer be compared with the new column type {}", + region_id, + column, + target_type + ))] + PartitionExprIncompatible { + region_id: RegionId, + column: String, + target_type: String, + #[snafu(implicit)] + location: Location, + }, + #[cfg(feature = "enterprise")] #[snafu(display("Too large duration"))] TooLargeDuration { @@ -1268,7 +1282,8 @@ impl ErrorExt for Error { | InvalidFileName { .. } | InvalidFlowRequestBody { .. } | InvalidFilePath { .. } - | ConflictingSchemaOptions { .. } => StatusCode::InvalidArguments, + | ConflictingSchemaOptions { .. } + | PartitionExprIncompatible { .. } => StatusCode::InvalidArguments, #[cfg(feature = "enterprise")] MissingInterval { .. } | NegativeDuration { .. } | TooLargeDuration { .. } => { diff --git a/src/partition/src/expr.rs b/src/partition/src/expr.rs index cb15f5ae5d7..121533acc60 100644 --- a/src/partition/src/expr.rs +++ b/src/partition/src/expr.rs @@ -463,6 +463,14 @@ impl Display for PartitionExpr { #[cfg(test)] mod tests { + use std::collections::HashMap; + + use common_meta::ddl::alter_table::validate_region_partition_bounds; + use common_meta::error::Error; + use common_meta::rpc::router::{Region, RegionRoute}; + use datatypes::data_type::ConcreteDataType; + use store_api::storage::RegionId; + use super::*; #[test] @@ -863,4 +871,113 @@ mod tests { assert!(columns.contains("b")); assert!(columns.contains("c")); } + + /// Real serialized expressions must survive the round trip through + /// `common-meta`'s persisted-expression mirror: whenever that mirror drifts + /// from these types, deserialization fails, validation is skipped and the + /// overflowing cases below stop producing `PartitionExprIncompatible`. + #[test] + fn test_serialized_expressions_validate_against_type_changes() { + let overflow_ms = serde_json::from_str::( + r#"{"Timestamp":{"value":32503680000000,"unit":"Millisecond"}}"#, + ) + .unwrap(); + let start_ms = + serde_json::from_str::(r#"{"Timestamp":{"value":0,"unit":"Millisecond"}}"#) + .unwrap(); + let ts_ms = ConcreteDataType::timestamp_millisecond_datatype(); + let ts_ns = ConcreteDataType::timestamp_nanosecond_datatype(); + + let routes = |expr: &PartitionExpr| { + vec![RegionRoute { + region: Region { + id: RegionId::new(1024, 1), + partition_expr: expr.as_json_str().unwrap(), + ..Default::default() + }, + ..Default::default() + }] + }; + let changed = |column: &str, old_type: ConcreteDataType, new_type: ConcreteDataType| { + HashMap::from([(column.to_string(), (old_type, new_type))]) + }; + + // Compile-time tripwire: exhaustively matching the operand and operator + // variants makes a new variant break this test at compile time, so the + // `partition_compat::Persisted*` mirror and these cases must be extended + // together instead of silently skipping validation. + fn walk(expr: &PartitionExpr) { + walk_operand(&expr.lhs); + walk_operand(&expr.rhs); + match expr.op { + RestrictedOp::Eq + | RestrictedOp::NotEq + | RestrictedOp::Lt + | RestrictedOp::LtEq + | RestrictedOp::Gt + | RestrictedOp::GtEq + | RestrictedOp::And + | RestrictedOp::Or => {} + } + } + fn walk_operand(operand: &Operand) { + match operand { + Operand::Column(_) | Operand::Value(_) => {} + Operand::Expr(expr) => walk(expr), + } + } + + // Every expression shape must be rejected while the bound overflows + // nanoseconds, with the error naming the offending column. + let simple = col("ts").lt(overflow_ms.clone()); + let nested_and = col("ts") + .gt_eq(start_ms.clone()) + .and(col("ts").lt(overflow_ms.clone())); + let nested_or = PartitionExpr::new( + Operand::Expr(col("host").eq(Value::String("a".into()))), + RestrictedOp::Or, + Operand::Expr(col("ts").lt(overflow_ms.clone())), + ); + for expr in [&simple, &nested_and, &nested_or] { + walk(expr); + let err = validate_region_partition_bounds( + &changed("ts", ts_ms.clone(), ts_ns.clone()), + &routes(expr), + ) + .unwrap_err(); + assert!( + matches!( + err, + Error::PartitionExprIncompatible { ref column, .. } if column == "ts" + ), + "unexpected error: {err}" + ); + } + + // A representable bound passes. + let compatible = col("ts").lt(start_ms); + validate_region_partition_bounds(&changed("ts", ts_ms, ts_ns), &routes(&compatible)) + .unwrap(); + + // Narrowing a field column must be rejected as well: the region's write + // path compares the persisted Int64 bound against Int32 rows directly + // and panics, even though the scan would coerce both sides. + let narrowing = col("v").lt(Value::Int64(3_000_000_000)); + let err = validate_region_partition_bounds( + &changed( + "v", + ConcreteDataType::int64_datatype(), + ConcreteDataType::int32_datatype(), + ), + &routes(&narrowing), + ) + .unwrap_err(); + assert!( + matches!( + err, + Error::PartitionExprIncompatible { ref column, .. } if column == "v" + ), + "unexpected error: {err}" + ); + } } diff --git a/tests/cases/standalone/common/alter/change_col_type.result b/tests/cases/standalone/common/alter/change_col_type.result index 348d3b2145d..d362215792b 100644 --- a/tests/cases/standalone/common/alter/change_col_type.result +++ b/tests/cases/standalone/common/alter/change_col_type.result @@ -386,3 +386,70 @@ DROP TABLE ts_overflow; Affected Rows: 0 +-- Issue 9454: persisted partition bounds keep their original unit. Widening the +-- time index unit must be rejected before any region is altered when a persisted +-- bound overflows the new unit. +CREATE TABLE ts_bound_overflow (host STRING, ts TIMESTAMP TIME INDEX) +PARTITION ON COLUMNS (ts) ( + ts < '2024-01-01 00:00:00', + ts >= '2024-01-01 00:00:00' AND ts < '3000-01-01 00:00:00', + ts >= '3000-01-01 00:00:00' +); + +Affected Rows: 0 + +INSERT INTO ts_bound_overflow VALUES ("a", "2024-06-01 12:00:00"); + +Affected Rows: 1 + +-- 3000-01-01 fits milliseconds but overflows nanoseconds, so the widening is +-- rejected before any region is altered. +-- SQLNESS REPLACE \d+\(\d+,\s+\d+\) REDACTED +ALTER TABLE ts_bound_overflow MODIFY COLUMN ts TIMESTAMP_NS; + +Error: 1004(InvalidArguments), Partition expression of region REDACTED is incompatible with the altered column 'ts': the persisted bound can no longer be compared with the new column type TimestampNanosecond + +-- The rejected ALTER leaves the schema and the stored data unchanged. +DESCRIBE ts_bound_overflow; + ++--------+----------------------+-----+------+---------+---------------+ +| Column | Type | Key | Null | Default | Semantic Type | ++--------+----------------------+-----+------+---------+---------------+ +| host | String | | YES | | FIELD | +| ts | TimestampMillisecond | PRI | NO | | TIMESTAMP | ++--------+----------------------+-----+------+---------+---------------+ + +SELECT * FROM ts_bound_overflow ORDER BY ts; + ++------+---------------------+ +| host | ts | ++------+---------------------+ +| a | 2024-06-01T12:00:00 | ++------+---------------------+ + +-- Widening to microseconds keeps every persisted bound representable. +ALTER TABLE ts_bound_overflow MODIFY COLUMN ts TIMESTAMP_US; + +Affected Rows: 0 + +DESCRIBE ts_bound_overflow; + ++--------+----------------------+-----+------+---------+---------------+ +| Column | Type | Key | Null | Default | Semantic Type | ++--------+----------------------+-----+------+---------+---------------+ +| host | String | | YES | | FIELD | +| ts | TimestampMicrosecond | PRI | NO | | TIMESTAMP | ++--------+----------------------+-----+------+---------+---------------+ + +SELECT * FROM ts_bound_overflow ORDER BY ts; + ++------+---------------------+ +| host | ts | ++------+---------------------+ +| a | 2024-06-01T12:00:00 | ++------+---------------------+ + +DROP TABLE ts_bound_overflow; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/alter/change_col_type.sql b/tests/cases/standalone/common/alter/change_col_type.sql index bb16af8847e..c7d206eaf7c 100644 --- a/tests/cases/standalone/common/alter/change_col_type.sql +++ b/tests/cases/standalone/common/alter/change_col_type.sql @@ -136,3 +136,34 @@ ALTER TABLE ts_overflow MODIFY COLUMN ts TIMESTAMP_NS; DESCRIBE ts_overflow; DROP TABLE ts_overflow; + +-- Issue 9454: persisted partition bounds keep their original unit. Widening the +-- time index unit must be rejected before any region is altered when a persisted +-- bound overflows the new unit. +CREATE TABLE ts_bound_overflow (host STRING, ts TIMESTAMP TIME INDEX) +PARTITION ON COLUMNS (ts) ( + ts < '2024-01-01 00:00:00', + ts >= '2024-01-01 00:00:00' AND ts < '3000-01-01 00:00:00', + ts >= '3000-01-01 00:00:00' +); + +INSERT INTO ts_bound_overflow VALUES ("a", "2024-06-01 12:00:00"); + +-- 3000-01-01 fits milliseconds but overflows nanoseconds, so the widening is +-- rejected before any region is altered. +-- SQLNESS REPLACE \d+\(\d+,\s+\d+\) REDACTED +ALTER TABLE ts_bound_overflow MODIFY COLUMN ts TIMESTAMP_NS; + +-- The rejected ALTER leaves the schema and the stored data unchanged. +DESCRIBE ts_bound_overflow; + +SELECT * FROM ts_bound_overflow ORDER BY ts; + +-- Widening to microseconds keeps every persisted bound representable. +ALTER TABLE ts_bound_overflow MODIFY COLUMN ts TIMESTAMP_US; + +DESCRIBE ts_bound_overflow; + +SELECT * FROM ts_bound_overflow ORDER BY ts; + +DROP TABLE ts_bound_overflow; From 7397ab6b230cc617bf5f5b7530a703df5b4c82a7 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Mon, 19 Oct 2026 14:57:15 +0000 Subject: [PATCH 042/239] perf(datanode): send the mito puts of one request to each worker at once (#9692) * perf(datanode): send the mito puts of one request to each worker at once The region server handled the puts of a request to several regions as separate requests. A Mito worker woke on the first of them and wrote it alone; in the durable object store WAL mode its append then blocked the worker until the next seal, so the rest of the request went into the next object, and clients writing a partitioned table locked into one phase that wastes a tick. When every put of a request goes to a ready mito region, the region server now hands them to `MitoEngine::put_regions_batch`, which submits the puts of each worker in one message, so the worker writes its regions of the request in one batch and one WAL append. Requests to one region and to other engines are handled as before. Signed-off-by: jeremyhi * perf(datanode): await each region of a batched mito put on its own task The batched path waited for the regions of a request in order and sent the worker messages one after another, so an error or a full worker queue of one region held back the others, and its metrics counted the request as a whole. `MitoEngine::put_regions_batch` now sends the worker messages concurrently on their own task and returns each region's result. The region server awaits and accounts every region on its own ingest task, as the per-region path does: the first error returns at once, and each region records its latency, changed rows or failure and is named in its error. A request with several puts to one region keeps the per-region path. Signed-off-by: jeremyhi * perf(datanode): spawn the accounting of batched mito puts from the batch task The request spawned the per-region accounting tasks only after the batch task returned, so a request dropped in between left accepted writes unaccounted. The batch task now spawns them and the request only joins their handles. Also state in the architecture invariants that the region server's batched puts to the metric and mito engines are the one request dispatch that does not go through the RegionEngine trait. Signed-off-by: jeremyhi * perf(datanode): batch mito puts through the RegionEngine trait Add `RegionEngine::handle_batch_put_requests`, which returns each region's response of the puts of one request. The default handles each put with `handle_request`, so other engines keep their behaviour; Mito overrides it with the one message per worker submission. The region server calls the trait method instead of downcasting to `MitoEngine`, and the architecture invariants no longer need an exception. Signed-off-by: jeremyhi * perf(datanode): trace each region of a batched mito put Attach the per-region `RegionServer::handle_region_request` span of the per-region path to each region's task of a batched mito put, so the region ids stay in distributed traces. Signed-off-by: jeremyhi * perf(datanode): let a batch put response answer for several regions `RegionEngine::handle_batch_put_requests` now returns each response with the regions it answers for. The default and Mito answer each region on its own; an engine that writes several puts together, such as the metric engine's batched puts, can answer them with one response. The region server accounts a response for each of its regions. Signed-off-by: jeremyhi * perf(datanode): batch any request with several mito puts Take the batch path for a request with at least two puts to ready mito regions, also when they all target one region: a request of many puts to one region otherwise still sends one worker message per put. Each put keeps its own response, affected rows and error. Signed-off-by: jeremyhi * perf(datanode): account batched mito puts in one task Await the responses of a batched mito put in the batch task with a FuturesUnordered instead of one spawned task per put. The task accounts each response as it completes, hands the first error back to the caller at once and keeps accounting the rest, so a request of many puts no longer spawns a task per put. Signed-off-by: jeremyhi * refactor(mito2): name the batched worker message BatchWrite Signed-off-by: jeremyhi * refactor(mito2): submit batched writes on the ingest runtime Grouping and submitting the writes of a request is part of ingestion, so the task runs on the ingest runtime and takes the write task class when the workload scheduler is enabled. Signed-off-by: jeremyhi * refactor(datanode): share ingest result accounting across put paths The per-region path, the metric batch path and the mito batch path each recorded the changed rows or the failures of an ingest response. They now call one helper, so a metric or label change happens in one place. A batch put response lists the region of each put it answers for, so the failure count stays per put when two puts go to one region; the trait documentation now says so. Signed-off-by: jeremyhi * fix(mito2): give batched puts of a stopped worker the worker's error A batched write message that could not reach a stopped worker was dropped, so its puts failed with "Failed to recv result" while a put submitted on its own fails with "Worker is stopped". The worker now returns the message it could not take, and each of its puts receives the WorkerStopped error. Signed-off-by: jeremyhi * refactor(datanode): look up the batch latency histogram once The label of the request latency histogram is the same for every put of a batch, so the batch task looks it up once instead of per observation. Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- src/datanode/src/region_server.rs | 467 +++++++++++++++++++++++++++-- src/mito2/AGENTS.md | 3 +- src/mito2/src/engine.rs | 53 +++- src/mito2/src/request.rs | 3 + src/mito2/src/worker.rs | 64 +++- src/store-api/src/region_engine.rs | 41 ++- 6 files changed, 587 insertions(+), 44 deletions(-) diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index 0b7c7645e00..0b7d36e9f2c 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -27,7 +27,7 @@ use std::pin::Pin; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, RwLock}; use std::task::{Context, Poll}; -use std::time::Duration; +use std::time::{Duration, Instant}; use api::region::RegionResponse; use api::v1::meta::TopicStat; @@ -59,6 +59,7 @@ use datafusion_common::tree_node::TreeNode; use datatypes::schema::SchemaRef; use either::Either; use futures_util::future::try_join_all; +use futures_util::stream::FuturesUnordered; use futures_util::{Stream, StreamExt}; use metric_engine::engine::MetricEngine; use mito2::engine::{MITO_ENGINE_NAME, MitoEngine}; @@ -95,7 +96,7 @@ use store_api::region_request::{ RegionOpenRequest, RegionRequest, }; use store_api::storage::RegionId; -use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc}; +use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc, oneshot}; use tokio::time::{self, timeout}; use tonic::{Request, Response, Result as TonicResult}; @@ -551,6 +552,10 @@ impl RegionServer { let requests = RegionRequest::try_from_request_body(request).context(BuildRegionRequestsSnafu)?; + let requests = match self.try_handle_mito_batch_puts(requests).await? { + Either::Left(response) => return Ok(response), + Either::Right(requests) => requests, + }; // Try to optimize batch Put requests for metric engine // Returns either Some(response) or None(requests_back) match self.try_handle_metric_batch_puts(requests).await? { @@ -704,24 +709,126 @@ impl RegionServer { .map_err(BoxedError::new) .context(HandleRegionRequestSnafu { region_id: first_region_id, + }) + .map(RegionResponse::new); + record_ingest_result(request_type, batch_size, &result); + result.map(Either::Left) + } + + /// Puts into ready mito regions together through + /// [`RegionEngine::handle_batch_put_requests`], so that each mito worker gets the puts of the + /// request to its regions at once and writes them in one batch. One ingest task awaits every + /// response at once and accounts each, for the regions it answers, as the per-region path + /// does; it hands the first error back at once and keeps accounting the rest, also when the + /// caller goes away. + /// + /// Returns Either::Right(original_requests) unless + /// [`mito_engine_for_batch_puts`](Self::mito_engine_for_batch_puts) accepts the requests. + async fn try_handle_mito_batch_puts( + &self, + requests: Vec<(RegionId, RegionRequest)>, + ) -> Result>> { + let Some(engine) = self.mito_engine_for_batch_puts(&requests) else { + return Ok(Either::Right(requests)); + }; + + let request_type = requests[0].1.request_type(); + let puts = requests + .into_iter() + .map(|(region_id, request)| match request { + RegionRequest::Put(put) => (region_id, put), + _ => unreachable!("Already checked all are Put"), + }) + .collect(); + let started = Instant::now(); + let tracing_context = TracingContext::from_current_span(); + let (result_tx, result_rx) = oneshot::channel(); + let batch = common_runtime::spawn_ingest(async move { + let mut pending = engine + .handle_batch_put_requests(puts) + .into_iter() + .map(|(region_ids, response)| { + let span = tracing_context.attach(info_span!( + "RegionServer::handle_region_request", + region_id = region_ids[0].to_string() + )); + async move { (region_ids, response.await) }.trace(span) + }) + .collect::>(); + let histogram = + crate::metrics::HANDLE_REGION_REQUEST_ELAPSED.with_label_values(&[request_type]); + let mut result_tx = Some(result_tx); + let mut affected_rows = 0; + let mut extensions = HashMap::new(); + while let Some((region_ids, result)) = pending.next().await { + let elapsed = started.elapsed().as_secs_f64(); + for _ in ®ion_ids { + histogram.observe(elapsed); + } + let result = result.context(HandleRegionRequestSnafu { + region_id: region_ids[0], + }); + record_ingest_result(request_type, region_ids.len(), &result); + match result { + Ok(response) => { + affected_rows += response.affected_rows; + extensions.extend(response.extensions); + } + Err(err) => { + if let Some(result_tx) = result_tx.take() { + let _ = result_tx.send(Err(err)); + } + } + } + } + if let Some(result_tx) = result_tx { + let _ = result_tx.send(Ok(RegionResponse { + affected_rows, + extensions, + metadata: Vec::new(), + })); + } }); - match result { - Ok(total_affected) => { - crate::metrics::REGION_CHANGED_ROW_COUNT - .with_label_values(&[request_type]) - .inc_by(total_affected as u64); - Ok(Either::Left(RegionResponse::new(total_affected))) - } - Err(err) => { - crate::metrics::REGION_SERVER_INSERT_FAIL_COUNT - .with_label_values(&[request_type, err.status_code().as_ref()]) - .inc_by(batch_size as u64); - Err(err) + match result_rx.await { + Ok(result) => result.map(Either::Left), + // The batch task answers on every path it completes, so it panicked or was cancelled. + Err(_) => { + batch.await.context(RuntimeJoinSnafu { request_type })?; + UnexpectedSnafu { + violated: "the batch put task completed without answering", + } + .fail() } } } + /// Returns the mito engine when there are at least two requests, every one a put into a ready + /// mito region. + fn mito_engine_for_batch_puts( + &self, + requests: &[(RegionId, RegionRequest)], + ) -> Option { + if requests.len() < 2 { + return None; + } + let first_region_id = requests[0].0; + let all_puts_to_ready_mito_regions = requests.iter().all(|(region_id, request)| { + matches!(request, RegionRequest::Put(_)) + && matches!( + self.inner.region_map.get(region_id).as_deref(), + Some(RegionEngineWithStatus::Ready(engine)) if engine.name() == MITO_ENGINE_NAME + ) + }); + if !all_puts_to_ready_mito_regions { + return None; + } + match self.inner.region_map.get(&first_region_id).as_deref() { + Some(RegionEngineWithStatus::Ready(engine)) => Some(engine.clone()), + _ => None, + } + } + async fn handle_sync_region_request(&self, request: &SyncRequest) -> Result { let region_id = RegionId::from_u64(request.region_id); let manifest_info = request @@ -1698,18 +1805,15 @@ impl RegionServerInner { // Sets corresponding region status to registering/deregistering before the operation. self.set_region_status_not_ready(region_id, &engine, ®ion_change); - match engine + let result = engine .handle_request(region_id, request) .await - .with_context(|_| HandleRegionRequestSnafu { region_id }) - { + .with_context(|_| HandleRegionRequestSnafu { region_id }); + if matches!(region_change, RegionChange::Ingest) { + record_ingest_result(request_type, 1, &result); + } + match result { Ok(result) => { - // Update metrics - if matches!(region_change, RegionChange::Ingest) { - crate::metrics::REGION_CHANGED_ROW_COUNT - .with_label_values(&[request_type]) - .inc_by(result.affected_rows as u64); - } // Sets corresponding region status to ready. self.set_region_status_ready(region_id, engine.clone(), region_change) .await?; @@ -1721,11 +1825,6 @@ impl RegionServerInner { }) } Err(err) => { - if matches!(region_change, RegionChange::Ingest) { - crate::metrics::REGION_SERVER_INSERT_FAIL_COUNT - .with_label_values(&[request_type, err.status_code().as_ref()]) - .inc(); - } // Removes the region status if the operation fails. self.unset_region_status(region_id, &engine, region_change); Err(err) @@ -2052,6 +2151,19 @@ fn record_region_open_failures(engine: &str, status_code: StatusCode, count: usi .inc_by(count as u64); } +/// Records the result of `count` ingest requests answered by one response: the rows they +/// changed, or their failures. +fn record_ingest_result(request_type: &str, count: usize, result: &Result) { + match result { + Ok(response) => crate::metrics::REGION_CHANGED_ROW_COUNT + .with_label_values(&[request_type]) + .inc_by(response.affected_rows as u64), + Err(err) => crate::metrics::REGION_SERVER_INSERT_FAIL_COUNT + .with_label_values(&[request_type, err.status_code().as_ref()]) + .inc_by(count as u64), + } +} + fn parse_region_attribute( engine: &str, options: &HashMap, @@ -2094,15 +2206,22 @@ mod tests { use std::collections::HashMap; use std::sync::Arc; + use api::v1::region::{InsertRequest, InsertRequests}; use api::v1::{Rows, SemanticType}; use common_error::ext::{ErrorExt, RetryHint}; use common_recordbatch::RecordBatches; use common_recordbatch::adapter::{RecordBatchMetrics, RegionWatermarkEntry}; + use common_wal::config::object_store::{ObjectStoreWalConfig, STANDALONE_GENERATION}; + use common_wal::options::{ObjectStoreWalOptions, WAL_OPTIONS_KEY, WalOptions}; use datatypes::prelude::{ConcreteDataType, VectorRef}; use datatypes::schema::{ColumnSchema, Schema}; use datatypes::vectors::Int32Vector; use futures_util::StreamExt; - use mito2::test_util::CreateRequestBuilder; + use log_store::ObjectStoreLogStore; + use mito2::config::MitoConfig; + use mito2::test_util::{CreateRequestBuilder, TestEnv, build_rows, rows_schema}; + use object_store::ObjectStore; + use object_store::services::Memory; use query::options::FLOW_RETURN_REGION_SEQ; use store_api::metadata::{ColumnMetadata, RegionMetadata, RegionMetadataBuilder}; use store_api::region_engine::RegionEngine; @@ -3059,4 +3178,294 @@ mod tests { .await .unwrap_err(); } + + /// Returns a region server with a Mito engine on a durable object store WAL, the WAL, and the + /// WAL options of a region on it. + async fn region_server_on_object_store_wal( + env: &mut TestEnv, + config: MitoConfig, + ) -> (RegionServer, Arc, String) { + let wal_config = ObjectStoreWalConfig::default(); + let log_store = ObjectStoreLogStore::try_new( + ObjectStore::new(Memory::default()).unwrap(), + &wal_config, + 0, + STANDALONE_GENERATION, + ) + .await + .unwrap(); + let engine = env + .create_engine_with_log_store(config, log_store.clone()) + .await; + let mut region_server = mock_region_server(); + region_server.register_engine(Arc::new(engine)); + let wal_options = serde_json::to_string(&WalOptions::ObjectStore( + ObjectStoreWalOptions::new(wal_config.node_prefix(0, STANDALONE_GENERATION)), + )) + .unwrap(); + (region_server, log_store, wal_options) + } + + /// Creates a region with `wal_options` and returns an insert of three rows into it. + async fn create_region( + region_server: &RegionServer, + region_id: RegionId, + wal_options: &str, + ) -> InsertRequest { + let request = CreateRequestBuilder::new() + .insert_option(WAL_OPTIONS_KEY, wal_options) + .build(); + let rows = Rows { + schema: rows_schema(&request), + rows: build_rows(0, 3), + }; + region_server + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + InsertRequest { + region_id: region_id.as_u64(), + rows: Some(rows), + partition_expr_version: None, + skip_wal: false, + } + } + + #[tokio::test] + async fn test_mito_puts_of_one_request_share_a_wal_append() { + // One worker that takes one message at a time, so that only puts that + // arrive in one message share a WAL append. + let mut env = TestEnv::with_prefix("region-server-mito-batch-puts").await; + let config = MitoConfig { + num_workers: 1, + worker_request_batch_size: 1, + ..Default::default() + }; + let (region_server, log_store, wal_options) = + region_server_on_object_store_wal(&mut env, config).await; + let mut requests = Vec::new(); + for region_number in 1..=2 { + let region_id = RegionId::new(1, region_number); + requests.push(create_region(®ion_server, region_id, &wal_options).await); + } + + let response = region_server + .handle_requests_in_parallel(region_request::Body::Inserts(InsertRequests { requests })) + .await + .unwrap(); + + assert_eq!(6, response.affected_rows); + // The worker got both puts at once and appended them together. + log_store.wait_for_admitted_appends(1).await.unwrap(); + assert!( + tokio::time::timeout(Duration::ZERO, log_store.wait_for_admitted_appends(2)) + .await + .is_err() + ); + } + + #[tokio::test] + async fn test_mito_puts_to_one_region_share_a_wal_append() { + // As above: only puts that arrive in one message share a WAL append. + let mut env = TestEnv::with_prefix("region-server-mito-batch-puts-one-region").await; + let config = MitoConfig { + num_workers: 1, + worker_request_batch_size: 1, + ..Default::default() + }; + let (region_server, log_store, wal_options) = + region_server_on_object_store_wal(&mut env, config).await; + let insert = create_region(®ion_server, RegionId::new(1, 1), &wal_options).await; + let requests = vec![insert.clone(), insert]; + + let response = region_server + .handle_requests_in_parallel(region_request::Body::Inserts(InsertRequests { requests })) + .await + .unwrap(); + + assert_eq!(6, response.affected_rows); + log_store.wait_for_admitted_appends(1).await.unwrap(); + assert!( + tokio::time::timeout(Duration::ZERO, log_store.wait_for_admitted_appends(2)) + .await + .is_err() + ); + } + + #[tokio::test] + async fn test_mito_batch_puts_fail_without_waiting_for_other_regions() { + let mut env = TestEnv::with_prefix("region-server-mito-batch-puts-error").await; + let config = MitoConfig { + num_workers: 2, + ..Default::default() + }; + let (region_server, log_store, wal_options) = + region_server_on_object_store_wal(&mut env, config).await; + // Regions 1 and 2 of table 1 belong to different workers. The put to region 1 waits for a + // WAL object whose create is held; Mito does not know region 2, so its put fails at once. + let held = RegionId::new(1, 1); + let unknown = RegionId::new(1, 2); + let held_insert = create_region(®ion_server, held, &wal_options).await; + let unknown_insert = InsertRequest { + region_id: unknown.as_u64(), + ..held_insert.clone() + }; + let mito = region_server + .inner + .region_map + .get(&held) + .unwrap() + .clone() + .into_engine(); + region_server + .inner + .region_map + .insert(unknown, RegionEngineWithStatus::Ready(mito)); + log_store.hold_creates(); + + let requests = vec![held_insert, unknown_insert]; + let err = tokio::time::timeout( + Duration::from_secs(10), + region_server.handle_requests_in_parallel(region_request::Body::Inserts( + InsertRequests { requests }, + )), + ) + .await + .unwrap() + .unwrap_err(); + + assert_eq!(StatusCode::RegionNotFound, err.status_code()); + log_store.release_creates(); + } + + #[tokio::test] + async fn test_mito_batch_puts_to_a_stopped_worker_fail_as_a_single_put_does() { + let mut env = TestEnv::with_prefix("region-server-mito-batch-puts-stopped").await; + let (region_server, _log_store, wal_options) = + region_server_on_object_store_wal(&mut env, MitoConfig::default()).await; + let region_id = RegionId::new(1, 1); + let insert = create_region(®ion_server, region_id, &wal_options).await; + let mito = region_server + .inner + .region_map + .get(®ion_id) + .unwrap() + .clone() + .into_engine(); + mito.stop().await.unwrap(); + + // One put takes the per-region path and two take the batch path. + for requests in [vec![insert.clone()], vec![insert.clone(), insert]] { + let err = region_server + .handle_requests_in_parallel(region_request::Body::Inserts(InsertRequests { + requests, + })) + .await + .unwrap_err(); + assert!(format!("{err:?}").contains("is stopped"), "{err:?}"); + } + } + + #[tokio::test] + async fn test_default_batch_puts_handle_each_put_as_a_request() { + let failing = RegionId::new(1, 2); + let (engine, _receiver) = MockRegionEngine::with_mock_fn( + "mock", + Box::new(move |region_id, request| { + assert!(matches!(request, RegionRequest::Put(_))); + if region_id == failing { + return error::UnexpectedSnafu { + violated: "test".to_string(), + } + .fail(); + } + Ok(region_id.region_number() as usize) + }), + ); + let region_ids = [1, 2, 3].map(|n| RegionId::new(1, n)); + let puts = region_ids + .iter() + .map(|region_id| { + let request = RegionPutRequest { + rows: Rows { + schema: Vec::new(), + rows: Vec::new(), + }, + hint: None, + skip_wal: false, + partition_expr_version: None, + }; + (*region_id, request) + }) + .collect(); + + let (answered, responses): (Vec<_>, Vec<_>) = + engine.handle_batch_put_requests(puts).into_iter().unzip(); + let responses = futures_util::future::join_all(responses).await; + + let answered_regions = region_ids.map(|region_id| vec![region_id]); + assert_eq!(answered_regions.to_vec(), answered); + assert_eq!(1, responses[0].as_ref().unwrap().affected_rows); + assert!(responses[1].is_err()); + assert_eq!(3, responses[2].as_ref().unwrap().affected_rows); + } + + #[tokio::test] + async fn test_only_several_puts_into_ready_mito_regions_are_batched() { + let mut env = TestEnv::with_prefix("region-server-mito-batch-puts-guard").await; + let mito: RegionEngineRef = Arc::new(env.create_engine(MitoConfig::default()).await); + let (other, _receiver) = MockRegionEngine::new("mock"); + let region_server = mock_region_server(); + let [mito_a, mito_b, registering, other_region] = [1, 2, 3, 4].map(|n| RegionId::new(1, n)); + let region_map = ®ion_server.inner.region_map; + region_map.insert(mito_a, RegionEngineWithStatus::Ready(mito.clone())); + region_map.insert(mito_b, RegionEngineWithStatus::Ready(mito.clone())); + region_map.insert(registering, RegionEngineWithStatus::Registering(mito)); + region_map.insert(other_region, RegionEngineWithStatus::Ready(other)); + let rows = || Rows { + schema: Vec::new(), + rows: Vec::new(), + }; + let put = |region_id| { + let request = RegionPutRequest { + rows: rows(), + hint: None, + skip_wal: false, + partition_expr_version: None, + }; + (region_id, RegionRequest::Put(request)) + }; + let delete = |region_id| { + let request = RegionDeleteRequest { + rows: rows(), + hint: None, + partition_expr_version: None, + }; + (region_id, RegionRequest::Delete(request)) + }; + + for requests in [ + vec![put(mito_a), put(mito_b)], + vec![put(mito_a), put(mito_a)], + ] { + assert!( + region_server + .mito_engine_for_batch_puts(&requests) + .is_some() + ); + } + for requests in [ + vec![put(mito_a)], + vec![put(mito_a), delete(mito_b)], + vec![put(mito_a), put(registering)], + vec![put(mito_a), put(other_region)], + vec![put(other_region), put(mito_a)], + ] { + assert!( + region_server + .mito_engine_for_batch_puts(&requests) + .is_none() + ); + } + } } diff --git a/src/mito2/AGENTS.md b/src/mito2/AGENTS.md index 08600cd00e4..c41aa649111 100644 --- a/src/mito2/AGENTS.md +++ b/src/mito2/AGENTS.md @@ -38,7 +38,8 @@ snapshot isolation). It implements the `RegionEngine` trait from `store-api`. ## Write path -`MitoEngine::handle_request` (`engine.rs`) → worker loop +`MitoEngine::handle_request` (`engine.rs`), or `handle_batch_put_requests` for the puts +of one request to several regions (one message per worker) → worker loop (`worker/handle_write.rs`) → sequence + WAL assembly (`region_write_ctx.rs`) → `wal.rs` → memtable (`memtable/`) → when buffer pressure trips, `flush.rs` writes SSTs via `access_layer.rs` and appends a `RegionEdit` to the manifest diff --git a/src/mito2/src/engine.rs b/src/mito2/src/engine.rs index cdfc78ff0c4..715357a39b3 100644 --- a/src/mito2/src/engine.rs +++ b/src/mito2/src/engine.rs @@ -108,7 +108,8 @@ use common_recordbatch::{QueryMemoryTracker, SendableRecordBatchStream}; use common_stat::get_total_memory_bytes; use common_telemetry::{debug, info, tracing, warn}; use common_wal::options::WalOptions; -use futures::future::{join_all, try_join_all}; +use futures::FutureExt; +use futures::future::{self, join_all, try_join_all}; use futures::stream::{self, Stream, StreamExt}; use object_store::manager::ObjectStoreManagerRef; use region_hook::RegionHookRef; @@ -122,14 +123,14 @@ use store_api::metric_engine_consts::{ MANIFEST_INFO_EXTENSION_KEY, TABLE_COLUMN_METADATA_EXTENSION_KEY, }; use store_api::region_engine::{ - BatchResponses, MitoCopyRegionFromRequest, MitoCopyRegionFromResponse, RegionEngine, - RegionManifestInfo, RegionRole, RegionScannerRef, RegionStatistic, RemapManifestsRequest, - RemapManifestsResponse, SetRegionRoleStateResponse, SettableRegionRoleState, - SyncRegionFromRequest, SyncRegionFromResponse, + BatchPutResponse, BatchResponses, MitoCopyRegionFromRequest, MitoCopyRegionFromResponse, + RegionEngine, RegionManifestInfo, RegionRole, RegionScannerRef, RegionStatistic, + RemapManifestsRequest, RemapManifestsResponse, SetRegionRoleStateResponse, + SettableRegionRoleState, SyncRegionFromRequest, SyncRegionFromResponse, }; use store_api::region_info::RegionInfoEntry; use store_api::region_request::{ - AffectedRows, RegionCatchupRequest, RegionOpenRequest, RegionRequest, + AffectedRows, RegionCatchupRequest, RegionOpenRequest, RegionPutRequest, RegionRequest, }; use store_api::sst_entry::{ManifestSstEntry, PuffinIndexMetaEntry, StorageSstEntry}; use store_api::storage::{FileId, FileRefsManifest, RegionId, ScanRequest, SequenceNumber}; @@ -1407,6 +1408,46 @@ impl RegionEngine for MitoEngine { .map_err(BoxedError::new) } + /// Sends the puts to regions of the same worker in one message, so that the worker writes + /// them in one batch and one WAL append. The messages are sent on their own task, so every + /// response can be awaited at once. + fn handle_batch_put_requests( + self: Arc, + requests: Vec<(RegionId, RegionPutRequest)>, + ) -> Vec { + let mut writes = Vec::with_capacity(requests.len()); + let mut responses = Vec::with_capacity(requests.len()); + for (region_id, request) in requests { + let request = RegionRequest::Put(request); + let timer = HANDLE_REQUEST_ELAPSED + .with_label_values(&[request.request_type()]) + .start_timer(); + let region_metadata = self.inner.get_metadata(region_id).ok(); + let response = + match WorkerRequest::try_from_region_request(region_id, request, region_metadata) { + Ok((WorkerRequest::Write(write), receiver)) => { + writes.push(write); + async move { + let _timer = timer; + receiver + .await + .context(RecvSnafu) + .flatten() + .map(RegionResponse::new) + .map_err(BoxedError::new) + } + .boxed() + } + Ok(_) => unreachable!("a put converts to a write"), + Err(e) => future::ready(Err(BoxedError::new(e))).boxed(), + }; + responses.push((vec![region_id], response)); + } + let inner = self.inner.clone(); + common_runtime::spawn_ingest(async move { inner.workers.submit_writes(writes).await }); + responses + } + #[tracing::instrument(skip_all)] async fn handle_request( &self, diff --git a/src/mito2/src/request.rs b/src/mito2/src/request.rs index 4cf475c7278..d084fc82c7e 100644 --- a/src/mito2/src/request.rs +++ b/src/mito2/src/request.rs @@ -758,6 +758,9 @@ pub(crate) enum WorkerRequest { /// Write to a region. Write(SenderWriteRequest), + /// Writes of one request to regions of the same worker. + BatchWrite(Vec), + /// Ddl request to a region. Ddl(SenderDdlRequest), diff --git a/src/mito2/src/worker.rs b/src/mito2/src/worker.rs index 763a068e106..556f4a0757f 100644 --- a/src/mito2/src/worker.rs +++ b/src/mito2/src/worker.rs @@ -44,12 +44,12 @@ use common_meta::key::SchemaMetadataManagerRef; use common_runtime::JoinHandle; use common_stat::get_total_memory_bytes; use common_telemetry::{error, info, warn}; -use futures::future::try_join_all; +use futures::future::{join_all, try_join_all}; use object_store::ObjectStore; use object_store::manager::ObjectStoreManagerRef; use prometheus::{Histogram, IntGauge}; use rand::{Rng, rng}; -use snafu::{ResultExt, ensure}; +use snafu::ResultExt; use store_api::logstore::LogStore; use store_api::region_engine::{ RegionRole, SetRegionRoleStateResponse, SetRegionRoleStateSuccess, SettableRegionRoleState, @@ -318,6 +318,40 @@ impl WorkerGroup { self.worker(region_id).submit_request(request).await } + /// Submits the writes of one request, those of each worker in one message, so that the + /// worker handles them in one batch. The messages are sent concurrently, so a full worker + /// queue holds back no other worker. A write whose worker has stopped receives the error a + /// write submitted on its own gets. + pub(crate) async fn submit_writes(&self, writes: Vec) { + let mut writes_by_worker: HashMap> = HashMap::new(); + for write in writes { + let index = region_id_to_index(write.request.region_id, self.workers.len()); + writes_by_worker.entry(index).or_default().push(write); + } + join_all( + writes_by_worker + .into_iter() + .map(|(index, writes)| async move { + let worker = &self.workers[index]; + let Err(request) = worker + .try_submit_request(WorkerRequest::BatchWrite(writes)) + .await + else { + return; + }; + let WorkerRequest::BatchWrite(writes) = request else { + unreachable!("the worker returns the request it was given"); + }; + for write in writes { + write + .sender + .send(WorkerStoppedSnafu { id: worker.id }.fail()); + } + }), + ) + .await; + } + /// Returns true if the specific region exists. pub(crate) fn is_region_exists(&self, region_id: RegionId) -> bool { self.worker(region_id).is_region_exists(region_id) @@ -738,9 +772,23 @@ pub(crate) struct RegionWorker { impl RegionWorker { /// Submits request to background worker thread. async fn submit_request(&self, request: WorkerRequest) -> Result<()> { - ensure!(self.is_running(), WorkerStoppedSnafu { id: self.id }); + self.try_submit_request(request) + .await + .map_err(|_| WorkerStoppedSnafu { id: self.id }.build()) + } + + /// Submits request to background worker thread, or returns it if the worker is stopped. + async fn try_submit_request( + &self, + request: WorkerRequest, + ) -> std::result::Result<(), WorkerRequest> { + if !self.is_running() { + return Err(request); + } let request_with_time = WorkerRequestWithTime::new(request); - if self.sender.send(request_with_time).await.is_err() { + if let Err(mpsc::error::SendError(request_with_time)) = + self.sender.send(request_with_time).await + { warn!( "Worker {} is already exited but the running flag is still true", self.id @@ -750,7 +798,7 @@ impl RegionWorker { if let Some(state) = &self.series_index_task_state { state.stop(); } - return WorkerStoppedSnafu { id: self.id }.fail(); + return Err(request_with_time.request); } Ok(()) @@ -1057,6 +1105,7 @@ impl RegionWorkerLoop { match request_with_time.request { WorkerRequest::Write(sender_req) => write_req_buffer.push(sender_req), + WorkerRequest::BatchWrite(sender_reqs) => write_req_buffer.extend(sender_reqs), WorkerRequest::Ddl(sender_req) => ddl_req_buffer.push(sender_req), WorkerRequest::BulkInserts(bulk_insert) => { bulk_insert_req_num += 1; @@ -1110,6 +1159,9 @@ impl RegionWorkerLoop { match request_with_time.request { WorkerRequest::Write(sender_req) => write_req_buffer.push(sender_req), + WorkerRequest::BatchWrite(sender_reqs) => { + write_req_buffer.extend(sender_reqs) + } WorkerRequest::Ddl(sender_req) => ddl_req_buffer.push(sender_req), WorkerRequest::BulkInserts(bulk_insert) => { bulk_insert_req_num += 1; @@ -1184,7 +1236,7 @@ impl RegionWorkerLoop { ) { for worker_req in general_requests.drain(..) { match worker_req { - WorkerRequest::Write(_) | WorkerRequest::Ddl(_) => { + WorkerRequest::Write(_) | WorkerRequest::BatchWrite(_) | WorkerRequest::Ddl(_) => { // These requests are categorized before dispatching general requests. continue; } diff --git a/src/store-api/src/region_engine.rs b/src/store-api/src/region_engine.rs index 36b8a471394..c95bf6d3b21 100644 --- a/src/store-api/src/region_engine.rs +++ b/src/store-api/src/region_engine.rs @@ -30,14 +30,15 @@ use datafusion_execution::memory_pool::{MemoryPool, UnboundedMemoryPool}; use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; use datafusion_physical_plan::{DisplayAs, DisplayFormatType, PhysicalExpr}; use datatypes::schema::SchemaRef; -use futures::future::join_all; +use futures::FutureExt; +use futures::future::{BoxFuture, join_all}; use serde::{Deserialize, Serialize}; use tokio::sync::Semaphore; use crate::logstore::entry; use crate::metadata::RegionMetadataRef; use crate::region_request::{ - BatchRegionDdlRequest, RegionCatchupRequest, RegionOpenRequest, RegionRequest, + BatchRegionDdlRequest, RegionCatchupRequest, RegionOpenRequest, RegionPutRequest, RegionRequest, }; use crate::storage::{FileId, RegionId, ScanRequest, SequenceNumber}; @@ -548,6 +549,13 @@ pub type RegionScannerRef = Box; pub type BatchResponses = Vec<(RegionId, Result)>; +/// A response of [`RegionEngine::handle_batch_put_requests`] and the region of each put it +/// answers for. +pub type BatchPutResponse = ( + Vec, + BoxFuture<'static, Result>, +); + /// Represents the statistics of a region. #[derive(Debug, Deserialize, Serialize, Default)] pub struct RegionStatistic { @@ -981,6 +989,35 @@ pub trait RegionEngine: Send + Sync { Ok(join_all(tasks).await) } + /// Handles the puts of one request into regions of this engine and returns the responses + /// to await on their own, each with the region of every put it answers for: two puts to one + /// region are two entries. + /// + /// The default handles each put with [`handle_request`](Self::handle_request) and answers + /// each put on its own. An engine that writes several puts together may answer them with + /// one response whose affected rows count all of them. + fn handle_batch_put_requests( + self: Arc, + requests: Vec<(RegionId, RegionPutRequest)>, + ) -> Vec + where + Self: 'static, + { + requests + .into_iter() + .map(|(region_id, request)| { + let engine = self.clone(); + let response = async move { + engine + .handle_request(region_id, RegionRequest::Put(request)) + .await + } + .boxed(); + (vec![region_id], response) + }) + .collect() + } + async fn handle_batch_ddl_requests( &self, request: BatchRegionDdlRequest, From da88ff86534a95496dc11f359e1e2daa8f1d7cfa Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Mon, 19 Oct 2026 15:22:13 +0000 Subject: [PATCH 043/239] fix(partition): route NULLs like row inserts and accept signed-minimum bounds (#9697) * fix(partition): route NULLs like row inserts and accept signed-minimum bounds Signed-off-by: Dennis Zhuang * fix(partition): keep equality two-valued for NULL rows and accept unsigned -0 Signed-off-by: Dennis Zhuang * fix(partition): treat NULL masks as false in batch routing and keep comparisons bare Signed-off-by: Dennis Zhuang * test(compat): update partition filter in analyze verbose expected output Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/common/sql/src/convert.rs | 68 +++++++++- src/partition/src/expr.rs | 119 +++++++----------- src/partition/src/multi_dim.rs | 105 +++++++++++++++- .../cases/standalone/common/partition.result | 30 +++++ tests/cases/standalone/common/partition.sql | 16 +++ .../verify.result | 2 +- 6 files changed, 261 insertions(+), 79 deletions(-) diff --git a/src/common/sql/src/convert.rs b/src/common/sql/src/convert.rs index fdd90ebe3df..8b84d591ef0 100644 --- a/src/common/sql/src/convert.rs +++ b/src/common/sql/src/convert.rs @@ -128,12 +128,29 @@ pub fn sql_value_to_value( column_schema: &ColumnSchema, sql_val: &SqlValue, timezone: Option<&Timezone>, - unary_op: Option, + mut unary_op: Option, auto_string_to_numeric: bool, ) -> Result { let column_name = &column_schema.name; let data_type = &column_schema.data_type; let mut value = match sql_val { + // Parse the sign with the digits, otherwise the minimum of a signed integer + // or timestamp overflows before it is negated. Other types keep the negation + // below: it accepts unsigned `-0`, and the decimal parser depends on length. + SqlValue::Number(n, _) + if unary_op == Some(UnaryOperator::Minus) + && matches!( + data_type, + ConcreteDataType::Int8(_) + | ConcreteDataType::Int16(_) + | ConcreteDataType::Int32(_) + | ConcreteDataType::Int64(_) + | ConcreteDataType::Timestamp(_) + ) => + { + unary_op = None; + sql_number_to_value(data_type, &format!("-{n}"))? + } SqlValue::Number(n, _) => sql_number_to_value(data_type, n)?, SqlValue::Null => Value::Null, SqlValue::Boolean(b) => { @@ -946,6 +963,55 @@ mod test { ); } + #[test] + fn test_negative_number_at_type_minimum() { + for (data_type, digits, expected) in [ + ( + ConcreteDataType::int8_datatype(), + "128", + Value::Int8(i8::MIN), + ), + ( + ConcreteDataType::int32_datatype(), + "2147483648", + Value::Int32(i32::MIN), + ), + ( + ConcreteDataType::int64_datatype(), + "9223372036854775808", + Value::Int64(i64::MIN), + ), + ] { + let value = call_sql_value_to_value!( + "test", + data_type, + &SqlValue::Number(digits.into(), false), + unary_op = UnaryOperator::Minus + ) + .unwrap(); + assert_eq!(value, expected); + } + let value = call_sql_value_to_value!( + "test", + ConcreteDataType::uint64_datatype(), + &SqlValue::Number("0".into(), false), + unary_op = UnaryOperator::Minus + ) + .unwrap(); + assert_eq!(value, Value::UInt64(0)); + + // 28 characters without the sign, 29 with it: the sign must not move the + // literal to the other decimal parser. + let value = call_sql_value_to_value!( + "test", + ConcreteDataType::decimal128_datatype(38, 30), + &SqlValue::Number("0.01234567890123456789012345".into(), false), + unary_op = UnaryOperator::Minus + ) + .unwrap(); + assert_eq!(value.to_string(), "-0.012345678901234567890123450000"); + } + #[test] fn test_auto_string_to_numeric() { // Test with auto_string_to_numeric=true diff --git a/src/partition/src/expr.rs b/src/partition/src/expr.rs index 121533acc60..560dc517060 100644 --- a/src/partition/src/expr.rs +++ b/src/partition/src/expr.rs @@ -278,76 +278,43 @@ impl PartitionExpr { } pub fn try_as_logical_expr(&self) -> error::Result { - // Special handling for null equality. - // `col = NULL` -> `col IS NULL` to match SQL (DataFusion) semantics. - let lhs_is_null = matches!(self.lhs.as_ref(), Operand::Value(Value::Null)); - let rhs_is_null = matches!(self.rhs.as_ref(), Operand::Value(Value::Null)); - - match (self.op.clone(), lhs_is_null, rhs_is_null) { - (RestrictedOp::Eq, _, true) => { - return Ok(self.lhs.try_as_logical_expr()?.is_null()); + // Row routing compares `Value`s, which order NULL before every other value. + // Batch routing and region filters treat a NULL result as false, so only the + // comparisons that hold for NULL need an explicit `IS NULL`. + let comparison = match (self.lhs.as_ref(), self.rhs.as_ref()) { + (column @ Operand::Column(_), Operand::Value(value)) => { + Some((column, value, self.op.clone())) } - (RestrictedOp::Eq, true, _) => { - return Ok(self.rhs.try_as_logical_expr()?.is_null()); + (Operand::Value(value), column @ Operand::Column(_)) => { + Some((column, value, self.op.invert_for_swap())) } - (RestrictedOp::NotEq, _, true) => { - return Ok(self.lhs.try_as_logical_expr()?.is_not_null()); - } - (RestrictedOp::NotEq, true, _) => { - return Ok(self.rhs.try_as_logical_expr()?.is_not_null()); - } - _ => {} - } - - if matches!( - self.op, - RestrictedOp::Lt | RestrictedOp::LtEq | RestrictedOp::Gt | RestrictedOp::GtEq - ) { - // Keep filtering semantics aligned with direct PartitionExpr evaluation (null-first ordering). - // In DataFusion SQL semantics, range comparisons with NULL yield NULL, so we inject - // `OR col IS NULL` on the null-first side of the comparison. - if matches!(self.lhs.as_ref(), Operand::Column(_)) { - let column_expr = self.lhs.try_as_logical_expr()?; - let other_expr = self.rhs.try_as_logical_expr()?; - let base = match self.op { - RestrictedOp::Lt => { - column_expr.clone().lt(other_expr).or(column_expr.is_null()) + _ => None, + }; + if let Some((column, value, op)) = comparison { + let column = column.try_as_logical_expr()?; + if matches!(value, Value::Null) { + return match op { + RestrictedOp::Eq | RestrictedOp::LtEq => Ok(column.is_null()), + RestrictedOp::NotEq | RestrictedOp::Gt => Ok(column.is_not_null()), + RestrictedOp::Lt => Ok(datafusion_expr::lit(false)), + RestrictedOp::GtEq => Ok(datafusion_expr::lit(true)), + RestrictedOp::And | RestrictedOp::Or => { + error::InvalidExprSnafu { expr: self.clone() }.fail() } - RestrictedOp::LtEq => column_expr - .clone() - .lt_eq(other_expr) - .or(column_expr.is_null()), - RestrictedOp::Gt => column_expr - .clone() - .gt(other_expr) - .and(column_expr.is_not_null()), - RestrictedOp::GtEq => column_expr - .clone() - .gt_eq(other_expr) - .and(column_expr.is_not_null()), - _ => unreachable!(), }; - return Ok(base); - } else if matches!(self.rhs.as_ref(), Operand::Column(_)) { - let other_expr = self.lhs.try_as_logical_expr()?; - let column_expr = self.rhs.try_as_logical_expr()?; - let base = match self.op { - RestrictedOp::Lt => other_expr - .lt(column_expr.clone()) - .and(column_expr.is_not_null()), - RestrictedOp::LtEq => other_expr - .lt_eq(column_expr.clone()) - .and(column_expr.is_not_null()), - RestrictedOp::Gt => { - other_expr.gt(column_expr.clone()).or(column_expr.is_null()) - } - RestrictedOp::GtEq => other_expr - .gt_eq(column_expr.clone()) - .or(column_expr.is_null()), - _ => unreachable!(), - }; - return Ok(base); } + let bound = Operand::Value(value.clone()).try_as_logical_expr()?; + return match op { + RestrictedOp::Eq => Ok(column.eq(bound)), + RestrictedOp::NotEq => Ok(column.clone().not_eq(bound).or(column.is_null())), + RestrictedOp::Lt => Ok(column.clone().lt(bound).or(column.is_null())), + RestrictedOp::LtEq => Ok(column.clone().lt_eq(bound).or(column.is_null())), + RestrictedOp::Gt => Ok(column.gt(bound)), + RestrictedOp::GtEq => Ok(column.gt_eq(bound)), + RestrictedOp::And | RestrictedOp::Or => { + error::InvalidExprSnafu { expr: self.clone() }.fail() + } + }; } // Normal cases handling, without NULL @@ -589,7 +556,7 @@ mod tests { .try_as_logical_expr() .unwrap() .to_string(), - "a > Int64(10) AND a IS NOT NULL" + "a > Int64(10)" ); // Test Gt with column on LHS @@ -600,7 +567,7 @@ mod tests { ); assert_eq!( gt_expr.try_as_logical_expr().unwrap().to_string(), - "a > Int64(10) AND a IS NOT NULL" + "a > Int64(10)" ); // Test Gt with column on RHS @@ -625,7 +592,7 @@ mod tests { ); assert_eq!( gteq_expr.try_as_logical_expr().unwrap().to_string(), - "a >= Int64(10) AND a IS NOT NULL" + "a >= Int64(10)" ); // Test LtEq with column on LHS @@ -662,7 +629,7 @@ mod tests { .try_as_logical_expr() .unwrap() .to_string(), - "a >= Int64(10) AND a IS NOT NULL" + "a >= Int64(10)" ); let and_expr = PartitionExpr::new( @@ -680,7 +647,7 @@ mod tests { ); assert_eq!( and_expr.try_as_logical_expr().unwrap().to_string(), - "(a <= Int64(10) OR a IS NULL) AND b > Int64(5) AND b IS NOT NULL" + "(a <= Int64(10) OR a IS NULL) AND b > Int64(5)" ); let and_expr = PartitionExpr::new( @@ -698,7 +665,7 @@ mod tests { ); assert_eq!( and_expr.try_as_logical_expr().unwrap().to_string(), - "(a <= Int64(10) OR a IS NULL) AND a > Int64(5) AND a IS NOT NULL" + "(a <= Int64(10) OR a IS NULL) AND a > Int64(5)" ); let and_expr_strict_lower = PartitionExpr::new( @@ -719,7 +686,7 @@ mod tests { .try_as_logical_expr() .unwrap() .to_string(), - "(a < Int64(10) OR a IS NULL) AND a >= Int64(5) AND a IS NOT NULL" + "(a < Int64(10) OR a IS NULL) AND a >= Int64(5)" ); let and_expr_rhs_column = PartitionExpr::new( @@ -740,7 +707,7 @@ mod tests { .try_as_logical_expr() .unwrap() .to_string(), - "(a <= Int64(10) OR a IS NULL) AND a > Int64(5) AND a IS NOT NULL" + "(a <= Int64(10) OR a IS NULL) AND a > Int64(5)" ); let or_expr_same_column = PartitionExpr::new( @@ -761,7 +728,7 @@ mod tests { .try_as_logical_expr() .unwrap() .to_string(), - "a <= Int64(10) OR a IS NULL OR a > Int64(5) AND a IS NOT NULL" + "a <= Int64(10) OR a IS NULL OR a > Int64(5)" ); } @@ -777,7 +744,7 @@ mod tests { .try_as_logical_expr() .unwrap() .to_string(), - "Int64(10) > a OR a IS NULL" + "a < Int64(10) OR a IS NULL" ); let gteq_expr_rhs_column = PartitionExpr { @@ -790,7 +757,7 @@ mod tests { .try_as_logical_expr() .unwrap() .to_string(), - "Int64(10) >= a OR a IS NULL" + "a <= Int64(10) OR a IS NULL" ); } diff --git a/src/partition/src/multi_dim.rs b/src/partition/src/multi_dim.rs index df7b5291e90..29e45126343 100644 --- a/src/partition/src/multi_dim.rs +++ b/src/partition/src/multi_dim.rs @@ -21,8 +21,9 @@ use std::sync::{Arc, RwLock}; use datafusion_expr::ColumnarValue; use datafusion_physical_expr::PhysicalExpr; use datatypes::arrow; -use datatypes::arrow::array::{BooleanArray, BooleanBufferBuilder, RecordBatch}; +use datatypes::arrow::array::{Array, BooleanArray, BooleanBufferBuilder, RecordBatch}; use datatypes::arrow::buffer::BooleanBuffer; +use datatypes::arrow::compute::prep_null_mask_filter; use datatypes::arrow::datatypes::Schema; use datatypes::prelude::Value; use datatypes::vectors::{Helper, VectorRef}; @@ -272,6 +273,9 @@ impl MultiDimPartitionRule { } }; let array = match columnar_value_to_boolean_array(col_val, num_rows) { + // A NULL result means the row is not in this region. Masks are merged + // with non-Kleene `or` below, where a NULL would drop the row. + Ok(array) if array.null_count() > 0 => prep_null_mask_filter(&array), Ok(array) => array, Err(e) => { return Some(Err(e)); @@ -977,6 +981,105 @@ mod test_split_record_batch { assert_eq!(total_rows, batch.num_rows()); } + #[test] + fn test_comparison_batch_matches_row() { + use datatypes::arrow::array::Array; + use datatypes::arrow::datatypes::{DataType, Field}; + + let inputs = [None, Some(4), Some(5), Some(6)]; + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, true)])), + vec![Arc::new(Int64Array::from(inputs.to_vec()))], + ) + .unwrap(); + let mut mismatches = vec![]; + for op in [ + RestrictedOp::Eq, + RestrictedOp::NotEq, + RestrictedOp::Lt, + RestrictedOp::LtEq, + RestrictedOp::Gt, + RestrictedOp::GtEq, + ] { + for bound in [Value::Null, Value::Int64(5)] { + // Both operand orders, since a rule restored without canonicalization may + // keep the literal on the left. + for expr in [ + PartitionExpr::new(col("a"), op.clone(), bound.clone().into()), + PartitionExpr { + lhs: Box::new(bound.clone().into()), + op: op.clone(), + rhs: Box::new(col("a")), + }, + ] { + let rule = MultiDimPartitionRule::try_new( + vec!["a".to_string()], + vec![1], + vec![expr.clone()], + false, + ) + .unwrap(); + let by_batch = match expr + .try_as_physical_expr(batch.schema_ref()) + .unwrap() + .evaluate(&batch) + { + Ok(by_batch) => { + columnar_value_to_boolean_array(by_batch, batch.num_rows()).unwrap() + } + Err(err) => { + mismatches.push(format!("{expr}: batch error {err}")); + continue; + } + }; + for (row, input) in inputs.iter().enumerate() { + let value = input.map_or(Value::Null, Value::Int64); + let by_row = rule.evaluate_expr(&expr, &[value]).unwrap(); + let by_batch = by_batch.is_valid(row) && by_batch.value(row); + if by_row != by_batch { + mismatches.push(format!( + "{expr} with a = {input:?}: row {by_row}, batch {by_batch}" + )); + } + } + } + } + } + assert!(mismatches.is_empty(), "{mismatches:#?}"); + } + + #[test] + fn test_split_null_rows_match_row_routing() { + use datatypes::arrow::datatypes::{DataType, Field}; + + let rule = MultiDimPartitionRule::try_new( + vec!["host".to_string()], + vec![0, 1], + vec![ + col("host").not_eq(Value::from("a")), + col("host").eq(Value::from("a")), + ], + true, + ) + .unwrap(); + for hosts in [vec![Some("a"), None, Some("z")], vec![Some("a"), None]] { + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("host", DataType::Utf8, true)])), + vec![Arc::new(StringArray::from(hosts.clone()))], + ) + .unwrap(); + let by_row = rule.split_record_batch_naive(&batch).unwrap(); + let by_batch = rule.split_record_batch(&batch).unwrap(); + for (region, expected) in by_row { + let actual = by_batch + .get(®ion) + .map(|mask| mask.array().clone()) + .unwrap_or_else(|| BooleanArray::from(vec![false; batch.num_rows()])); + assert_eq!(actual, expected, "{hosts:?}, region {region}"); + } + } + } + #[test] fn test_columnar_value_to_boolean_array_scalar_false() { let result = columnar_value_to_boolean_array( diff --git a/tests/cases/standalone/common/partition.result b/tests/cases/standalone/common/partition.result index fe6f99df591..132f2a325f6 100644 --- a/tests/cases/standalone/common/partition.result +++ b/tests/cases/standalone/common/partition.result @@ -463,3 +463,33 @@ DROP TABLE `molestiAe`; Affected Rows: 0 +-- The minimum of a signed type is a valid bound. +CREATE TABLE int_min_bound ( + a INT PRIMARY KEY, + ts TIMESTAMP TIME INDEX, +) PARTITION ON COLUMNS (a) ( + a < -2147483648, + a >= -2147483648 AND a < 0, + a >= 0 +); + +Affected Rows: 0 + +INSERT INTO int_min_bound VALUES (-2147483648, 1), (NULL, 2), (0, 3); + +Affected Rows: 3 + +SELECT * FROM int_min_bound ORDER BY ts; + ++-------------+-------------------------+ +| a | ts | ++-------------+-------------------------+ +| -2147483648 | 1970-01-01T00:00:00.001 | +| | 1970-01-01T00:00:00.002 | +| 0 | 1970-01-01T00:00:00.003 | ++-------------+-------------------------+ + +DROP TABLE int_min_bound; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/partition.sql b/tests/cases/standalone/common/partition.sql index 7b5e8ea7750..41627bf4d09 100644 --- a/tests/cases/standalone/common/partition.sql +++ b/tests/cases/standalone/common/partition.sql @@ -204,3 +204,19 @@ INSERT INTO `molestiAe` VALUES SELECT * FROM `molestiAe`; DROP TABLE `molestiAe`; + +-- The minimum of a signed type is a valid bound. +CREATE TABLE int_min_bound ( + a INT PRIMARY KEY, + ts TIMESTAMP TIME INDEX, +) PARTITION ON COLUMNS (a) ( + a < -2147483648, + a >= -2147483648 AND a < 0, + a >= 0 +); + +INSERT INTO int_min_bound VALUES (-2147483648, 1), (NULL, 2), (0, 3); + +SELECT * FROM int_min_bound ORDER BY ts; + +DROP TABLE int_min_bound; diff --git a/tests/compatibility/cases/analyze_verbose_remote_metrics_extension/verify.result b/tests/compatibility/cases/analyze_verbose_remote_metrics_extension/verify.result index 7394ca8faee..2bdcb907ebf 100644 --- a/tests/compatibility/cases/analyze_verbose_remote_metrics_extension/verify.result +++ b/tests/compatibility/cases/analyze_verbose_remote_metrics_extension/verify.result @@ -28,7 +28,7 @@ EXPLAIN ANALYZE VERBOSE SELECT count(*) FROM t_analyze_verbose_remote_metrics_ex | 1_| 1_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(t_analyze_verbose_remote_REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__count_state(t_analyze_verbose_remote_REDACTED -|_|_|_SeqScan: region=REDACTED, {"partition_count":REDACTED, "projection": ["ts"], "filters": ["host >= Utf8(\"m\") AND host IS NOT NULL"], "flat_format": REDACTED, "REDACTED +|_|_|_SeqScan: region=REDACTED, {"partition_count":REDACTED, "projection": ["ts"], "filters": ["host >= Utf8(\"m\")"], "flat_format": REDACTED, "REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ From 93c3b116da1462f3174c229d93ec9e9046670d95 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Tue, 20 Oct 2026 02:30:10 +0000 Subject: [PATCH 044/239] fix(mito2): re-check a filled write as sent after a schema change (#9701) * fix(mito2): re-check a filled write as sent after a schema change The engine fills the columns a write leaves out with their defaults when it accepts the write, appending them to the rows without adding them to `name_to_index` and `has_null`. If the region schema changes before the worker handles the write, the worker checks the request again and `check_schema` looks up a filled column in `name_to_index`, which panics and stops the region worker. If the filled column was dropped, the write fails with an unknown column that the client never sent. Remove the filled columns before checking, so the request is checked as sent and filled again for the current schema. Nothing is filled yet when the engine checks the request, so that check does no extra work. Signed-off-by: jeremyhi * test(mito2): rename put_request helper to build_put_request Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- src/mito2/src/engine/alter_test.rs | 163 ++++++++++++++++++++++++++++- src/mito2/src/request.rs | 9 ++ 2 files changed, 169 insertions(+), 3 deletions(-) diff --git a/src/mito2/src/engine/alter_test.rs b/src/mito2/src/engine/alter_test.rs index cf039b52278..7fb9f86c122 100644 --- a/src/mito2/src/engine/alter_test.rs +++ b/src/mito2/src/engine/alter_test.rs @@ -34,7 +34,9 @@ use datafusion_expr::{col, lit}; use datatypes::arrow::array::{ArrayRef, Float64Array, StringArray, TimestampMillisecondArray}; use datatypes::arrow::datatypes::{DataType, Field, Schema, TimeUnit}; use datatypes::prelude::ConcreteDataType; -use datatypes::schema::{ColumnSchema, FulltextAnalyzer, FulltextBackend, FulltextOptions}; +use datatypes::schema::{ + ColumnDefaultConstraint, ColumnSchema, FulltextAnalyzer, FulltextBackend, FulltextOptions, +}; use store_api::codec::PrimaryKeyEncoding; use store_api::logstore::provider::Provider; use store_api::metadata::{ColumnMetadata, RegionMetadata}; @@ -42,8 +44,8 @@ use store_api::metric_engine_consts::{PRIMARY_KEY_ENCODING, TABLE_COLUMN_METADAT use store_api::region_engine::{RegionEngine, RegionManifestInfo, RegionRole}; use store_api::region_request::{ AddColumn, AddColumnLocation, AlterKind, ModifyColumnType, PathType, RegionAlterRequest, - RegionBulkInsertsRequest, RegionCompactRequest, RegionOpenRequest, RegionPutRequest, - RegionRequest, SetIndexOption, SetRegionOption, UnsetRegionOption, + RegionBulkInsertsRequest, RegionCompactRequest, RegionCreateRequest, RegionOpenRequest, + RegionPutRequest, RegionRequest, SetIndexOption, SetRegionOption, UnsetRegionOption, }; use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME; use store_api::storage::{ColumnId, RegionId, ScanRequest}; @@ -52,6 +54,7 @@ use crate::config::MitoConfig; use crate::engine::MitoEngine; use crate::engine::listener::{AlterFlushListener, NotifyRegionChangeResultListener}; use crate::error; +use crate::request::WorkerRequest; use crate::sst::FormatType; use crate::test_util::batch_util::sort_batches_and_print; use crate::test_util::sst_util::{new_sparse_primary_key, sst_region_metadata_with_encoding}; @@ -3714,3 +3717,157 @@ async fn test_alter_time_index_widen_strict_window_compaction() { // Each instant exactly once, rescaled to microseconds. assert_eq!(vec![0, 1_000_000, 2_000_000, 3_000_000, 4_000_000], values); } + +/// Returns a create request with a field `field_1` that defaults to 42, and the +/// schema of rows that leave `field_1` out, so the engine fills it on puts. +fn create_request_with_default_field() -> (RegionCreateRequest, Vec) { + let mut request = CreateRequestBuilder::new().build(); + let column_schemas = rows_schema(&request); + request.column_metadatas.push(ColumnMetadata { + column_schema: ColumnSchema::new("field_1", ConcreteDataType::float64_datatype(), true) + .with_default_constraint(Some(ColumnDefaultConstraint::Value( + datatypes::value::Value::Float64(42.0.into()), + ))) + .unwrap(), + semantic_type: SemanticType::Field, + column_id: 3, + }); + (request, column_schemas) +} + +fn add_nullable_field2() -> RegionAlterRequest { + RegionAlterRequest { + kind: AlterKind::AddColumns { + columns: vec![AddColumn { + column_metadata: ColumnMetadata { + column_schema: ColumnSchema::new( + "field_2", + ConcreteDataType::float64_datatype(), + true, + ), + semantic_type: SemanticType::Field, + column_id: 4, + }, + location: None, + }], + }, + } +} + +fn build_put_request( + schema: Vec, + start: usize, + end: usize, +) -> RegionRequest { + RegionRequest::Put(RegionPutRequest { + rows: Rows { + schema, + rows: build_rows(start, end), + }, + hint: None, + skip_wal: false, + partition_expr_version: None, + }) +} + +/// Sends a put that the engine boundary checked and filled with `metadata`, +/// the region metadata before an alter, and returns the worker's answer. +async fn put_filled_with_stale_metadata( + engine: &MitoEngine, + region_id: RegionId, + metadata: store_api::metadata::RegionMetadataRef, + schema: Vec, +) -> error::Result { + let (request, receiver) = WorkerRequest::try_from_region_request( + region_id, + build_put_request(schema, 0, 2), + Some(metadata), + ) + .unwrap(); + engine + .inner + .workers + .submit_to_worker(region_id, request) + .await + .unwrap(); + receiver + .await + .expect("the worker dropped the request without an answer") +} + +/// A put filled with a default value at the engine boundary is checked again by +/// the worker after a schema change; the filled column must not break the check. +#[tokio::test] +async fn test_put_filled_at_boundary_rechecked_after_add_column() { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::new().await; + let engine = env.create_engine(MitoConfig::default()).await; + let region_id = RegionId::new(1, 1); + let (request, column_schemas) = create_request_with_default_field(); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + let metadata = engine.get_region(region_id).unwrap().metadata(); + engine + .handle_request(region_id, RegionRequest::Alter(add_nullable_field2())) + .await + .unwrap(); + + let result = + put_filled_with_stale_metadata(&engine, region_id, metadata, column_schemas.clone()).await; + assert_eq!(2, result.unwrap()); + + // The worker still serves the region. + let response = engine + .handle_request(region_id, build_put_request(column_schemas, 2, 3)) + .await + .unwrap(); + assert_eq!(1, response.affected_rows); + let scanner = engine + .scanner(region_id, ScanRequest::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(scanner.scan().await.unwrap()) + .await + .unwrap(); + let expected = "\ ++-------+---------+---------------------+---------+---------+ +| tag_0 | field_0 | ts | field_1 | field_2 | ++-------+---------+---------------------+---------+---------+ +| 0 | 0.0 | 1970-01-01T00:00:00 | 42.0 | | +| 1 | 1.0 | 1970-01-01T00:00:01 | 42.0 | | +| 2 | 2.0 | 1970-01-01T00:00:02 | 42.0 | | ++-------+---------+---------------------+---------+---------+"; + assert_eq!(expected, batches.pretty_print().unwrap()); +} + +/// A put that left out a column which is dropped before the worker handles the +/// put is still valid. +#[tokio::test] +async fn test_put_filled_at_boundary_rechecked_after_drop_column() { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::new().await; + let engine = env.create_engine(MitoConfig::default()).await; + let region_id = RegionId::new(1, 1); + let (request, column_schemas) = create_request_with_default_field(); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + let metadata = engine.get_region(region_id).unwrap().metadata(); + engine + .handle_request( + region_id, + RegionRequest::Alter(RegionAlterRequest { + kind: AlterKind::DropColumns { + names: vec!["field_1".to_string()], + }, + }), + ) + .await + .unwrap(); + + let result = put_filled_with_stale_metadata(&engine, region_id, metadata, column_schemas).await; + assert_eq!(2, result.unwrap()); +} diff --git a/src/mito2/src/request.rs b/src/mito2/src/request.rs index d084fc82c7e..5b9aff4f82d 100644 --- a/src/mito2/src/request.rs +++ b/src/mito2/src/request.rs @@ -241,6 +241,15 @@ impl WriteRequest { /// Checks the schema and fill missing columns. pub(crate) fn maybe_fill_missing_columns(&mut self, metadata: &RegionMetadata) -> Result<()> { + // Removes the columns filled for an older schema, so the request is checked + // as sent and filled again for this schema. + let num_columns = self.name_to_index.len(); + if self.rows.schema.len() > num_columns { + self.rows.schema.truncate(num_columns); + for row in &mut self.rows.rows { + row.values.truncate(num_columns); + } + } if let Err(e) = self.check_schema(metadata) { if e.is_fill_default() { // TODO(yingwen): Add metrics for this case. From a436213bf32991c386d3d759f4634e1c4d0f3531 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Tue, 20 Oct 2026 02:31:42 +0000 Subject: [PATCH 045/239] ci: compare several images in one object store WAL S3 benchmark run (#9691) * ci: compare several images in one object store WAL S3 benchmark run A new `images` input takes a JSON list of 2..4 {label, image}. Every cell then runs once per image, back to back, and the image order reverses with the configuration order every other round. The cell id gains the label. The single-image inputs keep their meaning and cell ids. The validate job pulls every image before an instance is allocated. It pins each image to its digest and records the version and commit that the image reports. Each cell starts the image of its label. An optional `warmup` runs the first load cell once per image as round 0, and no statistic uses it. The plan step estimates the benchmark job's minutes on c6id.2xlarge from the cell times of earlier runs. It refuses a plan that does not fit the job timeout. The summary keeps the Raft Engine ratios within each image. With several images it also reports, per group, concurrency and config: - every round's rows/s and client p50/p99; - each image's ratio to the first, with the geometric mean, minimum and maximum over the rounds. Every cell records the host load averages at its start and end. Signed-off-by: jeremyhi * ci: scale the S3 benchmark time estimate with the batch size The per-cell rates were measured with 1000-row batches. A smaller batch keeps about the request rate, not the row rate, so the estimate now shrinks rows/s with the batch. A one-row durable plan is refused instead of passing a timeout it cannot fit. Signed-off-by: jeremyhi * ci: pin the S3 benchmark runtime image without pulling it in validate An anonymous pull from ECR Public on a hosted runner can exceed the registry's data limit, and it failed the validate job of several dispatches with "toomanyrequests: Data limit exceeded". Validate now reads only the registry manifest to pin the runtime image's digest. A reference that names a digest must be served under that digest. The check that the image holds the loader and the generator moves to the benchmark host. It runs right after the runtime pull and before any GreptimeDB image is pulled or any cell starts. Signed-off-by: jeremyhi * test: run the S3 benchmark runtime content check against both files The fake docker now runs the pull step's own check in a directory that stands in for the runtime image. It covers both files present, the loader missing and the generator missing, so dropping either file check fails the test. Signed-off-by: jeremyhi * ci: pin a digest-named S3 benchmark runtime image without the registry ECR Public's anonymous data limit now refuses hosted runners even the manifest read, and validate failed in "Pin the runtime image digest". A runtime reference that names its sha256 digest, like the default one, is already pinned, so validate takes that digest with no network call. The format checks stay. Only a tag-only reference is resolved through the registry. The content check on the benchmark host is unchanged. Signed-off-by: jeremyhi * ci: check the full S3 benchmark runtime reference before pinning it offline A named digest no longer reaches the registry, so nothing parsed the rest of the reference. Validate now checks the whole reference against Docker's reference grammar first. It refuses an empty repository, a doubled '@' or a bad tag before any pin is emitted or any instance is allocated. Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- .github/scripts/oswal-s3-bench-run.sh | 81 ++++-- .github/scripts/oswal-s3-bench-summary.py | 105 ++++++-- .github/scripts/oswal-s3-bench.py | 160 ++++++++--- .github/workflows/oswal-s3-bench.yml | 34 ++- tests/perf/test_oswal_s3_bench_scripts.py | 308 +++++++++++++++++++++- 5 files changed, 600 insertions(+), 88 deletions(-) diff --git a/.github/scripts/oswal-s3-bench-run.sh b/.github/scripts/oswal-s3-bench-run.sh index f822a4a9b6f..d158ab14d60 100644 --- a/.github/scripts/oswal-s3-bench-run.sh +++ b/.github/scripts/oswal-s3-bench-run.sh @@ -10,24 +10,64 @@ runtime() { --security-opt no-new-privileges --cap-drop ALL "$@" } +# Pin the runtime image to a digest without pulling it: ECR Public's anonymous data limit can refuse a +# hosted runner even a manifest read. A reference that names its digest is pinned as given; only a tag +# is resolved from the registry's manifest. The benchmark host pulls the image and checks the contents. check_runtime() { - local image=$1 resolved + local image=$1 repository digest attempt component='[a-z0-9]+(([._]|__|-+)[a-z0-9]+)*' [[ "$image" == public.ecr.aws/* ]] || { echo 'The runtime image must come from AWS ECR Public' >&2; return 1; } - bash "$here/pull-benchmark-image.sh" "$image" - resolved=$(docker image inspect "$image" | jq -er '.[0].RepoDigests[0]') - docker run --rm --entrypoint sh "$resolved" -c \ - 'test -f scripts/agent-observability/greptimedb.sh && test -f scripts/generate_agent_observations.py' \ - || { echo 'The runtime image lacks the agent observability loader or generator' >&2; return 1; } - printf 'runtime_image=%s\n' "$resolved" >> "${GITHUB_OUTPUT:-/dev/stdout}" + # Docker's reference grammar, checked here because a named digest never reaches the registry. + local grammar="^public\.ecr\.aws(/$component)+(:[A-Za-z0-9_][A-Za-z0-9_.-]{0,127})?(@sha256:[0-9a-f]{64})?\$" + [[ "$image" =~ $grammar ]] || { echo "$image is not a full ECR Public image reference" >&2; return 1; } + if [[ "$image" == *@* ]]; then + digest=${image##*@} + else + for attempt in 1 2 3; do + digest=$(docker buildx imagetools inspect --format '{{json .Manifest}}' "$image" | jq -er .digest) && break + ((attempt < 3)) || { echo "Cannot read the manifest of $image" >&2; return 1; } + sleep $((attempt * 10)) + done + fi + [[ "$digest" =~ ^sha256:[0-9a-f]{64}$ ]] || { echo "$image has no sha256 digest" >&2; return 1; } + repository=${image%@*} + [[ "${repository##*/}" != *:* ]] || repository=${repository%:*} + printf 'runtime_image=%s@%s\n' "$repository" "$digest" >> "${GITHUB_OUTPUT:-/dev/stdout}" +} + +# Pull every GreptimeDB image of the run, pin it to its digest and record the version and +# commit it reports, before any instance is allocated. +check_images() { + local entry image resolved info version commit checked=() + while IFS= read -r entry; do + image=$(jq -r .image <<< "$entry") + bash "$here/pull-benchmark-image.sh" "$image" + resolved=$(docker image inspect "$image" | jq -er '[.[0].RepoDigests[]? | select(test("@sha256:[0-9a-f]{64}$"))][0]') \ + || { echo "$image has no sha256 repository digest" >&2; return 1; } + info=$(docker run --rm --entrypoint greptime "$resolved" --version) + version=$(sed -n 's/^version: //p' <<< "$info") commit=$(sed -n 's/^commit: //p' <<< "$info") + [[ -n "$version" && -n "$commit" ]] || { echo "$image reports no version or commit" >&2; return 1; } + echo "${image}: $resolved, version $version, commit $commit" + checked+=("$(jq -c --arg image "$resolved" --arg version "$version" --arg commit "$commit" \ + '{label, image: $image, requested: .image, version: $version, commit: $commit}' <<< "$entry")") + done < <(jq -c '.[]' <<< "$1") + ((${#checked[@]})) || { echo 'No GreptimeDB image to check' >&2; return 1; } + printf 'greptimedb_images=%s\n' "$(jq -sc . <<< "${checked[*]}")" >> "${GITHUB_OUTPUT:-/dev/stdout}" } pull_images() { - : "${RESULTS:?}" "${RUNTIME_IMAGE:?}" "${GREPTIMEDB_REF:?}" + : "${RESULTS:?}" "${RUNTIME_IMAGE:?}" "${GREPTIMEDB_IMAGES:?}" + local image images=() mkdir -p "$RESULTS/images" bash "$here/pull-benchmark-image.sh" "$RUNTIME_IMAGE" - bash "$here/pull-benchmark-image.sh" "$GREPTIMEDB_REF" + docker run --rm --entrypoint sh "$RUNTIME_IMAGE" -c \ + 'test -f scripts/agent-observability/greptimedb.sh && test -f scripts/generate_agent_observations.py' \ + || { echo 'The runtime image lacks the agent observability loader or generator' >&2; return 1; } + while IFS= read -r image; do + bash "$here/pull-benchmark-image.sh" "$image" + images+=("$image") + done < <(jq -r '.[].image' <<< "$GREPTIMEDB_IMAGES") docker image inspect "$RUNTIME_IMAGE" > "$RESULTS/images/runtime.json" - docker image inspect "$GREPTIMEDB_REF" > "$RESULTS/images/greptimedb.json" + docker image inspect "${images[@]}" > "$RESULTS/images/greptimedb.json" } wal_disk() { @@ -92,6 +132,8 @@ elapsed() { awk -v a="$1" -v b="$(now)" 'BEGIN {printf "%.3f", b - a}'; } # CPU seconds the server's cgroup used, and busy and total jiffies of the host. server_cpu() { docker exec "$cid" cat /sys/fs/cgroup/cpu.stat 2>/dev/null | awk '/^usage_usec / {printf "%.3f", $2 / 1e6}' || true; } host_cpu() { awk '/^cpu / {print $2 + $3 + $4 + $7 + $8 + $9, $2 + $3 + $4 + $5 + $6 + $7 + $8 + $9}' /proc/stat 2>/dev/null || true; } +# The host's 1, 5 and 15 minute load averages as a JSON list. +host_load() { awk '{printf "[%s, %s, %s]", $1, $2, $3}' "${OSWAL_LOADAVG:-/proc/loadavg}" 2>/dev/null || true; } # A server's log is named after the phase it started in, whatever phase stops it. start_server() { @@ -294,6 +336,7 @@ finish() { [[ -z "$wal_dir" ]] || sudo -n rm -rf "$wal_dir" [[ -z "$home_dir" ]] || sudo -n rm -rf "$home_dir" jq -n --argjson spec "$spec" --arg device "$wal_device" --arg version "$version" \ + --arg host_load_start "$host_load_start" --arg host_load_end "$(host_load)" \ --arg load_seconds "$load_seconds" --argjson exit_code "$rc" --argjson scraped "$metrics_scraped" \ --argjson restart_scraped "$restart_scraped" --arg restart_seconds "$restart_seconds" \ --arg restart_query_seconds "$restart_query_seconds" --arg restart_count "$restart_count" \ @@ -306,21 +349,26 @@ finish() { restart_scraped: $restart_scraped, restart_seconds: ($restart_seconds | tonumber? // null), restart_query_seconds: ($restart_query_seconds | tonumber? // null), restart_count: ($restart_count | tonumber? // null), - indexed_before_kill: ($indexed_before_kill | tonumber? // null)}' > "$dir/cell.json" + indexed_before_kill: ($indexed_before_kill | tonumber? // null), + host_load_start: ($host_load_start | fromjson? // null), host_load_end: ($host_load_end | fromjson? // null)}' \ + > "$dir/cell.json" exit "$rc" } # One cell in its own process, so errexit holds inside it. The cell's state is # global because the EXIT trap runs after this function returns. cell() { - : "${PLAN:?}" "${RESULTS:?}" "${CORPUS_DIR:?}" "${OWNER:?}" "${RUNTIME_IMAGE:?}" "${GREPTIMEDB_IMAGE:?}" + : "${PLAN:?}" "${RESULTS:?}" "${CORPUS_DIR:?}" "${OWNER:?}" "${RUNTIME_IMAGE:?}" "${GREPTIMEDB_IMAGES:?}" : "${RUN_PREFIX:?}" "${WAL_ROOT:?}" "${S3_BUCKET:?}" "${S3_REGION:?}" "${AWS_ACCESS_KEY_ID:?}" "${AWS_SECRET_ACCESS_KEY:?}" spec=$1 cid= wal_dir= home_dir= version= load_seconds= metrics_scraped=false id= phase= server_phase= restart_scraped=false restart_seconds= restart_query_seconds= restart_count= server_cpu_seconds= host_cpu_busy= - indexed_before_kill= - local concurrency setting started cpu_before host_before host_after corpus partitions load script - id=$(jq -r .id <<< "$spec"); concurrency=$(jq -r .concurrency <<< "$spec") + indexed_before_kill= host_load_start=$(host_load) + local concurrency setting started cpu_before host_before host_after corpus partitions load script label + id=$(jq -r .id <<< "$spec"); concurrency=$(jq -r .concurrency <<< "$spec") label=$(jq -r .image <<< "$spec") dir="$RESULTS/cells/$id" name="$OWNER-$id" + # The digest the image of the cell's label was pinned to before the run. + GREPTIMEDB_IMAGE=$(jq -er --arg label "$label" '.[] | select(.label == $label) | .image' <<< "$GREPTIMEDB_IMAGES") \ + || { echo "$id: no image labelled '$label'" >&2; exit 1; } mkdir -p "$dir" # The root holds data/ and the WAL's default prefix wal/ side by side. local settings=(STORAGE__TYPE=S3 "STORAGE__BUCKET=$S3_BUCKET" "STORAGE__REGION=$S3_REGION" @@ -419,6 +467,7 @@ scrub() { case "${1:-}" in check-runtime) check_runtime "${2:?}" ;; + check-images) check_images "${2:?}" ;; pull) pull_images ;; wal-disk) wal_disk "${2:?}" ;; generate) generate ;; @@ -434,5 +483,5 @@ case "${1:-}" in done < <(jq -c '.cells[]' "$PLAN") exit "$failed" ;; scrub) scrub ;; - *) echo "usage: $0 check-runtime IMAGE | pull | wal-disk DIR | generate | cells | scrub" >&2; exit 2 ;; + *) echo "usage: $0 check-runtime IMAGE | check-images JSON | pull | wal-disk DIR | generate | cells | scrub" >&2; exit 2 ;; esac diff --git a/.github/scripts/oswal-s3-bench-summary.py b/.github/scripts/oswal-s3-bench-summary.py index 305a647a9a4..e693b5f63db 100644 --- a/.github/scripts/oswal-s3-bench-summary.py +++ b/.github/scripts/oswal-s3-bench-summary.py @@ -192,6 +192,7 @@ def load_cell(path: Path, spec: dict | None = None) -> dict: wall, cpu = cell.get("load_wall_seconds"), cell.get("server_cpu_seconds") row = { "id": cell["id"], "group": cell.get("group", ""), "mode": cell.get("mode", "load"), + "image": cell.get("image", ""), "warmup": cell.get("warmup", False), "round": cell["round"], "concurrency": cell["concurrency"], "config": cell["config"], "exit_code": cell["exit_code"], "wal_device": cell.get("wal_device"), "missing": cell.get("missing", False), "env": cell.get("env", {}), @@ -200,6 +201,7 @@ def load_cell(path: Path, spec: dict | None = None) -> dict: "rows_per_second": result.get("rows_per_second") if cell["exit_code"] == 0 else None, "server_cpu_cores": cpu / wall if cpu is not None and wall else None, "host_cpu_busy": cell.get("host_cpu_busy"), "host_cpus": cell.get("host_cpus"), + "host_load_start": cell.get("host_load_start"), "host_load_end": cell.get("host_load_end"), } for name, key in (("append_ack_seconds", "append_ack"), ("seal_to_durable_seconds", "seal_to_durable"), ("object_bytes", "object_bytes")): @@ -264,13 +266,14 @@ def load_cell(path: Path, spec: dict | None = None) -> dict: def ratios(rows: list) -> list: - """Each object store WAL cell over each Raft Engine control per round, and their geometric mean over the rounds, per concurrency.""" - rows = [r for r in rows if r["mode"] == "load"] - speed = {(r["group"], r["round"], r["concurrency"], r["config"]): r["rows_per_second"] + """Each object store WAL cell over each Raft Engine control of the same image per round, and their + geometric mean over the rounds, per concurrency; warm-up cells are left out.""" + rows = [r for r in rows if r["mode"] == "load" and not r["warmup"]] + speed = {(r["group"], r["image"], r["round"], r["concurrency"], r["config"]): r["rows_per_second"] for r in rows} result = [] - for group in dict.fromkeys(r["group"] for r in rows): - members = [r for r in rows if r["group"] == group] + for group, image in dict.fromkeys((r["group"], r["image"]) for r in rows): + members = [r for r in rows if (r["group"], r["image"]) == (group, image)] series = sorted({r["config"] for r in members if r["config"] not in BASELINES}) baselines = [b for b in BASELINES if any(r["config"] == b for r in members)] for concurrency in sorted({r["concurrency"] for r in members}): @@ -278,16 +281,42 @@ def ratios(rows: list) -> list: for config, baseline in ((s, b) for s in series for b in baselines): per_round = [] for round_number in rounds: - base = speed.get((group, round_number, concurrency, baseline)) - measured = speed.get((group, round_number, concurrency, config)) + base = speed.get((group, image, round_number, concurrency, baseline)) + measured = speed.get((group, image, round_number, concurrency, config)) per_round.append(measured / base if base and measured else None) complete = [ratio for ratio in per_round if ratio is not None] - result.append({"group": group, "concurrency": concurrency, "ratio": f"{config}/{baseline}", - "per_round": per_round, + result.append({"group": group, "image": image, "concurrency": concurrency, + "ratio": f"{config}/{baseline}", "per_round": per_round, "geomean": geomean(complete) if len(complete) == len(per_round) else None}) return result +def image_ratios(rows: list, images: list) -> list: + """Every load cell of every image per round, and each image's rows/s over the first image's in the + same round, with their geometric mean, minimum and maximum over the rounds; warm-up cells are left out.""" + rows = [r for r in rows if r["mode"] == "load" and not r["warmup"]] + cells = {(r["group"], r["image"], r["round"], r["concurrency"], r["config"]): r for r in rows} + result = [] + for group, concurrency, config in dict.fromkeys((r["group"], r["concurrency"], r["config"]) for r in rows): + rounds = sorted({r["round"] for r in rows if (r["group"], r["concurrency"], r["config"]) == (group, concurrency, config)}) + for image in images: + measured = [cells.get((group, image, n, concurrency, config)) or {} for n in rounds] + reference = [cells.get((group, images[0], n, concurrency, config)) or {} for n in rounds] + per_round = [ratio(m.get("rows_per_second"), r.get("rows_per_second")) for m, r in zip(measured, reference)] + complete = image != images[0] and None not in per_round + result.append({ + "group": group, "concurrency": concurrency, "config": config, "image": image, + "rounds": [{"round": n, "id": m.get("id"), "rows_per_second": m.get("rows_per_second"), + "client_p50": m.get("client_p50"), "client_p99": m.get("client_p99"), + "host_load_start": m.get("host_load_start"), "host_load_end": m.get("host_load_end")} + for n, m in zip(rounds, measured)], + "ratio_to": None if image == images[0] else images[0], + "per_round": None if image == images[0] else per_round, + "geomean": geomean(per_round) if complete else None, + "min": min(per_round) if complete else None, "max": max(per_round) if complete else None}) + return result + + def operations(counts: dict) -> str: return ", ".join(f"{op} {n:,.0f}" for op, n in sorted(counts.items())) or "—" @@ -304,18 +333,31 @@ def fmt(value, digits=0, scale=1.0) -> str: return "—" if value is None else f"{value * scale:,.{digits}f}" -def render(run: dict, rows: list, pairs: list) -> str: +def load_average(value) -> str: + """The 1-minute load average of a recorded [1, 5, 15] minute triple.""" + return "—" if not value else f"{value[0]:.2f}" + + +def render(run: dict, rows: list, pairs: list, compared: list = ()) -> str: lines = [f"## Object store WAL against S3 ({'smoke, MinIO' if run.get('smoke') else run.get('instance_type')})", ""] - lines += [f"GreptimeDB `{run.get('greptimedb_image')}` ({next((r['version'] for r in rows if r['version']), '—')}); " - f"{fmt(run.get('rows'))} rows in {run.get('batch_rows')}-row batches; bucket region {run.get('region')}.", ""] + images = run.get("greptimedb_images") or [] + if len(images) > 1: + lines += [f"{fmt(run.get('rows'))} rows in {run.get('batch_rows')}-row batches; bucket region {run.get('region')}; " + f"each image is compared with the first, `{images[0]['label']}`.", "", + "| Image | Reference | Version | Commit |", "| --- | --- | --- | --- |"] + lines += [f"| {i['label']} | `{i['image']}` | {i['version']} | `{i['commit']}` |" for i in images] + [""] + else: + lines += [f"GreptimeDB `{run.get('greptimedb_image')}` ({next((r['version'] for r in rows if r['version']), '—')}); " + f"{fmt(run.get('rows'))} rows in {run.get('batch_rows')}-row batches; bucket region {run.get('region')}.", ""] lines += ["| Cell | Overrides | Rows/s | Load s | Wall s | Ack p50 / p99 ms | Seal to durable p50 / p99 ms " "| Object p50 / p99 MiB | Creates | Creates/min | Stalled | Poisoned | Create failures | Conflicts " - "| Server CPUs | Host busy | WAL on |", + "| Server CPUs | Host busy | WAL on | Load average at start / end |", "| --- | --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: " - "| ---: | ---: | --- |"] + "| ---: | ---: | --- | ---: |"] def label(r): - return f"{r['id']} (not run)" if r["missing"] else r["id"] if r["exit_code"] == 0 \ - else f"{r['id']} (failed, exit {r['exit_code']})" + cell = f"{r['id']} (warm-up)" if r["warmup"] else r["id"] + return f"{cell} (not run)" if r["missing"] else cell if r["exit_code"] == 0 \ + else f"{cell} (failed, exit {r['exit_code']})" loads = [r for r in rows if r["mode"] == "load"] for r in loads: cell = label(r) @@ -327,7 +369,8 @@ def render(run: dict, rows: list, pairs: list) -> str: *(fmt(r[c]) for c in COUNTERS[:1]), fmt(r["creates_per_minute"], 1), *(fmt(r[c]) for c in COUNTERS[1:]), fmt(r["server_cpu_cores"], 2), "—" if r["host_cpu_busy"] is None else f"{r['host_cpu_busy']:.0%} of {r['host_cpus'] or '?'}", - f"`{r['wal_device'] or '—'}` |"])) + f"`{r['wal_device'] or '—'}`", + f"{load_average(r['host_load_start'])} / {load_average(r['host_load_end'])} |"])) lines += ["", "| Cell | WAL MB/s | Objects/s | Mean object MiB | Sealed on size | OpenDAL S3 operations during the load " "| OpenDAL S3 operations of the restart | Idle s | Idle creates | Idle OpenDAL S3 operations |", "| --- | ---: | ---: | ---: | --- | --- | --- | ---: | ---: | --- |"] @@ -431,11 +474,22 @@ def render(run: dict, rows: list, pairs: list) -> str: fmt(x["wall_seconds"], 2), fmt(x["query_seconds"], 2), fmt(x["recovery_seconds"], 3), fmt(x["recovered_objects"]), f"{fmt(x['rows'])} |"])) - lines += ["", "| Group | Concurrency | Ratio | Per round | Geometric mean |", - "| --- | ---: | --- | --- | ---: |"] + lines += ["", "| Group | Image | Concurrency | Ratio | Per round | Geometric mean |", + "| --- | --- | ---: | --- | --- | ---: |"] for p in pairs: - lines.append(f"| {p['group'] or '—'} | {p['concurrency']} | {p['ratio']} " + lines.append(f"| {p['group'] or '—'} | {p['image'] or '—'} | {p['concurrency']} | {p['ratio']} " f"| {', '.join(fmt(v, 3) for v in p['per_round'])} | {fmt(p['geomean'], 3)} |") + if compared: + lines += ["", f"| Group | Concurrency | Config | Image | Rows/s per round | Client p50 / p99 ms per round " + f"| Ratio to {compared[0]['image']} per round | Geometric mean | Min | Max |", + "| --- | ---: | --- | --- | --- | --- | --- | ---: | ---: | ---: |"] + for c in compared: + lines.append(" | ".join([ + f"| {c['group'] or '—'}", str(c["concurrency"]), c["config"], c["image"], + "; ".join(fmt(r["rows_per_second"]) for r in c["rounds"]), + "; ".join(pair(r["client_p50"], r["client_p99"], 1, 1000) for r in c["rounds"]), + "—" if c["per_round"] is None else ", ".join(fmt(v, 3) for v in c["per_round"]), + fmt(c["geomean"], 3), fmt(c["min"], 3), fmt(c["max"], 3)]) + " |") return "\n".join(lines) + "\n" @@ -447,17 +501,20 @@ def main() -> None: run_path = args.root / "run.json" run = json.loads(run_path.read_text()) if run_path.exists() else {} plan_path = args.root / "plan.json" - planned = json.loads(plan_path.read_text())["cells"] if plan_path.exists() else [] + plan = json.loads(plan_path.read_text()) if plan_path.exists() else {} + planned = plan.get("cells", []) # The plan, not what happened to run, defines the expected cells and rounds. rows = [load_cell(args.root / "cells" / spec["id"], spec) for spec in planned] rows += [load_cell(path) for path in sorted((args.root / "cells").glob("*")) if path.name not in {spec["id"] for spec in planned} and (path / "cell.json").exists()] pairs = ratios(rows) + images = plan.get("images", [""]) + compared = image_ratios(rows, images) if len(images) > 1 else [] missing = [r["id"] for r in rows if r["missing"]] (args.root / "summary.json").write_text(json.dumps( - {"run": run, "planned_cells": len(planned), "missing_cells": missing, "cells": rows, "ratios": pairs}, - indent=2) + "\n") - print(render(run, rows, pairs) if rows else "No benchmark cell completed.\n") + {"run": run, "planned_cells": len(planned), "missing_cells": missing, "cells": rows, "ratios": pairs, + "image_ratios": compared}, indent=2) + "\n") + print(render(run, rows, pairs, compared) if rows else "No benchmark cell completed.\n") if __name__ == "__main__": diff --git a/.github/scripts/oswal-s3-bench.py b/.github/scripts/oswal-s3-bench.py index 9d674400fe2..eedd906ffef 100644 --- a/.github/scripts/oswal-s3-bench.py +++ b/.github/scripts/oswal-s3-bench.py @@ -24,6 +24,7 @@ from __future__ import annotations import argparse import json +import math import os import re import sys @@ -80,6 +81,7 @@ WITH( 'compaction.twcs.time_window' = '2h' )""" GROUP_NAME = re.compile(r"[a-z][a-z0-9]{0,11}") +IMAGE_LABEL = re.compile(r"[a-z][a-z0-9]{0,7}") RESERVED_ENV = ("GREPTIMEDB_STANDALONE__STORAGE__", "GREPTIMEDB_STANDALONE__WAL__") ENV_NAME = re.compile(r"[A-Za-z_][A-Za-z0-9_]*") # Docker's reference grammar: path components separated by ".", "_", "__" or runs of "-". @@ -88,6 +90,18 @@ IMAGE = re.compile(r"[a-z0-9]+(?:(?:[._]|__|-+)[a-z0-9]+)*(?::[0-9]+)?(?:/[a-z0- NIGHTLY = re.compile(r"v\d+\.\d+\.\d+(?:-[0-9A-Za-z.]+)?-nightly-(\d{8})") TAG = re.compile(r"[A-Za-z0-9_][A-Za-z0-9_.-]{0,127}") DIGEST = re.compile(r".+@sha256:[0-9a-f]{64}") +# Rows/s of a load cell in 1000-row batches on c6id.2xlarge at the concurrency levels earlier runs +# measured, interpolated on a log scale between them; durable writes slow down with several regions. +# Smaller batches are taken to keep the request rate, so their rows/s shrink with the batch. +MEASURED_BATCH_ROWS = 1000 +LOAD_RATES = {"durable": [(1, 9_000), (8, 34_000), (32, 55_000)], + "durable, partitioned": [(1, 7_000), (8, 26_000), (32, 55_000)], + "other": [(1, 28_000), (8, 60_000)]} +# Starting, scraping and stopping a server, and the restart of a restart cell or a crash cycle. +CELL_SECONDS, RESTART_SECONDS, CYCLE_SECONDS = 5, 60, 15 +# Checkout, dependencies, image pulls and the bucket probe; corpus generation per million rows; and the +# deletion of the run's objects one at a time, which took 12-17 % of the cells' time in earlier runs. +JOB_SECONDS, GENERATE_SECONDS_PER_MILLION, CLEANUP_SHARE = 180, 30, 0.2 def append_github(variable: str, name: str, value: str) -> None: @@ -115,16 +129,18 @@ def check_regions(bucket: str, runner: str) -> None: "a cross-region bucket measures the WAN, not S3") -def build_manifest(plan: dict, greptimedb_inspect: list, runtime_image: str, +def build_manifest(plan: dict, images: list, runtime_image: str, instance_type: str, region: str, prefix: str) -> dict: - image = next((d for d in greptimedb_inspect[0].get("RepoDigests") or [] if DIGEST.fullmatch(d)), None) - if image is None: - raise ValueError("the GreptimeDB image has no sha256 repository digest") + """The run's record; `images` are the checked images, each pinned to its digest, the first the reference.""" + if [image["label"] for image in images] != plan["images"]: + raise ValueError(f"the checked images {[i['label'] for i in images]} are not the plan's {plan['images']}") + if not all(DIGEST.fullmatch(image["image"]) for image in images): + raise ValueError("every GreptimeDB image must be pinned to a sha256 digest") return {"smoke": plan["smoke"], "rows": plan["rows"], "rounds": plan["rounds"], "concurrency": plan["concurrency"], "batch_rows": plan["batch_rows"], "configs": plan["configs"], "extra_env": plan["extra_env"], - "restart": plan["restart"], "groups": plan["groups"], - "greptimedb_image": image, "runtime_image": runtime_image, + "restart": plan["restart"], "groups": plan["groups"], "warmup": plan["warmup"], + "greptimedb_image": images[0]["image"], "greptimedb_images": images, "runtime_image": runtime_image, "instance_type": instance_type, "region": region, "prefix": prefix} @@ -166,8 +182,32 @@ def partitioned_table(summary: dict, partitions: int) -> str: return AGENT_TABLE.format(partitions=clause) -def group_cells(group: dict, smoke: bool) -> list: - """The cells of one group, in run order; a group without a name is the top-level inputs.""" +def parse_images(text: str) -> list: + """The images to compare: a JSON list of 2..4 {label, image}, the first the reference. Empty + means the one image of greptimedb_tag or greptimedb_image, with an empty label.""" + if not text.strip(): + return [{"label": "", "image": ""}] + images = json.loads(text) + if not isinstance(images, list) or not 2 <= len(images) <= 4 \ + or not all(isinstance(i, dict) and set(i) == {"label", "image"} for i in images): + raise ValueError("images must be a JSON list of 2..4 objects, each with only a label and an image") + labels = [image["label"] for image in images] + for image in images: + if not isinstance(image["label"], str) or not IMAGE_LABEL.fullmatch(image["label"]): + raise ValueError(f"an image label is up to 8 lowercase letters and digits, got {image['label']!r}") + full_reference("an image", image["image"]) + if len(set(labels)) != len(labels): + raise ValueError(f"image labels repeat: {labels}") + return images + + +def cell_id(group: str, round_number: int, level: int, config: str, label: str) -> str: + return (f"{group}-" if group else "") + f"{round_number}-{level}-{config}" + (f"-{label}" if label else "") + + +def group_cells(group: dict, smoke: bool, labels: list) -> list: + """The cells of one group, in run order; a group without a name is the top-level inputs. + Each cell runs once per image label, back to back.""" name, mode = group.get("name", ""), group.get("mode", "load") if mode not in ("load", "crash", "drop"): raise ValueError(f"group {name!r}: mode must be load, crash or drop, got {mode!r}") @@ -219,14 +259,16 @@ def group_cells(group: dict, smoke: bool) -> list: series = [config for config in WAL_CONFIGS if config in chosen] cells = [] for round_number in range(1, rounds + 1): - # Reverse the order every other round so no configuration always runs first. + # Reverse the order every other round so no configuration or image always runs first. + forward = round_number % 2 for level in levels: - for config in series if round_number % 2 else reversed(series): + for config in series if forward else reversed(series): wal = WAL_CONFIGS[config] | ({"FLUSH_INTERVAL": RESTART_FLUSH_INTERVAL} if restart and config not in BASELINES else {}) - cells.append(common | { - "id": (f"{name}-" if name else "") + f"{round_number}-{level}-{config}", - "round": round_number, "concurrency": level, "config": config, "wal": wal, "env": env, - }) + for label in labels if forward else reversed(labels): + cells.append(common | { + "id": cell_id(name, round_number, level, config, label), "round": round_number, + "concurrency": level, "config": config, "image": label, "wal": wal, "env": env, + }) return cells @@ -251,11 +293,20 @@ def parse_groups(text: str, defaults: dict) -> list: def build_plan(concurrency: str, rows: str, batch_rows: str, rounds: str, smoke: bool, - configs: str = "R,D,E", extra_env: str = "", restart: bool = False, groups: str = "") -> dict: + configs: str = "R,D,E", extra_env: str = "", restart: bool = False, groups: str = "", + images: str = "", warmup: bool = False) -> dict: defaults = {"concurrency": concurrency, "rows": rows, "rounds": rounds, "configs": configs, "extra_env": extra_env, "restart": restart} resolved = parse_groups(groups, defaults) - cells = [cell for group in resolved for cell in group_cells(group, smoke)] + labels = [image["label"] for image in parse_images(images)] + cells = [cell for group in resolved for cell in group_cells(group, smoke, labels)] + if warmup: + # Round 0 runs the first load cell once per image before anything else; no statistic uses it. + first = next((cell for cell in cells if cell["mode"] == "load"), None) + if first is None: + raise ValueError("warmup needs a load group") + cells = [first | {"id": cell_id(first["group"], 0, first["concurrency"], first["config"], label), + "round": 0, "image": label, "warmup": True} for label in labels] + cells if len({cell["id"] for cell in cells}) != len(cells): raise ValueError("two cells share an id") corpora = [] @@ -270,10 +321,36 @@ def build_plan(concurrency: str, rows: str, batch_rows: str, rounds: str, smoke: "concurrency": sorted({cell["concurrency"] for cell in cells}), "configs": list(dict.fromkeys(cell["config"] for cell in cells)), "extra_env": parse_extra_env(extra_env), "groups": [g.get("name", "") for g in resolved], - "restart": restart, "corpora": corpora, "cells": cells, + "restart": restart, "images": labels, "warmup": warmup, "corpora": corpora, "cells": cells, } +def load_rate(cell: dict) -> float: + durable = cell["wal"].get("ACK_MODE") == "durable" + points = LOAD_RATES["other" if not durable else "durable, partitioned" if cell["partitions"] > 1 else "durable"] + level = cell["concurrency"] + for (low, low_rate), (high, high_rate) in zip(points, points[1:]): + if level < high: + return low_rate + (high_rate - low_rate) * math.log(max(level, low) / low) / math.log(high / low) + return points[-1][1] + + +def cell_seconds(cell: dict, batch_rows: int) -> float: + if cell["mode"] == "crash": + return CELL_SECONDS + cell["cycles"] * (sum(cell["kill_after"]) / 2 + CYCLE_SECONDS) + if cell["mode"] == "drop": + return CELL_SECONDS + cell["write_seconds"] + RESTART_SECONDS / 2 + rate = load_rate(cell) * min(1.0, batch_rows / MEASURED_BATCH_ROWS) + return CELL_SECONDS + cell["rows"] / rate + cell["idle_seconds"] + (RESTART_SECONDS if cell["restart"] else 0) + + +def estimate_minutes(plan: dict) -> int: + """The benchmark job's minutes on c6id.2xlarge, from the times of earlier runs.""" + generate = GENERATE_SECONDS_PER_MILLION * sum(corpus["rows"] for corpus in plan["corpora"]) / 1e6 + cells = sum(cell_seconds(cell, plan["batch_rows"]) for cell in plan["cells"]) + return math.ceil((JOB_SECONDS + generate + cells * (1 + CLEANUP_SHARE)) / 60) + + def newest_nightly(tags: list[str]) -> str: dated = [(match.group(1), tag) for tag in tags if (match := NIGHTLY.fullmatch(tag))] if not dated: @@ -281,15 +358,30 @@ def newest_nightly(tags: list[str]) -> str: return max(dated)[1] +def full_reference(name: str, image: str) -> str: + if not isinstance(image, str) or not IMAGE.fullmatch(image) \ + or not re.search(r"(:[^/]+|@sha256:[0-9a-f]{64})$", image): + raise ValueError(f"{name} must be a full reference with a tag or digest: {image!r}") + return image + + def resolve_image(image: str, tag: str) -> str: """The full image reference when given, else greptime/greptimedb at the resolved tag.""" if image: - if not IMAGE.fullmatch(image) or not re.search(r"(:[^/]+|@sha256:[0-9a-f]{64})$", image): - raise ValueError(f"greptimedb_image must be a full reference with a tag or digest: {image!r}") - return image + return full_reference("greptimedb_image", image) return f"greptime/greptimedb:{resolve_tag(tag)}" +def resolve_images(images: str, image: str, tag: str) -> list: + """The images to pull and check: the list to compare, or the one image of the other two inputs.""" + entries = parse_images(images) + if entries[0]["label"]: + if image: + raise ValueError("set either images or greptimedb_image, not both") + return entries + return [{"label": "", "image": resolve_image(image, tag)}] + + def resolve_tag(requested: str) -> str: if requested != "nightly": if not TAG.fullmatch(requested): @@ -347,13 +439,15 @@ def main() -> int: plan.add_argument("--smoke", choices=("true", "false"), required=True) plan.add_argument("--restart", choices=("true", "false"), default="false") plan.add_argument("--configs", default="R,D,E") - for name in ("extra-env", "groups"): + plan.add_argument("--warmup", choices=("true", "false"), default="false") + for name in ("extra-env", "groups", "images"): plan.add_argument(f"--{name}", default="") for name in ("timeout-minutes", "ttl-hours", "output"): plan.add_argument(f"--{name}") image = commands.add_parser("image") image.add_argument("--tag", required=True) image.add_argument("--image", default="") + image.add_argument("--images", default="") regions = commands.add_parser("regions") for name in ("bucket", "runner"): regions.add_argument(f"--{name}", required=True) @@ -361,7 +455,7 @@ def main() -> int: table.add_argument("--summary", type=Path, required=True) table.add_argument("--partitions", type=int, required=True) manifest = commands.add_parser("manifest") - for name in ("results", "runtime-image", "instance-type", "region", "prefix"): + for name in ("results", "images", "runtime-image", "instance-type", "region", "prefix"): manifest.add_argument(f"--{name}", required=True) for name in ("probe", "cleanup"): bucket = commands.add_parser(name) @@ -373,15 +467,21 @@ def main() -> int: if args.timeout_minutes or args.ttl_hours: validate_budgets(args.timeout_minutes or "", args.ttl_hours or "") result = build_plan(args.concurrency, args.rows, args.batch_rows, args.rounds, args.smoke == "true", - args.configs, args.extra_env, args.restart == "true", args.groups) + args.configs, args.extra_env, args.restart == "true", args.groups, + args.images, args.warmup == "true") + minutes = estimate_minutes(result) + if args.timeout_minutes and minutes > int(args.timeout_minutes): + raise ValueError(f"the plan needs about {minutes} minutes on c6id.2xlarge, " + f"more than the {args.timeout_minutes}-minute benchmark timeout") if args.output: with open(args.output, "w") as output: json.dump(result, output, indent=2) - print(f"{len(result['cells'])} cells; corpora {[c['rows'] for c in result['corpora']]}") + print(f"{len(result['cells'])} cells; corpora {[c['rows'] for c in result['corpora']]}; " + f"about {minutes} minutes on c6id.2xlarge") elif args.command == "image": - reference = resolve_image(args.image, args.tag) - append_github("GITHUB_OUTPUT", "greptimedb_ref", reference) - print(f"GreptimeDB image: {reference}") + images = resolve_images(args.images, args.image, args.tag) + append_github("GITHUB_OUTPUT", "requested_images", json.dumps(images)) + print("GreptimeDB images: " + ", ".join(f"{i['label']} {i['image']}".strip() for i in images)) elif args.command == "regions": check_regions(args.bucket, args.runner) print(f"bucket and runner both in {args.runner}") @@ -389,12 +489,10 @@ def main() -> int: print(partitioned_table(json.loads(args.summary.read_text()), args.partitions)) elif args.command == "manifest": results = Path(args.results) - manifest = build_manifest(json.loads((results / "plan.json").read_text()), - json.loads((results / "images" / "greptimedb.json").read_text()), + manifest = build_manifest(json.loads((results / "plan.json").read_text()), json.loads(args.images), args.runtime_image, args.instance_type, args.region, args.prefix) (results / "run.json").write_text(json.dumps(manifest, indent=2) + "\n") - append_github("GITHUB_ENV", "GREPTIMEDB_IMAGE", manifest["greptimedb_image"]) - print(f"GreptimeDB image: {manifest['greptimedb_image']}") + print("GreptimeDB images: " + ", ".join(i["image"] for i in manifest["greptimedb_images"])) else: if not re.fullmatch(r"oswal-bench/[0-9]+-[0-9]+", args.prefix): raise ValueError(f"refusing a prefix outside oswal-bench/: {args.prefix!r}") diff --git a/.github/workflows/oswal-s3-bench.yml b/.github/workflows/oswal-s3-bench.yml index f5e92d2cad5..e61754cb1f1 100644 --- a/.github/workflows/oswal-s3-bench.yml +++ b/.github/workflows/oswal-s3-bench.yml @@ -14,6 +14,14 @@ on: description: 'Full image reference with a tag or digest; used instead of greptimedb_tag when set' type: string default: '' + images: + description: 'JSON list of 2..4 {"label", "image"} to compare on one instance instead of greptimedb_tag and greptimedb_image: every cell runs once per image, in an order that reverses every other round; a label is up to 8 lowercase letters and digits, an image a full reference with a tag or digest, the first the reference of the ratios' + type: string + default: '' + warmup: + description: Before round 1, run the first load cell once per image and leave it out of every statistic + type: boolean + default: false configs: description: 'WAL configurations: R (Raft Engine, sync_write), R0 (Raft Engine, standalone default), D (durable), E (enqueued)' type: string @@ -47,7 +55,7 @@ on: type: number default: 1000 rounds: - description: Rounds; the configuration order alternates between rounds (R, D, E then E, D, R) + description: Rounds; the configuration and image order alternates between rounds (R, D, E then E, D, R) type: number default: 2 runtime_image: @@ -68,6 +76,8 @@ env: SMOKE: ${{ inputs.smoke }} GREPTIMEDB_TAG: ${{ inputs.greptimedb_tag }} GREPTIMEDB_IMAGE_INPUT: ${{ inputs.greptimedb_image }} + OSWAL_IMAGES: ${{ inputs.images }} + WARMUP: ${{ inputs.warmup }} CONFIGS: ${{ inputs.configs }} EXTRA_ENV: ${{ inputs.extra_env }} RESTART: ${{ inputs.restart }} @@ -86,7 +96,7 @@ jobs: timeout-minutes: 15 outputs: instance_type: ${{ steps.runner.outputs.instance_type }} - greptimedb_ref: ${{ steps.plan.outputs.greptimedb_ref }} + greptimedb_images: ${{ steps.images.outputs.greptimedb_images }} runtime_image: ${{ steps.runtime.outputs.runtime_image }} steps: - uses: actions/checkout@v4 @@ -103,8 +113,15 @@ jobs: python3 greptimedb/.github/scripts/oswal-s3-bench.py plan --smoke "$SMOKE" \ --concurrency "$CONCURRENCY" --rows "$ROWS" --batch-rows "$BATCH_ROWS" --rounds "$ROUNDS" \ --configs "$CONFIGS" --extra-env "$EXTRA_ENV" --restart "$RESTART" --groups "$OSWAL_GROUPS" \ + --images "$OSWAL_IMAGES" --warmup "$WARMUP" \ --timeout-minutes "$BENCHMARK_TIMEOUT_MINUTES" --ttl-hours "$JANITOR_TTL_HOURS" - python3 greptimedb/.github/scripts/oswal-s3-bench.py image --tag "$GREPTIMEDB_TAG" --image "$GREPTIMEDB_IMAGE_INPUT" + python3 greptimedb/.github/scripts/oswal-s3-bench.py image --tag "$GREPTIMEDB_TAG" --image "$GREPTIMEDB_IMAGE_INPUT" \ + --images "$OSWAL_IMAGES" + - name: Pull every GreptimeDB image and record its version + id: images + env: + REQUESTED_IMAGES: ${{ steps.plan.outputs.requested_images }} + run: bash greptimedb/.github/scripts/oswal-s3-bench-run.sh check-images "$REQUESTED_IMAGES" - name: Resolve CI runner id: runner run: python3 greptimedb/.github/scripts/ci-runner-config.py --provider AWS --instance-type "$REQUESTED_INSTANCE_TYPE" @@ -113,7 +130,7 @@ jobs: BUCKET_REGION: ${{ vars.AWS_CI_TEST_BUCKET_REGION }} RUNNER_REGION: ${{ vars.EC2_RUNNER_REGION }} run: python3 greptimedb/.github/scripts/oswal-s3-bench.py regions --bucket "$BUCKET_REGION" --runner "$RUNNER_REGION" - - name: Check the runtime image + - name: Pin the runtime image digest id: runtime run: bash greptimedb/.github/scripts/oswal-s3-bench-run.sh check-runtime "$RUNTIME_IMAGE" provision: @@ -158,7 +175,7 @@ jobs: timeout-minutes: ${{ fromJSON(format('{0}', inputs.benchmark_timeout_minutes)) }} env: INSTANCE_TYPE: ${{ inputs.smoke && 'ubuntu-latest' || needs.validate.outputs.instance_type }} - GREPTIMEDB_REF: ${{ needs.validate.outputs.greptimedb_ref }} + GREPTIMEDB_IMAGES: ${{ needs.validate.outputs.greptimedb_images }} RUNTIME_IMAGE: ${{ needs.validate.outputs.runtime_image }} RUN_PREFIX: oswal-bench/${{ github.run_id }}-${{ github.run_attempt }} OWNER: o11ybench-oswal-${{ github.run_id }}-${{ github.run_attempt }} @@ -183,13 +200,14 @@ jobs: run: | bash greptimedb/.github/scripts/pull-benchmark-image.sh greptime/minio:2024 docker compose -f greptimedb/tests-integration/fixtures/docker-compose.yml up -d --wait minio - - name: Pull images and write the plan + - name: Pull and check images, write the plan run: | bash greptimedb/.github/scripts/oswal-s3-bench-run.sh pull python3 greptimedb/.github/scripts/oswal-s3-bench.py plan --smoke "$SMOKE" --output "$PLAN" \ --concurrency "$CONCURRENCY" --rows "$ROWS" --batch-rows "$BATCH_ROWS" --rounds "$ROUNDS" \ - --configs "$CONFIGS" --extra-env "$EXTRA_ENV" --restart "$RESTART" --groups "$OSWAL_GROUPS" - python3 greptimedb/.github/scripts/oswal-s3-bench.py manifest --results "$RESULTS" \ + --configs "$CONFIGS" --extra-env "$EXTRA_ENV" --restart "$RESTART" --groups "$OSWAL_GROUPS" \ + --images "$OSWAL_IMAGES" --warmup "$WARMUP" + python3 greptimedb/.github/scripts/oswal-s3-bench.py manifest --results "$RESULTS" --images "$GREPTIMEDB_IMAGES" \ --runtime-image "$RUNTIME_IMAGE" --instance-type "$INSTANCE_TYPE" --region "$S3_REGION" --prefix "$RUN_PREFIX" - name: Probe the bucket env: diff --git a/tests/perf/test_oswal_s3_bench_scripts.py b/tests/perf/test_oswal_s3_bench_scripts.py index 6934ff0db13..d95b29f7a52 100644 --- a/tests/perf/test_oswal_s3_bench_scripts.py +++ b/tests/perf/test_oswal_s3_bench_scripts.py @@ -224,15 +224,106 @@ class PlanTest(unittest.TestCase): with self.subTest(bucket=bucket), self.assertRaisesRegex(ValueError, "cross-region"): bench.check_regions(bucket, runner) - def test_manifest_takes_the_sha256_digest(self): + def test_manifest_records_the_checked_images_pinned_to_their_digests(self): plan = bench.build_plan("1,8", "2000000", "1000", "2", smoke=False) - inspect = [{"RepoDigests": ["greptime/greptimedb@sha256:" + "a" * 64]}] - manifest = bench.build_manifest(plan, inspect, "runtime@sha256:" + "b" * 64, "c6id.2xlarge", + checked = [{"label": "", "image": "greptime/greptimedb@sha256:" + "a" * 64, + "requested": "greptime/greptimedb:v1.4.0-nightly-20261012", "version": "1.4.0", "commit": "abc"}] + manifest = bench.build_manifest(plan, checked, "runtime@sha256:" + "b" * 64, "c6id.2xlarge", "ap-southeast-1", "oswal-bench/1-1") self.assertEqual(manifest["greptimedb_image"], "greptime/greptimedb@sha256:" + "a" * 64) + self.assertEqual(manifest["greptimedb_images"], checked) self.assertEqual((manifest["rows"], manifest["rounds"], manifest["concurrency"]), (2_000_000, 2, [1, 8])) - with self.assertRaises(ValueError): - bench.build_manifest(plan, [{"RepoDigests": []}], "", "", "", "") + for images in ([checked[0] | {"image": "greptime/greptimedb:v1.4.0"}], [checked[0] | {"label": "base"}]): + with self.subTest(images=images), self.assertRaises(ValueError): + bench.build_manifest(plan, images, "", "", "", "") + + def test_every_cell_runs_once_per_image_in_an_order_that_reverses_every_other_round(self): + images = json.dumps([{"label": label, "image": f"docker.io/greptime/greptimedb-dev:{label}"} + for label in ("base", "pr1", "pr2")]) + plan = bench.build_plan("8", "2000000", "1000", "3", smoke=False, configs="D", images=images) + self.assertEqual([c["id"] for c in plan["cells"]], [ + "1-8-D-base", "1-8-D-pr1", "1-8-D-pr2", "2-8-D-pr2", "2-8-D-pr1", "2-8-D-base", + "3-8-D-base", "3-8-D-pr1", "3-8-D-pr2"]) + self.assertEqual(plan["images"], ["base", "pr1", "pr2"]) + # The configurations reverse with the images, and a warm-up runs the first load cell per image. + groups = json.dumps([{"name": "crash", "mode": "crash", "configs": "E", "rounds": 1, "cycles": 1}, + {"name": "single", "configs": "D,E"}]) + plan = bench.build_plan("8", "2000000", "1000", "2", smoke=False, groups=groups, + images=json.dumps(json.loads(images)[:2]), warmup=True) + self.assertEqual([c["id"] for c in plan["cells"]], [ + "single-0-8-D-base", "single-0-8-D-pr1", "crash-1-8-E-base", "crash-1-8-E-pr1", + "single-1-8-D-base", "single-1-8-D-pr1", "single-1-8-E-base", "single-1-8-E-pr1", + "single-2-8-E-pr1", "single-2-8-E-base", "single-2-8-D-pr1", "single-2-8-D-base"]) + warmups = [c for c in plan["cells"] if c.get("warmup")] + self.assertEqual([(c["round"], c["image"], c["mode"]) for c in warmups], [(0, "base", "load"), (0, "pr1", "load")]) + self.assertEqual((plan["rounds"], plan["warmup"]), (2, True)) + # Without images every cell keeps its id and runs the one image. + single = bench.build_plan("8", "2000000", "1000", "1", smoke=False, configs="D") + self.assertEqual([(c["id"], c["image"]) for c in single["cells"]], [("1-8-D", "")]) + self.assertEqual((single["images"], single["warmup"]), ([""], False)) + + def test_invalid_images_and_warmups_are_rejected(self): + ref = "docker.io/greptime/greptimedb-dev:dev" + for images in ([{"label": "base", "image": ref}], [{"label": f"i{n}", "image": ref} for n in range(5)], + [{"label": "base", "image": ref}, {"label": "base", "image": ref}], + [{"label": "Base", "image": ref}, {"label": "pr1", "image": ref}], + [{"label": "ninechars", "image": ref}, {"label": "pr1", "image": ref}], + [{"label": 1, "image": ref}, {"label": "pr1", "image": ref}], + [{"label": "base", "image": "greptimedb"}, {"label": "pr1", "image": ref}], + [{"label": "base", "image": ref, "env": "A=1"}, {"label": "pr1", "image": ref}], + {"base": ref}): + with self.subTest(images=images), self.assertRaises(ValueError): + bench.build_plan("8", "6000", "1000", "1", smoke=False, images=json.dumps(images)) + crash_only = json.dumps([{"name": "c", "mode": "crash", "configs": "D"}]) + with self.assertRaisesRegex(ValueError, "warmup needs a load group"): + bench.build_plan("8", "6000", "1000", "1", smoke=False, groups=crash_only, warmup=True) + # One image twice compares it with itself. + pair = [{"label": "base", "image": ref}, {"label": "again", "image": ref}] + self.assertEqual(bench.resolve_images(json.dumps(pair), "", "nightly"), pair) + with self.assertRaisesRegex(ValueError, "not both"): + bench.resolve_images(json.dumps(pair), ref, "nightly") + self.assertEqual(bench.resolve_images("", ref, "nightly"), [{"label": "", "image": ref}]) + + def test_a_plan_longer_than_the_benchmark_timeout_is_refused(self): + groups = json.dumps([{"name": "single", "configs": "D,E"}, {"name": "multi", "configs": "D,E", "partitions": 8}]) + images = json.dumps([{"label": label, "image": f"docker.io/greptime/greptimedb-dev:{label}"} + for label in ("base", "pr1", "pr2")]) + plan = bench.build_plan("8", "2000000", "1000", "5", smoke=False, groups=groups, images=images, warmup=True) + self.assertEqual(len(plan["cells"]), 63) + self.assertEqual(bench.estimate_minutes(plan), 75) + cells = {c["id"]: c for c in plan["cells"]} + # Durable at concurrency 8 writes 34,000 rows/s to one region and 26,000 to eight. + self.assertAlmostEqual(bench.cell_seconds(cells["single-1-8-D-base"], 1000), 5 + 2_000_000 / 34_000) + self.assertAlmostEqual(bench.cell_seconds(cells["multi-1-8-D-pr1"], 1000), 5 + 2_000_000 / 26_000) + self.assertAlmostEqual(bench.cell_seconds(cells["multi-1-8-E-pr2"], 4000), 5 + 2_000_000 / 60_000) + # A smaller batch keeps the request rate: one-row batches write a thousandth of the rows. + self.assertAlmostEqual(bench.cell_seconds(cells["single-1-8-D-base"], 1), 5 + 2_000_000 / 34) + durable = cells["single-1-8-D-base"] + self.assertAlmostEqual(bench.load_rate(durable | {"concurrency": 4}), 9_000 + 25_000 * 2 / 3) + self.assertEqual(bench.load_rate(durable | {"concurrency": 64}), 55_000) + crash = bench.build_plan("8", "6000", "1000", "1", smoke=False, groups=json.dumps( + [{"name": "c", "mode": "crash", "configs": "D", "cycles": 5, "kill_after": [5, 20]}]))["cells"][0] + self.assertEqual(bench.cell_seconds(crash, 1), 5 + 5 * (12.5 + 15)) + command = [sys.executable, str(SCRIPTS / "oswal-s3-bench.py"), "plan", "--smoke", "false", "--concurrency", + "8", "--rows", "2000000", "--rounds", "5", "--groups", groups, + "--images", images, "--warmup", "true", "--ttl-hours", "8"] + for batch_rows, timeout, expected in (("1000", "75", "about 75 minutes"), ("1000", "74", "more than the 74-minute"), + ("1", "360", "more than the 360-minute")): + done = subprocess.run(command + ["--batch-rows", batch_rows, "--timeout-minutes", timeout], + capture_output=True, text=True, timeout=60) + with self.subTest(batch_rows=batch_rows, timeout=timeout): + self.assertEqual(done.returncode, 0 if expected.startswith("about") else 1, done.stderr) + self.assertIn(expected, done.stdout + done.stderr) + + def test_the_workflow_checks_every_image_before_an_instance_is_allocated(self): + workflow = (SCRIPTS.parent / "workflows" / "oswal-s3-bench.yml").read_text() + self.assertIn("OSWAL_IMAGES: ${{ inputs.images }}", workflow) + self.assertEqual(workflow.count('--images "$OSWAL_IMAGES"'), 3) + self.assertEqual(workflow.count('--warmup "$WARMUP"'), 2) + validate = workflow[workflow.index(" validate:"):workflow.index(" provision:")] + self.assertIn("greptimedb_images: ${{ steps.images.outputs.greptimedb_images }}", validate) + self.assertIn("check-images \"$REQUESTED_IMAGES\"", validate) + self.assertIn("GREPTIMEDB_IMAGES: ${{ needs.validate.outputs.greptimedb_images }}", workflow) def write_lines(path, lines): @@ -425,11 +516,12 @@ class SummaryTest(unittest.TestCase): self.assertNotIn("kafka_client_produce_elapsed_count", summary.parse_metrics(ACK_BUCKETS)["samples"]) def write_cell(self, root, cell_id, rows_per_second, exit_code=0, metrics="", **extra): - round_number, concurrency, config = cell_id.split("-") + round_number, concurrency, config, *image = cell_id.split("-") path = root / "cells" / cell_id (path / "greptimedb-1").mkdir(parents=True) (path / "cell.json").write_text(json.dumps({ "id": cell_id, "round": int(round_number), "concurrency": int(concurrency), "config": config, + "image": image[0] if image else "", "wal": bench.WAL_CONFIGS[config], "metrics_scraped": True, "env": {}, "exit_code": exit_code, "wal_device": "/dev/nvme1n1", "version": "1.4.0", "load_wall_seconds": 61.0} | extra)) if exit_code == 0: @@ -664,6 +756,51 @@ class SummaryTest(unittest.TestCase): self.assertIn("1-1-D (failed, exit 1)", summary.render({}, rows, summary.ratios(rows))) + def test_images_are_compared_with_the_first_in_every_round(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + images = [{"label": label, "image": f"docker.io/greptime/greptimedb-dev:{label}"} for label in ("base", "pr1")] + plan = bench.build_plan("8", "6000", "1000", "2", smoke=False, configs="R,D", images=json.dumps(images), + warmup=True) + (root / "plan.json").write_text(json.dumps(plan)) + (root / "run.json").write_text(json.dumps({"rows": 6000, "batch_rows": 1000, "greptimedb_images": [ + i | {"image": f"greptime/greptimedb-dev@sha256:{n * 64}", "version": "1.4.0", "commit": f"{n * 7}"} + for i, n in zip(images, "ab")]})) + # The warm-up cells are far off and must not move any ratio. + speeds = {"0-8-R-base": 10.0, "0-8-R-pr1": 999.0, + "1-8-R-base": 100.0, "1-8-R-pr1": 100.0, "1-8-D-base": 50.0, "1-8-D-pr1": 40.0, + "2-8-D-pr1": 66.0, "2-8-D-base": 55.0, "2-8-R-pr1": 100.0, "2-8-R-base": 110.0} + for cell_id, speed in speeds.items(): + self.write_cell(root, cell_id, speed, warmup=cell_id.startswith("0-"), + host_load_start=[0.5, 0.4, 0.3], host_load_end=[1.25, 0.5, 0.3]) + with open(root / "cells" / "2-8-D-pr1" / "greptimedb-1" / "load-transcript.jsonl", "w") as out: + for seconds in (0.2, 0.3): + out.write(json.dumps({"step": "grpc_write_batch", "status": "success", "duration_seconds": seconds}) + "\n") + output = subprocess.run([sys.executable, str(SCRIPTS / "oswal-s3-bench-summary.py"), "--root", str(root)], + check=True, capture_output=True, text=True).stdout + result = json.loads((root / "summary.json").read_text()) + compared = {(c["config"], c["image"]): c for c in result["image_ratios"]} + self.assertEqual(list(compared), [("R", "base"), ("R", "pr1"), ("D", "base"), ("D", "pr1")]) + durable = compared[("D", "pr1")] + self.assertEqual(durable["ratio_to"], "base") + self.assertEqual(durable["per_round"], [0.8, 1.2]) + self.assertAlmostEqual(durable["geomean"], 0.96 ** 0.5) + self.assertEqual((durable["min"], durable["max"]), (0.8, 1.2)) + self.assertEqual([r["rows_per_second"] for r in durable["rounds"]], [40.0, 66.0]) + self.assertEqual((durable["rounds"][1]["client_p50"], durable["rounds"][1]["client_p99"]), (0.2, 0.3)) + self.assertEqual(durable["rounds"][0]["host_load_end"], [1.25, 0.5, 0.3]) + self.assertEqual((compared[("D", "base")]["per_round"], compared[("D", "base")]["geomean"]), (None, None)) + # The ratios to the Raft Engine control stay within each image. + pairs = {(r["image"], r["ratio"]): r["per_round"] for r in result["ratios"]} + self.assertEqual(pairs, {("base", "D/R"): [0.5, 0.5], ("pr1", "D/R"): [0.4, 0.66]}) + self.assertIn("| base | `greptime/greptimedb-dev@sha256:" + "a" * 64 + "` | 1.4.0 | `aaaaaaa` |", output) + self.assertIn("| — | 8 | D | pr1 | 40; 66 | — / —; 200.0 / 300.0 | 0.800, 1.200 | 0.980 | 0.800 | 1.200 |", output) + self.assertIn("| — | 8 | D | base | 50; 55 | — / —; — / — | — | — | — | — |", output) + self.assertIn("| 0-8-R-pr1 (warm-up) |", output) + self.assertIn("| `/dev/nvme1n1` | 0.50 / 1.25 |", output) + self.assertIn("| — | pr1 | 8 | D/R | 0.400, 0.660 | 0.514 |", output) + + FAKE_DOCKER = r"""#!/usr/bin/env bash printf '%s\n' "docker $*" >> "$FAKE_LOG" case "$1" in @@ -677,6 +814,7 @@ case "$1" in echo "{\"count_before\": 0, \"count_after\": $FAKE_ROWS, \"load_seconds\": 2, \"rows_per_second\": 3000, \"table\": \"agent_observations\"}" \ > "$results/greptimedb-1/load-result.json" ;; kill) touch "${FAKE_KILL_FLAG:-/dev/null}" ;; + rm) [[ -z "${FAKE_LOAD_AFTER:-}" ]] || echo "$FAKE_LOAD_AFTER" > "$OSWAL_LOADAVG" ;; inspect) echo true ;; exec) [[ "$*" == *cpu.stat* ]] && echo 'usage_usec 5000000' || echo 'version: 1.4.0-test' ;; logs) echo 'server log' @@ -723,14 +861,16 @@ done class CellLifecycleTest(unittest.TestCase): # the run script against fake docker and curl - def run_cells(self, root, load_rc, metrics_rc=0, restart=False, count=6000, groups=""): + def run_cells(self, root, load_rc, metrics_rc=0, restart=False, count=6000, groups="", images="", checked=None): bin_dir = root / "bin" bin_dir.mkdir() for name, body in (("docker", FAKE_DOCKER), ("curl", FAKE_CURL)): (bin_dir / name).write_text(body) (bin_dir / name).chmod(0o755) plan = bench.build_plan("8", "6000", "1000", "1", smoke=False, configs="D,E", restart=restart, - extra_env="RUST_LOG=info", groups=groups) + extra_env="RUST_LOG=info", groups=groups, images=images) + if checked is None: + checked = [{"label": label, "image": f"greptimedb{label}@sha256:" + "a" * 64} for label in plan["images"]] (root / "results").mkdir() (root / "results" / "plan.json").write_text(json.dumps(plan)) env = os.environ | { @@ -738,7 +878,7 @@ class CellLifecycleTest(unittest.TestCase): # the run script against fake docke "FAKE_LOAD_RC": str(load_rc), "FAKE_METRICS_RC": str(metrics_rc), "FAKE_ROWS": "6000", "FAKE_COUNT": str(count), "PLAN": str(root / "results" / "plan.json"), "RESULTS": str(root / "results"), "CORPUS_DIR": str(root / "corpus"), "OWNER": "o11ybench-oswal-1-1", - "RUNTIME_IMAGE": "runtime@sha256:" + "b" * 64, "GREPTIMEDB_IMAGE": "greptimedb@sha256:" + "a" * 64, + "RUNTIME_IMAGE": "runtime@sha256:" + "b" * 64, "GREPTIMEDB_IMAGES": json.dumps(checked), "RUN_PREFIX": "oswal-bench/1-1", "WAL_ROOT": str(root / "wal"), "S3_BUCKET": "bucket", "S3_REGION": "ap-southeast-1", "AWS_ACCESS_KEY_ID": "AKIAFAKEKEY", "AWS_SECRET_ACCESS_KEY": "fake-secret-value", } @@ -1070,5 +1210,155 @@ class CellLifecycleTest(unittest.TestCase): # the run script against fake docke self.assertEqual((recorded["exit_code"], recorded["restart_count"]), (5, 5999)) + def test_every_cell_starts_its_own_image_and_records_the_host_load(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + (root / "loadavg").write_text("0.50 0.40 0.30 1/200 99\n") + images = json.dumps([{"label": label, "image": f"docker.io/greptime/greptimedb-dev:{label}"} + for label in ("base", "pr1")]) + old = os.environ.copy() + os.environ.update({"OSWAL_LOADAVG": str(root / "loadavg"), "FAKE_LOAD_AFTER": "1.25 1.00 0.75 1/200 99"}) + try: + cells, _, log, _ = self.run_cells(root, load_rc=0, images=images) + finally: + os.environ.clear() + os.environ.update(old) + self.assertEqual(cells.returncode, 0, cells.stdout + cells.stderr) + starts = [line for line in log.splitlines() if line.startswith("docker run -d")] + self.assertEqual([re.search(r"--name o11ybench-oswal-1-1-(\S+) ", line).group(1) for line in starts], + ["1-8-D-base", "1-8-D-pr1", "1-8-E-base", "1-8-E-pr1"]) + for line in starts: + label = line.split("--name ")[1].split()[0].rsplit("-", 1)[1] + self.assertIn(f" greptimedb{label}@sha256:{'a' * 64} standalone start", line) + self.assertIn(f"TARGET_IMAGE=greptimedbpr1@sha256:{'a' * 64}", log) + first, second = (json.loads((root / "results" / "cells" / i / "cell.json").read_text()) + for i in ("1-8-D-base", "1-8-D-pr1")) + self.assertEqual((first["image"], first["host_load_start"], first["host_load_end"]), + ("base", [0.5, 0.4, 0.3], [1.25, 1.0, 0.75])) + self.assertEqual(second["host_load_start"], [1.25, 1.0, 0.75]) + + def test_a_cell_without_a_checked_image_fails(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + images = json.dumps([{"label": label, "image": f"docker.io/greptime/greptimedb-dev:{label}"} + for label in ("base", "pr1")]) + cells, _, log, _ = self.run_cells(root, load_rc=0, images=images, + checked=[{"label": "base", "image": "greptimedb@sha256:" + "a" * 64}]) + self.assertNotEqual(cells.returncode, 0) + self.assertIn("1-8-D-pr1: no image labelled 'pr1'", cells.stderr) + self.assertEqual(log.count("docker run -d"), 2) + + def test_check_images_pins_every_image_to_its_digest_and_records_its_version(self): + fake_docker = r"""#!/usr/bin/env bash +printf '%s\n' "docker $*" >> "$FAKE_LOG" +case "$1 $2" in + "image inspect") printf '[{"RepoDigests": ["%s@sha256:%064d"]}]\n' "${3%:*}" "$(printf '%s' "$3" | cksum | cut -d' ' -f1)" ;; + "run --rm") [[ "$*" == *nocommit* ]] || commit="commit: 0123456789abcdef$(printf '%s' "$4" | cksum | cut -c1-4)" + printf 'greptime\nbranch: main\n%s\nclean: true\nversion: 1.4.0\n' "${commit:-}" ;; +esac +""" + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + for name, body in (("docker", fake_docker), ("timeout", '#!/usr/bin/env bash\nshift 2\nexec "$@"\n')): + (root / name).write_text(body) + (root / name).chmod(0o755) + env = os.environ | {"PATH": f"{root}:{os.environ['PATH']}", "FAKE_LOG": str(root / "fake.log"), + "GITHUB_OUTPUT": str(root / "output")} + requested = [{"label": label, "image": f"docker.io/greptime/greptimedb-dev:{label}"} for label in ("base", "pr1")] + done = subprocess.run(["bash", str(SCRIPTS / "oswal-s3-bench-run.sh"), "check-images", json.dumps(requested)], + env=env, capture_output=True, text=True, timeout=60) + self.assertEqual(done.returncode, 0, done.stderr) + name, _, value = (root / "output").read_text().strip().partition("=") + checked = json.loads(value) + self.assertEqual(name, "greptimedb_images") + self.assertEqual([(c["label"], c["requested"], c["version"]) for c in checked], + [(r["label"], r["image"], "1.4.0") for r in requested]) + for c in checked: + self.assertRegex(c["image"], r"^docker\.io/greptime/greptimedb-dev@sha256:[0-9]{64}$") + self.assertTrue(c["commit"].startswith("0123456789abcdef"), c) + self.assertIn(f"docker run --rm --entrypoint greptime {c['image']} --version", (root / "fake.log").read_text()) + self.assertNotEqual(checked[0]["image"], checked[1]["image"]) + requested[1]["image"] = "docker.io/greptime/nocommit:pr1" + done = subprocess.run(["bash", str(SCRIPTS / "oswal-s3-bench-run.sh"), "check-images", json.dumps(requested)], + env=env, capture_output=True, text=True, timeout=60) + self.assertNotEqual(done.returncode, 0) + self.assertIn("reports no version or commit", done.stderr) + + + def test_check_runtime_pins_a_named_digest_offline_and_resolves_only_a_tag(self): + digest = "sha256:" + "c" * 64 + fake_docker = f"""#!/usr/bin/env bash +printf '%s\\n' "docker $*" >> "$FAKE_LOG" +[[ "$1 $2 $3" == "buildx imagetools inspect" ]] || exit 9 +printf '{{"mediaType": "application/vnd.oci.image.index.v1+json", "digest": "{digest}", "size": 856}}\\n' +""" + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + (root / "docker").write_text(fake_docker) + (root / "docker").chmod(0o755) + env = os.environ | {"PATH": f"{root}:{os.environ['PATH']}", "FAKE_LOG": str(root / "fake.log"), + "GITHUB_OUTPUT": str(root / "output")} + repository, named = "public.ecr.aws/i8k6a5e1/greptime/o11ybenchruntime", "sha256:" + "d" * 64 + # Expected pin, or error, and whether the registry is asked. + for image, expected, asks in ((f"{repository}:77a84e1c8136", f"{repository}@{digest}", True), + (f"{repository}:77a84e1c8136@{named}", f"{repository}@{named}", False), + (f"{repository}@{named}", f"{repository}@{named}", False), + (f"{repository}:77a84e1c8136@sha256:abc", "not a full ECR Public", False), + (f"public.ecr.aws/@{named}", "not a full ECR Public", False), + (f"{repository}@@{named}", "not a full ECR Public", False), + (f"{repository}:bad tag", "not a full ECR Public", False), + ("docker.io/greptime/o11ybenchruntime:1", "must come from AWS ECR Public", False)): + with self.subTest(image=image): + (root / "output").write_text("") + (root / "fake.log").write_text("") + done = subprocess.run(["bash", str(SCRIPTS / "oswal-s3-bench-run.sh"), "check-runtime", image], + env=env, capture_output=True, text=True, timeout=60) + if expected.startswith(repository): + self.assertEqual(done.returncode, 0, done.stderr) + self.assertEqual((root / "output").read_text(), f"runtime_image={expected}\n") + else: + self.assertNotEqual(done.returncode, 0) + self.assertIn(expected, done.stderr) + self.assertEqual((root / "output").read_text(), "") + calls = (root / "fake.log").read_text().splitlines() + self.assertEqual(calls, [f"docker buildx imagetools inspect --format {{{{json .Manifest}}}} {image}"] + if asks else []) + + def test_the_benchmark_host_checks_the_runtime_contents_before_pulling_greptimedb(self): + # The fake runs the step's own check in a directory standing in for the image's working directory. + fake_docker = r"""#!/usr/bin/env bash +printf '%s\n' "docker $*" >> "$FAKE_LOG" +case "$1" in + run) cd "$FAKE_RUNTIME_ROOT" && exec sh -c "$7" ;; + image) echo '[{}]' ;; +esac +""" + loader, generator = "scripts/agent-observability/greptimedb.sh", "scripts/generate_agent_observations.py" + for present, passes in (((loader, generator), True), ((generator,), False), ((loader,), False)): + with self.subTest(present=present), tempfile.TemporaryDirectory() as directory: + root = Path(directory) + for name, body in (("docker", fake_docker), ("timeout", '#!/usr/bin/env bash\nshift 2\nexec "$@"\n')): + (root / name).write_text(body) + (root / name).chmod(0o755) + for path in present: + (root / "image" / path).parent.mkdir(parents=True, exist_ok=True) + (root / "image" / path).write_text("") + runtime = "public.ecr.aws/i8k6a5e1/greptime/o11ybenchruntime@sha256:" + "c" * 64 + env = os.environ | {"PATH": f"{root}:{os.environ['PATH']}", "FAKE_LOG": str(root / "fake.log"), + "FAKE_RUNTIME_ROOT": str(root / "image"), "RESULTS": str(root / "results"), + "RUNTIME_IMAGE": runtime, + "GREPTIMEDB_IMAGES": json.dumps([{"label": "", "image": "greptimedb@sha256:" + "a" * 64}])} + done = subprocess.run(["bash", str(SCRIPTS / "oswal-s3-bench-run.sh"), "pull"], env=env, + capture_output=True, text=True, timeout=60) + log = (root / "fake.log").read_text().splitlines() + self.assertTrue(log[1].startswith(f"docker run --rm --entrypoint sh {runtime} -c "), log[1]) + if passes: + self.assertEqual(done.returncode, 0, done.stderr) + self.assertEqual([line.split()[1] for line in log], ["pull", "run", "pull", "image", "image"]) + else: + self.assertNotEqual(done.returncode, 0) + self.assertIn("lacks the agent observability loader or generator", done.stderr) + self.assertEqual([line.split()[1] for line in log], ["pull", "run"]) + if __name__ == "__main__": unittest.main() From 9c7e1e1605b831451470d83b9e2281e7965e46b3 Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Tue, 20 Oct 2026 03:00:54 +0000 Subject: [PATCH 046/239] fix: record slow INSERT SELECT queries (#9733) Signed-off-by: shuiyisong --- src/frontend/src/instance.rs | 75 ++++++++++++++++++++++++++++-------- 1 file changed, 59 insertions(+), 16 deletions(-) diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index d72d0677552..cbbf2c7babf 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -291,7 +291,7 @@ impl Instance { stmt: &Statement, schema_name: String, ) -> Option { - if !stmt.is_readonly() || !self.slow_query_options.enable { + if !should_track_statement_process(stmt) || !self.slow_query_options.enable { return None; } @@ -1032,24 +1032,19 @@ impl Instance { .map(|s| s.to_string()) .unwrap_or_else(|| plan.display_indent().to_string()); - let plan_is_readonly = is_readonly_plan(&plan); let result = if should_track_plan_process(stmt.as_ref(), &plan) { let catalog_name = query_ctx.current_catalog().to_string(); let schema_name = query_ctx.current_schema(); - let slow_query_timer = if plan_is_readonly { - self.slow_query_options.enable.then(|| { - SlowQueryTimer::new( - CatalogQueryStatement::Plan(query.clone()), - schema_name.clone(), - self.slow_query_options.threshold, - self.slow_query_options.sample_ratio, - self.slow_query_options.record_type, - self.slow_query_recorder.clone(), - ) - }) - } else { - None - }; + let slow_query_timer = self.slow_query_options.enable.then(|| { + SlowQueryTimer::new( + CatalogQueryStatement::Plan(query.clone()), + schema_name.clone(), + self.slow_query_options.threshold, + self.slow_query_options.sample_ratio, + self.slow_query_options.record_type, + self.slow_query_recorder.clone(), + ) + }); let timeout_recorder = stmt .as_ref() @@ -2060,6 +2055,7 @@ mod tests { #[derive(Debug, Default)] struct RecordingSlowQueryEventRecorder { + queries: std::sync::Mutex>, payloads: std::sync::Mutex>, } @@ -2069,6 +2065,7 @@ mod tests { .as_any() .downcast_ref::() .expect("expected a slow query event"); + self.queries.lock().unwrap().push(event.query.clone()); self.payloads.lock().unwrap().push(event.payload.clone()); } @@ -3992,6 +3989,52 @@ mod tests { assert!(!should_track_plan_process(None, &insert_dml_plan())); } + #[tokio::test] + async fn test_insert_select_slow_query() -> TestResult<()> { + let mut instance = + test_instance_with_tables(test_table(1024, "source")?, test_table(1025, "target")?) + .await?; + let recorder = Arc::new(RecordingSlowQueryEventRecorder::default()); + instance.slow_query_recorder = recorder.clone(); + let sql = "INSERT INTO target SELECT * FROM source"; + + for prepared in [false, true] { + for (enabled, threshold, should_record) in [ + (true, Duration::ZERO, true), + (false, Duration::ZERO, false), + (true, Duration::from_secs(3600), false), + ] { + instance.slow_query_options.enable = enabled; + instance.slow_query_options.threshold = threshold; + let query_ctx = test_query_ctx(42); + let output = if prepared { + let stmt = parse_one_sql(sql); + let described = instance + .do_describe_inner(stmt.clone(), query_ctx.clone()) + .await + .unwrap() + .unwrap(); + instance + .do_exec_plan_inner(described.logical_plan, Some(stmt), query_ctx) + .await + .unwrap() + } else { + execute_one_sql(&instance, sql, query_ctx).await? + }; + assert!(matches!(output.data, OutputData::AffectedRows(0))); + + let queries = std::mem::take(&mut *recorder.queries.lock().unwrap()); + let expected = if should_record { vec![sql] } else { vec![] }; + assert_eq!( + queries, expected, + "prepared={prepared}, enabled={enabled}, threshold={threshold:?}" + ); + } + } + + Ok(()) + } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn test_insert_select_is_visible_in_show_processlist() -> TestResult<()> { let insert_sql = "INSERT INTO target SELECT * FROM source"; From 4645ab9ed74c9433bbc0a99a1208243c9df7019f Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Tue, 20 Oct 2026 03:41:18 +0000 Subject: [PATCH 047/239] perf(mito2): gather last rows with one take per input batch (#9737) * perf(mito2): gather last rows with one take per input batch FlatLastTimestampSelector buffered a single-row slice per series and concatenated up to 8192 of them. Concatenating slices with dictionary columns merges the dictionaries once per slice, which made high-cardinality instant PromQL queries using the LastRow scan path (#9034) slower. Select the rows of series that end inside a batch with one take instead. Signed-off-by: Dennis Zhuang * test(mito2): cover duplicate last timestamps in the take path Signed-off-by: Dennis Zhuang * fix(mito2): group last rows by primary key bytes Flat batches may assign different dictionary keys to the same primary key, for example when binary keys are wrapped with identity keys. Merge adjacent ranges with equal key bytes before selecting last rows. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/mito2/AGENTS.md | 3 + src/mito2/src/read/last_row.rs | 167 +++++++++++++++++++++++++++------ 2 files changed, 140 insertions(+), 30 deletions(-) diff --git a/src/mito2/AGENTS.md b/src/mito2/AGENTS.md index c41aa649111..811632b1e26 100644 --- a/src/mito2/AGENTS.md +++ b/src/mito2/AGENTS.md @@ -98,6 +98,9 @@ Tests live next to the code as `*_test.rs` (e.g. `src/mito2/src/engine/flush_tes only close/drop unregisters them. See `compaction/scheduler/state.rs`. - Local compaction units use the snapshot-free `CompactionContext`. Retaining a full `CompactionRegion` in a unit pins other units' inputs and delays SST purge. +- In flat batches, equal `__primary_key` dictionary keys imply the same series, but + not the reverse: `wrap_pk_binary_to_dict` gives every row its own key. Compare key + bytes when grouping rows by series. ## Maintenance contract diff --git a/src/mito2/src/read/last_row.rs b/src/mito2/src/read/last_row.rs index 2bafe02dbbe..b7da3e1f70c 100644 --- a/src/mito2/src/read/last_row.rs +++ b/src/mito2/src/read/last_row.rs @@ -16,8 +16,8 @@ use std::sync::Arc; -use datatypes::arrow::array::{Array, BinaryArray}; -use datatypes::arrow::compute::concat_batches; +use datatypes::arrow::array::{Array, BinaryArray, UInt32Array}; +use datatypes::arrow::compute::{concat_batches, take_record_batch}; use datatypes::arrow::record_batch::RecordBatch; use futures::{Stream, TryStreamExt}; use snafu::ResultExt; @@ -155,6 +155,11 @@ impl BatchBuffer { self.batches.extend_from_slice(batches); } + fn push(&mut self, batch: RecordBatch) { + self.num_rows += batch.num_rows(); + self.batches.push(batch); + } + /// Returns true if the buffer has no batches. fn is_empty(&self) -> bool { self.batches.is_empty() @@ -306,41 +311,70 @@ impl FlatLastTimestampSelector { return Ok(()); } + // Dictionary keys may differ for the same primary key, e.g. readers that wrap binary + // keys into a dictionary with identity keys, so merge adjacent ranges by key bytes. + let mut series_ranges: Vec<(usize, usize)> = Vec::with_capacity(offsets.len() - 1); + for window in offsets.windows(2) { + let (start, end) = (window[0], window[1]); + match series_ranges.last_mut() { + Some(last) + if primary_key_bytes_at(&batch, pk_col_idx, last.0) + == primary_key_bytes_at(&batch, pk_col_idx, start) => + { + last.1 = end; + } + _ => series_ranges.push((start, end)), + } + } + let ts_values = timestamp_array_to_i64_slice(batch.column(ts_col_idx)); - for i in 0..offsets.len() - 1 { - let range_start = offsets[i]; - let range_end = offsets[i + 1]; + let num_ranges = series_ranges.len(); + // Rows of keys that start and end inside this batch. They are gathered with one + // `take` instead of buffering a slice per key: concatenating thousands of single-row + // slices merges their dictionaries once per slice. + let mut completed_rows = Vec::new(); + for (i, &(range_start, range_end)) in series_ranges.iter().enumerate() { let range_key = primary_key_bytes_at(&batch, pk_col_idx, range_start); let range_last_ts = ts_values[range_end - 1]; let range_last_ts_start = last_timestamp_start(ts_values, range_start, range_end); - let range_slice = batch.slice(range_last_ts_start, range_end - range_last_ts_start); + let is_last_range = i + 1 == num_ranges; - match self.current_key.as_mut() { - Some(state) if state.key.as_slice() == range_key => { - if range_last_ts > state.last_timestamp { - state.last_timestamp = range_last_ts; - state.slices.clear(); - state.slices.push(range_slice); - } else if range_last_ts == state.last_timestamp { - state.slices.push(range_slice); - } + // Only the first range can continue the key carried over from previous batches. + if let Some(state) = self.current_key.as_mut() + && state.key.as_slice() == range_key + { + let range_slice = batch.slice(range_last_ts_start, range_end - range_last_ts_start); + if range_last_ts > state.last_timestamp { + state.last_timestamp = range_last_ts; + state.slices.clear(); + state.slices.push(range_slice); + } else if range_last_ts == state.last_timestamp { + state.slices.push(range_slice); } - Some(_) => { + if !is_last_range { self.flush_current_key(output_buffer); - self.current_key = Some(LastKeyState::new( - range_key.to_vec(), - range_last_ts, - range_slice, - )); - } - None => { - self.current_key = Some(LastKeyState::new( - range_key.to_vec(), - range_last_ts, - range_slice, - )); } + continue; } + + // Flushes the carried-over key before rows of this batch to keep the output sorted. + self.flush_current_key(output_buffer); + if is_last_range { + // The key may continue in the next batch. + self.current_key = Some(LastKeyState::new( + range_key.to_vec(), + range_last_ts, + batch.slice(range_last_ts_start, range_end - range_last_ts_start), + )); + } else { + completed_rows.extend(range_last_ts_start as u32..range_end as u32); + } + } + + if !completed_rows.is_empty() { + let indices = UInt32Array::from(completed_rows); + let selected = take_record_batch(&batch, &indices).context(ComputeArrowSnafu)?; + output_buffer.push(selected); } Ok(()) @@ -429,8 +463,8 @@ mod tests { use std::sync::Arc; use datatypes::arrow::array::{ - ArrayRef, BinaryDictionaryBuilder, Int64Array, TimestampMillisecondArray, UInt8Array, - UInt64Array, + ArrayRef, BinaryDictionaryBuilder, DictionaryArray, Int64Array, TimestampMillisecondArray, + UInt8Array, UInt64Array, }; use datatypes::arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit, UInt32Type}; use datatypes::arrow::record_batch::RecordBatch; @@ -630,6 +664,79 @@ mod tests { assert_eq!(vec![(b"k1".to_vec(), 4), (b"k1".to_vec(), 4)], results); } + #[test] + fn test_flat_concat_across_parent_dictionaries() { + let mut selector = FlatLastTimestampSelector::default(); + let mut output_buffer = BatchBuffer::new(); + let batches = vec![ + new_flat_batch(&[b"a", b"b", b"b", b"c"], &[1, 1, 2, 3], &[10, 20, 21, 30]), + new_flat_batch( + &[b"c", b"d", b"d", b"e", b"e"], + &[3, 4, 4, 5, 5], + &[31, 40, 41, 50, 51], + ), + new_flat_batch(&[b"f", b"g"], &[6, 7], &[60, 70]), + ]; + for batch in batches { + selector.on_next(batch, &mut output_buffer).unwrap(); + } + selector.finish(&mut output_buffer).unwrap(); + let merged = output_buffer.concat().unwrap(); + + let mut rows = Vec::new(); + extract_flat_rows(&merged, &mut rows); + let fields = merged + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec(); + assert_eq!( + vec![ + (b"a".to_vec(), 1), + (b"b".to_vec(), 2), + (b"c".to_vec(), 3), + (b"c".to_vec(), 3), + (b"d".to_vec(), 4), + (b"d".to_vec(), 4), + (b"e".to_vec(), 5), + (b"e".to_vec(), 5), + (b"f".to_vec(), 6), + (b"g".to_vec(), 7), + ], + rows + ); + assert_eq!(vec![10, 21, 30, 31, 40, 41, 50, 51, 60, 70], fields); + } + + /// Rebuilds the primary key column with one dictionary key per row, as + /// `wrap_pk_binary_to_dict` does for binary primary keys. + fn with_identity_pk_keys(batch: RecordBatch) -> RecordBatch { + let pk_col_idx = primary_key_column_index(batch.num_columns()); + let values = (0..batch.num_rows()) + .map(|i| Some(primary_key_bytes_at(&batch, pk_col_idx, i))) + .collect::(); + let keys = UInt32Array::from_iter_values(0..batch.num_rows() as u32); + let mut columns = batch.columns().to_vec(); + columns[pk_col_idx] = Arc::new(DictionaryArray::new(keys, Arc::new(values))); + RecordBatch::try_new(batch.schema(), columns).unwrap() + } + + #[test] + fn test_flat_identity_dictionary_keys() { + let mut selector = FlatLastTimestampSelector::default(); + let batches = vec![ + with_identity_pk_keys(new_flat_batch(&[b"a", b"a", b"b"], &[1, 2, 3], &[0; 3])), + with_identity_pk_keys(new_flat_batch(&[b"b", b"c", b"c"], &[4, 5, 6], &[0; 3])), + ]; + let results = collect_flat_results(&mut selector, batches); + assert_eq!( + vec![(b"a".to_vec(), 2), (b"b".to_vec(), 4), (b"c".to_vec(), 6)], + results + ); + } + #[test] fn test_flat_finish_is_one_shot() { let mut selector = FlatLastTimestampSelector::default(); From 3e2751d71ded602486f09b7ad925b27492af3c8e Mon Sep 17 00:00:00 2001 From: Yingwen Date: Tue, 20 Oct 2026 05:06:32 +0000 Subject: [PATCH 048/239] feat: improve SQL analyze scan diagnostics (#9565) * feat: improve SQL analyze scan diagnostics Signed-off-by: evenyag * test: redact dynamic filters in sqlness results Signed-off-by: evenyag * test: stabilize analyze snapshots for merge queue Signed-off-by: evenyag * test: remove obsolete scan_details assertion Signed-off-by: evenyag * test: refresh explain snapshots after partition filter changes Signed-off-by: evenyag --------- Signed-off-by: evenyag --- src/mito2/src/read/scan_region.rs | 148 +++++++++++++++--- tests-integration/tests/http.rs | 128 +++++++++++++++ .../cases/distributed/explain/analyze.result | 3 +- tests/cases/distributed/explain/analyze.sql | 1 + .../cases/distributed/explain/order_by.result | 6 +- tests/cases/distributed/explain/order_by.sql | 2 + .../distributed/explain/step_aggr.result | 28 ++-- tests/cases/distributed/explain/step_aggr.sql | 4 + .../explain/step_aggr_advance.result | 114 +++++++++++--- .../distributed/explain/step_aggr_advance.sql | 28 ++++ .../explain/step_aggr_basic.result | 35 +++-- .../distributed/explain/step_aggr_basic.sql | 5 + .../explain/step_aggr_massive.result | 66 ++++---- .../distributed/explain/step_aggr_massive.sql | 2 + .../cases/distributed/optimizer/count.result | 27 ++-- .../optimizer/first_value_advance.result | 56 ++++--- .../optimizer/last_value_advance.result | 56 ++++--- .../optimizer/range_select_projection.result | 3 +- .../optimizer/range_select_projection.sql | 1 + .../optimizer/windowed_sort.result | 12 +- .../distributed/optimizer/windowed_sort.sql | 4 + .../optimizer/windowed_sort_advance.result | 6 +- .../optimizer/windowed_sort_advance.sql | 2 + .../optimizer/windowed_sort_nanos.result | 3 +- .../optimizer/windowed_sort_nanos.sql | 1 + .../tql-explain-analyze/analyze.result | 22 ++- .../tql-explain-analyze/analyze.sql | 6 + .../common/aggregate/distinct.result | 16 +- .../standalone/common/aggregate/distinct.sql | 4 + .../aggregate/max_abs_regression.result | 6 +- .../common/aggregate/max_abs_regression.sql | 2 + .../common/aggregate/multi_regions.result | 15 +- .../common/aggregate/multi_regions.sql | 3 + .../common/filter/hash_join_dyn_filter.result | 10 +- .../common/filter/hash_join_dyn_filter.sql | 2 + .../filter/hash_join_topk_dyn_filter.result | 17 +- .../filter/hash_join_topk_dyn_filter.sql | 3 + .../common/filter/topk_dyn_filter.result | 3 +- .../common/filter/topk_dyn_filter.sql | 1 + .../standalone/common/order/order_by.result | 5 +- .../standalone/common/order/order_by.sql | 1 + .../common/order/windowed_sort.result | 26 ++- .../standalone/common/order/windowed_sort.sql | 10 ++ .../cases/standalone/common/partition.result | 50 +++--- tests/cases/standalone/common/partition.sql | 6 + .../common/promql/delta_temporality.result | 3 +- .../common/promql/delta_temporality.sql | 1 + .../promql/histogram_multi_partition.result | 5 +- .../promql/histogram_multi_partition.sql | 1 + .../promql/native_time_selection.result | 9 +- .../common/promql/native_time_selection.sql | 3 + .../standalone/common/promql/regex.result | 10 +- .../cases/standalone/common/promql/regex.sql | 4 + .../promql/tsid_binary_join_regression.result | 78 +++++---- .../promql/tsid_binary_join_regression.sql | 22 +++ .../tsid_histogram_quantile_regression.result | 4 +- .../tsid_histogram_quantile_regression.sql | 2 + .../cases/standalone/common/range/nest.result | 3 +- tests/cases/standalone/common/range/nest.sql | 1 + .../standalone/common/select/prune.result | 6 +- .../cases/standalone/common/select/prune.sql | 2 + .../common/select/skipping_index.result | 6 +- .../common/select/skipping_index.sql | 2 + .../common/select/tql_filter.result | 9 +- .../standalone/common/select/tql_filter.sql | 3 + .../standalone/common/tql/partition.result | 13 +- .../cases/standalone/common/tql/partition.sql | 3 + tests/cases/standalone/optimizer/count.result | 27 ++-- tests/cases/standalone/optimizer/count.sql | 9 ++ .../optimizer/filter_column_pruning.result | 81 ++++++---- .../optimizer/filter_column_pruning.sql | 27 ++++ .../optimizer/first_value_advance.result | 56 ++++--- .../optimizer/first_value_advance.sql | 16 ++ .../standalone/optimizer/last_value.result | 6 +- .../cases/standalone/optimizer/last_value.sql | 2 + .../optimizer/last_value_advance.result | 56 ++++--- .../optimizer/last_value_advance.sql | 16 ++ .../standalone/optimizer/order_by.result | 6 +- tests/cases/standalone/optimizer/order_by.sql | 2 + .../time_index_filter_pushdown.result | 3 +- .../optimizer/time_index_filter_pushdown.sql | 1 + .../standalone/optimizer/windowed_sort.result | 15 +- .../standalone/optimizer/windowed_sort.sql | 5 + .../optimizer/windowed_sort_advance.result | 6 +- .../optimizer/windowed_sort_advance.sql | 2 + .../optimizer/windowed_sort_nanos.result | 3 +- .../optimizer/windowed_sort_nanos.sql | 1 + .../tql-explain-analyze/analyze.result | 34 ++-- .../tql-explain-analyze/analyze.sql | 10 ++ .../tql-explain-analyze/column_pruning.result | 6 +- .../tql-explain-analyze/column_pruning.sql | 2 + .../tql-explain-analyze/tsid_column.result | 20 ++- .../tql-explain-analyze/tsid_column.sql | 8 + .../verify.result | 5 +- .../verify.sql | 1 + 95 files changed, 1169 insertions(+), 396 deletions(-) diff --git a/src/mito2/src/read/scan_region.rs b/src/mito2/src/read/scan_region.rs index d2bc53167be..bcd79cc6960 100644 --- a/src/mito2/src/read/scan_region.rs +++ b/src/mito2/src/read/scan_region.rs @@ -1964,6 +1964,41 @@ pub(crate) fn exact_sequence_range( Ok((files, sequence_range)) } +/// Metadata totals for distinct SST files, not post-filter rows or actual read bytes. +/// Older files may have unknown row and row-group counts represented by zero. +#[derive(Default, serde::Serialize)] +struct FileSummary { + files: usize, + rows: u64, + row_groups: u64, + size: u64, + index_size: u64, +} + +impl FileSummary { + fn new<'a>(files: impl IntoIterator) -> Self { + let mut summary = Self::default(); + for file in files { + summary.files += 1; + summary.rows += file.meta_ref().num_rows; + summary.row_groups += file.meta_ref().num_row_groups; + summary.size += file.size(); + summary.index_size += file.index_size(); + } + summary + } +} + +/// Keeps the existing debug-list spacing while escaping strings as JSON, including +/// control characters for which Rust's Debug escaping is not valid JSON. +struct ExplainString<'a>(&'a str); + +impl fmt::Debug for ExplainString<'_> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(&serde_json::to_string(self.0).map_err(|_| fmt::Error)?) + } +} + /// Context shared by different streams from a scanner. /// It contains the input and ranges to scan. pub struct StreamContext { @@ -2064,6 +2099,7 @@ impl StreamContext { write!(f, ", \"distribution\":\"{}\"", distribution)?; } + self.format_scan_details(f)?; if verbose { self.format_verbose_content(f)?; } @@ -2071,6 +2107,94 @@ impl StreamContext { Ok(()) } + fn format_scan_details(&self, f: &mut fmt::Formatter) -> fmt::Result { + let output_schema = self.input.mapper.output_schema(); + if !output_schema.is_empty() { + write!(f, ", \"projection\": ")?; + f.debug_list() + .entries( + output_schema + .column_schemas() + .iter() + .map(|col| ExplainString(&col.name)), + ) + .finish()?; + } + if let Some(predicate) = self.input.predicate.predicate() { + if !predicate.exprs().is_empty() { + let exprs: Vec<_> = predicate.exprs().iter().map(|e| e.to_string()).collect(); + write!(f, ", \"filters\": ")?; + f.debug_list() + .entries(exprs.iter().map(|e| ExplainString(e))) + .finish()?; + } + if !predicate.dyn_filters().is_empty() { + let filters: Vec<_> = predicate + .dyn_filters() + .iter() + .map(|filter| filter.to_string()) + .collect(); + write!(f, ", \"dyn_filters\": ")?; + f.debug_list() + .entries(filters.iter().map(|filter| ExplainString(filter))) + .finish()?; + } + } + if let Some(time_range) = &self.input.time_range { + write!( + f, + ", \"time_range\": {:?}", + ExplainString(&time_range.to_string()) + )?; + } + let summary = FileSummary::new(&self.input.files); + write!( + f, + ", \"file_summary\": {}", + serde_json::to_string(&summary).map_err(|_| fmt::Error)? + )?; + + // Source indices are unique within a range even when a file is split into + // row groups. Memtables and extension sources must not count as files. + let file_offset = self.input.num_memtables(); + let file_end = file_offset + self.input.num_files(); + let mut max_range = None; + let mut max_files = 0; + for range in &self.ranges { + let files = range + .indices + .iter() + .filter(|source| (file_offset..file_end).contains(&source.index)) + .count(); + if files > max_files { + max_files = files; + max_range = Some(range); + } + } + if let Some(range) = max_range { + let summary = FileSummary::new(range.indices.iter().filter_map(|source| { + source + .index + .checked_sub(file_offset) + .and_then(|index| self.input.files.get(index)) + })); + let (start, end) = range.time_range; + // RangeMeta bounds are inclusive, just like the verbose file list. + let summary = serde_json::json!({ + "time_range_start": format!("{}::{}", start.value(), start.unit()), + "time_range_end": format!("{}::{}", end.value(), end.unit()), + "files": summary.files, + "rows": summary.rows, + "row_groups": summary.row_groups, + "size": summary.size, + "index_size": summary.index_size, + "estimated_rows": range.num_rows, + }); + write!(f, ", \"max_part_range\": {summary}")?; + } + Ok(()) + } + fn format_verbose_content(&self, f: &mut fmt::Formatter) -> fmt::Result { struct FileWrapper<'a> { file: &'a FileHandle, @@ -2118,30 +2242,6 @@ impl StreamContext { impl fmt::Debug for InputWrapper<'_> { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - let output_schema = self.input.mapper.output_schema(); - if !output_schema.is_empty() { - let names: Vec<_> = output_schema - .column_schemas() - .iter() - .map(|col| &col.name) - .collect(); - write!(f, ", \"projection\": {:?}", names)?; - } - if let Some(predicate) = &self.input.predicate.predicate() { - if !predicate.exprs().is_empty() { - let exprs: Vec<_> = - predicate.exprs().iter().map(|e| e.to_string()).collect(); - write!(f, ", \"filters\": {:?}", exprs)?; - } - if !predicate.dyn_filters().is_empty() { - let dyn_filters: Vec<_> = predicate - .dyn_filters() - .iter() - .map(|f| format!("{}", f)) - .collect(); - write!(f, ", \"dyn_filters\": {:?}", dyn_filters)?; - } - } if !self.input.files.is_empty() { write!(f, ", \"files\": ")?; f.debug_list() diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 0bf6228f4c9..f4e1d16de49 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -163,6 +163,7 @@ macro_rules! http_tests { test_http_auth, test_sql_api, test_http_analyze_stream_tql, + test_http_analyze_stream_scan_diagnostics, test_http_sql_slow_query, test_prometheus_promql_api, test_promql_over_non_millisecond_physical_tables, @@ -1338,6 +1339,133 @@ async fn test_sql_format_api() { guard.remove_all().await; } +pub async fn test_http_analyze_stream_scan_diagnostics(store_type: StorageType) { + // Check both levels of JSON: the analyze node schema and the verbose scan + // object inside its string-valued `param`. + fn scan_details(node: &Value, scans: &mut Vec) { + assert!(node["name"].is_string()); + assert!(node["param"].is_string()); + assert!(node["output_rows"].is_u64()); + assert!(node["elapsed_compute"].is_u64()); + assert!(node["metrics"].is_object()); + if matches!( + node["name"].as_str(), + Some("SeqScan" | "UnorderedScan" | "SeriesScan") + ) { + let (_, fields) = node["param"].as_str().unwrap().split_once("), ").unwrap(); + let fields = if fields.starts_with('{') { + fields.to_string() + } else { + format!("{{{fields}}}") + }; + let details: Value = serde_json::from_str(&fields).unwrap(); + assert_eq!(details["projection"], json!(["reading", "ts"])); + assert!( + details["filters"] + .as_array() + .is_some_and(|filters| !filters.is_empty()) + ); + assert_eq!(details["time_range"], "TimestampRange{[0ms,6000ms)}"); + assert_eq!(details["file_summary"]["files"], 2); + assert_eq!(details["file_summary"]["rows"], 4); + assert_eq!(details["file_summary"]["row_groups"], 2); + assert_eq!(details["max_part_range"]["files"], 2); + assert!(details["max_part_range"].get("identifier").is_none()); + scans.push(details); + } + for child in node["children"].as_array().unwrap() { + scan_details(child, scans); + } + } + fn assert_verbose(scans: &[Value]) { + assert!(!scans.is_empty()); + for scan in scans { + let files = scan["files"].as_array().unwrap(); + assert_eq!(files.len(), 2); + for key in ["rows", "size", "index_size"] { + assert_eq!( + scan["file_summary"][key], + files + .iter() + .map(|file| file[key].as_u64().unwrap()) + .sum::() + ); + } + } + } + + let (app, mut guard) = setup_test_http_app_with_frontend(store_type, "scan_diagnostics").await; + let client = TestClient::new(app).await; + for sql in [ + "CREATE TABLE scan_diagnostics (host STRING PRIMARY KEY, reading DOUBLE, ts TIMESTAMP TIME INDEX)", + "INSERT INTO scan_diagnostics VALUES ('a', 1, 0), ('b', 2, 3000)", + "ADMIN FLUSH_TABLE('scan_diagnostics')", + "INSERT INTO scan_diagnostics VALUES ('c', 3, 2000), ('d', 4, 5000)", + "ADMIN FLUSH_TABLE('scan_diagnostics')", + "INSERT INTO scan_diagnostics VALUES ('outside', 5, 10000)", + "ADMIN FLUSH_TABLE('scan_diagnostics')", + "INSERT INTO scan_diagnostics VALUES ('mem', 6, 1000)", + ] { + let response = execute_sql(&client, sql).await; + let status = response.status(); + let body = response.text().await; + assert_eq!(status, StatusCode::OK, "{sql}: {body}"); + let body: Value = serde_json::from_str(&body).unwrap(); + assert!(body.get("error").is_none(), "{body}"); + } + let query = "SELECT reading, ts FROM scan_diagnostics WHERE ts >= 0 AND ts < 6000"; + let sql = format!("EXPLAIN ANALYZE FORMAT JSON {query}"); + let response = execute_sql(&client, &sql).await; + let status = response.status(); + let body = response.text().await; + assert_eq!(status, StatusCode::OK, "{body}"); + let body: Value = serde_json::from_str(&body).unwrap(); + let mut explain_scans = vec![]; + for row in body["output"][0]["records"]["rows"].as_array().unwrap() { + if row[0].is_null() { + assert_eq!(row[2], "Total rows: 5"); + continue; + } + let plan: Value = serde_json::from_str(row[2].as_str().unwrap()).unwrap(); + scan_details(&plan, &mut explain_scans); + } + assert!(!explain_scans.is_empty()); + assert!(explain_scans.iter().all(|scan| scan.get("files").is_none())); + + let sql = format!("EXPLAIN ANALYZE VERBOSE FORMAT JSON {query}"); + let response = client + .post("/v1/sql/analyze/stream?snapshot_interval_ms=100") + .header("Accept", "text/event-stream") + .form(&BTreeMap::from([("sql", sql)])) + .send() + .await; + assert_eq!(response.status(), StatusCode::OK); + let body = response.text().await; + let mut first_scan_seen = false; + let mut final_seen = false; + for event in body.split("\n\n") { + let Some(data) = event.lines().find_map(|line| line.strip_prefix("data: ")) else { + continue; + }; + let payload: Value = serde_json::from_str(data).unwrap(); + let mut scans = vec![]; + for stage in payload["metrics"].as_array().unwrap() { + assert!(stage["stage"].is_u64()); + assert!(stage["node"].is_u64()); + scan_details(&stage["plan"], &mut scans); + } + // Every available scan is checked, including its first appearance. A fast + // query may finish before the first periodic snapshot. + first_scan_seen |= !scans.is_empty(); + if payload["state"] == "final" { + assert_verbose(&scans); + final_seen = true; + } + } + assert!(first_scan_seen && final_seen, "{body}"); + guard.remove_all().await; +} + pub async fn test_http_analyze_stream_tql(store_type: StorageType) { common_telemetry::init_default_ut_logging(); let (app, mut guard) = diff --git a/tests/cases/distributed/explain/analyze.result b/tests/cases/distributed/explain/analyze.result index 77f6cff8b59..0fcdab8b704 100644 --- a/tests/cases/distributed/explain/analyze.result +++ b/tests/cases/distributed/explain/analyze.result @@ -33,6 +33,7 @@ SELECT count(*) FROM system_metrics; -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze SELECT count(*) FROM system_metrics; +-+-+-+ @@ -45,7 +46,7 @@ explain analyze SELECT count(*) FROM system_metrics; |_|_|_AggregateExec: mode=Final, gby=[], aggr=[count(system_REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[count(system_REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["ts"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ diff --git a/tests/cases/distributed/explain/analyze.sql b/tests/cases/distributed/explain/analyze.sql index bdf51cbca49..1e52a00236c 100644 --- a/tests/cases/distributed/explain/analyze.sql +++ b/tests/cases/distributed/explain/analyze.sql @@ -23,6 +23,7 @@ SELECT count(*) FROM system_metrics; -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze SELECT count(*) FROM system_metrics; drop table system_metrics; diff --git a/tests/cases/distributed/explain/order_by.result b/tests/cases/distributed/explain/order_by.result index e5781659521..062b5b26ad0 100644 --- a/tests/cases/distributed/explain/order_by.result +++ b/tests/cases/distributed/explain/order_by.result @@ -116,6 +116,7 @@ SELECT i, t AS alias_ts FROM test_pk ORDER BY alias_ts DESC LIMIT 5; -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT i, t AS alias_ts FROM test_pk ORDER BY t DESC LIMIT 5; +-+-+-+ @@ -127,7 +128,7 @@ EXPLAIN ANALYZE SELECT i, t AS alias_ts FROM test_pk ORDER BY t DESC LIMIT 5; | 1_| 0_|_ProjectionExec: expr=[i@0 as i, t@1 as alias_ts] REDACTED |_|_|_SortPreservingMergeExec: [t@1 DESC], fetch=5 REDACTED |_|_|_SortExec: TopK(fetch=5), expr=[t@1 DESC], preserve_partitioning=[true], filter=[t@1 IS NULL OR t@1 > 2] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["i", "t"], "dyn_filters":REDACTED |_|_|_| |_|_| Total rows: 5_| +-+-+-+ @@ -138,6 +139,7 @@ EXPLAIN ANALYZE SELECT i, t AS alias_ts FROM test_pk ORDER BY t DESC LIMIT 5; -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT i, t AS alias_ts FROM test_pk ORDER BY alias_ts DESC LIMIT 5; +-+-+-+ @@ -149,7 +151,7 @@ EXPLAIN ANALYZE SELECT i, t AS alias_ts FROM test_pk ORDER BY alias_ts DESC LIMI | 1_| 0_|_ProjectionExec: expr=[i@0 as i, t@1 as alias_ts] REDACTED |_|_|_SortPreservingMergeExec: [t@1 DESC], fetch=5 REDACTED |_|_|_SortExec: TopK(fetch=5), expr=[t@1 DESC], preserve_partitioning=[true], filter=[t@1 IS NULL OR t@1 > 2] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["i", "t"], "dyn_filters":REDACTED |_|_|_| |_|_| Total rows: 5_| +-+-+-+ diff --git a/tests/cases/distributed/explain/order_by.sql b/tests/cases/distributed/explain/order_by.sql index 3819e8061fa..152b70147f6 100644 --- a/tests/cases/distributed/explain/order_by.sql +++ b/tests/cases/distributed/explain/order_by.sql @@ -44,6 +44,7 @@ SELECT i, t AS alias_ts FROM test_pk ORDER BY alias_ts DESC LIMIT 5; -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT i, t AS alias_ts FROM test_pk ORDER BY t DESC LIMIT 5; -- SQLNESS REPLACE (-+) - @@ -52,6 +53,7 @@ EXPLAIN ANALYZE SELECT i, t AS alias_ts FROM test_pk ORDER BY t DESC LIMIT 5; -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT i, t AS alias_ts FROM test_pk ORDER BY alias_ts DESC LIMIT 5; DROP TABLE test_pk; diff --git a/tests/cases/distributed/explain/step_aggr.result b/tests/cases/distributed/explain/step_aggr.result index a69a82352b3..ed5dad360b2 100644 --- a/tests/cases/distributed/explain/step_aggr.result +++ b/tests/cases/distributed/explain/step_aggr.result @@ -81,6 +81,7 @@ FROM -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT count(i), @@ -103,17 +104,17 @@ FROM | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(integers.i), __sum_state(integers.i), __uddsketch_state_state(Int64(128),Float64(0.01),integers.i), __hll_state(integers.i)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__count_state(integers.i), __sum_state(integers.i), __uddsketch_state_state(Int64(128),Float64(0.01),integers.i), __hll_state(integers.i)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["i"], "filters": ["host < Utf8(\"550-A\") OR host IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(integers.i), __sum_state(integers.i), __uddsketch_state_state(Int64(128),Float64(0.01),integers.i), __hll_state(integers.i)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__count_state(integers.i), __sum_state(integers.i), __uddsketch_state_state(Int64(128),Float64(0.01),integers.i), __hll_state(integers.i)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["i"], "filters": ["host >= Utf8(\"550-A\") AND (host < Utf8(\"550-W\") OR host IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 2_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(integers.i), __sum_state(integers.i), __uddsketch_state_state(Int64(128),Float64(0.01),integers.i), __hll_state(integers.i)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__count_state(integers.i), __sum_state(integers.i), __uddsketch_state_state(Int64(128),Float64(0.01),integers.i), __hll_state(integers.i)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["i"], "filters": ["host >= Utf8(\"550-W\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ @@ -164,6 +165,7 @@ FROM -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT avg(i) FROM @@ -181,17 +183,17 @@ FROM | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__avg_state(integers.i)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__avg_state(integers.i)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["i"], "filters": ["host < Utf8(\"550-A\") OR host IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=Final, gby=[], aggr=[__avg_state(integers.i)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__avg_state(integers.i)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["i"], "filters": ["host >= Utf8(\"550-A\") AND (host < Utf8(\"550-W\") OR host IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 2_|_AggregateExec: mode=Final, gby=[], aggr=[__avg_state(integers.i)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__avg_state(integers.i)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["i"], "filters": ["host >= Utf8(\"550-W\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ @@ -267,6 +269,7 @@ ORDER BY -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT ts, @@ -296,17 +299,17 @@ ORDER BY | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[ts@0 as ts], aggr=[__count_state(integers.i), __sum_state(integers.i), __uddsketch_state_state(Int64(128),Float64(0.01),integers.i), __hll_state(integers.i)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[ts@1 as ts], aggr=[__count_state(integers.i), __sum_state(integers.i), __uddsketch_state_state(Int64(128),Float64(0.01),integers.i), __hll_state(integers.i)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["i", "ts"], "filters": ["host < Utf8(\"550-A\") OR host IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[ts@0 as ts], aggr=[__count_state(integers.i), __sum_state(integers.i), __uddsketch_state_state(Int64(128),Float64(0.01),integers.i), __hll_state(integers.i)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[ts@1 as ts], aggr=[__count_state(integers.i), __sum_state(integers.i), __uddsketch_state_state(Int64(128),Float64(0.01),integers.i), __hll_state(integers.i)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["i", "ts"], "filters": ["host >= Utf8(\"550-A\") AND (host < Utf8(\"550-W\") OR host IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 2_|_AggregateExec: mode=FinalPartitioned, gby=[ts@0 as ts], aggr=[__count_state(integers.i), __sum_state(integers.i), __uddsketch_state_state(Int64(128),Float64(0.01),integers.i), __hll_state(integers.i)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[ts@1 as ts], aggr=[__count_state(integers.i), __sum_state(integers.i), __uddsketch_state_state(Int64(128),Float64(0.01),integers.i), __hll_state(integers.i)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["i", "ts"], "filters": ["host >= Utf8(\"550-W\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -382,6 +385,7 @@ ORDER BY -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT date_bin('2s'::INTERVAL, ts) as time_window, @@ -411,17 +415,17 @@ ORDER BY | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("2 seconds"),integers.ts)@0 as date_bin(Utf8("2 seconds"),integers.ts)], aggr=[__count_state(integers.i), __sum_state(integers.i), __uddsketch_state_state(Int64(128),Float64(0.01),integers.i), __hll_state(integers.i)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 2000000000 }, ts@1) as date_bin(Utf8("2 seconds"),integers.ts)], aggr=[__count_state(integers.i), __sum_state(integers.i), __uddsketch_state_state(Int64(128),Float64(0.01),integers.i), __hll_state(integers.i)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["i", "ts"], "filters": ["host < Utf8(\"550-A\") OR host IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("2 seconds"),integers.ts)@0 as date_bin(Utf8("2 seconds"),integers.ts)], aggr=[__count_state(integers.i), __sum_state(integers.i), __uddsketch_state_state(Int64(128),Float64(0.01),integers.i), __hll_state(integers.i)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 2000000000 }, ts@1) as date_bin(Utf8("2 seconds"),integers.ts)], aggr=[__count_state(integers.i), __sum_state(integers.i), __uddsketch_state_state(Int64(128),Float64(0.01),integers.i), __hll_state(integers.i)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["i", "ts"], "filters": ["host >= Utf8(\"550-A\") AND (host < Utf8(\"550-W\") OR host IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 2_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("2 seconds"),integers.ts)@0 as date_bin(Utf8("2 seconds"),integers.ts)], aggr=[__count_state(integers.i), __sum_state(integers.i), __uddsketch_state_state(Int64(128),Float64(0.01),integers.i), __hll_state(integers.i)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 2000000000 }, ts@1) as date_bin(Utf8("2 seconds"),integers.ts)], aggr=[__count_state(integers.i), __sum_state(integers.i), __uddsketch_state_state(Int64(128),Float64(0.01),integers.i), __hll_state(integers.i)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["i", "ts"], "filters": ["host >= Utf8(\"550-W\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ diff --git a/tests/cases/distributed/explain/step_aggr.sql b/tests/cases/distributed/explain/step_aggr.sql index cb6cb3b8834..b8b8b0d6ffb 100644 --- a/tests/cases/distributed/explain/step_aggr.sql +++ b/tests/cases/distributed/explain/step_aggr.sql @@ -53,6 +53,7 @@ FROM -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT count(i), @@ -86,6 +87,7 @@ FROM -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT avg(i) FROM @@ -132,6 +134,7 @@ ORDER BY -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT ts, @@ -188,6 +191,7 @@ ORDER BY -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT date_bin('2s'::INTERVAL, ts) as time_window, diff --git a/tests/cases/distributed/explain/step_aggr_advance.result b/tests/cases/distributed/explain/step_aggr_advance.result index 11adf22f866..4186cb18c3a 100644 --- a/tests/cases/distributed/explain/step_aggr_advance.result +++ b/tests/cases/distributed/explain/step_aggr_advance.result @@ -49,6 +49,7 @@ tql explain (1752591864, 1752592164, '30s') max by (a, b, c) (max_over_time(aggr -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED tql analyze (1752591864, 1752592164, '30s') max by (a, b, c) (max_over_time(aggr_optimize_not [2m])); +-+-+-+ @@ -68,7 +69,7 @@ tql analyze (1752591864, 1752592164, '30s') max by (a, b, c) (max_over_time(aggr |_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["a", "b", "c", "d", "greptime_timestamp", "greptime_value"], "filters": ["greptime_timestamp >= TimestampMillisecond(1752591744001, None)", "greptime_timestamp <= TimestampMillisecond(1752592164000, None)", "a < Utf8(\"b\") OR a IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_ProjectionExec: expr=[a@0 as a, b@1 as b, c@2 as c, greptime_timestamp@3 as greptime_timestamp, prom_max(prom_max_over_time(greptime_timestamp_range,greptime_value))@4 as max(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED |_|_|_SortPreservingMergeExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] REDACTED @@ -81,7 +82,7 @@ tql analyze (1752591864, 1752592164, '30s') max by (a, b, c) (max_over_time(aggr |_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["a", "b", "c", "d", "greptime_timestamp", "greptime_value"], "filters": ["greptime_timestamp >= TimestampMillisecond(1752591744001, None)", "greptime_timestamp <= TimestampMillisecond(1752592164000, None)", "a >= Utf8(\"b\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -125,6 +126,7 @@ tql explain (1752591864, 1752592164, '30s') sum by (a, b) (max_over_time(aggr_op -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED tql analyze (1752591864, 1752592164, '30s') sum by (a, b) (max_over_time(aggr_optimize_not [2m])); +-+-+-+ @@ -146,7 +148,7 @@ tql analyze (1752591864, 1752592164, '30s') sum by (a, b) (max_over_time(aggr_op |_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["a", "b", "c", "d", "greptime_timestamp", "greptime_value"], "filters": ["greptime_timestamp >= TimestampMillisecond(1752591744001, None)", "greptime_timestamp <= TimestampMillisecond(1752592164000, None)", "a < Utf8(\"b\") OR a IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b, greptime_timestamp@2 as greptime_timestamp], aggr=[__sum_state(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED @@ -156,7 +158,7 @@ tql analyze (1752591864, 1752592164, '30s') sum by (a, b) (max_over_time(aggr_op |_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["a", "b", "c", "d", "greptime_timestamp", "greptime_value"], "filters": ["greptime_timestamp >= TimestampMillisecond(1752591744001, None)", "greptime_timestamp <= TimestampMillisecond(1752592164000, None)", "a >= Utf8(\"b\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -200,6 +202,7 @@ tql explain (1752591864, 1752592164, '30s') avg by (a) (max_over_time(aggr_optim -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED tql analyze (1752591864, 1752592164, '30s') avg by (a) (max_over_time(aggr_optimize_not [2m])); +-+-+-+ @@ -221,7 +224,7 @@ tql analyze (1752591864, 1752592164, '30s') avg by (a) (max_over_time(aggr_optim |_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["a", "b", "c", "d", "greptime_timestamp", "greptime_value"], "filters": ["greptime_timestamp >= TimestampMillisecond(1752591744001, None)", "greptime_timestamp <= TimestampMillisecond(1752592164000, None)", "a < Utf8(\"b\") OR a IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, greptime_timestamp@1 as greptime_timestamp], aggr=[__avg_state(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED @@ -231,7 +234,7 @@ tql analyze (1752591864, 1752592164, '30s') avg by (a) (max_over_time(aggr_optim |_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["a", "b", "c", "d", "greptime_timestamp", "greptime_value"], "filters": ["greptime_timestamp >= TimestampMillisecond(1752591744001, None)", "greptime_timestamp <= TimestampMillisecond(1752592164000, None)", "a >= Utf8(\"b\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -270,6 +273,7 @@ tql explain (1752591864, 1752592164, '30s') count by (a, b, c, d) (max_over_time -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED tql analyze (1752591864, 1752592164, '30s') count by (a, b, c, d) (max_over_time(aggr_optimize_not [2m])); +-+-+-+ @@ -286,7 +290,7 @@ tql analyze (1752591864, 1752592164, '30s') count by (a, b, c, d) (max_over_time |_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["a", "b", "c", "d", "greptime_timestamp", "greptime_value"], "filters": ["greptime_timestamp >= TimestampMillisecond(1752591744001, None)", "greptime_timestamp <= TimestampMillisecond(1752592164000, None)", "a < Utf8(\"b\") OR a IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_SortPreservingMergeExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, d@3 ASC NULLS LAST, greptime_timestamp@4 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, d@3 ASC NULLS LAST, greptime_timestamp@4 ASC NULLS LAST], preserve_partitioning=[true] REDACTED @@ -296,7 +300,7 @@ tql analyze (1752591864, 1752592164, '30s') count by (a, b, c, d) (max_over_time |_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["a", "b", "c", "d", "greptime_timestamp", "greptime_value"], "filters": ["greptime_timestamp >= TimestampMillisecond(1752591744001, None)", "greptime_timestamp <= TimestampMillisecond(1752592164000, None)", "a >= Utf8(\"b\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -340,6 +344,7 @@ tql explain (1752591864, 1752592164, '30s') min by (b, c, d) (max_over_time(aggr -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED tql analyze (1752591864, 1752592164, '30s') min by (b, c, d) (max_over_time(aggr_optimize_not [2m])); +-+-+-+ @@ -361,7 +366,7 @@ tql analyze (1752591864, 1752592164, '30s') min by (b, c, d) (max_over_time(aggr |_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["a", "b", "c", "d", "greptime_timestamp", "greptime_value"], "filters": ["greptime_timestamp >= TimestampMillisecond(1752591744001, None)", "greptime_timestamp <= TimestampMillisecond(1752592164000, None)", "a < Utf8(\"b\") OR a IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[b@0 as b, c@1 as c, d@2 as d, greptime_timestamp@3 as greptime_timestamp], aggr=[__prom_min_state(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED @@ -371,7 +376,7 @@ tql analyze (1752591864, 1752592164, '30s') min by (b, c, d) (max_over_time(aggr |_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["a", "b", "c", "d", "greptime_timestamp", "greptime_value"], "filters": ["greptime_timestamp >= TimestampMillisecond(1752591744001, None)", "greptime_timestamp <= TimestampMillisecond(1752592164000, None)", "a >= Utf8(\"b\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -412,6 +417,7 @@ tql explain sum(aggr_optimize_not); -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED tql analyze sum(aggr_optimize_not); +-+-+-+ @@ -431,7 +437,7 @@ tql analyze sum(aggr_optimize_not); |_|_|_ProjectionExec: expr=[greptime_timestamp@4 as greptime_timestamp, greptime_value@5 as greptime_value] REDACTED |_|_|_PromInstantManipulateExec: range=[0..0], lookback=[300000], interval=[300000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "selector":"LastRow { after_merge: true }", "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "selector":"LastRow { after_merge: true }", "distribution":"PerSeries", "projection": ["a", "b", "c", "d", "greptime_timestamp", "greptime_value"], "filters": ["greptime_timestamp >= TimestampMillisecond(-299999, None)", "greptime_timestamp <= TimestampMillisecond(0, None)", "a < Utf8(\"b\") OR a IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[greptime_timestamp@0 as greptime_timestamp], aggr=[__sum_state(aggr_optimize_not.greptime_value)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED @@ -439,7 +445,7 @@ tql analyze sum(aggr_optimize_not); |_|_|_ProjectionExec: expr=[greptime_timestamp@4 as greptime_timestamp, greptime_value@5 as greptime_value] REDACTED |_|_|_PromInstantManipulateExec: range=[0..0], lookback=[300000], interval=[300000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "selector":"LastRow { after_merge: true }", "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "selector":"LastRow { after_merge: true }", "distribution":"PerSeries", "projection": ["a", "b", "c", "d", "greptime_timestamp", "greptime_value"], "filters": ["greptime_timestamp >= TimestampMillisecond(-299999, None)", "greptime_timestamp <= TimestampMillisecond(0, None)", "a >= Utf8(\"b\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -459,6 +465,24 @@ CREATE TABLE IF NOT EXISTS aggr_optimize_not_count ( Affected Rows: 0 -- Case 6: Test average rate (sum/count like) +-- Populate both partitions on both sides so the join cannot short-circuit on +-- empty input and omit remote ANALYZE plans depending on execution timing. +INSERT INTO aggr_optimize_not VALUES + ('a', 'b', 'c', 'd', 1752591840000, 10), + ('a', 'b', 'c', 'd', 1752591860000, 20), + ('c', 'b', 'c', 'd', 1752591840000, 30), + ('c', 'b', 'c', 'd', 1752591860000, 40); + +Affected Rows: 4 + +INSERT INTO aggr_optimize_not_count VALUES + ('a', 'b', 'c', 'd', 1752591840000, 1), + ('a', 'b', 'c', 'd', 1752591860000, 2), + ('c', 'b', 'c', 'd', 1752591840000, 3), + ('c', 'b', 'c', 'd', 1752591860000, 4); + +Affected Rows: 4 + -- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED @@ -524,6 +548,7 @@ tql explain (1752591864, 1752592164, '30s') sum by (a, b, c) (rate(aggr_optimize -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED tql analyze (1752591864, 1752592164, '30s') sum by (a, b, c) (rate(aggr_optimize_not [2m])) / sum by (a, b, c) (rate(aggr_optimize_not_count [2m])); +-+-+-+ @@ -549,15 +574,49 @@ tql analyze (1752591864, 1752592164, '30s') sum by (a, b, c) (rate(aggr_optimize |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_CooperativeExec REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["a", "b", "c", "d", "greptime_timestamp", "greptime_value"], "filters": ["greptime_timestamp >= TimestampMillisecond(1752591744001, None)", "greptime_timestamp <= TimestampMillisecond(1752592164000, None)", "a < Utf8(\"b\") OR a IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_CooperativeExec REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["a", "b", "c", "d", "greptime_timestamp", "greptime_value"], "filters": ["greptime_timestamp >= TimestampMillisecond(1752591744001, None)", "greptime_timestamp <= TimestampMillisecond(1752592164000, None)", "a >= Utf8(\"b\")"], "time_range":REDACTED |_|_|_| -|_|_| Total rows: 0_| +| 1_| 0_|_ProjectionExec: expr=[a@0 as a, b@1 as b, c@2 as c, greptime_timestamp@3 as greptime_timestamp, sum(prom_rate(greptime_timestamp_range,greptime_value,aggr_optimize_not.greptime_timestamp,Int64(120000)))@4 as sum(prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000)))] REDACTED +|_|_|_SortPreservingMergeExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] REDACTED +|_|_|_SortExec: expr=[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST], preserve_partitioning=[true] REDACTED +|_|_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b, c@2 as c, greptime_timestamp@3 as greptime_timestamp], aggr=[sum(prom_rate(greptime_timestamp_range,greptime_value,aggr_optimize_not.greptime_timestamp,Int64(120000)))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED +|_|_|_AggregateExec: mode=Partial, gby=[a@2 as a, b@3 as b, c@4 as c, greptime_timestamp@0 as greptime_timestamp], aggr=[sum(prom_rate(greptime_timestamp_range,greptime_value,aggr_optimize_not.greptime_timestamp,Int64(120000)))] REDACTED +|_|_|_FilterExec: prom_rate(greptime_timestamp_range,greptime_value,aggr_optimize_not.greptime_timestamp,Int64(120000))@1 IS NOT NULL REDACTED +|_|_|_ProjectionExec: expr=[greptime_timestamp@4 as greptime_timestamp, prom_rate(greptime_timestamp_range@6, greptime_value@5, greptime_timestamp@4, 120000) as prom_rate(greptime_timestamp_range,greptime_value,aggr_optimize_not.greptime_timestamp,Int64(120000)), a@0 as a, b@1 as b, c@2 as c] REDACTED +|_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED +|_|_|_PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] REDACTED +|_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["a", "b", "c", "d", "greptime_timestamp", "greptime_value"], "filters": ["greptime_timestamp >= TimestampMillisecond(1752591744001, None)", "greptime_timestamp <= TimestampMillisecond(1752592164000, None)", "a < Utf8(\"b\") OR a IS NULL"], "time_range":REDACTED +|_|_|_| +| 1_| 1_|_ProjectionExec: expr=[a@0 as a, b@1 as b, c@2 as c, greptime_timestamp@3 as greptime_timestamp, sum(prom_rate(greptime_timestamp_range,greptime_value,aggr_optimize_not.greptime_timestamp,Int64(120000)))@4 as sum(prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000)))] REDACTED +|_|_|_SortPreservingMergeExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] REDACTED +|_|_|_SortExec: expr=[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST], preserve_partitioning=[true] REDACTED +|_|_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b, c@2 as c, greptime_timestamp@3 as greptime_timestamp], aggr=[sum(prom_rate(greptime_timestamp_range,greptime_value,aggr_optimize_not.greptime_timestamp,Int64(120000)))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED +|_|_|_AggregateExec: mode=Partial, gby=[a@2 as a, b@3 as b, c@4 as c, greptime_timestamp@0 as greptime_timestamp], aggr=[sum(prom_rate(greptime_timestamp_range,greptime_value,aggr_optimize_not.greptime_timestamp,Int64(120000)))] REDACTED +|_|_|_FilterExec: prom_rate(greptime_timestamp_range,greptime_value,aggr_optimize_not.greptime_timestamp,Int64(120000))@1 IS NOT NULL REDACTED +|_|_|_ProjectionExec: expr=[greptime_timestamp@4 as greptime_timestamp, prom_rate(greptime_timestamp_range@6, greptime_value@5, greptime_timestamp@4, 120000) as prom_rate(greptime_timestamp_range,greptime_value,aggr_optimize_not.greptime_timestamp,Int64(120000)), a@0 as a, b@1 as b, c@2 as c] REDACTED +|_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED +|_|_|_PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] REDACTED +|_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["a", "b", "c", "d", "greptime_timestamp", "greptime_value"], "filters": ["greptime_timestamp >= TimestampMillisecond(1752591744001, None)", "greptime_timestamp <= TimestampMillisecond(1752592164000, None)", "a >= Utf8(\"b\")"], "time_range":REDACTED +|_|_|_| +|_|_| Total rows: 8_| +-+-+-+ -- Case 7: aggregate without sort should be pushed down. This one push down for include all partition columns. +TRUNCATE TABLE aggr_optimize_not; + +Affected Rows: 0 + +TRUNCATE TABLE aggr_optimize_not_count; + +Affected Rows: 0 + -- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED @@ -591,6 +650,7 @@ GROUP BY -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT min(greptime_value) @@ -611,13 +671,13 @@ GROUP BY |_|_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b, c@2 as c], aggr=[min(aggr_optimize_not.greptime_value)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[a@0 as a, b@1 as b, c@2 as c], aggr=[min(aggr_optimize_not.greptime_value)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "projection": ["a", "b", "c", "greptime_value"], "filters": ["a < Utf8(\"b\") OR a IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_ProjectionExec: expr=[min(aggr_optimize_not.greptime_value)@3 as min(aggr_optimize_not.greptime_value)] REDACTED |_|_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b, c@2 as c], aggr=[min(aggr_optimize_not.greptime_value)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[a@0 as a, b@1 as b, c@2 as c], aggr=[min(aggr_optimize_not.greptime_value)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "projection": ["a", "b", "c", "greptime_value"], "filters": ["a >= Utf8(\"b\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -657,6 +717,7 @@ GROUP BY -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT min(greptime_value) @@ -678,13 +739,13 @@ GROUP BY |_|_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b, c@2 as c, d@3 as d], aggr=[min(aggr_optimize_not.greptime_value)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[a@0 as a, b@1 as b, c@2 as c, d@3 as d], aggr=[min(aggr_optimize_not.greptime_value)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "projection": ["a", "b", "c", "d", "greptime_value"], "filters": ["a < Utf8(\"b\") OR a IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_ProjectionExec: expr=[min(aggr_optimize_not.greptime_value)@4 as min(aggr_optimize_not.greptime_value)] REDACTED |_|_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b, c@2 as c, d@3 as d], aggr=[min(aggr_optimize_not.greptime_value)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[a@0 as a, b@1 as b, c@2 as c, d@3 as d], aggr=[min(aggr_optimize_not.greptime_value)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "projection": ["a", "b", "c", "d", "greptime_value"], "filters": ["a >= Utf8(\"b\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -727,6 +788,7 @@ GROUP BY -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT min(greptime_value) @@ -749,12 +811,12 @@ GROUP BY | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b], aggr=[__min_state(aggr_optimize_not.greptime_value)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[a@0 as a, b@1 as b], aggr=[__min_state(aggr_optimize_not.greptime_value)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "projection": ["a", "b", "greptime_value"], "filters": ["a < Utf8(\"b\") OR a IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b], aggr=[__min_state(aggr_optimize_not.greptime_value)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[a@0 as a, b@1 as b], aggr=[__min_state(aggr_optimize_not.greptime_value)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "projection": ["a", "b", "greptime_value"], "filters": ["a >= Utf8(\"b\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -797,6 +859,7 @@ GROUP BY -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT min(greptime_value) + max(greptime_value) @@ -819,12 +882,12 @@ GROUP BY | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b], aggr=[__min_state(aggr_optimize_not.greptime_value), __max_state(aggr_optimize_not.greptime_value)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[a@0 as a, b@1 as b], aggr=[__min_state(aggr_optimize_not.greptime_value), __max_state(aggr_optimize_not.greptime_value)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "projection": ["a", "b", "greptime_value"], "filters": ["a < Utf8(\"b\") OR a IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b], aggr=[__min_state(aggr_optimize_not.greptime_value), __max_state(aggr_optimize_not.greptime_value)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[a@0 as a, b@1 as b], aggr=[__min_state(aggr_optimize_not.greptime_value), __max_state(aggr_optimize_not.greptime_value)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "projection": ["a", "b", "greptime_value"], "filters": ["a >= Utf8(\"b\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -876,6 +939,7 @@ GROUP BY -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT a, @@ -907,12 +971,12 @@ GROUP BY | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[__min_state(aggr_optimize_not.greptime_value)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[__min_state(aggr_optimize_not.greptime_value)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "projection": ["a", "greptime_value"], "filters": ["a < Utf8(\"b\") OR a IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[__min_state(aggr_optimize_not.greptime_value)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[__min_state(aggr_optimize_not.greptime_value)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "projection": ["a", "greptime_value"], "filters": ["a >= Utf8(\"b\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ diff --git a/tests/cases/distributed/explain/step_aggr_advance.sql b/tests/cases/distributed/explain/step_aggr_advance.sql index e9a6a3363e3..adf932d9aa4 100644 --- a/tests/cases/distributed/explain/step_aggr_advance.sql +++ b/tests/cases/distributed/explain/step_aggr_advance.sql @@ -26,6 +26,7 @@ tql explain (1752591864, 1752592164, '30s') max by (a, b, c) (max_over_time(aggr -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED tql analyze (1752591864, 1752592164, '30s') max by (a, b, c) (max_over_time(aggr_optimize_not [2m])); -- Case 1: group by columns are prefix of partition columns. @@ -41,6 +42,7 @@ tql explain (1752591864, 1752592164, '30s') sum by (a, b) (max_over_time(aggr_op -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED tql analyze (1752591864, 1752592164, '30s') sum by (a, b) (max_over_time(aggr_optimize_not [2m])); -- Case 2: group by columns are prefix of partition columns. @@ -56,6 +58,7 @@ tql explain (1752591864, 1752592164, '30s') avg by (a) (max_over_time(aggr_optim -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED tql analyze (1752591864, 1752592164, '30s') avg by (a) (max_over_time(aggr_optimize_not [2m])); -- Case 3: group by columns are superset of partition columns. @@ -71,6 +74,7 @@ tql explain (1752591864, 1752592164, '30s') count by (a, b, c, d) (max_over_time -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED tql analyze (1752591864, 1752592164, '30s') count by (a, b, c, d) (max_over_time(aggr_optimize_not [2m])); -- Case 4: group by columns are not prefix of partition columns. @@ -86,6 +90,7 @@ tql explain (1752591864, 1752592164, '30s') min by (b, c, d) (max_over_time(aggr -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED tql analyze (1752591864, 1752592164, '30s') min by (b, c, d) (max_over_time(aggr_optimize_not [2m])); -- Case 5: a simple sum @@ -101,6 +106,7 @@ tql explain sum(aggr_optimize_not); -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED tql analyze sum(aggr_optimize_not); -- TODO(discord9): more cases for aggr push down interacting with partitioning&tql @@ -116,6 +122,19 @@ CREATE TABLE IF NOT EXISTS aggr_optimize_not_count ( ) PARTITION ON COLUMNS (a, b, c) (a < 'b', a >= 'b',); -- Case 6: Test average rate (sum/count like) +-- Populate both partitions on both sides so the join cannot short-circuit on +-- empty input and omit remote ANALYZE plans depending on execution timing. +INSERT INTO aggr_optimize_not VALUES + ('a', 'b', 'c', 'd', 1752591840000, 10), + ('a', 'b', 'c', 'd', 1752591860000, 20), + ('c', 'b', 'c', 'd', 1752591840000, 30), + ('c', 'b', 'c', 'd', 1752591860000, 40); +INSERT INTO aggr_optimize_not_count VALUES + ('a', 'b', 'c', 'd', 1752591840000, 1), + ('a', 'b', 'c', 'd', 1752591860000, 2), + ('c', 'b', 'c', 'd', 1752591840000, 3), + ('c', 'b', 'c', 'd', 1752591860000, 4); + -- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED @@ -128,9 +147,13 @@ tql explain (1752591864, 1752592164, '30s') sum by (a, b, c) (rate(aggr_optimize -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED tql analyze (1752591864, 1752592164, '30s') sum by (a, b, c) (rate(aggr_optimize_not [2m])) / sum by (a, b, c) (rate(aggr_optimize_not_count [2m])); -- Case 7: aggregate without sort should be pushed down. This one push down for include all partition columns. +TRUNCATE TABLE aggr_optimize_not; +TRUNCATE TABLE aggr_optimize_not_count; + -- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED @@ -151,6 +174,7 @@ GROUP BY -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT min(greptime_value) @@ -183,6 +207,7 @@ GROUP BY -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT min(greptime_value) @@ -214,6 +239,7 @@ GROUP BY -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT min(greptime_value) @@ -243,6 +269,7 @@ GROUP BY -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT min(greptime_value) + max(greptime_value) @@ -283,6 +310,7 @@ GROUP BY -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT a, diff --git a/tests/cases/distributed/explain/step_aggr_basic.result b/tests/cases/distributed/explain/step_aggr_basic.result index 39880450875..282d0f441da 100644 --- a/tests/cases/distributed/explain/step_aggr_basic.result +++ b/tests/cases/distributed/explain/step_aggr_basic.result @@ -74,6 +74,7 @@ FROM -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT count(i) @@ -92,17 +93,17 @@ FROM | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(integers.i)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__count_state(integers.i)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["i"], "filters": ["host < Utf8(\"550-A\") OR host IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(integers.i)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__count_state(integers.i)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["i"], "filters": ["host >= Utf8(\"550-A\") AND (host < Utf8(\"550-W\") OR host IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 2_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(integers.i)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__count_state(integers.i)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["i"], "filters": ["host >= Utf8(\"550-W\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ @@ -172,6 +173,7 @@ ORDER BY -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT ts, @@ -198,17 +200,17 @@ ORDER BY | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[ts@0 as ts], aggr=[__count_state(integers.i)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[ts@1 as ts], aggr=[__count_state(integers.i)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["i", "ts"], "filters": ["host < Utf8(\"550-A\") OR host IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[ts@0 as ts], aggr=[__count_state(integers.i)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[ts@1 as ts], aggr=[__count_state(integers.i)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["i", "ts"], "filters": ["host >= Utf8(\"550-A\") AND (host < Utf8(\"550-W\") OR host IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 2_|_AggregateExec: mode=FinalPartitioned, gby=[ts@0 as ts], aggr=[__count_state(integers.i)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[ts@1 as ts], aggr=[__count_state(integers.i)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["i", "ts"], "filters": ["host >= Utf8(\"550-W\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -280,6 +282,7 @@ ORDER BY -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT date_bin('1 hour' :: INTERVAL, ts) as time_window, @@ -307,17 +310,17 @@ ORDER BY | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("1 hour"),integers.ts)@0 as date_bin(Utf8("1 hour"),integers.ts)], aggr=[__count_state(integers.i)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 3600000000000 }, ts@1) as date_bin(Utf8("1 hour"),integers.ts)], aggr=[__count_state(integers.i)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["i", "ts"], "filters": ["host < Utf8(\"550-A\") OR host IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("1 hour"),integers.ts)@0 as date_bin(Utf8("1 hour"),integers.ts)], aggr=[__count_state(integers.i)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 3600000000000 }, ts@1) as date_bin(Utf8("1 hour"),integers.ts)], aggr=[__count_state(integers.i)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["i", "ts"], "filters": ["host >= Utf8(\"550-A\") AND (host < Utf8(\"550-W\") OR host IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 2_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("1 hour"),integers.ts)@0 as date_bin(Utf8("1 hour"),integers.ts)], aggr=[__count_state(integers.i)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 3600000000000 }, ts@1) as date_bin(Utf8("1 hour"),integers.ts)], aggr=[__count_state(integers.i)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["i", "ts"], "filters": ["host >= Utf8(\"550-W\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -397,6 +400,7 @@ ORDER BY -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT ts::BIGINT + 1, @@ -426,17 +430,17 @@ ORDER BY | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[integers.ts + Int64(1)@0 as integers.ts + Int64(1), integers.i / Int64(2)@1 as integers.i / Int64(2)], aggr=[__count_state(integers.i)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[CAST(ts@1 AS Int64) + 1 as integers.ts + Int64(1), i@0 / 2 as integers.i / Int64(2)], aggr=[__count_state(integers.i)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["i", "ts"], "filters": ["host < Utf8(\"550-A\") OR host IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[integers.ts + Int64(1)@0 as integers.ts + Int64(1), integers.i / Int64(2)@1 as integers.i / Int64(2)], aggr=[__count_state(integers.i)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[CAST(ts@1 AS Int64) + 1 as integers.ts + Int64(1), i@0 / 2 as integers.i / Int64(2)], aggr=[__count_state(integers.i)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["i", "ts"], "filters": ["host >= Utf8(\"550-A\") AND (host < Utf8(\"550-W\") OR host IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 2_|_AggregateExec: mode=FinalPartitioned, gby=[integers.ts + Int64(1)@0 as integers.ts + Int64(1), integers.i / Int64(2)@1 as integers.i / Int64(2)], aggr=[__count_state(integers.i)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[CAST(ts@1 AS Int64) + 1 as integers.ts + Int64(1), i@0 / 2 as integers.i / Int64(2)], aggr=[__count_state(integers.i)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["i", "ts"], "filters": ["host >= Utf8(\"550-W\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 8_| +-+-+-+ @@ -522,6 +526,7 @@ FROM -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT uddsketch_calc(0.5, uddsketch_merge(128, 0.01, udd_state)) as udd_result, @@ -542,17 +547,17 @@ FROM | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__uddsketch_merge_state(Int64(128),Float64(0.01),sink_table.udd_state), __hll_merge_state(sink_table.hll_state)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__uddsketch_merge_state(Int64(128),Float64(0.01),sink_table.udd_state), __hll_merge_state(sink_table.hll_state)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["udd_state", "hll_state"], "filters": ["host < Utf8(\"550-A\") OR host IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=Final, gby=[], aggr=[__uddsketch_merge_state(Int64(128),Float64(0.01),sink_table.udd_state), __hll_merge_state(sink_table.hll_state)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__uddsketch_merge_state(Int64(128),Float64(0.01),sink_table.udd_state), __hll_merge_state(sink_table.hll_state)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["udd_state", "hll_state"], "filters": ["host >= Utf8(\"550-A\") AND (host < Utf8(\"550-W\") OR host IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 2_|_AggregateExec: mode=Final, gby=[], aggr=[__uddsketch_merge_state(Int64(128),Float64(0.01),sink_table.udd_state), __hll_merge_state(sink_table.hll_state)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__uddsketch_merge_state(Int64(128),Float64(0.01),sink_table.udd_state), __hll_merge_state(sink_table.hll_state)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["udd_state", "hll_state"], "filters": ["host >= Utf8(\"550-W\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ diff --git a/tests/cases/distributed/explain/step_aggr_basic.sql b/tests/cases/distributed/explain/step_aggr_basic.sql index 9a0d226a130..d693ccbaabd 100644 --- a/tests/cases/distributed/explain/step_aggr_basic.sql +++ b/tests/cases/distributed/explain/step_aggr_basic.sql @@ -48,6 +48,7 @@ FROM -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT count(i) @@ -91,6 +92,7 @@ ORDER BY -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT ts, @@ -140,6 +142,7 @@ ORDER BY -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT date_bin('1 hour' :: INTERVAL, ts) as time_window, @@ -195,6 +198,7 @@ ORDER BY -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT ts::BIGINT + 1, @@ -264,6 +268,7 @@ FROM -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT uddsketch_calc(0.5, uddsketch_merge(128, 0.01, udd_state)) as udd_result, diff --git a/tests/cases/distributed/explain/step_aggr_massive.result b/tests/cases/distributed/explain/step_aggr_massive.result index 5aa969c6741..3a8e9634653 100644 --- a/tests/cases/distributed/explain/step_aggr_massive.result +++ b/tests/cases/distributed/explain/step_aggr_massive.result @@ -269,6 +269,7 @@ GROUP BY -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT env, @@ -466,112 +467,112 @@ GROUP BY |_|_|_AggregateExec: mode=Partial, gby=[env@7 as env, service_name@8 as service_name, city@9 as city, page@10 as page, CAST(date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }, time@6) AS Timestamp(s)) as arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))], aggr=[__uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(base_table.shard_key)] REDACTED |_|_|_ProjectionExec: expr=[CASE WHEN lcp@5 > 0 AND lcp@5 < 3000000 THEN lcp@5 END as __common_expr_1, CASE WHEN fmp@6 > 0 AND fmp@6 < 3000000 THEN fmp@6 END as __common_expr_2, CASE WHEN fcp@7 > 0 AND fcp@7 < 3000000 THEN fcp@7 END as __common_expr_3, CASE WHEN fp@8 > 0 AND fp@8 < 3000000 THEN fp@8 END as __common_expr_4, CASE WHEN tti@9 > 0 AND tti@9 < 3000000 THEN tti@9 END as __common_expr_5, CASE WHEN fid@10 > 0 AND fid@10 < 3000000 THEN fid@10 END as __common_expr_6, time@0 as time, env@1 as env, service_name@2 as service_name, city@3 as city, page@4 as page, shard_key@11 as shard_key] REDACTED |_|_|_FilterExec: lcp@5 > 0 AND lcp@5 < 3000000 OR fmp@6 > 0 AND fmp@6 < 3000000 OR fcp@7 > 0 AND fcp@7 < 3000000 OR fp@8 > 0 AND fp@8 < 3000000 OR tti@9 > 0 AND tti@9 < 3000000 OR fid@10 > 0 AND fid@10 < 3000000 REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED, "projection": ["time", "env", "service_name", "city", "page", "lcp", "fmp", "fcp", "fp", "tti", "fid", "shard_key"], "filters": ["lcp > Int64(0) AND lcp < Int64(3000000) OR fmp > Int64(0) AND fmp < Int64(3000000) OR fcp > Int64(0) AND fcp < Int64(3000000) OR fp > Int64(0) AND fp < Int64(3000000) OR tti > Int64(0) AND tti < Int64(3000000) OR fid > Int64(0) AND fid < Int64(3000000)", "time >= TimestampMillisecond(0, None)", "shard_key < Int64(4) OR shard_key IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[env@0 as env, service_name@1 as service_name, city@2 as city, page@3 as page, arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))@4 as arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))], aggr=[__uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(base_table.shard_key)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[env@7 as env, service_name@8 as service_name, city@9 as city, page@10 as page, CAST(date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }, time@6) AS Timestamp(s)) as arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))], aggr=[__uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(base_table.shard_key)] REDACTED |_|_|_ProjectionExec: expr=[CASE WHEN lcp@5 > 0 AND lcp@5 < 3000000 THEN lcp@5 END as __common_expr_1, CASE WHEN fmp@6 > 0 AND fmp@6 < 3000000 THEN fmp@6 END as __common_expr_2, CASE WHEN fcp@7 > 0 AND fcp@7 < 3000000 THEN fcp@7 END as __common_expr_3, CASE WHEN fp@8 > 0 AND fp@8 < 3000000 THEN fp@8 END as __common_expr_4, CASE WHEN tti@9 > 0 AND tti@9 < 3000000 THEN tti@9 END as __common_expr_5, CASE WHEN fid@10 > 0 AND fid@10 < 3000000 THEN fid@10 END as __common_expr_6, time@0 as time, env@1 as env, service_name@2 as service_name, city@3 as city, page@4 as page, shard_key@11 as shard_key] REDACTED |_|_|_FilterExec: lcp@5 > 0 AND lcp@5 < 3000000 OR fmp@6 > 0 AND fmp@6 < 3000000 OR fcp@7 > 0 AND fcp@7 < 3000000 OR fp@8 > 0 AND fp@8 < 3000000 OR tti@9 > 0 AND tti@9 < 3000000 OR fid@10 > 0 AND fid@10 < 3000000 REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED, "projection": ["time", "env", "service_name", "city", "page", "lcp", "fmp", "fcp", "fp", "tti", "fid", "shard_key"], "filters": ["lcp > Int64(0) AND lcp < Int64(3000000) OR fmp > Int64(0) AND fmp < Int64(3000000) OR fcp > Int64(0) AND fcp < Int64(3000000) OR fp > Int64(0) AND fp < Int64(3000000) OR tti > Int64(0) AND tti < Int64(3000000) OR fid > Int64(0) AND fid < Int64(3000000)", "time >= TimestampMillisecond(0, None)", "shard_key >= Int64(4) AND (shard_key < Int64(8) OR shard_key IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 2_|_AggregateExec: mode=FinalPartitioned, gby=[env@0 as env, service_name@1 as service_name, city@2 as city, page@3 as page, arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))@4 as arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))], aggr=[__uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(base_table.shard_key)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[env@7 as env, service_name@8 as service_name, city@9 as city, page@10 as page, CAST(date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }, time@6) AS Timestamp(s)) as arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))], aggr=[__uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(base_table.shard_key)] REDACTED |_|_|_ProjectionExec: expr=[CASE WHEN lcp@5 > 0 AND lcp@5 < 3000000 THEN lcp@5 END as __common_expr_1, CASE WHEN fmp@6 > 0 AND fmp@6 < 3000000 THEN fmp@6 END as __common_expr_2, CASE WHEN fcp@7 > 0 AND fcp@7 < 3000000 THEN fcp@7 END as __common_expr_3, CASE WHEN fp@8 > 0 AND fp@8 < 3000000 THEN fp@8 END as __common_expr_4, CASE WHEN tti@9 > 0 AND tti@9 < 3000000 THEN tti@9 END as __common_expr_5, CASE WHEN fid@10 > 0 AND fid@10 < 3000000 THEN fid@10 END as __common_expr_6, time@0 as time, env@1 as env, service_name@2 as service_name, city@3 as city, page@4 as page, shard_key@11 as shard_key] REDACTED |_|_|_FilterExec: lcp@5 > 0 AND lcp@5 < 3000000 OR fmp@6 > 0 AND fmp@6 < 3000000 OR fcp@7 > 0 AND fcp@7 < 3000000 OR fp@8 > 0 AND fp@8 < 3000000 OR tti@9 > 0 AND tti@9 < 3000000 OR fid@10 > 0 AND fid@10 < 3000000 REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED, "projection": ["time", "env", "service_name", "city", "page", "lcp", "fmp", "fcp", "fp", "tti", "fid", "shard_key"], "filters": ["lcp > Int64(0) AND lcp < Int64(3000000) OR fmp > Int64(0) AND fmp < Int64(3000000) OR fcp > Int64(0) AND fcp < Int64(3000000) OR fp > Int64(0) AND fp < Int64(3000000) OR tti > Int64(0) AND tti < Int64(3000000) OR fid > Int64(0) AND fid < Int64(3000000)", "time >= TimestampMillisecond(0, None)", "shard_key >= Int64(8) AND (shard_key < Int64(12) OR shard_key IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 3_|_AggregateExec: mode=FinalPartitioned, gby=[env@0 as env, service_name@1 as service_name, city@2 as city, page@3 as page, arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))@4 as arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))], aggr=[__uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(base_table.shard_key)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[env@7 as env, service_name@8 as service_name, city@9 as city, page@10 as page, CAST(date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }, time@6) AS Timestamp(s)) as arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))], aggr=[__uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(base_table.shard_key)] REDACTED |_|_|_ProjectionExec: expr=[CASE WHEN lcp@5 > 0 AND lcp@5 < 3000000 THEN lcp@5 END as __common_expr_1, CASE WHEN fmp@6 > 0 AND fmp@6 < 3000000 THEN fmp@6 END as __common_expr_2, CASE WHEN fcp@7 > 0 AND fcp@7 < 3000000 THEN fcp@7 END as __common_expr_3, CASE WHEN fp@8 > 0 AND fp@8 < 3000000 THEN fp@8 END as __common_expr_4, CASE WHEN tti@9 > 0 AND tti@9 < 3000000 THEN tti@9 END as __common_expr_5, CASE WHEN fid@10 > 0 AND fid@10 < 3000000 THEN fid@10 END as __common_expr_6, time@0 as time, env@1 as env, service_name@2 as service_name, city@3 as city, page@4 as page, shard_key@11 as shard_key] REDACTED |_|_|_FilterExec: lcp@5 > 0 AND lcp@5 < 3000000 OR fmp@6 > 0 AND fmp@6 < 3000000 OR fcp@7 > 0 AND fcp@7 < 3000000 OR fp@8 > 0 AND fp@8 < 3000000 OR tti@9 > 0 AND tti@9 < 3000000 OR fid@10 > 0 AND fid@10 < 3000000 REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED, "projection": ["time", "env", "service_name", "city", "page", "lcp", "fmp", "fcp", "fp", "tti", "fid", "shard_key"], "filters": ["lcp > Int64(0) AND lcp < Int64(3000000) OR fmp > Int64(0) AND fmp < Int64(3000000) OR fcp > Int64(0) AND fcp < Int64(3000000) OR fp > Int64(0) AND fp < Int64(3000000) OR tti > Int64(0) AND tti < Int64(3000000) OR fid > Int64(0) AND fid < Int64(3000000)", "time >= TimestampMillisecond(0, None)", "shard_key >= Int64(12) AND (shard_key < Int64(16) OR shard_key IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 4_|_AggregateExec: mode=FinalPartitioned, gby=[env@0 as env, service_name@1 as service_name, city@2 as city, page@3 as page, arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))@4 as arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))], aggr=[__uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(base_table.shard_key)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[env@7 as env, service_name@8 as service_name, city@9 as city, page@10 as page, CAST(date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }, time@6) AS Timestamp(s)) as arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))], aggr=[__uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(base_table.shard_key)] REDACTED |_|_|_ProjectionExec: expr=[CASE WHEN lcp@5 > 0 AND lcp@5 < 3000000 THEN lcp@5 END as __common_expr_1, CASE WHEN fmp@6 > 0 AND fmp@6 < 3000000 THEN fmp@6 END as __common_expr_2, CASE WHEN fcp@7 > 0 AND fcp@7 < 3000000 THEN fcp@7 END as __common_expr_3, CASE WHEN fp@8 > 0 AND fp@8 < 3000000 THEN fp@8 END as __common_expr_4, CASE WHEN tti@9 > 0 AND tti@9 < 3000000 THEN tti@9 END as __common_expr_5, CASE WHEN fid@10 > 0 AND fid@10 < 3000000 THEN fid@10 END as __common_expr_6, time@0 as time, env@1 as env, service_name@2 as service_name, city@3 as city, page@4 as page, shard_key@11 as shard_key] REDACTED |_|_|_FilterExec: lcp@5 > 0 AND lcp@5 < 3000000 OR fmp@6 > 0 AND fmp@6 < 3000000 OR fcp@7 > 0 AND fcp@7 < 3000000 OR fp@8 > 0 AND fp@8 < 3000000 OR tti@9 > 0 AND tti@9 < 3000000 OR fid@10 > 0 AND fid@10 < 3000000 REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED, "projection": ["time", "env", "service_name", "city", "page", "lcp", "fmp", "fcp", "fp", "tti", "fid", "shard_key"], "filters": ["lcp > Int64(0) AND lcp < Int64(3000000) OR fmp > Int64(0) AND fmp < Int64(3000000) OR fcp > Int64(0) AND fcp < Int64(3000000) OR fp > Int64(0) AND fp < Int64(3000000) OR tti > Int64(0) AND tti < Int64(3000000) OR fid > Int64(0) AND fid < Int64(3000000)", "time >= TimestampMillisecond(0, None)", "shard_key >= Int64(16) AND (shard_key < Int64(20) OR shard_key IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 5_|_AggregateExec: mode=FinalPartitioned, gby=[env@0 as env, service_name@1 as service_name, city@2 as city, page@3 as page, arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))@4 as arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))], aggr=[__uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(base_table.shard_key)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[env@7 as env, service_name@8 as service_name, city@9 as city, page@10 as page, CAST(date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }, time@6) AS Timestamp(s)) as arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))], aggr=[__uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(base_table.shard_key)] REDACTED |_|_|_ProjectionExec: expr=[CASE WHEN lcp@5 > 0 AND lcp@5 < 3000000 THEN lcp@5 END as __common_expr_1, CASE WHEN fmp@6 > 0 AND fmp@6 < 3000000 THEN fmp@6 END as __common_expr_2, CASE WHEN fcp@7 > 0 AND fcp@7 < 3000000 THEN fcp@7 END as __common_expr_3, CASE WHEN fp@8 > 0 AND fp@8 < 3000000 THEN fp@8 END as __common_expr_4, CASE WHEN tti@9 > 0 AND tti@9 < 3000000 THEN tti@9 END as __common_expr_5, CASE WHEN fid@10 > 0 AND fid@10 < 3000000 THEN fid@10 END as __common_expr_6, time@0 as time, env@1 as env, service_name@2 as service_name, city@3 as city, page@4 as page, shard_key@11 as shard_key] REDACTED |_|_|_FilterExec: lcp@5 > 0 AND lcp@5 < 3000000 OR fmp@6 > 0 AND fmp@6 < 3000000 OR fcp@7 > 0 AND fcp@7 < 3000000 OR fp@8 > 0 AND fp@8 < 3000000 OR tti@9 > 0 AND tti@9 < 3000000 OR fid@10 > 0 AND fid@10 < 3000000 REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED, "projection": ["time", "env", "service_name", "city", "page", "lcp", "fmp", "fcp", "fp", "tti", "fid", "shard_key"], "filters": ["lcp > Int64(0) AND lcp < Int64(3000000) OR fmp > Int64(0) AND fmp < Int64(3000000) OR fcp > Int64(0) AND fcp < Int64(3000000) OR fp > Int64(0) AND fp < Int64(3000000) OR tti > Int64(0) AND tti < Int64(3000000) OR fid > Int64(0) AND fid < Int64(3000000)", "time >= TimestampMillisecond(0, None)", "shard_key >= Int64(20) AND (shard_key < Int64(24) OR shard_key IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 6_|_AggregateExec: mode=FinalPartitioned, gby=[env@0 as env, service_name@1 as service_name, city@2 as city, page@3 as page, arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))@4 as arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))], aggr=[__uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(base_table.shard_key)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[env@7 as env, service_name@8 as service_name, city@9 as city, page@10 as page, CAST(date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }, time@6) AS Timestamp(s)) as arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))], aggr=[__uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(base_table.shard_key)] REDACTED |_|_|_ProjectionExec: expr=[CASE WHEN lcp@5 > 0 AND lcp@5 < 3000000 THEN lcp@5 END as __common_expr_1, CASE WHEN fmp@6 > 0 AND fmp@6 < 3000000 THEN fmp@6 END as __common_expr_2, CASE WHEN fcp@7 > 0 AND fcp@7 < 3000000 THEN fcp@7 END as __common_expr_3, CASE WHEN fp@8 > 0 AND fp@8 < 3000000 THEN fp@8 END as __common_expr_4, CASE WHEN tti@9 > 0 AND tti@9 < 3000000 THEN tti@9 END as __common_expr_5, CASE WHEN fid@10 > 0 AND fid@10 < 3000000 THEN fid@10 END as __common_expr_6, time@0 as time, env@1 as env, service_name@2 as service_name, city@3 as city, page@4 as page, shard_key@11 as shard_key] REDACTED |_|_|_FilterExec: lcp@5 > 0 AND lcp@5 < 3000000 OR fmp@6 > 0 AND fmp@6 < 3000000 OR fcp@7 > 0 AND fcp@7 < 3000000 OR fp@8 > 0 AND fp@8 < 3000000 OR tti@9 > 0 AND tti@9 < 3000000 OR fid@10 > 0 AND fid@10 < 3000000 REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED, "projection": ["time", "env", "service_name", "city", "page", "lcp", "fmp", "fcp", "fp", "tti", "fid", "shard_key"], "filters": ["lcp > Int64(0) AND lcp < Int64(3000000) OR fmp > Int64(0) AND fmp < Int64(3000000) OR fcp > Int64(0) AND fcp < Int64(3000000) OR fp > Int64(0) AND fp < Int64(3000000) OR tti > Int64(0) AND tti < Int64(3000000) OR fid > Int64(0) AND fid < Int64(3000000)", "time >= TimestampMillisecond(0, None)", "shard_key >= Int64(24) AND (shard_key < Int64(28) OR shard_key IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 7_|_AggregateExec: mode=FinalPartitioned, gby=[env@0 as env, service_name@1 as service_name, city@2 as city, page@3 as page, arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))@4 as arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))], aggr=[__uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(base_table.shard_key)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[env@7 as env, service_name@8 as service_name, city@9 as city, page@10 as page, CAST(date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }, time@6) AS Timestamp(s)) as arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))], aggr=[__uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(base_table.shard_key)] REDACTED |_|_|_ProjectionExec: expr=[CASE WHEN lcp@5 > 0 AND lcp@5 < 3000000 THEN lcp@5 END as __common_expr_1, CASE WHEN fmp@6 > 0 AND fmp@6 < 3000000 THEN fmp@6 END as __common_expr_2, CASE WHEN fcp@7 > 0 AND fcp@7 < 3000000 THEN fcp@7 END as __common_expr_3, CASE WHEN fp@8 > 0 AND fp@8 < 3000000 THEN fp@8 END as __common_expr_4, CASE WHEN tti@9 > 0 AND tti@9 < 3000000 THEN tti@9 END as __common_expr_5, CASE WHEN fid@10 > 0 AND fid@10 < 3000000 THEN fid@10 END as __common_expr_6, time@0 as time, env@1 as env, service_name@2 as service_name, city@3 as city, page@4 as page, shard_key@11 as shard_key] REDACTED |_|_|_FilterExec: lcp@5 > 0 AND lcp@5 < 3000000 OR fmp@6 > 0 AND fmp@6 < 3000000 OR fcp@7 > 0 AND fcp@7 < 3000000 OR fp@8 > 0 AND fp@8 < 3000000 OR tti@9 > 0 AND tti@9 < 3000000 OR fid@10 > 0 AND fid@10 < 3000000 REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED, "projection": ["time", "env", "service_name", "city", "page", "lcp", "fmp", "fcp", "fp", "tti", "fid", "shard_key"], "filters": ["lcp > Int64(0) AND lcp < Int64(3000000) OR fmp > Int64(0) AND fmp < Int64(3000000) OR fcp > Int64(0) AND fcp < Int64(3000000) OR fp > Int64(0) AND fp < Int64(3000000) OR tti > Int64(0) AND tti < Int64(3000000) OR fid > Int64(0) AND fid < Int64(3000000)", "time >= TimestampMillisecond(0, None)", "shard_key >= Int64(28) AND (shard_key < Int64(32) OR shard_key IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 8_|_AggregateExec: mode=FinalPartitioned, gby=[env@0 as env, service_name@1 as service_name, city@2 as city, page@3 as page, arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))@4 as arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))], aggr=[__uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(base_table.shard_key)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[env@7 as env, service_name@8 as service_name, city@9 as city, page@10 as page, CAST(date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }, time@6) AS Timestamp(s)) as arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))], aggr=[__uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(base_table.shard_key)] REDACTED |_|_|_ProjectionExec: expr=[CASE WHEN lcp@5 > 0 AND lcp@5 < 3000000 THEN lcp@5 END as __common_expr_1, CASE WHEN fmp@6 > 0 AND fmp@6 < 3000000 THEN fmp@6 END as __common_expr_2, CASE WHEN fcp@7 > 0 AND fcp@7 < 3000000 THEN fcp@7 END as __common_expr_3, CASE WHEN fp@8 > 0 AND fp@8 < 3000000 THEN fp@8 END as __common_expr_4, CASE WHEN tti@9 > 0 AND tti@9 < 3000000 THEN tti@9 END as __common_expr_5, CASE WHEN fid@10 > 0 AND fid@10 < 3000000 THEN fid@10 END as __common_expr_6, time@0 as time, env@1 as env, service_name@2 as service_name, city@3 as city, page@4 as page, shard_key@11 as shard_key] REDACTED |_|_|_FilterExec: lcp@5 > 0 AND lcp@5 < 3000000 OR fmp@6 > 0 AND fmp@6 < 3000000 OR fcp@7 > 0 AND fcp@7 < 3000000 OR fp@8 > 0 AND fp@8 < 3000000 OR tti@9 > 0 AND tti@9 < 3000000 OR fid@10 > 0 AND fid@10 < 3000000 REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED, "projection": ["time", "env", "service_name", "city", "page", "lcp", "fmp", "fcp", "fp", "tti", "fid", "shard_key"], "filters": ["lcp > Int64(0) AND lcp < Int64(3000000) OR fmp > Int64(0) AND fmp < Int64(3000000) OR fcp > Int64(0) AND fcp < Int64(3000000) OR fp > Int64(0) AND fp < Int64(3000000) OR tti > Int64(0) AND tti < Int64(3000000) OR fid > Int64(0) AND fid < Int64(3000000)", "time >= TimestampMillisecond(0, None)", "shard_key >= Int64(32) AND (shard_key < Int64(36) OR shard_key IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 9_|_AggregateExec: mode=FinalPartitioned, gby=[env@0 as env, service_name@1 as service_name, city@2 as city, page@3 as page, arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))@4 as arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))], aggr=[__uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(base_table.shard_key)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[env@7 as env, service_name@8 as service_name, city@9 as city, page@10 as page, CAST(date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }, time@6) AS Timestamp(s)) as arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))], aggr=[__uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(base_table.shard_key)] REDACTED |_|_|_ProjectionExec: expr=[CASE WHEN lcp@5 > 0 AND lcp@5 < 3000000 THEN lcp@5 END as __common_expr_1, CASE WHEN fmp@6 > 0 AND fmp@6 < 3000000 THEN fmp@6 END as __common_expr_2, CASE WHEN fcp@7 > 0 AND fcp@7 < 3000000 THEN fcp@7 END as __common_expr_3, CASE WHEN fp@8 > 0 AND fp@8 < 3000000 THEN fp@8 END as __common_expr_4, CASE WHEN tti@9 > 0 AND tti@9 < 3000000 THEN tti@9 END as __common_expr_5, CASE WHEN fid@10 > 0 AND fid@10 < 3000000 THEN fid@10 END as __common_expr_6, time@0 as time, env@1 as env, service_name@2 as service_name, city@3 as city, page@4 as page, shard_key@11 as shard_key] REDACTED |_|_|_FilterExec: lcp@5 > 0 AND lcp@5 < 3000000 OR fmp@6 > 0 AND fmp@6 < 3000000 OR fcp@7 > 0 AND fcp@7 < 3000000 OR fp@8 > 0 AND fp@8 < 3000000 OR tti@9 > 0 AND tti@9 < 3000000 OR fid@10 > 0 AND fid@10 < 3000000 REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED, "projection": ["time", "env", "service_name", "city", "page", "lcp", "fmp", "fcp", "fp", "tti", "fid", "shard_key"], "filters": ["lcp > Int64(0) AND lcp < Int64(3000000) OR fmp > Int64(0) AND fmp < Int64(3000000) OR fcp > Int64(0) AND fcp < Int64(3000000) OR fp > Int64(0) AND fp < Int64(3000000) OR tti > Int64(0) AND tti < Int64(3000000) OR fid > Int64(0) AND fid < Int64(3000000)", "time >= TimestampMillisecond(0, None)", "shard_key >= Int64(36) AND (shard_key < Int64(40) OR shard_key IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 10_|_AggregateExec: mode=FinalPartitioned, gby=[env@0 as env, service_name@1 as service_name, city@2 as city, page@3 as page, arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))@4 as arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))], aggr=[__uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(base_table.shard_key)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[env@7 as env, service_name@8 as service_name, city@9 as city, page@10 as page, CAST(date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }, time@6) AS Timestamp(s)) as arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))], aggr=[__uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(base_table.shard_key)] REDACTED |_|_|_ProjectionExec: expr=[CASE WHEN lcp@5 > 0 AND lcp@5 < 3000000 THEN lcp@5 END as __common_expr_1, CASE WHEN fmp@6 > 0 AND fmp@6 < 3000000 THEN fmp@6 END as __common_expr_2, CASE WHEN fcp@7 > 0 AND fcp@7 < 3000000 THEN fcp@7 END as __common_expr_3, CASE WHEN fp@8 > 0 AND fp@8 < 3000000 THEN fp@8 END as __common_expr_4, CASE WHEN tti@9 > 0 AND tti@9 < 3000000 THEN tti@9 END as __common_expr_5, CASE WHEN fid@10 > 0 AND fid@10 < 3000000 THEN fid@10 END as __common_expr_6, time@0 as time, env@1 as env, service_name@2 as service_name, city@3 as city, page@4 as page, shard_key@11 as shard_key] REDACTED |_|_|_FilterExec: lcp@5 > 0 AND lcp@5 < 3000000 OR fmp@6 > 0 AND fmp@6 < 3000000 OR fcp@7 > 0 AND fcp@7 < 3000000 OR fp@8 > 0 AND fp@8 < 3000000 OR tti@9 > 0 AND tti@9 < 3000000 OR fid@10 > 0 AND fid@10 < 3000000 REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED, "projection": ["time", "env", "service_name", "city", "page", "lcp", "fmp", "fcp", "fp", "tti", "fid", "shard_key"], "filters": ["lcp > Int64(0) AND lcp < Int64(3000000) OR fmp > Int64(0) AND fmp < Int64(3000000) OR fcp > Int64(0) AND fcp < Int64(3000000) OR fp > Int64(0) AND fp < Int64(3000000) OR tti > Int64(0) AND tti < Int64(3000000) OR fid > Int64(0) AND fid < Int64(3000000)", "time >= TimestampMillisecond(0, None)", "shard_key >= Int64(40) AND (shard_key < Int64(44) OR shard_key IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 11_|_AggregateExec: mode=FinalPartitioned, gby=[env@0 as env, service_name@1 as service_name, city@2 as city, page@3 as page, arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))@4 as arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))], aggr=[__uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(base_table.shard_key)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[env@7 as env, service_name@8 as service_name, city@9 as city, page@10 as page, CAST(date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }, time@6) AS Timestamp(s)) as arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))], aggr=[__uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(base_table.shard_key)] REDACTED |_|_|_ProjectionExec: expr=[CASE WHEN lcp@5 > 0 AND lcp@5 < 3000000 THEN lcp@5 END as __common_expr_1, CASE WHEN fmp@6 > 0 AND fmp@6 < 3000000 THEN fmp@6 END as __common_expr_2, CASE WHEN fcp@7 > 0 AND fcp@7 < 3000000 THEN fcp@7 END as __common_expr_3, CASE WHEN fp@8 > 0 AND fp@8 < 3000000 THEN fp@8 END as __common_expr_4, CASE WHEN tti@9 > 0 AND tti@9 < 3000000 THEN tti@9 END as __common_expr_5, CASE WHEN fid@10 > 0 AND fid@10 < 3000000 THEN fid@10 END as __common_expr_6, time@0 as time, env@1 as env, service_name@2 as service_name, city@3 as city, page@4 as page, shard_key@11 as shard_key] REDACTED |_|_|_FilterExec: lcp@5 > 0 AND lcp@5 < 3000000 OR fmp@6 > 0 AND fmp@6 < 3000000 OR fcp@7 > 0 AND fcp@7 < 3000000 OR fp@8 > 0 AND fp@8 < 3000000 OR tti@9 > 0 AND tti@9 < 3000000 OR fid@10 > 0 AND fid@10 < 3000000 REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED, "projection": ["time", "env", "service_name", "city", "page", "lcp", "fmp", "fcp", "fp", "tti", "fid", "shard_key"], "filters": ["lcp > Int64(0) AND lcp < Int64(3000000) OR fmp > Int64(0) AND fmp < Int64(3000000) OR fcp > Int64(0) AND fcp < Int64(3000000) OR fp > Int64(0) AND fp < Int64(3000000) OR tti > Int64(0) AND tti < Int64(3000000) OR fid > Int64(0) AND fid < Int64(3000000)", "time >= TimestampMillisecond(0, None)", "shard_key >= Int64(44) AND (shard_key < Int64(48) OR shard_key IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 12_|_AggregateExec: mode=FinalPartitioned, gby=[env@0 as env, service_name@1 as service_name, city@2 as city, page@3 as page, arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))@4 as arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))], aggr=[__uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(base_table.shard_key)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[env@7 as env, service_name@8 as service_name, city@9 as city, page@10 as page, CAST(date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }, time@6) AS Timestamp(s)) as arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))], aggr=[__uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(base_table.shard_key)] REDACTED |_|_|_ProjectionExec: expr=[CASE WHEN lcp@5 > 0 AND lcp@5 < 3000000 THEN lcp@5 END as __common_expr_1, CASE WHEN fmp@6 > 0 AND fmp@6 < 3000000 THEN fmp@6 END as __common_expr_2, CASE WHEN fcp@7 > 0 AND fcp@7 < 3000000 THEN fcp@7 END as __common_expr_3, CASE WHEN fp@8 > 0 AND fp@8 < 3000000 THEN fp@8 END as __common_expr_4, CASE WHEN tti@9 > 0 AND tti@9 < 3000000 THEN tti@9 END as __common_expr_5, CASE WHEN fid@10 > 0 AND fid@10 < 3000000 THEN fid@10 END as __common_expr_6, time@0 as time, env@1 as env, service_name@2 as service_name, city@3 as city, page@4 as page, shard_key@11 as shard_key] REDACTED |_|_|_FilterExec: lcp@5 > 0 AND lcp@5 < 3000000 OR fmp@6 > 0 AND fmp@6 < 3000000 OR fcp@7 > 0 AND fcp@7 < 3000000 OR fp@8 > 0 AND fp@8 < 3000000 OR tti@9 > 0 AND tti@9 < 3000000 OR fid@10 > 0 AND fid@10 < 3000000 REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED, "projection": ["time", "env", "service_name", "city", "page", "lcp", "fmp", "fcp", "fp", "tti", "fid", "shard_key"], "filters": ["lcp > Int64(0) AND lcp < Int64(3000000) OR fmp > Int64(0) AND fmp < Int64(3000000) OR fcp > Int64(0) AND fcp < Int64(3000000) OR fp > Int64(0) AND fp < Int64(3000000) OR tti > Int64(0) AND tti < Int64(3000000) OR fid > Int64(0) AND fid < Int64(3000000)", "time >= TimestampMillisecond(0, None)", "shard_key >= Int64(48) AND (shard_key < Int64(52) OR shard_key IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 13_|_AggregateExec: mode=FinalPartitioned, gby=[env@0 as env, service_name@1 as service_name, city@2 as city, page@3 as page, arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))@4 as arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))], aggr=[__uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(base_table.shard_key)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[env@7 as env, service_name@8 as service_name, city@9 as city, page@10 as page, CAST(date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }, time@6) AS Timestamp(s)) as arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))], aggr=[__uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(base_table.shard_key)] REDACTED |_|_|_ProjectionExec: expr=[CASE WHEN lcp@5 > 0 AND lcp@5 < 3000000 THEN lcp@5 END as __common_expr_1, CASE WHEN fmp@6 > 0 AND fmp@6 < 3000000 THEN fmp@6 END as __common_expr_2, CASE WHEN fcp@7 > 0 AND fcp@7 < 3000000 THEN fcp@7 END as __common_expr_3, CASE WHEN fp@8 > 0 AND fp@8 < 3000000 THEN fp@8 END as __common_expr_4, CASE WHEN tti@9 > 0 AND tti@9 < 3000000 THEN tti@9 END as __common_expr_5, CASE WHEN fid@10 > 0 AND fid@10 < 3000000 THEN fid@10 END as __common_expr_6, time@0 as time, env@1 as env, service_name@2 as service_name, city@3 as city, page@4 as page, shard_key@11 as shard_key] REDACTED |_|_|_FilterExec: lcp@5 > 0 AND lcp@5 < 3000000 OR fmp@6 > 0 AND fmp@6 < 3000000 OR fcp@7 > 0 AND fcp@7 < 3000000 OR fp@8 > 0 AND fp@8 < 3000000 OR tti@9 > 0 AND tti@9 < 3000000 OR fid@10 > 0 AND fid@10 < 3000000 REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED, "projection": ["time", "env", "service_name", "city", "page", "lcp", "fmp", "fcp", "fp", "tti", "fid", "shard_key"], "filters": ["lcp > Int64(0) AND lcp < Int64(3000000) OR fmp > Int64(0) AND fmp < Int64(3000000) OR fcp > Int64(0) AND fcp < Int64(3000000) OR fp > Int64(0) AND fp < Int64(3000000) OR tti > Int64(0) AND tti < Int64(3000000) OR fid > Int64(0) AND fid < Int64(3000000)", "time >= TimestampMillisecond(0, None)", "shard_key >= Int64(52) AND (shard_key < Int64(56) OR shard_key IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 14_|_AggregateExec: mode=FinalPartitioned, gby=[env@0 as env, service_name@1 as service_name, city@2 as city, page@3 as page, arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))@4 as arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))], aggr=[__uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(base_table.shard_key)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[env@7 as env, service_name@8 as service_name, city@9 as city, page@10 as page, CAST(date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }, time@6) AS Timestamp(s)) as arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))], aggr=[__uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(base_table.shard_key)] REDACTED |_|_|_ProjectionExec: expr=[CASE WHEN lcp@5 > 0 AND lcp@5 < 3000000 THEN lcp@5 END as __common_expr_1, CASE WHEN fmp@6 > 0 AND fmp@6 < 3000000 THEN fmp@6 END as __common_expr_2, CASE WHEN fcp@7 > 0 AND fcp@7 < 3000000 THEN fcp@7 END as __common_expr_3, CASE WHEN fp@8 > 0 AND fp@8 < 3000000 THEN fp@8 END as __common_expr_4, CASE WHEN tti@9 > 0 AND tti@9 < 3000000 THEN tti@9 END as __common_expr_5, CASE WHEN fid@10 > 0 AND fid@10 < 3000000 THEN fid@10 END as __common_expr_6, time@0 as time, env@1 as env, service_name@2 as service_name, city@3 as city, page@4 as page, shard_key@11 as shard_key] REDACTED |_|_|_FilterExec: lcp@5 > 0 AND lcp@5 < 3000000 OR fmp@6 > 0 AND fmp@6 < 3000000 OR fcp@7 > 0 AND fcp@7 < 3000000 OR fp@8 > 0 AND fp@8 < 3000000 OR tti@9 > 0 AND tti@9 < 3000000 OR fid@10 > 0 AND fid@10 < 3000000 REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED, "projection": ["time", "env", "service_name", "city", "page", "lcp", "fmp", "fcp", "fp", "tti", "fid", "shard_key"], "filters": ["lcp > Int64(0) AND lcp < Int64(3000000) OR fmp > Int64(0) AND fmp < Int64(3000000) OR fcp > Int64(0) AND fcp < Int64(3000000) OR fp > Int64(0) AND fp < Int64(3000000) OR tti > Int64(0) AND tti < Int64(3000000) OR fid > Int64(0) AND fid < Int64(3000000)", "time >= TimestampMillisecond(0, None)", "shard_key >= Int64(56) AND (shard_key < Int64(60) OR shard_key IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 15_|_AggregateExec: mode=FinalPartitioned, gby=[env@0 as env, service_name@1 as service_name, city@2 as city, page@3 as page, arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))@4 as arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))], aggr=[__uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(base_table.shard_key)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[env@7 as env, service_name@8 as service_name, city@9 as city, page@10 as page, CAST(date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }, time@6) AS Timestamp(s)) as arrow_cast(date_bin(Utf8("60 seconds"),base_table.time),Utf8("Timestamp(s)"))], aggr=[__uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.lcp > Int64(0) AND base_table.lcp < Int64(3000000) THEN base_table.lcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fmp > Int64(0) AND base_table.fmp < Int64(3000000) THEN base_table.fmp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fcp > Int64(0) AND base_table.fcp < Int64(3000000) THEN base_table.fcp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fp > Int64(0) AND base_table.fp < Int64(3000000) THEN base_table.fp ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.tti > Int64(0) AND base_table.tti < Int64(3000000) THEN base_table.tti ELSE Int64(NULL) END), __uddsketch_state_state(Int64(128),Float64(0.01),CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __min_state(CASE WHEN base_table.fid > Int64(0) AND base_table.fid < Int64(3000000) THEN base_table.fid ELSE Int64(NULL) END), __max_state(base_table.shard_key)] REDACTED |_|_|_ProjectionExec: expr=[CASE WHEN lcp@5 > 0 AND lcp@5 < 3000000 THEN lcp@5 END as __common_expr_1, CASE WHEN fmp@6 > 0 AND fmp@6 < 3000000 THEN fmp@6 END as __common_expr_2, CASE WHEN fcp@7 > 0 AND fcp@7 < 3000000 THEN fcp@7 END as __common_expr_3, CASE WHEN fp@8 > 0 AND fp@8 < 3000000 THEN fp@8 END as __common_expr_4, CASE WHEN tti@9 > 0 AND tti@9 < 3000000 THEN tti@9 END as __common_expr_5, CASE WHEN fid@10 > 0 AND fid@10 < 3000000 THEN fid@10 END as __common_expr_6, time@0 as time, env@1 as env, service_name@2 as service_name, city@3 as city, page@4 as page, shard_key@11 as shard_key] REDACTED |_|_|_FilterExec: lcp@5 > 0 AND lcp@5 < 3000000 OR fmp@6 > 0 AND fmp@6 < 3000000 OR fcp@7 > 0 AND fcp@7 < 3000000 OR fp@8 > 0 AND fp@8 < 3000000 OR tti@9 > 0 AND tti@9 < 3000000 OR fid@10 > 0 AND fid@10 < 3000000 REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED, "projection": ["time", "env", "service_name", "city", "page", "lcp", "fmp", "fcp", "fp", "tti", "fid", "shard_key"], "filters": ["lcp > Int64(0) AND lcp < Int64(3000000) OR fmp > Int64(0) AND fmp < Int64(3000000) OR fcp > Int64(0) AND fcp < Int64(3000000) OR fp > Int64(0) AND fp < Int64(3000000) OR tti > Int64(0) AND tti < Int64(3000000) OR fid > Int64(0) AND fid < Int64(3000000)", "time >= TimestampMillisecond(0, None)", "shard_key >= Int64(60)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -616,6 +617,7 @@ where -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT count(*) @@ -636,82 +638,82 @@ where | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(base_table.time)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__count_state(base_table.time)] REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED, "projection": ["time"], "filters": ["time >= TimestampMillisecond(0, None)", "shard_key < Int64(4) OR shard_key IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(base_table.time)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__count_state(base_table.time)] REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED, "projection": ["time"], "filters": ["time >= TimestampMillisecond(0, None)", "shard_key >= Int64(4) AND (shard_key < Int64(8) OR shard_key IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 2_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(base_table.time)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__count_state(base_table.time)] REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED, "projection": ["time"], "filters": ["time >= TimestampMillisecond(0, None)", "shard_key >= Int64(8) AND (shard_key < Int64(12) OR shard_key IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 3_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(base_table.time)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__count_state(base_table.time)] REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED, "projection": ["time"], "filters": ["time >= TimestampMillisecond(0, None)", "shard_key >= Int64(12) AND (shard_key < Int64(16) OR shard_key IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 4_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(base_table.time)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__count_state(base_table.time)] REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED, "projection": ["time"], "filters": ["time >= TimestampMillisecond(0, None)", "shard_key >= Int64(16) AND (shard_key < Int64(20) OR shard_key IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 5_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(base_table.time)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__count_state(base_table.time)] REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED, "projection": ["time"], "filters": ["time >= TimestampMillisecond(0, None)", "shard_key >= Int64(20) AND (shard_key < Int64(24) OR shard_key IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 6_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(base_table.time)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__count_state(base_table.time)] REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED, "projection": ["time"], "filters": ["time >= TimestampMillisecond(0, None)", "shard_key >= Int64(24) AND (shard_key < Int64(28) OR shard_key IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 7_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(base_table.time)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__count_state(base_table.time)] REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED, "projection": ["time"], "filters": ["time >= TimestampMillisecond(0, None)", "shard_key >= Int64(28) AND (shard_key < Int64(32) OR shard_key IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 8_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(base_table.time)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__count_state(base_table.time)] REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED, "projection": ["time"], "filters": ["time >= TimestampMillisecond(0, None)", "shard_key >= Int64(32) AND (shard_key < Int64(36) OR shard_key IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 9_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(base_table.time)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__count_state(base_table.time)] REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED, "projection": ["time"], "filters": ["time >= TimestampMillisecond(0, None)", "shard_key >= Int64(36) AND (shard_key < Int64(40) OR shard_key IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 10_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(base_table.time)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__count_state(base_table.time)] REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED, "projection": ["time"], "filters": ["time >= TimestampMillisecond(0, None)", "shard_key >= Int64(40) AND (shard_key < Int64(44) OR shard_key IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 11_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(base_table.time)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__count_state(base_table.time)] REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED, "projection": ["time"], "filters": ["time >= TimestampMillisecond(0, None)", "shard_key >= Int64(44) AND (shard_key < Int64(48) OR shard_key IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 12_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(base_table.time)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__count_state(base_table.time)] REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED, "projection": ["time"], "filters": ["time >= TimestampMillisecond(0, None)", "shard_key >= Int64(48) AND (shard_key < Int64(52) OR shard_key IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 13_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(base_table.time)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__count_state(base_table.time)] REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED, "projection": ["time"], "filters": ["time >= TimestampMillisecond(0, None)", "shard_key >= Int64(52) AND (shard_key < Int64(56) OR shard_key IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 14_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(base_table.time)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__count_state(base_table.time)] REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED, "projection": ["time"], "filters": ["time >= TimestampMillisecond(0, None)", "shard_key >= Int64(56) AND (shard_key < Int64(60) OR shard_key IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 15_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(base_table.time)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__count_state(base_table.time)] REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED, "projection": ["time"], "filters": ["time >= TimestampMillisecond(0, None)", "shard_key >= Int64(60)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ diff --git a/tests/cases/distributed/explain/step_aggr_massive.sql b/tests/cases/distributed/explain/step_aggr_massive.sql index e30bf79a670..490ee255602 100644 --- a/tests/cases/distributed/explain/step_aggr_massive.sql +++ b/tests/cases/distributed/explain/step_aggr_massive.sql @@ -249,6 +249,7 @@ GROUP BY -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT env, @@ -454,6 +455,7 @@ where -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT count(*) diff --git a/tests/cases/distributed/optimizer/count.result b/tests/cases/distributed/optimizer/count.result index 7c2d419b8cf..706340006cf 100644 --- a/tests/cases/distributed/optimizer/count.result +++ b/tests/cases/distributed/optimizer/count.result @@ -89,6 +89,7 @@ Affected Rows: 5 -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (\{count\[count\]:)\d+(\}) {count[count]:REDACTED} -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select count(1) from count_where_bug; @@ -126,6 +127,7 @@ select count(1) from count_where_bug; -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select count(1) from count_where_bug where `tag` = 'b'; @@ -139,7 +141,7 @@ select count(1) from count_where_bug where `tag` = 'b'; |_|_|_AggregateExec: mode=Final, gby=[], aggr=[count(count_where_bug.ts)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[count(count_where_bug.ts)] REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["ts"], "filters": ["tag = Dictionary(UInt32, Utf8(\"b\"))"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ @@ -168,6 +170,7 @@ select count(1) from count_where_bug where `tag` = 'b'; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (\{count\[count\]:)\d+(\}) {count[count]:REDACTED} -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select count(1) from count_where_bug where ts > '2024-09-06T06:00:04Z'; @@ -181,7 +184,7 @@ select count(1) from count_where_bug where ts > '2024-09-06T06:00:04Z'; |_|_|_AggregateExec: mode=Final, gby=[], aggr=[count(count_where_bug.ts)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[count(count_where_bug.ts)] REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED, "projection": ["ts"], "filters": ["ts > TimestampMillisecond(1725602404000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ @@ -208,6 +211,7 @@ select count(1) from count_where_bug where ts > '2024-09-06T06:00:04Z'; -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select count(1) from count_where_bug where num != 3; @@ -222,7 +226,7 @@ select count(1) from count_where_bug where num != 3; |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[count(count_where_bug.ts)] REDACTED |_|_|_FilterExec: num@1 != 3, projection=[ts@0] REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["ts", "num"], "filters": ["num != Int64(3)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ @@ -272,6 +276,7 @@ Affected Rows: 5 -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (\{count\[count\]:)\d+(\}) {count[count]:REDACTED} -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select count(1) from count_where_bug; @@ -308,6 +313,7 @@ select count(1) from count_where_bug; -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select count(1) from count_where_bug where `tag` = 'b'; @@ -323,7 +329,7 @@ select count(1) from count_where_bug where `tag` = 'b'; | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(count_where_bug.ts)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__count_state(count_where_bug.ts)] REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["ts"], "filters": ["tag = Dictionary(UInt32, Utf8(\"b\"))", "tag > Utf8(\"a\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ @@ -335,6 +341,7 @@ select count(1) from count_where_bug where `tag` = 'b'; -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select count(1) from count_where_bug where `tag` = 'b'; @@ -350,7 +357,7 @@ select count(1) from count_where_bug where `tag` = 'b'; | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(count_where_bug.ts)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__count_state(count_where_bug.ts)] REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["ts"], "filters": ["tag = Dictionary(UInt32, Utf8(\"b\"))", "tag > Utf8(\"a\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ @@ -372,6 +379,7 @@ select count(1) from count_where_bug where `tag` = 'b'; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (\{count\[count]:)\d+(\}) {count[count]:REDACTED} -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select count(1) from count_where_bug where ts > '2024-09-06T06:00:04Z'; @@ -387,12 +395,12 @@ select count(1) from count_where_bug where ts > '2024-09-06T06:00:04Z'; | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(count_where_bug.ts)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__count_state(count_where_bug.ts)] REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED, "projection": ["ts"], "filters": ["ts > TimestampMillisecond(1725602404000, None)", "tag <= Utf8(\"a\") OR tag IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(count_where_bug.ts)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__count_state(count_where_bug.ts)] REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED, "projection": ["ts"], "filters": ["ts > TimestampMillisecond(1725602404000, None)", "tag > Utf8(\"a\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ @@ -412,6 +420,7 @@ select count(1) from count_where_bug where ts > '2024-09-06T06:00:04Z'; -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select count(1) from count_where_bug where num != 3; @@ -428,13 +437,13 @@ select count(1) from count_where_bug where num != 3; |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__count_state(count_where_bug.ts)] REDACTED |_|_|_FilterExec: num@1 != 3, projection=[ts@0] REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["ts", "num"], "filters": ["num != Int64(3)", "tag <= Utf8(\"a\") OR tag IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(count_where_bug.ts)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__count_state(count_where_bug.ts)] REDACTED |_|_|_FilterExec: num@1 != 3, projection=[ts@0] REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["ts", "num"], "filters": ["num != Int64(3)", "tag > Utf8(\"a\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ diff --git a/tests/cases/distributed/optimizer/first_value_advance.result b/tests/cases/distributed/optimizer/first_value_advance.result index 78aa8bcf6ab..59de599197c 100644 --- a/tests/cases/distributed/optimizer/first_value_advance.result +++ b/tests/cases/distributed/optimizer/first_value_advance.result @@ -95,6 +95,8 @@ explain select -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE (partitioning.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select first_value(host order by ts), @@ -113,7 +115,7 @@ explain analyze |_|_|_AggregateExec: mode=FinalPartitioned, gby=[host@0 as host], aggr=[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST], first_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[host@1 as host], aggr=[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST], first_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1} REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["ts", "host", "not_pk", "val"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -153,6 +155,8 @@ explain select first_value(ts order by ts) from t; -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE (partitioning.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select first_value(ts order by ts) from t; @@ -165,7 +169,7 @@ explain analyze | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[first_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[first_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1} REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["ts"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ @@ -255,6 +259,8 @@ order by ordered_host; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select first_value(host order by ts) as ordered_host, @@ -277,7 +283,7 @@ order by ordered_host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[host@0 as host], aggr=[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST], first_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[host@1 as host], aggr=[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST], first_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["ts", "host", "not_pk", "val"], "filters": ["host < Utf8(\"b\") OR host IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_ProjectionExec: expr=[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST]@0 as ordered_host, first_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST]@1 as first_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST]@2 as first_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_SortPreservingMergeExec: [first_value(t.host) ORDER BY [t.ts ASC NULLS LAST]@0 ASC NULLS LAST] REDACTED @@ -286,7 +292,7 @@ order by ordered_host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[host@0 as host], aggr=[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST], first_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[host@1 as host], aggr=[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST], first_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["ts", "host", "not_pk", "val"], "filters": ["host >= Utf8(\"b\") AND (host < Utf8(\"d\") OR host IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 2_|_ProjectionExec: expr=[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST]@0 as ordered_host, first_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST]@1 as first_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST]@2 as first_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_SortPreservingMergeExec: [first_value(t.host) ORDER BY [t.ts ASC NULLS LAST]@0 ASC NULLS LAST] REDACTED @@ -295,7 +301,7 @@ order by ordered_host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[host@0 as host], aggr=[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST], first_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[host@1 as host], aggr=[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST], first_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["ts", "host", "not_pk", "val"], "filters": ["host >= Utf8(\"d\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -340,6 +346,8 @@ explain select first_value(ts order by ts) from t; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select first_value(ts order by ts) from t; @@ -355,17 +363,17 @@ explain analyze | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__first_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__first_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["ts"], "filters": ["host < Utf8(\"b\") OR host IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=Final, gby=[], aggr=[__first_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__first_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["ts"], "filters": ["host >= Utf8(\"b\") AND (host < Utf8(\"d\") OR host IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 2_|_AggregateExec: mode=Final, gby=[], aggr=[__first_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__first_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["ts"], "filters": ["host >= Utf8(\"d\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ @@ -431,6 +439,8 @@ order by time_window, ordered_host; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select first_value(host order by ts) as ordered_host, @@ -456,17 +466,17 @@ order by time_window, ordered_host; | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[__first_value_state(t.host) ORDER BY [t.ts ASC NULLS LAST], __first_value_state(t.val) ORDER BY [t.ts ASC NULLS LAST], __first_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 5000000 }, ts@0) as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[__first_value_state(t.host) ORDER BY [t.ts ASC NULLS LAST], __first_value_state(t.val) ORDER BY [t.ts ASC NULLS LAST], __first_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["ts", "host", "val"], "filters": ["host < Utf8(\"b\") OR host IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[__first_value_state(t.host) ORDER BY [t.ts ASC NULLS LAST], __first_value_state(t.val) ORDER BY [t.ts ASC NULLS LAST], __first_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 5000000 }, ts@0) as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[__first_value_state(t.host) ORDER BY [t.ts ASC NULLS LAST], __first_value_state(t.val) ORDER BY [t.ts ASC NULLS LAST], __first_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["ts", "host", "val"], "filters": ["host >= Utf8(\"b\") AND (host < Utf8(\"d\") OR host IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 2_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[__first_value_state(t.host) ORDER BY [t.ts ASC NULLS LAST], __first_value_state(t.val) ORDER BY [t.ts ASC NULLS LAST], __first_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 5000000 }, ts@0) as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[__first_value_state(t.host) ORDER BY [t.ts ASC NULLS LAST], __first_value_state(t.val) ORDER BY [t.ts ASC NULLS LAST], __first_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["ts", "host", "val"], "filters": ["host >= Utf8(\"d\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 2_| +-+-+-+ @@ -544,6 +554,8 @@ explain -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select first_value(ts order by ts) from t1; @@ -559,17 +571,17 @@ explain analyze | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__first_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__first_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["ts"], "filters": ["__table_id=UInt32(REDACTED)", "host < Utf8(\"b\") OR host IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=Final, gby=[], aggr=[__first_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__first_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["ts"], "filters": ["__table_id=UInt32(REDACTED)", "host >= Utf8(\"b\") AND (host < Utf8(\"d\") OR host IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 2_|_AggregateExec: mode=Final, gby=[], aggr=[__first_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__first_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["ts"], "filters": ["__table_id=UInt32(REDACTED)", "host >= Utf8(\"d\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ @@ -627,6 +639,8 @@ order by ordered_host; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select first_value(host order by ts) as ordered_host, @@ -648,7 +662,7 @@ order by ordered_host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[host@0 as host], aggr=[first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[host@0 as host], aggr=[first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["host", "ts", "val"], "filters": ["__table_id=UInt32(REDACTED)", "host < Utf8(\"b\") OR host IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_ProjectionExec: expr=[first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST]@0 as ordered_host, first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]@1 as first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_SortPreservingMergeExec: [first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST]@0 ASC NULLS LAST] REDACTED @@ -657,7 +671,7 @@ order by ordered_host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[host@0 as host], aggr=[first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[host@0 as host], aggr=[first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["host", "ts", "val"], "filters": ["__table_id=UInt32(REDACTED)", "host >= Utf8(\"b\") AND (host < Utf8(\"d\") OR host IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 2_|_ProjectionExec: expr=[first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST]@0 as ordered_host, first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]@1 as first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_SortPreservingMergeExec: [first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST]@0 ASC NULLS LAST] REDACTED @@ -666,7 +680,7 @@ order by ordered_host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[host@0 as host], aggr=[first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[host@0 as host], aggr=[first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["host", "ts", "val"], "filters": ["__table_id=UInt32(REDACTED)", "host >= Utf8(\"d\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -732,6 +746,8 @@ order by time_window, ordered_host; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select first_value(host order by ts) as ordered_host, @@ -757,17 +773,17 @@ order by time_window, ordered_host; | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[__first_value_state(t1.host) ORDER BY [t1.ts ASC NULLS LAST], __first_value_state(t1.val) ORDER BY [t1.ts ASC NULLS LAST], __first_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 5000000 }, ts@1) as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[__first_value_state(t1.host) ORDER BY [t1.ts ASC NULLS LAST], __first_value_state(t1.val) ORDER BY [t1.ts ASC NULLS LAST], __first_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["host", "ts", "val"], "filters": ["__table_id=UInt32(REDACTED)", "host < Utf8(\"b\") OR host IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[__first_value_state(t1.host) ORDER BY [t1.ts ASC NULLS LAST], __first_value_state(t1.val) ORDER BY [t1.ts ASC NULLS LAST], __first_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 5000000 }, ts@1) as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[__first_value_state(t1.host) ORDER BY [t1.ts ASC NULLS LAST], __first_value_state(t1.val) ORDER BY [t1.ts ASC NULLS LAST], __first_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["host", "ts", "val"], "filters": ["__table_id=UInt32(REDACTED)", "host >= Utf8(\"b\") AND (host < Utf8(\"d\") OR host IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 2_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[__first_value_state(t1.host) ORDER BY [t1.ts ASC NULLS LAST], __first_value_state(t1.val) ORDER BY [t1.ts ASC NULLS LAST], __first_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 5000000 }, ts@1) as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[__first_value_state(t1.host) ORDER BY [t1.ts ASC NULLS LAST], __first_value_state(t1.val) ORDER BY [t1.ts ASC NULLS LAST], __first_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["host", "ts", "val"], "filters": ["__table_id=UInt32(REDACTED)", "host >= Utf8(\"d\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 2_| +-+-+-+ diff --git a/tests/cases/distributed/optimizer/last_value_advance.result b/tests/cases/distributed/optimizer/last_value_advance.result index 795f34a582e..c9c1f3887ef 100644 --- a/tests/cases/distributed/optimizer/last_value_advance.result +++ b/tests/cases/distributed/optimizer/last_value_advance.result @@ -95,6 +95,8 @@ explain select -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE (partitioning.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select last_value(host order by ts), @@ -113,7 +115,7 @@ explain analyze |_|_|_AggregateExec: mode=FinalPartitioned, gby=[host@0 as host], aggr=[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST], last_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[host@1 as host], aggr=[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST], last_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "selector":"LastRow { after_merge: false }" REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "selector":"LastRow { after_merge: false }", "projection": ["ts", "host", "not_pk", "val"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -153,6 +155,8 @@ explain select last_value(ts order by ts) from t; -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE (partitioning.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select last_value(ts order by ts) from t; @@ -165,7 +169,7 @@ explain analyze | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[last_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[last_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "selector":"LastRow { after_merge: false }" REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "selector":"LastRow { after_merge: false }", "projection": ["ts"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ @@ -255,6 +259,8 @@ order by ordered_host; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select last_value(host order by ts) as ordered_host, @@ -277,7 +283,7 @@ order by ordered_host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[host@0 as host], aggr=[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST], last_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[host@1 as host], aggr=[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST], last_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }" REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }", "projection": ["ts", "host", "not_pk", "val"], "filters": ["host < Utf8(\"b\") OR host IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_ProjectionExec: expr=[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST]@0 as ordered_host, last_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST]@1 as last_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST]@2 as last_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_SortPreservingMergeExec: [last_value(t.host) ORDER BY [t.ts ASC NULLS LAST]@0 ASC NULLS LAST] REDACTED @@ -286,7 +292,7 @@ order by ordered_host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[host@0 as host], aggr=[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST], last_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[host@1 as host], aggr=[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST], last_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }" REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }", "projection": ["ts", "host", "not_pk", "val"], "filters": ["host >= Utf8(\"b\") AND (host < Utf8(\"d\") OR host IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 2_|_ProjectionExec: expr=[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST]@0 as ordered_host, last_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST]@1 as last_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST]@2 as last_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_SortPreservingMergeExec: [last_value(t.host) ORDER BY [t.ts ASC NULLS LAST]@0 ASC NULLS LAST] REDACTED @@ -295,7 +301,7 @@ order by ordered_host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[host@0 as host], aggr=[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST], last_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[host@1 as host], aggr=[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST], last_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }" REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }", "projection": ["ts", "host", "not_pk", "val"], "filters": ["host >= Utf8(\"d\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -340,6 +346,8 @@ explain select last_value(ts order by ts) from t; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select last_value(ts order by ts) from t; @@ -355,17 +363,17 @@ explain analyze | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__last_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__last_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }" REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }", "projection": ["ts"], "filters": ["host < Utf8(\"b\") OR host IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=Final, gby=[], aggr=[__last_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__last_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }" REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }", "projection": ["ts"], "filters": ["host >= Utf8(\"b\") AND (host < Utf8(\"d\") OR host IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 2_|_AggregateExec: mode=Final, gby=[], aggr=[__last_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__last_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }" REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }", "projection": ["ts"], "filters": ["host >= Utf8(\"d\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ @@ -431,6 +439,8 @@ order by time_window, ordered_host; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select last_value(host order by ts) as ordered_host, @@ -456,17 +466,17 @@ order by time_window, ordered_host; | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[__last_value_state(t.host) ORDER BY [t.ts ASC NULLS LAST], __last_value_state(t.val) ORDER BY [t.ts ASC NULLS LAST], __last_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 5000000 }, ts@0) as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[__last_value_state(t.host) ORDER BY [t.ts ASC NULLS LAST], __last_value_state(t.val) ORDER BY [t.ts ASC NULLS LAST], __last_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["ts", "host", "val"], "filters": ["host < Utf8(\"b\") OR host IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[__last_value_state(t.host) ORDER BY [t.ts ASC NULLS LAST], __last_value_state(t.val) ORDER BY [t.ts ASC NULLS LAST], __last_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 5000000 }, ts@0) as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[__last_value_state(t.host) ORDER BY [t.ts ASC NULLS LAST], __last_value_state(t.val) ORDER BY [t.ts ASC NULLS LAST], __last_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["ts", "host", "val"], "filters": ["host >= Utf8(\"b\") AND (host < Utf8(\"d\") OR host IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 2_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[__last_value_state(t.host) ORDER BY [t.ts ASC NULLS LAST], __last_value_state(t.val) ORDER BY [t.ts ASC NULLS LAST], __last_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 5000000 }, ts@0) as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[__last_value_state(t.host) ORDER BY [t.ts ASC NULLS LAST], __last_value_state(t.val) ORDER BY [t.ts ASC NULLS LAST], __last_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["ts", "host", "val"], "filters": ["host >= Utf8(\"d\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 2_| +-+-+-+ @@ -544,6 +554,8 @@ explain -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select last_value(ts order by ts) from t1; @@ -559,17 +571,17 @@ explain analyze | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__last_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__last_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }" REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }", "projection": ["ts"], "filters": ["__table_id=UInt32(REDACTED)", "host < Utf8(\"b\") OR host IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=Final, gby=[], aggr=[__last_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__last_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }" REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }", "projection": ["ts"], "filters": ["__table_id=UInt32(REDACTED)", "host >= Utf8(\"b\") AND (host < Utf8(\"d\") OR host IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 2_|_AggregateExec: mode=Final, gby=[], aggr=[__last_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__last_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }" REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }", "projection": ["ts"], "filters": ["__table_id=UInt32(REDACTED)", "host >= Utf8(\"d\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ @@ -627,6 +639,8 @@ order by ordered_host; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select last_value(host order by ts) as ordered_host, @@ -648,7 +662,7 @@ order by ordered_host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[host@0 as host], aggr=[last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[host@0 as host], aggr=[last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }" REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }", "projection": ["host", "ts", "val"], "filters": ["__table_id=UInt32(REDACTED)", "host < Utf8(\"b\") OR host IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_ProjectionExec: expr=[last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST]@0 as ordered_host, last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]@1 as last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_SortPreservingMergeExec: [last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST]@0 ASC NULLS LAST] REDACTED @@ -657,7 +671,7 @@ order by ordered_host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[host@0 as host], aggr=[last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[host@0 as host], aggr=[last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }" REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }", "projection": ["host", "ts", "val"], "filters": ["__table_id=UInt32(REDACTED)", "host >= Utf8(\"b\") AND (host < Utf8(\"d\") OR host IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 2_|_ProjectionExec: expr=[last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST]@0 as ordered_host, last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]@1 as last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_SortPreservingMergeExec: [last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST]@0 ASC NULLS LAST] REDACTED @@ -666,7 +680,7 @@ order by ordered_host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[host@0 as host], aggr=[last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[host@0 as host], aggr=[last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }" REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }", "projection": ["host", "ts", "val"], "filters": ["__table_id=UInt32(REDACTED)", "host >= Utf8(\"d\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -732,6 +746,8 @@ order by time_window, ordered_host; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select last_value(host order by ts) as ordered_host, @@ -757,17 +773,17 @@ order by time_window, ordered_host; | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[__last_value_state(t1.host) ORDER BY [t1.ts ASC NULLS LAST], __last_value_state(t1.val) ORDER BY [t1.ts ASC NULLS LAST], __last_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 5000000 }, ts@1) as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[__last_value_state(t1.host) ORDER BY [t1.ts ASC NULLS LAST], __last_value_state(t1.val) ORDER BY [t1.ts ASC NULLS LAST], __last_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["host", "ts", "val"], "filters": ["__table_id=UInt32(REDACTED)", "host < Utf8(\"b\") OR host IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[__last_value_state(t1.host) ORDER BY [t1.ts ASC NULLS LAST], __last_value_state(t1.val) ORDER BY [t1.ts ASC NULLS LAST], __last_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 5000000 }, ts@1) as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[__last_value_state(t1.host) ORDER BY [t1.ts ASC NULLS LAST], __last_value_state(t1.val) ORDER BY [t1.ts ASC NULLS LAST], __last_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["host", "ts", "val"], "filters": ["__table_id=UInt32(REDACTED)", "host >= Utf8(\"b\") AND (host < Utf8(\"d\") OR host IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 2_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[__last_value_state(t1.host) ORDER BY [t1.ts ASC NULLS LAST], __last_value_state(t1.val) ORDER BY [t1.ts ASC NULLS LAST], __last_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 5000000 }, ts@1) as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[__last_value_state(t1.host) ORDER BY [t1.ts ASC NULLS LAST], __last_value_state(t1.val) ORDER BY [t1.ts ASC NULLS LAST], __last_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["host", "ts", "val"], "filters": ["__table_id=UInt32(REDACTED)", "host >= Utf8(\"d\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 2_| +-+-+-+ diff --git a/tests/cases/distributed/optimizer/range_select_projection.result b/tests/cases/distributed/optimizer/range_select_projection.result index 28d8fff50a3..7541d929073 100644 --- a/tests/cases/distributed/optimizer/range_select_projection.result +++ b/tests/cases/distributed/optimizer/range_select_projection.result @@ -29,6 +29,7 @@ Affected Rows: 5 -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED -- SQLNESS REPLACE Dictionary\(UInt32,\s*Utf8\(([^)]*)\)\) Utf8($1) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT ts, station, @@ -52,7 +53,7 @@ ORDER BY station, "channel", ts; |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_CooperativeExec REDACTED -|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["ts", "station", "channel", "value_a", "value_b"], "filters": ["station = Utf8(\"station-a\")", "unused_tag = Utf8(\"unused-a\")", "ts >= TimestampMillisecond(0, None)", "ts < TimestampMillisecond(15000, None)"], "flat_format": true, "REDACTED +|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["ts", "station", "channel", "value_a", "value_b"], "filters": ["station = Utf8(\"station-a\")", "unused_tag = Utf8(\"unused-a\")", "ts >= TimestampMillisecond(0, None)", "ts < TimestampMillisecond(15000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: REDACTED_| +-+-+-+ diff --git a/tests/cases/distributed/optimizer/range_select_projection.sql b/tests/cases/distributed/optimizer/range_select_projection.sql index 4b15a30c675..5d82804aafe 100644 --- a/tests/cases/distributed/optimizer/range_select_projection.sql +++ b/tests/cases/distributed/optimizer/range_select_projection.sql @@ -25,6 +25,7 @@ INSERT INTO range_select_projection VALUES -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED -- SQLNESS REPLACE Dictionary\(UInt32,\s*Utf8\(([^)]*)\)\) Utf8($1) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT ts, station, diff --git a/tests/cases/distributed/optimizer/windowed_sort.result b/tests/cases/distributed/optimizer/windowed_sort.result index 0c4d1931c05..6c2944f866c 100644 --- a/tests/cases/distributed/optimizer/windowed_sort.result +++ b/tests/cases/distributed/optimizer/windowed_sort.result @@ -56,6 +56,7 @@ ORDER BY -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT collect_time_utc AS collect_time, peak_current, @@ -73,7 +74,7 @@ ORDER BY | 1_| 0_|_ProjectionExec: expr=[collect_time_utc@0 as collect_time, peak_current@1 as peak_current] REDACTED |_|_|_SortPreservingMergeExec: [collect_time_utc@0 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[collect_time_utc@0 ASC NULLS LAST], preserve_partitioning=[true] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["collect_time_utc", "peak_current"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 8_| +-+-+-+ @@ -106,6 +107,7 @@ ORDER BY -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT collect_time_utc AS collect_time_0, peak_current, @@ -123,7 +125,7 @@ ORDER BY | 1_| 0_|_ProjectionExec: expr=[collect_time_utc@0 as collect_time_0, peak_current@1 as peak_current] REDACTED |_|_|_SortPreservingMergeExec: [collect_time_utc@0 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[collect_time_utc@0 ASC NULLS LAST], preserve_partitioning=[true] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["collect_time_utc", "peak_current"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 8_| +-+-+-+ @@ -157,6 +159,7 @@ ORDER BY -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT collect_time AS true_collect_time, collect_time_utc AS collect_time, @@ -176,7 +179,7 @@ ORDER BY |_|_|_SortPreservingMergeExec: [collect_time@0 DESC] REDACTED |_|_|_WindowedSortExec: expr=collect_time@0 DESC num_ranges=REDACTED REDACTED |_|_|_PartSortExec: expr=collect_time@0 DESC num_ranges=REDACTED REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["collect_time", "collect_time_utc", "peak_current"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 8_| +-+-+-+ @@ -210,6 +213,7 @@ ORDER BY -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT collect_time_utc AS collect_time, collect_time AS true_collect_time, @@ -230,7 +234,7 @@ ORDER BY |_|_|_ProjectionExec: expr=[collect_time_utc@1 as collect_time_utc, collect_time@0 as collect_time, peak_current@2 as peak_current] REDACTED |_|_|_WindowedSortExec: expr=collect_time@0 DESC num_ranges=REDACTED REDACTED |_|_|_PartSortExec: expr=collect_time@0 DESC num_ranges=REDACTED REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["collect_time", "collect_time_utc", "peak_current"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 8_| +-+-+-+ diff --git a/tests/cases/distributed/optimizer/windowed_sort.sql b/tests/cases/distributed/optimizer/windowed_sort.sql index eefdacd51a1..ffdfde853e2 100644 --- a/tests/cases/distributed/optimizer/windowed_sort.sql +++ b/tests/cases/distributed/optimizer/windowed_sort.sql @@ -39,6 +39,7 @@ ORDER BY -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT collect_time_utc AS collect_time, peak_current, @@ -62,6 +63,7 @@ ORDER BY -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT collect_time_utc AS collect_time_0, peak_current, @@ -86,6 +88,7 @@ ORDER BY -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT collect_time AS true_collect_time, collect_time_utc AS collect_time, @@ -111,6 +114,7 @@ ORDER BY -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT collect_time_utc AS collect_time, collect_time AS true_collect_time, diff --git a/tests/cases/distributed/optimizer/windowed_sort_advance.result b/tests/cases/distributed/optimizer/windowed_sort_advance.result index 1b122e61c8d..3c695f3879c 100644 --- a/tests/cases/distributed/optimizer/windowed_sort_advance.result +++ b/tests/cases/distributed/optimizer/windowed_sort_advance.result @@ -56,6 +56,7 @@ select ts as ts, status, value from `a` where ts >= '2026-03-12T08:00:00+08:00' -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE select ts as ts, status, value from `a` where ts >= '2026-03-12T08:00:00+08:00' and ts < '2026-03-12T08:02:01+08:00' order by ts asc; +-+-+-+ @@ -67,7 +68,7 @@ EXPLAIN ANALYZE select ts as ts, status, value from `a` where ts >= '2026-03-12T | 1_| 0_|_SortPreservingMergeExec: [ts@0 ASC NULLS LAST] REDACTED |_|_|_ProjectionExec: expr=[ts@2 as ts, status@1 as status, value@0 as value] REDACTED |_|_|_WindowedSortExec: expr=ts@2 ASC NULLS LAST num_ranges=REDACTED REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["value", "status", "ts"], "filters": ["ts >= TimestampNanosecond(1773273600000000000, None)", "ts < TimestampNanosecond(1773273721000000000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 10_| +-+-+-+ @@ -95,6 +96,7 @@ select to_timestamp_millis(ts) as ts, status, value from `a` where ts >= '2026-0 -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE select to_timestamp_millis(ts) as ts, status, value from `a` where ts >= '2026-03-12T08:00:00+08:00' and ts < '2026-03-12T08:02:01+08:00' order by ts asc; +-+-+-+ @@ -107,7 +109,7 @@ EXPLAIN ANALYZE select to_timestamp_millis(ts) as ts, status, value from `a` whe |_|_|_SortPreservingMergeExec: [to_timestamp_millis(a.ts)@0 ASC NULLS LAST] REDACTED |_|_|_WindowedSortExec: expr=to_timestamp_millis(a.ts)@0 ASC NULLS LAST num_ranges=REDACTED REDACTED |_|_|_ProjectionExec: expr=[to_timestamp_millis(ts@2) as to_timestamp_millis(a.ts), status@1 as status, value@0 as value] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["value", "status", "ts"], "filters": ["ts >= TimestampNanosecond(1773273600000000000, None)", "ts < TimestampNanosecond(1773273721000000000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 10_| +-+-+-+ diff --git a/tests/cases/distributed/optimizer/windowed_sort_advance.sql b/tests/cases/distributed/optimizer/windowed_sort_advance.sql index fa73b97933c..167a42d591c 100644 --- a/tests/cases/distributed/optimizer/windowed_sort_advance.sql +++ b/tests/cases/distributed/optimizer/windowed_sort_advance.sql @@ -22,6 +22,7 @@ select ts as ts, status, value from `a` where ts >= '2026-03-12T08:00:00+08:00' -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE select ts as ts, status, value from `a` where ts >= '2026-03-12T08:00:00+08:00' and ts < '2026-03-12T08:02:01+08:00' order by ts asc; select to_timestamp_millis(ts) as ts, status, value from `a` where ts >= '2026-03-12T08:00:00+08:00' and ts < '2026-03-12T08:02:01+08:00' order by ts asc; @@ -32,6 +33,7 @@ select to_timestamp_millis(ts) as ts, status, value from `a` where ts >= '2026-0 -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE select to_timestamp_millis(ts) as ts, status, value from `a` where ts >= '2026-03-12T08:00:00+08:00' and ts < '2026-03-12T08:02:01+08:00' order by ts asc; DROP TABLE `a`; diff --git a/tests/cases/distributed/optimizer/windowed_sort_nanos.result b/tests/cases/distributed/optimizer/windowed_sort_nanos.result index 2310ec8f284..52d264073ae 100644 --- a/tests/cases/distributed/optimizer/windowed_sort_nanos.result +++ b/tests/cases/distributed/optimizer/windowed_sort_nanos.result @@ -56,6 +56,7 @@ select to_timestamp_nanos(ts) as ts, status, value from `a_ms` where ts >= '2026 -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE select to_timestamp_nanos(ts) as ts, status, value from `a_ms` where ts >= '2026-03-12T08:00:00.000+08:00' and ts < '2026-03-12T08:02:01.000+08:00' order by ts asc; +-+-+-+ @@ -68,7 +69,7 @@ EXPLAIN ANALYZE select to_timestamp_nanos(ts) as ts, status, value from `a_ms` w |_|_|_SortPreservingMergeExec: [to_timestamp_nanos(a_ms.ts)@0 ASC NULLS LAST] REDACTED |_|_|_WindowedSortExec: expr=to_timestamp_nanos(a_ms.ts)@0 ASC NULLS LAST num_ranges=REDACTED REDACTED |_|_|_ProjectionExec: expr=[to_timestamp_nanos(ts@2) as to_timestamp_nanos(a_ms.ts), status@1 as status, value@0 as value] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["value", "status", "ts"], "filters": ["ts >= TimestampMillisecond(1773273600000, None)", "ts < TimestampMillisecond(1773273721000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 10_| +-+-+-+ diff --git a/tests/cases/distributed/optimizer/windowed_sort_nanos.sql b/tests/cases/distributed/optimizer/windowed_sort_nanos.sql index 6a6cf51c4b6..27abd0bc805 100644 --- a/tests/cases/distributed/optimizer/windowed_sort_nanos.sql +++ b/tests/cases/distributed/optimizer/windowed_sort_nanos.sql @@ -22,6 +22,7 @@ select to_timestamp_nanos(ts) as ts, status, value from `a_ms` where ts >= '2026 -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE select to_timestamp_nanos(ts) as ts, status, value from `a_ms` where ts >= '2026-03-12T08:00:00.000+08:00' and ts < '2026-03-12T08:02:01.000+08:00' order by ts asc; DROP TABLE `a_ms`; diff --git a/tests/cases/distributed/tql-explain-analyze/analyze.result b/tests/cases/distributed/tql-explain-analyze/analyze.result index 059cf55df64..695781e2c0f 100644 --- a/tests/cases/distributed/tql-explain-analyze/analyze.result +++ b/tests/cases/distributed/tql-explain-analyze/analyze.result @@ -15,6 +15,7 @@ Affected Rows: 3 -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 10, '5s') test; +-+-+-+ @@ -26,7 +27,7 @@ TQL ANALYZE (0, 10, '5s') test; | 1_| 0_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[j] REDACTED |_|_|_PromSeriesDivideExec: tags=["k"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["i", "j", "k"], "filters": ["j >= TimestampMillisecond(-299999, None)", "j <= TimestampMillisecond(10000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -40,6 +41,7 @@ TQL ANALYZE (0, 10, '5s') test; -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 10, '1s', '2s') test; +-+-+-+ @@ -51,7 +53,7 @@ TQL ANALYZE (0, 10, '1s', '2s') test; | 1_| 0_|_PromInstantManipulateExec: range=[0..10000], lookback=[2000], interval=[1000], time index=[j] REDACTED |_|_|_PromSeriesDivideExec: tags=["k"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["i", "j", "k"], "filters": ["j >= TimestampMillisecond(-1999, None)", "j <= TimestampMillisecond(10000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -64,6 +66,7 @@ TQL ANALYZE (0, 10, '1s', '2s') test; -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE ('1970-01-01T00:00:00'::timestamp, '1970-01-01T00:00:00'::timestamp + '10 seconds'::interval, '5s') test; +-+-+-+ @@ -75,7 +78,7 @@ TQL ANALYZE ('1970-01-01T00:00:00'::timestamp, '1970-01-01T00:00:00'::timestamp | 1_| 0_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[j] REDACTED |_|_|_PromSeriesDivideExec: tags=["k"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["i", "j", "k"], "filters": ["j >= TimestampMillisecond(-299999, None)", "j <= TimestampMillisecond(10000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -91,6 +94,7 @@ TQL ANALYZE ('1970-01-01T00:00:00'::timestamp, '1970-01-01T00:00:00'::timestamp -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED -- SQLNESS REPLACE (flat_format.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE VERBOSE (0, 10, '5s') test; +-+-+-+ @@ -102,7 +106,7 @@ TQL ANALYZE VERBOSE (0, 10, '5s') test; | 1_| 0_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[j] REDACTED |_|_|_PromSeriesDivideExec: tags=["k"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["i", "j", "k"], "filters": ["j >= TimestampMillisecond(-299999, None)", "j <= TimestampMillisecond(10000, None)"], "REDACTED +|_|_|_SeriesScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["i", "j", "k"], "filters": ["j >= TimestampMillisecond(-299999, None)", "j <= TimestampMillisecond(10000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -123,6 +127,7 @@ Affected Rows: 0 -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 10, '5s') test; +-+-+-+ @@ -134,12 +139,12 @@ TQL ANALYZE (0, 10, '5s') test; | 1_| 0_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[j] REDACTED |_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["i", "j", "k", "l"], "filters": ["j >= TimestampMillisecond(-299999, None)", "j <= TimestampMillisecond(10000, None)", "k < Utf8(\"a\") OR k IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[j] REDACTED |_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["i", "j", "k", "l"], "filters": ["j >= TimestampMillisecond(-299999, None)", "j <= TimestampMillisecond(10000, None)", "k >= Utf8(\"a\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -151,6 +156,7 @@ TQL ANALYZE (0, 10, '5s') test; -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 10, '5s') rate(test[10s]); +-+-+-+ @@ -166,7 +172,7 @@ TQL ANALYZE (0, 10, '5s') rate(test[10s]); |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[j], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["i", "j", "k", "l"], "filters": ["j >= TimestampMillisecond(-9999, None)", "j <= TimestampMillisecond(10000, None)", "k < Utf8(\"a\") OR k IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_ProjectionExec: expr=[j@0 as j, prom_rate(j_range,i,test.j,Int64(10000))@1 as prom_rate(j_range,i,j,Int64(10000)), k@2 as k, l@3 as l] REDACTED |_|_|_FilterExec: prom_rate(j_range,i,test.j,Int64(10000))@1 IS NOT NULL REDACTED @@ -175,7 +181,7 @@ TQL ANALYZE (0, 10, '5s') rate(test[10s]); |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[j], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["i", "j", "k", "l"], "filters": ["j >= TimestampMillisecond(-9999, None)", "j <= TimestampMillisecond(10000, None)", "k >= Utf8(\"a\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ diff --git a/tests/cases/distributed/tql-explain-analyze/analyze.sql b/tests/cases/distributed/tql-explain-analyze/analyze.sql index e8bdf92838a..7dd4ce70f1d 100644 --- a/tests/cases/distributed/tql-explain-analyze/analyze.sql +++ b/tests/cases/distributed/tql-explain-analyze/analyze.sql @@ -11,6 +11,7 @@ INSERT INTO test VALUES (1, 1, "a"), (1, 1, "b"), (2, 2, "a"); -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 10, '5s') test; -- 'lookback' parameter is not fully supported, the test has to be updated @@ -22,6 +23,7 @@ TQL ANALYZE (0, 10, '5s') test; -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 10, '1s', '2s') test; -- analyze at 0s, 5s and 10s. No point at 0s. @@ -32,6 +34,7 @@ TQL ANALYZE (0, 10, '1s', '2s') test; -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE ('1970-01-01T00:00:00'::timestamp, '1970-01-01T00:00:00'::timestamp + '10 seconds'::interval, '5s') test; -- analyze verbose at 0s, 5s and 10s. No point at 0s. @@ -45,6 +48,7 @@ TQL ANALYZE ('1970-01-01T00:00:00'::timestamp, '1970-01-01T00:00:00'::timestamp -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED -- SQLNESS REPLACE (flat_format.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE VERBOSE (0, 10, '5s') test; DROP TABLE test; @@ -59,6 +63,7 @@ CREATE TABLE test(i DOUBLE, j TIMESTAMP TIME INDEX, k STRING, l STRING, PRIMARY -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 10, '5s') test; -- SQLNESS REPLACE (metrics.*) REDACTED @@ -68,6 +73,7 @@ TQL ANALYZE (0, 10, '5s') test; -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 10, '5s') rate(test[10s]); drop table test; diff --git a/tests/cases/standalone/common/aggregate/distinct.result b/tests/cases/standalone/common/aggregate/distinct.result index 5c10ee22239..da377f8562d 100644 --- a/tests/cases/standalone/common/aggregate/distinct.result +++ b/tests/cases/standalone/common/aggregate/distinct.result @@ -85,6 +85,7 @@ SELECT DISTINCT ON (a) * FROM test ORDER BY a, t DESC; -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT DISTINCT a FROM test ORDER BY a; +-+-+-+ @@ -98,7 +99,7 @@ EXPLAIN ANALYZE SELECT DISTINCT a FROM test ORDER BY a; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["a"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 2_| +-+-+-+ @@ -110,6 +111,7 @@ EXPLAIN ANALYZE SELECT DISTINCT a FROM test ORDER BY a; -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT DISTINCT a, b FROM test ORDER BY a; +-+-+-+ @@ -123,7 +125,7 @@ EXPLAIN ANALYZE SELECT DISTINCT a, b FROM test ORDER BY a; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b], aggr=[] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[a@0 as a, b@1 as b], aggr=[] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["a", "b"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 3_| +-+-+-+ @@ -162,6 +164,7 @@ SELECT DISTINCT a, b FROM test ORDER BY a, b; -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT DISTINCT a FROM test ORDER BY a; +-+-+-+ @@ -178,12 +181,12 @@ EXPLAIN ANALYZE SELECT DISTINCT a FROM test ORDER BY a; | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["a"], "filters": ["a <= Int32(10) OR a IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["a"], "filters": ["a > Int32(10)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 2_| +-+-+-+ @@ -195,6 +198,7 @@ EXPLAIN ANALYZE SELECT DISTINCT a FROM test ORDER BY a; -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT DISTINCT a, b FROM test ORDER BY a; +-+-+-+ @@ -211,12 +215,12 @@ EXPLAIN ANALYZE SELECT DISTINCT a, b FROM test ORDER BY a; | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b], aggr=[] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[a@0 as a, b@1 as b], aggr=[] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["a", "b"], "filters": ["a <= Int32(10) OR a IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b], aggr=[] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[a@0 as a, b@1 as b], aggr=[] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["a", "b"], "filters": ["a > Int32(10)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ diff --git a/tests/cases/standalone/common/aggregate/distinct.sql b/tests/cases/standalone/common/aggregate/distinct.sql index cdf9de5e6d1..adbabd694ab 100644 --- a/tests/cases/standalone/common/aggregate/distinct.sql +++ b/tests/cases/standalone/common/aggregate/distinct.sql @@ -25,6 +25,7 @@ SELECT DISTINCT ON (a) * FROM test ORDER BY a, t DESC; -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT DISTINCT a FROM test ORDER BY a; -- SQLNESS REPLACE (metrics.*) REDACTED @@ -34,6 +35,7 @@ EXPLAIN ANALYZE SELECT DISTINCT a FROM test ORDER BY a; -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT DISTINCT a, b FROM test ORDER BY a; DROP TABLE test; @@ -55,6 +57,7 @@ SELECT DISTINCT a, b FROM test ORDER BY a, b; -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT DISTINCT a FROM test ORDER BY a; -- SQLNESS REPLACE (metrics.*) REDACTED @@ -64,6 +67,7 @@ EXPLAIN ANALYZE SELECT DISTINCT a FROM test ORDER BY a; -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT DISTINCT a, b FROM test ORDER BY a; DROP TABLE test; diff --git a/tests/cases/standalone/common/aggregate/max_abs_regression.result b/tests/cases/standalone/common/aggregate/max_abs_regression.result index a0e0b394667..70f3dc06f60 100644 --- a/tests/cases/standalone/common/aggregate/max_abs_regression.result +++ b/tests/cases/standalone/common/aggregate/max_abs_regression.result @@ -67,6 +67,7 @@ SELECT MAX(v), MAX(ts) FROM t; -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE ,\s"dyn_filters":\s\["DynamicFilter\s\[[^"]*\]"\] , "dyn_filters": ["DynamicFilter [ REDACTED ]"] -- SQLNESS REPLACE metrics=REDACTED\s*\| metrics=REDACTED_| +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT MAX(ABS(v)), MAX(ts) FROM t; +-+-+-+ @@ -78,7 +79,7 @@ EXPLAIN ANALYZE VERBOSE SELECT MAX(ABS(v)), MAX(ts) FROM t; | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[max(abs(t.v)), max(t.ts)] metrics=REDACTED_| |_|_|_CoalescePartitionsExec metrics=REDACTED_| |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[max(abs(t.v)), max(t.ts)] metrics=REDACTED_| -|_|_|_SeqScan: region=REDACTED, {"partition_count":REDACTED, "projection": ["ts", "v"], "flat_format":REDACTED, "metrics_per_partition": REDACTED metrics=REDACTED_| +|_|_|_SeqScan: region=REDACTED, {"partition_count":REDACTED, "projection": ["ts", "v"], "time_range":REDACTED |_|_|_| |_|_| Total rows: REDACTED_| +-+-+-+ @@ -101,6 +102,7 @@ EXPLAIN ANALYZE VERBOSE SELECT MAX(ABS(v)), MAX(ts) FROM t; -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE ,\s"dyn_filters":\s\["DynamicFilter\s\[[^"]*\]"\] , "dyn_filters": ["DynamicFilter [ REDACTED ]"] -- SQLNESS REPLACE metrics=REDACTED\s*\| metrics=REDACTED_| +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT MAX(v), MAX(ts) FROM t; +-+-+-+ @@ -112,7 +114,7 @@ EXPLAIN ANALYZE VERBOSE SELECT MAX(v), MAX(ts) FROM t; | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[max(t.v), max(t.ts)] metrics=REDACTED_| |_|_|_CoalescePartitionsExec metrics=REDACTED_| |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[max(t.v), max(t.ts)] metrics=REDACTED_| -|_|_|_SeqScan: region=REDACTED, {"partition_count":REDACTED, "projection": ["ts", "v"], "dyn_filters": ["DynamicFilter [ REDACTED ]"], "flat_format":REDACTED, "metrics_per_partition": REDACTED metrics=REDACTED_| +|_|_|_SeqScan: region=REDACTED, {"partition_count":REDACTED, "projection": ["ts", "v"], "dyn_filters":REDACTED |_|_|_| |_|_| Total rows: REDACTED_| +-+-+-+ diff --git a/tests/cases/standalone/common/aggregate/max_abs_regression.sql b/tests/cases/standalone/common/aggregate/max_abs_regression.sql index c9705f6086e..76a862eaa34 100644 --- a/tests/cases/standalone/common/aggregate/max_abs_regression.sql +++ b/tests/cases/standalone/common/aggregate/max_abs_regression.sql @@ -40,6 +40,7 @@ SELECT MAX(v), MAX(ts) FROM t; -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE ,\s"dyn_filters":\s\["DynamicFilter\s\[[^"]*\]"\] , "dyn_filters": ["DynamicFilter [ REDACTED ]"] -- SQLNESS REPLACE metrics=REDACTED\s*\| metrics=REDACTED_| +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT MAX(ABS(v)), MAX(ts) FROM t; -- A column-only multiple-MAX query remains a dynamic-filter positive control. @@ -60,6 +61,7 @@ EXPLAIN ANALYZE VERBOSE SELECT MAX(ABS(v)), MAX(ts) FROM t; -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE ,\s"dyn_filters":\s\["DynamicFilter\s\[[^"]*\]"\] , "dyn_filters": ["DynamicFilter [ REDACTED ]"] -- SQLNESS REPLACE metrics=REDACTED\s*\| metrics=REDACTED_| +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT MAX(v), MAX(ts) FROM t; DROP TABLE t; diff --git a/tests/cases/standalone/common/aggregate/multi_regions.result b/tests/cases/standalone/common/aggregate/multi_regions.result index bdfc5ed2dd1..7a212c5be91 100644 --- a/tests/cases/standalone/common/aggregate/multi_regions.result +++ b/tests/cases/standalone/common/aggregate/multi_regions.result @@ -21,6 +21,7 @@ Affected Rows: 0 -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select sum(val) from t group by host; @@ -34,13 +35,13 @@ select sum(val) from t group by host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[host@0 as host], aggr=[sum(t.val)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[host@1 as host], aggr=[sum(t.val)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["val", "host"], "filters": ["host < Utf8(\"1024\") OR host IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_ProjectionExec: expr=[sum(t.val)@1 as sum(t.val)] REDACTED |_|_|_AggregateExec: mode=FinalPartitioned, gby=[host@0 as host], aggr=[sum(t.val)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[host@1 as host], aggr=[sum(t.val)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["val", "host"], "filters": ["host >= Utf8(\"1024\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -53,6 +54,7 @@ select sum(val) from t group by host; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select sum(val) from t; @@ -68,12 +70,12 @@ select sum(val) from t; | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__sum_state(t.val)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__sum_state(t.val)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["val"], "filters": ["host < Utf8(\"1024\") OR host IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=Final, gby=[], aggr=[__sum_state(t.val)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__sum_state(t.val)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["val"], "filters": ["host >= Utf8(\"1024\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ @@ -87,6 +89,7 @@ select sum(val) from t; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select sum(val) from t group by idc; @@ -103,12 +106,12 @@ select sum(val) from t group by idc; | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[idc@0 as idc], aggr=[__sum_state(t.val)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[idc@1 as idc], aggr=[__sum_state(t.val)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["val", "idc"], "filters": ["host < Utf8(\"1024\") OR host IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[idc@0 as idc], aggr=[__sum_state(t.val)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[idc@1 as idc], aggr=[__sum_state(t.val)] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["val", "idc"], "filters": ["host >= Utf8(\"1024\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ diff --git a/tests/cases/standalone/common/aggregate/multi_regions.sql b/tests/cases/standalone/common/aggregate/multi_regions.sql index 9c8d8af4490..db6b9c02392 100644 --- a/tests/cases/standalone/common/aggregate/multi_regions.sql +++ b/tests/cases/standalone/common/aggregate/multi_regions.sql @@ -19,6 +19,7 @@ partition on columns (host) ( -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select sum(val) from t group by host; @@ -30,6 +31,7 @@ select sum(val) from t group by host; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select sum(val) from t; @@ -42,6 +44,7 @@ select sum(val) from t; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select sum(val) from t group by idc; diff --git a/tests/cases/standalone/common/filter/hash_join_dyn_filter.result b/tests/cases/standalone/common/filter/hash_join_dyn_filter.result index b1402bd4d4a..e534f55791c 100644 --- a/tests/cases/standalone/common/filter/hash_join_dyn_filter.result +++ b/tests/cases/standalone/common/filter/hash_join_dyn_filter.result @@ -78,6 +78,7 @@ WHERE c.tier = 'gold'; -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE ,\s"dyn_filters":\s\["DynamicFilter\s\[[^"]*\]"\] , "dyn_filters": ["DynamicFilter [ REDACTED ]"] -- SQLNESS REPLACE metrics=REDACTED\s*\| metrics=REDACTED_| +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT o."id", o.amount, c."name", c.tier FROM orders o JOIN customers c ON o.customer_id = c.customer_id @@ -96,11 +97,11 @@ WHERE c.tier = 'gold'; |_|_|_| | 1_| 0_|_FilterExec: customer_id@1 IS NOT NULL metrics=REDACTED_| |_|_|_CooperativeExec metrics=REDACTED_| -|_|_|_SeqScan: region=REDACTED, {"partition_count":REDACTED, "projection": ["id", "customer_id", "amount", "ts"], "filters": ["customer_id IS NOT NULL"], "flat_format":REDACTED, "metrics_per_partition": REDACTED metrics=REDACTED_| +|_|_|_SeqScan: region=REDACTED, {"partition_count":REDACTED, "projection": ["id", "customer_id", "amount", "ts"], "filters": ["customer_id IS NOT NULL"], "time_range":REDACTED |_|_|_| | 1_| 0_|_FilterExec: customer_id@0 IS NOT NULL AND tier@2 = gold metrics=REDACTED_| |_|_|_CooperativeExec metrics=REDACTED_| -|_|_|_SeqScan: region=REDACTED, {"partition_count":REDACTED, "projection": ["customer_id", "name", "tier", "ts"], "filters": ["customer_id IS NOT NULL", "tier = Utf8(\"gold\")"], "dyn_filters": ["DynamicFilter [ REDACTED ]"], "flat_format":REDACTED, "metrics_per_partition": REDACTED metrics=REDACTED_| +|_|_|_SeqScan: region=REDACTED, {"partition_count":REDACTED, "projection": ["customer_id", "name", "tier", "ts"], "filters": ["customer_id IS NOT NULL", "tier = Utf8(\"gold\")"], "dyn_filters":REDACTED |_|_|_| |_|_| Total rows: REDACTED_| +-+-+-+ @@ -206,6 +207,7 @@ WHERE c.tier IN ('gold', 'silver'); -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE ,\s"dyn_filters":\s\["DynamicFilter\s\[[^"]*\]"\] , "dyn_filters": ["DynamicFilter [ REDACTED ]"] -- SQLNESS REPLACE metrics=REDACTED\s*\| metrics=REDACTED_| +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT o."id", o.amount, c."name", c.tier FROM (SELECT "id", customer_id as cid, amount, ts FROM orders) o JOIN customers c ON o.cid = c.customer_id @@ -225,11 +227,11 @@ WHERE c.tier IN ('gold', 'silver'); | 1_| 0_|_ProjectionExec: expr=[id@0 as id, customer_id@1 as cid, amount@2 as amount, ts@3 as ts] metrics=REDACTED_| |_|_|_FilterExec: customer_id@1 IS NOT NULL metrics=REDACTED_| |_|_|_CooperativeExec metrics=REDACTED_| -|_|_|_SeqScan: region=REDACTED, {"partition_count":REDACTED, "projection": ["id", "customer_id", "amount", "ts"], "filters": ["customer_id IS NOT NULL"], "flat_format":REDACTED, "metrics_per_partition": REDACTED metrics=REDACTED_| +|_|_|_SeqScan: region=REDACTED, {"partition_count":REDACTED, "projection": ["id", "customer_id", "amount", "ts"], "filters": ["customer_id IS NOT NULL"], "time_range":REDACTED |_|_|_| | 1_| 0_|_FilterExec: customer_id@0 IS NOT NULL AND (tier@2 = gold OR tier@2 = silver) metrics=REDACTED_| |_|_|_CooperativeExec metrics=REDACTED_| -|_|_|_SeqScan: region=REDACTED, {"partition_count":REDACTED, "projection": ["customer_id", "name", "tier", "ts"], "filters": ["customer_id IS NOT NULL", "tier = Utf8(\"gold\") OR tier = Utf8(\"silver\")"], "dyn_filters": ["DynamicFilter [ REDACTED ]"], "flat_format":REDACTED, "metrics_per_partition": REDACTED metrics=REDACTED_| +|_|_|_SeqScan: region=REDACTED, {"partition_count":REDACTED, "projection": ["customer_id", "name", "tier", "ts"], "filters": ["customer_id IS NOT NULL", "tier = Utf8(\"gold\") OR tier = Utf8(\"silver\")"], "dyn_filters":REDACTED |_|_|_| |_|_| Total rows: REDACTED_| +-+-+-+ diff --git a/tests/cases/standalone/common/filter/hash_join_dyn_filter.sql b/tests/cases/standalone/common/filter/hash_join_dyn_filter.sql index d6079bc496a..405d1bc0530 100644 --- a/tests/cases/standalone/common/filter/hash_join_dyn_filter.sql +++ b/tests/cases/standalone/common/filter/hash_join_dyn_filter.sql @@ -43,6 +43,7 @@ WHERE c.tier = 'gold'; -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE ,\s"dyn_filters":\s\["DynamicFilter\s\[[^"]*\]"\] , "dyn_filters": ["DynamicFilter [ REDACTED ]"] -- SQLNESS REPLACE metrics=REDACTED\s*\| metrics=REDACTED_| +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT o."id", o.amount, c."name", c.tier FROM orders o JOIN customers c ON o.customer_id = c.customer_id @@ -100,6 +101,7 @@ WHERE c.tier IN ('gold', 'silver'); -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE ,\s"dyn_filters":\s\["DynamicFilter\s\[[^"]*\]"\] , "dyn_filters": ["DynamicFilter [ REDACTED ]"] -- SQLNESS REPLACE metrics=REDACTED\s*\| metrics=REDACTED_| +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT o."id", o.amount, c."name", c.tier FROM (SELECT "id", customer_id as cid, amount, ts FROM orders) o JOIN customers c ON o.cid = c.customer_id diff --git a/tests/cases/standalone/common/filter/hash_join_topk_dyn_filter.result b/tests/cases/standalone/common/filter/hash_join_topk_dyn_filter.result index 0ba0f3c4e00..1c2d05bbcbb 100644 --- a/tests/cases/standalone/common/filter/hash_join_topk_dyn_filter.result +++ b/tests/cases/standalone/common/filter/hash_join_topk_dyn_filter.result @@ -94,6 +94,7 @@ WHERE c.tier IN ('gold', 'bronze'); -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE ,\s"dyn_filters":\s\["DynamicFilter\s\[[^"]*\]"\] , "dyn_filters": ["DynamicFilter [ REDACTED ]"] -- SQLNESS REPLACE metrics=REDACTED\s*\| metrics=REDACTED_| +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT top_orders."id", top_orders.amount, c."name", c.tier FROM ( SELECT "id", customer_id, amount, ts @@ -119,11 +120,11 @@ WHERE c.tier IN ('gold', 'bronze'); |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_SortPreservingMergeExec: [amount@2 DESC], fetch=5 metrics=REDACTED_| |_|_|_SortExec: TopK(fetch=5), expr=[amount@2 DESC], preserve_partitioning=[true] metrics=REDACTED_| -|_|_|_SeqScan: region=REDACTED, {"partition_count":REDACTED, "projection": ["id", "customer_id", "amount", "ts"], "dyn_filters": ["DynamicFilter [ REDACTED ]"], "flat_format":REDACTED, "metrics_per_partition": REDACTED metrics=REDACTED_| +|_|_|_SeqScan: region=REDACTED, {"partition_count":REDACTED, "projection": ["id", "customer_id", "amount", "ts"], "dyn_filters":REDACTED |_|_|_| | 1_| 0_|_FilterExec: customer_id@0 IS NOT NULL AND (tier@2 = gold OR tier@2 = bronze) metrics=REDACTED_| |_|_|_CooperativeExec metrics=REDACTED_| -|_|_|_SeqScan: region=REDACTED, {"partition_count":REDACTED, "projection": ["customer_id", "name", "tier", "ts"], "filters": ["customer_id IS NOT NULL", "tier = Utf8(\"gold\") OR tier = Utf8(\"bronze\")"], "dyn_filters": ["DynamicFilter [ REDACTED ]"], "flat_format":REDACTED, "metrics_per_partition": REDACTED metrics=REDACTED_| +|_|_|_SeqScan: region=REDACTED, {"partition_count":REDACTED, "projection": ["customer_id", "name", "tier", "ts"], "filters": ["customer_id IS NOT NULL", "tier = Utf8(\"gold\") OR tier = Utf8(\"bronze\")"], "dyn_filters":REDACTED |_|_|_| |_|_| Total rows: REDACTED_| +-+-+-+ @@ -251,6 +252,7 @@ LIMIT 4; -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE ,\s"dyn_filters":\s\["DynamicFilter\s\[[^"]*\]"\] , "dyn_filters": ["DynamicFilter [ REDACTED ]"] -- SQLNESS REPLACE metrics=REDACTED\s*\| metrics=REDACTED_| +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT "id", customer_id, "name", tier, amount FROM ( SELECT o."id", o.customer_id, c."name", c.tier, o.amount @@ -276,11 +278,11 @@ LIMIT 4; |_|_|_| | 1_| 0_|_FilterExec: customer_id@1 IS NOT NULL metrics=REDACTED_| |_|_|_CooperativeExec metrics=REDACTED_| -|_|_|_SeqScan: region=REDACTED, {"partition_count":REDACTED, "projection": ["id", "customer_id", "amount", "ts"], "filters": ["customer_id IS NOT NULL"], "dyn_filters": ["DynamicFilter [ REDACTED ]"], "flat_format":REDACTED, "metrics_per_partition": REDACTED metrics=REDACTED_| +|_|_|_SeqScan: region=REDACTED, {"partition_count":REDACTED, "projection": ["id", "customer_id", "amount", "ts"], "filters": ["customer_id IS NOT NULL"], "dyn_filters":REDACTED |_|_|_| | 1_| 0_|_FilterExec: customer_id@0 IS NOT NULL AND (tier@2 = gold OR tier@2 = silver) metrics=REDACTED_| |_|_|_CooperativeExec metrics=REDACTED_| -|_|_|_SeqScan: region=REDACTED, {"partition_count":REDACTED, "projection": ["customer_id", "name", "tier", "ts"], "filters": ["customer_id IS NOT NULL", "tier = Utf8(\"gold\") OR tier = Utf8(\"silver\")"], "dyn_filters": ["DynamicFilter [ REDACTED ]"], "flat_format":REDACTED, "metrics_per_partition": REDACTED metrics=REDACTED_| +|_|_|_SeqScan: region=REDACTED, {"partition_count":REDACTED, "projection": ["customer_id", "name", "tier", "ts"], "filters": ["customer_id IS NOT NULL", "tier = Utf8(\"gold\") OR tier = Utf8(\"silver\")"], "dyn_filters":REDACTED |_|_|_| |_|_| Total rows: REDACTED_| +-+-+-+ @@ -426,6 +428,7 @@ WHERE c.tier IN ('gold', 'silver') -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE ,\s"dyn_filters":\s\["DynamicFilter\s\[[^"]*\]"\] , "dyn_filters": ["DynamicFilter [ REDACTED ]"] -- SQLNESS REPLACE metrics=REDACTED\s*\| metrics=REDACTED_| +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT o."id", o.amount, c."name", c.tier, p."name" as product_name, p."category" FROM orders o JOIN customers c ON o.customer_id = c.customer_id @@ -452,15 +455,15 @@ WHERE c.tier IN ('gold', 'silver') |_|_|_| | 1_| 0_|_FilterExec: customer_id@1 IS NOT NULL AND product_id@2 IS NOT NULL metrics=REDACTED_| |_|_|_CooperativeExec metrics=REDACTED_| -|_|_|_SeqScan: region=REDACTED, {"partition_count":REDACTED, "projection": ["id", "customer_id", "product_id", "amount", "ts"], "filters": ["customer_id IS NOT NULL", "product_id IS NOT NULL"], "flat_format":REDACTED, "metrics_per_partition": REDACTED metrics=REDACTED_| +|_|_|_SeqScan: region=REDACTED, {"partition_count":REDACTED, "projection": ["id", "customer_id", "product_id", "amount", "ts"], "filters": ["customer_id IS NOT NULL", "product_id IS NOT NULL"], "time_range":REDACTED |_|_|_| | 1_| 0_|_FilterExec: customer_id@0 IS NOT NULL AND (tier@2 = gold OR tier@2 = silver) metrics=REDACTED_| |_|_|_CooperativeExec metrics=REDACTED_| -|_|_|_SeqScan: region=REDACTED, {"partition_count":REDACTED, "projection": ["customer_id", "name", "tier", "ts"], "filters": ["customer_id IS NOT NULL", "tier = Utf8(\"gold\") OR tier = Utf8(\"silver\")"], "dyn_filters": ["DynamicFilter [ REDACTED ]"], "flat_format":REDACTED, "metrics_per_partition": REDACTED metrics=REDACTED_| +|_|_|_SeqScan: region=REDACTED, {"partition_count":REDACTED, "projection": ["customer_id", "name", "tier", "ts"], "filters": ["customer_id IS NOT NULL", "tier = Utf8(\"gold\") OR tier = Utf8(\"silver\")"], "dyn_filters":REDACTED |_|_|_| | 1_| 0_|_FilterExec: product_id@0 IS NOT NULL AND category@2 = electronics metrics=REDACTED_| |_|_|_CooperativeExec metrics=REDACTED_| -|_|_|_SeqScan: region=REDACTED, {"partition_count":REDACTED, "projection": ["product_id", "name", "category", "ts"], "filters": ["product_id IS NOT NULL", "category = Utf8(\"electronics\")"], "dyn_filters": ["DynamicFilter [ REDACTED ]"], "flat_format":REDACTED, "metrics_per_partition": REDACTED metrics=REDACTED_| +|_|_|_SeqScan: region=REDACTED, {"partition_count":REDACTED, "projection": ["product_id", "name", "category", "ts"], "filters": ["product_id IS NOT NULL", "category = Utf8(\"electronics\")"], "dyn_filters":REDACTED |_|_|_| |_|_| Total rows: REDACTED_| +-+-+-+ diff --git a/tests/cases/standalone/common/filter/hash_join_topk_dyn_filter.sql b/tests/cases/standalone/common/filter/hash_join_topk_dyn_filter.sql index 97f0c1231d7..d04c3b21bca 100644 --- a/tests/cases/standalone/common/filter/hash_join_topk_dyn_filter.sql +++ b/tests/cases/standalone/common/filter/hash_join_topk_dyn_filter.sql @@ -56,6 +56,7 @@ WHERE c.tier IN ('gold', 'bronze'); -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE ,\s"dyn_filters":\s\["DynamicFilter\s\[[^"]*\]"\] , "dyn_filters": ["DynamicFilter [ REDACTED ]"] -- SQLNESS REPLACE metrics=REDACTED\s*\| metrics=REDACTED_| +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT top_orders."id", top_orders.amount, c."name", c.tier FROM ( SELECT "id", customer_id, amount, ts @@ -136,6 +137,7 @@ LIMIT 4; -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE ,\s"dyn_filters":\s\["DynamicFilter\s\[[^"]*\]"\] , "dyn_filters": ["DynamicFilter [ REDACTED ]"] -- SQLNESS REPLACE metrics=REDACTED\s*\| metrics=REDACTED_| +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT "id", customer_id, "name", tier, amount FROM ( SELECT o."id", o.customer_id, c."name", c.tier, o.amount @@ -218,6 +220,7 @@ WHERE c.tier IN ('gold', 'silver') -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE ,\s"dyn_filters":\s\["DynamicFilter\s\[[^"]*\]"\] , "dyn_filters": ["DynamicFilter [ REDACTED ]"] -- SQLNESS REPLACE metrics=REDACTED\s*\| metrics=REDACTED_| +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT o."id", o.amount, c."name", c.tier, p."name" as product_name, p."category" FROM orders o JOIN customers c ON o.customer_id = c.customer_id diff --git a/tests/cases/standalone/common/filter/topk_dyn_filter.result b/tests/cases/standalone/common/filter/topk_dyn_filter.result index ac905994ea3..14639053b92 100644 --- a/tests/cases/standalone/common/filter/topk_dyn_filter.result +++ b/tests/cases/standalone/common/filter/topk_dyn_filter.result @@ -87,6 +87,7 @@ LIMIT 3; -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE DynamicFilter\s\[[^\]]*\] DynamicFilter [ REDACTED ] -- SQLNESS REPLACE "files":\s\[[^\]]*\] "files": REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT "id", customer_id, amount FROM orders ORDER BY amount DESC @@ -100,7 +101,7 @@ LIMIT 3; |_|_|_| | 1_| 0_|_SortPreservingMergeExec: [amount@2 DESC], fetch=3 metrics=REDACTED_| |_|_|_SortExec: TopK(fetch=3), expr=[amount@2 DESC], preserve_partitioning=[true] metrics=REDACTED_| -|_|_|_SeqScan: region=REDACTED, {"partition_count":REDACTED, "projection": ["id", "customer_id", "amount"], "dyn_filters": ["DynamicFilter [ REDACTED ]"], "files": REDACTED, "flat_format":REDACTED, "metrics_per_partition": REDACTED metrics=REDACTED | +|_|_|_SeqScan: region=REDACTED, {"partition_count":REDACTED, "projection": ["id", "customer_id", "amount"], "dyn_filters":REDACTED |_|_|_| |_|_| Total rows: REDACTED_| +-+-+-+ diff --git a/tests/cases/standalone/common/filter/topk_dyn_filter.sql b/tests/cases/standalone/common/filter/topk_dyn_filter.sql index ffce7d258e8..d77d59da4d7 100644 --- a/tests/cases/standalone/common/filter/topk_dyn_filter.sql +++ b/tests/cases/standalone/common/filter/topk_dyn_filter.sql @@ -54,6 +54,7 @@ LIMIT 3; -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE DynamicFilter\s\[[^\]]*\] DynamicFilter [ REDACTED ] -- SQLNESS REPLACE "files":\s\[[^\]]*\] "files": REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT "id", customer_id, amount FROM orders ORDER BY amount DESC diff --git a/tests/cases/standalone/common/order/order_by.result b/tests/cases/standalone/common/order/order_by.result index da47fa663c6..12f0a60f679 100644 --- a/tests/cases/standalone/common/order/order_by.result +++ b/tests/cases/standalone/common/order/order_by.result @@ -291,6 +291,7 @@ select tag from t where num > 6 order by ts; -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE ,\sfilter=\[[^]]+\] -- SQLNESS REPLACE DynamicFilter\s\[[^\]]*\] DynamicFilter [ REDACTED ] +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select tag from t where num > 6 order by ts desc limit 2; +-+-+-+ @@ -306,14 +307,14 @@ explain analyze select tag from t where num > 6 order by ts desc limit 2; |_|_|_SortPreservingMergeExec: [ts@1 DESC], fetch=2 REDACTED |_|_|_SortExec: TopK(fetch=2), expr=[ts@1 DESC], preserve_partitioning=[true] REDACTED |_|_|_FilterExec: num@2 > 6, projection=[tag@0, ts@1] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["tag", "ts", "num"], "filters": ["num > Int64(6)", "tag <= Utf8(\"z\") OR tag IS NULL"], "dyn_filters":REDACTED |_|_|_| | 1_| 1_|_FilterExec: DynamicFilter [ REDACTED ] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_SortPreservingMergeExec: [ts@1 DESC], fetch=2 REDACTED |_|_|_SortExec: TopK(fetch=2), expr=[ts@1 DESC], preserve_partitioning=[true] REDACTED |_|_|_FilterExec: num@2 > 6, projection=[tag@0, ts@1] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["tag", "ts", "num"], "filters": ["num > Int64(6)", "tag > Utf8(\"z\")"], "dyn_filters":REDACTED |_|_|_| |_|_| Total rows: 2_| +-+-+-+ diff --git a/tests/cases/standalone/common/order/order_by.sql b/tests/cases/standalone/common/order/order_by.sql index 51ffbbc4d8e..1f17be41fee 100644 --- a/tests/cases/standalone/common/order/order_by.sql +++ b/tests/cases/standalone/common/order/order_by.sql @@ -97,6 +97,7 @@ select tag from t where num > 6 order by ts; -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE ,\sfilter=\[[^]]+\] -- SQLNESS REPLACE DynamicFilter\s\[[^\]]*\] DynamicFilter [ REDACTED ] +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select tag from t where num > 6 order by ts desc limit 2; drop table t; diff --git a/tests/cases/standalone/common/order/windowed_sort.result b/tests/cases/standalone/common/order/windowed_sort.result index f85aa8c04e6..a80ff6714ad 100644 --- a/tests/cases/standalone/common/order/windowed_sort.result +++ b/tests/cases/standalone/common/order/windowed_sort.result @@ -61,6 +61,7 @@ SELECT * FROM test ORDER BY t LIMIT 5; -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT * FROM test ORDER BY t LIMIT 5; +-+-+-+ @@ -71,7 +72,7 @@ EXPLAIN ANALYZE SELECT * FROM test ORDER BY t LIMIT 5; |_|_|_| | 1_| 0_|_SortPreservingMergeExec: [t@1 ASC NULLS LAST], fetch=5 REDACTED |_|_|_SortExec: TopK(fetch=5), expr=[t@1 ASC NULLS LAST], preserve_partitioning=[true] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":4, "mem_ranges":1, "files":3, "file_ranges":3} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":4, "mem_ranges":1, "files":3, "file_ranges":3}, "projection": ["i", "t"], "dyn_filters":REDACTED |_|_|_| |_|_| Total rows: 5_| +-+-+-+ @@ -94,6 +95,7 @@ SELECT * FROM test ORDER BY t DESC LIMIT 5; -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT * FROM test ORDER BY t DESC LIMIT 5; +-+-+-+ @@ -104,7 +106,7 @@ EXPLAIN ANALYZE SELECT * FROM test ORDER BY t DESC LIMIT 5; |_|_|_| | 1_| 0_|_SortPreservingMergeExec: [t@1 DESC], fetch=5 REDACTED |_|_|_SortExec: TopK(fetch=5), expr=[t@1 DESC], preserve_partitioning=[true] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":4, "mem_ranges":1, "files":3, "file_ranges":3} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":4, "mem_ranges":1, "files":3, "file_ranges":3}, "projection": ["i", "t"], "dyn_filters":REDACTED |_|_|_| |_|_| Total rows: 5_| +-+-+-+ @@ -127,6 +129,7 @@ SELECT * FROM test where i > 2 ORDER BY t LIMIT 4; -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT * FROM test where i > 2 ORDER BY t LIMIT 4; +-+-+-+ @@ -138,7 +141,7 @@ EXPLAIN ANALYZE SELECT * FROM test where i > 2 ORDER BY t LIMIT 4; | 1_| 0_|_SortPreservingMergeExec: [t@1 ASC NULLS LAST], fetch=4 REDACTED |_|_|_SortExec: TopK(fetch=4), expr=[t@1 ASC NULLS LAST], preserve_partitioning=[true] REDACTED |_|_|_FilterExec: i@0 > 2 REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":4, "mem_ranges":1, "files":3, "file_ranges":3} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":4, "mem_ranges":1, "files":3, "file_ranges":3}, "projection": ["i", "t"], "filters": ["i > Int32(2)"], "dyn_filters":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -161,6 +164,7 @@ SELECT * FROM test where i > 2 ORDER BY t DESC LIMIT 4; -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT * FROM test where i > 2 ORDER BY t DESC LIMIT 4; +-+-+-+ @@ -172,7 +176,7 @@ EXPLAIN ANALYZE SELECT * FROM test where i > 2 ORDER BY t DESC LIMIT 4; | 1_| 0_|_SortPreservingMergeExec: [t@1 DESC], fetch=4 REDACTED |_|_|_SortExec: TopK(fetch=4), expr=[t@1 DESC], preserve_partitioning=[true] REDACTED |_|_|_FilterExec: i@0 > 2 REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":4, "mem_ranges":1, "files":3, "file_ranges":3} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":4, "mem_ranges":1, "files":3, "file_ranges":3}, "projection": ["i", "t"], "filters": ["i > Int32(2)"], "dyn_filters":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -195,6 +199,7 @@ SELECT * FROM test where t > 8 ORDER BY t DESC LIMIT 4; -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT * FROM test where t > 8 ORDER BY t DESC LIMIT 4; +-+-+-+ @@ -205,7 +210,7 @@ EXPLAIN ANALYZE SELECT * FROM test where t > 8 ORDER BY t DESC LIMIT 4; |_|_|_| | 1_| 0_|_SortPreservingMergeExec: [t@1 DESC], fetch=4 REDACTED |_|_|_SortExec: TopK(fetch=4), expr=[t@1 DESC], preserve_partitioning=[true] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":2, "mem_ranges":1, "files":1, "file_ranges":1} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":2, "mem_ranges":1, "files":1, "file_ranges":1}, "projection": ["i", "t"], "filters": ["t > TimestampMillisecond(8, None)"], "dyn_filters":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -277,6 +282,7 @@ SELECT * FROM test_pk ORDER BY t LIMIT 5; -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT * FROM test_pk ORDER BY t LIMIT 5; +-+-+-+ @@ -287,7 +293,7 @@ EXPLAIN ANALYZE SELECT * FROM test_pk ORDER BY t LIMIT 5; |_|_|_| | 1_| 0_|_SortPreservingMergeExec: [t@2 ASC NULLS LAST], fetch=5 REDACTED |_|_|_SortExec: TopK(fetch=5), expr=[t@2 ASC NULLS LAST], preserve_partitioning=[true] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":4, "mem_ranges":1, "files":3, "file_ranges":3} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":4, "mem_ranges":1, "files":3, "file_ranges":3}, "projection": ["pk", "i", "t"], "dyn_filters":REDACTED |_|_|_| |_|_| Total rows: 5_| +-+-+-+ @@ -299,6 +305,7 @@ EXPLAIN ANALYZE SELECT * FROM test_pk ORDER BY t LIMIT 5; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (files.*) REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT * FROM test_pk ORDER BY t LIMIT 5; +-+-+-+ @@ -332,6 +339,7 @@ SELECT * FROM test_pk ORDER BY t DESC LIMIT 5; -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT * FROM test_pk ORDER BY t DESC LIMIT 5; +-+-+-+ @@ -342,7 +350,7 @@ EXPLAIN ANALYZE SELECT * FROM test_pk ORDER BY t DESC LIMIT 5; |_|_|_| | 1_| 0_|_SortPreservingMergeExec: [t@2 DESC], fetch=5 REDACTED |_|_|_SortExec: TopK(fetch=5), expr=[t@2 DESC], preserve_partitioning=[true] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":4, "mem_ranges":1, "files":3, "file_ranges":3} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":4, "mem_ranges":1, "files":3, "file_ranges":3}, "projection": ["pk", "i", "t"], "dyn_filters":REDACTED |_|_|_| |_|_| Total rows: 5_| +-+-+-+ @@ -366,6 +374,7 @@ SELECT * FROM test_pk where pk > 7 ORDER BY t LIMIT 5; -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT * FROM test_pk where pk > 7 ORDER BY t LIMIT 5; +-+-+-+ @@ -376,7 +385,7 @@ EXPLAIN ANALYZE SELECT * FROM test_pk where pk > 7 ORDER BY t LIMIT 5; |_|_|_| | 1_| 0_|_SortPreservingMergeExec: [t@2 ASC NULLS LAST], fetch=5 REDACTED |_|_|_SortExec: TopK(fetch=5), expr=[t@2 ASC NULLS LAST], preserve_partitioning=[true] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":4, "mem_ranges":1, "files":3, "file_ranges":3} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":4, "mem_ranges":1, "files":3, "file_ranges":3}, "projection": ["pk", "i", "t"], "filters": ["pk > Int32(7)"], "dyn_filters":REDACTED |_|_|_| |_|_| Total rows: 5_| +-+-+-+ @@ -388,6 +397,7 @@ EXPLAIN ANALYZE SELECT * FROM test_pk where pk > 7 ORDER BY t LIMIT 5; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (files.*) REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT * FROM test_pk where pk > 7 ORDER BY t LIMIT 5; +-+-+-+ diff --git a/tests/cases/standalone/common/order/windowed_sort.sql b/tests/cases/standalone/common/order/windowed_sort.sql index 1349e4d96ca..7815693e237 100644 --- a/tests/cases/standalone/common/order/windowed_sort.sql +++ b/tests/cases/standalone/common/order/windowed_sort.sql @@ -23,6 +23,7 @@ SELECT * FROM test ORDER BY t LIMIT 5; -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT * FROM test ORDER BY t LIMIT 5; SELECT * FROM test ORDER BY t DESC LIMIT 5; @@ -33,6 +34,7 @@ SELECT * FROM test ORDER BY t DESC LIMIT 5; -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT * FROM test ORDER BY t DESC LIMIT 5; -- Filter on a field. @@ -44,6 +46,7 @@ SELECT * FROM test where i > 2 ORDER BY t LIMIT 4; -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT * FROM test where i > 2 ORDER BY t LIMIT 4; -- Filter on a field. @@ -55,6 +58,7 @@ SELECT * FROM test where i > 2 ORDER BY t DESC LIMIT 4; -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT * FROM test where i > 2 ORDER BY t DESC LIMIT 4; -- Filter on the time index. @@ -66,6 +70,7 @@ SELECT * FROM test where t > 8 ORDER BY t DESC LIMIT 4; -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT * FROM test where t > 8 ORDER BY t DESC LIMIT 4; DROP TABLE test; @@ -95,6 +100,7 @@ SELECT * FROM test_pk ORDER BY t LIMIT 5; -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT * FROM test_pk ORDER BY t LIMIT 5; -- SQLNESS REPLACE (-+) - @@ -104,6 +110,7 @@ EXPLAIN ANALYZE SELECT * FROM test_pk ORDER BY t LIMIT 5; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (files.*) REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT * FROM test_pk ORDER BY t LIMIT 5; SELECT * FROM test_pk ORDER BY t DESC LIMIT 5; @@ -114,6 +121,7 @@ SELECT * FROM test_pk ORDER BY t DESC LIMIT 5; -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT * FROM test_pk ORDER BY t DESC LIMIT 5; -- Filter on a pk column. @@ -125,6 +133,7 @@ SELECT * FROM test_pk where pk > 7 ORDER BY t LIMIT 5; -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT * FROM test_pk where pk > 7 ORDER BY t LIMIT 5; -- SQLNESS REPLACE (-+) - @@ -134,6 +143,7 @@ EXPLAIN ANALYZE SELECT * FROM test_pk where pk > 7 ORDER BY t LIMIT 5; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (files.*) REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT * FROM test_pk where pk > 7 ORDER BY t LIMIT 5; DROP TABLE test_pk; diff --git a/tests/cases/standalone/common/partition.result b/tests/cases/standalone/common/partition.result index 132f2a325f6..810899f1452 100644 --- a/tests/cases/standalone/common/partition.result +++ b/tests/cases/standalone/common/partition.result @@ -235,6 +235,7 @@ Affected Rows: 0 -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT * FROM valid_rule; +-+-+-+ @@ -244,19 +245,19 @@ EXPLAIN ANALYZE SELECT * FROM valid_rule; |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_CooperativeExec REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "projection": ["a", "b", "ts"], "filters": ["a < Int32(10) OR a IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_CooperativeExec REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "projection": ["a", "b", "ts"], "filters": ["a = Int32(10) AND (b < Utf8(\"a\") OR b IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 2_|_CooperativeExec REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "projection": ["a", "b", "ts"], "filters": ["a = Int32(10) AND b >= Utf8(\"a\") AND (b < Utf8(\"o\") OR b IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 3_|_CooperativeExec REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "projection": ["a", "b", "ts"], "filters": ["a = Int32(10) AND b >= Utf8(\"o\")"], "time_range":REDACTED |_|_|_| | 1_| 4_|_CooperativeExec REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "projection": ["a", "b", "ts"], "filters": ["a > Int32(10)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -268,6 +269,7 @@ EXPLAIN ANALYZE SELECT * FROM valid_rule; -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT * FROM valid_rule WHERE a < 10; @@ -278,7 +280,7 @@ WHERE a < 10; |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_CooperativeExec REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "projection": ["a", "b", "ts"], "filters": ["a < Int32(10)", "a < Int32(10) OR a IS NULL"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -290,6 +292,7 @@ WHERE a < 10; -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT * FROM valid_rule WHERE a = 10 AND b= 'z'; @@ -300,7 +303,7 @@ WHERE a = 10 AND b= 'z'; |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_CooperativeExec REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "projection": ["a", "b", "ts"], "filters": ["a = Int32(10)", "b = Dictionary(UInt32, Utf8(\"z\"))", "a = Int32(10) AND b >= Utf8(\"o\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -312,6 +315,7 @@ WHERE a = 10 AND b= 'z'; -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT * FROM valid_rule WHERE a = 10 OR b= 'z'; @@ -323,23 +327,23 @@ WHERE a = 10 OR b= 'z'; |_|_|_| | 1_| 0_|_FilterExec: a@0 = 10 OR b@1 = z REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "projection": ["a", "b", "ts"], "filters": ["a = Int32(10) OR b = Dictionary(UInt32, Utf8(\"z\"))", "a < Int32(10) OR a IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_FilterExec: a@0 = 10 OR b@1 = z REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "projection": ["a", "b", "ts"], "filters": ["a = Int32(10) OR b = Dictionary(UInt32, Utf8(\"z\"))", "a = Int32(10) AND (b < Utf8(\"a\") OR b IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 2_|_FilterExec: a@0 = 10 OR b@1 = z REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "projection": ["a", "b", "ts"], "filters": ["a = Int32(10) OR b = Dictionary(UInt32, Utf8(\"z\"))", "a = Int32(10) AND b >= Utf8(\"a\") AND (b < Utf8(\"o\") OR b IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 3_|_FilterExec: a@0 = 10 OR b@1 = z REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "projection": ["a", "b", "ts"], "filters": ["a = Int32(10) OR b = Dictionary(UInt32, Utf8(\"z\"))", "a = Int32(10) AND b >= Utf8(\"o\")"], "time_range":REDACTED |_|_|_| | 1_| 4_|_FilterExec: a@0 = 10 OR b@1 = z REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "projection": ["a", "b", "ts"], "filters": ["a = Int32(10) OR b = Dictionary(UInt32, Utf8(\"z\"))", "a > Int32(10)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -351,6 +355,7 @@ WHERE a = 10 OR b= 'z'; -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT * FROM valid_rule WHERE a = 10 OR ts > 1; @@ -362,23 +367,23 @@ WHERE a = 10 OR ts > 1; |_|_|_| | 1_| 0_|_FilterExec: a@0 = 10 OR ts@2 > 1 REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "projection": ["a", "b", "ts"], "filters": ["a = Int32(10) OR ts > TimestampMillisecond(1, None)", "a < Int32(10) OR a IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_FilterExec: a@0 = 10 OR ts@2 > 1 REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "projection": ["a", "b", "ts"], "filters": ["a = Int32(10) OR ts > TimestampMillisecond(1, None)", "a = Int32(10) AND (b < Utf8(\"a\") OR b IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 2_|_FilterExec: a@0 = 10 OR ts@2 > 1 REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "projection": ["a", "b", "ts"], "filters": ["a = Int32(10) OR ts > TimestampMillisecond(1, None)", "a = Int32(10) AND b >= Utf8(\"a\") AND (b < Utf8(\"o\") OR b IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 3_|_FilterExec: a@0 = 10 OR ts@2 > 1 REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "projection": ["a", "b", "ts"], "filters": ["a = Int32(10) OR ts > TimestampMillisecond(1, None)", "a = Int32(10) AND b >= Utf8(\"o\")"], "time_range":REDACTED |_|_|_| | 1_| 4_|_FilterExec: a@0 = 10 OR ts@2 > 1 REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "projection": ["a", "b", "ts"], "filters": ["a = Int32(10) OR ts > TimestampMillisecond(1, None)", "a > Int32(10)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -390,6 +395,7 @@ WHERE a = 10 OR ts > 1; -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT * FROM valid_rule WHERE a = 10 OR (ts > 1 AND b ='h'); @@ -401,23 +407,23 @@ WHERE a = 10 OR (ts > 1 AND b ='h'); |_|_|_| | 1_| 0_|_FilterExec: a@0 = 10 OR ts@2 > 1 AND b@1 = h REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "projection": ["a", "b", "ts"], "filters": ["a = Int32(10) OR ts > TimestampMillisecond(1, None) AND b = Dictionary(UInt32, Utf8(\"h\"))", "a < Int32(10) OR a IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_FilterExec: a@0 = 10 OR ts@2 > 1 AND b@1 = h REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "projection": ["a", "b", "ts"], "filters": ["a = Int32(10) OR ts > TimestampMillisecond(1, None) AND b = Dictionary(UInt32, Utf8(\"h\"))", "a = Int32(10) AND (b < Utf8(\"a\") OR b IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 2_|_FilterExec: a@0 = 10 OR ts@2 > 1 AND b@1 = h REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "projection": ["a", "b", "ts"], "filters": ["a = Int32(10) OR ts > TimestampMillisecond(1, None) AND b = Dictionary(UInt32, Utf8(\"h\"))", "a = Int32(10) AND b >= Utf8(\"a\") AND (b < Utf8(\"o\") OR b IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 3_|_FilterExec: a@0 = 10 OR ts@2 > 1 AND b@1 = h REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "projection": ["a", "b", "ts"], "filters": ["a = Int32(10) OR ts > TimestampMillisecond(1, None) AND b = Dictionary(UInt32, Utf8(\"h\"))", "a = Int32(10) AND b >= Utf8(\"o\")"], "time_range":REDACTED |_|_|_| | 1_| 4_|_FilterExec: a@0 = 10 OR ts@2 > 1 AND b@1 = h REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "projection": ["a", "b", "ts"], "filters": ["a = Int32(10) OR ts > TimestampMillisecond(1, None) AND b = Dictionary(UInt32, Utf8(\"h\"))", "a > Int32(10)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ diff --git a/tests/cases/standalone/common/partition.sql b/tests/cases/standalone/common/partition.sql index 41627bf4d09..9728a34a0c1 100644 --- a/tests/cases/standalone/common/partition.sql +++ b/tests/cases/standalone/common/partition.sql @@ -128,6 +128,7 @@ PARTITION ON COLUMNS (a, b) ( -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT * FROM valid_rule; -- SQLNESS REPLACE (metrics.*) REDACTED @@ -137,6 +138,7 @@ EXPLAIN ANALYZE SELECT * FROM valid_rule; -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT * FROM valid_rule WHERE a < 10; @@ -147,6 +149,7 @@ WHERE a < 10; -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT * FROM valid_rule WHERE a = 10 AND b= 'z'; @@ -157,6 +160,7 @@ WHERE a = 10 AND b= 'z'; -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT * FROM valid_rule WHERE a = 10 OR b= 'z'; @@ -167,6 +171,7 @@ WHERE a = 10 OR b= 'z'; -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT * FROM valid_rule WHERE a = 10 OR ts > 1; @@ -177,6 +182,7 @@ WHERE a = 10 OR ts > 1; -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT * FROM valid_rule WHERE a = 10 OR (ts > 1 AND b ='h'); diff --git a/tests/cases/standalone/common/promql/delta_temporality.result b/tests/cases/standalone/common/promql/delta_temporality.result index ac32a907cb9..4d3b1f93f17 100644 --- a/tests/cases/standalone/common/promql/delta_temporality.result +++ b/tests/cases/standalone/common/promql/delta_temporality.result @@ -54,6 +54,7 @@ TQL EVAL (180, 180, '1m') rate(delta_temporality[3m]); -- SQLNESS REPLACE input_partitions=\d+ input_partitions=REDACTED -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (180, 180, '1m') rate(delta_temporality[3m]); +-+-+-+ @@ -67,7 +68,7 @@ TQL ANALYZE (180, 180, '1m') rate(delta_temporality[3m]); |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[ts], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["series", "otlp_aggregation_temporality"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "projection": ["ts", "greptime_value", "series", "otlp_aggregation_temporality"], "filters": ["ts >= TimestampMillisecond(1, None)", "ts <= TimestampMillisecond(180000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 3_| +-+-+-+ diff --git a/tests/cases/standalone/common/promql/delta_temporality.sql b/tests/cases/standalone/common/promql/delta_temporality.sql index b021f5f6b00..0ea9328e445 100644 --- a/tests/cases/standalone/common/promql/delta_temporality.sql +++ b/tests/cases/standalone/common/promql/delta_temporality.sql @@ -34,6 +34,7 @@ TQL EVAL (180, 180, '1m') rate(delta_temporality[3m]); -- SQLNESS REPLACE input_partitions=\d+ input_partitions=REDACTED -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (180, 180, '1m') rate(delta_temporality[3m]); -- The reserved temporality marker treats NULL as the absent cumulative state. diff --git a/tests/cases/standalone/common/promql/histogram_multi_partition.result b/tests/cases/standalone/common/promql/histogram_multi_partition.result index 64e2f1e8551..d4e8090b314 100644 --- a/tests/cases/standalone/common/promql/histogram_multi_partition.result +++ b/tests/cases/standalone/common/promql/histogram_multi_partition.result @@ -36,6 +36,7 @@ Affected Rows: 12 -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE Hash\(\[ts@1\],.* Hash([ts@1],REDACTED -- SQLNESS REPLACE Hash\(\[le@0,\sts@1\],.* Hash([le@0, ts@1],REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED tql analyze (0, 10, '10s') histogram_quantile(0.5, sum by (le) (histogram_gap_bucket)); +-+-+-+ @@ -56,7 +57,7 @@ tql analyze (0, 10, '10s') histogram_quantile(0.5, sum by (le) (histogram_gap_bu |_|_|_ProjectionExec: expr=[ts@0 as ts, le@1 as le, val@3 as val] REDACTED |_|_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[10000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["shard", "le"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["ts", "le", "shard", "val"], "filters": ["ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(10000, None)", "shard < Utf8(\"n\") OR shard IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[le@0 as le, ts@1 as ts], aggr=[__sum_state(histogram_gap_bucket.val)] REDACTED |_|_|_RepartitionExec: partitioning=Hash([le@0, ts@1],REDACTED @@ -64,7 +65,7 @@ tql analyze (0, 10, '10s') histogram_quantile(0.5, sum by (le) (histogram_gap_bu |_|_|_ProjectionExec: expr=[ts@0 as ts, le@1 as le, val@3 as val] REDACTED |_|_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[10000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["shard", "le"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["ts", "le", "shard", "val"], "filters": ["ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(10000, None)", "shard >= Utf8(\"n\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 2_| +-+-+-+ diff --git a/tests/cases/standalone/common/promql/histogram_multi_partition.sql b/tests/cases/standalone/common/promql/histogram_multi_partition.sql index b360999fcf7..3bc11f98b99 100644 --- a/tests/cases/standalone/common/promql/histogram_multi_partition.sql +++ b/tests/cases/standalone/common/promql/histogram_multi_partition.sql @@ -32,6 +32,7 @@ insert into histogram_gap_bucket values -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE Hash\(\[ts@1\],.* Hash([ts@1],REDACTED -- SQLNESS REPLACE Hash\(\[le@0,\sts@1\],.* Hash([le@0, ts@1],REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED tql analyze (0, 10, '10s') histogram_quantile(0.5, sum by (le) (histogram_gap_bucket)); -- SQLNESS SORT_RESULT 2 1 diff --git a/tests/cases/standalone/common/promql/native_time_selection.result b/tests/cases/standalone/common/promql/native_time_selection.result index 945189e68f6..710f8705a85 100644 --- a/tests/cases/standalone/common/promql/native_time_selection.result +++ b/tests/cases/standalone/common/promql/native_time_selection.result @@ -65,6 +65,7 @@ TQL EXPLAIN (1, 1, '1s', '1ms') native_time_us{series="exact"}; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (flat_format.*) REDACTED -- SQLNESS REPLACE (elapsed_compute.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE VERBOSE (1, 1, '1s', '1ms') native_time_us{series="exact"}; +-+-+-+ @@ -77,7 +78,7 @@ TQL ANALYZE VERBOSE (1, 1, '1s', '1ms') native_time_us{series="exact"}; |_|_|_PromSeriesDivideExec: tags=["series"] metrics=[output_rows: 1, REDACTED |_|_|_ProjectionExec: expr=[val@2 as val, series@1 as series, ts@0 as ts] metrics=[output_rows: 1, REDACTED |_|_|_CooperativeExec metrics=[]_| -|_|_|_SeriesScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "selector":"LastRow { after_merge: true }", "distribution":"PerSeries", "projection": ["ts", "series", "val"], "filters": ["series = Dictionary(UInt32, Utf8(\"exact\"))", "ts > TimestampMicrosecond(999000, None)", "ts <= TimestampMicrosecond(1000000, None)"], "REDACTED +|_|_|_SeriesScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "selector":"LastRow { after_merge: true }", "distribution":"PerSeries", "projection": ["ts", "series", "val"], "filters": ["series = Dictionary(UInt32, Utf8(\"exact\"))", "ts > TimestampMicrosecond(999000, None)", "ts <= TimestampMicrosecond(1000000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ @@ -312,6 +313,7 @@ TQL EXPLAIN (1, 1, '1s', '1ms') native_time_ns{series="exact"}; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (flat_format.*) REDACTED -- SQLNESS REPLACE (elapsed_compute.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE VERBOSE (1, 1, '1s', '1ms') native_time_ns{series="exact"}; +-+-+-+ @@ -324,7 +326,7 @@ TQL ANALYZE VERBOSE (1, 1, '1s', '1ms') native_time_ns{series="exact"}; |_|_|_PromSeriesDivideExec: tags=["series"] metrics=[output_rows: 1, REDACTED |_|_|_ProjectionExec: expr=[val@2 as val, series@1 as series, ts@0 as ts] metrics=[output_rows: 1, REDACTED |_|_|_CooperativeExec metrics=[]_| -|_|_|_SeriesScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "selector":"LastRow { after_merge: true }", "distribution":"PerSeries", "projection": ["ts", "series", "val"], "filters": ["series = Dictionary(UInt32, Utf8(\"exact\"))", "ts > TimestampNanosecond(999000000, None)", "ts <= TimestampNanosecond(1000000000, None)"], "REDACTED +|_|_|_SeriesScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "selector":"LastRow { after_merge: true }", "distribution":"PerSeries", "projection": ["ts", "series", "val"], "filters": ["series = Dictionary(UInt32, Utf8(\"exact\"))", "ts > TimestampNanosecond(999000000, None)", "ts <= TimestampNanosecond(1000000000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ @@ -541,6 +543,7 @@ TQL EXPLAIN (0, 0, '1s', '2d') native_time_ns_lower_overflow{series="exact"} off -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (flat_format.*) REDACTED -- SQLNESS REPLACE (elapsed_compute.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE VERBOSE (0, 0, '1s', '2d') native_time_ns_lower_overflow{series="exact"} offset 106750d; +-+-+-+ @@ -554,7 +557,7 @@ TQL ANALYZE VERBOSE (0, 0, '1s', '2d') native_time_ns_lower_overflow{series="exa |_|_|_PromSeriesDivideExec: tags=["series"] metrics=[output_rows: 1, REDACTED |_|_|_ProjectionExec: expr=[val@2 as val, series@1 as series, ts@0 as ts] metrics=[output_rows: 1, REDACTED |_|_|_CooperativeExec metrics=[]_| -|_|_|_SeriesScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "selector":"LastRow { after_merge: true }", "distribution":"PerSeries", "projection": ["ts", "series", "val"], "filters": ["series = Dictionary(UInt32, Utf8(\"exact\"))", "ts <= TimestampNanosecond(-9223200000000000000, None)"], "REDACTED +|_|_|_SeriesScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "selector":"LastRow { after_merge: true }", "distribution":"PerSeries", "projection": ["ts", "series", "val"], "filters": ["series = Dictionary(UInt32, Utf8(\"exact\"))", "ts <= TimestampNanosecond(-9223200000000000000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ diff --git a/tests/cases/standalone/common/promql/native_time_selection.sql b/tests/cases/standalone/common/promql/native_time_selection.sql index 204b047ef23..fb6e05d74b2 100644 --- a/tests/cases/standalone/common/promql/native_time_selection.sql +++ b/tests/cases/standalone/common/promql/native_time_selection.sql @@ -46,6 +46,7 @@ TQL EXPLAIN (1, 1, '1s', '1ms') native_time_us{series="exact"}; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (flat_format.*) REDACTED -- SQLNESS REPLACE (elapsed_compute.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE VERBOSE (1, 1, '1s', '1ms') native_time_us{series="exact"}; -- The same-series future tick is in the memtable, while exact 1s remains selected. @@ -136,6 +137,7 @@ TQL EXPLAIN (1, 1, '1s', '1ms') native_time_ns{series="exact"}; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (flat_format.*) REDACTED -- SQLNESS REPLACE (elapsed_compute.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE VERBOSE (1, 1, '1s', '1ms') native_time_ns{series="exact"}; -- The same-series future tick is in the memtable, while exact 1s remains selected. @@ -206,6 +208,7 @@ TQL EXPLAIN (0, 0, '1s', '2d') native_time_ns_lower_overflow{series="exact"} off -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (flat_format.*) REDACTED -- SQLNESS REPLACE (elapsed_compute.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE VERBOSE (0, 0, '1s', '2d') native_time_ns_lower_overflow{series="exact"} offset 106750d; -- The representable upper bound selects only the exact row. diff --git a/tests/cases/standalone/common/promql/regex.result b/tests/cases/standalone/common/promql/regex.result index c2a10626ff8..be289cdacb8 100644 --- a/tests/cases/standalone/common/promql/regex.result +++ b/tests/cases/standalone/common/promql/regex.result @@ -66,6 +66,7 @@ TQL EVAL (0, 100, '15s') test{host=~"(10\\.0\\.160\\.237:8080|10\\.0\\.160\\.237 -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (flat_format.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE VERBOSE (0, 0, '1s') test{host=~".*"}; +-+-+-+ @@ -77,7 +78,7 @@ TQL ANALYZE VERBOSE (0, 0, '1s') test{host=~".*"}; | 1_| 0_|_PromInstantManipulateExec: range=[0..0], lookback=[300000], interval=[1000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["host"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "selector":"LastRow { after_merge: true }", "distribution":"PerSeries", "projection": ["ts", "host", "val"], "filters": ["ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(0, None)"], "REDACTED +|_|_|_SeriesScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "selector":"LastRow { after_merge: true }", "distribution":"PerSeries", "projection": ["ts", "host", "val"], "filters": ["ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(0, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 2_| +-+-+-+ @@ -90,6 +91,7 @@ TQL ANALYZE VERBOSE (0, 0, '1s') test{host=~".*"}; -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (flat_format.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE VERBOSE (0, 0, '1s') test{host=~".+"}; +-+-+-+ @@ -101,7 +103,7 @@ TQL ANALYZE VERBOSE (0, 0, '1s') test{host=~".+"}; | 1_| 0_|_PromInstantManipulateExec: range=[0..0], lookback=[300000], interval=[1000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["host"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "selector":"LastRow { after_merge: true }", "distribution":"PerSeries", "projection": ["ts", "host", "val"], "filters": ["host != Dictionary(UInt32, Utf8(\"\"))", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(0, None)"], "REDACTED +|_|_|_SeriesScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "selector":"LastRow { after_merge: true }", "distribution":"PerSeries", "projection": ["ts", "host", "val"], "filters": ["host != Dictionary(UInt32, Utf8(\"\"))", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(0, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 2_| +-+-+-+ @@ -114,6 +116,7 @@ TQL ANALYZE VERBOSE (0, 0, '1s') test{host=~".+"}; -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (flat_format.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE VERBOSE (0, 0, '1s') test{host!~".*"}; +-+-+-+ @@ -136,6 +139,7 @@ TQL ANALYZE VERBOSE (0, 0, '1s') test{host!~".*"}; -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (flat_format.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE VERBOSE (0, 0, '1s') test{host!~".+"}; +-+-+-+ @@ -147,7 +151,7 @@ TQL ANALYZE VERBOSE (0, 0, '1s') test{host!~".+"}; | 1_| 0_|_PromInstantManipulateExec: range=[0..0], lookback=[300000], interval=[1000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["host"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "selector":"LastRow { after_merge: true }", "distribution":"PerSeries", "projection": ["ts", "host", "val"], "filters": ["host IS NULL OR host = Dictionary(UInt32, Utf8(\"\"))", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(0, None)"], "REDACTED +|_|_|_SeriesScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "selector":"LastRow { after_merge: true }", "distribution":"PerSeries", "projection": ["ts", "host", "val"], "filters": ["host IS NULL OR host = Dictionary(UInt32, Utf8(\"\"))", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(0, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ diff --git a/tests/cases/standalone/common/promql/regex.sql b/tests/cases/standalone/common/promql/regex.sql index 80a8f3ed381..3ad88c6ea8c 100644 --- a/tests/cases/standalone/common/promql/regex.sql +++ b/tests/cases/standalone/common/promql/regex.sql @@ -27,6 +27,7 @@ TQL EVAL (0, 100, '15s') test{host=~"(10\\.0\\.160\\.237:8080|10\\.0\\.160\\.237 -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (flat_format.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE VERBOSE (0, 0, '1s') test{host=~".*"}; -- SQLNESS REPLACE (metrics.*) REDACTED @@ -37,6 +38,7 @@ TQL ANALYZE VERBOSE (0, 0, '1s') test{host=~".*"}; -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (flat_format.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE VERBOSE (0, 0, '1s') test{host=~".+"}; -- SQLNESS REPLACE (metrics.*) REDACTED @@ -47,6 +49,7 @@ TQL ANALYZE VERBOSE (0, 0, '1s') test{host=~".+"}; -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (flat_format.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE VERBOSE (0, 0, '1s') test{host!~".*"}; -- SQLNESS REPLACE (metrics.*) REDACTED @@ -57,6 +60,7 @@ TQL ANALYZE VERBOSE (0, 0, '1s') test{host!~".*"}; -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (flat_format.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE VERBOSE (0, 0, '1s') test{host!~".+"}; INSERT INTO test VALUES (0, NULL, 2), (0, '', 3), (0, 'ext4', 4), (0, 'tmpfs', 5); diff --git a/tests/cases/standalone/common/promql/tsid_binary_join_regression.result b/tests/cases/standalone/common/promql/tsid_binary_join_regression.result index 0288a045390..144f9defc35 100644 --- a/tests/cases/standalone/common/promql/tsid_binary_join_regression.result +++ b/tests/cases/standalone/common/promql/tsid_binary_join_regression.result @@ -110,6 +110,8 @@ Affected Rows: 4 -- SQLNESS REPLACE input_partitions=\d+ input_partitions=REDACTED -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 5, '5s') tsid_binary_join_left / tsid_binary_join_right; +-+-+-+ @@ -127,13 +129,13 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left / tsid_binary_join_right; |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "projection": ["ts", "greptime_value", "__tsid", "host", "job"], "filters": ["__table_id=UInt32(REDACTED)", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(5000, None)"], "time_range":REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "projection": ["ts", "greptime_value", "__tsid", "host", "job"], "filters": ["__table_id=UInt32(REDACTED)", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(5000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -150,6 +152,8 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left / tsid_binary_join_right; -- SQLNESS REPLACE input_partitions=\d+ input_partitions=REDACTED -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left + tsid_binary_join_right) / tsid_binary_join_left; +-+-+-+ @@ -167,13 +171,13 @@ TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left + tsid_binary_join_right) / tsid |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "projection": ["ts", "greptime_value", "__tsid", "host", "job"], "filters": ["__table_id=UInt32(REDACTED)", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(5000, None)"], "time_range":REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "projection": ["ts", "greptime_value", "__tsid", "host", "job"], "filters": ["__table_id=UInt32(REDACTED)", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(5000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -189,6 +193,8 @@ TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left + tsid_binary_join_right) / tsid -- SQLNESS REPLACE input_partitions=\d+ input_partitions=REDACTED -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 5, '5s') ((tsid_binary_join_left + tsid_binary_join_right) * (tsid_binary_join_left - tsid_binary_join_third)) / (tsid_binary_join_left + 2); +-+-+-+ @@ -210,19 +216,19 @@ TQL ANALYZE (0, 5, '5s') ((tsid_binary_join_left + tsid_binary_join_right) * (ts |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "projection": ["ts", "greptime_value", "__tsid", "host", "job"], "filters": ["__table_id=UInt32(REDACTED)", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(5000, None)"], "time_range":REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "projection": ["ts", "greptime_value", "__tsid", "host", "job"], "filters": ["__table_id=UInt32(REDACTED)", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(5000, None)"], "time_range":REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "projection": ["ts", "greptime_value", "__tsid", "host", "job"], "filters": ["__table_id=UInt32(REDACTED)", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(5000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -238,6 +244,8 @@ TQL ANALYZE (0, 5, '5s') ((tsid_binary_join_left + tsid_binary_join_right) * (ts -- SQLNESS REPLACE input_partitions=\d+ input_partitions=REDACTED -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 5, '5s') tsid_binary_join_left / ignoring(host) tsid_binary_join_right; +-+-+-+ @@ -262,13 +270,13 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left / ignoring(host) tsid_binary_join |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "projection": ["ts", "greptime_value", "__tsid", "host", "job"], "filters": ["__table_id=UInt32(REDACTED)", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(5000, None)"], "time_range":REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "projection": ["ts", "greptime_value", "__tsid", "host", "job"], "filters": ["__table_id=UInt32(REDACTED)", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(5000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -284,6 +292,8 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left / ignoring(host) tsid_binary_join -- SQLNESS REPLACE input_partitions=\d+ input_partitions=REDACTED -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 5, '5s') tsid_binary_join_left / on(job) tsid_binary_join_right_by_job; +-+-+-+ @@ -308,13 +318,13 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left / on(job) tsid_binary_join_right_ |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "projection": ["ts", "greptime_value", "__tsid", "host", "job"], "filters": ["__table_id=UInt32(REDACTED)", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(5000, None)"], "time_range":REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, job@3 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "projection": ["ts", "greptime_value", "__tsid", "job"], "filters": ["__table_id=UInt32(REDACTED)", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(5000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -331,6 +341,8 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left / on(job) tsid_binary_join_right_ -- SQLNESS REPLACE input_partitions=\d+ input_partitions=REDACTED -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 5, '5s') tsid_binary_join_left > tsid_binary_join_right; +-+-+-+ @@ -347,13 +359,13 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left > tsid_binary_join_right; |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "projection": ["ts", "greptime_value", "__tsid", "host", "job"], "filters": ["__table_id=UInt32(REDACTED)", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(5000, None)"], "time_range":REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "projection": ["ts", "greptime_value", "__tsid", "host", "job"], "filters": ["__table_id=UInt32(REDACTED)", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(5000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -369,6 +381,8 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left > tsid_binary_join_right; -- SQLNESS REPLACE input_partitions=\d+ input_partitions=REDACTED -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 5, '5s') tsid_binary_join_left > bool tsid_binary_join_right; +-+-+-+ @@ -386,13 +400,13 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left > bool tsid_binary_join_right; |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "projection": ["ts", "greptime_value", "__tsid", "host", "job"], "filters": ["__table_id=UInt32(REDACTED)", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(5000, None)"], "time_range":REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "projection": ["ts", "greptime_value", "__tsid", "host", "job"], "filters": ["__table_id=UInt32(REDACTED)", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(5000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -408,6 +422,8 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left > bool tsid_binary_join_right; -- SQLNESS REPLACE input_partitions=\d+ input_partitions=REDACTED -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 5, '5s') ((tsid_binary_join_left > tsid_binary_join_right) / tsid_binary_join_left) * 100; +-+-+-+ @@ -430,19 +446,19 @@ TQL ANALYZE (0, 5, '5s') ((tsid_binary_join_left > tsid_binary_join_right) / tsi |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "projection": ["ts", "greptime_value", "__tsid", "host", "job"], "filters": ["__table_id=UInt32(REDACTED)", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(5000, None)"], "time_range":REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "projection": ["ts", "greptime_value", "__tsid", "host", "job"], "filters": ["__table_id=UInt32(REDACTED)", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(5000, None)"], "time_range":REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "projection": ["ts", "greptime_value", "__tsid", "host", "job"], "filters": ["__table_id=UInt32(REDACTED)", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(5000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -457,6 +473,8 @@ TQL ANALYZE (0, 5, '5s') ((tsid_binary_join_left > tsid_binary_join_right) / tsi -- SQLNESS REPLACE input_partitions=\d+ input_partitions=REDACTED -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 5, '5s') ((tsid_binary_join_left > bool tsid_binary_join_right) + tsid_binary_join_left) / tsid_binary_join_left; +-+-+-+ @@ -485,25 +503,25 @@ TQL ANALYZE (0, 5, '5s') ((tsid_binary_join_left > bool tsid_binary_join_right) |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "projection": ["ts", "greptime_value", "__tsid", "host", "job"], "filters": ["__table_id=UInt32(REDACTED)", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(5000, None)"], "time_range":REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "projection": ["ts", "greptime_value", "__tsid", "host", "job"], "filters": ["__table_id=UInt32(REDACTED)", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(5000, None)"], "time_range":REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "projection": ["ts", "greptime_value", "__tsid", "host", "job"], "filters": ["__table_id=UInt32(REDACTED)", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(5000, None)"], "time_range":REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "projection": ["ts", "greptime_value", "__tsid", "host", "job"], "filters": ["__table_id=UInt32(REDACTED)", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(5000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -519,6 +537,8 @@ TQL ANALYZE (0, 5, '5s') ((tsid_binary_join_left > bool tsid_binary_join_right) -- SQLNESS REPLACE input_partitions=\d+ input_partitions=REDACTED -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left or tsid_binary_join_right) / tsid_binary_join_left; +-+-+-+ @@ -541,7 +561,7 @@ TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left or tsid_binary_join_right) / tsi |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "projection": ["ts", "greptime_value", "__tsid", "host", "job"], "filters": ["__table_id=UInt32(REDACTED)", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(5000, None)"], "time_range":REDACTED |_|_|_| | 1_| 0_|_ProjectionExec: expr=[ts@2 as ts, __tsid@3 as __tsid, greptime_value@4 as greptime_value, host@5 as host, job@6 as job, CASE WHEN __common_expr_1@0 IS NOT NULL THEN __common_expr_1@0 ELSE_END as __promql_or_match_0, CASE WHEN __common_expr_2@1 IS NOT NULL THEN __common_expr_2@1 ELSE_END as __promql_or_match_1] REDACTED |_|_|_ProjectionExec: expr=[CAST(host@1 AS Utf8) as __common_expr_1, CAST(job@2 AS Utf8) as __common_expr_2, ts@4 as ts, __tsid@3 as __tsid, greptime_value@0 as greptime_value, host@1 as host, job@2 as job] REDACTED @@ -549,13 +569,13 @@ TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left or tsid_binary_join_right) / tsi |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "projection": ["ts", "greptime_value", "__tsid", "host", "job"], "filters": ["__table_id=UInt32(REDACTED)", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(5000, None)"], "time_range":REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "projection": ["ts", "greptime_value", "__tsid", "host", "job"], "filters": ["__table_id=UInt32(REDACTED)", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(5000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -571,6 +591,8 @@ TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left or tsid_binary_join_right) / tsi -- SQLNESS REPLACE input_partitions=\d+ input_partitions=REDACTED -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left / ignoring(host) group_left tsid_binary_join_right) / tsid_binary_join_left; +-+-+-+ @@ -597,19 +619,19 @@ TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left / ignoring(host) group_left tsid |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "projection": ["ts", "greptime_value", "__tsid", "host", "job"], "filters": ["__table_id=UInt32(REDACTED)", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(5000, None)"], "time_range":REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "projection": ["ts", "greptime_value", "__tsid", "host", "job"], "filters": ["__table_id=UInt32(REDACTED)", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(5000, None)"], "time_range":REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "projection": ["ts", "greptime_value", "__tsid", "host", "job"], "filters": ["__table_id=UInt32(REDACTED)", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(5000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ diff --git a/tests/cases/standalone/common/promql/tsid_binary_join_regression.sql b/tests/cases/standalone/common/promql/tsid_binary_join_regression.sql index 627680330aa..07c0d5bb606 100644 --- a/tests/cases/standalone/common/promql/tsid_binary_join_regression.sql +++ b/tests/cases/standalone/common/promql/tsid_binary_join_regression.sql @@ -93,6 +93,8 @@ INSERT INTO tsid_binary_join_third (host, job, ts, greptime_value) VALUES -- SQLNESS REPLACE input_partitions=\d+ input_partitions=REDACTED -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 5, '5s') tsid_binary_join_left / tsid_binary_join_right; -- Repeated operands in a safe arithmetic island should be planned once and reused @@ -107,6 +109,8 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left / tsid_binary_join_right; -- SQLNESS REPLACE input_partitions=\d+ input_partitions=REDACTED -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left + tsid_binary_join_right) / tsid_binary_join_left; -- A larger arithmetic island should still plan each distinct vector selector only once @@ -120,6 +124,8 @@ TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left + tsid_binary_join_right) / tsid -- SQLNESS REPLACE input_partitions=\d+ input_partitions=REDACTED -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 5, '5s') ((tsid_binary_join_left + tsid_binary_join_right) * (tsid_binary_join_left - tsid_binary_join_third)) / (tsid_binary_join_left + 2); -- Label modifiers must disable the TSID shortcut and keep matching on the remaining labels. @@ -133,6 +139,8 @@ TQL ANALYZE (0, 5, '5s') ((tsid_binary_join_left + tsid_binary_join_right) * (ts -- SQLNESS REPLACE input_partitions=\d+ input_partitions=REDACTED -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 5, '5s') tsid_binary_join_left / ignoring(host) tsid_binary_join_right; -- `on(job)` must stay label-based when only the left side has extra row-key labels. @@ -146,6 +154,8 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left / ignoring(host) tsid_binary_join -- SQLNESS REPLACE input_partitions=\d+ input_partitions=REDACTED -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 5, '5s') tsid_binary_join_left / on(job) tsid_binary_join_right_by_job; -- Comparison filters can join on `__tsid`, but the filtered result must still behave like @@ -160,6 +170,8 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left / on(job) tsid_binary_join_right_ -- SQLNESS REPLACE input_partitions=\d+ input_partitions=REDACTED -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 5, '5s') tsid_binary_join_left > tsid_binary_join_right; -- `bool` comparison should follow the same TSID-backed matching path. @@ -173,6 +185,8 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left > tsid_binary_join_right; -- SQLNESS REPLACE input_partitions=\d+ input_partitions=REDACTED -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 5, '5s') tsid_binary_join_left > bool tsid_binary_join_right; -- Comparison filters are a barrier for binary island coalescing because they filter the @@ -186,6 +200,8 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left > bool tsid_binary_join_right; -- SQLNESS REPLACE input_partitions=\d+ input_partitions=REDACTED -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 5, '5s') ((tsid_binary_join_left > tsid_binary_join_right) / tsid_binary_join_left) * 100; -- Bool comparisons are intentionally outside the first island optimization version. @@ -198,6 +214,8 @@ TQL ANALYZE (0, 5, '5s') ((tsid_binary_join_left > tsid_binary_join_right) / tsi -- SQLNESS REPLACE input_partitions=\d+ input_partitions=REDACTED -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 5, '5s') ((tsid_binary_join_left > bool tsid_binary_join_right) + tsid_binary_join_left) / tsid_binary_join_left; -- Set operators are a barrier because they have distinct matching and output-domain @@ -211,6 +229,8 @@ TQL ANALYZE (0, 5, '5s') ((tsid_binary_join_left > bool tsid_binary_join_right) -- SQLNESS REPLACE input_partitions=\d+ input_partitions=REDACTED -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left or tsid_binary_join_right) / tsid_binary_join_left; -- Group modifiers are many-to-one/one-to-many matching barriers and must stay on the @@ -224,6 +244,8 @@ TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left or tsid_binary_join_right) / tsi -- SQLNESS REPLACE input_partitions=\d+ input_partitions=REDACTED -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left / ignoring(host) group_left tsid_binary_join_right) / tsid_binary_join_left; -- A group modifier that keeps the many side's whole tag set keeps its `__tsid` too, so the diff --git a/tests/cases/standalone/common/promql/tsid_histogram_quantile_regression.result b/tests/cases/standalone/common/promql/tsid_histogram_quantile_regression.result index 10b31cac1e1..e128850b241 100644 --- a/tests/cases/standalone/common/promql/tsid_histogram_quantile_regression.result +++ b/tests/cases/standalone/common/promql/tsid_histogram_quantile_regression.result @@ -103,6 +103,8 @@ TQL EVAL (0, 10, '5s') avg_over_time(tsid_no_aggr_metric[5s]); -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 10, '5s') histogram_quantile(0.5, tsid_no_aggr_histogram_bucket); +-+-+-+ @@ -118,7 +120,7 @@ TQL ANALYZE (0, 10, '5s') histogram_quantile(0.5, tsid_no_aggr_histogram_bucket) |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[val@1 as val, job@3 as job, le@4 as le, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["ts", "val", "__tsid", "job", "le"], "filters": ["__table_id=UInt32(REDACTED)", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(10000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 3_| +-+-+-+ diff --git a/tests/cases/standalone/common/promql/tsid_histogram_quantile_regression.sql b/tests/cases/standalone/common/promql/tsid_histogram_quantile_regression.sql index 2c450d3bc23..7faea98007e 100644 --- a/tests/cases/standalone/common/promql/tsid_histogram_quantile_regression.sql +++ b/tests/cases/standalone/common/promql/tsid_histogram_quantile_regression.sql @@ -70,6 +70,8 @@ TQL EVAL (0, 10, '5s') avg_over_time(tsid_no_aggr_metric[5s]); -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 10, '5s') histogram_quantile(0.5, tsid_no_aggr_histogram_bucket); -- SQLNESS SORT_RESULT 3 1 diff --git a/tests/cases/standalone/common/range/nest.result b/tests/cases/standalone/common/range/nest.result index e283c457526..2409155486c 100644 --- a/tests/cases/standalone/common/range/nest.result +++ b/tests/cases/standalone/common/range/nest.result @@ -147,6 +147,7 @@ EXPLAIN SELECT ts, host, min(val) RANGE '5s' FROM host ALIGN '5s'; -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT ts, host, min(val) RANGE '5s' FROM host ALIGN '5s'; +-+-+-+ @@ -157,7 +158,7 @@ EXPLAIN ANALYZE SELECT ts, host, min(val) RANGE '5s' FROM host ALIGN '5s'; |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_CooperativeExec REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["ts", "host", "val"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 10_| +-+-+-+ diff --git a/tests/cases/standalone/common/range/nest.sql b/tests/cases/standalone/common/range/nest.sql index d92f9f1345f..49be120b887 100644 --- a/tests/cases/standalone/common/range/nest.sql +++ b/tests/cases/standalone/common/range/nest.sql @@ -84,6 +84,7 @@ EXPLAIN SELECT ts, host, min(val) RANGE '5s' FROM host ALIGN '5s'; -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT ts, host, min(val) RANGE '5s' FROM host ALIGN '5s'; DROP TABLE host; diff --git a/tests/cases/standalone/common/select/prune.result b/tests/cases/standalone/common/select/prune.result index 5e219a195db..36b636a514b 100644 --- a/tests/cases/standalone/common/select/prune.result +++ b/tests/cases/standalone/common/select/prune.result @@ -84,6 +84,7 @@ select * from demo where collector='disk' order by ts; -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select * from demo where idc='idc1'; +-+-+-+ @@ -93,7 +94,7 @@ explain analyze select * from demo where idc='idc1'; |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_CooperativeExec REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["ts", "value", "host", "idc", "collector"], "filters": ["idc = Dictionary(UInt32, Utf8(\"idc1\"))"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 2_| +-+-+-+ @@ -112,6 +113,7 @@ SELECT * FROM demo where host in ('test1'); -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze SELECT * FROM demo where host in ('test1'); +-+-+-+ @@ -121,7 +123,7 @@ explain analyze SELECT * FROM demo where host in ('test1'); |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_CooperativeExec REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["ts", "value", "host", "idc", "collector"], "filters": ["host = Dictionary(UInt32, Utf8(\"test1\"))"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ diff --git a/tests/cases/standalone/common/select/prune.sql b/tests/cases/standalone/common/select/prune.sql index 226c1cef69c..7a938e9c43e 100644 --- a/tests/cases/standalone/common/select/prune.sql +++ b/tests/cases/standalone/common/select/prune.sql @@ -25,6 +25,7 @@ select * from demo where collector='disk' order by ts; -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select * from demo where idc='idc1'; SELECT * FROM demo where host in ('test1'); @@ -35,6 +36,7 @@ SELECT * FROM demo where host in ('test1'); -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze SELECT * FROM demo where host in ('test1'); drop table demo; diff --git a/tests/cases/standalone/common/select/skipping_index.result b/tests/cases/standalone/common/select/skipping_index.result index 3802e1d7df5..3131e68c4c7 100644 --- a/tests/cases/standalone/common/select/skipping_index.result +++ b/tests/cases/standalone/common/select/skipping_index.result @@ -51,6 +51,7 @@ ADMIN FLUSH_TABLE('skipping_table'); -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT * FROM skipping_table WHERE id = 'id2' ORDER BY `name`; +-+-+-+ @@ -62,7 +63,7 @@ EXPLAIN ANALYZE SELECT * FROM skipping_table WHERE id = 'id2' ORDER BY `name`; | 1_| 0_|_SortPreservingMergeExec: [name@2 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[name@2 ASC NULLS LAST], preserve_partitioning=[true] REDACTED |_|_|_FilterExec: id@1 = id2 REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":3, "mem_ranges":0, "files":3, "file_ranges":3} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":3, "mem_ranges":0, "files":3, "file_ranges":3}, "projection": ["ts", "id", "name"], "filters": ["id = Utf8(\"id2\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ @@ -73,6 +74,7 @@ EXPLAIN ANALYZE SELECT * FROM skipping_table WHERE id = 'id2' ORDER BY `name`; -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT * FROM skipping_table WHERE id = 'id5' ORDER BY `name`; +-+-+-+ @@ -84,7 +86,7 @@ EXPLAIN ANALYZE SELECT * FROM skipping_table WHERE id = 'id5' ORDER BY `name`; | 1_| 0_|_SortPreservingMergeExec: [name@2 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[name@2 ASC NULLS LAST], preserve_partitioning=[true] REDACTED |_|_|_FilterExec: id@1 = id5 REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":3, "mem_ranges":0, "files":3, "file_ranges":3} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":3, "mem_ranges":0, "files":3, "file_ranges":3}, "projection": ["ts", "id", "name"], "filters": ["id = Utf8(\"id5\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ diff --git a/tests/cases/standalone/common/select/skipping_index.sql b/tests/cases/standalone/common/select/skipping_index.sql index b17d70f598f..a2ac365a525 100644 --- a/tests/cases/standalone/common/select/skipping_index.sql +++ b/tests/cases/standalone/common/select/skipping_index.sql @@ -25,6 +25,7 @@ ADMIN FLUSH_TABLE('skipping_table'); -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT * FROM skipping_table WHERE id = 'id2' ORDER BY `name`; -- SQLNESS REPLACE (metrics.*) REDACTED @@ -33,6 +34,7 @@ EXPLAIN ANALYZE SELECT * FROM skipping_table WHERE id = 'id2' ORDER BY `name`; -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT * FROM skipping_table WHERE id = 'id5' ORDER BY `name`; DROP TABLE IF EXISTS skipping_table; diff --git a/tests/cases/standalone/common/select/tql_filter.result b/tests/cases/standalone/common/select/tql_filter.result index 45f85ead00d..d15fb075755 100644 --- a/tests/cases/standalone/common/select/tql_filter.result +++ b/tests/cases/standalone/common/select/tql_filter.result @@ -13,6 +13,7 @@ Affected Rows: 3 -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED tql analyze (1, 3, '1s') t1{ a = "a" }; +-+-+-+ @@ -24,7 +25,7 @@ tql analyze (1, 3, '1s') t1{ a = "a" }; | 1_| 0_|_PromInstantManipulateExec: range=[1000..3000], lookback=[300000], interval=[1000], time index=[b] REDACTED |_|_|_PromSeriesDivideExec: tags=["a"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["a", "b", "c"], "filters": ["a = Dictionary(UInt32, Utf8(\"a\"))", "b >= TimestampMillisecond(-298999, None)", "b <= TimestampMillisecond(3000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 3_| +-+-+-+ @@ -36,6 +37,7 @@ tql analyze (1, 3, '1s') t1{ a = "a" }; -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED tql analyze (1, 3, '1s') t1{ a =~ ".*" }; +-+-+-+ @@ -47,7 +49,7 @@ tql analyze (1, 3, '1s') t1{ a =~ ".*" }; | 1_| 0_|_PromInstantManipulateExec: range=[1000..3000], lookback=[300000], interval=[1000], time index=[b] REDACTED |_|_|_PromSeriesDivideExec: tags=["a"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["a", "b", "c"], "filters": ["b >= TimestampMillisecond(-298999, None)", "b <= TimestampMillisecond(3000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 6_| +-+-+-+ @@ -59,6 +61,7 @@ tql analyze (1, 3, '1s') t1{ a =~ ".*" }; -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED tql analyze (1, 3, '1s') t1{ a =~ "a.*" }; +-+-+-+ @@ -70,7 +73,7 @@ tql analyze (1, 3, '1s') t1{ a =~ "a.*" }; | 1_| 0_|_PromInstantManipulateExec: range=[1000..3000], lookback=[300000], interval=[1000], time index=[b] REDACTED |_|_|_PromSeriesDivideExec: tags=["a"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["a", "b", "c"], "filters": ["b >= TimestampMillisecond(-298999, None)", "b <= TimestampMillisecond(3000, None)", "a ~ Utf8(\"^(?:a.*)$\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 3_| +-+-+-+ diff --git a/tests/cases/standalone/common/select/tql_filter.sql b/tests/cases/standalone/common/select/tql_filter.sql index 442efbdd145..325fdd06d9a 100644 --- a/tests/cases/standalone/common/select/tql_filter.sql +++ b/tests/cases/standalone/common/select/tql_filter.sql @@ -9,6 +9,7 @@ insert into t1 values ("a", 1000, 1.0), ("b", 2000, 2.0), ("c", 3000, 3.0); -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED tql analyze (1, 3, '1s') t1{ a = "a" }; -- SQLNESS REPLACE (metrics.*) REDACTED @@ -18,6 +19,7 @@ tql analyze (1, 3, '1s') t1{ a = "a" }; -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED tql analyze (1, 3, '1s') t1{ a =~ ".*" }; -- SQLNESS REPLACE (metrics.*) REDACTED @@ -27,6 +29,7 @@ tql analyze (1, 3, '1s') t1{ a =~ ".*" }; -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED tql analyze (1, 3, '1s') t1{ a =~ "a.*" }; drop table t1; diff --git a/tests/cases/standalone/common/tql/partition.result b/tests/cases/standalone/common/tql/partition.result index ac3072342d7..04efb194222 100644 --- a/tests/cases/standalone/common/tql/partition.result +++ b/tests/cases/standalone/common/tql/partition.result @@ -14,6 +14,7 @@ Affected Rows: 0 -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED tql analyze (0, 10, '1s') 100 - (avg by (k) (irate(t[1m])) * 100); +-+-+-+ @@ -32,7 +33,7 @@ tql analyze (0, 10, '1s') 100 - (avg by (k) (irate(t[1m])) * 100); |_|_|_PromRangeManipulateExec: req range=[0..10000], interval=[1000], eval range=[60000], time index=[j] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[j], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["k"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["i", "j", "k"], "filters": ["j >= TimestampMillisecond(-59999, None)", "j <= TimestampMillisecond(10000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -59,6 +60,7 @@ Affected Rows: 0 -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED tql analyze (0, 10, '1s') 100 - (avg by (k) (irate(t[1m])) * 100); +-+-+-+ @@ -82,7 +84,7 @@ tql analyze (0, 10, '1s') 100 - (avg by (k) (irate(t[1m])) * 100); |_|_|_PromRangeManipulateExec: req range=[0..10000], interval=[1000], eval range=[60000], time index=[j] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[j], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["i", "j", "k", "l"], "filters": ["j >= TimestampMillisecond(-59999, None)", "j <= TimestampMillisecond(10000, None)", "k < Utf8(\"a\") OR k IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[k@0 as k, j@1 as j], aggr=[__avg_state(prom_irate(j_range,i))] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED @@ -92,7 +94,7 @@ tql analyze (0, 10, '1s') 100 - (avg by (k) (irate(t[1m])) * 100); |_|_|_PromRangeManipulateExec: req range=[0..10000], interval=[1000], eval range=[60000], time index=[j] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[j], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["i", "j", "k", "l"], "filters": ["j >= TimestampMillisecond(-59999, None)", "j <= TimestampMillisecond(10000, None)", "k >= Utf8(\"a\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -119,6 +121,7 @@ Affected Rows: 0 -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED tql analyze (0, 10, '1s') 100 - (avg by (k) (irate(t[1m])) * 100); +-+-+-+ @@ -141,10 +144,10 @@ tql analyze (0, 10, '1s') 100 - (avg by (k) (irate(t[1m])) * 100); |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_SortPreservingMergeExec: [k@2 ASC, l@3 ASC, j@1 ASC] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["i", "j", "k", "l"], "filters": ["j >= TimestampMillisecond(-59999, None)", "j <= TimestampMillisecond(10000, None)", "i < Float64(1) OR i IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_SortPreservingMergeExec: [k@2 ASC, l@3 ASC, j@1 ASC] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["i", "j", "k", "l"], "filters": ["j >= TimestampMillisecond(-59999, None)", "j <= TimestampMillisecond(10000, None)", "i >= Float64(1)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ diff --git a/tests/cases/standalone/common/tql/partition.sql b/tests/cases/standalone/common/tql/partition.sql index 332caa4252e..d68ce313e87 100644 --- a/tests/cases/standalone/common/tql/partition.sql +++ b/tests/cases/standalone/common/tql/partition.sql @@ -12,6 +12,7 @@ create table t ( -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED tql analyze (0, 10, '1s') 100 - (avg by (k) (irate(t[1m])) * 100); drop table t; @@ -32,6 +33,7 @@ create table t ( -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED tql analyze (0, 10, '1s') 100 - (avg by (k) (irate(t[1m])) * 100); drop table t; @@ -52,6 +54,7 @@ create table t ( -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED tql analyze (0, 10, '1s') 100 - (avg by (k) (irate(t[1m])) * 100); drop table t; diff --git a/tests/cases/standalone/optimizer/count.result b/tests/cases/standalone/optimizer/count.result index 925e89eeb23..1ffeca4b69e 100644 --- a/tests/cases/standalone/optimizer/count.result +++ b/tests/cases/standalone/optimizer/count.result @@ -89,6 +89,7 @@ Affected Rows: 5 -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (\{count\[count\]:)\d+(\}) {count[count]:REDACTED} -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select count(1) from count_where_bug; @@ -126,6 +127,7 @@ select count(1) from count_where_bug; -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select count(1) from count_where_bug where `tag` = 'b'; @@ -138,7 +140,7 @@ select count(1) from count_where_bug where `tag` = 'b'; | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[count(count_where_bug.ts) as count(Int64(1))] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[count(count_where_bug.ts) as count(Int64(1))] REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["ts"], "filters": ["tag = Dictionary(UInt32, Utf8(\"b\"))"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ @@ -167,6 +169,7 @@ select count(1) from count_where_bug where `tag` = 'b'; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (\{count\[count\]:)\d+(\}) {count[count]:REDACTED} -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select count(1) from count_where_bug where ts > '2024-09-06T06:00:04Z'; @@ -179,7 +182,7 @@ select count(1) from count_where_bug where ts > '2024-09-06T06:00:04Z'; | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[count(count_where_bug.ts) as count(Int64(1))] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[count(count_where_bug.ts) as count(Int64(1))] REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED, "projection": ["ts"], "filters": ["ts > TimestampMillisecond(1725602404000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ @@ -206,6 +209,7 @@ select count(1) from count_where_bug where ts > '2024-09-06T06:00:04Z'; -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select count(1) from count_where_bug where num != 3; @@ -219,7 +223,7 @@ select count(1) from count_where_bug where num != 3; |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[count(count_where_bug.ts) as count(Int64(1))] REDACTED |_|_|_FilterExec: num@1 != 3, projection=[ts@0] REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["ts", "num"], "filters": ["num != Int64(3)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ @@ -269,6 +273,7 @@ Affected Rows: 5 -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (\{count\[count\]:)\d+(\}) {count[count]:REDACTED} -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select count(1) from count_where_bug; @@ -305,6 +310,7 @@ select count(1) from count_where_bug; -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select count(1) from count_where_bug where `tag` = 'b'; @@ -320,7 +326,7 @@ select count(1) from count_where_bug where `tag` = 'b'; | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(count_where_bug.ts)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__count_state(count_where_bug.ts)] REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["ts"], "filters": ["tag = Dictionary(UInt32, Utf8(\"b\"))", "tag > Utf8(\"a\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ @@ -332,6 +338,7 @@ select count(1) from count_where_bug where `tag` = 'b'; -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select count(1) from count_where_bug where `tag` = 'b'; @@ -347,7 +354,7 @@ select count(1) from count_where_bug where `tag` = 'b'; | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(count_where_bug.ts)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__count_state(count_where_bug.ts)] REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["ts"], "filters": ["tag = Dictionary(UInt32, Utf8(\"b\"))", "tag > Utf8(\"a\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ @@ -369,6 +376,7 @@ select count(1) from count_where_bug where `tag` = 'b'; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (\{count\[count]:)\d+(\}) {count[count]:REDACTED} -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select count(1) from count_where_bug where ts > '2024-09-06T06:00:04Z'; @@ -384,12 +392,12 @@ select count(1) from count_where_bug where ts > '2024-09-06T06:00:04Z'; | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(count_where_bug.ts)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__count_state(count_where_bug.ts)] REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED, "projection": ["ts"], "filters": ["ts > TimestampMillisecond(1725602404000, None)", "tag <= Utf8(\"a\") OR tag IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(count_where_bug.ts)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__count_state(count_where_bug.ts)] REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":REDACTED, "projection": ["ts"], "filters": ["ts > TimestampMillisecond(1725602404000, None)", "tag > Utf8(\"a\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ @@ -409,6 +417,7 @@ select count(1) from count_where_bug where ts > '2024-09-06T06:00:04Z'; -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select count(1) from count_where_bug where num != 3; @@ -425,13 +434,13 @@ select count(1) from count_where_bug where num != 3; |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__count_state(count_where_bug.ts)] REDACTED |_|_|_FilterExec: num@1 != 3, projection=[ts@0] REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["ts", "num"], "filters": ["num != Int64(3)", "tag <= Utf8(\"a\") OR tag IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(count_where_bug.ts)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__count_state(count_where_bug.ts)] REDACTED |_|_|_FilterExec: num@1 != 3, projection=[ts@0] REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["ts", "num"], "filters": ["num != Int64(3)", "tag > Utf8(\"a\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ diff --git a/tests/cases/standalone/optimizer/count.sql b/tests/cases/standalone/optimizer/count.sql index caa961b1f01..fc2494589f8 100644 --- a/tests/cases/standalone/optimizer/count.sql +++ b/tests/cases/standalone/optimizer/count.sql @@ -50,6 +50,7 @@ values ('a', '2024-09-06T06:00:01Z', 1), -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (\{count\[count\]:)\d+(\}) {count[count]:REDACTED} -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select count(1) from count_where_bug; @@ -69,6 +70,7 @@ select count(1) from count_where_bug; -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select count(1) from count_where_bug where `tag` = 'b'; @@ -90,6 +92,7 @@ select count(1) from count_where_bug where `tag` = 'b'; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (\{count\[count\]:)\d+(\}) {count[count]:REDACTED} -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select count(1) from count_where_bug where ts > '2024-09-06T06:00:04Z'; @@ -109,6 +112,7 @@ select count(1) from count_where_bug where ts > '2024-09-06T06:00:04Z'; -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select count(1) from count_where_bug where num != 3; @@ -145,6 +149,7 @@ values ('a', '2024-09-06T06:00:01Z', 1), -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (\{count\[count\]:)\d+(\}) {count[count]:REDACTED} -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select count(1) from count_where_bug; @@ -157,6 +162,7 @@ select count(1) from count_where_bug; -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select count(1) from count_where_bug where `tag` = 'b'; @@ -167,6 +173,7 @@ select count(1) from count_where_bug where `tag` = 'b'; -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select count(1) from count_where_bug where `tag` = 'b'; @@ -181,6 +188,7 @@ select count(1) from count_where_bug where `tag` = 'b'; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (\{count\[count]:)\d+(\}) {count[count]:REDACTED} -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select count(1) from count_where_bug where ts > '2024-09-06T06:00:04Z'; @@ -193,6 +201,7 @@ select count(1) from count_where_bug where ts > '2024-09-06T06:00:04Z'; -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select count(1) from count_where_bug where num != 3; diff --git a/tests/cases/standalone/optimizer/filter_column_pruning.result b/tests/cases/standalone/optimizer/filter_column_pruning.result index 3a16bef50df..f966b0504e4 100644 --- a/tests/cases/standalone/optimizer/filter_column_pruning.result +++ b/tests/cases/standalone/optimizer/filter_column_pruning.result @@ -44,6 +44,7 @@ SELECT cpu_usage FROM filter_prune_test WHERE host = 'host1' ORDER BY ts; -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT cpu_usage FROM filter_prune_test WHERE host = 'host1' ORDER BY ts; +-+-+-+ @@ -57,7 +58,7 @@ EXPLAIN ANALYZE VERBOSE SELECT cpu_usage FROM filter_prune_test WHERE host = 'ho |_|_|_ProjectionExec: expr=[cpu_usage@1 as cpu_usage, ts@0 as ts] REDACTED |_|_|_WindowedSortExec: expr=ts@0 ASC NULLS LAST num_ranges=REDACTED REDACTED |_|_|_PartSortExec: expr=ts@0 ASC NULLS LAST num_ranges=REDACTED REDACTED -|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["ts", "cpu_usage"], "filters": ["host = Dictionary(UInt32, Utf8(\"host1\"))"], "flat_format": REDACTED, "REDACTED +|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["ts", "cpu_usage"], "filters": ["host = Dictionary(UInt32, Utf8(\"host1\"))"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 3_| +-+-+-+ @@ -80,6 +81,7 @@ SELECT mem_usage, cpu_usage FROM filter_prune_test WHERE host = 'host1' AND `reg -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT mem_usage, cpu_usage FROM filter_prune_test WHERE host = 'host1' AND `region` = 'us-east' ORDER BY ts; +-+-+-+ @@ -93,7 +95,7 @@ EXPLAIN ANALYZE VERBOSE SELECT mem_usage, cpu_usage FROM filter_prune_test WHERE |_|_|_ProjectionExec: expr=[mem_usage@2 as mem_usage, cpu_usage@1 as cpu_usage, ts@0 as ts] REDACTED |_|_|_WindowedSortExec: expr=ts@0 ASC NULLS LAST num_ranges=REDACTED REDACTED |_|_|_PartSortExec: expr=ts@0 ASC NULLS LAST num_ranges=REDACTED REDACTED -|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["ts", "cpu_usage", "mem_usage"], "filters": ["host = Dictionary(UInt32, Utf8(\"host1\"))", "region = Dictionary(UInt32, Utf8(\"us-east\"))"], "flat_format": REDACTED, "REDACTED +|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["ts", "cpu_usage", "mem_usage"], "filters": ["host = Dictionary(UInt32, Utf8(\"host1\"))", "region = Dictionary(UInt32, Utf8(\"us-east\"))"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 2_| +-+-+-+ @@ -118,6 +120,7 @@ SELECT host, `region` FROM filter_prune_test WHERE cpu_usage > 20.0 ORDER BY ts; -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT host, `region` FROM filter_prune_test WHERE cpu_usage > 20.0 ORDER BY ts; +-+-+-+ @@ -131,7 +134,7 @@ EXPLAIN ANALYZE VERBOSE SELECT host, `region` FROM filter_prune_test WHERE cpu_u |_|_|_WindowedSortExec: expr=ts@2 ASC NULLS LAST num_ranges=REDACTED REDACTED |_|_|_PartSortExec: expr=ts@2 ASC NULLS LAST num_ranges=REDACTED REDACTED |_|_|_FilterExec: cpu_usage@3 > 20, projection=[host@1, region@2, ts@0] REDACTED -|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["ts", "host", "region", "cpu_usage"], "filters": ["cpu_usage > Float64(20)"], "flat_format": REDACTED, "REDACTED +|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["ts", "host", "region", "cpu_usage"], "filters": ["cpu_usage > Float64(20)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -156,6 +159,7 @@ SELECT host, cpu_usage FROM filter_prune_test WHERE ts > 2000 ORDER BY ts; -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT host, cpu_usage FROM filter_prune_test WHERE ts > 2000 ORDER BY ts; +-+-+-+ @@ -169,7 +173,7 @@ EXPLAIN ANALYZE VERBOSE SELECT host, cpu_usage FROM filter_prune_test WHERE ts > |_|_|_ProjectionExec: expr=[host@1 as host, cpu_usage@2 as cpu_usage, ts@0 as ts] REDACTED |_|_|_WindowedSortExec: expr=ts@0 ASC NULLS LAST num_ranges=REDACTED REDACTED |_|_|_PartSortExec: expr=ts@0 ASC NULLS LAST num_ranges=REDACTED REDACTED -|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["ts", "host", "cpu_usage"], "filters": ["ts > TimestampMillisecond(2000, None)"], "flat_format": REDACTED, "REDACTED +|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["ts", "host", "cpu_usage"], "filters": ["ts > TimestampMillisecond(2000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -192,6 +196,7 @@ SELECT mem_usage FROM filter_prune_test WHERE host = 'host1' AND cpu_usage > 10. -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT mem_usage FROM filter_prune_test WHERE host = 'host1' AND cpu_usage > 10.0 AND `region` = 'us-east' ORDER BY ts; +-+-+-+ @@ -205,7 +210,7 @@ EXPLAIN ANALYZE VERBOSE SELECT mem_usage FROM filter_prune_test WHERE host = 'ho |_|_|_WindowedSortExec: expr=ts@1 ASC NULLS LAST num_ranges=REDACTED REDACTED |_|_|_PartSortExec: expr=ts@1 ASC NULLS LAST num_ranges=REDACTED REDACTED |_|_|_FilterExec: cpu_usage@1 > 10, projection=[mem_usage@2, ts@0] REDACTED -|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["ts", "cpu_usage", "mem_usage"], "filters": ["host = Dictionary(UInt32, Utf8(\"host1\"))", "cpu_usage > Float64(10)", "region = Dictionary(UInt32, Utf8(\"us-east\"))"], "flat_format": REDACTED, "REDACTED +|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["ts", "cpu_usage", "mem_usage"], "filters": ["host = Dictionary(UInt32, Utf8(\"host1\"))", "cpu_usage > Float64(10)", "region = Dictionary(UInt32, Utf8(\"us-east\"))"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 2_| +-+-+-+ @@ -227,6 +232,7 @@ SELECT SUM(cpu_usage) as total_cpu FROM filter_prune_test WHERE host = 'host1'; -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT SUM(cpu_usage) as total_cpu FROM filter_prune_test WHERE host = 'host1'; +-+-+-+ @@ -239,7 +245,7 @@ EXPLAIN ANALYZE VERBOSE SELECT SUM(cpu_usage) as total_cpu FROM filter_prune_tes |_|_|_AggregateExec: mode=Final, gby=[], aggr=[sum(filter_prune_test.cpu_usage)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[sum(filter_prune_test.cpu_usage)] REDACTED -|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["cpu_usage"], "filters": ["host = Dictionary(UInt32, Utf8(\"host1\"))"], "flat_format": REDACTED, "REDACTED +|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["cpu_usage"], "filters": ["host = Dictionary(UInt32, Utf8(\"host1\"))"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ @@ -263,6 +269,7 @@ SELECT `region`, AVG(cpu_usage) as avg_cpu FROM filter_prune_test WHERE mem_usag -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT `region`, AVG(cpu_usage) as avg_cpu FROM filter_prune_test WHERE mem_usage > 25.0 GROUP BY `region` ORDER BY `region`; +-+-+-+ @@ -278,7 +285,7 @@ EXPLAIN ANALYZE VERBOSE SELECT `region`, AVG(cpu_usage) as avg_cpu FROM filter_p |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[region@0 as region], aggr=[avg(filter_prune_test.cpu_usage)] REDACTED |_|_|_FilterExec: mem_usage@2 > 25, projection=[region@0, cpu_usage@1] REDACTED -|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["region", "cpu_usage", "mem_usage"], "filters": ["mem_usage > Float64(25)"], "flat_format": REDACTED, "REDACTED +|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["region", "cpu_usage", "mem_usage"], "filters": ["mem_usage > Float64(25)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 3_| +-+-+-+ @@ -304,6 +311,7 @@ SELECT cpu_usage FROM filter_prune_test WHERE host IN ('host1', 'host2') ORDER B -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT cpu_usage FROM filter_prune_test WHERE host IN ('host1', 'host2') ORDER BY ts; +-+-+-+ @@ -317,7 +325,7 @@ EXPLAIN ANALYZE VERBOSE SELECT cpu_usage FROM filter_prune_test WHERE host IN (' |_|_|_ProjectionExec: expr=[cpu_usage@1 as cpu_usage, ts@0 as ts] REDACTED |_|_|_WindowedSortExec: expr=ts@0 ASC NULLS LAST num_ranges=REDACTED REDACTED |_|_|_PartSortExec: expr=ts@0 ASC NULLS LAST num_ranges=REDACTED REDACTED -|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["ts", "cpu_usage"], "filters": ["host = Dictionary(UInt32, Utf8(\"host1\")) OR host = Dictionary(UInt32, Utf8(\"host2\"))"], "flat_format": REDACTED, "REDACTED +|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["ts", "cpu_usage"], "filters": ["host = Dictionary(UInt32, Utf8(\"host1\")) OR host = Dictionary(UInt32, Utf8(\"host2\"))"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 5_| +-+-+-+ @@ -341,6 +349,7 @@ SELECT host FROM filter_prune_test WHERE cpu_usage BETWEEN 15.0 AND 30.0 ORDER B -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT host FROM filter_prune_test WHERE cpu_usage BETWEEN 15.0 AND 30.0 ORDER BY ts; +-+-+-+ @@ -354,7 +363,7 @@ EXPLAIN ANALYZE VERBOSE SELECT host FROM filter_prune_test WHERE cpu_usage BETWE |_|_|_WindowedSortExec: expr=ts@1 ASC NULLS LAST num_ranges=REDACTED REDACTED |_|_|_PartSortExec: expr=ts@1 ASC NULLS LAST num_ranges=REDACTED REDACTED |_|_|_FilterExec: cpu_usage@2 >= 15 AND cpu_usage@2 <= 30, projection=[host@1, ts@0] REDACTED -|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["ts", "host", "cpu_usage"], "filters": ["cpu_usage >= Float64(15)", "cpu_usage <= Float64(30)"], "flat_format": REDACTED, "REDACTED +|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["ts", "host", "cpu_usage"], "filters": ["cpu_usage >= Float64(15)", "cpu_usage <= Float64(30)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 3_| +-+-+-+ @@ -380,6 +389,7 @@ SELECT cpu_usage FROM filter_prune_test WHERE host LIKE 'host%' AND `region` LIK -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT cpu_usage FROM filter_prune_test WHERE host LIKE 'host%' AND `region` LIKE 'us-%' ORDER BY ts; +-+-+-+ @@ -393,7 +403,7 @@ EXPLAIN ANALYZE VERBOSE SELECT cpu_usage FROM filter_prune_test WHERE host LIKE |_|_|_WindowedSortExec: expr=ts@1 ASC NULLS LAST num_ranges=REDACTED REDACTED |_|_|_PartSortExec: expr=ts@1 ASC NULLS LAST num_ranges=REDACTED REDACTED |_|_|_FilterExec: host@1 LIKE host% AND region@2 LIKE us-%, projection=[cpu_usage@3, ts@0] REDACTED -|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["ts", "host", "region", "cpu_usage"], "filters": ["host LIKE Dictionary(UInt32, Utf8(\"host%\"))", "region LIKE Dictionary(UInt32, Utf8(\"us-%\"))"], "flat_format": REDACTED, "REDACTED +|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["ts", "host", "region", "cpu_usage"], "filters": ["host LIKE Dictionary(UInt32, Utf8(\"host%\"))", "region LIKE Dictionary(UInt32, Utf8(\"us-%\"))"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 5_| +-+-+-+ @@ -418,6 +428,7 @@ SELECT cpu_usage FROM filter_prune_test WHERE host = 'host1' OR `region` = 'eu-w -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT cpu_usage FROM filter_prune_test WHERE host = 'host1' OR `region` = 'eu-west' ORDER BY ts; +-+-+-+ @@ -431,7 +442,7 @@ EXPLAIN ANALYZE VERBOSE SELECT cpu_usage FROM filter_prune_test WHERE host = 'ho |_|_|_WindowedSortExec: expr=ts@1 ASC NULLS LAST num_ranges=REDACTED REDACTED |_|_|_PartSortExec: expr=ts@1 ASC NULLS LAST num_ranges=REDACTED REDACTED |_|_|_FilterExec: host@1 = host1 OR region@2 = eu-west, projection=[cpu_usage@3, ts@0] REDACTED -|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["ts", "host", "region", "cpu_usage"], "filters": ["host = Dictionary(UInt32, Utf8(\"host1\")) OR region = Dictionary(UInt32, Utf8(\"eu-west\"))"], "flat_format": REDACTED, "REDACTED +|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["ts", "host", "region", "cpu_usage"], "filters": ["host = Dictionary(UInt32, Utf8(\"host1\")) OR region = Dictionary(UInt32, Utf8(\"eu-west\"))"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -454,6 +465,7 @@ SELECT cpu_usage FROM (SELECT * FROM filter_prune_test WHERE host = 'host1') sub -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT cpu_usage FROM (SELECT * FROM filter_prune_test WHERE host = 'host1') sub WHERE `region` = 'us-east' ORDER BY ts; +-+-+-+ @@ -467,7 +479,7 @@ EXPLAIN ANALYZE VERBOSE SELECT cpu_usage FROM (SELECT * FROM filter_prune_test W |_|_|_ProjectionExec: expr=[cpu_usage@1 as cpu_usage, ts@0 as ts] REDACTED |_|_|_WindowedSortExec: expr=ts@0 ASC NULLS LAST num_ranges=REDACTED REDACTED |_|_|_PartSortExec: expr=ts@0 ASC NULLS LAST num_ranges=REDACTED REDACTED -|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["ts", "cpu_usage"], "filters": ["host = Dictionary(UInt32, Utf8(\"host1\"))", "region = Dictionary(UInt32, Utf8(\"us-east\"))"], "flat_format": REDACTED, "REDACTED +|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["ts", "cpu_usage"], "filters": ["host = Dictionary(UInt32, Utf8(\"host1\"))", "region = Dictionary(UInt32, Utf8(\"us-east\"))"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 2_| +-+-+-+ @@ -501,6 +513,7 @@ SELECT cpu_usage FROM filter_prune_test WHERE host = 'host1' ORDER BY ts; -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT cpu_usage FROM filter_prune_test WHERE host = 'host1' ORDER BY ts; +-+-+-+ @@ -514,7 +527,7 @@ EXPLAIN ANALYZE VERBOSE SELECT cpu_usage FROM filter_prune_test WHERE host = 'ho |_|_|_ProjectionExec: expr=[cpu_usage@1 as cpu_usage, ts@0 as ts] REDACTED |_|_|_WindowedSortExec: expr=ts@0 ASC NULLS LAST num_ranges=REDACTED REDACTED |_|_|_PartSortExec: expr=ts@0 ASC NULLS LAST num_ranges=REDACTED REDACTED -|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["ts", "cpu_usage"], "filters": ["host = Dictionary(UInt32, Utf8(\"host1\"))"], \"file\":REDACTED +|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["ts", "cpu_usage"], "filters": ["host = Dictionary(UInt32, Utf8(\"host1\"))"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 3_| +-+-+-+ @@ -538,6 +551,7 @@ SELECT mem_usage, cpu_usage FROM filter_prune_test WHERE host = 'host1' AND `reg -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT mem_usage, cpu_usage FROM filter_prune_test WHERE host = 'host1' AND `region` = 'us-east' ORDER BY ts; +-+-+-+ @@ -551,7 +565,7 @@ EXPLAIN ANALYZE VERBOSE SELECT mem_usage, cpu_usage FROM filter_prune_test WHERE |_|_|_ProjectionExec: expr=[mem_usage@2 as mem_usage, cpu_usage@1 as cpu_usage, ts@0 as ts] REDACTED |_|_|_WindowedSortExec: expr=ts@0 ASC NULLS LAST num_ranges=REDACTED REDACTED |_|_|_PartSortExec: expr=ts@0 ASC NULLS LAST num_ranges=REDACTED REDACTED -|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["ts", "cpu_usage", "mem_usage"], "filters": ["host = Dictionary(UInt32, Utf8(\"host1\"))", "region = Dictionary(UInt32, Utf8(\"us-east\"))"], \"file\":REDACTED +|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["ts", "cpu_usage", "mem_usage"], "filters": ["host = Dictionary(UInt32, Utf8(\"host1\"))", "region = Dictionary(UInt32, Utf8(\"us-east\"))"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 2_| +-+-+-+ @@ -577,6 +591,7 @@ SELECT host, `region` FROM filter_prune_test WHERE cpu_usage > 20.0 ORDER BY ts; -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT host, `region` FROM filter_prune_test WHERE cpu_usage > 20.0 ORDER BY ts; +-+-+-+ @@ -590,7 +605,7 @@ EXPLAIN ANALYZE VERBOSE SELECT host, `region` FROM filter_prune_test WHERE cpu_u |_|_|_WindowedSortExec: expr=ts@2 ASC NULLS LAST num_ranges=REDACTED REDACTED |_|_|_PartSortExec: expr=ts@2 ASC NULLS LAST num_ranges=REDACTED REDACTED |_|_|_FilterExec: cpu_usage@3 > 20, projection=[host@1, region@2, ts@0] REDACTED -|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["ts", "host", "region", "cpu_usage"], "filters": ["cpu_usage > Float64(20)"], \"file\":REDACTED +|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["ts", "host", "region", "cpu_usage"], "filters": ["cpu_usage > Float64(20)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -616,6 +631,7 @@ SELECT host, cpu_usage FROM filter_prune_test WHERE ts > 2000 ORDER BY ts; -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT host, cpu_usage FROM filter_prune_test WHERE ts > 2000 ORDER BY ts; +-+-+-+ @@ -629,7 +645,7 @@ EXPLAIN ANALYZE VERBOSE SELECT host, cpu_usage FROM filter_prune_test WHERE ts > |_|_|_ProjectionExec: expr=[host@1 as host, cpu_usage@2 as cpu_usage, ts@0 as ts] REDACTED |_|_|_WindowedSortExec: expr=ts@0 ASC NULLS LAST num_ranges=REDACTED REDACTED |_|_|_PartSortExec: expr=ts@0 ASC NULLS LAST num_ranges=REDACTED REDACTED -|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["ts", "host", "cpu_usage"], "filters": ["ts > TimestampMillisecond(2000, None)"], \"file\":REDACTED +|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["ts", "host", "cpu_usage"], "filters": ["ts > TimestampMillisecond(2000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -653,6 +669,7 @@ SELECT mem_usage FROM filter_prune_test WHERE host = 'host1' AND cpu_usage > 10. -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT mem_usage FROM filter_prune_test WHERE host = 'host1' AND cpu_usage > 10.0 AND `region` = 'us-east' ORDER BY ts; +-+-+-+ @@ -666,7 +683,7 @@ EXPLAIN ANALYZE VERBOSE SELECT mem_usage FROM filter_prune_test WHERE host = 'ho |_|_|_WindowedSortExec: expr=ts@1 ASC NULLS LAST num_ranges=REDACTED REDACTED |_|_|_PartSortExec: expr=ts@1 ASC NULLS LAST num_ranges=REDACTED REDACTED |_|_|_FilterExec: cpu_usage@1 > 10, projection=[mem_usage@2, ts@0] REDACTED -|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["ts", "cpu_usage", "mem_usage"], "filters": ["host = Dictionary(UInt32, Utf8(\"host1\"))", "cpu_usage > Float64(10)", "region = Dictionary(UInt32, Utf8(\"us-east\"))"], \"file\":REDACTED +|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["ts", "cpu_usage", "mem_usage"], "filters": ["host = Dictionary(UInt32, Utf8(\"host1\"))", "cpu_usage > Float64(10)", "region = Dictionary(UInt32, Utf8(\"us-east\"))"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 2_| +-+-+-+ @@ -689,6 +706,7 @@ SELECT SUM(cpu_usage) as total_cpu FROM filter_prune_test WHERE host = 'host1'; -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT SUM(cpu_usage) as total_cpu FROM filter_prune_test WHERE host = 'host1'; +-+-+-+ @@ -701,7 +719,7 @@ EXPLAIN ANALYZE VERBOSE SELECT SUM(cpu_usage) as total_cpu FROM filter_prune_tes |_|_|_AggregateExec: mode=Final, gby=[], aggr=[sum(filter_prune_test.cpu_usage)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[sum(filter_prune_test.cpu_usage)] REDACTED -|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["cpu_usage"], "filters": ["host = Dictionary(UInt32, Utf8(\"host1\"))"], \"file\":REDACTED +|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["cpu_usage"], "filters": ["host = Dictionary(UInt32, Utf8(\"host1\"))"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ @@ -726,6 +744,7 @@ SELECT `region`, AVG(cpu_usage) as avg_cpu FROM filter_prune_test WHERE mem_usag -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT `region`, AVG(cpu_usage) as avg_cpu FROM filter_prune_test WHERE mem_usage > 25.0 GROUP BY `region` ORDER BY `region`; +-+-+-+ @@ -741,7 +760,7 @@ EXPLAIN ANALYZE VERBOSE SELECT `region`, AVG(cpu_usage) as avg_cpu FROM filter_p |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[region@0 as region], aggr=[avg(filter_prune_test.cpu_usage)] REDACTED |_|_|_FilterExec: mem_usage@2 > 25, projection=[region@0, cpu_usage@1] REDACTED -|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["region", "cpu_usage", "mem_usage"], "filters": ["mem_usage > Float64(25)"], \"file\":REDACTED +|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["region", "cpu_usage", "mem_usage"], "filters": ["mem_usage > Float64(25)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 3_| +-+-+-+ @@ -768,6 +787,7 @@ SELECT cpu_usage FROM filter_prune_test WHERE host IN ('host1', 'host2') ORDER B -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT cpu_usage FROM filter_prune_test WHERE host IN ('host1', 'host2') ORDER BY ts; +-+-+-+ @@ -781,7 +801,7 @@ EXPLAIN ANALYZE VERBOSE SELECT cpu_usage FROM filter_prune_test WHERE host IN (' |_|_|_ProjectionExec: expr=[cpu_usage@1 as cpu_usage, ts@0 as ts] REDACTED |_|_|_WindowedSortExec: expr=ts@0 ASC NULLS LAST num_ranges=REDACTED REDACTED |_|_|_PartSortExec: expr=ts@0 ASC NULLS LAST num_ranges=REDACTED REDACTED -|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["ts", "cpu_usage"], "filters": ["host = Dictionary(UInt32, Utf8(\"host1\")) OR host = Dictionary(UInt32, Utf8(\"host2\"))"], \"file\":REDACTED +|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["ts", "cpu_usage"], "filters": ["host = Dictionary(UInt32, Utf8(\"host1\")) OR host = Dictionary(UInt32, Utf8(\"host2\"))"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 5_| +-+-+-+ @@ -806,6 +826,7 @@ SELECT host FROM filter_prune_test WHERE cpu_usage BETWEEN 15.0 AND 30.0 ORDER B -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT host FROM filter_prune_test WHERE cpu_usage BETWEEN 15.0 AND 30.0 ORDER BY ts; +-+-+-+ @@ -819,7 +840,7 @@ EXPLAIN ANALYZE VERBOSE SELECT host FROM filter_prune_test WHERE cpu_usage BETWE |_|_|_WindowedSortExec: expr=ts@1 ASC NULLS LAST num_ranges=REDACTED REDACTED |_|_|_PartSortExec: expr=ts@1 ASC NULLS LAST num_ranges=REDACTED REDACTED |_|_|_FilterExec: cpu_usage@2 >= 15 AND cpu_usage@2 <= 30, projection=[host@1, ts@0] REDACTED -|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["ts", "host", "cpu_usage"], "filters": ["cpu_usage >= Float64(15)", "cpu_usage <= Float64(30)"], \"file\":REDACTED +|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["ts", "host", "cpu_usage"], "filters": ["cpu_usage >= Float64(15)", "cpu_usage <= Float64(30)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 3_| +-+-+-+ @@ -846,6 +867,7 @@ SELECT cpu_usage FROM filter_prune_test WHERE host LIKE 'host%' AND `region` LIK -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT cpu_usage FROM filter_prune_test WHERE host LIKE 'host%' AND `region` LIKE 'us-%' ORDER BY ts; +-+-+-+ @@ -859,7 +881,7 @@ EXPLAIN ANALYZE VERBOSE SELECT cpu_usage FROM filter_prune_test WHERE host LIKE |_|_|_WindowedSortExec: expr=ts@1 ASC NULLS LAST num_ranges=REDACTED REDACTED |_|_|_PartSortExec: expr=ts@1 ASC NULLS LAST num_ranges=REDACTED REDACTED |_|_|_FilterExec: host@1 LIKE host% AND region@2 LIKE us-%, projection=[cpu_usage@3, ts@0] REDACTED -|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["ts", "host", "region", "cpu_usage"], "filters": ["host LIKE Dictionary(UInt32, Utf8(\"host%\"))", "region LIKE Dictionary(UInt32, Utf8(\"us-%\"))"], \"file\":REDACTED +|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["ts", "host", "region", "cpu_usage"], "filters": ["host LIKE Dictionary(UInt32, Utf8(\"host%\"))", "region LIKE Dictionary(UInt32, Utf8(\"us-%\"))"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 5_| +-+-+-+ @@ -885,6 +907,7 @@ SELECT cpu_usage FROM filter_prune_test WHERE host = 'host1' OR `region` = 'eu-w -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT cpu_usage FROM filter_prune_test WHERE host = 'host1' OR `region` = 'eu-west' ORDER BY ts; +-+-+-+ @@ -898,7 +921,7 @@ EXPLAIN ANALYZE VERBOSE SELECT cpu_usage FROM filter_prune_test WHERE host = 'ho |_|_|_WindowedSortExec: expr=ts@1 ASC NULLS LAST num_ranges=REDACTED REDACTED |_|_|_PartSortExec: expr=ts@1 ASC NULLS LAST num_ranges=REDACTED REDACTED |_|_|_FilterExec: host@1 = host1 OR region@2 = eu-west, projection=[cpu_usage@3, ts@0] REDACTED -|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["ts", "host", "region", "cpu_usage"], "filters": ["host = Dictionary(UInt32, Utf8(\"host1\")) OR region = Dictionary(UInt32, Utf8(\"eu-west\"))"], \"file\":REDACTED +|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["ts", "host", "region", "cpu_usage"], "filters": ["host = Dictionary(UInt32, Utf8(\"host1\")) OR region = Dictionary(UInt32, Utf8(\"eu-west\"))"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -922,6 +945,7 @@ SELECT cpu_usage FROM (SELECT * FROM filter_prune_test WHERE host = 'host1') sub -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT cpu_usage FROM (SELECT * FROM filter_prune_test WHERE host = 'host1') sub WHERE `region` = 'us-east' ORDER BY ts; +-+-+-+ @@ -935,7 +959,7 @@ EXPLAIN ANALYZE VERBOSE SELECT cpu_usage FROM (SELECT * FROM filter_prune_test W |_|_|_ProjectionExec: expr=[cpu_usage@1 as cpu_usage, ts@0 as ts] REDACTED |_|_|_WindowedSortExec: expr=ts@0 ASC NULLS LAST num_ranges=REDACTED REDACTED |_|_|_PartSortExec: expr=ts@0 ASC NULLS LAST num_ranges=REDACTED REDACTED -|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["ts", "cpu_usage"], "filters": ["host = Dictionary(UInt32, Utf8(\"host1\"))", "region = Dictionary(UInt32, Utf8(\"us-east\"))"], \"file\":REDACTED +|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["ts", "cpu_usage"], "filters": ["host = Dictionary(UInt32, Utf8(\"host1\"))", "region = Dictionary(UInt32, Utf8(\"us-east\"))"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 2_| +-+-+-+ @@ -1011,6 +1035,7 @@ SELECT field1 FROM filter_prune_files WHERE tag_key = 'a' ORDER BY ts; -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT field1 FROM filter_prune_files WHERE tag_key = 'a' ORDER BY ts; +-+-+-+ @@ -1024,7 +1049,7 @@ EXPLAIN ANALYZE VERBOSE SELECT field1 FROM filter_prune_files WHERE tag_key = 'a |_|_|_ProjectionExec: expr=[field1@1 as field1, ts@0 as ts] REDACTED |_|_|_WindowedSortExec: expr=ts@0 ASC NULLS LAST num_ranges=REDACTED REDACTED |_|_|_PartSortExec: expr=ts@0 ASC NULLS LAST num_ranges=REDACTED REDACTED -|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":3, "mem_ranges":1, "files":2, "file_ranges":2}, "projection": ["ts", "field1"], "filters": ["tag_key = Dictionary(UInt32, Utf8(\"a\"))"], \"file\":REDACTED +|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":3, "mem_ranges":1, "files":2, "file_ranges":2}, "projection": ["ts", "field1"], "filters": ["tag_key = Dictionary(UInt32, Utf8(\"a\"))"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 2_| +-+-+-+ @@ -1051,6 +1076,7 @@ SELECT field1 FROM filter_prune_files WHERE field2 > 5.0 ORDER BY ts; -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT field1 FROM filter_prune_files WHERE field2 > 5.0 ORDER BY ts; +-+-+-+ @@ -1064,7 +1090,7 @@ EXPLAIN ANALYZE VERBOSE SELECT field1 FROM filter_prune_files WHERE field2 > 5.0 |_|_|_WindowedSortExec: expr=ts@1 ASC NULLS LAST num_ranges=REDACTED REDACTED |_|_|_PartSortExec: expr=ts@1 ASC NULLS LAST num_ranges=REDACTED REDACTED |_|_|_FilterExec: field2@2 > 5, projection=[field1@1, ts@0] REDACTED -|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":3, "mem_ranges":1, "files":2, "file_ranges":2}, "projection": ["ts", "field1", "field2"], "filters": ["field2 > Float64(5)"], \"file\":REDACTED +|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":3, "mem_ranges":1, "files":2, "file_ranges":2}, "projection": ["ts", "field1", "field2"], "filters": ["field2 > Float64(5)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 5_| +-+-+-+ @@ -1088,6 +1114,7 @@ SELECT field3 FROM filter_prune_files WHERE tag_key = 'b' AND field2 > 7.0 ORDER -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT field3 FROM filter_prune_files WHERE tag_key = 'b' AND field2 > 7.0 ORDER BY ts; +-+-+-+ @@ -1101,7 +1128,7 @@ EXPLAIN ANALYZE VERBOSE SELECT field3 FROM filter_prune_files WHERE tag_key = 'b |_|_|_WindowedSortExec: expr=ts@1 ASC NULLS LAST num_ranges=REDACTED REDACTED |_|_|_PartSortExec: expr=ts@1 ASC NULLS LAST num_ranges=REDACTED REDACTED |_|_|_FilterExec: field2@1 > 7, projection=[field3@2, ts@0] REDACTED -|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":3, "mem_ranges":1, "files":2, "file_ranges":2}, "projection": ["ts", "field2", "field3"], "filters": ["tag_key = Dictionary(UInt32, Utf8(\"b\"))", "field2 > Float64(7)"], \"file\":REDACTED +|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":3, "mem_ranges":1, "files":2, "file_ranges":2}, "projection": ["ts", "field2", "field3"], "filters": ["tag_key = Dictionary(UInt32, Utf8(\"b\"))", "field2 > Float64(7)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 2_| +-+-+-+ diff --git a/tests/cases/standalone/optimizer/filter_column_pruning.sql b/tests/cases/standalone/optimizer/filter_column_pruning.sql index a71db303631..5c8f73c0fff 100644 --- a/tests/cases/standalone/optimizer/filter_column_pruning.sql +++ b/tests/cases/standalone/optimizer/filter_column_pruning.sql @@ -33,6 +33,7 @@ SELECT cpu_usage FROM filter_prune_test WHERE host = 'host1' ORDER BY ts; -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT cpu_usage FROM filter_prune_test WHERE host = 'host1' ORDER BY ts; -- Filter on multiple columns (host, region) but only select value columns @@ -46,6 +47,7 @@ SELECT mem_usage, cpu_usage FROM filter_prune_test WHERE host = 'host1' AND `reg -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT mem_usage, cpu_usage FROM filter_prune_test WHERE host = 'host1' AND `region` = 'us-east' ORDER BY ts; -- Filter on value column (cpu_usage) not in projection @@ -59,6 +61,7 @@ SELECT host, `region` FROM filter_prune_test WHERE cpu_usage > 20.0 ORDER BY ts; -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT host, `region` FROM filter_prune_test WHERE cpu_usage > 20.0 ORDER BY ts; -- Filter on time index but time index not in projection @@ -72,6 +75,7 @@ SELECT host, cpu_usage FROM filter_prune_test WHERE ts > 2000 ORDER BY ts; -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT host, cpu_usage FROM filter_prune_test WHERE ts > 2000 ORDER BY ts; -- Complex filter: multiple columns in WHERE, only one in SELECT @@ -85,6 +89,7 @@ SELECT mem_usage FROM filter_prune_test WHERE host = 'host1' AND cpu_usage > 10. -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT mem_usage FROM filter_prune_test WHERE host = 'host1' AND cpu_usage > 10.0 AND `region` = 'us-east' ORDER BY ts; -- Aggregation with filter on non-projected column @@ -98,6 +103,7 @@ SELECT SUM(cpu_usage) as total_cpu FROM filter_prune_test WHERE host = 'host1'; -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT SUM(cpu_usage) as total_cpu FROM filter_prune_test WHERE host = 'host1'; -- GROUP BY with filter on column not in projection or grouping @@ -111,6 +117,7 @@ SELECT `region`, AVG(cpu_usage) as avg_cpu FROM filter_prune_test WHERE mem_usag -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT `region`, AVG(cpu_usage) as avg_cpu FROM filter_prune_test WHERE mem_usage > 25.0 GROUP BY `region` ORDER BY `region`; -- Filter with IN clause on non-projected column @@ -124,6 +131,7 @@ SELECT cpu_usage FROM filter_prune_test WHERE host IN ('host1', 'host2') ORDER B -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT cpu_usage FROM filter_prune_test WHERE host IN ('host1', 'host2') ORDER BY ts; -- Filter with BETWEEN on non-projected column @@ -137,6 +145,7 @@ SELECT host FROM filter_prune_test WHERE cpu_usage BETWEEN 15.0 AND 30.0 ORDER B -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT host FROM filter_prune_test WHERE cpu_usage BETWEEN 15.0 AND 30.0 ORDER BY ts; -- Filter with LIKE on non-projected column @@ -150,6 +159,7 @@ SELECT cpu_usage FROM filter_prune_test WHERE host LIKE 'host%' AND `region` LIK -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT cpu_usage FROM filter_prune_test WHERE host LIKE 'host%' AND `region` LIKE 'us-%' ORDER BY ts; -- Filter with OR conditions on non-projected columns @@ -163,6 +173,7 @@ SELECT cpu_usage FROM filter_prune_test WHERE host = 'host1' OR `region` = 'eu-w -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT cpu_usage FROM filter_prune_test WHERE host = 'host1' OR `region` = 'eu-west' ORDER BY ts; -- Subquery with filter on non-projected column @@ -176,6 +187,7 @@ SELECT cpu_usage FROM (SELECT * FROM filter_prune_test WHERE host = 'host1') sub -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT cpu_usage FROM (SELECT * FROM filter_prune_test WHERE host = 'host1') sub WHERE `region` = 'us-east' ORDER BY ts; -- Repeat with data flushed to verify column pruning works with on-disk data @@ -194,6 +206,7 @@ SELECT cpu_usage FROM filter_prune_test WHERE host = 'host1' ORDER BY ts; -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT cpu_usage FROM filter_prune_test WHERE host = 'host1' ORDER BY ts; -- Filter on multiple columns (host, region) but only select value columns @@ -208,6 +221,7 @@ SELECT mem_usage, cpu_usage FROM filter_prune_test WHERE host = 'host1' AND `reg -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT mem_usage, cpu_usage FROM filter_prune_test WHERE host = 'host1' AND `region` = 'us-east' ORDER BY ts; -- Filter on value column (cpu_usage) not in projection @@ -222,6 +236,7 @@ SELECT host, `region` FROM filter_prune_test WHERE cpu_usage > 20.0 ORDER BY ts; -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT host, `region` FROM filter_prune_test WHERE cpu_usage > 20.0 ORDER BY ts; -- Filter on time index but time index not in projection @@ -236,6 +251,7 @@ SELECT host, cpu_usage FROM filter_prune_test WHERE ts > 2000 ORDER BY ts; -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT host, cpu_usage FROM filter_prune_test WHERE ts > 2000 ORDER BY ts; -- Complex filter: multiple columns in WHERE, only one in SELECT @@ -250,6 +266,7 @@ SELECT mem_usage FROM filter_prune_test WHERE host = 'host1' AND cpu_usage > 10. -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT mem_usage FROM filter_prune_test WHERE host = 'host1' AND cpu_usage > 10.0 AND `region` = 'us-east' ORDER BY ts; -- Aggregation with filter on non-projected column @@ -264,6 +281,7 @@ SELECT SUM(cpu_usage) as total_cpu FROM filter_prune_test WHERE host = 'host1'; -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT SUM(cpu_usage) as total_cpu FROM filter_prune_test WHERE host = 'host1'; -- GROUP BY with filter on column not in projection or grouping @@ -278,6 +296,7 @@ SELECT `region`, AVG(cpu_usage) as avg_cpu FROM filter_prune_test WHERE mem_usag -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT `region`, AVG(cpu_usage) as avg_cpu FROM filter_prune_test WHERE mem_usage > 25.0 GROUP BY `region` ORDER BY `region`; -- Filter with IN clause on non-projected column @@ -292,6 +311,7 @@ SELECT cpu_usage FROM filter_prune_test WHERE host IN ('host1', 'host2') ORDER B -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT cpu_usage FROM filter_prune_test WHERE host IN ('host1', 'host2') ORDER BY ts; -- Filter with BETWEEN on non-projected column @@ -306,6 +326,7 @@ SELECT host FROM filter_prune_test WHERE cpu_usage BETWEEN 15.0 AND 30.0 ORDER B -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT host FROM filter_prune_test WHERE cpu_usage BETWEEN 15.0 AND 30.0 ORDER BY ts; -- Filter with LIKE on non-projected column @@ -320,6 +341,7 @@ SELECT cpu_usage FROM filter_prune_test WHERE host LIKE 'host%' AND `region` LIK -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT cpu_usage FROM filter_prune_test WHERE host LIKE 'host%' AND `region` LIKE 'us-%' ORDER BY ts; -- Filter with OR conditions on non-projected columns @@ -334,6 +356,7 @@ SELECT cpu_usage FROM filter_prune_test WHERE host = 'host1' OR `region` = 'eu-w -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT cpu_usage FROM filter_prune_test WHERE host = 'host1' OR `region` = 'eu-west' ORDER BY ts; -- Subquery with filter on non-projected column @@ -348,6 +371,7 @@ SELECT cpu_usage FROM (SELECT * FROM filter_prune_test WHERE host = 'host1') sub -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT cpu_usage FROM (SELECT * FROM filter_prune_test WHERE host = 'host1') sub WHERE `region` = 'us-east' ORDER BY ts; DROP TABLE filter_prune_test; @@ -392,6 +416,7 @@ SELECT field1 FROM filter_prune_files WHERE tag_key = 'a' ORDER BY ts; -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT field1 FROM filter_prune_files WHERE tag_key = 'a' ORDER BY ts; -- Query: filter on field2 (not in projection), select only field1 @@ -406,6 +431,7 @@ SELECT field1 FROM filter_prune_files WHERE field2 > 5.0 ORDER BY ts; -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT field1 FROM filter_prune_files WHERE field2 > 5.0 ORDER BY ts; -- Complex: filter on tag and field2, select only field3 @@ -420,6 +446,7 @@ SELECT field3 FROM filter_prune_files WHERE tag_key = 'b' AND field2 > 7.0 ORDER -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT field3 FROM filter_prune_files WHERE tag_key = 'b' AND field2 > 7.0 ORDER BY ts; DROP TABLE filter_prune_files; diff --git a/tests/cases/standalone/optimizer/first_value_advance.result b/tests/cases/standalone/optimizer/first_value_advance.result index f3b24db6d45..4ee877d14b7 100644 --- a/tests/cases/standalone/optimizer/first_value_advance.result +++ b/tests/cases/standalone/optimizer/first_value_advance.result @@ -95,6 +95,8 @@ explain select -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE (partitioning.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select first_value(host order by ts), @@ -113,7 +115,7 @@ explain analyze |_|_|_AggregateExec: mode=FinalPartitioned, gby=[host@0 as host], aggr=[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST], first_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[host@1 as host], aggr=[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST], first_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1} REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["ts", "host", "not_pk", "val"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -153,6 +155,8 @@ explain select first_value(ts order by ts) from t; -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE (partitioning.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select first_value(ts order by ts) from t; @@ -165,7 +169,7 @@ explain analyze | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[first_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[first_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_UnorderedScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1} REDACTED +|_|_|_UnorderedScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["ts"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ @@ -255,6 +259,8 @@ order by ordered_host; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select first_value(host order by ts) as ordered_host, @@ -276,7 +282,7 @@ order by ordered_host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[host@0 as host], aggr=[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST], first_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[host@1 as host], aggr=[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST], first_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["ts", "host", "not_pk", "val"], "filters": ["host < Utf8(\"b\") OR host IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[ordered_host@0 ASC NULLS LAST], preserve_REDACTED @@ -284,7 +290,7 @@ order by ordered_host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[host@0 as host], aggr=[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST], first_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[host@1 as host], aggr=[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST], first_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["ts", "host", "not_pk", "val"], "filters": ["host >= Utf8(\"b\") AND (host < Utf8(\"d\") OR host IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 2_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[ordered_host@0 ASC NULLS LAST], preserve_REDACTED @@ -292,7 +298,7 @@ order by ordered_host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[host@0 as host], aggr=[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST], first_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[host@1 as host], aggr=[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST], first_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["ts", "host", "not_pk", "val"], "filters": ["host >= Utf8(\"d\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -337,6 +343,8 @@ explain select first_value(ts order by ts) from t; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select first_value(ts order by ts) from t; @@ -352,17 +360,17 @@ explain analyze | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__first_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__first_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["ts"], "filters": ["host < Utf8(\"b\") OR host IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=Final, gby=[], aggr=[__first_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__first_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["ts"], "filters": ["host >= Utf8(\"b\") AND (host < Utf8(\"d\") OR host IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 2_|_AggregateExec: mode=Final, gby=[], aggr=[__first_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__first_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["ts"], "filters": ["host >= Utf8(\"d\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ @@ -428,6 +436,8 @@ order by time_window, ordered_host; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select first_value(host order by ts) as ordered_host, @@ -453,17 +463,17 @@ order by time_window, ordered_host; | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[__first_value_state(t.host) ORDER BY [t.ts ASC NULLS LAST], __first_value_state(t.val) ORDER BY [t.ts ASC NULLS LAST], __first_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 5000000 }, ts@0) as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[__first_value_state(t.host) ORDER BY [t.ts ASC NULLS LAST], __first_value_state(t.val) ORDER BY [t.ts ASC NULLS LAST], __first_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["ts", "host", "val"], "filters": ["host < Utf8(\"b\") OR host IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[__first_value_state(t.host) ORDER BY [t.ts ASC NULLS LAST], __first_value_state(t.val) ORDER BY [t.ts ASC NULLS LAST], __first_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 5000000 }, ts@0) as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[__first_value_state(t.host) ORDER BY [t.ts ASC NULLS LAST], __first_value_state(t.val) ORDER BY [t.ts ASC NULLS LAST], __first_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["ts", "host", "val"], "filters": ["host >= Utf8(\"b\") AND (host < Utf8(\"d\") OR host IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 2_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[__first_value_state(t.host) ORDER BY [t.ts ASC NULLS LAST], __first_value_state(t.val) ORDER BY [t.ts ASC NULLS LAST], __first_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 5000000 }, ts@0) as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[__first_value_state(t.host) ORDER BY [t.ts ASC NULLS LAST], __first_value_state(t.val) ORDER BY [t.ts ASC NULLS LAST], __first_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["ts", "host", "val"], "filters": ["host >= Utf8(\"d\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 2_| +-+-+-+ @@ -541,6 +551,8 @@ explain -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select first_value(ts order by ts) from t1; @@ -556,17 +568,17 @@ explain analyze | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__first_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__first_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["ts"], "filters": ["__table_id=UInt32(REDACTED)", "host < Utf8(\"b\") OR host IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=Final, gby=[], aggr=[__first_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__first_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["ts"], "filters": ["__table_id=UInt32(REDACTED)", "host >= Utf8(\"b\") AND (host < Utf8(\"d\") OR host IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 2_|_AggregateExec: mode=Final, gby=[], aggr=[__first_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__first_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["ts"], "filters": ["__table_id=UInt32(REDACTED)", "host >= Utf8(\"d\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ @@ -624,6 +636,8 @@ order by ordered_host; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select first_value(host order by ts) as ordered_host, @@ -644,7 +658,7 @@ order by ordered_host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[host@0 as host], aggr=[first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[host@0 as host], aggr=[first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["host", "ts", "val"], "filters": ["__table_id=UInt32(REDACTED)", "host < Utf8(\"b\") OR host IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[ordered_host@0 ASC NULLS LAST], preserve_REDACTED @@ -652,7 +666,7 @@ order by ordered_host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[host@0 as host], aggr=[first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[host@0 as host], aggr=[first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["host", "ts", "val"], "filters": ["__table_id=UInt32(REDACTED)", "host >= Utf8(\"b\") AND (host < Utf8(\"d\") OR host IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 2_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[ordered_host@0 ASC NULLS LAST], preserve_REDACTED @@ -660,7 +674,7 @@ order by ordered_host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[host@0 as host], aggr=[first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[host@0 as host], aggr=[first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["host", "ts", "val"], "filters": ["__table_id=UInt32(REDACTED)", "host >= Utf8(\"d\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -726,6 +740,8 @@ order by time_window, ordered_host; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select first_value(host order by ts) as ordered_host, @@ -751,17 +767,17 @@ order by time_window, ordered_host; | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[__first_value_state(t1.host) ORDER BY [t1.ts ASC NULLS LAST], __first_value_state(t1.val) ORDER BY [t1.ts ASC NULLS LAST], __first_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 5000000 }, ts@1) as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[__first_value_state(t1.host) ORDER BY [t1.ts ASC NULLS LAST], __first_value_state(t1.val) ORDER BY [t1.ts ASC NULLS LAST], __first_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["host", "ts", "val"], "filters": ["__table_id=UInt32(REDACTED)", "host < Utf8(\"b\") OR host IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[__first_value_state(t1.host) ORDER BY [t1.ts ASC NULLS LAST], __first_value_state(t1.val) ORDER BY [t1.ts ASC NULLS LAST], __first_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 5000000 }, ts@1) as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[__first_value_state(t1.host) ORDER BY [t1.ts ASC NULLS LAST], __first_value_state(t1.val) ORDER BY [t1.ts ASC NULLS LAST], __first_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["host", "ts", "val"], "filters": ["__table_id=UInt32(REDACTED)", "host >= Utf8(\"b\") AND (host < Utf8(\"d\") OR host IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 2_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[__first_value_state(t1.host) ORDER BY [t1.ts ASC NULLS LAST], __first_value_state(t1.val) ORDER BY [t1.ts ASC NULLS LAST], __first_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 5000000 }, ts@1) as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[__first_value_state(t1.host) ORDER BY [t1.ts ASC NULLS LAST], __first_value_state(t1.val) ORDER BY [t1.ts ASC NULLS LAST], __first_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["host", "ts", "val"], "filters": ["__table_id=UInt32(REDACTED)", "host >= Utf8(\"d\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 2_| +-+-+-+ diff --git a/tests/cases/standalone/optimizer/first_value_advance.sql b/tests/cases/standalone/optimizer/first_value_advance.sql index 2f8da34b99a..b683761221a 100644 --- a/tests/cases/standalone/optimizer/first_value_advance.sql +++ b/tests/cases/standalone/optimizer/first_value_advance.sql @@ -55,6 +55,8 @@ explain select -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE (partitioning.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select first_value(host order by ts), @@ -79,6 +81,8 @@ explain select first_value(ts order by ts) from t; -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE (partitioning.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select first_value(ts order by ts) from t; @@ -138,6 +142,8 @@ order by ordered_host; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select first_value(host order by ts) as ordered_host, @@ -165,6 +171,8 @@ explain select first_value(ts order by ts) from t; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select first_value(ts order by ts) from t; @@ -202,6 +210,8 @@ order by time_window, ordered_host; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select first_value(host order by ts) as ordered_host, @@ -256,6 +266,8 @@ explain -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select first_value(ts order by ts) from t1; @@ -289,6 +301,8 @@ order by ordered_host; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select first_value(host order by ts) as ordered_host, @@ -330,6 +344,8 @@ order by time_window, ordered_host; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select first_value(host order by ts) as ordered_host, diff --git a/tests/cases/standalone/optimizer/last_value.result b/tests/cases/standalone/optimizer/last_value.result index 55ecfc8bb56..d2765209006 100644 --- a/tests/cases/standalone/optimizer/last_value.result +++ b/tests/cases/standalone/optimizer/last_value.result @@ -69,6 +69,7 @@ explain select -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE (partitioning.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select last_value(host order by ts), @@ -87,7 +88,7 @@ explain analyze |_|_|_AggregateExec: mode=FinalPartitioned, gby=[host@0 as host], aggr=[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST], last_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[host@1 as host], aggr=[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST], last_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "selector":"LastRow { after_merge: false }" REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "selector":"LastRow { after_merge: false }", "projection": ["ts", "host", "not_pk", "val"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -127,6 +128,7 @@ explain select last_value(ts order by ts) from t; -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE (partitioning.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select last_value(ts order by ts) from t; @@ -139,7 +141,7 @@ explain analyze | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[last_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[last_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "selector":"LastRow { after_merge: false }" REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "selector":"LastRow { after_merge: false }", "projection": ["ts"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ diff --git a/tests/cases/standalone/optimizer/last_value.sql b/tests/cases/standalone/optimizer/last_value.sql index 760ee799d48..4b6ecd09ee2 100644 --- a/tests/cases/standalone/optimizer/last_value.sql +++ b/tests/cases/standalone/optimizer/last_value.sql @@ -43,6 +43,7 @@ explain select -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE (partitioning.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select last_value(host order by ts), @@ -67,6 +68,7 @@ explain select last_value(ts order by ts) from t; -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE (partitioning.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select last_value(ts order by ts) from t; diff --git a/tests/cases/standalone/optimizer/last_value_advance.result b/tests/cases/standalone/optimizer/last_value_advance.result index 6b4960be6a1..1b85741eb89 100644 --- a/tests/cases/standalone/optimizer/last_value_advance.result +++ b/tests/cases/standalone/optimizer/last_value_advance.result @@ -95,6 +95,8 @@ explain select -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE (partitioning.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select last_value(host order by ts), @@ -113,7 +115,7 @@ explain analyze |_|_|_AggregateExec: mode=FinalPartitioned, gby=[host@0 as host], aggr=[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST], last_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[host@1 as host], aggr=[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST], last_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "selector":"LastRow { after_merge: false }" REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "selector":"LastRow { after_merge: false }", "projection": ["ts", "host", "not_pk", "val"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -153,6 +155,8 @@ explain select last_value(ts order by ts) from t; -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE (partitioning.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select last_value(ts order by ts) from t; @@ -165,7 +169,7 @@ explain analyze | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[last_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[last_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "selector":"LastRow { after_merge: false }" REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "selector":"LastRow { after_merge: false }", "projection": ["ts"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ @@ -255,6 +259,8 @@ order by ordered_host; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select last_value(host order by ts) as ordered_host, @@ -276,7 +282,7 @@ order by ordered_host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[host@0 as host], aggr=[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST], last_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[host@1 as host], aggr=[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST], last_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }" REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }", "projection": ["ts", "host", "not_pk", "val"], "filters": ["host < Utf8(\"b\") OR host IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[ordered_host@0 ASC NULLS LAST], preserve_REDACTED @@ -284,7 +290,7 @@ order by ordered_host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[host@0 as host], aggr=[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST], last_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[host@1 as host], aggr=[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST], last_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }" REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }", "projection": ["ts", "host", "not_pk", "val"], "filters": ["host >= Utf8(\"b\") AND (host < Utf8(\"d\") OR host IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 2_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[ordered_host@0 ASC NULLS LAST], preserve_REDACTED @@ -292,7 +298,7 @@ order by ordered_host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[host@0 as host], aggr=[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST], last_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[host@1 as host], aggr=[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST], last_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }" REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }", "projection": ["ts", "host", "not_pk", "val"], "filters": ["host >= Utf8(\"d\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -337,6 +343,8 @@ explain select last_value(ts order by ts) from t; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select last_value(ts order by ts) from t; @@ -352,17 +360,17 @@ explain analyze | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__last_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__last_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }" REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }", "projection": ["ts"], "filters": ["host < Utf8(\"b\") OR host IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=Final, gby=[], aggr=[__last_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__last_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }" REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }", "projection": ["ts"], "filters": ["host >= Utf8(\"b\") AND (host < Utf8(\"d\") OR host IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 2_|_AggregateExec: mode=Final, gby=[], aggr=[__last_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__last_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }" REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }", "projection": ["ts"], "filters": ["host >= Utf8(\"d\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ @@ -428,6 +436,8 @@ order by time_window, ordered_host; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select last_value(host order by ts) as ordered_host, @@ -453,17 +463,17 @@ order by time_window, ordered_host; | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[__last_value_state(t.host) ORDER BY [t.ts ASC NULLS LAST], __last_value_state(t.val) ORDER BY [t.ts ASC NULLS LAST], __last_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 5000000 }, ts@0) as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[__last_value_state(t.host) ORDER BY [t.ts ASC NULLS LAST], __last_value_state(t.val) ORDER BY [t.ts ASC NULLS LAST], __last_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["ts", "host", "val"], "filters": ["host < Utf8(\"b\") OR host IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[__last_value_state(t.host) ORDER BY [t.ts ASC NULLS LAST], __last_value_state(t.val) ORDER BY [t.ts ASC NULLS LAST], __last_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 5000000 }, ts@0) as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[__last_value_state(t.host) ORDER BY [t.ts ASC NULLS LAST], __last_value_state(t.val) ORDER BY [t.ts ASC NULLS LAST], __last_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["ts", "host", "val"], "filters": ["host >= Utf8(\"b\") AND (host < Utf8(\"d\") OR host IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 2_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[__last_value_state(t.host) ORDER BY [t.ts ASC NULLS LAST], __last_value_state(t.val) ORDER BY [t.ts ASC NULLS LAST], __last_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 5000000 }, ts@0) as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[__last_value_state(t.host) ORDER BY [t.ts ASC NULLS LAST], __last_value_state(t.val) ORDER BY [t.ts ASC NULLS LAST], __last_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["ts", "host", "val"], "filters": ["host >= Utf8(\"d\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 2_| +-+-+-+ @@ -541,6 +551,8 @@ explain -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select last_value(ts order by ts) from t1; @@ -556,17 +568,17 @@ explain analyze | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__last_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__last_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }" REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }", "projection": ["ts"], "filters": ["__table_id=UInt32(REDACTED)", "host < Utf8(\"b\") OR host IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=Final, gby=[], aggr=[__last_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__last_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }" REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }", "projection": ["ts"], "filters": ["__table_id=UInt32(REDACTED)", "host >= Utf8(\"b\") AND (host < Utf8(\"d\") OR host IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 2_|_AggregateExec: mode=Final, gby=[], aggr=[__last_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__last_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }" REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }", "projection": ["ts"], "filters": ["__table_id=UInt32(REDACTED)", "host >= Utf8(\"d\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ @@ -624,6 +636,8 @@ order by ordered_host; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select last_value(host order by ts) as ordered_host, @@ -644,7 +658,7 @@ order by ordered_host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[host@0 as host], aggr=[last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[host@0 as host], aggr=[last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }" REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }", "projection": ["host", "ts", "val"], "filters": ["__table_id=UInt32(REDACTED)", "host < Utf8(\"b\") OR host IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[ordered_host@0 ASC NULLS LAST], preserve_REDACTED @@ -652,7 +666,7 @@ order by ordered_host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[host@0 as host], aggr=[last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[host@0 as host], aggr=[last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }" REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }", "projection": ["host", "ts", "val"], "filters": ["__table_id=UInt32(REDACTED)", "host >= Utf8(\"b\") AND (host < Utf8(\"d\") OR host IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 2_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[ordered_host@0 ASC NULLS LAST], preserve_REDACTED @@ -660,7 +674,7 @@ order by ordered_host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[host@0 as host], aggr=[last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[host@0 as host], aggr=[last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }" REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "selector":"LastRow { after_merge: false }", "projection": ["host", "ts", "val"], "filters": ["__table_id=UInt32(REDACTED)", "host >= Utf8(\"d\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -726,6 +740,8 @@ order by time_window, ordered_host; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select last_value(host order by ts) as ordered_host, @@ -751,17 +767,17 @@ order by time_window, ordered_host; | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[__last_value_state(t1.host) ORDER BY [t1.ts ASC NULLS LAST], __last_value_state(t1.val) ORDER BY [t1.ts ASC NULLS LAST], __last_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 5000000 }, ts@1) as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[__last_value_state(t1.host) ORDER BY [t1.ts ASC NULLS LAST], __last_value_state(t1.val) ORDER BY [t1.ts ASC NULLS LAST], __last_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["host", "ts", "val"], "filters": ["__table_id=UInt32(REDACTED)", "host < Utf8(\"b\") OR host IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[__last_value_state(t1.host) ORDER BY [t1.ts ASC NULLS LAST], __last_value_state(t1.val) ORDER BY [t1.ts ASC NULLS LAST], __last_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 5000000 }, ts@1) as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[__last_value_state(t1.host) ORDER BY [t1.ts ASC NULLS LAST], __last_value_state(t1.val) ORDER BY [t1.ts ASC NULLS LAST], __last_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["host", "ts", "val"], "filters": ["__table_id=UInt32(REDACTED)", "host >= Utf8(\"b\") AND (host < Utf8(\"d\") OR host IS NULL)"], "time_range":REDACTED |_|_|_| | 1_| 2_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[__last_value_state(t1.host) ORDER BY [t1.ts ASC NULLS LAST], __last_value_state(t1.val) ORDER BY [t1.ts ASC NULLS LAST], __last_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 5000000 }, ts@1) as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[__last_value_state(t1.host) ORDER BY [t1.ts ASC NULLS LAST], __last_value_state(t1.val) ORDER BY [t1.ts ASC NULLS LAST], __last_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":REDACTED, "projection": ["host", "ts", "val"], "filters": ["__table_id=UInt32(REDACTED)", "host >= Utf8(\"d\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 2_| +-+-+-+ diff --git a/tests/cases/standalone/optimizer/last_value_advance.sql b/tests/cases/standalone/optimizer/last_value_advance.sql index 9090b399b07..6e83b8751c5 100644 --- a/tests/cases/standalone/optimizer/last_value_advance.sql +++ b/tests/cases/standalone/optimizer/last_value_advance.sql @@ -55,6 +55,8 @@ explain select -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE (partitioning.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select last_value(host order by ts), @@ -79,6 +81,8 @@ explain select last_value(ts order by ts) from t; -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE (partitioning.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select last_value(ts order by ts) from t; @@ -138,6 +142,8 @@ order by ordered_host; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select last_value(host order by ts) as ordered_host, @@ -165,6 +171,8 @@ explain select last_value(ts order by ts) from t; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select last_value(ts order by ts) from t; @@ -202,6 +210,8 @@ order by time_window, ordered_host; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select last_value(host order by ts) as ordered_host, @@ -256,6 +266,8 @@ explain -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select last_value(ts order by ts) from t1; @@ -289,6 +301,8 @@ order by ordered_host; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select last_value(host order by ts) as ordered_host, @@ -330,6 +344,8 @@ order by time_window, ordered_host; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- might write to different partitions -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED explain analyze select last_value(host order by ts) as ordered_host, diff --git a/tests/cases/standalone/optimizer/order_by.result b/tests/cases/standalone/optimizer/order_by.result index 020e1dc28ac..ba5119d8a9c 100644 --- a/tests/cases/standalone/optimizer/order_by.result +++ b/tests/cases/standalone/optimizer/order_by.result @@ -132,6 +132,7 @@ SELECT i, t AS alias_ts FROM test_pk ORDER BY alias_ts DESC LIMIT 5; -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT i, t AS alias_ts FROM test_pk ORDER BY t DESC LIMIT 5; +-+-+-+ @@ -143,7 +144,7 @@ EXPLAIN ANALYZE SELECT i, t AS alias_ts FROM test_pk ORDER BY t DESC LIMIT 5; | 1_| 0_|_SortPreservingMergeExec: [alias_ts@1 DESC], fetch=5 REDACTED |_|_|_ProjectionExec: expr=[i@0 as i, t@1 as alias_ts] REDACTED |_|_|_SortExec: TopK(fetch=5), expr=[t@1 DESC], preserve_partitioning=[true], filter=[t@1 IS NULL OR t@1 > 2] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["i", "t"], "dyn_filters":REDACTED |_|_|_| |_|_| Total rows: 5_| +-+-+-+ @@ -154,6 +155,7 @@ EXPLAIN ANALYZE SELECT i, t AS alias_ts FROM test_pk ORDER BY t DESC LIMIT 5; -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT i, t AS alias_ts FROM test_pk ORDER BY alias_ts DESC LIMIT 5; +-+-+-+ @@ -165,7 +167,7 @@ EXPLAIN ANALYZE SELECT i, t AS alias_ts FROM test_pk ORDER BY alias_ts DESC LIMI | 1_| 0_|_SortPreservingMergeExec: [alias_ts@1 DESC], fetch=5 REDACTED |_|_|_ProjectionExec: expr=[i@0 as i, t@1 as alias_ts] REDACTED |_|_|_SortExec: TopK(fetch=5), expr=[t@1 DESC], preserve_partitioning=[true], filter=[t@1 IS NULL OR t@1 > 2] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["i", "t"], "dyn_filters":REDACTED |_|_|_| |_|_| Total rows: 5_| +-+-+-+ diff --git a/tests/cases/standalone/optimizer/order_by.sql b/tests/cases/standalone/optimizer/order_by.sql index ee6b6bc0bdd..b785af2fd4b 100644 --- a/tests/cases/standalone/optimizer/order_by.sql +++ b/tests/cases/standalone/optimizer/order_by.sql @@ -29,6 +29,7 @@ SELECT i, t AS alias_ts FROM test_pk ORDER BY alias_ts DESC LIMIT 5; -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT i, t AS alias_ts FROM test_pk ORDER BY t DESC LIMIT 5; -- SQLNESS REPLACE (-+) - @@ -37,6 +38,7 @@ EXPLAIN ANALYZE SELECT i, t AS alias_ts FROM test_pk ORDER BY t DESC LIMIT 5; -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT i, t AS alias_ts FROM test_pk ORDER BY alias_ts DESC LIMIT 5; DROP TABLE test_pk; diff --git a/tests/cases/standalone/optimizer/time_index_filter_pushdown.result b/tests/cases/standalone/optimizer/time_index_filter_pushdown.result index fd0d869a33f..50f39764ca2 100644 --- a/tests/cases/standalone/optimizer/time_index_filter_pushdown.result +++ b/tests/cases/standalone/optimizer/time_index_filter_pushdown.result @@ -123,6 +123,7 @@ Affected Rows: 3 -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT rack FROM @@ -138,7 +139,7 @@ WHERE |_|_|_| | 1_| 0_|_FilterExec: usage_small@1 >= 10, projection=[rack@0] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_UnorderedScan: region=REDACTED, {"partition_count":REDACTED, "projection": ["rack", "usage_small"], "filters": ["usage_small >= Int16(10)"], "flat_format": REDACTED, "REDACTED +|_|_|_UnorderedScan: region=REDACTED, {"partition_count":REDACTED, "projection": ["rack", "usage_small"], "filters": ["usage_small >= Int16(10)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 3_| +-+-+-+ diff --git a/tests/cases/standalone/optimizer/time_index_filter_pushdown.sql b/tests/cases/standalone/optimizer/time_index_filter_pushdown.sql index 7a7030d0925..1dc010e1bf4 100644 --- a/tests/cases/standalone/optimizer/time_index_filter_pushdown.sql +++ b/tests/cases/standalone/optimizer/time_index_filter_pushdown.sql @@ -81,6 +81,7 @@ VALUES -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE VERBOSE SELECT rack FROM diff --git a/tests/cases/standalone/optimizer/windowed_sort.result b/tests/cases/standalone/optimizer/windowed_sort.result index 9bffb3d58d3..0636e19ba34 100644 --- a/tests/cases/standalone/optimizer/windowed_sort.result +++ b/tests/cases/standalone/optimizer/windowed_sort.result @@ -56,6 +56,7 @@ ORDER BY -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT collect_time_utc AS collect_time, peak_current, @@ -73,7 +74,7 @@ ORDER BY | 1_| 0_|_SortPreservingMergeExec: [collect_time@0 ASC NULLS LAST] REDACTED |_|_|_ProjectionExec: expr=[collect_time_utc@0 as collect_time, peak_current@1 as peak_current] REDACTED |_|_|_SortExec: expr=[collect_time_utc@0 ASC NULLS LAST], preserve_partitioning=[true] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["collect_time_utc", "peak_current"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 8_| +-+-+-+ @@ -105,6 +106,7 @@ ORDER BY -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT collect_time_utc AS collect_time_0, peak_current, @@ -122,7 +124,7 @@ ORDER BY | 1_| 0_|_SortPreservingMergeExec: [collect_time_0@0 ASC NULLS LAST] REDACTED |_|_|_ProjectionExec: expr=[collect_time_utc@0 as collect_time_0, peak_current@1 as peak_current] REDACTED |_|_|_SortExec: expr=[collect_time_utc@0 ASC NULLS LAST], preserve_partitioning=[true] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["collect_time_utc", "peak_current"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 8_| +-+-+-+ @@ -156,6 +158,7 @@ ORDER BY -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT collect_time AS true_collect_time, collect_time_utc AS collect_time, @@ -175,7 +178,7 @@ ORDER BY |_|_|_ProjectionExec: expr=[collect_time@0 as true_collect_time, collect_time_utc@1 as collect_time, peak_current@2 as peak_current] REDACTED |_|_|_WindowedSortExec: expr=collect_time@0 DESC num_ranges=REDACTED REDACTED |_|_|_PartSortExec: expr=collect_time@0 DESC num_ranges=REDACTED REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["collect_time", "collect_time_utc", "peak_current"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 8_| +-+-+-+ @@ -209,6 +212,7 @@ ORDER BY -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT collect_time_utc AS collect_time, collect_time AS true_collect_time, @@ -228,7 +232,7 @@ ORDER BY |_|_|_ProjectionExec: expr=[collect_time_utc@1 as collect_time, collect_time@0 as true_collect_time, peak_current@2 as peak_current] REDACTED |_|_|_WindowedSortExec: expr=collect_time@0 DESC num_ranges=REDACTED REDACTED |_|_|_PartSortExec: expr=collect_time@0 DESC num_ranges=REDACTED REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["collect_time", "collect_time_utc", "peak_current"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 8_| +-+-+-+ @@ -274,6 +278,7 @@ TQL EVAL('2023-10-01 00:00:00.000'::TIMESTAMP, '2023-10-01 00:00:05.000'::TIMEST -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE('2023-10-01 00:00:00.000'::TIMESTAMP, '2023-10-01 00:00:05.000'::TIMESTAMP, '1s') sum(instance_job_metrics); +-+-+-+ @@ -290,7 +295,7 @@ TQL ANALYZE('2023-10-01 00:00:00.000'::TIMESTAMP, '2023-10-01 00:00:05.000'::TIM |_|_|_ProjectionExec: expr=[greptime_timestamp@0 as greptime_timestamp, greptime_value@1 as greptime_value] REDACTED |_|_|_PromInstantManipulateExec: range=[1696118400000..1696118405000], lookback=[300000], interval=[1000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesDivideExec: tags=["instance", "job"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["greptime_timestamp", "greptime_value", "instance", "job"], "filters": ["greptime_timestamp >= TimestampMillisecond(1696118100001, None)", "greptime_timestamp <= TimestampMillisecond(1696118405000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 5_| +-+-+-+ diff --git a/tests/cases/standalone/optimizer/windowed_sort.sql b/tests/cases/standalone/optimizer/windowed_sort.sql index b74db2b6b72..b19f06ec9ab 100644 --- a/tests/cases/standalone/optimizer/windowed_sort.sql +++ b/tests/cases/standalone/optimizer/windowed_sort.sql @@ -39,6 +39,7 @@ ORDER BY -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT collect_time_utc AS collect_time, peak_current, @@ -61,6 +62,7 @@ ORDER BY -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT collect_time_utc AS collect_time_0, peak_current, @@ -85,6 +87,7 @@ ORDER BY -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT collect_time AS true_collect_time, collect_time_utc AS collect_time, @@ -110,6 +113,7 @@ ORDER BY -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE SELECT collect_time_utc AS collect_time, collect_time AS true_collect_time, @@ -144,6 +148,7 @@ TQL EVAL('2023-10-01 00:00:00.000'::TIMESTAMP, '2023-10-01 00:00:05.000'::TIMEST -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE('2023-10-01 00:00:00.000'::TIMESTAMP, '2023-10-01 00:00:05.000'::TIMESTAMP, '1s') sum(instance_job_metrics); DROP TABLE IF EXISTS `instance_job_metrics`; diff --git a/tests/cases/standalone/optimizer/windowed_sort_advance.result b/tests/cases/standalone/optimizer/windowed_sort_advance.result index 46ff90a69f1..f8863341471 100644 --- a/tests/cases/standalone/optimizer/windowed_sort_advance.result +++ b/tests/cases/standalone/optimizer/windowed_sort_advance.result @@ -56,6 +56,7 @@ select ts as ts, status, value from `a` where ts >= '2026-03-12T08:00:00+08:00' -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE select ts as ts, status, value from `a` where ts >= '2026-03-12T08:00:00+08:00' and ts < '2026-03-12T08:02:01+08:00' order by ts asc; +-+-+-+ @@ -67,7 +68,7 @@ EXPLAIN ANALYZE select ts as ts, status, value from `a` where ts >= '2026-03-12T | 1_| 0_|_SortPreservingMergeExec: [ts@0 ASC NULLS LAST] REDACTED |_|_|_ProjectionExec: expr=[ts@2 as ts, status@1 as status, value@0 as value] REDACTED |_|_|_WindowedSortExec: expr=ts@2 ASC NULLS LAST num_ranges=REDACTED REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["value", "status", "ts"], "filters": ["ts >= TimestampNanosecond(1773273600000000000, None)", "ts < TimestampNanosecond(1773273721000000000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 10_| +-+-+-+ @@ -95,6 +96,7 @@ select to_timestamp_millis(ts) as ts, status, value from `a` where ts >= '2026-0 -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE select to_timestamp_millis(ts) as ts, status, value from `a` where ts >= '2026-03-12T08:00:00+08:00' and ts < '2026-03-12T08:02:01+08:00' order by ts asc; +-+-+-+ @@ -106,7 +108,7 @@ EXPLAIN ANALYZE select to_timestamp_millis(ts) as ts, status, value from `a` whe | 1_| 0_|_SortPreservingMergeExec: [ts@0 ASC NULLS LAST] REDACTED |_|_|_WindowedSortExec: expr=ts@0 ASC NULLS LAST num_ranges=REDACTED REDACTED |_|_|_ProjectionExec: expr=[to_timestamp_millis(ts@2) as ts, status@1 as status, value@0 as value] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["value", "status", "ts"], "filters": ["ts >= TimestampNanosecond(1773273600000000000, None)", "ts < TimestampNanosecond(1773273721000000000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 10_| +-+-+-+ diff --git a/tests/cases/standalone/optimizer/windowed_sort_advance.sql b/tests/cases/standalone/optimizer/windowed_sort_advance.sql index fa73b97933c..167a42d591c 100644 --- a/tests/cases/standalone/optimizer/windowed_sort_advance.sql +++ b/tests/cases/standalone/optimizer/windowed_sort_advance.sql @@ -22,6 +22,7 @@ select ts as ts, status, value from `a` where ts >= '2026-03-12T08:00:00+08:00' -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE select ts as ts, status, value from `a` where ts >= '2026-03-12T08:00:00+08:00' and ts < '2026-03-12T08:02:01+08:00' order by ts asc; select to_timestamp_millis(ts) as ts, status, value from `a` where ts >= '2026-03-12T08:00:00+08:00' and ts < '2026-03-12T08:02:01+08:00' order by ts asc; @@ -32,6 +33,7 @@ select to_timestamp_millis(ts) as ts, status, value from `a` where ts >= '2026-0 -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE select to_timestamp_millis(ts) as ts, status, value from `a` where ts >= '2026-03-12T08:00:00+08:00' and ts < '2026-03-12T08:02:01+08:00' order by ts asc; DROP TABLE `a`; diff --git a/tests/cases/standalone/optimizer/windowed_sort_nanos.result b/tests/cases/standalone/optimizer/windowed_sort_nanos.result index f6736363079..9577b711704 100644 --- a/tests/cases/standalone/optimizer/windowed_sort_nanos.result +++ b/tests/cases/standalone/optimizer/windowed_sort_nanos.result @@ -56,6 +56,7 @@ select to_timestamp_nanos(ts) as ts, status, value from `a_ms` where ts >= '2026 -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE select to_timestamp_nanos(ts) as ts, status, value from `a_ms` where ts >= '2026-03-12T08:00:00.000+08:00' and ts < '2026-03-12T08:02:01.000+08:00' order by ts asc; +-+-+-+ @@ -67,7 +68,7 @@ EXPLAIN ANALYZE select to_timestamp_nanos(ts) as ts, status, value from `a_ms` w | 1_| 0_|_SortPreservingMergeExec: [ts@0 ASC NULLS LAST] REDACTED |_|_|_WindowedSortExec: expr=ts@0 ASC NULLS LAST num_ranges=REDACTED REDACTED |_|_|_ProjectionExec: expr=[to_timestamp_nanos(ts@2) as ts, status@1 as status, value@0 as value] REDACTED -|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0} REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["value", "status", "ts"], "filters": ["ts >= TimestampMillisecond(1773273600000, None)", "ts < TimestampMillisecond(1773273721000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 10_| +-+-+-+ diff --git a/tests/cases/standalone/optimizer/windowed_sort_nanos.sql b/tests/cases/standalone/optimizer/windowed_sort_nanos.sql index 6a6cf51c4b6..27abd0bc805 100644 --- a/tests/cases/standalone/optimizer/windowed_sort_nanos.sql +++ b/tests/cases/standalone/optimizer/windowed_sort_nanos.sql @@ -22,6 +22,7 @@ select to_timestamp_nanos(ts) as ts, status, value from `a_ms` where ts >= '2026 -- SQLNESS REPLACE (metrics.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE num_ranges=\d+ num_ranges=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED EXPLAIN ANALYZE select to_timestamp_nanos(ts) as ts, status, value from `a_ms` where ts >= '2026-03-12T08:00:00.000+08:00' and ts < '2026-03-12T08:02:01.000+08:00' order by ts asc; DROP TABLE `a_ms`; diff --git a/tests/cases/standalone/tql-explain-analyze/analyze.result b/tests/cases/standalone/tql-explain-analyze/analyze.result index 7c0b9b16e6b..45942403722 100644 --- a/tests/cases/standalone/tql-explain-analyze/analyze.result +++ b/tests/cases/standalone/tql-explain-analyze/analyze.result @@ -15,6 +15,7 @@ Affected Rows: 3 -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 10, '5s') test; +-+-+-+ @@ -26,7 +27,7 @@ TQL ANALYZE (0, 10, '5s') test; | 1_| 0_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[j] REDACTED |_|_|_PromSeriesDivideExec: tags=["k"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["i", "j", "k"], "filters": ["j >= TimestampMillisecond(-299999, None)", "j <= TimestampMillisecond(10000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -40,6 +41,7 @@ TQL ANALYZE (0, 10, '5s') test; -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 10, '1s', '2s') test; +-+-+-+ @@ -51,7 +53,7 @@ TQL ANALYZE (0, 10, '1s', '2s') test; | 1_| 0_|_PromInstantManipulateExec: range=[0..10000], lookback=[2000], interval=[1000], time index=[j] REDACTED |_|_|_PromSeriesDivideExec: tags=["k"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["i", "j", "k"], "filters": ["j >= TimestampMillisecond(-1999, None)", "j <= TimestampMillisecond(10000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -64,6 +66,7 @@ TQL ANALYZE (0, 10, '1s', '2s') test; -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE ('1970-01-01T00:00:00'::timestamp, '1970-01-01T00:00:00'::timestamp + '10 seconds'::interval, '5s') test; +-+-+-+ @@ -75,7 +78,7 @@ TQL ANALYZE ('1970-01-01T00:00:00'::timestamp, '1970-01-01T00:00:00'::timestamp | 1_| 0_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[j] REDACTED |_|_|_PromSeriesDivideExec: tags=["k"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["i", "j", "k"], "filters": ["j >= TimestampMillisecond(-299999, None)", "j <= TimestampMillisecond(10000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -91,6 +94,7 @@ TQL ANALYZE ('1970-01-01T00:00:00'::timestamp, '1970-01-01T00:00:00'::timestamp -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED -- SQLNESS REPLACE (flat_format.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE VERBOSE (0, 10, '5s') test; +-+-+-+ @@ -102,7 +106,7 @@ TQL ANALYZE VERBOSE (0, 10, '5s') test; | 1_| 0_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[j] REDACTED |_|_|_PromSeriesDivideExec: tags=["k"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["i", "j", "k"], "filters": ["j >= TimestampMillisecond(-299999, None)", "j <= TimestampMillisecond(10000, None)"], "REDACTED +|_|_|_SeriesScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["i", "j", "k"], "filters": ["j >= TimestampMillisecond(-299999, None)", "j <= TimestampMillisecond(10000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -123,6 +127,7 @@ Affected Rows: 0 -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 10, '5s') test; +-+-+-+ @@ -134,12 +139,12 @@ TQL ANALYZE (0, 10, '5s') test; | 1_| 0_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[j] REDACTED |_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["i", "j", "k", "l"], "filters": ["j >= TimestampMillisecond(-299999, None)", "j <= TimestampMillisecond(10000, None)", "k < Utf8(\"a\") OR k IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[j] REDACTED |_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["i", "j", "k", "l"], "filters": ["j >= TimestampMillisecond(-299999, None)", "j <= TimestampMillisecond(10000, None)", "k >= Utf8(\"a\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -151,6 +156,7 @@ TQL ANALYZE (0, 10, '5s') test; -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 10, '5s') rate(test[10s]); +-+-+-+ @@ -165,7 +171,7 @@ TQL ANALYZE (0, 10, '5s') rate(test[10s]); |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[j], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["i", "j", "k", "l"], "filters": ["j >= TimestampMillisecond(-9999, None)", "j <= TimestampMillisecond(10000, None)", "k < Utf8(\"a\") OR k IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_FilterExec: prom_rate(j_range,i,j,Int64(10000))@1 IS NOT NULL REDACTED |_|_|_ProjectionExec: expr=[j@1 as j, prom_rate(j_range@4, i@0, j@1, 10000) as prom_rate(j_range,i,j,Int64(10000)), k@2 as k, l@3 as l] REDACTED @@ -173,7 +179,7 @@ TQL ANALYZE (0, 10, '5s') rate(test[10s]); |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[j], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["i", "j", "k", "l"], "filters": ["j >= TimestampMillisecond(-9999, None)", "j <= TimestampMillisecond(10000, None)", "k >= Utf8(\"a\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -187,6 +193,7 @@ TQL ANALYZE (0, 10, '5s') rate(test[10s]); -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE FORMAT JSON (0, 10, '5s') test; +-+-+-+ @@ -208,6 +215,7 @@ TQL ANALYZE FORMAT JSON (0, 10, '5s') test; -- SQLNESS REPLACE (Duration.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (flat_format.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE VERBOSE FORMAT JSON (0, 10, '5s') test; +-+-+-+ @@ -227,6 +235,7 @@ TQL ANALYZE VERBOSE FORMAT JSON (0, 10, '5s') test; -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE FORMAT TEXT (0, 10, '5s') test; +-+-+-+ @@ -238,12 +247,12 @@ TQL ANALYZE FORMAT TEXT (0, 10, '5s') test; | 1_| 0_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[j] REDACTED |_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["i", "j", "k", "l"], "filters": ["j >= TimestampMillisecond(-299999, None)", "j <= TimestampMillisecond(10000, None)", "k < Utf8(\"a\") OR k IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[j] REDACTED |_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["i", "j", "k", "l"], "filters": ["j >= TimestampMillisecond(-299999, None)", "j <= TimestampMillisecond(10000, None)", "k >= Utf8(\"a\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -280,6 +289,7 @@ TQL EVAL sum(test2); -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE sum(test2); +-+-+-+ @@ -299,7 +309,7 @@ TQL ANALYZE sum(test2); |_|_|_ProjectionExec: expr=[greptime_timestamp@0 as greptime_timestamp, greptime_value@1 as greptime_value] REDACTED |_|_|_PromInstantManipulateExec: range=[0..0], lookback=[300000], interval=[300000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesDivideExec: tags=["shard"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "selector":"LastRow { after_merge: true }", "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "selector":"LastRow { after_merge: true }", "distribution":"PerSeries", "projection": ["greptime_timestamp", "greptime_value", "shard"], "filters": ["greptime_timestamp >= TimestampMillisecond(-299999, None)", "greptime_timestamp <= TimestampMillisecond(0, None)", "shard <= Utf8(\"2\") OR shard IS NULL"], "time_range":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[greptime_timestamp@0 as greptime_timestamp], aggr=[__sum_state(test2.greptime_value)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED @@ -307,7 +317,7 @@ TQL ANALYZE sum(test2); |_|_|_ProjectionExec: expr=[greptime_timestamp@0 as greptime_timestamp, greptime_value@1 as greptime_value] REDACTED |_|_|_PromInstantManipulateExec: range=[0..0], lookback=[300000], interval=[300000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesDivideExec: tags=["shard"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "selector":"LastRow { after_merge: true }", "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "selector":"LastRow { after_merge: true }", "distribution":"PerSeries", "projection": ["greptime_timestamp", "greptime_value", "shard"], "filters": ["greptime_timestamp >= TimestampMillisecond(-299999, None)", "greptime_timestamp <= TimestampMillisecond(0, None)", "shard > Utf8(\"2\")"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ diff --git a/tests/cases/standalone/tql-explain-analyze/analyze.sql b/tests/cases/standalone/tql-explain-analyze/analyze.sql index 598cdf0aa98..cfcb615c289 100644 --- a/tests/cases/standalone/tql-explain-analyze/analyze.sql +++ b/tests/cases/standalone/tql-explain-analyze/analyze.sql @@ -11,6 +11,7 @@ INSERT INTO test VALUES (1, 1, "a"), (1, 1, "b"), (2, 2, "a"); -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 10, '5s') test; -- 'lookback' parameter is not fully supported, the test has to be updated @@ -22,6 +23,7 @@ TQL ANALYZE (0, 10, '5s') test; -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 10, '1s', '2s') test; -- analyze at 0s, 5s and 10s. No point at 0s. @@ -32,6 +34,7 @@ TQL ANALYZE (0, 10, '1s', '2s') test; -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE ('1970-01-01T00:00:00'::timestamp, '1970-01-01T00:00:00'::timestamp + '10 seconds'::interval, '5s') test; -- analyze verbose at 0s, 5s and 10s. No point at 0s. @@ -45,6 +48,7 @@ TQL ANALYZE ('1970-01-01T00:00:00'::timestamp, '1970-01-01T00:00:00'::timestamp -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED -- SQLNESS REPLACE (flat_format.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE VERBOSE (0, 10, '5s') test; DROP TABLE test; @@ -59,6 +63,7 @@ CREATE TABLE test(i DOUBLE, j TIMESTAMP TIME INDEX, k STRING, l STRING, PRIMARY -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 10, '5s') test; -- SQLNESS REPLACE (metrics.*) REDACTED @@ -68,6 +73,7 @@ TQL ANALYZE (0, 10, '5s') test; -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 10, '5s') rate(test[10s]); -- Test new FORMAT functionality for ANALYZE @@ -79,6 +85,7 @@ TQL ANALYZE (0, 10, '5s') rate(test[10s]); -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE FORMAT JSON (0, 10, '5s') test; -- analyze verbose with JSON format @@ -91,6 +98,7 @@ TQL ANALYZE FORMAT JSON (0, 10, '5s') test; -- SQLNESS REPLACE (Duration.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (flat_format.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE VERBOSE FORMAT JSON (0, 10, '5s') test; -- analyze with TEXT format (should be same as default) @@ -101,6 +109,7 @@ TQL ANALYZE VERBOSE FORMAT JSON (0, 10, '5s') test; -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE FORMAT TEXT (0, 10, '5s') test; drop table test; @@ -126,6 +135,7 @@ TQL EVAL sum(test2); -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE sum(test2); DROP TABLE test2; diff --git a/tests/cases/standalone/tql-explain-analyze/column_pruning.result b/tests/cases/standalone/tql-explain-analyze/column_pruning.result index cecb72fcee1..24ac84a7bc0 100644 --- a/tests/cases/standalone/tql-explain-analyze/column_pruning.result +++ b/tests/cases/standalone/tql-explain-analyze/column_pruning.result @@ -29,6 +29,7 @@ Affected Rows: 9 -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 10, '5s') sum(promql_column_pruning); +-+-+-+ @@ -45,7 +46,7 @@ TQL ANALYZE (0, 10, '5s') sum(promql_column_pruning); |_|_|_ProjectionExec: expr=[ts@0 as ts, greptime_value@4 as greptime_value] REDACTED |_|_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["job", "instance", "region"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["ts", "job", "instance", "region", "greptime_value"], "filters": ["ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(10000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 3_| +-+-+-+ @@ -57,6 +58,7 @@ TQL ANALYZE (0, 10, '5s') sum(promql_column_pruning); -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 10, '5s') sum(rate(promql_column_pruning[5s])); +-+-+-+ @@ -75,7 +77,7 @@ TQL ANALYZE (0, 10, '5s') sum(rate(promql_column_pruning[5s])); |_|_|_PromRangeManipulateExec: req range=[0..10000], interval=[5000], eval range=[5000], time index=[ts] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[ts], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["job", "instance", "region"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["ts", "job", "instance", "region", "greptime_value"], "filters": ["ts >= TimestampMillisecond(-4999, None)", "ts <= TimestampMillisecond(10000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ diff --git a/tests/cases/standalone/tql-explain-analyze/column_pruning.sql b/tests/cases/standalone/tql-explain-analyze/column_pruning.sql index 88ab7484ab7..74684136fb5 100644 --- a/tests/cases/standalone/tql-explain-analyze/column_pruning.sql +++ b/tests/cases/standalone/tql-explain-analyze/column_pruning.sql @@ -25,6 +25,7 @@ INSERT INTO promql_column_pruning VALUES -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 10, '5s') sum(promql_column_pruning); -- SQLNESS REPLACE (metrics.*) REDACTED @@ -34,6 +35,7 @@ TQL ANALYZE (0, 10, '5s') sum(promql_column_pruning); -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 10, '5s') sum(rate(promql_column_pruning[5s])); DROP TABLE promql_column_pruning; diff --git a/tests/cases/standalone/tql-explain-analyze/tsid_column.result b/tests/cases/standalone/tql-explain-analyze/tsid_column.result index 3b509cdd20b..723b1803237 100644 --- a/tests/cases/standalone/tql-explain-analyze/tsid_column.result +++ b/tests/cases/standalone/tql-explain-analyze/tsid_column.result @@ -37,6 +37,8 @@ Affected Rows: 6 -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 10, '5s') sum(tsid_metric); +-+-+-+ @@ -54,7 +56,7 @@ TQL ANALYZE (0, 10, '5s') sum(tsid_metric); |_|_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[val@1 as val, __tsid@2 as __tsid, ts@0 as ts] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["ts", "val", "__tsid"], "filters": ["__table_id=UInt32(REDACTED)", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(10000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 3_| +-+-+-+ @@ -66,6 +68,8 @@ TQL ANALYZE (0, 10, '5s') sum(tsid_metric); -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 10, '5s') sum by (job, instance) (tsid_metric); +-+-+-+ @@ -82,7 +86,7 @@ TQL ANALYZE (0, 10, '5s') sum by (job, instance) (tsid_metric); |_|_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[val@1 as val, instance@3 as instance, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["ts", "val", "__tsid", "instance", "job"], "filters": ["__table_id=UInt32(REDACTED)", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(10000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 6_| +-+-+-+ @@ -94,6 +98,8 @@ TQL ANALYZE (0, 10, '5s') sum by (job, instance) (tsid_metric); -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 10, '5s') sum(irate(tsid_metric[1h])) / scalar(count(count(tsid_metric) by (job))); +-+-+-+ @@ -120,7 +126,7 @@ TQL ANALYZE (0, 10, '5s') sum(irate(tsid_metric[1h])) / scalar(count(count(tsid |_|_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[val@1 as val, job@3 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["ts", "val", "__tsid", "job"], "filters": ["__table_id=UInt32(REDACTED)", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(10000, None)"], "time_range":REDACTED |_|_|_| | 1_| 0_|_SortPreservingMergeExec: [ts@0 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[ts@0 ASC NULLS LAST], preserve_partitioning=[true] REDACTED @@ -133,7 +139,7 @@ TQL ANALYZE (0, 10, '5s') sum(irate(tsid_metric[1h])) / scalar(count(count(tsid |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[ts], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[val@1 as val, __tsid@2 as __tsid, ts@0 as ts] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["ts", "val", "__tsid"], "filters": ["__table_id=UInt32(REDACTED)", "ts >= TimestampMillisecond(-3599999, None)", "ts <= TimestampMillisecond(10000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 2_| +-+-+-+ @@ -145,6 +151,8 @@ TQL ANALYZE (0, 10, '5s') sum(irate(tsid_metric[1h])) / scalar(count(count(tsid -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 10, '5s') sum(irate(tsid_metric[1h])) / scalar(count(sum(tsid_metric) by (job))); +-+-+-+ @@ -172,7 +180,7 @@ TQL ANALYZE (0, 10, '5s') sum(irate(tsid_metric[1h])) / scalar(count(sum(tsid_m |_|_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[val@1 as val, job@3 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["ts", "val", "__tsid", "job"], "filters": ["__table_id=UInt32(REDACTED)", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(10000, None)"], "time_range":REDACTED |_|_|_| | 1_| 0_|_SortPreservingMergeExec: [ts@0 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[ts@0 ASC NULLS LAST], preserve_partitioning=[true] REDACTED @@ -185,7 +193,7 @@ TQL ANALYZE (0, 10, '5s') sum(irate(tsid_metric[1h])) / scalar(count(sum(tsid_m |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[ts], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[val@1 as val, __tsid@2 as __tsid, ts@0 as ts] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["ts", "val", "__tsid"], "filters": ["__table_id=UInt32(REDACTED)", "ts >= TimestampMillisecond(-3599999, None)", "ts <= TimestampMillisecond(10000, None)"], "time_range":REDACTED |_|_|_| |_|_| Total rows: 2_| +-+-+-+ diff --git a/tests/cases/standalone/tql-explain-analyze/tsid_column.sql b/tests/cases/standalone/tql-explain-analyze/tsid_column.sql index dedce2dfb19..51cb584b6f4 100644 --- a/tests/cases/standalone/tql-explain-analyze/tsid_column.sql +++ b/tests/cases/standalone/tql-explain-analyze/tsid_column.sql @@ -31,6 +31,8 @@ INSERT INTO tsid_metric VALUES -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 10, '5s') sum(tsid_metric); -- SQLNESS REPLACE (metrics.*) REDACTED @@ -40,6 +42,8 @@ TQL ANALYZE (0, 10, '5s') sum(tsid_metric); -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 10, '5s') sum by (job, instance) (tsid_metric); -- SQLNESS REPLACE (metrics.*) REDACTED @@ -49,6 +53,8 @@ TQL ANALYZE (0, 10, '5s') sum by (job, instance) (tsid_metric); -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 10, '5s') sum(irate(tsid_metric[1h])) / scalar(count(count(tsid_metric) by (job))); -- SQLNESS REPLACE (metrics.*) REDACTED @@ -58,6 +64,8 @@ TQL ANALYZE (0, 10, '5s') sum(irate(tsid_metric[1h])) / scalar(count(count(tsid -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE __table_id\s*=\s*UInt32\(\d+\) __table_id=UInt32(REDACTED) +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED TQL ANALYZE (0, 10, '5s') sum(irate(tsid_metric[1h])) / scalar(count(sum(tsid_metric) by (job))); DROP TABLE tsid_metric; diff --git a/tests/compatibility/cases/analyze_verbose_remote_metrics_extension/verify.result b/tests/compatibility/cases/analyze_verbose_remote_metrics_extension/verify.result index 2bdcb907ebf..a3ed4dbba0e 100644 --- a/tests/compatibility/cases/analyze_verbose_remote_metrics_extension/verify.result +++ b/tests/compatibility/cases/analyze_verbose_remote_metrics_extension/verify.result @@ -8,6 +8,7 @@ -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE ("files":\s\[.*\],) "files": REDACTED, -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE "(time_range|file_summary|max_part_range)":.* "scan_details":REDACTED EXPLAIN ANALYZE VERBOSE SELECT count(*) FROM t_analyze_verbose_remote_metrics_extension; +-+-+-+ @@ -23,12 +24,12 @@ EXPLAIN ANALYZE VERBOSE SELECT count(*) FROM t_analyze_verbose_remote_metrics_ex | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(t_analyze_verbose_remote_REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__count_state(t_analyze_verbose_remote_REDACTED -|_|_|_SeqScan: region=REDACTED, {"partition_count":REDACTED, "projection": ["ts"], "filters": ["host < Utf8(\"m\") OR host IS NULL"], "files": REDACTED, "flat_format": REDACTED, "REDACTED +|_|_|_SeqScan: region=REDACTED, {"partition_count":REDACTED, "projection": ["ts"], "filters": ["host < Utf8(\"m\") OR host IS NULL"], "scan_details":REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(t_analyze_verbose_remote_REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[__count_state(t_analyze_verbose_remote_REDACTED -|_|_|_SeqScan: region=REDACTED, {"partition_count":REDACTED, "projection": ["ts"], "filters": ["host >= Utf8(\"m\")"], "flat_format": REDACTED, "REDACTED +|_|_|_SeqScan: region=REDACTED, {"partition_count":REDACTED, "projection": ["ts"], "filters": ["host >= Utf8(\"m\")"], "scan_details":REDACTED |_|_|_| |_|_| Total rows: 1_| +-+-+-+ diff --git a/tests/compatibility/cases/analyze_verbose_remote_metrics_extension/verify.sql b/tests/compatibility/cases/analyze_verbose_remote_metrics_extension/verify.sql index 81eb489e663..f50b4a29116 100644 --- a/tests/compatibility/cases/analyze_verbose_remote_metrics_extension/verify.sql +++ b/tests/compatibility/cases/analyze_verbose_remote_metrics_extension/verify.sql @@ -8,4 +8,5 @@ -- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED -- SQLNESS REPLACE ("files":\s\[.*\],) "files": REDACTED, -- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE "(time_range|file_summary|max_part_range)":.* "scan_details":REDACTED EXPLAIN ANALYZE VERBOSE SELECT count(*) FROM t_analyze_verbose_remote_metrics_extension; From c51293836e46d92bd362b844f293240d8b459b68 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Tue, 20 Oct 2026 06:17:32 +0000 Subject: [PATCH 049/239] fix(query): push LIKE filters down into MergeScan remote plans (#9648) Signed-off-by: Dennis Zhuang --- src/query/src/dist_plan/analyzer/test.rs | 66 +++++++++++++++++++ src/query/src/dist_plan/commutativity.rs | 5 +- .../function/fulltext_index_pruning.result | 21 ++++++ .../function/fulltext_index_pruning.sql | 8 +++ 4 files changed, 97 insertions(+), 3 deletions(-) diff --git a/src/query/src/dist_plan/analyzer/test.rs b/src/query/src/dist_plan/analyzer/test.rs index a689d94c282..207de7d43fd 100644 --- a/src/query/src/dist_plan/analyzer/test.rs +++ b/src/query/src/dist_plan/analyzer/test.rs @@ -530,6 +530,72 @@ fn dictionary_literals_in_values_keep_their_schema() { ); } +#[test] +fn like_filters_are_pushed_down() { + let test_table = TestTable::table_with_name(0, "t".to_string()); + let table_provider = Arc::new(DfTableProviderAdapter::new(test_table)); + let table_source = Arc::new(DefaultTableSource::new(table_provider.clone() as _)); + let ctx = SessionContext::new(); + ctx.register_table(TableReference::bare("t"), table_provider as _) + .unwrap(); + + let like = |negated, pattern: &str, escape_char, case_insensitive| { + Expr::Like(datafusion_expr::expr::Like::new( + negated, + Box::new(col("t.pk1")), + Box::new(lit(pattern.to_string())), + escape_char, + case_insensitive, + )) + }; + // Substrait has no negated LIKE, so `NOT LIKE` comes back as `NOT (.. LIKE ..)`. + for (predicate, decoded) in [ + ( + like(false, "% timeout %", None, false), + like(false, "% timeout %", None, false), + ), + ( + like(true, "% timeout %", None, false), + Expr::Not(Box::new(like(false, "% timeout %", None, false))), + ), + ( + like(false, "% timeout %", None, true), + like(false, "% timeout %", None, true), + ), + ( + like(false, "100!% done", Some('!'), false), + like(false, "100!% done", Some('!'), false), + ), + ] { + let plan = LogicalPlanBuilder::scan_with_filters("t", table_source.clone(), None, vec![]) + .unwrap() + .filter(predicate.clone()) + .unwrap() + .build() + .unwrap(); + + let result = DistPlannerAnalyzer {} + .analyze(plan, &ConfigOptions::default()) + .unwrap(); + let remote_input = find_merge_scan(&result).unwrap().input(); + assert!( + has_filter_above_table_scan(remote_input, "t", &predicate), + "{result}" + ); + + let bytes = DFLogicalSubstraitConvertor + .encode(remote_input, DefaultSerializer) + .unwrap(); + let decoded_plan = + futures::executor::block_on(DFLogicalSubstraitConvertor.decode(bytes, ctx.state())) + .unwrap(); + assert!( + has_filter_above_table_scan(&decoded_plan, "t", &decoded), + "{decoded_plan}" + ); + } +} + fn try_encode_decode_substrait(plan: &LogicalPlan, state: SessionState) { let sub_plan_bytes = substrait::DFLogicalSubstraitConvertor .encode(plan, crate::query_engine::DefaultSerializer) diff --git a/src/query/src/dist_plan/commutativity.rs b/src/query/src/dist_plan/commutativity.rs index d02b746a8bf..b69589a7cfe 100644 --- a/src/query/src/dist_plan/commutativity.rs +++ b/src/query/src/dist_plan/commutativity.rs @@ -113,7 +113,6 @@ impl Categorizer { } Commutativity::Commutative } - // TODO(ruihang): Change this to Commutative once Like is supported in substrait LogicalPlan::Filter(filter) => Self::check_expr(&filter.predicate), LogicalPlan::Window(_) => Commutativity::Unimplemented, LogicalPlan::Aggregate(aggr) => { @@ -276,6 +275,7 @@ impl Categorizer { | Expr::Between(_) | Expr::Exists(_) | Expr::InList(_) + | Expr::Like(_) | Expr::Case(_) => Commutativity::Commutative, // Annotation collection must not affect distribution. Keep scalar UDFs pushdownable; // TODO: Preserve datanode PromQL annotations through plan decoding, return them in @@ -283,8 +283,7 @@ impl Categorizer { Expr::ScalarFunction(_) => Commutativity::Commutative, Expr::AggregateFunction(_udaf) => Commutativity::Commutative, - Expr::Like(_) - | Expr::SimilarTo(_) + Expr::SimilarTo(_) | Expr::IsUnknown(_) | Expr::IsNotUnknown(_) | Expr::WindowFunction(_) diff --git a/tests/cases/standalone/common/function/fulltext_index_pruning.result b/tests/cases/standalone/common/function/fulltext_index_pruning.result index 10e4b56e2b6..b357d561984 100644 --- a/tests/cases/standalone/common/function/fulltext_index_pruning.result +++ b/tests/cases/standalone/common/function/fulltext_index_pruning.result @@ -219,6 +219,27 @@ SELECT ts, msg FROM ft_en_bloom WHERE msg LIKE '%user-123%' ORDER BY ts; | 1970-01-01T00:00:00.007 | 用户user-123登录 | +-------------------------+------------------+ +-- The `LIKE` filter runs on the datanode, after the index prunes the row group. +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE SeqScan:.*("filters":\s\[[^\]]*\]).*("rg_fulltext_filtered":\d+,\s"rows_fulltext_filtered":\d+).* SeqScan:_${1}_${2} +-- SQLNESS REPLACE (metrics=.*) REDACTED +EXPLAIN ANALYZE VERBOSE SELECT ts, msg FROM ft_en_bloom WHERE msg LIKE 'nomatch %'; + ++-+-+-+ +| stage | node | plan_| ++-+-+-+ +| 0_| 0_|_CooperativeExec REDACTED +|_|_|_MergeScanExec: REDACTED +|_|_|_| +| 1_| 0_|_FilterExec: msg@1 LIKE nomatch % REDACTED +|_|_|_CooperativeExec REDACTED +|_|_|_SeqScan:_"filters": ["msg LIKE Utf8(\"nomatch %\")"]_"rg_fulltext_filtered":1, "rows_fulltext_filtered":8 +|_|_|_| +|_|_| Total rows: 0_| ++-+-+-+ + SELECT ts, matches_term(lower(msg), 'quota') AS "matches_term(lower(msg), 'quota')", matches_term(lower(msg), 'vol') AS "matches_term(lower(msg), 'vol')", lower(msg) LIKE '%disk quota exceeded%' AS "lower(msg) LIKE '%disk quota exceeded%'", lower(msg) LIKE 'disk quota%' AS "lower(msg) LIKE 'disk quota%'" FROM ft_en_bloom ORDER BY ts; +-------------------------+-----------------------------------+---------------------------------+-----------------------------------------+-------------------------------+ diff --git a/tests/cases/standalone/common/function/fulltext_index_pruning.sql b/tests/cases/standalone/common/function/fulltext_index_pruning.sql index f4e711b94e0..b063c92d518 100644 --- a/tests/cases/standalone/common/function/fulltext_index_pruning.sql +++ b/tests/cases/standalone/common/function/fulltext_index_pruning.sql @@ -64,6 +64,14 @@ SELECT ts, msg FROM ft_en_bloom WHERE msg LIKE '连接timeout.%' ORDER BY ts; SELECT ts, msg FROM ft_en_bloom WHERE msg LIKE '%user-123%' ORDER BY ts; +-- The `LIKE` filter runs on the datanode, after the index prunes the row group. +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE SeqScan:.*("filters":\s\[[^\]]*\]).*("rg_fulltext_filtered":\d+,\s"rows_fulltext_filtered":\d+).* SeqScan:_${1}_${2} +-- SQLNESS REPLACE (metrics=.*) REDACTED +EXPLAIN ANALYZE VERBOSE SELECT ts, msg FROM ft_en_bloom WHERE msg LIKE 'nomatch %'; + SELECT ts, matches_term(lower(msg), 'quota') AS "matches_term(lower(msg), 'quota')", matches_term(lower(msg), 'vol') AS "matches_term(lower(msg), 'vol')", lower(msg) LIKE '%disk quota exceeded%' AS "lower(msg) LIKE '%disk quota exceeded%'", lower(msg) LIKE 'disk quota%' AS "lower(msg) LIKE 'disk quota%'" FROM ft_en_bloom ORDER BY ts; SELECT ts, msg FROM ft_en_bloom WHERE matches_term(lower(msg), 'quota') ORDER BY ts; From 41ec5e076303abde4c43c972d47eb1b56f9c376f Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Tue, 20 Oct 2026 07:40:54 +0000 Subject: [PATCH 050/239] fix: decouple slow-query logging from the global log level (#9765) Signed-off-by: shuiyisong --- config/config.md | 4 +- config/frontend.example.toml | 4 +- config/standalone.example.toml | 4 +- src/catalog/src/process_manager.rs | 2 +- src/common/telemetry/src/logging.rs | 346 +++++++++++++++++++--------- src/common/telemetry/src/macros.rs | 9 +- 6 files changed, 258 insertions(+), 111 deletions(-) diff --git a/config/config.md b/config/config.md index 8663ff1cc28..a549f5b78ad 100644 --- a/config/config.md +++ b/config/config.md @@ -262,7 +262,7 @@ | `logging.tracing_sample_ratio.default_ratio` | Float | `1.0` | -- | | `slow_query` | -- | -- | The slow query log options. | | `slow_query.enable` | Bool | `true` | Whether to enable slow query log. | -| `slow_query.record_type` | String | `system_table` | The record type of slow queries. It can be `system_table` or `log`.
If `system_table` is selected, the slow queries will be recorded in a system table `greptime_private.slow_queries`.
If `log` is selected, the slow queries will be logged in a log file `greptimedb-slow-queries.*`. | +| `slow_query.record_type` | String | `system_table` | The record type of slow queries. It can be `system_table` or `log`.
If `system_table` is selected, the slow queries will be recorded in a system table `greptime_private.slow_queries`.
If `log` is selected, slow queries are logged to stdout when `logging.append_stdout` is enabled,
and to `greptimedb-slow-queries.*` when `logging.enable_file_logging` is enabled and `logging.dir` is nonempty.
Slow query logging is independent of `logging.level` and uses `logging.log_format` and the logging retention settings. | | `slow_query.threshold` | String | `30s` | The threshold of slow query. It can be human readable time string, for example: `10s`, `100ms`, `1s`. | | `slow_query.sample_ratio` | Float | `1.0` | The sampling ratio of slow query log. The value should be in the range of (0, 1]. For example, `0.1` means 10% of the slow queries will be logged and `1.0` means all slow queries will be logged. | | `slow_query.ttl` | String | `90d` | The TTL of the `slow_queries` system table. Default is `90d` when `record_type` is `system_table`. | @@ -422,7 +422,7 @@ | `logging.tracing_sample_ratio.default_ratio` | Float | `1.0` | -- | | `slow_query` | -- | -- | The slow query log options. | | `slow_query.enable` | Bool | `true` | Whether to enable slow query log. | -| `slow_query.record_type` | String | `system_table` | The record type of slow queries. It can be `system_table` or `log`.
If `system_table` is selected, the slow queries will be recorded in a system table `greptime_private.slow_queries`.
If `log` is selected, the slow queries will be logged in a log file `greptimedb-slow-queries.*`. | +| `slow_query.record_type` | String | `system_table` | The record type of slow queries. It can be `system_table` or `log`.
If `system_table` is selected, the slow queries will be recorded in a system table `greptime_private.slow_queries`.
If `log` is selected, slow queries are logged to stdout when `logging.append_stdout` is enabled,
and to `greptimedb-slow-queries.*` when `logging.enable_file_logging` is enabled and `logging.dir` is nonempty.
Slow query logging is independent of `logging.level` and uses `logging.log_format` and the logging retention settings. | | `slow_query.threshold` | String | `30s` | The threshold of slow query. It can be human readable time string, for example: `10s`, `100ms`, `1s`. | | `slow_query.sample_ratio` | Float | `1.0` | The sampling ratio of slow query log. The value should be in the range of (0, 1]. For example, `0.1` means 10% of the slow queries will be logged and `1.0` means all slow queries will be logged. | | `slow_query.ttl` | String | `90d` | The TTL of the `slow_queries` system table. Default is `90d` when `record_type` is `system_table`. | diff --git a/config/frontend.example.toml b/config/frontend.example.toml index 97a939c9829..6d7e74824d1 100644 --- a/config/frontend.example.toml +++ b/config/frontend.example.toml @@ -482,7 +482,9 @@ enable = true ## The record type of slow queries. It can be `system_table` or `log`. ## If `system_table` is selected, the slow queries will be recorded in a system table `greptime_private.slow_queries`. -## If `log` is selected, the slow queries will be logged in a log file `greptimedb-slow-queries.*`. +## If `log` is selected, slow queries are logged to stdout when `logging.append_stdout` is enabled, +## and to `greptimedb-slow-queries.*` when `logging.enable_file_logging` is enabled and `logging.dir` is nonempty. +## Slow query logging is independent of `logging.level` and uses `logging.log_format` and the logging retention settings. record_type = "system_table" ## The threshold of slow query. It can be human readable time string, for example: `10s`, `100ms`, `1s`. diff --git a/config/standalone.example.toml b/config/standalone.example.toml index e16210d3644..3aedde0aec0 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -1092,7 +1092,9 @@ enable = true ## The record type of slow queries. It can be `system_table` or `log`. ## If `system_table` is selected, the slow queries will be recorded in a system table `greptime_private.slow_queries`. -## If `log` is selected, the slow queries will be logged in a log file `greptimedb-slow-queries.*`. +## If `log` is selected, slow queries are logged to stdout when `logging.append_stdout` is enabled, +## and to `greptimedb-slow-queries.*` when `logging.enable_file_logging` is enabled and `logging.dir` is nonempty. +## Slow query logging is independent of `logging.level` and uses `logging.log_format` and the logging retention settings. record_type = "system_table" ## The threshold of slow query. It can be human readable time string, for example: `10s`, `100ms`, `1s`. diff --git a/src/catalog/src/process_manager.rs b/src/catalog/src/process_manager.rs index dfe3f7f283d..17df06d7648 100644 --- a/src/catalog/src/process_manager.rs +++ b/src/catalog/src/process_manager.rs @@ -417,7 +417,7 @@ impl SlowQueryTimer { SlowQueriesRecordType::SystemTable => { self.recorder.record(Box::new(slow_query_event)); } - // Record the slow query in a specific logs file. + // Record the slow query through its dedicated logger. SlowQueriesRecordType::Log => { slow!( cost = slow_query_event.cost, diff --git a/src/common/telemetry/src/logging.rs b/src/common/telemetry/src/logging.rs index dd99e901f75..b1430bbd726 100644 --- a/src/common/telemetry/src/logging.rs +++ b/src/common/telemetry/src/logging.rs @@ -35,7 +35,7 @@ use opentelemetry_semantic_conventions::resource; use serde::{Deserialize, Serialize}; use tracing::callsite; use tracing::metadata::LevelFilter; -use tracing_appender::non_blocking::WorkerGuard; +use tracing_appender::non_blocking::{NonBlocking, WorkerGuard}; use tracing_log::LogTracer; use tracing_subscriber::filter::{FilterFn, Targets}; use tracing_subscriber::fmt::Layer; @@ -58,6 +58,10 @@ pub const DEFAULT_LOGGING_DIR: &str = "logs"; pub static LOG_RELOAD_HANDLE: OnceCell> = OnceCell::new(); +/// The independent dispatcher used by the `slow!` macro. +#[doc(hidden)] +pub static SLOW_QUERY_LOGGER: OnceCell = OnceCell::new(); + type DynSubscriber = Layered, Registry>; type OtelTraceLayer = tracing_opentelemetry::OpenTelemetryLayer; @@ -324,7 +328,8 @@ pub struct SlowQueryOptions { /// Whether to enable slow query log. pub enable: bool, - /// The record type of slow queries. + /// The record type of slow queries. Log records ignore the global log level and use + /// the configured stdout and file destinations, format, and retention settings. #[serde(deserialize_with = "empty_string_as_default")] pub record_type: SlowQueriesRecordType, @@ -359,7 +364,7 @@ pub enum SlowQueriesRecordType { /// Record the slow query in the system table. #[default] SystemTable, - /// Record the slow query in a specific logs file. + /// Record the slow query to the configured stdout and file destinations. Log, } @@ -462,32 +467,20 @@ pub fn init_global_logging( // Enable log compatible layer to convert log record to tracing span. LogTracer::init().expect("log tracer must be valid"); - // Configure the stdout logging layer. - let stdout_logging_layer = if opts.append_stdout { + let stdout_writer = if opts.append_stdout { let (writer, guard) = tracing_appender::non_blocking(std::io::stdout()); guards.push(guard); - - if opts.log_format == LogFormat::Json { - Some( - Layer::new() - .json() - .with_writer(writer) - .with_ansi(std::io::stdout().is_terminal()) - .boxed(), - ) - } else { - Some( - Layer::new() - .with_writer(writer) - .with_ansi(std::io::stdout().is_terminal()) - .boxed(), - ) - } + Some(writer) } else { None }; + let stdout_logging_layer = stdout_writer.clone().map(|writer| { + build_logging_layer(writer, opts.log_format, std::io::stdout().is_terminal()) + }); let file_logging_enabled = opts.enable_file_logging && !opts.dir.is_empty(); + let slow_query_logging_enabled = slow_query_opts + .is_some_and(|opts| opts.enable && opts.record_type == SlowQueriesRecordType::Log); let retention = file_logging_enabled .then(|| { @@ -502,17 +495,7 @@ pub fn init_global_logging( let (writer, guard) = tracing_appender::non_blocking(rolling_appender); guards.push(guard); - if opts.log_format == LogFormat::Json { - Some( - Layer::new() - .json() - .with_writer(writer) - .with_ansi(false) - .boxed(), - ) - } else { - Some(Layer::new().with_writer(writer).with_ansi(false).boxed()) - } + Some(build_logging_layer(writer, opts.log_format, false)) } else { None }; @@ -524,30 +507,24 @@ pub fn init_global_logging( let (writer, guard) = tracing_appender::non_blocking(rolling_appender); guards.push(guard); - if opts.log_format == LogFormat::Json { - Some( - Layer::new() - .json() - .with_writer(writer) - .with_ansi(false) - .with_filter(filter::LevelFilter::ERROR) - .boxed(), - ) - } else { - Some( - Layer::new() - .with_writer(writer) - .with_ansi(false) - .with_filter(filter::LevelFilter::ERROR) - .boxed(), - ) - } + Some( + build_logging_layer(writer, opts.log_format, false) + .with_filter(filter::LevelFilter::ERROR), + ) } else { None }; - let slow_query_logging_layer = - build_slow_query_logger(opts, slow_query_opts, retention.as_ref(), &mut guards); + if slow_query_logging_enabled { + SLOW_QUERY_LOGGER + .set(build_slow_query_logger( + opts, + stdout_writer, + retention.as_ref(), + &mut guards, + )) + .expect("slow query logger already initialized"); + } if let Some(retention) = &retention { retention.initialize(); @@ -625,7 +602,6 @@ pub fn init_global_logging( .with(stdout_logging_layer) .with(file_logging_layer) .with(err_file_logging_layer) - .with(slow_query_logging_layer) }; // consume the `tracing_opts` to avoid "unused" warnings. @@ -637,8 +613,7 @@ pub fn init_global_logging( .with(dyn_trace_layer) .with(stdout_logging_layer) .with(file_logging_layer) - .with(err_file_logging_layer) - .with(slow_query_logging_layer); + .with(err_file_logging_layer); global::set_text_map_propagator(TraceContextPropagator::new()); @@ -727,59 +702,52 @@ fn build_otlp_exporter(opts: &LoggingOptions) -> SpanExporter { } } -fn build_slow_query_logger( +fn build_slow_query_logger( opts: &LoggingOptions, - slow_query_opts: Option<&SlowQueryOptions>, + stdout_writer: Option, retention: Option<&DirectoryRetention>, guards: &mut Vec, -) -> Option + Send + Sync + 'static>> -where - S: tracing::Subscriber - + Send - + 'static - + for<'span> tracing_subscriber::registry::LookupSpan<'span>, -{ - if let Some(slow_query_opts) = slow_query_opts { - if opts.enable_file_logging - && !opts.dir.is_empty() - && slow_query_opts.enable - && slow_query_opts.record_type == SlowQueriesRecordType::Log - { - let rolling_appender = build_file_appender(opts, LogFileKind::SlowQuery, retention); - let (writer, guard) = tracing_appender::non_blocking(rolling_appender); - guards.push(guard); - - // Only logs if the field contains "slow". - let slow_query_filter = FilterFn::new(|metadata| { - metadata - .fields() - .iter() - .any(|field| field.name().contains("slow")) - }); - - if opts.log_format == LogFormat::Json { - Some( - Layer::new() - .json() - .with_writer(writer) - .with_ansi(false) - .with_filter(slow_query_filter) - .boxed(), - ) - } else { - Some( - Layer::new() - .with_writer(writer) - .with_ansi(false) - .with_filter(slow_query_filter) - .boxed(), - ) - } - } else { - None - } +) -> tracing::Dispatch { + let stdout_layer = stdout_writer.map(|writer| { + build_logging_layer(writer, opts.log_format, std::io::stdout().is_terminal()) + }); + let file_layer = if opts.enable_file_logging && !opts.dir.is_empty() { + let rolling_appender = build_file_appender(opts, LogFileKind::SlowQuery, retention); + let (writer, guard) = tracing_appender::non_blocking(rolling_appender); + guards.push(guard); + Some(build_logging_layer(writer, opts.log_format, false)) } else { None + }; + // Avoid enabling unrelated callsites through this independently registered dispatcher. + let filter = FilterFn::new(|metadata| { + *metadata.level() == tracing::Level::INFO && metadata.fields().field("slow").is_some() + }) + .with_max_level_hint(LevelFilter::INFO); + tracing::Dispatch::new( + Registry::default() + .with(filter) + .with(stdout_layer) + .with(file_layer), + ) +} + +fn build_logging_layer( + writer: NonBlocking, + log_format: LogFormat, + ansi: bool, +) -> Box + Send + Sync> +where + S: tracing::Subscriber + for<'span> tracing_subscriber::registry::LookupSpan<'span>, +{ + if log_format == LogFormat::Json { + Layer::new() + .json() + .with_writer(writer) + .with_ansi(ansi) + .boxed() + } else { + Layer::new().with_writer(writer).with_ansi(ansi).boxed() } } @@ -794,6 +762,176 @@ mod tests { use super::*; + #[test] + fn test_slow_query_logging() { + const CASE_ENV: &str = "GREPTIME_TEST_SLOW_QUERY_LOGGING"; + let Ok(case) = env::var(CASE_ENV) else { + // Each configuration needs a fresh process because logging initializes only once. + for case in [ + "text", + "json", + "disabled", + "system_table", + "file_disabled", + "stdout_only_json", + "stdout_disabled", + "no_directory", + "no_options", + ] { + let output = std::process::Command::new(env::current_exe().unwrap()) + .args([ + "--exact", + "logging::tests::test_slow_query_logging", + "--nocapture", + ]) + .env(CASE_ENV, case) + .output() + .unwrap(); + assert!( + output.status.success(), + "case {case} failed:\n{}\n{}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr), + ); + let stdout = String::from_utf8(output.stdout).unwrap(); + let expected_slow_queries = if matches!( + case, + "text" | "json" | "file_disabled" | "stdout_only_json" | "no_directory" + ) { + 4 + } else { + 0 + }; + assert_eq!( + stdout.matches("SELECT 1").count(), + expected_slow_queries, + "{case}: {stdout}" + ); + assert_eq!( + stdout.matches("SELECT 2").count(), + expected_slow_queries, + "{case}: {stdout}" + ); + assert_eq!( + stdout.matches("custom_slow_target").count(), + expected_slow_queries, + "{case}: {stdout}" + ); + if case != "stdout_disabled" { + assert_eq!( + stdout.matches("ordinary-info").count(), + 1, + "{case}: {stdout}" + ); + assert_eq!( + stdout.matches("ordinary-warning").count(), + 2, + "{case}: {stdout}" + ); + } else { + assert!(!stdout.contains("ordinary-"), "{stdout}"); + } + if matches!(case, "json" | "stdout_only_json") { + for line in stdout.lines().filter(|line| line.contains("SELECT")) { + let event: serde_json::Value = serde_json::from_str(line).unwrap(); + assert_eq!(event["fields"]["slow"], true); + assert_eq!(event["level"], "INFO"); + } + } + } + return; + }; + + let dir = tempfile::tempdir().unwrap(); + let opts = LoggingOptions { + dir: if case == "no_directory" { + String::new() + } else { + dir.path().to_str().unwrap().to_string() + }, + level: Some("warn".to_string()), + log_format: if matches!(case.as_str(), "json" | "stdout_only_json") { + LogFormat::Json + } else { + LogFormat::Text + }, + append_stdout: case != "stdout_disabled", + enable_file_logging: !matches!(case.as_str(), "file_disabled" | "stdout_only_json"), + max_log_dir_size: ReadableSize::mb(1), + ..Default::default() + }; + let slow_opts = SlowQueryOptions { + enable: case != "disabled", + record_type: if case == "system_table" { + SlowQueriesRecordType::SystemTable + } else { + SlowQueriesRecordType::Log + }, + ..Default::default() + }; + let guards = init_global_logging( + "slow-query-test", + &opts, + &TracingOptions::default(), + None, + (case != "no_options").then_some(&slow_opts), + ); + + for level in ["warn", "error", "off", "info"] { + LOG_RELOAD_HANDLE + .get() + .unwrap() + .reload(level.parse::().unwrap()) + .unwrap(); + crate::info!(slow_operation = true, "ordinary-info"); + crate::warn!("ordinary-warning"); + crate::slow!(query = "SELECT 1", cost = 31, threshold = 30); + crate::slow!(target: "custom_slow_target", query = "SELECT 2"); + } + // Dropping the guards flushes the asynchronous writers before reading the files. + drop(guards); + + let mut slow_logs = String::new(); + let mut regular_logs = String::new(); + for entry in std::fs::read_dir(dir.path()).unwrap() { + let entry = entry.unwrap(); + let name = entry.file_name(); + let name = name.to_str().unwrap(); + if name.starts_with("greptimedb-slow-queries.") { + slow_logs.push_str(&std::fs::read_to_string(entry.path()).unwrap()); + } else if name.starts_with("greptimedb.") { + regular_logs.push_str(&std::fs::read_to_string(entry.path()).unwrap()); + } + } + + if matches!(case.as_str(), "text" | "json" | "stdout_disabled") { + assert_eq!(slow_logs.lines().count(), 8, "{slow_logs}"); + assert_eq!(slow_logs.matches("SELECT 1").count(), 4); + assert_eq!(slow_logs.matches("SELECT 2").count(), 4); + assert_eq!(slow_logs.matches("custom_slow_target").count(), 4); + assert!(!slow_logs.contains("ordinary-")); + if case == "json" { + for line in slow_logs.lines() { + let event: serde_json::Value = serde_json::from_str(line).unwrap(); + assert_eq!(event["fields"]["slow"], true); + assert_eq!(event["level"], "INFO"); + } + } + } else { + assert!(slow_logs.is_empty(), "{slow_logs}"); + } + if matches!( + case.as_str(), + "file_disabled" | "stdout_only_json" | "no_directory" + ) { + assert!(regular_logs.is_empty(), "{regular_logs}"); + } else { + assert_eq!(regular_logs.matches("ordinary-info").count(), 1); + assert_eq!(regular_logs.matches("ordinary-warning").count(), 2); + assert!(!regular_logs.contains("SELECT"), "{regular_logs}"); + } + } + #[test] fn test_trace_switch_preserves_active_spans() { let provider = SdkTracerProvider::builder() diff --git a/src/common/telemetry/src/macros.rs b/src/common/telemetry/src/macros.rs index 2846bd5d205..d92d5c17594 100644 --- a/src/common/telemetry/src/macros.rs +++ b/src/common/telemetry/src/macros.rs @@ -152,14 +152,19 @@ macro_rules! trace { }; } +/// Records a slow query independently of the global log level. #[macro_export] macro_rules! slow { (target: $target:expr, $($arg:tt)+) => { - $crate::log!(target: $target, slow = true, $crate::tracing::Level::INFO, $($arg)+) + if let Some(logger) = $crate::logging::SLOW_QUERY_LOGGER.get() { + $crate::tracing::dispatcher::with_default(logger, || { + $crate::log!(target: $target, $crate::tracing::Level::INFO, slow = true, $($arg)+) + }); + } }; ($($arg:tt)+) => { - $crate::log!($crate::tracing::Level::INFO, slow = true, $($arg)+) + $crate::slow!(target: module_path!(), $($arg)+) }; } From 977e76d0a4100ce9b6f2ccd9e171769ba362428a Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Tue, 20 Oct 2026 07:46:56 +0000 Subject: [PATCH 051/239] fix(mito2): apply SST internal column encodings to bulk-encoded parts (#9735) * fix(mito2): apply SST internal column encodings to bulk-encoded parts Bulk memtable parts are encoded to parquet and flushed as SSTs as-is, but the bulk encoder did not apply the time index / __sequence DELTA_BINARY_PACKED encoding (with dictionary disabled) and the uncompressed op_type setting that the regular SST writer uses. Share these settings through one helper. Signed-off-by: Dennis Zhuang * chore(mito2): drop the answered bulk encoder compression TODO Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/mito2/src/memtable/bulk.rs | 42 ++++++++++++++++++++++++++++- src/mito2/src/memtable/bulk/part.rs | 4 +-- src/mito2/src/sst/parquet.rs | 19 +++++++++++++ src/mito2/src/sst/parquet/writer.rs | 16 ++--------- 4 files changed, 64 insertions(+), 17 deletions(-) diff --git a/src/mito2/src/memtable/bulk.rs b/src/mito2/src/memtable/bulk.rs index 7402eeb0582..847a12db5ae 100644 --- a/src/mito2/src/memtable/bulk.rs +++ b/src/mito2/src/memtable/bulk.rs @@ -1692,9 +1692,10 @@ mod tests { use datatypes::schema::ColumnSchema; use datatypes::types::json_type::{JsonNativeType, JsonObjectType}; use mito_codec::row_converter::build_primary_key_codec; - use parquet::basic::{Encoding, Type as PhysicalType}; + use parquet::basic::{Compression, Encoding, Type as PhysicalType}; use serde_json::json; use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder, RegionMetadataRef}; + use store_api::storage::consts::{OP_TYPE_COLUMN_NAME, SEQUENCE_COLUMN_NAME}; use super::*; use crate::memtable::bulk::part::BulkPartConverter; @@ -3060,6 +3061,45 @@ mod tests { Ok(()) } + #[test] + fn test_bulk_part_encoder_matches_sst_writer_internal_encodings() -> WhateverResult<()> { + let metadata = metadata_for_test(); + let bulk_part = create_bulk_part_with_converter( + "key", + 0, + vec![1000, 2000, 3000], + vec![Some(1.0), Some(2.0), Some(3.0)], + 100, + )?; + let encoder = BulkPartEncoder::new(metadata, DEFAULT_ROW_GROUP_SIZE)?; + let encoded_part = encoder.encode_part(&bulk_part)?.unwrap(); + let columns = encoded_part + .metadata() + .parquet_metadata + .row_group(0) + .columns(); + + for name in ["ts", SEQUENCE_COLUMN_NAME] { + let column = columns + .iter() + .find(|column| column.column_path().string() == name) + .unwrap(); + assert!( + column + .encodings() + .any(|encoding| encoding == Encoding::DELTA_BINARY_PACKED), + "{name}" + ); + assert!(column.dictionary_page_offset().is_none(), "{name}"); + } + let op_type = columns + .iter() + .find(|column| column.column_path().string() == OP_TYPE_COLUMN_NAME) + .unwrap(); + assert_eq!(Compression::UNCOMPRESSED, op_type.compression()); + Ok(()) + } + #[test] fn test_encoded_part_batch_size_uses_largest_uncompressed_row_group() { const NUM_ROWS: usize = 14; diff --git a/src/mito2/src/memtable/bulk/part.rs b/src/mito2/src/memtable/bulk/part.rs index 27fc7983513..9bcc4238b45 100644 --- a/src/mito2/src/memtable/bulk/part.rs +++ b/src/mito2/src/memtable/bulk/part.rs @@ -74,7 +74,7 @@ use crate::sst::parquet::flat_format::primary_key_column_index; use crate::sst::parquet::format::{PrimaryKeyArray, PrimaryKeyArrayBuilder}; use crate::sst::parquet::{ COLUMN_INDEX_TRUNCATE_LENGTH, PARQUET_METADATA_KEY, SstInfo, apply_column_compression, - apply_float_field_encoding, + apply_float_field_encoding, apply_internal_column_encoding, }; const INIT_DICT_VALUE_CAPACITY: usize = 8; @@ -1339,7 +1339,6 @@ impl BulkPartEncoder { let key_value_meta = parquet::file::metadata::KeyValue::new(PARQUET_METADATA_KEY.to_string(), json); - // TODO(yingwen): Do we need compression? let mut props = WriterProperties::builder() .set_key_value_metadata(Some(vec![key_value_meta])) .set_write_batch_size(row_group_size) @@ -1347,6 +1346,7 @@ impl BulkPartEncoder { .set_compression(Compression::ZSTD(ZstdLevel::default())) .set_column_index_truncate_length(COLUMN_INDEX_TRUNCATE_LENGTH) .set_statistics_truncate_length(None); + props = apply_internal_column_encoding(props, &metadata); props = apply_float_field_encoding(props, &metadata, float_field_encoding); props = apply_column_compression(props, &metadata)?; let writer_props = Some(props.build()); diff --git a/src/mito2/src/sst/parquet.rs b/src/mito2/src/sst/parquet.rs index 79ad957b75a..b1b459a595a 100644 --- a/src/mito2/src/sst/parquet.rs +++ b/src/mito2/src/sst/parquet.rs @@ -28,6 +28,7 @@ use parquet::schema::types::ColumnPath; use snafu::ResultExt; use store_api::metadata::RegionMetadataRef; use store_api::mito_engine_options::FloatFieldEncoding; +use store_api::storage::consts::{OP_TYPE_COLUMN_NAME, SEQUENCE_COLUMN_NAME}; use store_api::storage::{ColumnId, FileId}; use crate::error::{CompressionOptionsSnafu, Result, WriteParquetSnafu}; @@ -93,6 +94,24 @@ pub const DEFAULT_ROW_GROUP_SIZE: usize = 100 * 1024; pub(crate) const COLUMN_INDEX_TRUNCATE_LENGTH: Option = parquet::file::properties::DEFAULT_COLUMN_INDEX_TRUNCATE_LENGTH; +/// Applies the encodings of the time index and internal columns shared by every SST writer. +pub(crate) fn apply_internal_column_encoding( + builder: WriterPropertiesBuilder, + metadata: &RegionMetadataRef, +) -> WriterPropertiesBuilder { + let ts_col = ColumnPath::new(vec![ + metadata.time_index_column().column_schema.name.clone(), + ]); + let seq_col = ColumnPath::new(vec![SEQUENCE_COLUMN_NAME.to_string()]); + let op_type_col = ColumnPath::new(vec![OP_TYPE_COLUMN_NAME.to_string()]); + builder + .set_column_encoding(seq_col.clone(), Encoding::DELTA_BINARY_PACKED) + .set_column_dictionary_enabled(seq_col, false) + .set_column_encoding(ts_col.clone(), Encoding::DELTA_BINARY_PACKED) + .set_column_dictionary_enabled(ts_col, false) + .set_column_compression(op_type_col, Compression::UNCOMPRESSED) +} + /// Applies the per-column options declared with `COMPRESSION WITH (...)`. pub(crate) fn apply_column_compression( mut builder: WriterPropertiesBuilder, diff --git a/src/mito2/src/sst/parquet/writer.rs b/src/mito2/src/sst/parquet/writer.rs index 9d2defcceb3..69a89e5802b 100644 --- a/src/mito2/src/sst/parquet/writer.rs +++ b/src/mito2/src/sst/parquet/writer.rs @@ -39,11 +39,9 @@ use parquet::arrow::AsyncArrowWriter; use parquet::basic::{Compression, Encoding, ZstdLevel}; use parquet::file::metadata::KeyValue; use parquet::file::properties::WriterProperties; -use parquet::schema::types::ColumnPath; use smallvec::smallvec; use snafu::{OptionExt, ResultExt}; use store_api::metadata::RegionMetadataRef; -use store_api::storage::consts::{OP_TYPE_COLUMN_NAME, SEQUENCE_COLUMN_NAME}; use store_api::storage::{FileId, SequenceNumber}; use tokio::io::AsyncWrite; use tokio_util::compat::{Compat, FuturesAsyncWriteCompatExt}; @@ -63,7 +61,7 @@ use crate::sst::parquet::flat_format::{ use crate::sst::parquet::format::{PrimaryKeyArray, PrimaryKeyWriteFormat}; use crate::sst::parquet::{ COLUMN_INDEX_TRUNCATE_LENGTH, PARQUET_METADATA_KEY, SstInfo, WriteOptions, - apply_column_compression, apply_float_field_encoding, + apply_column_compression, apply_float_field_encoding, apply_internal_column_encoding, }; use crate::sst::{ DEFAULT_WRITE_CONCURRENCY, FlatSchemaOptions, SeriesEstimator, maybe_wrap_schema, @@ -547,17 +545,7 @@ where .set_max_row_group_row_count(Some(opts.row_group_size)) .set_column_index_truncate_length(COLUMN_INDEX_TRUNCATE_LENGTH) .set_statistics_truncate_length(None); - let ts_col = ColumnPath::new(vec![ - self.metadata.time_index_column().column_schema.name.clone(), - ]); - let seq_col = ColumnPath::new(vec![SEQUENCE_COLUMN_NAME.to_string()]); - let op_type_col = ColumnPath::new(vec![OP_TYPE_COLUMN_NAME.to_string()]); - let props_builder = props_builder - .set_column_encoding(seq_col.clone(), Encoding::DELTA_BINARY_PACKED) - .set_column_dictionary_enabled(seq_col, false) - .set_column_encoding(ts_col.clone(), Encoding::DELTA_BINARY_PACKED) - .set_column_dictionary_enabled(ts_col, false) - .set_column_compression(op_type_col, Compression::UNCOMPRESSED); + let props_builder = apply_internal_column_encoding(props_builder, &self.metadata); let props_builder = apply_float_field_encoding( props_builder, &self.metadata, From b9aa627180ac696b83de2e675a8ceb0910295216 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Tue, 20 Oct 2026 08:03:35 +0000 Subject: [PATCH 052/239] fix(query): forward scalar subquery filters to region scans (#9738) * fix(query): forward scalar subquery filters to region scans Since DataFusion 55, uncorrelated scalar subqueries run in ScalarSubqueryExec instead of being rewritten into joins, so filters such as ts = (SELECT ts ... ORDER BY ts DESC LIMIT 1) no longer reach region scans through a join dynamic filter. MergeScanExec now captures such filters as remote dynamic filters, keeps them in the parent FilterExec, and replaces the subquery with its result before querying regions. Signed-off-by: Dennis Zhuang * fix(query): check the remote dyn filter budget across pushdown phases Filters captured in the Pre and Post phases are sent in one request, but each phase validated the initial registration budget alone. When their sum exceeded the budget, serialization dropped every registration, including join and TopK filters. Validate the combined set and keep filters accepted by the other phase. Signed-off-by: Dennis Zhuang * fix(query): keep IN-bearing scalar subquery filters at the parent Oversized remote dyn filter payloads replace IN lists with true. That only widens a conjunction of bounds, but inside an arbitrary predicate such as (host IN (...)) = (SELECT flag ...) it can drop rows the parent filter would keep. Do not forward scalar subquery filters that contain IN lists. Signed-off-by: Dennis Zhuang * chore(query): note the IN fallback restriction for scalar subquery filters Signed-off-by: Dennis Zhuang * chore(query): document the scalar subquery dyn filter bridge removal condition Signed-off-by: Dennis Zhuang * test(query): keep only dyn_filters of region scans in the scalar subquery plan case #9565 added scan diagnostics to EXPLAIN ANALYZE VERBOSE and partition filter rendering changed; the case only checks the forwarded dyn filter. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/query/AGENTS.md | 3 + src/query/src/dist_plan/dyn_filter_bridge.rs | 375 ++++++++++++++++-- src/query/src/dist_plan/merge_scan.rs | 59 ++- .../time_index_filter_pushdown.result | 82 ++++ .../optimizer/time_index_filter_pushdown.sql | 33 ++ .../time_index_filter_pushdown.result | 82 ++++ .../optimizer/time_index_filter_pushdown.sql | 33 ++ 7 files changed, 619 insertions(+), 48 deletions(-) diff --git a/src/query/AGENTS.md b/src/query/AGENTS.md index 8dc03a22877..fe794dcfc73 100644 --- a/src/query/AGENTS.md +++ b/src/query/AGENTS.md @@ -24,6 +24,9 @@ side effects live in `operator`. - Optimizer or `MergeScan` changes usually touch rule ordering, physical planning, serialization, and distributed-plan tests together. - User-visible SQL/PromQL/log behavior needs planner coverage and a sqlness case. +- `MergeScanExec` sends remote dyn filters captured in every filter pushdown + phase in one request; the initial registration budget applies to all of them + together (`dist_plan/dyn_filter_bridge.rs`). - `RegionQueryHandler` changes require the frontend implementation to move with it. - DataFusion dependency changes follow the pinned-fork rule in the repo-wide diff --git a/src/query/src/dist_plan/dyn_filter_bridge.rs b/src/query/src/dist_plan/dyn_filter_bridge.rs index 4761c7aea93..fb49839b304 100644 --- a/src/query/src/dist_plan/dyn_filter_bridge.rs +++ b/src/query/src/dist_plan/dyn_filter_bridge.rs @@ -19,9 +19,16 @@ use common_query::request::{ DynFilterPayload, INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY, InitialDynFilterReg, InitialDynFilterRegs, InitialDynFilterSnapshot, REMOTE_DYN_FILTER_PAYLOAD_MAX_BYTES, }; +use datafusion::physical_plan::filter_pushdown::PushedDown; use datafusion_common::Result; +use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; +use datafusion_expr::ColumnarValue; use datafusion_physical_expr::PhysicalExpr; -use datafusion_physical_expr::expressions::DynamicFilterPhysicalExpr; +use datafusion_physical_expr::expressions::{ + Column, DynamicFilterPhysicalExpr, InListExpr, Literal, lit, +}; +use datafusion_physical_expr::scalar_subquery::ScalarSubqueryExpr; +use datafusion_physical_expr::utils::collect_columns; use session::context::{QueryContext, QueryContextRef}; use store_api::storage::RegionId; @@ -36,49 +43,82 @@ pub(crate) struct CapturedDynFilter { filter_id: FilterId, initial_registration: InitialDynFilterReg, pub(crate) alive_dyn_filter: Arc, + /// A parent filter that references uncorrelated scalar subqueries. Its value is only + /// known once `ScalarSubqueryExec` has run the subqueries, so `alive_dyn_filter` + /// starts as `true` and is replaced by the resolved filter before regions are queried. + // TODO(discord9): temporary bridge. DataFusion fills the scalar subquery results in + // `ScalarSubqueryExec` but does not expose predicates using them as + // `DynamicFilterPhysicalExpr`, so we keep the predicate and bind it ourselves. Once + // upstream does and the pinned DataFusion includes it, remove this field and the + // scalar subquery capture/resolve path in favor of the regular dyn filter capture. + // Keep the per-phase bookkeeping in `MergeScanExec`. + scalar_subquery_filter: Option>, } #[derive(Debug, Clone)] pub(crate) struct RemoteDynFilterPushdown { pub(crate) captured_dyn_filters: Vec, /// Preflight result per parent filter. - pub(crate) pushed_down: Vec, + pub(crate) pushed_down: Vec, } +/// Captures `parent_filters` as remote dyn filters. `accepted` holds filters captured by +/// another pushdown phase; they are sent in the same request, so the initial registration +/// budget is checked against both and they are kept if the new filters do not fit. pub(crate) fn capture_remote_dyn_filters_for_pushdown( remote_dyn_filter_producer_id: RemoteDynFilterProducerId, + first_producer_local_ordinal: usize, + accepted: &[CapturedDynFilter], parent_filters: Vec>, ) -> RemoteDynFilterPushdown { let mut pushed_down = Vec::with_capacity(parent_filters.len()); let mut captured_dyn_filters = Vec::new(); - for (producer_local_ordinal, filter) in parent_filters.into_iter().enumerate() { - let Some(alive_dyn_filter) = downcast_dynamic_filter(filter) else { - pushed_down.push(false); + for (index, filter) in parent_filters.into_iter().enumerate() { + let producer_local_ordinal = first_producer_local_ordinal + index; + let captured = if let Some(alive_dyn_filter) = downcast_dynamic_filter(filter.clone()) { + build_captured_dyn_filter( + remote_dyn_filter_producer_id, + producer_local_ordinal, + alive_dyn_filter, + ) + .map(|captured| (captured, PushedDown::Yes)) + } else if is_scalar_subquery_filter(&filter) { + // The filter is only forwarded to regions for pruning; the parent keeps + // evaluating it, so report it as not pushed down. + build_scalar_subquery_dyn_filter( + remote_dyn_filter_producer_id, + producer_local_ordinal, + filter, + ) + .map(|captured| (captured, PushedDown::No)) + } else { + pushed_down.push(PushedDown::No); continue; }; - match build_captured_dyn_filter( - remote_dyn_filter_producer_id, - producer_local_ordinal, - alive_dyn_filter, - ) { - Ok(captured_dyn_filter) => { - pushed_down.push(true); + match captured { + Ok((captured_dyn_filter, pushdown)) => { + pushed_down.push(pushdown); captured_dyn_filters.push(captured_dyn_filter); } Err(error) => { common_telemetry::warn!(error; "Remote dyn filter is not pushed down because initial registration cannot be built"); - pushed_down.push(false); + pushed_down.push(PushedDown::No); } } } - if let Err(error) = validate_initial_registrations_for_pushdown(&captured_dyn_filters) { + let combined = accepted + .iter() + .chain(&captured_dyn_filters) + .cloned() + .collect::>(); + if let Err(error) = validate_initial_registrations_for_pushdown(&combined) { common_telemetry::warn!(error; "Remote dyn filters are not pushed down because initial registrations are invalid"); return RemoteDynFilterPushdown { captured_dyn_filters: Vec::new(), - pushed_down: vec![false; pushed_down.len()], + pushed_down: vec![PushedDown::No; pushed_down.len()], }; } @@ -88,6 +128,105 @@ pub(crate) fn capture_remote_dyn_filters_for_pushdown( } } +/// Returns true if `filter` references at least one scalar subquery and is otherwise +/// built only from columns and literals, so it becomes a plain predicate once the +/// subqueries are resolved. +/// +/// Filters containing `IN` lists are rejected: an oversized payload replaces `IN` with +/// `true`, which only widens a conjunction but can narrow an arbitrary predicate such as +/// `(host IN (...)) = false`. +// TODO(dennis): make the bounds-only fallback in `DynFilterPayload` replace only `IN` lists +// in conjunctive positions, then drop this restriction. +fn is_scalar_subquery_filter(filter: &Arc) -> bool { + let mut has_scalar_subquery = false; + let mut supported = true; + let _ = filter.apply(|expr| { + if expr.is::() { + has_scalar_subquery = true; + } else if expr.is::() + || (expr.children().is_empty() && !expr.is::() && !expr.is::()) + { + supported = false; + return Ok(TreeNodeRecursion::Stop); + } + Ok(TreeNodeRecursion::Continue) + }); + has_scalar_subquery && supported +} + +fn build_scalar_subquery_dyn_filter( + remote_dyn_filter_producer_id: RemoteDynFilterProducerId, + producer_local_ordinal: usize, + filter: Arc, +) -> Result { + let mut columns = collect_columns(&filter).into_iter().collect::>(); + columns.sort_by_key(|column| column.index()); + let children = columns + .into_iter() + .map(|column| Arc::new(column) as Arc) + .collect(); + let alive_dyn_filter = Arc::new(DynamicFilterPhysicalExpr::new(children, lit(true))); + let mut captured = build_captured_dyn_filter( + remote_dyn_filter_producer_id, + producer_local_ordinal, + alive_dyn_filter, + )?; + captured.scalar_subquery_filter = Some(filter); + Ok(captured) +} + +/// Replaces the scalar subqueries in captured filters with their results. +/// +/// `ScalarSubqueryExec` runs its subqueries before executing its input, so the results +/// are available when `MergeScanExec` is executed. This runs on every execution because +/// `ScalarSubqueryExec::reset_state` clears the results for re-execution. A filter whose +/// subqueries are unresolved falls back to `true`, which only disables pruning. +pub(crate) fn resolve_scalar_subquery_dyn_filters(captured_dyn_filters: &[CapturedDynFilter]) { + for captured in captured_dyn_filters { + let Some(filter) = &captured.scalar_subquery_filter else { + continue; + }; + let resolved = replace_scalar_subqueries(filter).unwrap_or_else(|error| { + common_telemetry::debug!( + "Scalar subquery filter {} is not forwarded to regions: {}", + filter, + error + ); + lit(true) + }); + if captured + .alive_dyn_filter + .current() + .is_ok_and(|current| current.as_ref() == resolved.as_ref()) + { + continue; + } + if let Err(error) = captured.alive_dyn_filter.update(resolved) { + common_telemetry::warn!(error; "Failed to update scalar subquery remote dyn filter"); + } + } +} + +fn replace_scalar_subqueries(filter: &Arc) -> Result> { + let empty_batch = + arrow::record_batch::RecordBatch::new_empty(Arc::new(arrow_schema::Schema::empty())); + filter + .clone() + .transform_up(|expr| { + if !expr.is::() { + return Ok(Transformed::no(expr)); + } + // Fails if the subquery has not been executed yet. + match expr.evaluate(&empty_batch)? { + ColumnarValue::Scalar(value) => Ok(Transformed::yes(lit(value))), + ColumnarValue::Array(_) => Err(datafusion_common::DataFusionError::Internal( + "scalar subquery evaluated to an array".to_string(), + )), + } + }) + .map(|transformed| transformed.data) +} + fn downcast_dynamic_filter( expr: Arc, ) -> Option> { @@ -159,6 +298,7 @@ fn build_captured_dyn_filter( filter_id, initial_registration: attach_initial_snapshot(initial_registration, &alive_dyn_filter), alive_dyn_filter, + scalar_subquery_filter: None, }) } @@ -315,7 +455,9 @@ mod tests { use datafusion::execution::TaskContext; use datafusion_common::ScalarValue; use datafusion_expr::ColumnarValue; - use datafusion_physical_expr::expressions::{Column, lit}; + use datafusion_expr::Operator; + use datafusion_expr::physical_planning_context::{ScalarSubqueryResults, SubqueryIndex}; + use datafusion_physical_expr::expressions::BinaryExpr; use session::query_id::QueryId; use uuid::Uuid; @@ -461,6 +603,133 @@ mod tests { .to_string() } + fn pushed_down(pushdown: &RemoteDynFilterPushdown) -> Vec { + pushdown + .pushed_down + .iter() + .map(|pushed_down| matches!(pushed_down, PushedDown::Yes)) + .collect() + } + + fn host_eq_scalar_subquery(results: &ScalarSubqueryResults) -> Arc { + Arc::new(BinaryExpr::new( + Arc::new(Column::new("host", 0)), + Operator::Eq, + Arc::new(ScalarSubqueryExpr::new( + arrow_schema::DataType::Utf8, + true, + SubqueryIndex::new(0), + results.clone(), + )), + )) + } + + #[test] + fn scalar_subquery_filter_is_captured_but_left_to_parent() { + let results = ScalarSubqueryResults::new(1); + let parent_filters = vec![ + host_eq_scalar_subquery(&results), + Arc::new(BinaryExpr::new( + Arc::new(Column::new("host", 0)), + Operator::Eq, + lit("a"), + )) as Arc, + ]; + + let pushdown = capture_remote_dyn_filters_for_pushdown( + test_remote_dyn_filter_producer_id(42), + 0, + &[], + parent_filters, + ); + + assert_eq!(pushed_down(&pushdown), vec![false, false]); + assert_eq!(pushdown.captured_dyn_filters.len(), 1); + assert_eq!( + pushdown.captured_dyn_filters[0] + .filter_id + .producer_ordinal(), + 0 + ); + } + + #[test] + fn scalar_subquery_filter_with_in_list_is_left_to_parent() { + let results = ScalarSubqueryResults::new(1); + let schema = arrow_schema::Schema::new(vec![arrow_schema::Field::new( + "host", + arrow_schema::DataType::Utf8, + true, + )]); + let host_in_list = datafusion_physical_expr::expressions::in_list( + Arc::new(Column::new("host", 0)), + vec![lit("a"), lit("b")], + &false, + &schema, + ) + .unwrap(); + // `(host IN ('a', 'b')) = (SELECT flag ...)` + let filter = Arc::new(BinaryExpr::new( + host_in_list, + Operator::Eq, + Arc::new(ScalarSubqueryExpr::new( + arrow_schema::DataType::Boolean, + true, + SubqueryIndex::new(0), + results, + )), + )) as Arc; + + let pushdown = capture_remote_dyn_filters_for_pushdown( + test_remote_dyn_filter_producer_id(42), + 0, + &[], + vec![filter], + ); + + assert_eq!(pushed_down(&pushdown), vec![false]); + assert!(pushdown.captured_dyn_filters.is_empty()); + } + + #[test] + fn scalar_subquery_dyn_filter_follows_subquery_results() { + let results = ScalarSubqueryResults::new(1); + let captured = capture_remote_dyn_filters_for_pushdown( + test_remote_dyn_filter_producer_id(42), + 0, + &[], + vec![host_eq_scalar_subquery(&results)], + ) + .captured_dyn_filters; + let snapshot = |captured: &[CapturedDynFilter]| { + decode_snapshot( + refreshed_regs(captured).regs[0] + .initial_snapshot + .as_ref() + .unwrap(), + "host", + ) + }; + + // Not executed yet: regions must not be filtered. + resolve_scalar_subquery_dyn_filters(&captured); + assert_eq!(snapshot(&captured), "true"); + + results + .set( + SubqueryIndex::new(0), + ScalarValue::Utf8(Some("a".to_string())), + ) + .unwrap(); + resolve_scalar_subquery_dyn_filters(&captured); + assert_eq!(snapshot(&captured), "host@0 = a"); + + // `ScalarSubqueryExec::reset_state` clears the results before re-execution. + results.clear(); + resolve_scalar_subquery_dyn_filters(&captured); + assert_eq!(snapshot(&captured), "true"); + } + #[test] fn capture_remote_dyn_filters_for_pushdown_preserves_parent_filter_ordinals() { let parent_filters = vec![ @@ -477,9 +746,13 @@ mod tests { ]; let remote_dyn_filter_producer_id = test_remote_dyn_filter_producer_id(42); - let captured = - capture_remote_dyn_filters_for_pushdown(remote_dyn_filter_producer_id, parent_filters) - .captured_dyn_filters; + let captured = capture_remote_dyn_filters_for_pushdown( + remote_dyn_filter_producer_id, + 0, + &[], + parent_filters, + ) + .captured_dyn_filters; assert_eq!(captured.len(), 2); assert_eq!( @@ -506,10 +779,14 @@ mod tests { ]; let remote_dyn_filter_producer_id = test_remote_dyn_filter_producer_id(42); - let pushdown = - capture_remote_dyn_filters_for_pushdown(remote_dyn_filter_producer_id, parent_filters); + let pushdown = capture_remote_dyn_filters_for_pushdown( + remote_dyn_filter_producer_id, + 0, + &[], + parent_filters, + ); - assert_eq!(pushdown.pushed_down, vec![false, true, false]); + assert_eq!(pushed_down(&pushdown), vec![false, true, false]); assert_eq!(pushdown.captured_dyn_filters.len(), 1); assert_eq!( pushdown.captured_dyn_filters[0] @@ -541,10 +818,12 @@ mod tests { let pushdown = capture_remote_dyn_filters_for_pushdown( test_remote_dyn_filter_producer_id(42), + 0, + &[], parent_filters, ); - assert_eq!(pushdown.pushed_down, vec![false]); + assert_eq!(pushed_down(&pushdown), vec![false]); assert!(pushdown.captured_dyn_filters.is_empty()); } @@ -558,10 +837,12 @@ mod tests { let pushdown = capture_remote_dyn_filters_for_pushdown( test_remote_dyn_filter_producer_id(42), + 0, + &[], parent_filters, ); - assert_eq!(pushdown.pushed_down, vec![true]); + assert_eq!(pushed_down(&pushdown), vec![true]); assert!( pushdown.captured_dyn_filters[0] .initial_registration @@ -581,10 +862,12 @@ mod tests { let pushdown = capture_remote_dyn_filters_for_pushdown( test_remote_dyn_filter_producer_id(42), + 0, + &[], parent_filters, ); - assert_eq!(pushdown.pushed_down, vec![true]); + assert_eq!(pushed_down(&pushdown), vec![true]); let snapshot = pushdown.captured_dyn_filters[0] .initial_registration .initial_snapshot @@ -610,10 +893,42 @@ mod tests { let pushdown = capture_remote_dyn_filters_for_pushdown( test_remote_dyn_filter_producer_id(42), + 0, + &[], parent_filters, ); - assert_eq!(pushdown.pushed_down, vec![false, false]); + assert_eq!(pushed_down(&pushdown), vec![false, false]); + assert!(pushdown.captured_dyn_filters.is_empty()); + } + + #[test] + fn capture_remote_dyn_filters_for_pushdown_keeps_accepted_filters_within_budget() { + const MAX_INITIAL_REGS: usize = 64; + + let accepted = capture_remote_dyn_filters_for_pushdown( + test_remote_dyn_filter_producer_id(42), + 0, + &[], + (0..MAX_INITIAL_REGS) + .map(|ordinal| { + test_dyn_filter_with_snapshot_payload(&format!("host_{ordinal}"), ordinal, 1) + as Arc + }) + .collect(), + ) + .captured_dyn_filters; + assert_eq!(accepted.len(), MAX_INITIAL_REGS); + + let results = ScalarSubqueryResults::new(1); + let pushdown = capture_remote_dyn_filters_for_pushdown( + test_remote_dyn_filter_producer_id(42), + 1 << 31, + &accepted, + vec![host_eq_scalar_subquery(&results)], + ); + + assert_eq!(pushed_down(&pushdown), vec![false]); assert!(pushdown.captured_dyn_filters.is_empty()); } @@ -630,11 +945,13 @@ mod tests { let pushdown = capture_remote_dyn_filters_for_pushdown( test_remote_dyn_filter_producer_id(42), + 0, + &[], parent_filters, ); assert!(pushdown.captured_dyn_filters.is_empty()); - assert_eq!(pushdown.pushed_down, vec![false; TOO_MANY_INITIAL_REGS]); + assert_eq!(pushed_down(&pushdown), vec![false; TOO_MANY_INITIAL_REGS]); } #[test] @@ -652,11 +969,13 @@ mod tests { let pushdown = capture_remote_dyn_filters_for_pushdown( test_remote_dyn_filter_producer_id(42), + 0, + &[], parent_filters, ); assert!(pushdown.captured_dyn_filters.is_empty()); - assert_eq!(pushdown.pushed_down, vec![false; TOO_MANY_INITIAL_REGS]); + assert_eq!(pushed_down(&pushdown), vec![false; TOO_MANY_INITIAL_REGS]); } #[test] diff --git a/src/query/src/dist_plan/merge_scan.rs b/src/query/src/dist_plan/merge_scan.rs index 9717c436e0b..90a57bf1a5c 100644 --- a/src/query/src/dist_plan/merge_scan.rs +++ b/src/query/src/dist_plan/merge_scan.rs @@ -68,6 +68,7 @@ use crate::dist_plan::dyn_filter_bridge::{ CapturedDynFilter, capture_remote_dyn_filters_for_pushdown, query_context_with_refreshed_initial_dyn_filter_regs, register_dyn_filter_subscribers_for_region, register_remote_dyn_filters, + resolve_scalar_subquery_dyn_filters, }; use crate::dist_plan::{ FilterId, RemoteDynFilterProducerId, RemoteDynFilterRegistryLease, Subscriber, @@ -126,6 +127,22 @@ fn remote_plan_row_bound(plan: &LogicalPlan) -> Option { } } +fn pushdown_phase_slot(phase: FilterPushdownPhase) -> usize { + match phase { + FilterPushdownPhase::Pre => 0, + FilterPushdownPhase::Post => 1, + } +} + +/// Producer-local ordinals must be unique across phases because they are part of the +/// remote filter id; pre-phase filters use the upper half of the `u32` ordinal space. +fn first_producer_local_ordinal(phase: FilterPushdownPhase) -> usize { + match phase { + FilterPushdownPhase::Pre => 1 << 31, + FilterPushdownPhase::Post => 0, + } +} + fn remote_dyn_filter_enabled(query_ctx: &QueryContextRef) -> Result { remote_dyn_filter_pushdown_enabled_from_extensions(&query_ctx.extensions()) .map_err(|err| DataFusionError::External(Box::new(err))) @@ -483,7 +500,10 @@ pub struct MergeScanExec { query_ctx: QueryContextRef, /// Optional because RDF must fail open: missing ids skip RDF but keep normal query execution. remote_dyn_filter_producer_id: Option, - captured_remote_dyn_filters: Arc>>, + /// Filters captured by each filter pushdown phase, indexed by [`pushdown_phase_slot`]. + /// The pre-optimization phase carries `FilterExec` predicates (e.g. scalar subquery + /// filters) and the post-optimization phase carries dynamic filters from joins and TopK. + captured_remote_dyn_filters: Arc; 2]>>, target_partition: usize, enable_per_region_metrics: bool, } @@ -616,6 +636,7 @@ impl MergeScanExec { } else { Vec::new() }; + resolve_scalar_subquery_dyn_filters(&captured_remote_dyn_filters); let dbname = context.task_id().unwrap_or_default(); let tracing_context = TracingContext::from_json(context.session_id().as_str()); let current_channel = self.query_ctx.channel(); @@ -890,7 +911,7 @@ impl MergeScanExec { } fn captured_remote_dyn_filters(&self) -> Vec { - self.captured_remote_dyn_filters.lock().unwrap().clone() + self.captured_remote_dyn_filters.lock().unwrap().concat() } pub fn sub_stage_metrics(&self) -> Vec { @@ -1092,7 +1113,7 @@ impl ExecutionPlan for MergeScanExec { fn handle_child_pushdown_result( &self, - _phase: FilterPushdownPhase, + phase: FilterPushdownPhase, child_pushdown_result: ChildPushdownResult, _config: &datafusion::config::ConfigOptions, ) -> Result>> { @@ -1106,7 +1127,7 @@ impl ExecutionPlan for MergeScanExec { // Reject remote pushdown instead of pretending success: this keeps // DataFusion/local dynamic filter semantics intact while disabling // only FE -> DN remote dynamic filter propagation. - self.captured_remote_dyn_filters.lock().unwrap().clear(); + self.captured_remote_dyn_filters.lock().unwrap()[pushdown_phase_slot(phase)].clear(); let new_self = Arc::new(self.clone()); return Ok(FilterPushdownPropagation { @@ -1120,7 +1141,7 @@ impl ExecutionPlan for MergeScanExec { common_telemetry::warn!( "MergeScan remote dynamic filter producer id is not assigned; skipping remote dynamic filter pushdown" ); - self.captured_remote_dyn_filters.lock().unwrap().clear(); + self.captured_remote_dyn_filters.lock().unwrap()[pushdown_phase_slot(phase)].clear(); let new_self = Arc::new(self.clone()); return Ok(FilterPushdownPropagation { @@ -1128,24 +1149,20 @@ impl ExecutionPlan for MergeScanExec { updated_node: Some(new_self), }); }; - let remote_dyn_filter_pushdown = - capture_remote_dyn_filters_for_pushdown(remote_dyn_filter_producer_id, parent_filters); - *self.captured_remote_dyn_filters.lock().unwrap() = - remote_dyn_filter_pushdown.captured_dyn_filters; + let slot = pushdown_phase_slot(phase); + let mut captured_remote_dyn_filters = self.captured_remote_dyn_filters.lock().unwrap(); + let remote_dyn_filter_pushdown = capture_remote_dyn_filters_for_pushdown( + remote_dyn_filter_producer_id, + first_producer_local_ordinal(phase), + &captured_remote_dyn_filters[1 - slot], + parent_filters, + ); + captured_remote_dyn_filters[slot] = remote_dyn_filter_pushdown.captured_dyn_filters; + drop(captured_remote_dyn_filters); let new_self = Arc::new(self.clone()); Ok(FilterPushdownPropagation { - filters: remote_dyn_filter_pushdown - .pushed_down - .into_iter() - .map(|pushdown_ready| { - if pushdown_ready { - PushedDown::Yes - } else { - PushedDown::No - } - }) - .collect(), + filters: remote_dyn_filter_pushdown.pushed_down, updated_node: Some(new_self), }) } @@ -1713,6 +1730,8 @@ mod tests { )) as Arc; let captured = capture_remote_dyn_filters_for_pushdown( RemoteDynFilterProducerId::new(42), + 0, + &[], vec![dyn_filter], ); assert_eq!(captured.captured_dyn_filters.len(), 1); diff --git a/tests/cases/distributed/optimizer/time_index_filter_pushdown.result b/tests/cases/distributed/optimizer/time_index_filter_pushdown.result index f03818154c1..5530edbcead 100644 --- a/tests/cases/distributed/optimizer/time_index_filter_pushdown.result +++ b/tests/cases/distributed/optimizer/time_index_filter_pushdown.result @@ -160,6 +160,88 @@ WHERE | 8 | windows | 2023-06-12T02:10:00 | +------+---------+---------------------+ +-- The resolved scalar subquery reaches the region scans as a dynamic filter. +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (metrics.*) REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE (?m)^\|_\|_\|_SortExec:.*\n +-- SQLNESS REPLACE FilterExec:\sDynamicFilter\s\[[^\]]*\] FilterExec: DynamicFilter [ REDACTED ] +-- SQLNESS REPLACE "files":\s\[.*?\],\s +-- SQLNESS REPLACE IS\sNULL\sOR\sgreptime_timestamp@0\s>\s\d+ IS NULL OR greptime_timestamp@0 > REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE UnorderedScan:.*"dyn_filters":\s(\["DynamicFilter\s\[[^\]]*\]"\]).* UnorderedScan: REDACTED, "dyn_filters": $1 +EXPLAIN ANALYZE VERBOSE SELECT + rack, + os, + greptime_timestamp +FROM + cpu +WHERE + greptime_timestamp = ( + SELECT + greptime_timestamp + FROM + cpu + ORDER BY + greptime_timestamp DESC + LIMIT + 1 + ); + ++-+-+-+ +| stage | node | plan_| ++-+-+-+ +| 0_| 0_|_ScalarSubqueryExec REDACTED +|_|_|_FilterExec: greptime_timestamp@2 = scalar_subquery(1686535800000000000) REDACTED +|_|_|_ProjectionExec: expr=[rack@0 as rack, os@1 as os, greptime_timestamp@3 as greptime_timestamp] REDACTED +|_|_|_CooperativeExec REDACTED +|_|_|_MergeScanExec: REDACTED +|_|_|_MergeSortExec: [greptime_timestamp@0 DESC], fetch=1 REDACTED +|_|_|_MergeScanExec: REDACTED +|_|_|_| +| 1_| 0_|_CooperativeExec REDACTED +|_|_|_UnorderedScan: REDACTED, "dyn_filters": ["DynamicFilter [ greptime_timestamp@3 = 1686535800000000000 ]"] +|_|_|_| +| 1_| 1_|_CooperativeExec REDACTED +|_|_|_UnorderedScan: REDACTED, "dyn_filters": ["DynamicFilter [ greptime_timestamp@3 = 1686535800000000000 ]"] +|_|_|_| +| 1_| 2_|_CooperativeExec REDACTED +|_|_|_UnorderedScan: REDACTED, "dyn_filters": ["DynamicFilter [ greptime_timestamp@3 = 1686535800000000000 ]"] +|_|_|_| +| 1_| 3_|_CooperativeExec REDACTED +|_|_|_UnorderedScan: REDACTED, "dyn_filters": ["DynamicFilter [ greptime_timestamp@3 = 1686535800000000000 ]"] +|_|_|_| +| 1_| 4_|_CooperativeExec REDACTED +|_|_|_UnorderedScan: REDACTED, "dyn_filters": ["DynamicFilter [ greptime_timestamp@3 = 1686535800000000000 ]"] +|_|_|_| +| 1_| 0_|_FilterExec: DynamicFilter [ REDACTED ] REDACTED +|_|_|_SortPreservingMergeExec: [greptime_timestamp@0 DESC], fetch=1 REDACTED +|_|_|_UnorderedScan: REDACTED, "dyn_filters": ["DynamicFilter [ greptime_timestamp@0 IS NULL OR greptime_timestamp@0 > REDACTED ]"] +|_|_|_| +| 1_| 1_|_FilterExec: DynamicFilter [ REDACTED ] REDACTED +|_|_|_SortPreservingMergeExec: [greptime_timestamp@0 DESC], fetch=1 REDACTED +|_|_|_UnorderedScan: REDACTED, "dyn_filters": ["DynamicFilter [ greptime_timestamp@0 IS NULL OR greptime_timestamp@0 > REDACTED ]"] +|_|_|_| +| 1_| 2_|_FilterExec: DynamicFilter [ REDACTED ] REDACTED +|_|_|_SortPreservingMergeExec: [greptime_timestamp@0 DESC], fetch=1 REDACTED +|_|_|_UnorderedScan: REDACTED, "dyn_filters": ["DynamicFilter [ greptime_timestamp@0 IS NULL OR greptime_timestamp@0 > REDACTED ]"] +|_|_|_| +| 1_| 3_|_FilterExec: DynamicFilter [ REDACTED ] REDACTED +|_|_|_SortPreservingMergeExec: [greptime_timestamp@0 DESC], fetch=1 REDACTED +|_|_|_UnorderedScan: REDACTED, "dyn_filters": ["DynamicFilter [ greptime_timestamp@0 IS NULL OR greptime_timestamp@0 > REDACTED ]"] +|_|_|_| +| 1_| 4_|_FilterExec: DynamicFilter [ REDACTED ] REDACTED +|_|_|_SortPreservingMergeExec: [greptime_timestamp@0 DESC], fetch=1 REDACTED +|_|_|_UnorderedScan: REDACTED, "dyn_filters": ["DynamicFilter [ greptime_timestamp@0 IS NULL OR greptime_timestamp@0 > REDACTED ]"] +|_|_|_| +|_|_| Total rows: REDACTED_| ++-+-+-+ + drop table cpu; Affected Rows: 0 diff --git a/tests/cases/distributed/optimizer/time_index_filter_pushdown.sql b/tests/cases/distributed/optimizer/time_index_filter_pushdown.sql index c9c80c24a61..41e85a4f40a 100644 --- a/tests/cases/distributed/optimizer/time_index_filter_pushdown.sql +++ b/tests/cases/distributed/optimizer/time_index_filter_pushdown.sql @@ -102,4 +102,37 @@ WHERE 1 ); +-- The resolved scalar subquery reaches the region scans as a dynamic filter. +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (metrics.*) REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE (?m)^\|_\|_\|_SortExec:.*\n +-- SQLNESS REPLACE FilterExec:\sDynamicFilter\s\[[^\]]*\] FilterExec: DynamicFilter [ REDACTED ] +-- SQLNESS REPLACE "files":\s\[.*?\],\s +-- SQLNESS REPLACE IS\sNULL\sOR\sgreptime_timestamp@0\s>\s\d+ IS NULL OR greptime_timestamp@0 > REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE UnorderedScan:.*"dyn_filters":\s(\["DynamicFilter\s\[[^\]]*\]"\]).* UnorderedScan: REDACTED, "dyn_filters": $1 +EXPLAIN ANALYZE VERBOSE SELECT + rack, + os, + greptime_timestamp +FROM + cpu +WHERE + greptime_timestamp = ( + SELECT + greptime_timestamp + FROM + cpu + ORDER BY + greptime_timestamp DESC + LIMIT + 1 + ); + drop table cpu; diff --git a/tests/cases/standalone/optimizer/time_index_filter_pushdown.result b/tests/cases/standalone/optimizer/time_index_filter_pushdown.result index 50f39764ca2..106a78c49d2 100644 --- a/tests/cases/standalone/optimizer/time_index_filter_pushdown.result +++ b/tests/cases/standalone/optimizer/time_index_filter_pushdown.result @@ -223,6 +223,88 @@ WHERE | 8 | windows | 2023-06-12T02:10:00 | +------+---------+---------------------+ +-- The resolved scalar subquery reaches the region scans as a dynamic filter. +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (metrics.*) REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE (?m)^\|_\|_\|_SortExec:.*\n +-- SQLNESS REPLACE FilterExec:\sDynamicFilter\s\[[^\]]*\] FilterExec: DynamicFilter [ REDACTED ] +-- SQLNESS REPLACE "files":\s\[.*?\],\s +-- SQLNESS REPLACE IS\sNULL\sOR\sgreptime_timestamp@0\s>\s\d+ IS NULL OR greptime_timestamp@0 > REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE UnorderedScan:.*"dyn_filters":\s(\["DynamicFilter\s\[[^\]]*\]"\]).* UnorderedScan: REDACTED, "dyn_filters": $1 +EXPLAIN ANALYZE VERBOSE SELECT + rack, + os, + greptime_timestamp +FROM + cpu +WHERE + greptime_timestamp = ( + SELECT + greptime_timestamp + FROM + cpu + ORDER BY + greptime_timestamp DESC + LIMIT + 1 + ); + ++-+-+-+ +| stage | node | plan_| ++-+-+-+ +| 0_| 0_|_ScalarSubqueryExec REDACTED +|_|_|_FilterExec: greptime_timestamp@2 = scalar_subquery(1686535800000000000) REDACTED +|_|_|_ProjectionExec: expr=[rack@0 as rack, os@1 as os, greptime_timestamp@4 as greptime_timestamp] REDACTED +|_|_|_CooperativeExec REDACTED +|_|_|_MergeScanExec: REDACTED +|_|_|_MergeSortExec: [greptime_timestamp@0 DESC], fetch=1 REDACTED +|_|_|_MergeScanExec: REDACTED +|_|_|_| +| 1_| 0_|_CooperativeExec REDACTED +|_|_|_UnorderedScan: REDACTED, "dyn_filters": ["DynamicFilter [ greptime_timestamp@4 = 1686535800000000000 ]"] +|_|_|_| +| 1_| 1_|_CooperativeExec REDACTED +|_|_|_UnorderedScan: REDACTED, "dyn_filters": ["DynamicFilter [ greptime_timestamp@4 = 1686535800000000000 ]"] +|_|_|_| +| 1_| 2_|_CooperativeExec REDACTED +|_|_|_UnorderedScan: REDACTED, "dyn_filters": ["DynamicFilter [ greptime_timestamp@4 = 1686535800000000000 ]"] +|_|_|_| +| 1_| 3_|_CooperativeExec REDACTED +|_|_|_UnorderedScan: REDACTED, "dyn_filters": ["DynamicFilter [ greptime_timestamp@4 = 1686535800000000000 ]"] +|_|_|_| +| 1_| 4_|_CooperativeExec REDACTED +|_|_|_UnorderedScan: REDACTED, "dyn_filters": ["DynamicFilter [ greptime_timestamp@4 = 1686535800000000000 ]"] +|_|_|_| +| 1_| 0_|_FilterExec: DynamicFilter [ REDACTED ] REDACTED +|_|_|_SortPreservingMergeExec: [greptime_timestamp@0 DESC], fetch=1 REDACTED +|_|_|_UnorderedScan: REDACTED, "dyn_filters": ["DynamicFilter [ greptime_timestamp@0 IS NULL OR greptime_timestamp@0 > REDACTED ]"] +|_|_|_| +| 1_| 1_|_FilterExec: DynamicFilter [ REDACTED ] REDACTED +|_|_|_SortPreservingMergeExec: [greptime_timestamp@0 DESC], fetch=1 REDACTED +|_|_|_UnorderedScan: REDACTED, "dyn_filters": ["DynamicFilter [ greptime_timestamp@0 IS NULL OR greptime_timestamp@0 > REDACTED ]"] +|_|_|_| +| 1_| 2_|_FilterExec: DynamicFilter [ REDACTED ] REDACTED +|_|_|_SortPreservingMergeExec: [greptime_timestamp@0 DESC], fetch=1 REDACTED +|_|_|_UnorderedScan: REDACTED, "dyn_filters": ["DynamicFilter [ greptime_timestamp@0 IS NULL OR greptime_timestamp@0 > REDACTED ]"] +|_|_|_| +| 1_| 3_|_FilterExec: DynamicFilter [ REDACTED ] REDACTED +|_|_|_SortPreservingMergeExec: [greptime_timestamp@0 DESC], fetch=1 REDACTED +|_|_|_UnorderedScan: REDACTED, "dyn_filters": ["DynamicFilter [ greptime_timestamp@0 IS NULL OR greptime_timestamp@0 > REDACTED ]"] +|_|_|_| +| 1_| 4_|_FilterExec: DynamicFilter [ REDACTED ] REDACTED +|_|_|_SortPreservingMergeExec: [greptime_timestamp@0 DESC], fetch=1 REDACTED +|_|_|_UnorderedScan: REDACTED, "dyn_filters": ["DynamicFilter [ greptime_timestamp@0 IS NULL OR greptime_timestamp@0 > REDACTED ]"] +|_|_|_| +|_|_| Total rows: REDACTED_| ++-+-+-+ + drop table cpu; Affected Rows: 0 diff --git a/tests/cases/standalone/optimizer/time_index_filter_pushdown.sql b/tests/cases/standalone/optimizer/time_index_filter_pushdown.sql index 1dc010e1bf4..64f0fbc0e26 100644 --- a/tests/cases/standalone/optimizer/time_index_filter_pushdown.sql +++ b/tests/cases/standalone/optimizer/time_index_filter_pushdown.sql @@ -140,4 +140,37 @@ WHERE 1 ); +-- The resolved scalar subquery reaches the region scans as a dynamic filter. +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (metrics.*) REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE "flat_format":\s\w+, "flat_format": REDACTED, +-- SQLNESS REPLACE (?m)^\|_\|_\|_SortExec:.*\n +-- SQLNESS REPLACE FilterExec:\sDynamicFilter\s\[[^\]]*\] FilterExec: DynamicFilter [ REDACTED ] +-- SQLNESS REPLACE "files":\s\[.*?\],\s +-- SQLNESS REPLACE IS\sNULL\sOR\sgreptime_timestamp@0\s>\s\d+ IS NULL OR greptime_timestamp@0 > REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE UnorderedScan:.*"dyn_filters":\s(\["DynamicFilter\s\[[^\]]*\]"\]).* UnorderedScan: REDACTED, "dyn_filters": $1 +EXPLAIN ANALYZE VERBOSE SELECT + rack, + os, + greptime_timestamp +FROM + cpu +WHERE + greptime_timestamp = ( + SELECT + greptime_timestamp + FROM + cpu + ORDER BY + greptime_timestamp DESC + LIMIT + 1 + ); + drop table cpu; From a7230f7a92724b36fc6c5fa7b2708c983ba573f3 Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Tue, 20 Oct 2026 08:17:52 +0000 Subject: [PATCH 053/239] fix(partition): coerce partition exprs against the batch schema before routing (#9752) * fix(partition): coerce partition exprs against the batch schema before routing Batch partition routing (MultiDimPartitionRule::split_record_batch) builds physical expressions via PartitionExpr::try_as_physical_expr, which did no type coercion. After widening the time index unit (e.g. ms -> ns via ALTER TABLE ... MODIFY COLUMN ts TIMESTAMP_NS), the persisted partition bounds keep their original unit, so evaluating the bounds against the batch failed with: Invalid comparison operation: Timestamp(ns) < Timestamp(ms) Row inserts were unaffected because find_region compares Values and Timestamp::cmp converts units; scans were unaffected because PredicateGroup::new coerces the region partition expr against the region schema first. Coerce the logical partition expression against the evaluation schema in try_as_physical_expr before creating the physical expression, the same way PredicateGroup does. This also covers the region partition filter built for SST pruning (src/mito2/src/sst/parquet/reader.rs), which shares this conversion. Fixes #9747 Signed-off-by: Ning Sun * test: regenerate change_col_type result after comment update The ts_widen_part comment in change_col_type.sql was reworded after the result file was last generated; sqlness echoes SQL comments into the result, so the stale comment made the case fail. Regenerate it. Signed-off-by: Ning Sun --------- Signed-off-by: Ning Sun --- Cargo.lock | 2 + src/partition/Cargo.toml | 2 + src/partition/src/error.rs | 9 ++ src/partition/src/expr.rs | 17 ++- src/partition/src/multi_dim.rs | 44 ++++++- tests-integration/src/grpc/flight.rs | 112 +++++++++++++++++- .../common/alter/change_col_type.result | 35 ++++++ .../common/alter/change_col_type.sql | 10 ++ 8 files changed, 225 insertions(+), 6 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 946472d1df0..73bd260e9de 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10489,9 +10489,11 @@ dependencies = [ "common-error", "common-macro", "common-meta", + "common-time", "criterion 0.7.0", "datafusion-common", "datafusion-expr", + "datafusion-optimizer", "datafusion-physical-expr", "datatypes", "moka", diff --git a/src/partition/Cargo.toml b/src/partition/Cargo.toml index 6773cc5878b..43f6c8e792a 100644 --- a/src/partition/Cargo.toml +++ b/src/partition/Cargo.toml @@ -15,6 +15,7 @@ common-macro.workspace = true common-meta.workspace = true datafusion-common.workspace = true datafusion-expr.workspace = true +datafusion-optimizer.workspace = true datafusion-physical-expr.workspace = true datatypes.workspace = true moka = { workspace = true, features = ["future"] } @@ -27,6 +28,7 @@ store-api.workspace = true table.workspace = true [dev-dependencies] +common-time.workspace = true criterion.workspace = true rand = "0.8" diff --git a/src/partition/src/error.rs b/src/partition/src/error.rs index 065ff9d5250..83b9b767346 100644 --- a/src/partition/src/error.rs +++ b/src/partition/src/error.rs @@ -206,6 +206,14 @@ pub enum Error { location: Location, }, + #[snafu(display("Failed to coerce partition expr against the schema"))] + CoercePartitionExpr { + #[snafu(source)] + error: datafusion_common::error::DataFusionError, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Partition expr value is not supported: {:?}", value))] UnsupportedPartitionExprValue { value: Value, @@ -275,6 +283,7 @@ impl ErrorExt for Error { Error::UnexpectedColumnType { .. } => StatusCode::Internal, Error::ToDFSchema { .. } => StatusCode::Internal, Error::CreatePhysicalExpr { .. } => StatusCode::Internal, + Error::CoercePartitionExpr { .. } => StatusCode::Internal, Error::UnsupportedPartitionExprValue { .. } => StatusCode::InvalidArguments, Error::ConvertPartitionExprValue { .. } => StatusCode::InvalidArguments, } diff --git a/src/partition/src/expr.rs b/src/partition/src/expr.rs index 560dc517060..03e7249797b 100644 --- a/src/partition/src/expr.rs +++ b/src/partition/src/expr.rs @@ -21,6 +21,8 @@ use datafusion_common::{ScalarValue, ToDFSchema}; use datafusion_expr::Expr; use datafusion_expr::execution_props::ExecutionProps; use datafusion_expr::physical_planning_context::PhysicalPlanningContext; +use datafusion_expr::simplify::SimplifyContext; +use datafusion_optimizer::simplify_expressions::ExprSimplifier; use datafusion_physical_expr::{PhysicalExpr, create_physical_expr}; use datatypes::arrow; use datatypes::value::{ @@ -357,8 +359,21 @@ impl PartitionExpr { .clone() .to_dfschema_ref() .context(error::ToDFSchemaSnafu)?; - let execution_props = &ExecutionProps::default(); let expr = self.try_as_logical_expr()?; + // Persisted partition bounds keep the unit (and type) they were created + // with, so after a schema change like widening the time index unit the + // expression may mix types with the columns it is evaluated against. + // Coerce it against the schema so comparisons share a common type, the + // same way `PredicateGroup` does for region partition filters. + let simplifier = ExprSimplifier::new( + SimplifyContext::builder() + .with_schema(df_schema.clone()) + .build(), + ); + let expr = simplifier + .coerce(expr, &df_schema) + .context(error::CoercePartitionExprSnafu)?; + let execution_props = &ExecutionProps::default(); create_physical_expr( &expr, &df_schema, diff --git a/src/partition/src/multi_dim.rs b/src/partition/src/multi_dim.rs index 29e45126343..25dcc0fb89d 100644 --- a/src/partition/src/multi_dim.rs +++ b/src/partition/src/multi_dim.rs @@ -798,9 +798,10 @@ mod tests { mod test_split_record_batch { use std::sync::Arc; + use common_time::timestamp::Timestamp; use datafusion_common::ScalarValue; - use datatypes::arrow::array::{Int64Array, StringArray}; - use datatypes::arrow::datatypes::{DataType, Field, Schema}; + use datatypes::arrow::array::{Int64Array, StringArray, TimestampNanosecondArray}; + use datatypes::arrow::datatypes::{DataType, Field, Schema, TimeUnit}; use datatypes::arrow::record_batch::RecordBatch; use rand::Rng; @@ -857,6 +858,45 @@ mod test_split_record_batch { } } + #[test] + fn test_split_record_batch_widened_time_index_unit() { + // https://github.com/GreptimeTeam/greptimedb/issues/9747 + // The partition bounds keep the unit they were created with (millisecond) + // after the time index column is widened to nanosecond, so the batch + // schema and the bound values mix timestamp units. + let bound = Value::from(Timestamp::new_millisecond(1_000)); + let rule = MultiDimPartitionRule::try_new( + vec!["ts".to_string()], + vec![0, 1], + vec![col("ts").lt(bound.clone()), col("ts").gt_eq(bound)], + true, + ) + .unwrap(); + + let schema = Arc::new(Schema::new(vec![Field::new( + "ts", + DataType::Timestamp(TimeUnit::Nanosecond, None), + false, + )])); + let batch = RecordBatch::try_new( + schema, + vec![Arc::new(TimestampNanosecondArray::from(vec![ + 500_000_000, + 2_000_000_000, + ]))], + ) + .unwrap(); + + // 500_000_000ns < 1000ms and 2_000_000_000ns >= 1000ms: one row per region. + let expected = rule.split_record_batch_naive(&batch).unwrap(); + assert_eq!(expected.len(), 2); + let result = rule.split_record_batch(&batch).unwrap(); + assert_eq!(result.len(), expected.len()); + for (region, value) in &result { + assert_eq!(value.array(), expected.get(region).unwrap()); + } + } + #[test] fn test_split_record_batch_empty() { // Create a simple MultiDimPartitionRule diff --git a/tests-integration/src/grpc/flight.rs b/tests-integration/src/grpc/flight.rs index e74d4d5bf70..80030174424 100644 --- a/tests-integration/src/grpc/flight.rs +++ b/tests-integration/src/grpc/flight.rs @@ -52,7 +52,9 @@ mod test { }; use datatypes::prelude::{ConcreteDataType, ScalarVector, VectorRef}; use datatypes::schema::{ColumnSchema, Schema}; - use datatypes::vectors::{Int32Vector, StringVector, TimestampMillisecondVector}; + use datatypes::vectors::{ + Int32Vector, StringVector, TimestampMillisecondVector, TimestampNanosecondVector, + }; use futures_util::{Stream, StreamExt}; use hyper_util::rt::TokioIo; use itertools::Itertools; @@ -837,6 +839,82 @@ mod test { ); } + #[tokio::test(flavor = "multi_thread")] + async fn test_flight_bulk_widen_partitioned_time_index() { + // https://github.com/GreptimeTeam/greptimedb/issues/9747 + // Bulk writes route through `MultiDimPartitionRule::split_record_batch`, + // which evaluates the persisted partition bounds against the batch. The + // bounds keep the unit they were created with, so widening the time + // index unit must not mix timestamp units during that evaluation. + common_telemetry::init_default_ut_logging(); + + let (db, server) = setup_grpc_server( + StorageType::File, + "test_flight_bulk_widen_partitioned_time_index", + ) + .await; + let addr = server.bind_addr().unwrap().to_string(); + + let client = Client::with_urls(vec![addr]); + let client = Database::new_with_dbname("greptime-public", client); + + client + .sql( + "CREATE TABLE widen_ts (host STRING, ts TIMESTAMP TIME INDEX) \ + PARTITION ON COLUMNS (ts) (\ + ts < '2024-01-02 00:00:00', \ + ts >= '2024-01-02 00:00:00'\ + )", + ) + .await + .unwrap(); + client + .sql("ALTER TABLE widen_ts MODIFY COLUMN ts TIMESTAMP_NS") + .await + .unwrap(); + + // 2024-01-02 00:00:00 UTC == 1704153600000 ms; rows below and above the + // bound, expressed in the widened nanosecond unit ("c" only exists in ns). + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new("host", ConcreteDataType::string_datatype(), true), + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_nanosecond_datatype(), + false, + ) + .with_time_index(true), + ])); + let batch = RecordBatch::new( + schema, + vec![ + Arc::new(StringVector::from_vec(vec![ + "a".to_string(), + "b".to_string(), + "c".to_string(), + ])) as VectorRef, + Arc::new(TimestampNanosecondVector::from_vec(vec![ + 1_704_153_599_500_000_000, + 1_704_153_600_000_000_000, + 1_704_153_600_000_000_001, + ])) as VectorRef, + ], + ) + .unwrap(); + + test_put_record_batches_to_table(&client, "widen_ts", vec![batch]).await; + + let sql = "select host, ts from widen_ts order by ts"; + let expected = "\ ++------+-------------------------------+ +| host | ts | ++------+-------------------------------+ +| a | 2024-01-01T23:59:59.500 | +| b | 2024-01-02T00:00:00 | +| c | 2024-01-02T00:00:00.000000001 | ++------+-------------------------------+"; + query_and_expect(db.frontend().as_ref(), sql, expected).await; + } + #[tokio::test(flavor = "multi_thread")] async fn test_standalone_flight_do_put_missing_nullable_columns() { common_telemetry::init_default_ut_logging(); @@ -1124,15 +1202,36 @@ mod test { test_put_record_batches_with_hints(client, record_batches, &[]).await; } + async fn test_put_record_batches_to_table( + client: &Database, + table: &str, + record_batches: Vec, + ) { + test_put_record_batches_with_hints_and_table(client, table, record_batches, &[]).await; + } + async fn test_put_record_batches_with_hints( client: &Database, record_batches: Vec, hints: &[(&str, &str)], ) { + test_put_record_batches_with_hints_and_table(client, "foo", record_batches, hints).await + } + + async fn test_put_record_batches_with_hints_and_table( + client: &Database, + table: &str, + record_batches: Vec, + hints: &[(&str, &str)], + ) { + let table = table.to_string(); let requests_count = record_batches.len(); let schema = record_batches[0].schema.arrow_schema().clone(); - let stream = futures::stream::once(async move { encode_put_schema(schema.as_ref()) }) + let stream = + futures::stream::once( + async move { encode_put_schema_for_table(schema.as_ref(), &table) }, + ) .chain( tokio_stream::iter(record_batches) .enumerate() @@ -1169,11 +1268,18 @@ mod test { } fn encode_put_schema(schema: &datatypes::arrow::datatypes::Schema) -> FlightData { + encode_put_schema_for_table(schema, "foo") + } + + fn encode_put_schema_for_table( + schema: &datatypes::arrow::datatypes::Schema, + table: &str, + ) -> FlightData { let mut schema_data = FlightEncoder::default().encode_schema(schema); schema_data.app_metadata = serde_json::to_vec(&DoPutMetadata::new(0)).unwrap().into(); schema_data.flight_descriptor = Some(FlightDescriptor { r#type: arrow_flight::flight_descriptor::DescriptorType::Path as i32, - path: vec!["foo".to_string()], + path: vec![table.to_string()], ..Default::default() }); schema_data diff --git a/tests/cases/standalone/common/alter/change_col_type.result b/tests/cases/standalone/common/alter/change_col_type.result index d362215792b..92a986cd73b 100644 --- a/tests/cases/standalone/common/alter/change_col_type.result +++ b/tests/cases/standalone/common/alter/change_col_type.result @@ -341,6 +341,41 @@ SELECT * FROM ts_widen_part ORDER BY ts; | b | 2.0 | 2024-01-02T00:00:00.000250 | +------+-----+----------------------------+ +-- force the rows into SSTs so scans read old-unit data through the widened +-- schema of a time-index-partitioned table (https://github.com/GreptimeTeam/greptimedb/issues/9747) +ADMIN FLUSH_TABLE('ts_widen_part'); + ++------------------------------------+ +| ADMIN FLUSH_TABLE('ts_widen_part') | ++------------------------------------+ +| 0 | ++------------------------------------+ + +SELECT * FROM ts_widen_part ORDER BY ts; + ++------+-----+----------------------------+ +| host | val | ts | ++------+-----+----------------------------+ +| a | 1.0 | 2024-01-01T23:59:59.500 | +| b | 2.0 | 2024-01-02T00:00:00.000250 | ++------+-----+----------------------------+ + +SELECT * FROM ts_widen_part WHERE ts >= '2024-01-02 00:00:00' ORDER BY ts; + ++------+-----+----------------------------+ +| host | val | ts | ++------+-----+----------------------------+ +| b | 2.0 | 2024-01-02T00:00:00.000250 | ++------+-----+----------------------------+ + +SELECT * FROM ts_widen_part WHERE ts < '2024-01-02 00:00:00' ORDER BY ts; + ++------+-----+-------------------------+ +| host | val | ts | ++------+-----+-------------------------+ +| a | 1.0 | 2024-01-01T23:59:59.500 | ++------+-----+-------------------------+ + DROP TABLE ts_widen_part; Affected Rows: 0 diff --git a/tests/cases/standalone/common/alter/change_col_type.sql b/tests/cases/standalone/common/alter/change_col_type.sql index c7d206eaf7c..2da3094a6f1 100644 --- a/tests/cases/standalone/common/alter/change_col_type.sql +++ b/tests/cases/standalone/common/alter/change_col_type.sql @@ -119,6 +119,16 @@ SELECT * FROM ts_widen_part ORDER BY ts; INSERT INTO ts_widen_part VALUES ("b", 2, "2024-01-02 00:00:00.000250"); SELECT * FROM ts_widen_part ORDER BY ts; +-- force the rows into SSTs so scans read old-unit data through the widened +-- schema of a time-index-partitioned table (https://github.com/GreptimeTeam/greptimedb/issues/9747) +ADMIN FLUSH_TABLE('ts_widen_part'); + +SELECT * FROM ts_widen_part ORDER BY ts; + +SELECT * FROM ts_widen_part WHERE ts >= '2024-01-02 00:00:00' ORDER BY ts; + +SELECT * FROM ts_widen_part WHERE ts < '2024-01-02 00:00:00' ORDER BY ts; + DROP TABLE ts_widen_part; CREATE TABLE ts_overflow (host STRING, ts TIMESTAMP TIME INDEX); From 6f53d57d15655e5522cd7edfcfed43ad9df7c00b Mon Sep 17 00:00:00 2001 From: Aum Date: Tue, 20 Oct 2026 11:25:25 +0000 Subject: [PATCH 054/239] fix: avoid exact scan row counts for SSTs with unknown counts (#9759) fix: avoid exact scan row counts for legacy SSTs Treat selected SSTs with unknown row counts as non-exact before checking region partition visibility. Preserve exact statistics for known counts and cover legacy metadata mixed with known-count files. Fixes GreptimeTeam/greptimedb#9718 Signed-off-by: Aumnertic --- src/mito2/src/read/scan_region.rs | 43 +++++++++++++++++++++++++++++-- 1 file changed, 41 insertions(+), 2 deletions(-) diff --git a/src/mito2/src/read/scan_region.rs b/src/mito2/src/read/scan_region.rs index bcd79cc6960..da4e1c82bf9 100644 --- a/src/mito2/src/read/scan_region.rs +++ b/src/mito2/src/read/scan_region.rs @@ -1703,9 +1703,14 @@ impl ScanInput { ); } - /// Whether physical source row counts equal the rows visible to this region, - /// before applying query predicates. + /// Whether physical source row counts are known and equal the rows visible + /// to this region, before applying query predicates. pub(crate) fn total_rows_is_exact(&self) -> bool { + // Legacy SST metadata uses zero for unknown row counts. + if self.files.iter().any(|file| file.num_rows() == 0) { + return false; + } + if self.region_partition_expr.is_none() { return true; } @@ -2542,6 +2547,39 @@ mod tests { ); } + #[tokio::test] + async fn test_total_rows_is_exact_with_legacy_file() { + let file_meta = FileMeta { + num_rows: 10, + ..Default::default() + }; + let mut json = serde_json::to_value(&file_meta).unwrap(); + json.as_object_mut().unwrap().remove("num_rows"); + let legacy_meta = serde_json::from_value::(json).unwrap(); + assert_eq!(legacy_meta.num_rows, 0); + + let metadata = Arc::new(metadata_with_primary_key(vec![0, 1], false)); + for (file_metas, expected_rows, exact) in [ + (vec![file_meta.clone()], 10, true), + (vec![legacy_meta.clone()], 0, false), + (vec![file_meta, legacy_meta], 10, false), + ] { + let files = file_metas + .into_iter() + .map(|meta| { + FileHandle::new(meta, Arc::new(crate::sst::file_purger::NoopFilePurger)) + }) + .collect(); + let input = new_scan_input(metadata.clone(), vec![]) + .await + .with_append_mode(true) + .with_files(files) + .build(); + assert_eq!(input.total_rows(), expected_rows); + assert_eq!(input.total_rows_is_exact(), exact); + } + } + #[tokio::test] async fn test_total_rows_is_exact_after_partition_filter() { let expr = partition_col("k0").gt_eq(Value::String("foo".into())); @@ -2564,6 +2602,7 @@ mod tests { .map(|partition_expr| { FileHandle::new( FileMeta { + num_rows: 1, partition_expr, ..Default::default() }, From 9a2ac7f9de29efded3b8a74200f7b04340c1c59e Mon Sep 17 00:00:00 2001 From: LFC <990479+MichaelScofield@users.noreply.github.com> Date: Tue, 20 Oct 2026 11:57:26 +0000 Subject: [PATCH 055/239] feat: support dot access for JSON columns (#9775) Signed-off-by: luofucong --- src/query/src/datafusion/json_expr_planner.rs | 135 ++++++++++-------- .../common/types/json/json_dot_access.result | 97 +++++++++++++ .../common/types/json/json_dot_access.sql | 32 +++++ 3 files changed, 204 insertions(+), 60 deletions(-) create mode 100644 tests/cases/standalone/common/types/json/json_dot_access.result create mode 100644 tests/cases/standalone/common/types/json/json_dot_access.sql diff --git a/src/query/src/datafusion/json_expr_planner.rs b/src/query/src/datafusion/json_expr_planner.rs index 4837403c73f..a60e6155445 100644 --- a/src/query/src/datafusion/json_expr_planner.rs +++ b/src/query/src/datafusion/json_expr_planner.rs @@ -33,7 +33,8 @@ use datafusion_expr::{ Expr, ExprSchemable, GetFieldAccess, Operator, ScalarUDF, WindowFunctionDefinition, }; use datatypes::extension::json::{ - Json2ExtensionType, is_json2_extension_type, parse_legacy_json2_settings, + Json2ExtensionType, is_any_json_extension_type, is_json2_extension_type, + parse_legacy_json2_settings, }; use datatypes::types::json_type::JsonNativeType; use jsonb::jsonpath::Path; @@ -138,7 +139,13 @@ impl ExprPlanner for JsonExprPlanner { qualifier: Option<&TableReference>, nested_names: &[String], ) -> Result>> { - if !is_json2_extension_type(field) { + if !is_any_json_extension_type(field) + // SQL-created JSONB columns use Binary with greptime:type=Json, without an Arrow extension. + && field + .metadata() + .get(datatypes::schema::TYPE_KEY) + .is_none_or(|value| value != "Json") + { return Ok(PlannerResult::Original(Vec::new())); } @@ -514,9 +521,10 @@ mod tests { use datafusion_expr::WindowFrame; use datafusion_functions::core::coalesce; use datafusion_functions::math::{abs, power}; - use datatypes::extension::json::{Json2ExtensionType, JsonMetadata}; + use datatypes::extension::json::{Json2ExtensionType, JsonExtensionType, JsonMetadata}; use datatypes::json::{JsonSettings, JsonTypeHint}; use datatypes::prelude::ConcreteDataType; + use std::collections::HashMap; use super::*; @@ -654,78 +662,85 @@ mod tests { let qualifier = TableReference::bare("events"); let nested_names = vec!["payload".to_string(), "cpu".to_string()]; - let planned = planner.plan_compound_identifier( - &Field::new("labels", DataType::Struct(Fields::empty()), true) + for field in [ + Field::new("labels", DataType::Struct(Fields::empty()), true) .with_extension_type(Json2ExtensionType::default()), - Some(&qualifier), - &nested_names, - )?; + Field::new("labels", DataType::Binary, true).with_extension_type(JsonExtensionType), + Field::new("labels", DataType::Binary, true).with_metadata(HashMap::from([( + datatypes::schema::TYPE_KEY.to_string(), + "Json".to_string(), + )])), + ] { + let planned = + planner.plan_compound_identifier(&field, Some(&qualifier), &nested_names)?; - match planned { - PlannerResult::Planned(Expr::ScalarFunction(func)) => { - assert_eq!(func.func.name(), JsonGetWithType::NAME); - assert_eq!(func.args.len(), 2); - assert_eq!( - func.args[0], - Expr::Column(Column::new(Some(qualifier.clone()), "labels")) - ); + match planned { + PlannerResult::Planned(Expr::ScalarFunction(func)) => { + assert_eq!(func.func.name(), JsonGetWithType::NAME); + assert_eq!(func.args.len(), 2); + assert_eq!( + func.args[0], + Expr::Column(Column::new(Some(qualifier.clone()), "labels")) + ); + assert_eq!( + func.args[1], + Expr::Literal(ScalarValue::Utf8(Some("$.payload.cpu".to_string())), None) + ); + } + other => panic!("expected json_get scalar function, got {other:?}"), + } + + for (key, path) in [ + ("http.status_code", r#"$."http.status_code""#), + ("a\"b", r#"$."a\"b""#), + ("a\\b", r#"$."a\\b""#), + ] { + let planned = planner.plan_compound_identifier( + &field, + Some(&qualifier), + &[key.to_string()], + )?; + let PlannerResult::Planned(Expr::ScalarFunction(func)) = planned else { + unreachable!() + }; assert_eq!( func.args[1], - Expr::Literal(ScalarValue::Utf8(Some("$.payload.cpu".to_string())), None) + Expr::Literal(ScalarValue::Utf8(Some(path.to_string())), None) ); } - other => panic!("expected json_get scalar function, got {other:?}"), - } - for (key, path) in [ - ("http.status_code", r#"$."http.status_code""#), - ("a\"b", r#"$."a\"b""#), - ("a\\b", r#"$."a\\b""#), - ] { - let planned = planner.plan_compound_identifier( - &Field::new("labels", DataType::Struct(Fields::empty()), true) - .with_extension_type(Json2ExtensionType::default()), - Some(&qualifier), - &[key.to_string()], - )?; - let PlannerResult::Planned(Expr::ScalarFunction(func)) = planned else { + let PlannerResult::Planned(Expr::ScalarFunction(func)) = planner + .plan_compound_identifier( + &field, + Some(&qualifier), + &["resource".to_string(), "http.status_code".to_string()], + )? + else { unreachable!() }; assert_eq!( func.args[1], - Expr::Literal(ScalarValue::Utf8(Some(path.to_string())), None) + Expr::Literal( + ScalarValue::Utf8(Some(r#"$.resource."http.status_code""#.to_string())), + None + ) ); } - let PlannerResult::Planned(Expr::ScalarFunction(func)) = planner.plan_compound_identifier( - &Field::new("labels", DataType::Struct(Fields::empty()), true) - .with_extension_type(Json2ExtensionType::default()), - Some(&qualifier), - &["resource".to_string(), "http.status_code".to_string()], - )? - else { - unreachable!() - }; - assert_eq!( - func.args[1], - Expr::Literal( - ScalarValue::Utf8(Some(r#"$.resource."http.status_code""#.to_string())), - None - ) - ); + for data_type in [DataType::Utf8, DataType::Binary] { + let original = planner.plan_compound_identifier( + &Field::new("plain", data_type, true), + Some(&qualifier), + &nested_names, + )?; - let original = planner.plan_compound_identifier( - &Field::new("plain", DataType::Utf8, true), - Some(&qualifier), - &nested_names, - )?; - - match original { - PlannerResult::Original(exprs) => assert!(exprs.is_empty()), - other => panic!( - "expected original empty result for non-json field, got {:?}", - other, - ), + match original { + PlannerResult::Original(exprs) => assert!(exprs.is_empty()), + other => panic!( + "expected original empty result for non-json field, got {:?}", + other, + ), + } } Ok(()) diff --git a/tests/cases/standalone/common/types/json/json_dot_access.result b/tests/cases/standalone/common/types/json/json_dot_access.result new file mode 100644 index 00000000000..fe14dde72eb --- /dev/null +++ b/tests/cases/standalone/common/types/json/json_dot_access.result @@ -0,0 +1,97 @@ +CREATE TABLE json_dot_access (ts TIMESTAMP TIME INDEX, j JSON); + +Affected Rows: 0 + +INSERT INTO json_dot_access VALUES + (1, '{"a":{"b":2},"items":[{"name":"first"},{"name":"second"}],"a.b":"literal","UPPER":"quoted","enabled":true}'), + (2, '{"a":{"b":4},"items":[],"enabled":false}'), + (3, '{"a":null}'), + (4, '{}'), + (5, NULL); + +Affected Rows: 5 + +SELECT ts, j.a.b AS nested, j."a.b" AS dotted_key, j."UPPER" AS quoted_key, + j.missing AS missing, j.items[1].name AS item +FROM json_dot_access ORDER BY ts; + ++-------------------------+--------+------------+------------+---------+--------+ +| ts | nested | dotted_key | quoted_key | missing | item | ++-------------------------+--------+------------+------------+---------+--------+ +| 1970-01-01T00:00:00.001 | 2 | literal | quoted | | second | +| 1970-01-01T00:00:00.002 | 4 | | | | | +| 1970-01-01T00:00:00.003 | | | | | | +| 1970-01-01T00:00:00.004 | | | | | | +| 1970-01-01T00:00:00.005 | | | | | | ++-------------------------+--------+------------+------------+---------+--------+ + +SELECT ts, j.a.b = json_get(j, '$.a.b') AS same_value, + j.a.b + 1 AS incremented, j.a.b::BIGINT AS cast_value, + abs(j.a.b) AS absolute +FROM json_dot_access ORDER BY ts; + ++-------------------------+------------+-------------+------------+----------+ +| ts | same_value | incremented | cast_value | absolute | ++-------------------------+------------+-------------+------------+----------+ +| 1970-01-01T00:00:00.001 | true | 3 | 2 | 2.0 | +| 1970-01-01T00:00:00.002 | true | 5 | 4 | 4.0 | +| 1970-01-01T00:00:00.003 | | | | | +| 1970-01-01T00:00:00.004 | | | | | +| 1970-01-01T00:00:00.005 | | | | | ++-------------------------+------------+-------------+------------+----------+ + +SELECT t.ts, t.j.a.b AS qualified FROM json_dot_access AS t +WHERE t.j.a.b > 2 AND t.j.enabled = false ORDER BY t.ts; + ++-------------------------+-----------+ +| ts | qualified | ++-------------------------+-----------+ +| 1970-01-01T00:00:00.002 | 4 | ++-------------------------+-----------+ + +SELECT sum(j.a.b) AS total FROM json_dot_access; + ++-------+ +| total | ++-------+ +| 6.0 | ++-------+ + +SELECT ts, payload.a.b AS nested +FROM (SELECT ts, j AS payload FROM json_dot_access) ORDER BY ts; + ++-------------------------+--------+ +| ts | nested | ++-------------------------+--------+ +| 1970-01-01T00:00:00.001 | 2 | +| 1970-01-01T00:00:00.002 | 4 | +| 1970-01-01T00:00:00.003 | | +| 1970-01-01T00:00:00.004 | | +| 1970-01-01T00:00:00.005 | | ++-------------------------+--------+ + +ADMIN FLUSH_TABLE('json_dot_access'); + ++--------------------------------------+ +| ADMIN FLUSH_TABLE('json_dot_access') | ++--------------------------------------+ +| 0 | ++--------------------------------------+ + +SELECT ts, j.a.b AS nested, j.items[1].name AS item +FROM json_dot_access ORDER BY ts; + ++-------------------------+--------+--------+ +| ts | nested | item | ++-------------------------+--------+--------+ +| 1970-01-01T00:00:00.001 | 2 | second | +| 1970-01-01T00:00:00.002 | 4 | | +| 1970-01-01T00:00:00.003 | | | +| 1970-01-01T00:00:00.004 | | | +| 1970-01-01T00:00:00.005 | | | ++-------------------------+--------+--------+ + +DROP TABLE json_dot_access; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/types/json/json_dot_access.sql b/tests/cases/standalone/common/types/json/json_dot_access.sql new file mode 100644 index 00000000000..b8cbc86dd23 --- /dev/null +++ b/tests/cases/standalone/common/types/json/json_dot_access.sql @@ -0,0 +1,32 @@ +CREATE TABLE json_dot_access (ts TIMESTAMP TIME INDEX, j JSON); + +INSERT INTO json_dot_access VALUES + (1, '{"a":{"b":2},"items":[{"name":"first"},{"name":"second"}],"a.b":"literal","UPPER":"quoted","enabled":true}'), + (2, '{"a":{"b":4},"items":[],"enabled":false}'), + (3, '{"a":null}'), + (4, '{}'), + (5, NULL); + +SELECT ts, j.a.b AS nested, j."a.b" AS dotted_key, j."UPPER" AS quoted_key, + j.missing AS missing, j.items[1].name AS item +FROM json_dot_access ORDER BY ts; + +SELECT ts, j.a.b = json_get(j, '$.a.b') AS same_value, + j.a.b + 1 AS incremented, j.a.b::BIGINT AS cast_value, + abs(j.a.b) AS absolute +FROM json_dot_access ORDER BY ts; + +SELECT t.ts, t.j.a.b AS qualified FROM json_dot_access AS t +WHERE t.j.a.b > 2 AND t.j.enabled = false ORDER BY t.ts; + +SELECT sum(j.a.b) AS total FROM json_dot_access; + +SELECT ts, payload.a.b AS nested +FROM (SELECT ts, j AS payload FROM json_dot_access) ORDER BY ts; + +ADMIN FLUSH_TABLE('json_dot_access'); + +SELECT ts, j.a.b AS nested, j.items[1].name AS item +FROM json_dot_access ORDER BY ts; + +DROP TABLE json_dot_access; From a1b2362e5870f6ff52c29b17472232a3b0e3bb8e Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Wed, 21 Oct 2026 03:52:34 +0000 Subject: [PATCH 056/239] fix: avoid duplicate keys in update flow metadata txn (#9816) CREATE OR REPLACE FLOW deleted and then put the same flow route, flownode flow and table flow keys in a single txn when the replacement kept the partitions, flownodes or source tables, which etcd rejects with "duplicate key given in txn request". Only delete the old keys that are not put again. Closes #9810. Signed-off-by: Lei, HUANG --- src/common/meta/src/key/flow.rs | 181 +++++++++++++++++- src/common/meta/src/key/flow/flow_route.rs | 27 +-- src/common/meta/src/key/flow/flownode_flow.rs | 24 +-- src/common/meta/src/key/flow/table_flow.rs | 33 ++-- 4 files changed, 223 insertions(+), 42 deletions(-) diff --git a/src/common/meta/src/key/flow.rs b/src/common/meta/src/key/flow.rs index 50c7e0ac38a..5ef0bbf39fb 100644 --- a/src/common/meta/src/key/flow.rs +++ b/src/common/meta/src/key/flow.rs @@ -19,7 +19,7 @@ pub mod flow_state; mod flownode_addr_helper; pub(crate) mod flownode_flow; pub(crate) mod table_flow; -use std::collections::BTreeMap; +use std::collections::{BTreeMap, HashSet}; use std::ops::Deref; use std::sync::Arc; @@ -42,7 +42,7 @@ pub use crate::key::flow::table_flow::{TableFlowManager, TableFlowManagerRef}; use crate::key::txn_helper::TxnOpGetResponseSet; use crate::key::{DeserializedValueWithBytes, FlowId, FlowPartitionId, MetadataKey}; use crate::kv_backend::KvBackendRef; -use crate::kv_backend::txn::Txn; +use crate::kv_backend::txn::{Txn, TxnOp}; use crate::rpc::store::BatchDeleteRequest; pub const FLOW_KEY_PREFIX: &str = "__flow"; @@ -134,6 +134,31 @@ impl<'a, T: MetadataKey<'a, T>> MetadataKey<'a, FlowScoped> for FlowScoped } } +/// Builds txn ops that replace the `old_keys` with the `new_kvs`. +/// +/// Only deletes the old keys absent from `new_kvs`: etcd rejects a txn that +/// writes the same key more than once ("duplicate key given in txn request"). +fn build_replace_ops( + old_keys: impl IntoIterator>, + new_kvs: Vec<(Vec, Vec)>, +) -> Vec { + let new_keys = new_kvs + .iter() + .map(|(key, _)| key.as_slice()) + .collect::>(); + let mut ops = old_keys + .into_iter() + .filter(|key| !new_keys.contains(key.as_slice())) + .map(TxnOp::Delete) + .collect::>(); + ops.extend( + new_kvs + .into_iter() + .map(|(key, value)| TxnOp::Put(key, value)), + ); + ops +} + pub type FlowMetadataManagerRef = Arc; /// The manager of metadata, provides ability to: @@ -456,6 +481,7 @@ mod tests { use futures::TryStreamExt; use table::metadata::TableId; use table::table_name::TableName; + use uuid::Uuid; use super::*; use crate::FlownodeId; @@ -463,8 +489,10 @@ mod tests { use crate::key::flow::table_flow::TableFlowKey; use crate::key::node_address::{NodeAddressKey, NodeAddressValue}; use crate::key::{FlowPartitionId, MetadataValue}; - use crate::kv_backend::KvBackend; + use crate::kv_backend::chroot::ChrootKvBackend; + use crate::kv_backend::etcd::EtcdStore; use crate::kv_backend::memory::MemoryKvBackend; + use crate::kv_backend::{KvBackend, KvBackendRef}; use crate::peer::Peer; use crate::rpc::store::PutRequest; @@ -1189,4 +1217,151 @@ mod tests { err ); } + + #[test] + fn test_build_replace_ops() { + let ops = build_replace_ops( + vec![b"a".to_vec(), b"b".to_vec()], + vec![ + (b"b".to_vec(), b"1".to_vec()), + (b"c".to_vec(), b"2".to_vec()), + ], + ); + assert_eq!( + ops, + vec![ + TxnOp::Delete(b"a".to_vec()), + TxnOp::Put(b"b".to_vec(), b"1".to_vec()), + TxnOp::Put(b"c".to_vec(), b"2".to_vec()), + ] + ); + } + + /// Returns an etcd backend isolated by a chroot if `GT_ETCD_ENDPOINTS` is set, + /// because only etcd rejects txns that write the same key more than once. + async fn build_kv_backend() -> KvBackendRef { + let endpoints = std::env::var("GT_ETCD_ENDPOINTS").unwrap_or_default(); + if endpoints.is_empty() { + return Arc::new(MemoryKvBackend::default()); + } + let endpoints = endpoints.split(',').collect::>(); + let backend = EtcdStore::with_endpoints(endpoints, 128).await.unwrap(); + let chroot = format!("test_update_flow_metadata/{}/", Uuid::new_v4()); + Arc::new(ChrootKvBackend::new(chroot.into(), backend)) + } + + async fn assert_flow_metadata( + manager: &FlowMetadataManager, + flow_id: FlowId, + flow_info: &FlowInfoValue, + ) { + let got = manager.flow_info_manager().get(flow_id).await.unwrap(); + assert_eq!(got.as_ref(), Some(flow_info)); + + let routes = manager.flow_route_manager().routes(flow_id).await.unwrap(); + let expected_routes = flow_info + .flownode_ids() + .iter() + .map(|(partition_id, flownode_id)| { + ( + FlowRouteKey::new(flow_id, *partition_id), + FlowRouteValue { + peer: Peer::empty(*flownode_id), + }, + ) + }) + .collect::>(); + assert_eq!(routes, expected_routes); + + for flownode_id in 1..=4 { + let flows = manager + .flownode_flow_manager() + .flows(flownode_id) + .try_collect::>() + .await + .unwrap(); + let expected = flow_info + .flownode_ids() + .iter() + .filter(|(_, id)| **id == flownode_id) + .map(|(partition_id, _)| (flow_id, *partition_id)) + .collect::>(); + assert_eq!(flows, expected, "flownode {flownode_id}"); + } + + for table_id in 1024..=1027 { + let flows = manager + .table_flow_manager() + .flows(table_id) + .await + .unwrap() + .into_iter() + .map(|(key, _)| (key.flownode_id(), key.partition_id())) + .collect::>(); + let expected = if flow_info.source_table_ids().contains(&table_id) { + flow_info + .flownode_ids() + .iter() + .map(|(partition_id, flownode_id)| (*flownode_id, *partition_id)) + .collect() + } else { + vec![] + }; + assert_eq!(flows, expected, "table {table_id}"); + } + } + + #[tokio::test] + async fn test_update_flow_metadata_with_overlapping_keys() { + let flow_metadata_manager = FlowMetadataManager::new(build_kv_backend().await); + let flow_id = 10; + let to_routes = |flow_info: &FlowInfoValue| { + flow_info + .flownode_ids() + .iter() + .map(|(partition_id, flownode_id)| { + ( + *partition_id, + FlowRouteValue { + peer: Peer::empty(*flownode_id), + }, + ) + }) + .collect::>() + }; + + let mut current = + test_flow_info_value("flow", [(0, 1u64), (1, 2u64)].into(), vec![1024, 1025]); + flow_metadata_manager + .create_flow_metadata(flow_id, current.clone(), to_routes(¤t)) + .await + .unwrap(); + assert_flow_metadata(&flow_metadata_manager, flow_id, ¤t).await; + + let updates = [ + // Unchanged partitions, flownodes and source tables. + ([(0, 1u64), (1, 2u64)].into(), vec![1024, 1025]), + // Partially overlapping flownodes and source tables. + ([(0, 1u64), (1, 3u64)].into(), vec![1025, 1026]), + // Fully changed partitions, flownodes and source tables. + ([(2, 4u64)].into(), vec![1027]), + ]; + for (i, (flownode_ids, source_table_ids)) in updates.into_iter().enumerate() { + let mut new_flow_info = current.clone(); + new_flow_info.raw_sql = format!("update {i}"); + new_flow_info.flownode_ids = flownode_ids; + new_flow_info.source_table_ids = source_table_ids; + flow_metadata_manager + .update_flow_metadata( + flow_id, + &DeserializedValueWithBytes::from_inner(current.clone()), + &new_flow_info, + to_routes(&new_flow_info), + ) + .await + .unwrap(); + assert_flow_metadata(&flow_metadata_manager, flow_id, &new_flow_info).await; + current = new_flow_info; + } + } } diff --git a/src/common/meta/src/key/flow/flow_route.rs b/src/common/meta/src/key/flow/flow_route.rs index d402e1acef7..0cd62a3b909 100644 --- a/src/common/meta/src/key/flow/flow_route.rs +++ b/src/common/meta/src/key/flow/flow_route.rs @@ -20,7 +20,7 @@ use snafu::OptionExt; use crate::error::{self, Result}; use crate::key::flow::flow_info::FlowInfoValue; -use crate::key::flow::{FlowScoped, flownode_addr_helper}; +use crate::key::flow::{FlowScoped, build_replace_ops, flownode_addr_helper}; use crate::key::node_address::NodeAddressKey; use crate::key::{ BytesAdapter, DeserializedValueWithBytes, FlowId, FlowPartitionId, MetadataKey, MetadataValue, @@ -216,25 +216,26 @@ impl FlowRouteManager { /// Builds a update flow routes transaction. /// /// Puts `__flow/route/{flow_id}/{partition_id}` keys. - /// Also removes `__flow/route/{flow_id}/{old_partition_id}` keys. + /// Also removes the `__flow/route/{flow_id}/{old_partition_id}` keys that are not put. pub(crate) fn build_update_txn>( &self, flow_id: FlowId, current_flow_info: &DeserializedValueWithBytes, flow_routes: I, ) -> Result { - let del_txns = current_flow_info.flownode_ids().keys().map(|partition_id| { - let key = FlowRouteKey::new(flow_id, *partition_id).to_bytes(); - Ok(TxnOp::Delete(key)) - }); + let old_keys = current_flow_info + .flownode_ids() + .keys() + .map(|partition_id| FlowRouteKey::new(flow_id, *partition_id).to_bytes()); + let new_kvs = flow_routes + .into_iter() + .map(|(partition_id, route)| { + let key = FlowRouteKey::new(flow_id, partition_id).to_bytes(); + Ok((key, route.try_as_raw_value()?)) + }) + .collect::>>()?; - let put_txns = flow_routes.into_iter().map(|(partition_id, route)| { - let key = FlowRouteKey::new(flow_id, partition_id).to_bytes(); - - Ok(TxnOp::Put(key, route.try_as_raw_value()?)) - }); - let txns = del_txns.chain(put_txns).collect::>>()?; - Ok(Txn::new().and_then(txns)) + Ok(Txn::new().and_then(build_replace_ops(old_keys, new_kvs))) } async fn remap_flow_route_addresses( diff --git a/src/common/meta/src/key/flow/flownode_flow.rs b/src/common/meta/src/key/flow/flownode_flow.rs index 2ba690c167e..83d49229fac 100644 --- a/src/common/meta/src/key/flow/flownode_flow.rs +++ b/src/common/meta/src/key/flow/flownode_flow.rs @@ -20,8 +20,8 @@ use snafu::OptionExt; use crate::FlownodeId; use crate::error::{self, Result}; -use crate::key::flow::FlowScoped; use crate::key::flow::flow_info::FlowInfoValue; +use crate::key::flow::{FlowScoped, build_replace_ops}; use crate::key::{BytesAdapter, DeserializedValueWithBytes, FlowId, FlowPartitionId, MetadataKey}; use crate::kv_backend::KvBackendRef; use crate::kv_backend::txn::{Txn, TxnOp}; @@ -218,28 +218,30 @@ impl FlownodeFlowManager { /// Builds a update flownode flow transaction. /// /// Puts `__flownode_flow/{flownode_id}/{flow_id}/{partition_id}` keys. - /// Remove the old `__flownode_flow/{old_flownode_id}/{flow_id}/{old_partition_id}` keys. + /// Removes the old `__flownode_flow/{old_flownode_id}/{flow_id}/{old_partition_id}` + /// keys that are not put. pub(crate) fn build_update_txn>( &self, flow_id: FlowId, current_flow_info: &DeserializedValueWithBytes, flownode_ids: I, ) -> Txn { - let del_txns = + let old_keys = current_flow_info .flownode_ids() .iter() .map(|(partition_id, flownode_id)| { - let key = FlownodeFlowKey::new(*flownode_id, flow_id, *partition_id).to_bytes(); - TxnOp::Delete(key) + FlownodeFlowKey::new(*flownode_id, flow_id, *partition_id).to_bytes() }); - let put_txns = flownode_ids.into_iter().map(|(partition_id, flownode_id)| { - let key = FlownodeFlowKey::new(flownode_id, flow_id, partition_id).to_bytes(); - TxnOp::Put(key, vec![]) - }); - let txns = del_txns.chain(put_txns).collect::>(); + let new_kvs = flownode_ids + .into_iter() + .map(|(partition_id, flownode_id)| { + let key = FlownodeFlowKey::new(flownode_id, flow_id, partition_id).to_bytes(); + (key, vec![]) + }) + .collect(); - Txn::new().and_then(txns) + Txn::new().and_then(build_replace_ops(old_keys, new_kvs)) } } diff --git a/src/common/meta/src/key/flow/table_flow.rs b/src/common/meta/src/key/flow/table_flow.rs index d334ca02181..ae0fcb525f1 100644 --- a/src/common/meta/src/key/flow/table_flow.rs +++ b/src/common/meta/src/key/flow/table_flow.rs @@ -24,7 +24,7 @@ use table::metadata::TableId; use crate::FlownodeId; use crate::error::{self, Result}; use crate::key::flow::flow_info::FlowInfoValue; -use crate::key::flow::{FlowScoped, flownode_addr_helper}; +use crate::key::flow::{FlowScoped, build_replace_ops, flownode_addr_helper}; use crate::key::node_address::NodeAddressKey; use crate::key::{ BytesAdapter, DeserializedValueWithBytes, FlowId, FlowPartitionId, MetadataKey, MetadataValue, @@ -245,8 +245,9 @@ impl TableFlowManager { /// Builds a update table flow transaction. /// /// Puts `__flow/source_table/{table_id}/{node_id}/{flow_id}/{partition_id}` keys, - /// Also remove previous - /// `__flow/source_table/{table_id}/{old_node_id}/{flow_id}/{partition_id}` keys. + /// Also removes the previous + /// `__flow/source_table/{table_id}/{old_node_id}/{flow_id}/{partition_id}` keys + /// that are not put. pub fn build_update_txn( &self, flow_id: FlowId, @@ -254,22 +255,24 @@ impl TableFlowManager { table_flow_values: Vec<(FlowPartitionId, TableFlowValue)>, source_table_ids: &[TableId], ) -> Result { - let mut txns = Vec::with_capacity(2 * source_table_ids.len() * table_flow_values.len()); - - // first remove the old keys - for (part_id, node_id) in current_flow_info.flownode_ids() { - for source_table_id in current_flow_info.source_table_ids() { - txns.push(TxnOp::Delete( - TableFlowKey::new(*source_table_id, *node_id, flow_id, *part_id).to_bytes(), - )); - } - } + let old_keys = current_flow_info + .flownode_ids() + .iter() + .flat_map(|(part_id, node_id)| { + current_flow_info + .source_table_ids() + .iter() + .map(move |source_table_id| { + TableFlowKey::new(*source_table_id, *node_id, flow_id, *part_id).to_bytes() + }) + }); + let mut new_kvs = Vec::with_capacity(source_table_ids.len() * table_flow_values.len()); for (partition_id, table_flow_value) in table_flow_values { let flownode_id = table_flow_value.peer.id; let value = table_flow_value.try_as_raw_value()?; for source_table_id in source_table_ids { - txns.push(TxnOp::Put( + new_kvs.push(( TableFlowKey::new(*source_table_id, flownode_id, flow_id, partition_id) .to_bytes(), value.clone(), @@ -277,7 +280,7 @@ impl TableFlowManager { } } - Ok(Txn::new().and_then(txns)) + Ok(Txn::new().and_then(build_replace_ops(old_keys, new_kvs))) } async fn remap_table_flow_addresses( From f866ef2bd9853d990bcbb8623aa71713a82e84ff Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Wed, 21 Oct 2026 05:51:15 +0000 Subject: [PATCH 057/239] feat: return per-entry region write failures (#9728) * feat: return per-entry region write failures Signed-off-by: WenyXu * fix: preserve retry hints for single-entry region writes Signed-off-by: WenyXu * fix: preserve region write error compatibility Signed-off-by: WenyXu * refactor: simplify region partial result handling Signed-off-by: WenyXu * fix: align partial failure summaries with input order Signed-off-by: WenyXu * refactor: separate region write response handling Signed-off-by: WenyXu * test: fix partial result expectations and region map guard Signed-off-by: WenyXu * chore: limit lockfile changes to greptime proto Signed-off-by: WenyXu * fix: make partial region results an explicit client opt-in Signed-off-by: WenyXu * chore: pin greptime-proto to merged partial-result changes Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- Cargo.lock | 2 +- Cargo.toml | 2 +- src/api/src/region.rs | 29 +- src/client/src/error.rs | 56 +- src/client/src/lib.rs | 12 +- src/client/src/region.rs | 477 ++++++++- src/common/error/src/status_code.rs | 7 +- src/common/meta/src/error.rs | 13 + src/datanode/src/error.rs | 5 + src/datanode/src/region_server.rs | 1013 ++++++++++++++---- src/datanode/src/region_server/flow_write.rs | 1 + src/meta-srv/src/procedure/utils.rs | 2 + src/metric-engine/AGENTS.md | 1 + src/metric-engine/src/engine.rs | 17 +- src/metric-engine/src/engine/put.rs | 298 +++--- src/operator/src/write_dispatcher.rs | 5 +- src/operator/src/write_dispatcher/tests.rs | 19 +- src/servers/src/error.rs | 1 + src/servers/src/mysql/writer.rs | 4 +- src/servers/src/postgres/types/error.rs | 3 +- src/standalone/Cargo.toml | 2 + src/standalone/src/datanode_manager.rs | 69 +- src/store-api/src/region_request.rs | 74 +- tests/runner/src/client.rs | 1 + 24 files changed, 1647 insertions(+), 466 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 73bd260e9de..7a8da222c03 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6083,7 +6083,7 @@ dependencies = [ [[package]] name = "greptime-proto" version = "0.1.0" -source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=9533078d6fd4bdaddbca0750cb9bf7ad64b2e568#9533078d6fd4bdaddbca0750cb9bf7ad64b2e568" +source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=acb247415b8f4c2f0750783dab4e1f38c71f5860#acb247415b8f4c2f0750783dab4e1f38c71f5860" dependencies = [ "prost 0.14.1", "prost-types 0.14.1", diff --git a/Cargo.toml b/Cargo.toml index 3cb4e77fb77..20582d8eaa4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -162,7 +162,7 @@ fs2 = "0.4" fst = "0.4.7" futures = "0.3" futures-util = "0.3" -greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "9533078d6fd4bdaddbca0750cb9bf7ad64b2e568" } +greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "acb247415b8f4c2f0750783dab4e1f38c71f5860" } hdfs-native = "0.14.6" hex = "0.4" hostname = "0.4.0" diff --git a/src/api/src/region.rs b/src/api/src/region.rs index e7000cd7448..fb0d5b5e3f4 100644 --- a/src/api/src/region.rs +++ b/src/api/src/region.rs @@ -15,7 +15,19 @@ use std::collections::HashMap; use common_base::AffectedRows; -use greptime_proto::v1::region::RegionResponse as RegionResponseV1; +use greptime_proto::v1::region::{RegionRequestHeader, RegionResponse as RegionResponseV1}; + +/// Request extension advertising support for per-entry write failures. +pub const ACCEPT_PARTIAL_RESULT: &str = "accept_partial_result"; + +/// Advertises support for per-entry write failures in the request header. +pub fn set_accept_partial_result(header: &mut RegionRequestHeader) { + header + .query_context + .get_or_insert_with(Default::default) + .extensions + .insert(ACCEPT_PARTIAL_RESULT.to_string(), "true".to_string()); +} /// This result struct is derived from [RegionResponseV1] #[derive(Debug)] @@ -25,6 +37,21 @@ pub struct RegionResponse { pub metadata: Vec, } +/// Successful entries and failed-entry errors indexed by input order. +#[derive(Debug)] +pub struct BatchRegionResponse { + pub affected_rows: AffectedRows, + pub extensions: HashMap>, + pub failures: Vec<(usize, E)>, +} + +impl BatchRegionResponse { + /// Returns the failure with the smallest original request index. + pub fn first_failure(&self) -> Option<&(usize, E)> { + self.failures.iter().min_by_key(|(index, _)| *index) + } +} + impl RegionResponse { pub fn from_region_response(region_response: RegionResponseV1) -> Self { Self { diff --git a/src/client/src/error.rs b/src/client/src/error.rs index fdf0df9582f..d9fa8c96050 100644 --- a/src/client/src/error.rs +++ b/src/client/src/error.rs @@ -18,7 +18,7 @@ use common_error::define_from_tonic_status; use common_error::ext::{BoxedError, ErrorExt, RetryHint}; use common_error::status_code::StatusCode; use common_macro::stack_trace_debug; -use snafu::{Location, Snafu, location}; +use snafu::{Location, OptionExt, Snafu, location}; use tonic::Code; use tonic::metadata::errors::InvalidMetadataValue; @@ -26,6 +26,15 @@ use tonic::metadata::errors::InvalidMetadataValue; #[snafu(visibility(pub))] #[stack_trace_debug] pub enum Error { + #[snafu(display( + "Region request has {} failed entries; {}", + response.failures.len(), + response.first_failure().map(|(index, error)| format!("entry {index}: {}", error.output_msg())).unwrap_or_default() + ))] + PartialRegion { + response: Box>, + }, + #[snafu(display("Illegal Flight messages, reason: {}", reason))] IllegalFlightMessages { reason: String, @@ -109,11 +118,12 @@ pub enum Error { location: Location, }, - // Server error carried in Tonic Status's metadata. + // Server error carried in a protobuf response status. #[snafu(display("{}", msg))] Server { code: StatusCode, msg: String, + retry_hint: RetryHint, #[snafu(implicit)] location: Location, }, @@ -163,6 +173,9 @@ pub type Result = std::result::Result; impl ErrorExt for Error { fn status_code(&self) -> StatusCode { match self { + Error::PartialRegion { response } => response + .first_failure() + .map_or(StatusCode::Internal, |(_, error)| error.status_code()), Error::IllegalFlightMessages { .. } | Error::MissingField { .. } | Error::IllegalDatabaseResponse { .. } => StatusCode::Internal, @@ -188,7 +201,8 @@ impl ErrorExt for Error { fn retry_hint(&self) -> RetryHint { match self { - Error::Tonic { retry_hint, .. } => *retry_hint, + Error::PartialRegion { .. } => RetryHint::Retryable, + Error::Tonic { retry_hint, .. } | Error::Server { retry_hint, .. } => *retry_hint, Error::FlightGet { source, .. } | Error::FlightStream { source, .. } | Error::FlowServer { source, .. } @@ -220,9 +234,22 @@ define_from_tonic_status!(Error, Tonic); impl From for common_meta::error::Error { fn from(error: Error) -> Self { - Self::External { - source: BoxedError::new(error), - location: location!(), + match error { + Error::PartialRegion { response } => Self::PartialRegion { + response: Box::new(api::region::BatchRegionResponse { + affected_rows: response.affected_rows, + extensions: response.extensions, + failures: response + .failures + .into_iter() + .map(|(index, error)| (index, error.into())) + .collect(), + }), + }, + error => Self::External { + source: BoxedError::new(error), + location: location!(), + }, } } } @@ -258,3 +285,20 @@ impl Error { ) } } + +impl Error { + /// Decodes an application error carried in a protobuf response. + pub(crate) fn try_from_response_status(status: &api::v1::Status) -> Result { + let code = StatusCode::from_u32(status.status_code).with_context(|| { + IllegalDatabaseResponseSnafu { + err_msg: format!("Unknown server status: {}", status.status_code), + } + })?; + Ok(ServerSnafu { + code, + msg: status.err_msg.clone(), + retry_hint: status.retry_hint.parse().unwrap_or(RetryHint::NonRetryable), + } + .build()) + } +} diff --git a/src/client/src/lib.rs b/src/client/src/lib.rs index 4b9d543e35f..208a3233ff7 100644 --- a/src/client/src/lib.rs +++ b/src/client/src/lib.rs @@ -35,7 +35,7 @@ use snafu::OptionExt; pub use self::client::{Client, ClientOptions}; pub use self::database::{Database, OutputMetrics, OutputWithMetrics}; pub use self::error::{Error, Result}; -use crate::error::{IllegalDatabaseResponseSnafu, ServerSnafu}; +use crate::error::IllegalDatabaseResponseSnafu; pub fn from_grpc_response(response: GreptimeResponse) -> Result { let header = response.header.context(IllegalDatabaseResponseSnafu { @@ -53,14 +53,6 @@ pub fn from_grpc_response(response: GreptimeResponse) -> Result { Response::AffectedRows(AffectedRows { value }) => Ok(value), } } else { - let status_code = - StatusCode::from_u32(status.status_code).context(IllegalDatabaseResponseSnafu { - err_msg: format!("invalid status: {:?}", status), - })?; - ServerSnafu { - code: status_code, - msg: status.err_msg, - } - .fail() + Err(Error::try_from_response_status(&status)?) } } diff --git a/src/client/src/region.rs b/src/client/src/region.rs index f9ebbb0d0bf..2e517eb30d0 100644 --- a/src/client/src/region.rs +++ b/src/client/src/region.rs @@ -18,8 +18,9 @@ use std::borrow::Borrow; use std::sync::Arc; use std::time::Duration; -use api::region::RegionResponse; +use api::region::{BatchRegionResponse, RegionResponse}; use api::v1::ResponseHeader; +use api::v1::region::RegionResponse as PbRegionResponse; use api::v1::region::{ RegionRequest, RegionRequestHeader, RemoteDynFilterRequest, RemoteDynFilterUnregister, RemoteDynFilterUpdate, region_request, remote_dyn_filter_request, @@ -49,7 +50,7 @@ use tonic::codec::CompressionEncoding; use crate::error::{ self, FlightGetSnafu, IllegalDatabaseResponseSnafu, IllegalFlightMessagesSnafu, - MissingFieldSnafu, Result, ServerSnafu, + MissingFieldSnafu, Result, }; use crate::flight::{FlightMessageReader, decode_flight_data}; use crate::region::codec::RegionCodec; @@ -141,7 +142,8 @@ impl RegionRequester { /// /// Accepts both [`RegionRequest`] and [`Arc`]. Callers can retain a /// shared request for another attempt; this method does not retry or determine - /// whether replaying the request is safe. + /// whether replaying the request is safe. Callers accepting partial results in the + /// request header must inspect failed entries before retrying. pub async fn handle(&self, request: R) -> Result where R: Borrow + Send + Sync + 'static, @@ -180,6 +182,11 @@ impl RegionRequester { .body .as_ref() .context(MissingFieldSnafu { field: "body" })?; + let request_count = match request_body { + region_request::Body::Inserts(requests) => Some(requests.requests.len()), + region_request::Body::Deletes(requests) => Some(requests.requests.len()), + _ => None, + }; let is_insert = matches!( request_body, region_request::Body::Inserts(_) | region_request::Body::BulkInsert(_) @@ -199,6 +206,13 @@ impl RegionRequester { } } + let accept_partial_result = request + .borrow() + .header + .as_ref() + .and_then(|header| header.query_context.as_ref()) + .and_then(|context| context.extensions.get(api::region::ACCEPT_PARTIAL_RESULT)) + .is_some_and(|value| value == "true"); let response = async { client .ready() @@ -223,9 +237,10 @@ impl RegionRequester { .map_err(|status| map_region_error(status, addr))? .into_inner(); - check_response_header(&response.header)?; - - Ok(RegionResponse::from_region_response(response)) + match request_count { + Some(count) if accept_partial_result => decode_write_response(response, count), + _ => decode_region_response(response), + } } } @@ -385,6 +400,7 @@ fn build_remote_dyn_filter_request( } } +/// Checks that the response contains a successful application status. pub fn check_response_header(header: &Option) -> Result<()> { let status = header .as_ref() @@ -396,35 +412,107 @@ pub fn check_response_header(header: &Option) -> Result<()> { if StatusCode::is_success(status.status_code) { Ok(()) } else { - let code = - StatusCode::from_u32(status.status_code).context(IllegalDatabaseResponseSnafu { - err_msg: format!("unknown server status: {:?}", status), - })?; - ServerSnafu { - code, - msg: status.err_msg.clone(), - } - .fail() + Err(error::Error::try_from_response_status(status)?) } } +/// Decodes an ordinary region response without per-entry write failures. +pub fn decode_region_response(response: PbRegionResponse) -> Result { + if !response.failures.is_empty() { + return IllegalDatabaseResponseSnafu { + err_msg: "Failure details are only valid for inserts/deletes", + } + .fail(); + } + check_response_header(&response.header)?; + Ok(RegionResponse::from_region_response(response)) +} + +/// Decodes an inserts/deletes response, validating failures against the original entry count. +/// A failed single-entry request returns its underlying error. +fn decode_write_response(response: PbRegionResponse, entry_count: usize) -> Result { + if response.failures.is_empty() { + if response + .header + .as_ref() + .and_then(|header| header.status.as_ref()) + .is_some_and(|status| status.status_code == StatusCode::PartialFailure as u32) + { + return IllegalDatabaseResponseSnafu { + err_msg: "Partial failure has no failed entries", + } + .fail(); + } + return decode_region_response(response); + } + let status = response + .header + .as_ref() + .and_then(|header| header.status.as_ref()) + .context(IllegalDatabaseResponseSnafu { + err_msg: "Missing region response status", + })?; + if status.status_code != StatusCode::PartialFailure as u32 { + return IllegalDatabaseResponseSnafu { + err_msg: "Failure details require a partial-failure status", + } + .fail(); + } + let mut seen = std::collections::HashSet::with_capacity(response.failures.len()); + let mut failures = Vec::with_capacity(response.failures.len()); + for error in response.failures { + let index = error.request_index as usize; + if index >= entry_count || !seen.insert(index) { + return IllegalDatabaseResponseSnafu { + err_msg: "Invalid or duplicate batch error index", + } + .fail(); + } + let status = error.status.context(IllegalDatabaseResponseSnafu { + err_msg: "Missing batch error status", + })?; + if StatusCode::is_success(status.status_code) { + return IllegalDatabaseResponseSnafu { + err_msg: "Batch error contains a success status", + } + .fail(); + } + failures.push((index, error::Error::try_from_response_status(&status)?)); + } + // A single failed entry has no successful peers to preserve. + if entry_count == 1 + && let Some((_, error)) = failures.pop() + { + return Err(error); + } + Err(error::Error::PartialRegion { + response: Box::new(BatchRegionResponse { + affected_rows: response.affected_rows as _, + extensions: response.extensions, + failures, + }), + }) +} + #[cfg(test)] #[allow(deprecated)] mod tests { use std::collections::VecDeque; use std::sync::Mutex; + use api::v1::ResponseHeader; use api::v1::region::region_server::{Region, RegionServer}; use api::v1::region::{ - BulkInsertRequest, InsertRequest, InsertRequests, RegionResponse as PbRegionResponse, - RemoteDynFilterUnregister, RemoteDynFilterUpdate, region_request, - remote_dyn_filter_request, + BulkInsertRequest, InsertRequest, InsertRequests, RegionRequestFailure, + RegionResponse as PbRegionResponse, RemoteDynFilterUnregister, RemoteDynFilterUpdate, + region_request, remote_dyn_filter_request, }; use api::v1::value::ValueData; use api::v1::{ ColumnDataType, ColumnSchema as PbColumnSchema, Row, Rows, Status as PbStatus, Value, }; use common_error::ext::RetryHint; + use common_error::status_code::StatusCode; use common_grpc::channel_manager::{ChannelConfig, ChannelManager}; use common_recordbatch::adapter::RecordBatchMetrics; use datatypes::arrow::array::Int32Array; @@ -579,6 +667,20 @@ mod tests { } } + fn with_partial_result(mut request: RegionRequest) -> RegionRequest { + request + .header + .get_or_insert_with(Default::default) + .query_context + .get_or_insert_with(Default::default) + .extensions + .insert( + api::region::ACCEPT_PARTIAL_RESULT.to_string(), + "true".to_string(), + ); + request + } + #[tokio::test] async fn test_owned_and_shared_requests_preserve_payload_and_compression() { let server = RegionTestServer::start().await; @@ -636,6 +738,54 @@ mod tests { } } + #[tokio::test] + async fn test_partial_result_requires_explicit_acceptance() { + let server = RegionTestServer::start().await; + let requester = server.create_requester(ChannelConfig::default(), false, false); + for accepted in [false, true] { + let request = if accepted { + with_partial_result(build_insert_request()) + } else { + build_insert_request() + }; + let shared = Arc::new(request); + server + .service + .responses + .lock() + .unwrap() + .push_back(Ok(PbRegionResponse { + header: Some(ResponseHeader { + status: Some(PbStatus { + status_code: StatusCode::PartialFailure as u32, + ..Default::default() + }), + }), + failures: vec![RegionRequestFailure { + request_index: 0, + status: Some(PbStatus { + status_code: StatusCode::RegionNotFound as u32, + retry_hint: RetryHint::RetryAfterRouteRefresh.as_str().to_string(), + ..Default::default() + }), + }], + ..Default::default() + })); + let error = requester.handle(Arc::clone(&shared)).await.unwrap_err(); + if accepted { + assert_eq!(error.status_code(), StatusCode::RegionNotFound); + assert_eq!(error.retry_hint(), RetryHint::RetryAfterRouteRefresh); + } else { + assert!(matches!(error, Error::IllegalDatabaseResponse { .. })); + } + assert_eq!( + &server.service.requests.lock().unwrap().last().unwrap().0, + shared.as_ref() + ); + assert_eq!(Arc::strong_count(&shared), 1); + } + } + #[tokio::test] async fn test_shared_request_can_be_resent_after_transport_error() { let server = RegionTestServer::start().await; @@ -762,7 +912,10 @@ mod tests { assert_eq!(error.tonic_code(), baseline.tonic_code()); } let datanode: &dyn Datanode = &requester; - let error = datanode.handle(build_insert_request()).await.unwrap_err(); + let error = datanode + .handle_request(Arc::new(build_insert_request())) + .await + .unwrap_err(); assert_eq!(error.status_code(), StatusCode::RegionNotFound); assert_eq!(error.retry_hint(), hint); assert!(matches!(error, meta_error::Error::External { .. })); @@ -832,10 +985,16 @@ mod tests { body: Some(body), }; let shared = Arc::new(request); - let shared_error = datanode.handle(Arc::clone(&shared)).await.unwrap_err(); + let shared_error = datanode + .handle_request(Arc::clone(&shared)) + .await + .unwrap_err(); assert_eq!(Arc::strong_count(&shared), 1); let request = Arc::try_unwrap(shared).unwrap(); - let owned_error = datanode.handle(request).await.unwrap_err(); + let owned_error = datanode + .handle_request(Arc::new(request)) + .await + .unwrap_err(); for error in [owned_error, shared_error] { assert_eq!( error.retry_hint(), @@ -1053,8 +1212,281 @@ mod tests { .unwrap() } + #[tokio::test] + async fn test_partial_response_validation_and_compression() { + let server = RegionTestServer::start().await; + let requester = server.create_requester(ChannelConfig::default(), true, true); + let mut request = build_insert_request(); + if let Some(region_request::Body::Inserts(inserts)) = &mut request.body { + inserts.requests.push(inserts.requests[0].clone()); + } + let request = Arc::new(with_partial_result(request)); + let failure = RegionRequestFailure { + request_index: 0, + status: Some(PbStatus { + status_code: StatusCode::RegionNotFound as u32, + err_msg: "moved".to_string(), + retry_hint: "retry_after_route_refresh".to_string(), + }), + }; + for (failures, valid) in [ + (vec![], true), + (vec![failure.clone()], true), + (vec![failure.clone(), failure.clone()], false), + ( + vec![RegionRequestFailure { + request_index: u32::MAX, + ..failure.clone() + }], + false, + ), + ( + vec![RegionRequestFailure { + request_index: 0, + status: None, + }], + false, + ), + ( + vec![RegionRequestFailure { + request_index: 0, + status: Some(PbStatus::default()), + }], + false, + ), + ( + vec![RegionRequestFailure { + request_index: 0, + status: Some(PbStatus { + status_code: u32::MAX, + ..Default::default() + }), + }], + false, + ), + ] { + let has_error = !failures.is_empty(); + server + .service + .responses + .lock() + .unwrap() + .push_back(Ok(PbRegionResponse { + header: Some(ResponseHeader { + status: Some(PbStatus { + status_code: if has_error { + StatusCode::PartialFailure + } else { + StatusCode::Success + } as u32, + ..Default::default() + }), + }), + affected_rows: 3, + extensions: std::collections::HashMap::from([("test".to_string(), vec![1, 2])]), + failures, + ..Default::default() + })); + let response = requester.handle_inner(Arc::clone(&request)).await; + if valid { + if has_error { + let error = response.unwrap_err(); + assert!(error.output_msg().contains("entry 0")); + assert!(error.output_msg().contains("moved")); + let error: common_meta::error::Error = error.into(); + assert!(error.output_msg().contains("moved")); + let common_meta::error::Error::PartialRegion { response } = error else { + panic!("Expected partial region failure"); + }; + assert_eq!(response.affected_rows, 3); + assert_eq!(response.extensions["test"], vec![1, 2]); + assert_eq!( + response.failures[0].1.retry_hint(), + RetryHint::RetryAfterRouteRefresh + ); + } else { + let response = response.unwrap(); + assert_eq!(response.affected_rows, 3); + assert_eq!(response.extensions["test"], vec![1, 2]); + } + } else { + assert!(matches!( + response.unwrap_err(), + Error::IllegalDatabaseResponse { .. } + )); + } + } + for (received, send, accept) in server.service.requests.lock().unwrap().iter() { + assert_eq!(received, request.as_ref()); + assert!(*send && *accept); + } + assert_eq!(Arc::strong_count(&request), 1); + } + + #[test] + fn test_partial_response_envelope() { + let failure = RegionRequestFailure { + request_index: 0, + status: Some(PbStatus { + status_code: StatusCode::RegionNotFound as u32, + retry_hint: RetryHint::RetryAfterRouteRefresh.as_str().to_string(), + ..Default::default() + }), + }; + for (status, failures, count) in [ + (Some(StatusCode::Success), vec![failure.clone()], Some(2)), + (Some(StatusCode::PartialFailure), vec![], Some(2)), + ( + Some(StatusCode::PartialFailure), + vec![failure.clone()], + None, + ), + (None, vec![failure.clone()], Some(2)), + ] { + let response = PbRegionResponse { + header: status.map(|code| ResponseHeader { + status: Some(PbStatus { + status_code: code as u32, + ..Default::default() + }), + }), + failures, + ..Default::default() + }; + assert!(matches!( + match count { + Some(count) => decode_write_response(response, count), + None => decode_region_response(response), + }, + Err(Error::IllegalDatabaseResponse { .. }) + )); + } + let response = PbRegionResponse { + header: Some(ResponseHeader { + status: Some(PbStatus { + status_code: StatusCode::PartialFailure as u32, + ..Default::default() + }), + }), + failures: vec![failure], + ..Default::default() + }; + // A partial envelope must not be interpreted as success by the header-only decoder. + assert!(check_response_header(&response.header).is_err()); + let error = decode_write_response(response, 1).unwrap_err(); + assert_eq!(error.status_code(), StatusCode::RegionNotFound); + assert_eq!(error.retry_hint(), RetryHint::RetryAfterRouteRefresh); + } + + #[test] + fn test_partial_error_preserves_status_and_retry_hint() { + let expected_message = "Region request has 2 failed entries; entry 0: first input failed"; + for code in [StatusCode::StorageUnavailable, StatusCode::InvalidArguments] { + let response = PbRegionResponse { + header: Some(ResponseHeader { + status: Some(PbStatus { + status_code: StatusCode::PartialFailure as u32, + ..Default::default() + }), + }), + failures: vec![ + RegionRequestFailure { + request_index: 1, + status: Some(PbStatus { + status_code: StatusCode::Internal as u32, + err_msg: "later input failed".to_string(), + ..Default::default() + }), + }, + RegionRequestFailure { + request_index: 0, + status: Some(PbStatus { + status_code: code as u32, + err_msg: "first input failed".to_string(), + retry_hint: RetryHint::NonRetryable.as_str().to_string(), + }), + }, + ], + ..Default::default() + }; + let error = decode_write_response(response, 2).unwrap_err(); + assert_eq!(code, error.status_code()); + assert_eq!(expected_message, error.to_string()); + assert_eq!(RetryHint::Retryable, error.retry_hint()); + let error: common_meta::error::Error = error.into(); + assert_eq!(code, error.status_code()); + assert_eq!(expected_message, error.to_string()); + assert_eq!(RetryHint::Retryable, error.retry_hint()); + let common_meta::error::Error::PartialRegion { response } = error else { + panic!("expected partial result") + }; + assert!( + response + .failures + .iter() + .all(|(_, error)| error.retry_hint() == RetryHint::NonRetryable) + ); + } + } + + #[tokio::test] + async fn test_region_rpc_errors() { + let server = RegionTestServer::start().await; + let requester = server.create_requester(ChannelConfig::default(), false, false); + let request = Arc::new(build_insert_request()); + let mut business_error = Status::unimplemented("unsupported write"); + business_error.metadata_mut().insert( + common_error::GREPTIME_DB_HEADER_ERROR_CODE, + (StatusCode::Unsupported as u32) + .to_string() + .parse() + .unwrap(), + ); + for status in [ + Status::unimplemented("unknown method"), + business_error, + Status::unavailable("lost response"), + Status::deadline_exceeded("timeout"), + ] { + let code = status.code(); + server + .service + .responses + .lock() + .unwrap() + .push_back(Err(status)); + let error = requester + .handle_inner(Arc::clone(&request)) + .await + .unwrap_err(); + assert!(matches!(error, Error::RegionServer { code: actual, .. } if actual == code)); + } + assert_eq!(server.service.requests.lock().unwrap().len(), 4); + } + #[test] fn test_check_response_header() { + for (hint, expected) in [ + ("retryable", RetryHint::Retryable), + ( + "retry_after_route_refresh", + RetryHint::RetryAfterRouteRefresh, + ), + ("non_retryable", RetryHint::NonRetryable), + ("", RetryHint::NonRetryable), + ("future_hint", RetryHint::NonRetryable), + ] { + let error = check_response_header(&Some(ResponseHeader { + status: Some(PbStatus { + status_code: StatusCode::RegionNotFound as u32, + err_msg: "region moved".to_string(), + retry_hint: hint.to_string(), + }), + })) + .unwrap_err(); + assert_eq!(error.retry_hint(), expected, "{hint}"); + } + let result = check_response_header(&None); assert!(matches!( result.unwrap_err(), @@ -1071,6 +1503,7 @@ mod tests { status: Some(PbStatus { status_code: StatusCode::Success as u32, err_msg: String::default(), + ..Default::default() }), })); assert!(result.is_ok()); @@ -1079,6 +1512,7 @@ mod tests { status: Some(PbStatus { status_code: u32::MAX, err_msg: String::default(), + ..Default::default() }), })); assert!(matches!( @@ -1090,6 +1524,7 @@ mod tests { status: Some(PbStatus { status_code: StatusCode::Internal as u32, err_msg: "blabla".to_string(), + ..Default::default() }), })); let Error::Server { code, msg, .. } = result.unwrap_err() else { diff --git a/src/common/error/src/status_code.rs b/src/common/error/src/status_code.rs index b2c99c8e06b..592e40fec45 100644 --- a/src/common/error/src/status_code.rs +++ b/src/common/error/src/status_code.rs @@ -45,6 +45,8 @@ pub enum StatusCode { DeadlineExceeded = 1008, /// Service got suspended for various reason. For example, resources exceed limit. Suspended = 1009, + /// One or more entries failed; inspect per-entry results before retrying. + PartialFailure = 1010, // ====== End of common status code ================ // ====== Begin of SQL related status code ========= @@ -150,7 +152,8 @@ impl StatusCode { | StatusCode::RuntimeResourcesExhausted | StatusCode::External => true, - StatusCode::Success + StatusCode::PartialFailure + | StatusCode::Success | StatusCode::Unsupported | StatusCode::InvalidArguments | StatusCode::InvalidSyntax @@ -300,7 +303,7 @@ macro_rules! define_into_tonic_status { pub fn status_to_tonic_code(status_code: StatusCode) -> Code { match status_code { StatusCode::Success => Code::Ok, - StatusCode::Unknown | StatusCode::External => Code::Unknown, + StatusCode::Unknown | StatusCode::External | StatusCode::PartialFailure => Code::Unknown, StatusCode::Unsupported => Code::Unimplemented, StatusCode::Unexpected | StatusCode::IllegalState diff --git a/src/common/meta/src/error.rs b/src/common/meta/src/error.rs index 3de9e4a3a46..3b8f55e8a4c 100644 --- a/src/common/meta/src/error.rs +++ b/src/common/meta/src/error.rs @@ -41,6 +41,15 @@ pub use retry_hint::{retry_hint_from_postgres_error, retry_hint_from_postgres_po #[snafu(visibility(pub))] #[stack_trace_debug] pub enum Error { + #[snafu(display( + "Region request has {} failed entries; {}", + response.failures.len(), + response.first_failure().map(|(index, error)| format!("entry {index}: {}", error.output_msg())).unwrap_or_default() + ))] + PartialRegion { + response: Box>, + }, + #[snafu(display("Empty key is not allowed"))] EmptyKey { #[snafu(implicit)] @@ -1195,6 +1204,9 @@ impl ErrorExt for Error { fn status_code(&self) -> StatusCode { use Error::*; match self { + Error::PartialRegion { response } => response + .first_failure() + .map_or(StatusCode::Internal, |(_, error)| error.status_code()), IllegalServerState { .. } | EtcdTxnOpResponse { .. } | EtcdFailed { .. } @@ -1361,6 +1373,7 @@ impl ErrorExt for Error { use Error::*; match self { + Error::PartialRegion { .. } => RetryHint::Retryable, RetryLater { source, .. } => match source.retry_hint() { RetryHint::RetryAfterRouteRefresh => RetryHint::RetryAfterRouteRefresh, RetryHint::Retryable | RetryHint::NonRetryable => RetryHint::Retryable, diff --git a/src/datanode/src/error.rs b/src/datanode/src/error.rs index 3f61d4b3d76..aaf1ff4d8df 100644 --- a/src/datanode/src/error.rs +++ b/src/datanode/src/error.rs @@ -30,6 +30,9 @@ use tokio::time::error::Elapsed; #[snafu(visibility(pub))] #[stack_trace_debug] pub enum Error { + #[snafu(display("Failed to write a batch group"))] + BatchWrite { source: std::sync::Arc }, + #[snafu(display("Failed to execute async task"))] AsyncTaskExecute { #[snafu(implicit)] @@ -479,6 +482,7 @@ impl ErrorExt for Error { fn status_code(&self) -> StatusCode { use Error::*; match self { + Self::BatchWrite { source } => source.status_code(), NewPlanDecoder { source, .. } | ExecuteLogicalPlan { source, .. } => { source.status_code() } @@ -561,6 +565,7 @@ impl ErrorExt for Error { use Error::*; match self { + Self::BatchWrite { source } => source.retry_hint(), RegionNotReady { .. } | RegionNotFound { .. } => RetryHint::RetryAfterRouteRefresh, RegionBusy { .. } | ConcurrentQueryLimiterClosed { .. } diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index 0b7d36e9f2c..15bf786c389 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -29,12 +29,12 @@ use std::sync::{Arc, RwLock}; use std::task::{Context, Poll}; use std::time::{Duration, Instant}; -use api::region::RegionResponse; +use api::region::{BatchRegionResponse, RegionResponse}; use api::v1::meta::TopicStat; use api::v1::region::sync_request::ManifestInfo; use api::v1::region::{ - ListMetadataRequest, RegionRequest as RegionRequestV1, RegionResponse as RegionResponseV1, - SyncRequest, region_request, + ListMetadataRequest, RegionRequest as RegionRequestV1, RegionRequestFailure, + RegionResponse as RegionResponseV1, SyncRequest, region_request, }; use api::v1::{ResponseHeader, Status}; use arrow_flight::{FlightData, Ticket}; @@ -58,7 +58,6 @@ use datafusion::datasource::TableProvider; use datafusion_common::tree_node::TreeNode; use datatypes::schema::SchemaRef; use either::Either; -use futures_util::future::try_join_all; use futures_util::stream::FuturesUnordered; use futures_util::{Stream, StreamExt}; use metric_engine::engine::MetricEngine; @@ -82,7 +81,7 @@ use servers::grpc::region_server::RegionServerHandler; use session::context::{ FLIGHT_METRICS_HEARTBEAT_INTERVAL, QueryContext, QueryContextBuilder, QueryContextRef, }; -use snafu::{OptionExt, ResultExt, ensure}; +use snafu::{IntoError, OptionExt, ResultExt, ensure}; use store_api::metric_engine_consts::{ FILE_ENGINE_NAME, LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME, }; @@ -92,11 +91,11 @@ use store_api::region_engine::{ SyncRegionFromRequest, }; use store_api::region_request::{ - AffectedRows, BatchRegionDdlRequest, RegionCatchupRequest, RegionCloseRequest, - RegionOpenRequest, RegionRequest, + self as store_region_request, AffectedRows, BatchRegionDdlRequest, RegionCatchupRequest, + RegionCloseRequest, RegionOpenRequest, RegionRequest, }; use store_api::storage::RegionId; -use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc, oneshot}; +use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc}; use tokio::time::{self, timeout}; use tonic::{Request, Response, Result as TonicResult}; @@ -540,7 +539,7 @@ impl RegionServer { let span = tracing_context.attach(info_span!("RegionServer::handle_batch_ddl_requests")); self.inner - .handle_batch_request(batch_request) + .handle_batch_ddl_requests(batch_request) .trace(span) .await } @@ -548,53 +547,50 @@ impl RegionServer { async fn handle_requests_in_parallel( &self, request: region_request::Body, - ) -> Result { - let requests = - RegionRequest::try_from_request_body(request).context(BuildRegionRequestsSnafu)?; - + ) -> Result> { + let requests = prepare_batch_write_requests(request); let requests = match self.try_handle_mito_batch_puts(requests).await? { Either::Left(response) => return Ok(response), Either::Right(requests) => requests, }; - // Try to optimize batch Put requests for metric engine - // Returns either Some(response) or None(requests_back) - match self.try_handle_metric_batch_puts(requests).await? { - Either::Left(response) => Ok(response), - Either::Right(requests) => { - // Fallback: original parallel processing - let tracing_context = TracingContext::from_current_span(); - let join_tasks = - requests - .into_iter() - .map(|(region_id, req): (RegionId, RegionRequest)| { - let self_to_move = self; - let span = tracing_context.attach(info_span!( - "RegionServer::handle_region_request", - region_id = region_id.to_string() - )); - async move { - self_to_move - .handle_request(region_id, req) - .trace(span) - .await - } - }); - - let results = try_join_all(join_tasks).await?; - let mut affected_rows = 0; - let mut extensions = HashMap::new(); - for result in results { - affected_rows += result.affected_rows; - extensions.extend(result.extensions); + let requests = match self.try_handle_metric_batch_puts(requests).await? { + Either::Left(response) => return Ok(response), + Either::Right(requests) => requests, + }; + let tracing_context = TracingContext::from_current_span(); + let results = futures::future::join_all(requests.into_iter().map(|request| { + let span = tracing_context.attach(info_span!( + "RegionServer::handle_region_request", + region_id = request.region_id.to_string() + )); + async move { + ( + request.index, + self.handle_request(request.region_id, request.request) + .trace(span) + .await, + ) + } + })) + .await; + let mut response = BatchRegionResponse { + affected_rows: 0, + extensions: HashMap::new(), + failures: Vec::new(), + }; + for (index, result) in results { + match result { + Ok(result) => { + response.affected_rows += result.affected_rows; + response.extensions.extend(result.extensions); + } + Err(error) => { + log_batch_write_error(&error, index); + response.failures.push((index, error)); } - - Ok(RegionResponse { - affected_rows, - extensions, - metadata: Vec::new(), - }) } } + Ok(response) } async fn handle_requests_in_serial( @@ -625,65 +621,35 @@ impl RegionServer { }) } - /// Attempts to optimize batch Put requests for metric engine. - /// - /// Returns Either::Left(response) if optimization succeeded, - /// or Either::Right(original_requests) to fall back to parallel processing. - /// - /// This avoids cloning requests when optimization cannot be applied. + /// Attempts metric batch writes, returning the requests when ordinary execution is needed. async fn try_handle_metric_batch_puts( &self, - requests: Vec<(RegionId, RegionRequest)>, - ) -> Result>> { - if requests.is_empty() { + requests: Vec, + ) -> Result, Vec>> { + let Some(first) = requests.first() else { return Ok(Either::Right(requests)); - } - - // Quick check: verify first request is Put and is metric engine - if !matches!(requests[0].1, RegionRequest::Put(_)) { - return Ok(Either::Right(requests)); - } - let first_region_id = requests[0].0; - let request_type = requests[0].1.request_type(); - - // SAFETY: If the first request belongs to metric engine, then ALL requests - // in this batch are guaranteed to belong to metric engine. This invariant - // is maintained by the request batching logic upstream. - let engine = match self - .inner - .get_engine(first_region_id, &RegionChange::None)? - { - CurrentEngine::Engine(e) => e, - _ => return Ok(Either::Right(requests)), }; - - if engine.name() != METRIC_ENGINE_NAME { + if !matches!(first.request, RegionRequest::Put(_)) { return Ok(Either::Right(requests)); } - // Check if ALL requests are Put (now we know it's worth checking) - let mut all_puts = true; - for (_, req) in &requests { - if !matches!(req, RegionRequest::Put(_)) { - all_puts = false; - break; - } - } + let first_region_id = first.region_id; + let request_type = first.request.request_type(); - if !all_puts { + // Upstream batching keeps requests from the same engine together. + // Ordinary execution reports engine lookup failures per entry. + let Ok(CurrentEngine::Engine(engine)) = + self.inner.get_engine(first_region_id, &RegionChange::None) + else { + return Ok(Either::Right(requests)); + }; + if engine.name() != METRIC_ENGINE_NAME + || !requests + .iter() + .all(|request| matches!(request.request, RegionRequest::Put(_))) + { return Ok(Either::Right(requests)); } - - // Now extract Put requests by consuming ownership (zero clone!) - let put_requests = requests.into_iter().map(|(region_id, req)| { - if let RegionRequest::Put(put) = req { - (region_id, put) - } else { - unreachable!("Already checked all are Put") - } - }); - - // Downcast to MetricEngine and call batch API let metric_engine = engine .as_any() .downcast_ref::() @@ -692,75 +658,115 @@ impl RegionServer { })? .clone(); + let indices: Vec<_> = requests + .iter() + .map(|request| (request.index, request.region_id)) + .collect(); + let put_requests = requests.into_iter().map(|request| match request.request { + RegionRequest::Put(put) => (request.region_id, put), + _ => unreachable!("Already checked all are Put"), + }); let tracing_context = TracingContext::from_current_span(); let batch_size = put_requests.len(); let span = tracing_context.attach(info_span!( "RegionServer::handle_metric_batch_puts", - batch_size = batch_size, + batch_size = batch_size )); - let result = common_runtime::spawn_ingest(async move { + let results = common_runtime::spawn_ingest(async move { metric_engine .put_regions_batch(put_requests) .trace(span) .await }) .await - .context(RuntimeJoinSnafu { request_type })? - .map_err(BoxedError::new) - .context(HandleRegionRequestSnafu { - region_id: first_region_id, - }) - .map(RegionResponse::new); - record_ingest_result(request_type, batch_size, &result); - result.map(Either::Left) + .context(RuntimeJoinSnafu { request_type })?; + let mut response = BatchRegionResponse { + affected_rows: 0, + extensions: HashMap::new(), + failures: Vec::new(), + }; + for group in results { + let region_id = indices[group.request_indices[0]].1; + let result = group + .result + .map_err(BoxedError::new) + .context(HandleRegionRequestSnafu { region_id }) + .map(RegionResponse::new); + record_ingest_result(request_type, group.request_indices.len(), &result); + match result { + Ok(result) => { + response.affected_rows += result.affected_rows; + } + Err(error) => { + log_batch_write_error(&error, indices[group.request_indices[0]].0); + let error = Arc::new(error); + response + .failures + .extend(group.request_indices.into_iter().map(|index| { + ( + indices[index].0, + error::BatchWriteSnafu.into_error(Arc::clone(&error)), + ) + })); + } + } + } + Ok(Either::Left(response)) } - /// Puts into ready mito regions together through - /// [`RegionEngine::handle_batch_put_requests`], so that each mito worker gets the puts of the - /// request to its regions at once and writes them in one batch. One ingest task awaits every - /// response at once and accounts each, for the regions it answers, as the per-region path - /// does; it hands the first error back at once and keeps accounting the rest, also when the - /// caller goes away. - /// - /// Returns Either::Right(original_requests) unless - /// [`mito_engine_for_batch_puts`](Self::mito_engine_for_batch_puts) accepts the requests. + /// Batches puts into ready mito regions and retains every entry's outcome. async fn try_handle_mito_batch_puts( &self, - requests: Vec<(RegionId, RegionRequest)>, - ) -> Result>> { + requests: Vec, + ) -> Result, Vec>> { let Some(engine) = self.mito_engine_for_batch_puts(&requests) else { return Ok(Either::Right(requests)); }; - - let request_type = requests[0].1.request_type(); + let request_type = requests[0].request.request_type(); + let mut indices: HashMap> = HashMap::new(); let puts = requests .into_iter() - .map(|(region_id, request)| match request { - RegionRequest::Put(put) => (region_id, put), - _ => unreachable!("Already checked all are Put"), + .map(|request| { + indices + .entry(request.region_id) + .or_default() + .push_back(request.index); + match request.request { + RegionRequest::Put(put) => (request.region_id, put), + _ => unreachable!("Already checked all are Put"), + } }) .collect(); let started = Instant::now(); let tracing_context = TracingContext::from_current_span(); - let (result_tx, result_rx) = oneshot::channel(); - let batch = common_runtime::spawn_ingest(async move { + common_runtime::spawn_ingest(async move { let mut pending = engine .handle_batch_put_requests(puts) .into_iter() .map(|(region_ids, response)| { + let request_indices = region_ids + .iter() + .filter_map(|region_id| { + indices + .get_mut(region_id) + .and_then(|indices| indices.pop_front()) + }) + .collect::>(); let span = tracing_context.attach(info_span!( "RegionServer::handle_region_request", region_id = region_ids[0].to_string() )); - async move { (region_ids, response.await) }.trace(span) + async move { (region_ids, request_indices, response.await) }.trace(span) }) .collect::>(); let histogram = crate::metrics::HANDLE_REGION_REQUEST_ELAPSED.with_label_values(&[request_type]); - let mut result_tx = Some(result_tx); - let mut affected_rows = 0; - let mut extensions = HashMap::new(); - while let Some((region_ids, result)) = pending.next().await { + let mut response = BatchRegionResponse { + affected_rows: 0, + extensions: HashMap::new(), + failures: Vec::new(), + }; + while let Some((region_ids, request_indices, result)) = pending.next().await { let elapsed = started.elapsed().as_secs_f64(); for _ in ®ion_ids { histogram.observe(elapsed); @@ -770,53 +776,43 @@ impl RegionServer { }); record_ingest_result(request_type, region_ids.len(), &result); match result { - Ok(response) => { - affected_rows += response.affected_rows; - extensions.extend(response.extensions); + Ok(result) => { + response.affected_rows += result.affected_rows; + response.extensions.extend(result.extensions); } - Err(err) => { - if let Some(result_tx) = result_tx.take() { - let _ = result_tx.send(Err(err)); - } + Err(error) => { + log_batch_write_error(&error, request_indices[0]); + let error = Arc::new(error); + response + .failures + .extend(request_indices.into_iter().map(|index| { + (index, error::BatchWriteSnafu.into_error(Arc::clone(&error))) + })); } } } - if let Some(result_tx) = result_tx { - let _ = result_tx.send(Ok(RegionResponse { - affected_rows, - extensions, - metadata: Vec::new(), - })); - } - }); - - match result_rx.await { - Ok(result) => result.map(Either::Left), - // The batch task answers on every path it completes, so it panicked or was cancelled. - Err(_) => { - batch.await.context(RuntimeJoinSnafu { request_type })?; - UnexpectedSnafu { - violated: "the batch put task completed without answering", - } - .fail() - } - } + response.failures.sort_unstable_by_key(|(index, _)| *index); + response + }) + .await + .context(RuntimeJoinSnafu { request_type }) + .map(Either::Left) } /// Returns the mito engine when there are at least two requests, every one a put into a ready /// mito region. fn mito_engine_for_batch_puts( &self, - requests: &[(RegionId, RegionRequest)], + requests: &[IndexedRegionRequest], ) -> Option { if requests.len() < 2 { return None; } - let first_region_id = requests[0].0; - let all_puts_to_ready_mito_regions = requests.iter().all(|(region_id, request)| { - matches!(request, RegionRequest::Put(_)) + let first_region_id = requests[0].region_id; + let all_puts_to_ready_mito_regions = requests.iter().all(|request| { + matches!(request.request, RegionRequest::Put(_)) && matches!( - self.inner.region_map.get(region_id).as_deref(), + self.inner.region_map.get(&request.region_id).as_deref(), Some(RegionEngineWithStatus::Ready(engine)) if engine.name() == MITO_ENGINE_NAME ) }); @@ -1062,22 +1058,45 @@ impl RegionServerHandler for RegionServer { { return self.handle_flow_write_request(request).await; } + let accept_partial_result = request + .header + .as_ref() + .and_then(|header| header.query_context.as_ref()) + .and_then(|context| context.extensions.get(api::region::ACCEPT_PARTIAL_RESULT)) + .is_some_and(|value| value == "true"); let body = request.body.context(servers_error::InvalidQuerySnafu { reason: "Expecting non-empty region request body.", })?; - self.handle(body).await + self.handle_inner(body, accept_partial_result).await } async fn handle(&self, request: region_request::Body) -> ServerResult { + self.handle_inner(request, false).await + } +} + +impl RegionServer { + async fn handle_inner( + &self, + request: region_request::Body, + accept_partial_result: bool, + ) -> ServerResult { + let request_count = match &request { + region_request::Body::Inserts(requests) => Some(requests.requests.len()), + region_request::Body::Deletes(requests) => Some(requests.requests.len()), + _ => None, + }; + if let Some(request_count) = request_count { + return self + .handle_write_request(request, request_count, accept_partial_result) + .await; + } let failed_requests_cnt = crate::metrics::REGION_SERVER_REQUEST_FAILURE_COUNT .with_label_values(&[request.as_ref()]); let response = match &request { region_request::Body::Creates(_) | region_request::Body::Drops(_) | region_request::Body::Alters(_) => self.handle_batch_ddl_requests(request).await, - region_request::Body::Inserts(_) | region_request::Body::Deletes(_) => { - self.handle_requests_in_parallel(request).await - } region_request::Body::Sync(sync_request) => { self.handle_sync_region_request(sync_request).await } @@ -1092,11 +1111,8 @@ impl RegionServerHandler for RegionServer { _ => self.handle_requests_in_serial(request).await, } .map_err(BoxedError::new) - .inspect_err(|_| { - failed_requests_cnt.inc(); - }) + .inspect_err(|_| failed_requests_cnt.inc()) .context(ExecuteGrpcRequestSnafu)?; - Ok(RegionResponseV1 { header: Some(ResponseHeader { status: Some(Status { @@ -1107,8 +1123,35 @@ impl RegionServerHandler for RegionServer { affected_rows: response.affected_rows as _, extensions: response.extensions, metadata: response.metadata, + failures: Vec::new(), }) } + + async fn handle_write_request( + &self, + request: region_request::Body, + entry_count: usize, + accept_partial_result: bool, + ) -> ServerResult { + let failed_requests_cnt = crate::metrics::REGION_SERVER_REQUEST_FAILURE_COUNT + .with_label_values(&[request.as_ref()]); + let mut result = self + .handle_requests_in_parallel(request) + .await + .map_err(BoxedError::new) + .inspect_err(|_| failed_requests_cnt.inc()) + .context(ExecuteGrpcRequestSnafu)?; + if !result.failures.is_empty() { + failed_requests_cnt.inc(); + // Old clients and single-entry writes retain tonic errors and retry metadata. + if (!accept_partial_result || entry_count == 1) + && let Some((_, error)) = result.failures.drain(..).min_by_key(|(index, _)| *index) + { + return Err(BoxedError::new(error)).context(ExecuteGrpcRequestSnafu); + } + } + Ok(encode_batch_write_response(result)) + } } #[async_trait] @@ -1166,6 +1209,97 @@ impl FlightCraft for RegionServer { } } +/// A decoded write retains its protobuf entry index even when preceding entries are empty. +struct IndexedRegionRequest { + index: usize, + region_id: RegionId, + request: RegionRequest, +} + +fn prepare_batch_write_requests(body: region_request::Body) -> Vec { + match body { + region_request::Body::Inserts(inserts) => inserts + .requests + .into_iter() + .enumerate() + .filter_map(|(index, request)| { + store_region_request::make_region_put(request).map(|(region_id, request)| { + IndexedRegionRequest { + index, + region_id, + request: RegionRequest::Put(request), + } + }) + }) + .collect(), + region_request::Body::Deletes(deletes) => deletes + .requests + .into_iter() + .enumerate() + .filter_map(|(index, request)| { + store_region_request::make_region_delete(request).map(|(region_id, request)| { + IndexedRegionRequest { + index, + region_id, + request: RegionRequest::Delete(request), + } + }) + }) + .collect(), + _ => unreachable!("Only inserts/deletes are prepared as batch writes"), + } +} + +fn encode_batch_write_response(result: BatchRegionResponse) -> RegionResponseV1 { + RegionResponseV1 { + header: Some(ResponseHeader { + status: Some(Status { + status_code: if result.failures.is_empty() { + StatusCode::Success + } else { + StatusCode::PartialFailure + } as u32, + err_msg: if let Some((index, error)) = result.first_failure() { + format!( + "{} region requests failed; entry {index}: {}", + result.failures.len(), + error.output_msg() + ) + } else { + String::new() + }, + ..Default::default() + }), + }), + metadata: Vec::new(), + affected_rows: result.affected_rows as _, + extensions: result.extensions, + failures: result + .failures + .into_iter() + .map(|(index, error)| RegionRequestFailure { + request_index: index as u32, + status: Some(Status { + status_code: error.status_code() as u32, + err_msg: error.output_msg(), + retry_hint: error.retry_hint().as_str().to_string(), + }), + }) + .collect(), + } +} + +fn log_batch_write_error(error: &error::Error, index: usize) { + if error.status_code().should_log_error() { + error!(error; "Failed to handle batch entry, request_index: {}", index); + } else { + debug!( + "Failed to handle batch entry, request_index: {}, err: {}", + index, error + ); + } +} + #[derive(Clone)] enum RegionEngineWithStatus { // An opening, or creating region. @@ -1679,7 +1813,7 @@ impl RegionServerInner { // Handle requests in batch. // // limitation: all create requests must be in the same engine. - pub async fn handle_batch_request( + pub async fn handle_batch_ddl_requests( &self, batch_request: BatchRegionDdlRequest, ) -> Result { @@ -2207,7 +2341,7 @@ mod tests { use std::sync::Arc; use api::v1::region::{InsertRequest, InsertRequests}; - use api::v1::{Rows, SemanticType}; + use api::v1::{PartitionExprVersion, Rows, SemanticType}; use common_error::ext::{ErrorExt, RetryHint}; use common_recordbatch::RecordBatches; use common_recordbatch::adapter::{RecordBatchMetrics, RegionWatermarkEntry}; @@ -2617,6 +2751,170 @@ mod tests { assert_eq!(err.status_code(), StatusCode::RegionNotReady); } + #[test] + fn test_batch_write_summary_uses_first_input_failure() { + let first = error::RegionNotFoundSnafu { + region_id: RegionId::new(1, 0), + } + .build(); + let expected = format!("2 region requests failed; entry 0: {}", first.output_msg()); + let later = error::RegionBusySnafu { + region_id: RegionId::new(1, 1), + } + .build(); + let response = encode_batch_write_response(BatchRegionResponse { + affected_rows: 0, + extensions: HashMap::new(), + failures: vec![(1, later), (0, first)], + }); + let status = response.header.unwrap().status.unwrap(); + assert_eq!(status.status_code, StatusCode::PartialFailure as u32); + assert_eq!(status.err_msg, expected); + assert_eq!(response.failures[0].request_index, 1); + assert_eq!(response.failures[1].request_index, 0); + } + + #[tokio::test] + async fn test_write_failure_response_compatibility() { + let server = mock_region_server(); + let missing = RegionId::new(1, 1); + let registering = RegionId::new(1, 2); + let (engine, _receiver) = MockRegionEngine::new(MITO_ENGINE_NAME); + server + .inner + .region_map + .insert(registering, RegionEngineWithStatus::Registering(engine)); + for (region_id, code) in [ + (missing, StatusCode::RegionNotFound), + (registering, StatusCode::RegionNotReady), + ] { + for delete in [false, true] { + for accept_partial in [false, true] { + // A preceding empty entry must still count toward the original envelope size. + for with_empty_entry in [false, true] { + let rows = Some(Rows { + schema: Vec::new(), + rows: Vec::new(), + }); + let body = if delete { + let request = api::v1::region::DeleteRequest { + region_id: region_id.as_u64(), + rows, + ..Default::default() + }; + region_request::Body::Deletes(api::v1::region::DeleteRequests { + requests: if with_empty_entry { + vec![Default::default(), request] + } else { + vec![request] + }, + }) + } else { + let request = InsertRequest { + region_id: region_id.as_u64(), + rows, + ..Default::default() + }; + region_request::Body::Inserts(InsertRequests { + requests: if with_empty_entry { + vec![Default::default(), request] + } else { + vec![request] + }, + }) + }; + let mut header = api::v1::region::RegionRequestHeader::default(); + if accept_partial { + header + .query_context + .get_or_insert_with(Default::default) + .extensions + .insert( + api::region::ACCEPT_PARTIAL_RESULT.to_string(), + "true".to_string(), + ); + } + let result = server + .handle_region_request(RegionRequestV1 { + header: Some(header), + body: Some(body), + }) + .await; + if accept_partial && with_empty_entry { + let response = result.unwrap(); + assert_eq!( + StatusCode::PartialFailure as u32, + response.header.unwrap().status.unwrap().status_code + ); + assert_eq!(1, response.failures.len()); + assert_eq!(1, response.failures[0].request_index); + let status = response.failures[0].status.as_ref().unwrap(); + assert_eq!(code as u32, status.status_code); + assert_eq!( + RetryHint::RetryAfterRouteRefresh.as_str(), + status.retry_hint + ); + } else { + let status: tonic::Status = result.unwrap_err().into(); + assert_eq!( + common_error::status_code::status_to_tonic_code(code), + status.code() + ); + assert_eq!( + (code as u32).to_string(), + status + .metadata() + .get(common_error::GREPTIME_DB_HEADER_ERROR_CODE) + .unwrap() + .to_str() + .unwrap() + ); + assert_eq!( + RetryHint::RetryAfterRouteRefresh.as_str(), + status + .metadata() + .get(common_error::GREPTIME_DB_HEADER_ERROR_RETRY_HINT) + .unwrap() + .to_str() + .unwrap() + ); + } + } + } + } + } + } + + #[tokio::test] + async fn test_legacy_batch_returns_first_input_error() { + let server = mock_region_server(); + let first = RegionId::new(1, 1); + let second = RegionId::new(1, 2); + let (engine, _receiver) = MockRegionEngine::with_custom_apply_fn("mock", |engine| { + engine.handle_request_delay = Some(Duration::from_millis(10)); + engine.handle_request_mock_fn = Some(Box::new(|region_id, _| { + error::RegionBusySnafu { region_id }.fail() + })); + }); + server.register_test_region(first, engine); + let requests = [first, second] + .map(|region_id| InsertRequest { + region_id: region_id.as_u64(), + rows: Some(Rows::default()), + ..Default::default() + }) + .to_vec(); + let error = server + .handle_region_request(RegionRequestV1 { + header: Some(Default::default()), + body: Some(region_request::Body::Inserts(InsertRequests { requests })), + }) + .await + .unwrap_err(); + assert_eq!(StatusCode::RegionBusy, error.status_code()); + assert_eq!(RetryHint::Retryable, error.retry_hint()); + } + #[tokio::test] async fn test_region_request_failed() { common_telemetry::init_default_ut_logging(); @@ -2642,8 +2940,7 @@ mod tests { .await .unwrap_err(); - let status = mock_region_server.inner.region_map.get(®ion_id); - assert!(status.is_none()); + assert!(!mock_region_server.inner.region_map.contains_key(®ion_id)); mock_region_server .inner @@ -2662,8 +2959,79 @@ mod tests { .await .unwrap_err(); - let status = mock_region_server.inner.region_map.get(®ion_id); - assert!(status.is_some()); + assert!(mock_region_server.inner.region_map.contains_key(®ion_id)); + + let insert = api::v1::region::InsertRequest { + region_id: region_id.as_u64(), + rows: Some(Rows { + schema: Vec::new(), + rows: Vec::new(), + }), + ..Default::default() + }; + let success_region = RegionId::new(1, 2); + let (success_engine, _receiver) = + MockRegionEngine::with_custom_apply_fn("mock", |engine| { + engine.handle_request_delay = Some(std::time::Duration::from_millis(10)); + engine.handle_request_mock_fn = Some(Box::new(|_, _| Ok(5))); + }); + mock_region_server.register_test_region(success_region, success_engine); + let successful_insert = api::v1::region::InsertRequest { + region_id: success_region.as_u64(), + ..insert.clone() + }; + let response = mock_region_server + .handle_region_request(RegionRequestV1 { + header: Some(api::v1::region::RegionRequestHeader { + query_context: Some(api::v1::QueryContext { + extensions: HashMap::from([( + api::region::ACCEPT_PARTIAL_RESULT.to_string(), + "true".to_string(), + )]), + ..Default::default() + }), + ..Default::default() + }), + body: Some(region_request::Body::Inserts(InsertRequests { + requests: vec![ + Default::default(), + insert.clone(), + successful_insert, + insert, + ], + })), + }) + .await + .unwrap(); + assert_eq!(response.affected_rows, 5); + assert_eq!( + response + .header + .as_ref() + .unwrap() + .status + .as_ref() + .unwrap() + .status_code, + StatusCode::PartialFailure as u32 + ); + assert_eq!( + response + .failures + .iter() + .map(|error| error.request_index) + .collect::>(), + vec![1, 3] + ); + assert!( + response + .failures + .iter() + .all(|error| error.status.as_ref().is_some_and(|status| { + status.status_code == StatusCode::Unexpected as u32 + && status.retry_hint == RetryHint::NonRetryable.as_str() + })) + ); } /// A region that fails to close is reported in the result of the stop, @@ -3231,6 +3599,114 @@ mod tests { } } + #[tokio::test] + async fn test_metric_batch_failure_indices() { + let mut env = TestEnv::with_prefix("region-server-metric-results").await; + let mito = env.create_engine(MitoConfig::default()).await; + let metric = Arc::new( + MetricEngine::try_new(mito, metric_engine::config::EngineConfig::default()).unwrap(), + ); + let mut server = mock_region_server(); + server.register_engine(metric.clone()); + let mut inserts = Vec::new(); + for table in [10, 20] { + let physical = RegionId::new(table, 0); + let logical = RegionId::new(table + 1, 0); + let mut create = CreateRequestBuilder::new().build(); + create.engine = METRIC_ENGINE_NAME.to_string(); + create.column_metadatas = vec![ + store_api::metadata::ColumnMetadata { + column_id: 0, + semantic_type: SemanticType::Timestamp, + column_schema: ColumnSchema::new( + "greptime_timestamp", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + }, + store_api::metadata::ColumnMetadata { + column_id: 1, + semantic_type: SemanticType::Field, + column_schema: ColumnSchema::new( + "greptime_value", + ConcreteDataType::float64_datatype(), + false, + ), + }, + ]; + create.primary_key.clear(); + create.table_dir = format!("metric-results/{table}/"); + create.options = HashMap::from([( + store_api::metric_engine_consts::PHYSICAL_TABLE_METADATA_KEY.to_string(), + String::new(), + )]); + server + .handle_request(physical, RegionRequest::Create(create.clone())) + .await + .unwrap(); + create.table_dir = format!("metric-results/{}/", table + 1); + create.options = HashMap::from([( + store_api::metric_engine_consts::LOGICAL_TABLE_METADATA_KEY.to_string(), + physical.as_u64().to_string(), + )]); + server + .handle_request(logical, RegionRequest::Create(create.clone())) + .await + .unwrap(); + inserts.push(InsertRequest { + region_id: logical.as_u64(), + rows: Some(Rows { + schema: rows_schema(&create), + rows: vec![api::v1::Row { + values: vec![ + api::v1::Value { + value_data: Some( + api::v1::value::ValueData::TimestampMillisecondValue(1), + ), + }, + api::v1::Value { + value_data: Some(api::v1::value::ValueData::F64Value(1.0)), + }, + ], + }], + }), + ..Default::default() + }); + } + // Both entries of the second physical group fail at write_data. + metric + .set_region_role(RegionId::new(20, 0), RegionRole::Follower) + .unwrap(); + let response = server + .handle_requests_in_parallel(region_request::Body::Inserts(InsertRequests { + requests: vec![ + Default::default(), + inserts[0].clone(), + inserts[1].clone(), + inserts[1].clone(), + ], + })) + .await + .unwrap(); + assert_eq!(1, response.affected_rows); + let mut failures = response.failures; + failures.sort_unstable_by_key(|(index, _)| *index); + assert_eq!( + vec![2, 3], + failures.iter().map(|(index, _)| *index).collect::>() + ); + assert!( + failures + .iter() + .all(|(_, error)| error.status_code() == StatusCode::RegionNotReady) + ); + assert!( + failures + .iter() + .all(|(_, error)| error.retry_hint() == RetryHint::RetryAfterRouteRefresh) + ); + } + #[tokio::test] async fn test_mito_puts_of_one_request_share_a_wal_append() { // One worker that takes one message at a time, so that only puts that @@ -3293,49 +3769,126 @@ mod tests { } #[tokio::test] - async fn test_mito_batch_puts_fail_without_waiting_for_other_regions() { - let mut env = TestEnv::with_prefix("region-server-mito-batch-puts-error").await; - let config = MitoConfig { - num_workers: 2, - ..Default::default() - }; - let (region_server, log_store, wal_options) = - region_server_on_object_store_wal(&mut env, config).await; - // Regions 1 and 2 of table 1 belong to different workers. The put to region 1 waits for a - // WAL object whose create is held; Mito does not know region 2, so its put fails at once. - let held = RegionId::new(1, 1); - let unknown = RegionId::new(1, 2); - let held_insert = create_region(®ion_server, held, &wal_options).await; - let unknown_insert = InsertRequest { - region_id: unknown.as_u64(), - ..held_insert.clone() - }; - let mito = region_server - .inner - .region_map - .get(&held) - .unwrap() - .clone() - .into_engine(); - region_server - .inner - .region_map - .insert(unknown, RegionEngineWithStatus::Ready(mito)); - log_store.hold_creates(); + async fn test_mito_batch_failure_indices_for_same_region() { + for failed_index in [0, 1] { + let mut env = TestEnv::with_prefix("region-server-mito-same-region-failure").await; + let (server, _log_store, wal_options) = + region_server_on_object_store_wal(&mut env, MitoConfig::default()).await; + let region_id = RegionId::new(1, 1); + let insert = create_region(&server, region_id, &wal_options).await; + let mut requests = vec![insert.clone(), insert]; + requests[failed_index].partition_expr_version = Some(PartitionExprVersion { value: 1 }); + let response = server + .handle_requests_in_parallel(region_request::Body::Inserts(InsertRequests { + requests, + })) + .await + .unwrap(); + assert_eq!(response.affected_rows, 3); + assert_eq!(response.failures.len(), 1); + assert_eq!(response.failures[0].0, failed_index); + assert_eq!( + response.failures[0].1.status_code(), + StatusCode::InvalidArguments + ); + let engine = server + .inner + .region_map + .get(®ion_id) + .unwrap() + .clone() + .into_engine(); + let mito = engine.as_any().downcast_ref::().unwrap(); + let stream = mito + .scan_to_stream(region_id, Default::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + assert_eq!( + batches.iter().map(|batch| batch.num_rows()).sum::(), + 3 + ); + } + } - let requests = vec![held_insert, unknown_insert]; - let err = tokio::time::timeout( - Duration::from_secs(10), - region_server.handle_requests_in_parallel(region_request::Body::Inserts( - InsertRequests { requests }, - )), - ) - .await - .unwrap() - .unwrap_err(); + #[tokio::test] + async fn test_mito_batch_puts_retain_successful_regions() { + for accept_partial_result in [false, true] { + let mut env = TestEnv::with_prefix("region-server-mito-batch-puts-error").await; + let config = MitoConfig { + num_workers: 2, + ..Default::default() + }; + let (region_server, log_store, wal_options) = + region_server_on_object_store_wal(&mut env, config).await; + // Regions 1 and 2 of table 1 belong to different workers. The put to region 1 waits for a + // WAL object whose create is held; Mito does not know region 2, so its put fails at once. + let held = RegionId::new(1, 1); + let unknown = RegionId::new(1, 2); + let held_insert = create_region(®ion_server, held, &wal_options).await; + let unknown_insert = InsertRequest { + region_id: unknown.as_u64(), + ..held_insert.clone() + }; + let mito = region_server + .inner + .region_map + .get(&held) + .unwrap() + .clone() + .into_engine(); + region_server + .inner + .region_map + .insert(unknown, RegionEngineWithStatus::Ready(mito.clone())); + log_store.hold_creates(); - assert_eq!(StatusCode::RegionNotFound, err.status_code()); - log_store.release_creates(); + let requests = vec![Default::default(), held_insert, unknown_insert]; + let mut header = api::v1::region::RegionRequestHeader::default(); + if accept_partial_result { + api::region::set_accept_partial_result(&mut header); + } + let write = region_server.handle_region_request(RegionRequestV1 { + header: Some(header), + body: Some(region_request::Body::Inserts(InsertRequests { requests })), + }); + tokio::pin!(write); + assert!( + tokio::time::timeout(Duration::from_millis(100), &mut write) + .await + .is_err() + ); + log_store.release_creates(); + let result = tokio::time::timeout(Duration::from_secs(10), write) + .await + .unwrap(); + if accept_partial_result { + let response = result.unwrap(); + assert_eq!(3, response.affected_rows); + assert_eq!( + StatusCode::PartialFailure as u32, + response.header.unwrap().status.unwrap().status_code + ); + assert_eq!(1, response.failures.len()); + assert_eq!(2, response.failures[0].request_index); + assert_eq!( + StatusCode::RegionNotFound as u32, + response.failures[0].status.as_ref().unwrap().status_code + ); + } else { + assert_eq!( + StatusCode::RegionNotFound, + result.unwrap_err().status_code() + ); + } + let mito = mito.as_any().downcast_ref::().unwrap(); + let stream = mito.scan_to_stream(held, Default::default()).await.unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + assert_eq!( + batches.iter().map(|batch| batch.num_rows()).sum::(), + 3 + ); + } } #[tokio::test] @@ -3356,13 +3909,19 @@ mod tests { // One put takes the per-region path and two take the batch path. for requests in [vec![insert.clone()], vec![insert.clone(), insert]] { - let err = region_server + let count = requests.len(); + let response = region_server .handle_requests_in_parallel(region_request::Body::Inserts(InsertRequests { requests, })) .await - .unwrap_err(); - assert!(format!("{err:?}").contains("is stopped"), "{err:?}"); + .unwrap(); + assert_eq!(0, response.affected_rows); + assert_eq!(count, response.failures.len()); + for (index, (request_index, error)) in response.failures.iter().enumerate() { + assert_eq!(index, *request_index); + assert!(format!("{error:?}").contains("is stopped"), "{error:?}"); + } } } @@ -3433,7 +3992,11 @@ mod tests { skip_wal: false, partition_expr_version: None, }; - (region_id, RegionRequest::Put(request)) + IndexedRegionRequest { + index: 0, + region_id, + request: RegionRequest::Put(request), + } }; let delete = |region_id| { let request = RegionDeleteRequest { @@ -3441,7 +4004,11 @@ mod tests { hint: None, partition_expr_version: None, }; - (region_id, RegionRequest::Delete(request)) + IndexedRegionRequest { + index: 0, + region_id, + request: RegionRequest::Delete(request), + } }; for requests in [ diff --git a/src/datanode/src/region_server/flow_write.rs b/src/datanode/src/region_server/flow_write.rs index adb9d88568a..6bcb112bc78 100644 --- a/src/datanode/src/region_server/flow_write.rs +++ b/src/datanode/src/region_server/flow_write.rs @@ -129,6 +129,7 @@ fn response(reply: Reply) -> ServerResult { affected_rows: 0, extensions, metadata: Vec::new(), + failures: Vec::new(), }) } diff --git a/src/meta-srv/src/procedure/utils.rs b/src/meta-srv/src/procedure/utils.rs index ff88d1b38ba..9375e1a4392 100644 --- a/src/meta-srv/src/procedure/utils.rs +++ b/src/meta-srv/src/procedure/utils.rs @@ -356,11 +356,13 @@ pub mod mock { status: Some(PbStatus { status_code: 0, err_msg: String::default(), + ..Default::default() }), }), affected_rows: 0, extensions: Default::default(), metadata: Vec::new(), + failures: Vec::new(), }) } } diff --git a/src/metric-engine/AGENTS.md b/src/metric-engine/AGENTS.md index d1049478978..9c3820a0cc3 100644 --- a/src/metric-engine/AGENTS.md +++ b/src/metric-engine/AGENTS.md @@ -51,6 +51,7 @@ Mito2 data region. Queries against a physical region pass straight through. ## Public surface - Entry: `MetricEngine` in `src/metric-engine/src/engine.rs`, built via `try_new(mito, config)`. +- `MetricEngine::put_regions_batch` returns each physical group's outcome with its input indices; a failed group does not stop other groups. - Trait: `impl RegionEngine for MetricEngine` (name = `"metric"`). - Depends on `mito2`, `store-api`, and `mito-codec` (sparse primary key codec). diff --git a/src/metric-engine/src/engine.rs b/src/metric-engine/src/engine.rs index ec180839526..d01ca4650e6 100644 --- a/src/metric-engine/src/engine.rs +++ b/src/metric-engine/src/engine.rs @@ -488,6 +488,17 @@ impl RegionEngine for MetricEngine { } } +/// Outcome of a group of input requests. +/// +/// Requests rejected before physical-region grouping produce singleton failures. +#[derive(Debug)] +pub struct PutGroupResult { + /// Positions in the input batch, including separate entries for the same region. + pub request_indices: Vec, + /// Aggregate rows on success; no other physical group's outcome is implied. + pub result: Result, +} + impl MetricEngine { pub fn try_new(mito: MitoEngine, mut config: EngineConfig) -> Result { let metadata_region = MetadataRegion::new(mito.clone()); @@ -520,13 +531,11 @@ impl MetricEngine { self.inner.mito.clone() } - /// Batch put operation for multiple logical regions. - /// Requests are grouped by physical region; a failure can leave earlier - /// physical-region groups committed. + /// Writes all physical-region groups and retains each group's outcome. pub async fn put_regions_batch( &self, requests: impl ExactSizeIterator, - ) -> Result { + ) -> Vec { self.inner.put_regions_batch(requests).await } diff --git a/src/metric-engine/src/engine/put.rs b/src/metric-engine/src/engine/put.rs index b709d99d847..27ee3fba56c 100644 --- a/src/metric-engine/src/engine/put.rs +++ b/src/metric-engine/src/engine/put.rs @@ -29,7 +29,7 @@ use store_api::region_request::{ }; use store_api::storage::{RegionId, TableId}; -use crate::engine::MetricEngineInner; +use crate::engine::{MetricEngineInner, PutGroupResult}; use crate::error::{ ColumnNotFoundSnafu, CreateDefaultSnafu, ForbiddenPhysicalWriteSnafu, InvalidRequestSnafu, LogicalRegionNotFoundSnafu, PhysicalRegionNotFoundSnafu, Result, UnexpectedRequestSnafu, @@ -39,6 +39,12 @@ use crate::metrics::{FORBIDDEN_OPERATION_COUNT, MITO_OPERATION_ELAPSED}; use crate::row_modifier::{RowsIter, TableIdInput}; use crate::utils::to_data_region_id; +#[derive(Default)] +struct PutGroup { + request_indices: Vec, + requests: Vec<(RegionId, RegionPutRequest)>, +} + impl MetricEngineInner { /// Dispatch region put request pub async fn put_region( @@ -61,67 +67,63 @@ impl MetricEngineInner { } } - /// Batch write multiple logical regions to the same physical region. - /// - /// Dispatch region put requests in batch. - /// - /// Requests may span multiple physical regions. We group them by physical - /// region and write sequentially. This method fails fast on validation or - /// preparation errors within a group and stops at the first failure. - /// Writes in earlier physical-region groups are not rolled back if a later - /// group fails. + /// Retains every group's outcome instead of returning on the first failure. pub async fn put_regions_batch( &self, - requests: impl ExactSizeIterator, - ) -> Result { + mut requests: impl ExactSizeIterator, + ) -> Vec { let len = requests.len(); - if len == 0 { - return Ok(0); + return Vec::new(); } - let _timer = MITO_OPERATION_ELAPSED .with_label_values(&["put_batch"]) .start_timer(); - - // Fast path: single request, no batching overhead if len == 1 { - let (region_id, req) = requests.into_iter().next().unwrap(); - let is_putting_physical_region = - self.state.read().unwrap().exist_physical_region(region_id); - if is_putting_physical_region { - FORBIDDEN_OPERATION_COUNT.inc(); - return ForbiddenPhysicalWriteSnafu.fail(); + let (region_id, request) = requests.next().unwrap(); + return vec![PutGroupResult { + request_indices: vec![0], + result: self.put_region(region_id, request).await, + }]; + } + let mut groups = HashMap::new(); + let mut results = Vec::new(); + for (index, (region_id, request)) in requests.enumerate() { + if let Err(error) = self.group_put_request(&mut groups, index, region_id, request) { + results.push(PutGroupResult { + request_indices: vec![index], + result: Err(error), + }); } - - return self.put_logical_region(region_id, req).await; } - - let mut requests_per_physical: HashMap> = - HashMap::new(); - for (region_id, request) in requests { - let is_putting_physical_region = - self.state.read().unwrap().exist_physical_region(region_id); - if is_putting_physical_region { - FORBIDDEN_OPERATION_COUNT.inc(); - return ForbiddenPhysicalWriteSnafu.fail(); - } - let physical_region_id = self.find_physical_region_id(region_id)?; - requests_per_physical - .entry(physical_region_id) - .or_default() - .push((region_id, request)); + for (physical_region_id, group) in groups { + let result = self + .put_regions_batch_single_physical(physical_region_id, group.requests) + .await; + results.push(PutGroupResult { + request_indices: group.request_indices, + result, + }); } + results + } - let mut total_affected_rows: AffectedRows = 0; - for (physical_region_id, requests) in requests_per_physical { - let affected_rows = self - .put_regions_batch_single_physical(physical_region_id, requests) - .await?; - total_affected_rows += affected_rows; + fn group_put_request( + &self, + groups: &mut HashMap, + index: usize, + region_id: RegionId, + request: RegionPutRequest, + ) -> Result<()> { + if self.state.read().unwrap().exist_physical_region(region_id) { + FORBIDDEN_OPERATION_COUNT.inc(); + return ForbiddenPhysicalWriteSnafu.fail(); } - - Ok(total_affected_rows) + let physical_region_id = self.find_physical_region_id(region_id)?; + let group = groups.entry(physical_region_id).or_default(); + group.request_indices.push(index); + group.requests.push((region_id, request)); + Ok(()) } /// Write a batch of requests that all belong to the same physical region. @@ -1001,7 +1003,9 @@ mod tests { }, ) }); - let affected_rows = engine.inner.put_regions_batch(requests).await.unwrap(); + let results = engine.put_regions_batch(requests).await; + assert_eq!(results.len(), 1); + let affected_rows = results.into_iter().next().unwrap().result.unwrap(); assert_eq!(affected_rows, 6); assert_eq!( scan_timestamp_values(&engine, logical_region_id).await, @@ -1376,12 +1380,13 @@ mod tests { for ((_, request), skip_wal) in mixed_requests.iter_mut().zip(policies) { request.skip_wal = skip_wal; } - let err = env + let results = env .metric() - .inner .put_regions_batch(mixed_requests.into_iter()) - .await - .unwrap_err(); + .await; + assert_eq!(results.len(), 1); + assert_eq!(results[0].request_indices, vec![0, 1]); + let err = results.into_iter().next().unwrap().result.unwrap_err(); assert!(err.to_string().contains("inconsistent WAL policy in batch")); for logical_region_id in [logical_region_1, logical_region_2] { assert!( @@ -1392,12 +1397,12 @@ mod tests { } } - let affected_rows = env + let results = env .metric() - .inner .put_regions_batch(build_requests().into_iter()) - .await - .unwrap(); + .await; + assert_eq!(results.len(), 1); + let affected_rows = results.into_iter().next().unwrap().result.unwrap(); assert_eq!(affected_rows, 5); let request = ScanRequest::default(); @@ -1736,11 +1741,9 @@ mod tests { ]; // Batch write - let affected_rows = engine - .inner - .put_regions_batch(requests.into_iter()) - .await - .unwrap(); + let results = engine.put_regions_batch(requests.into_iter()).await; + assert_eq!(results.len(), 1); + let affected_rows = results.into_iter().next().unwrap().result.unwrap(); assert_eq!(affected_rows, 10); // Verify physical region contains data from all logical regions @@ -1760,72 +1763,70 @@ mod tests { async fn test_batch_write_with_partial_failure() { let env = TestEnv::new().await; env.init_metric_region().await; - let engine = env.metric(); - let physical_region_id = env.default_physical_region_id(); - let logical_region_1 = env.default_logical_region_id(); - let logical_region_2 = RegionId::new(1024, 1); - let nonexistent_region = RegionId::new(9999, 9999); - - env.create_logical_region(physical_region_id, logical_region_2) + let physical = env.default_physical_region_id(); + let logical = env.default_logical_region_id(); + let other_physical = RegionId::new(4096, 0); + let other_logical = RegionId::new(4097, 0); + env.create_physical_region(other_physical, "batch_results", vec![]) .await; - - // Prepare batch with one invalid region - let schema = test_util::row_schema_with_tags(&["job"]); - let requests = vec![ - ( - logical_region_1, - RegionPutRequest { - skip_wal: false, - rows: Rows { - schema: schema.clone(), - rows: test_util::build_rows(1, 3), - }, - hint: None, - partition_expr_version: None, + env.create_logical_region(other_physical, other_logical) + .await; + let request = |index: i64, version| { + let mut rows = test_util::build_rows(1, 1); + rows[0].values[0].value_data = + Some(api::v1::value::ValueData::TimestampMillisecondValue(index)); + RegionPutRequest { + rows: Rows { + schema: test_util::row_schema_with_tags(&["job"]), + rows, }, - ), - ( - nonexistent_region, - RegionPutRequest { - skip_wal: false, - rows: Rows { - schema: schema.clone(), - rows: test_util::build_rows(1, 2), - }, - hint: None, - partition_expr_version: None, - }, - ), - ( - logical_region_2, - RegionPutRequest { - skip_wal: false, - rows: Rows { - schema: schema.clone(), - rows: test_util::build_rows(1, 5), - }, - hint: None, - partition_expr_version: None, - }, - ), - ]; - - // Batch write - let result = engine.inner.put_regions_batch(requests.into_iter()).await; - assert!(result.is_err()); - - // Invalid region is detected before any write, so the physical region remains empty. - // Fail-fast is per physical-region group; cross-group partial success is possible. - let request = ScanRequest::default(); - let stream = env + skip_wal: false, + hint: None, + partition_expr_version: version, + } + }; + let mut results = env .metric() - .scan_to_stream(physical_region_id, request) - .await - .unwrap(); - let batches = RecordBatches::try_collect(stream).await.unwrap(); - - assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 0); + .put_regions_batch( + vec![ + (logical, request(0, None)), + (RegionId::new(9999, 0), request(1, None)), + (other_logical, request(2, Some(1))), + (other_logical, request(3, None)), + (logical, request(4, None)), + (physical, request(5, None)), + ] + .into_iter(), + ) + .await; + results.sort_by_key(|result| result.request_indices[0]); + assert_eq!( + results + .iter() + .map(|result| result.request_indices.clone()) + .collect::>(), + vec![vec![0, 4], vec![1], vec![2, 3], vec![5]] + ); + assert_eq!(results[0].result.as_ref().unwrap(), &2); + assert!(results[1].result.is_err()); + assert!(results[2].result.is_err()); + assert!(matches!( + results[3].result, + Err(crate::error::Error::ForbiddenPhysicalWrite { .. }) + )); + for (region, expected) in [(physical, 2), (other_physical, 0)] { + let stream = env + .metric() + .scan_to_stream(region, ScanRequest::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + assert_eq!( + batches.iter().map(|batch| batch.num_rows()).sum::(), + expected + ); + } } #[tokio::test] @@ -1849,11 +1850,10 @@ mod tests { }, )]; - let err = engine - .inner - .put_regions_batch(requests.into_iter()) - .await - .unwrap_err(); + let results = engine.put_regions_batch(requests.into_iter()).await; + assert_eq!(results.len(), 1); + assert_eq!(results[0].request_indices, vec![0]); + let err = results.into_iter().next().unwrap().result.unwrap_err(); assert!(matches!( err, @@ -1897,15 +1897,15 @@ mod tests { ), ]; - let err = engine - .inner - .put_regions_batch(requests.into_iter()) - .await - .unwrap_err(); - + let mut results = engine.put_regions_batch(requests.into_iter()).await; + results.sort_by_key(|group| group.request_indices[0]); + assert_eq!(results.len(), 2); + assert_eq!(results[0].request_indices, vec![0]); + assert_eq!(results[0].result.as_ref().unwrap(), &1); + assert_eq!(results[1].request_indices, vec![1]); assert!(matches!( - err, - crate::error::Error::ForbiddenPhysicalWrite { .. } + results[1].result, + Err(crate::error::Error::ForbiddenPhysicalWrite { .. }) )); } @@ -1932,11 +1932,10 @@ mod tests { }, )]; - let affected_rows = engine - .inner - .put_regions_batch(requests.into_iter()) - .await - .unwrap(); + let results = engine.put_regions_batch(requests.into_iter()).await; + assert_eq!(results.len(), 1); + assert_eq!(results[0].request_indices, vec![0]); + let affected_rows = results.into_iter().next().unwrap().result.unwrap(); assert_eq!(affected_rows, 5); } @@ -1946,15 +1945,12 @@ mod tests { env.init_metric_region().await; let engine = env.metric(); - // Empty batch should return zero affected rows - let requests = vec![]; - let affected_rows = engine - .inner - .put_regions_batch(requests.into_iter()) - .await - .unwrap(); - - assert_eq!(affected_rows, 0); + assert!( + engine + .put_regions_batch(vec![].into_iter()) + .await + .is_empty() + ); } #[tokio::test] diff --git a/src/operator/src/write_dispatcher.rs b/src/operator/src/write_dispatcher.rs index 9a042cfb4e7..d8a6ac9905e 100644 --- a/src/operator/src/write_dispatcher.rs +++ b/src/operator/src/write_dispatcher.rs @@ -203,10 +203,13 @@ impl WriteDispatcher { pub(crate) async fn send( &self, peer: PeerRef, - request: RegionRequest, + mut request: RegionRequest, histogram: Option<&Histogram>, ) -> WriteAttempt { if let Some(retry) = &self.retry { + if matches!(&request.body, Some(Body::Inserts(_) | Body::Deletes(_))) { + api::region::set_accept_partial_result(request.header.get_or_insert_default()); + } return retry .send( RoutedWrite { diff --git a/src/operator/src/write_dispatcher/tests.rs b/src/operator/src/write_dispatcher/tests.rs index 0f9af3b7190..a249a8dfe8e 100644 --- a/src/operator/src/write_dispatcher/tests.rs +++ b/src/operator/src/write_dispatcher/tests.rs @@ -356,6 +356,15 @@ async fn test_write_retry_disabled_and_enabled_ablation() { ); assert_eq!(result.is_ok(), max_retries > 0); assert_eq!(attempts, max_retries + 1); + for request in server.service.requests.lock().unwrap().iter() { + let accepts_partial = request + .header + .as_ref() + .and_then(|header| header.query_context.as_ref()) + .and_then(|context| context.extensions.get(api::region::ACCEPT_PARTIAL_RESULT)) + .is_some_and(|value| value == "true"); + assert_eq!(accepts_partial, max_retries > 0); + } assert!(server.invalidator.calls.lock().unwrap().is_empty()); } } @@ -1709,11 +1718,11 @@ async fn check_write_retry_rebuilds_logical_rows_in_mixed_table_envelope(kind: W ); let mut requests = server.service.requests.lock().unwrap().clone(); requests.sort_by_key(|request| write_region_id(request).unwrap().as_u64()); - assert_eq!( - requests, - vec![physical, rebuilt_upper, logical, rebuilt_lower], - "{kind:?}" - ); + let mut expected = vec![physical, rebuilt_upper, logical, rebuilt_lower]; + for request in &mut expected { + api::region::set_accept_partial_result(request.header.get_or_insert_default()); + } + assert_eq!(requests, expected, "{kind:?}"); let calls = invalidator.calls.lock().unwrap(); assert_eq!(calls.len(), 2, "{kind:?}"); assert_eq!(calls[0].len(), 2, "{kind:?}"); diff --git a/src/servers/src/error.rs b/src/servers/src/error.rs index 6e558c8234b..2e17179206e 100644 --- a/src/servers/src/error.rs +++ b/src/servers/src/error.rs @@ -1014,6 +1014,7 @@ pub fn status_code_to_http_status(status_code: &StatusCode) -> HttpStatusCode { StatusCode::Internal | StatusCode::Unexpected | StatusCode::IllegalState + | StatusCode::PartialFailure | StatusCode::Unknown | StatusCode::RuntimeResourcesExhausted | StatusCode::EngineExecuteQuery => HttpStatusCode::INTERNAL_SERVER_ERROR, diff --git a/src/servers/src/mysql/writer.rs b/src/servers/src/mysql/writer.rs index d9253c5bea7..049f925c650 100644 --- a/src/servers/src/mysql/writer.rs +++ b/src/servers/src/mysql/writer.rs @@ -512,7 +512,9 @@ pub fn create_mysql_column_def(schema: &SchemaRef) -> Result> { fn mysql_error_kind(status_code: &StatusCode) -> ErrorKind { match status_code { StatusCode::Success => ErrorKind::ER_YES, - StatusCode::Unknown | StatusCode::External => ErrorKind::ER_UNKNOWN_ERROR, + StatusCode::PartialFailure | StatusCode::Unknown | StatusCode::External => { + ErrorKind::ER_UNKNOWN_ERROR + } StatusCode::Unsupported => ErrorKind::ER_NOT_SUPPORTED_YET, StatusCode::Cancelled | StatusCode::DeadlineExceeded => ErrorKind::ER_QUERY_INTERRUPTED, StatusCode::RuntimeResourcesExhausted => ErrorKind::ER_OUT_OF_RESOURCES, diff --git a/src/servers/src/postgres/types/error.rs b/src/servers/src/postgres/types/error.rs index a76bb4362c4..00d1bbbb1c1 100644 --- a/src/servers/src/postgres/types/error.rs +++ b/src/servers/src/postgres/types/error.rs @@ -382,7 +382,8 @@ impl From for PgErrorCode { StatusCode::DeadlineExceeded => PgErrorCode::Ec57000, StatusCode::External => PgErrorCode::Ec58000, - StatusCode::Unknown + StatusCode::PartialFailure + | StatusCode::Unknown | StatusCode::Unexpected | StatusCode::IllegalState | StatusCode::Internal => PgErrorCode::EcXX000, diff --git a/src/standalone/Cargo.toml b/src/standalone/Cargo.toml index a36983203ee..d89366d2e7e 100644 --- a/src/standalone/Cargo.toml +++ b/src/standalone/Cargo.toml @@ -46,5 +46,7 @@ tokio.workspace = true url.workspace = true [dev-dependencies] +datanode = { workspace = true, features = ["testing"] } +mito2 = { workspace = true, features = ["test"] } temp-env.workspace = true toml.workspace = true diff --git a/src/standalone/src/datanode_manager.rs b/src/standalone/src/datanode_manager.rs index 14104eacc2f..6eb704ed03c 100644 --- a/src/standalone/src/datanode_manager.rs +++ b/src/standalone/src/datanode_manager.rs @@ -17,7 +17,7 @@ use std::sync::Arc; use api::region::RegionResponse; use api::v1::region::{RegionRequest, RegionResponse as RegionResponseV1}; use async_trait::async_trait; -use client::region::check_response_header; +use client::region; use common_error::ext::BoxedError; use common_meta::error::{self as meta_error, Result as MetaResult}; use common_meta::node_manager::{ @@ -96,10 +96,7 @@ impl Datanode for RegionInvoker { .await .map_err(BoxedError::new) .context(meta_error::ExternalSnafu)?; - check_response_header(&response.header) - .map_err(BoxedError::new) - .context(meta_error::ExternalSnafu)?; - Ok(RegionResponse::from_region_response(response)) + region::decode_region_response(response).map_err(Into::into) } async fn handle_query(&self, request: QueryRequest) -> MetaResult { @@ -121,3 +118,65 @@ impl Datanode for RegionInvoker { .context(meta_error::ExternalSnafu) } } + +#[cfg(test)] +mod tests { + use api::v1::Rows; + use api::v1::region::{InsertRequest, InsertRequests, region_request}; + use common_error::ext::ErrorExt; + use common_error::mock::MockError; + use common_error::status_code::StatusCode; + use datanode::tests::{MockRegionEngine, mock_region_server}; + use mito2::engine::MITO_ENGINE_NAME; + use mito2::test_util::CreateRequestBuilder; + use store_api::region_request::RegionRequest as StoreRegionRequest; + use store_api::storage::RegionId; + + use super::*; + + #[tokio::test] + async fn test_write_preserves_original_error() { + for code in [StatusCode::StorageUnavailable, StatusCode::InvalidArguments] { + let mut server = mock_region_server(); + let (engine, _receiver) = MockRegionEngine::with_mock_fn( + MITO_ENGINE_NAME, + Box::new(move |region_id, request| { + if matches!(request, StoreRegionRequest::Put(_)) { + Err(BoxedError::new(MockError::new(code))) + .context(datanode::error::HandleRegionRequestSnafu { region_id }) + } else { + Ok(0) + } + }), + ); + server.register_engine(engine); + let mut requests = Vec::new(); + for number in 0..2 { + let region_id = RegionId::new(1, number); + server + .handle_request( + region_id, + StoreRegionRequest::Create(CreateRequestBuilder::new().build()), + ) + .await + .unwrap(); + requests.push(InsertRequest { + region_id: region_id.as_u64(), + rows: Some(Rows::default()), + ..Default::default() + }); + } + let datanode: DatanodeRef = RegionInvoker::arc(server); + let error = datanode + .handle(RegionRequest { + header: Some(Default::default()), + body: Some(region_request::Body::Inserts(InsertRequests { requests })), + }) + .await + .unwrap_err(); + assert_eq!(code, error.status_code()); + assert_eq!(MockError::new(code).retry_hint(), error.retry_hint()); + assert!(!matches!(error, meta_error::Error::PartialRegion { .. })); + } + } +} diff --git a/src/store-api/src/region_request.rs b/src/store-api/src/region_request.rs index ebb2f7c3684..f0650f8d7dc 100644 --- a/src/store-api/src/region_request.rs +++ b/src/store-api/src/region_request.rs @@ -213,47 +213,55 @@ impl RegionRequest { } } +/// Converts one insert, treating absent rows as a no-op like vector inserts. +pub fn make_region_put( + request: api::v1::region::InsertRequest, +) -> Option<(RegionId, RegionPutRequest)> { + request.rows.map(|rows| { + ( + request.region_id.into(), + RegionPutRequest { + rows, + hint: None, + skip_wal: request.skip_wal, + partition_expr_version: request.partition_expr_version.map(|version| version.value), + }, + ) + }) +} + fn make_region_puts(inserts: InsertRequests) -> Result> { - let requests = inserts + Ok(inserts .requests .into_iter() - .filter_map(|r| { - let region_id = r.region_id.into(); - r.rows.map(|rows| { - ( - region_id, - RegionRequest::Put(RegionPutRequest { - rows, - hint: None, - skip_wal: r.skip_wal, - partition_expr_version: r.partition_expr_version.map(|v| v.value), - }), - ) - }) - }) - .collect(); - Ok(requests) + .filter_map(make_region_put) + .map(|(region_id, request)| (region_id, RegionRequest::Put(request))) + .collect()) +} + +/// Converts one delete, treating absent rows as a no-op like vector deletes. +pub fn make_region_delete( + request: api::v1::region::DeleteRequest, +) -> Option<(RegionId, RegionDeleteRequest)> { + request.rows.map(|rows| { + ( + request.region_id.into(), + RegionDeleteRequest { + rows, + hint: None, + partition_expr_version: request.partition_expr_version.map(|version| version.value), + }, + ) + }) } fn make_region_deletes(deletes: DeleteRequests) -> Result> { - let requests = deletes + Ok(deletes .requests .into_iter() - .filter_map(|r| { - let region_id = r.region_id.into(); - r.rows.map(|rows| { - ( - region_id, - RegionRequest::Delete(RegionDeleteRequest { - rows, - hint: None, - partition_expr_version: r.partition_expr_version.map(|v| v.value), - }), - ) - }) - }) - .collect(); - Ok(requests) + .filter_map(make_region_delete) + .map(|(region_id, request)| (region_id, RegionRequest::Delete(request))) + .collect()) } fn parse_region_create(create: CreateRequest) -> Result<(RegionId, RegionCreateRequest)> { diff --git a/tests/runner/src/client.rs b/tests/runner/src/client.rs index a595c877940..03f7f7ed1f5 100644 --- a/tests/runner/src/client.rs +++ b/tests/runner/src/client.rs @@ -169,6 +169,7 @@ impl MultiProtocolClient { result = ServerSnafu { code: status_code, msg, + retry_hint: e.retry_hint(), } .fail(); } From e6fa1a0b212116263e097d8ac15d8273977c7f42 Mon Sep 17 00:00:00 2001 From: Yingwen Date: Wed, 21 Oct 2026 06:07:58 +0000 Subject: [PATCH 058/239] fix(ci): increase release integration test stack size (#9817) Signed-off-by: evenyag --- .github/actions/build-linux-artifacts/action.yml | 1 + .github/actions/build-windows-artifacts/action.yml | 1 + Makefile | 1 + 3 files changed, 3 insertions(+) diff --git a/.github/actions/build-linux-artifacts/action.yml b/.github/actions/build-linux-artifacts/action.yml index c5212c28998..26028485061 100644 --- a/.github/actions/build-linux-artifacts/action.yml +++ b/.github/actions/build-linux-artifacts/action.yml @@ -57,6 +57,7 @@ runs: # Without GT_S3_*, every remote object-store case is skipped and the # release only ever exercises the local file backend. env: + RUST_MIN_STACK: 8388608 # 8 MiB, aligned with Rust CI tests. GT_S3_BUCKET: ${{ inputs.aws-ci-test-bucket }} GT_S3_REGION: ${{ inputs.aws-region }} GT_S3_ACCESS_KEY_ID: ${{ inputs.aws-access-key-id }} diff --git a/.github/actions/build-windows-artifacts/action.yml b/.github/actions/build-windows-artifacts/action.yml index 2c4ae1f7db7..e5da671bad8 100644 --- a/.github/actions/build-windows-artifacts/action.yml +++ b/.github/actions/build-windows-artifacts/action.yml @@ -48,6 +48,7 @@ runs: run: make test sqlness-test env: RUST_BACKTRACE: 1 + RUST_MIN_STACK: 8388608 # 8 MiB, aligned with Rust CI tests. SQLNESS_OPTS: "--preserve-state" - name: Upload sqlness logs diff --git a/Makefile b/Makefile index 7f92098262f..21c87cc3707 100644 --- a/Makefile +++ b/Makefile @@ -262,6 +262,7 @@ run-it-in-container: start-etcd ## Run integration tests in dev-builder. docker run --network=host \ -v ${PWD}:/greptimedb -v ${CARGO_REGISTRY_CACHE}:/root/.cargo/registry -v ${CARGO_GIT_CACHE}:/root/.cargo/git -v /tmp:/tmp \ -e GT_S3_BUCKET -e GT_S3_ACCESS_KEY_ID -e GT_S3_ACCESS_KEY -e GT_S3_REGION \ + -e RUST_MIN_STACK \ -w /greptimedb ${IMAGE_REGISTRY}/${IMAGE_NAMESPACE}/dev-builder-${BASE_IMAGE}:${DEV_BUILDER_IMAGE_TAG} \ make test sqlness-test BUILD_JOBS=${BUILD_JOBS} From 9f9b35c270885acdec1381711f489bdffb626f4d Mon Sep 17 00:00:00 2001 From: Chenglong Zheng Date: Wed, 21 Oct 2026 07:57:16 +0000 Subject: [PATCH 059/239] feat: support wait and timeout options for TRUNCATE TABLE (#9753) * feat: support wait and timeout options for TRUNCATE TABLE * fix: address review feedback for truncate DDL options * fix: address review feedback for truncate DDL options --- src/common/meta/src/ddl_manager.rs | 143 +++++++++-- src/common/meta/src/error.rs | 12 + src/common/meta/src/key.rs | 1 + src/common/meta/src/kv_backend/test_util.rs | 12 +- src/operator/src/statement.rs | 3 +- src/operator/src/statement/ddl.rs | 78 +++++- src/sql/src/parsers/truncate_parser.rs | 83 ++++++- src/sql/src/parsers/utils.rs | 23 +- src/sql/src/statements/truncate.rs | 47 +++- tests-integration/tests/main.rs | 1 + tests-integration/tests/truncate.rs | 225 ++++++++++++++++++ .../common/truncate/truncate.result | 28 ++- .../standalone/common/truncate/truncate.sql | 14 +- 13 files changed, 620 insertions(+), 50 deletions(-) create mode 100644 tests-integration/tests/truncate.rs diff --git a/src/common/meta/src/ddl_manager.rs b/src/common/meta/src/ddl_manager.rs index 7b490895a4b..05100341112 100644 --- a/src/common/meta/src/ddl_manager.rs +++ b/src/common/meta/src/ddl_manager.rs @@ -247,19 +247,21 @@ pub trait RepartitionProcedureFactory: Send + Sync { /// The options for DDL tasks. /// /// Note: These options may not be utilized by all procedures. -/// At present, they are specifically applied in `RepartitionProcedure`. +/// At present, they are specifically applied in `RepartitionProcedure` and truncate tasks. #[derive(Debug, Clone, Copy)] pub struct DdlOptions { /// The timeout will be passed to the procedure. /// /// Note: Each procedure may implement its own timeout handling mechanism. + /// For truncate tasks, the timeout bounds submission and waiting rather than procedure execution. + /// A timeout does not cancel a submitted procedure. pub timeout: Duration, /// The flag that controls whether to wait for the procedure to complete. /// /// If wait is `true`, the procedure will wait for completion(success or failure) and the result will be returned. /// Otherwise, the procedure will be submitted and return the [ProcedureId](common_procedure::ProcedureId) immediately. /// - /// Note: The value of `wait` is independent of the `timeout` option. If a procedure ignores the `timeout` and `wait` is set to true, the operation returns until the procedure completes. + /// Note: Except for truncate tasks, the value of `wait` is independent of the `timeout` option. If a procedure ignores the `timeout` and `wait` is set to true, the operation returns until the procedure completes. pub wait: bool, } @@ -454,13 +456,8 @@ impl DdlManager { .context(PersistRepartitionGcRequirementSnafu)?; let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); - if wait { - self.execute_procedure_and_wait(procedure_with_id).await - } else { - self.submit_procedure(procedure_with_id) - .await - .map(|p| (p, None)) - } + self.execute_procedure_with_options(procedure_with_id, DdlOptions { wait, timeout }) + .await } /// Submits and executes an alter table task. @@ -870,13 +867,14 @@ impl DdlManager { self.execute_procedure_and_wait(procedure_with_id).await } - /// Submits and executes a truncate table task. + /// Submits a truncate table task and optionally waits for completion. #[tracing::instrument(skip_all)] pub async fn submit_truncate_table_task( &self, truncate_table_task: TruncateTableTask, table_info_value: DeserializedValueWithBytes, procedure_context: ProcedureContext, + ddl_options: DdlOptions, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); let procedure = TruncateTableProcedure::new(truncate_table_task, table_info_value, context); @@ -884,7 +882,8 @@ impl DdlManager { let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); - self.execute_procedure_and_wait(procedure_with_id).await + self.execute_procedure_with_options(procedure_with_id, ddl_options) + .await } /// Submits and executes a comment on task. @@ -908,6 +907,24 @@ impl DdlManager { self.execute_procedure_and_wait(procedure_with_id).await } + /// Submits a procedure and optionally waits for completion. + /// + /// The truncate task handler bounds its whole request, including metadata reads. + /// The manager drives submitted procedures independently of this future. + async fn execute_procedure_with_options( + &self, + procedure_with_id: ProcedureWithId, + ddl_options: DdlOptions, + ) -> Result<(ProcedureId, Option)> { + if ddl_options.wait { + self.execute_procedure_and_wait(procedure_with_id).await + } else { + self.submit_procedure(procedure_with_id) + .await + .map(|id| (id, None)) + } + } + /// Executes a procedure and waits for the result. async fn execute_procedure_and_wait( &self, @@ -1001,7 +1018,13 @@ impl DdlManager { .await } TruncateTable(truncate_table_task) => { - handle_truncate_table_task(self, truncate_table_task, procedure_context).await + handle_truncate_table_task( + self, + truncate_table_task, + procedure_context, + ddl_options, + ) + .await } CreateLogicalTables(create_table_tasks) => { handle_create_logical_table_tasks(self, create_table_tasks, procedure_context) @@ -1084,6 +1107,33 @@ async fn handle_truncate_table_task( ddl_manager: &DdlManager, truncate_table_task: TruncateTableTask, procedure_context: ProcedureContext, + ddl_options: DdlOptions, +) -> Result { + // Share one timeout across metadata checks, submission and optional waiting. + // Dropping this future does not cancel an already submitted procedure. + tokio::time::timeout( + ddl_options.timeout, + handle_truncate_table_task_inner( + ddl_manager, + truncate_table_task, + procedure_context, + ddl_options, + ), + ) + .await + .map_err(|_| { + error::DdlTaskTimeoutSnafu { + timeout: ddl_options.timeout, + } + .build() + })? +} + +async fn handle_truncate_table_task_inner( + ddl_manager: &DdlManager, + truncate_table_task: TruncateTableTask, + procedure_context: ProcedureContext, + ddl_options: DdlOptions, ) -> Result { let table_id = truncate_table_task.table_id; let table_metadata_manager = &ddl_manager.table_metadata_manager(); @@ -1108,10 +1158,18 @@ async fn handle_truncate_table_task( ); let (id, _) = ddl_manager - .submit_truncate_table_task(truncate_table_task, table_info_value, procedure_context) + .submit_truncate_table_task( + truncate_table_task, + table_info_value, + procedure_context, + ddl_options, + ) .await?; - info!("Table: {table_id} is truncated via procedure_id {id:?}"); + info!( + "Submitted truncate for table: {table_id} via procedure_id {id:?}, wait: {}", + ddl_options.wait + ); Ok(SubmitDdlTaskResponse { key: id.to_string().into(), @@ -1621,9 +1679,7 @@ mod tests { #[cfg(feature = "enterprise")] use common_base::protocol::Channel; use common_error::ext::BoxedError; - #[cfg(feature = "enterprise")] use common_error::ext::ErrorExt; - #[cfg(feature = "enterprise")] use common_error::status_code::StatusCode; #[cfg(feature = "enterprise")] use common_event_recorder::{PersistentEventContext, ProcedureEventInput, TriggerReason}; @@ -1656,6 +1712,7 @@ mod tests { use crate::key::TableMetadataManager; use crate::key::flow::FlowMetadataManager; use crate::kv_backend::memory::MemoryKvBackend; + use crate::kv_backend::test_util::MockKvBackendBuilder; use crate::node_manager::{DatanodeManager, DatanodeRef, FlownodeManager, FlownodeRef}; use crate::peer::Peer; use crate::procedure_executor::ExecutorContext; @@ -2221,6 +2278,60 @@ mod tests { } } + #[tokio::test] + async fn test_truncate_timeout_covers_metadata_checks() { + let mut context = new_ddl_context(Arc::new(DummyDatanodeManager)); + let kv_backend = MockKvBackendBuilder::default() + .range_async_fn(Arc::new(|_| Box::pin(std::future::pending()))) + .build() + .unwrap(); + context.table_metadata_manager = Arc::new(TableMetadataManager::new(Arc::new(kv_backend))); + let ddl_manager = ddl_manager_with_context(context).await; + + for wait in [true, false] { + let timeout = Duration::from_millis(20); + let request = SubmitDdlTaskRequest { + task: DdlTask::new_truncate_table( + "greptime".to_string(), + "public".to_string(), + "metrics".to_string(), + 1024, + Vec::new(), + ), + wait, + timeout, + }; + let err = tokio::time::timeout( + Duration::from_secs(1), + ddl_manager.submit_ddl_task( + ExecutorContext { + query_context: Some(QueryContext::default()), + ..Default::default() + }, + request, + ), + ) + .await + .expect("metadata reads must be bounded by the request timeout") + .unwrap_err(); + assert!(matches!( + err, + crate::error::Error::DdlTaskTimeout { timeout: actual, .. } if actual == timeout + )); + assert_eq!(err.status_code(), StatusCode::Cancelled); + assert!( + ddl_manager + .procedure_manager + .list_procedures() + .await + .unwrap() + .is_empty() + ); + } + + ddl_manager.procedure_manager.stop().await.unwrap(); + } + async fn ddl_manager_with_context(ddl_context: DdlContext) -> DdlManager { let kv_backend = Arc::new(MemoryKvBackend::new()); let state_store = Arc::new(KvStateStore::new(kv_backend.clone())); diff --git a/src/common/meta/src/error.rs b/src/common/meta/src/error.rs index 3b8f55e8a4c..c048f499dff 100644 --- a/src/common/meta/src/error.rs +++ b/src/common/meta/src/error.rs @@ -209,6 +209,15 @@ pub enum Error { source: common_procedure::Error, }, + #[snafu(display( + "DDL request timed out after {timeout:?}; a submitted procedure may still be running" + ))] + DdlTaskTimeout { + timeout: std::time::Duration, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Failed to start procedure manager"))] StartProcedureManager { #[snafu(implicit)] @@ -1343,6 +1352,9 @@ impl ErrorExt for Error { | InvalidRole { .. } | EmptyDdlTasks { .. } => StatusCode::InvalidArguments, + // DeadlineExceeded is treated as unreachable by MetaClient and + // retried. A wait timeout must not resubmit a running procedure. + DdlTaskTimeout { .. } => StatusCode::Cancelled, LoadTlsCertificate { .. } => StatusCode::Internal, #[cfg(feature = "pg_kvbackend")] diff --git a/src/common/meta/src/key.rs b/src/common/meta/src/key.rs index 644ba472233..1ed8f8df6f0 100644 --- a/src/common/meta/src/key.rs +++ b/src/common/meta/src/key.rs @@ -3318,6 +3318,7 @@ mod tests { Ok(RangeResponse { kvs, more: false }) }) }), + range_async_fn: None, batch_get_fn: Some({ let batch_get_calls = batch_get_calls.clone(); Arc::new(move |req| { diff --git a/src/common/meta/src/kv_backend/test_util.rs b/src/common/meta/src/kv_backend/test_util.rs index ce502c33321..589c6c699ef 100644 --- a/src/common/meta/src/kv_backend/test_util.rs +++ b/src/common/meta/src/kv_backend/test_util.rs @@ -16,6 +16,7 @@ use std::any::Any; use std::sync::Arc; use derive_builder::Builder; +use futures::future::BoxFuture; use crate::error::Result; use crate::kv_backend::txn::{Txn, TxnResponse}; @@ -27,11 +28,18 @@ use crate::kv_backend::{ pub type MockFn = Arc Result + Send + Sync>; +/// An asynchronous callback for mock backend operations. +pub type AsyncMockFn = + Arc BoxFuture<'static, Result> + Send + Sync>; + /// A mock kv backend for testing. #[derive(Builder)] pub struct MockKvBackend { #[builder(setter(strip_option), default)] pub range_fn: Option>, + /// Takes precedence over the synchronous range callback when set. + #[builder(setter(strip_option), default)] + pub range_async_fn: Option>, #[builder(setter(strip_option), default)] pub put_fn: Option>, #[builder(setter(strip_option), default)] @@ -76,7 +84,9 @@ impl KvBackend for MockKvBackend { } async fn range(&self, req: RangeRequest) -> Result { - if let Some(f) = &self.range_fn { + if let Some(f) = &self.range_async_fn { + f(req).await + } else if let Some(f) = &self.range_fn { f(req) } else { unimplemented!() diff --git a/src/operator/src/statement.rs b/src/operator/src/statement.rs index 6122a9f4281..b52cb814cbb 100644 --- a/src/operator/src/statement.rs +++ b/src/operator/src/statement.rs @@ -421,6 +421,7 @@ impl StatementExecutor { .await } Statement::TruncateTable(stmt) => { + let ddl_options = ddl::parse_ddl_options(stmt.options())?; let (catalog, schema, table) = table_idents_to_full_name(stmt.table_name(), &query_ctx) .map_err(BoxedError::new) @@ -429,7 +430,7 @@ impl StatementExecutor { let time_ranges = self .convert_truncate_time_ranges(&table_name, stmt.time_ranges(), &query_ctx) .await?; - self.truncate_table(table_name, time_ranges, query_ctx) + self.truncate_table_with_options(table_name, time_ranges, ddl_options, query_ctx) .await } Statement::CreateDatabase(stmt) => { diff --git a/src/operator/src/statement/ddl.rs b/src/operator/src/statement/ddl.rs index fa9da9f322e..19cc294d1bf 100644 --- a/src/operator/src/statement/ddl.rs +++ b/src/operator/src/statement/ddl.rs @@ -124,12 +124,22 @@ use crate::statement::StatementExecutor; use crate::statement::show::create_partitions_stmt; use crate::utils::{to_executor_context, to_executor_context_with_origin_frontend}; +/// Options shared by SQL DDL submission paths. #[derive(Debug, Clone, Copy)] -struct DdlSubmitOptions { +pub(crate) struct DdlSubmitOptions { wait: bool, timeout: Duration, } +impl Default for DdlSubmitOptions { + fn default() -> Self { + Self { + wait: SubmitDdlTaskRequest::default_wait(), + timeout: SubmitDdlTaskRequest::default_timeout(), + } + } +} + #[cfg(not(feature = "enterprise"))] const ALLOWED_FLOW_OPTIONS: [&str; 2] = [ DEFER_ON_MISSING_SOURCE_KEY, @@ -151,7 +161,9 @@ fn build_procedure_id_output(procedure_id: Vec) -> Result { Ok(Output::new_with_record_batches(batches)) } -fn parse_ddl_options(options: &OptionMap) -> Result { +/// Parses DDL submission options using the existing request defaults. +pub(crate) fn parse_ddl_options(options: &OptionMap) -> Result { + let defaults = DdlSubmitOptions::default(); let wait = match options.get(DDL_WAIT) { Some(value) => value.parse::().map_err(|_| { InvalidSqlSnafu { @@ -159,7 +171,7 @@ fn parse_ddl_options(options: &OptionMap) -> Result { } .build() })?, - None => SubmitDdlTaskRequest::default_wait(), + None => defaults.wait, }; let timeout = match options.get(DDL_TIMEOUT) { @@ -169,7 +181,7 @@ fn parse_ddl_options(options: &OptionMap) -> Result { } .build() })?, - None => SubmitDdlTaskRequest::default_timeout(), + None => defaults.timeout, }; Ok(DdlSubmitOptions { wait, timeout }) @@ -1717,12 +1729,29 @@ impl StatementExecutor { } } - #[tracing::instrument(skip_all)] pub async fn truncate_table( &self, table_name: TableName, time_ranges: Vec<(Timestamp, Timestamp)>, query_context: QueryContextRef, + ) -> Result { + self.truncate_table_with_options( + table_name, + time_ranges, + DdlSubmitOptions::default(), + query_context, + ) + .await + } + + /// Truncates a table, optionally waiting for the submitted procedure to finish. + #[tracing::instrument(name = "truncate_table", skip_all)] + pub(crate) async fn truncate_table_with_options( + &self, + table_name: TableName, + time_ranges: Vec<(Timestamp, Timestamp)>, + ddl_options: DdlSubmitOptions, + query_context: QueryContextRef, ) -> Result { ensure_table_writable(&table_name.schema_name, &table_name.table_name)?; @@ -1740,9 +1769,20 @@ impl StatementExecutor { table_name: table_name.to_string(), })?; let table_id = table.table_info().table_id(); - self.truncate_table_procedure(&table_name, table_id, time_ranges, query_context) + let response = self + .truncate_table_procedure( + &table_name, + table_id, + time_ranges, + ddl_options, + query_context, + ) .await?; + if !ddl_options.wait { + return build_procedure_id_output(response.key); + } + Ok(Output::new_with_affected_rows(0)) } @@ -2353,16 +2393,19 @@ impl StatementExecutor { table_name: &TableName, table_id: TableId, time_ranges: Vec<(Timestamp, Timestamp)>, + ddl_options: DdlSubmitOptions, query_context: QueryContextRef, ) -> Result { let executor_context = to_executor_context(query_context, TriggerReason::Manual); - let request = SubmitDdlTaskRequest::new(DdlTask::new_truncate_table( + let mut request = SubmitDdlTaskRequest::new(DdlTask::new_truncate_table( table_name.catalog_name.clone(), table_name.schema_name.clone(), table_name.table_name.clone(), table_id, time_ranges, )); + request.wait = ddl_options.wait; + request.timeout = ddl_options.timeout; self.procedure_executor .submit_ddl_task(executor_context, request) @@ -3272,6 +3315,10 @@ mod test { #[test] fn test_parse_ddl_options() { + let defaults = parse_ddl_options(&OptionMap::default()).unwrap(); + assert!(defaults.wait); + assert_eq!(defaults.timeout, SubmitDdlTaskRequest::default_timeout()); + let options = OptionMap::from([ ("timeout".to_string(), "5m".to_string()), ("wait".to_string(), "false".to_string()), @@ -3279,6 +3326,23 @@ mod test { let ddl_options = parse_ddl_options(&options).unwrap(); assert!(!ddl_options.wait); assert_eq!(Duration::from_secs(300), ddl_options.timeout); + + let options = OptionMap::from([("timeout".to_string(), "1500ms".to_string())]); + assert_eq!( + parse_ddl_options(&options).unwrap().timeout, + Duration::from_millis(1500) + ); + for (key, value) in [ + ("wait", "yes"), + ("timeout", "not-a-duration"), + ("timeout", "-1s"), + ] { + let options = OptionMap::from([(key.to_string(), value.to_string())]); + assert!(matches!( + parse_ddl_options(&options), + Err(error::Error::InvalidSql { .. }) + )); + } } #[test] diff --git a/src/sql/src/parsers/truncate_parser.rs b/src/sql/src/parsers/truncate_parser.rs index 1fe04865751..a5e33ef486a 100644 --- a/src/sql/src/parsers/truncate_parser.rs +++ b/src/sql/src/parsers/truncate_parser.rs @@ -18,10 +18,11 @@ use sqlparser::tokenizer::Token; use crate::error::{self, InvalidSqlSnafu, InvalidTableNameSnafu, Result, UnexpectedTokenSnafu}; use crate::parser::ParserContext; +use crate::parsers::utils; use crate::statements::statement::Statement; use crate::statements::truncate::TruncateTable; -/// `TRUNCATE [TABLE] table_name;` +/// `TRUNCATE [TABLE] table_name [FILE RANGE ...] [WITH (wait = ..., timeout = ...)];` impl ParserContext<'_> { pub(crate) fn parse_truncate(&mut self) -> Result { let _ = self.parser.next_token(); @@ -46,7 +47,10 @@ impl ParserContext<'_> { // if no range is specified, we just truncate the table if !have_range { - return Ok(Statement::TruncateTable(TruncateTable::new(table_ident))); + let options = utils::parse_ddl_options(&mut self.parser)?; + return Ok(Statement::TruncateTable( + TruncateTable::new(table_ident).with_options(options), + )); } // parse a list of time ranges consist of (Timestamp, Timestamp),? @@ -100,6 +104,7 @@ impl ParserContext<'_> { let peek = self.parser.peek_token().token; match peek { + Token::Word(word) if word.keyword == Keyword::WITH => break, sqlparser::tokenizer::Token::EOF | Token::SemiColon => { if time_ranges.is_empty() { return Err(InvalidSqlSnafu { @@ -112,24 +117,26 @@ impl ParserContext<'_> { Token::Comma => { self.parser.next_token(); // Consume the comma let next_peek = self.parser.peek_token().token; // Peek the token after the comma - if matches!(next_peek, Token::EOF | Token::SemiColon) { + if matches!(next_peek, Token::EOF | Token::SemiColon) + || matches!(next_peek, Token::Word(word) if word.keyword == Keyword::WITH) + { break; // Trailing comma, end of statement } // Otherwise, continue to parse next range continue; } _ => UnexpectedTokenSnafu { - expected: "a comma or end of statement", + expected: "a comma, WITH or end of statement", actual: self.peek_token_as_string(), } .fail()?, } } - Ok(Statement::TruncateTable(TruncateTable::new_with_ranges( - table_ident, - time_ranges, - ))) + let options = utils::parse_ddl_options(&mut self.parser)?; + Ok(Statement::TruncateTable( + TruncateTable::new_with_ranges(table_ident, time_ranges).with_options(options), + )) } } @@ -141,6 +148,64 @@ mod tests { use crate::dialect::GreptimeDbDialect; use crate::parser::ParseOptions; + #[test] + fn test_parse_truncate_ddl_options() { + for prefix in [ + "TRUNCATE foo", + "TRUNCATE TABLE foo", + "TRUNCATE TABLE foo FILE RANGE (0, 20)", + "TRUNCATE TABLE foo FILE RANGE (0, 20), (30, 40),", + ] { + let sql = format!("{prefix} WITH (WAIT = false, TIMEOUT = '30s');"); + let statements = ParserContext::create_with_dialect( + &sql, + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .unwrap(); + let Statement::TruncateTable(truncate) = &statements[0] else { + unreachable!() + }; + assert_eq!(truncate.options().get("wait"), Some("false")); + assert_eq!(truncate.options().get("timeout"), Some("30s")); + + let formatted = truncate.to_string(); + let reparsed = ParserContext::create_with_dialect( + &formatted, + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .unwrap(); + let Statement::TruncateTable(reparsed) = &reparsed[0] else { + unreachable!() + }; + assert_eq!(truncate.table_name(), reparsed.table_name()); + assert_eq!(truncate.time_ranges(), reparsed.time_ranges()); + assert_eq!( + truncate.options().to_str_map(), + reparsed.options().to_str_map() + ); + } + + for options in [ + "ttl = '1d'", + "unknown = true", + "wait = NULL", + "timeout = ['30s']", + ] { + let sql = format!("TRUNCATE TABLE foo WITH ({options})"); + assert!( + ParserContext::create_with_dialect( + &sql, + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .is_err(), + "{sql}" + ); + } + } + #[test] pub fn test_parse_truncate_with_ranges() { let sql = r#"TRUNCATE foo FILE RANGE (0, 20)"#; @@ -451,7 +516,7 @@ mod tests { assert!( result.is_err() && format!("{result:?}") - .contains("expected: 'a comma or end of statement', found: ("), + .contains("expected: 'a comma, WITH or end of statement', found: ("), "result is: {result:?}" ); diff --git a/src/sql/src/parsers/utils.rs b/src/sql/src/parsers/utils.rs index 01b136a997a..093bff02a0d 100644 --- a/src/sql/src/parsers/utils.rs +++ b/src/sql/src/parsers/utils.rs @@ -39,7 +39,9 @@ use snafu::{ResultExt, ensure}; use sqlparser::dialect::Dialect; use sqlparser::keywords::Keyword; use sqlparser::parser::Parser; -use table::requests::{SEMANTIC_PREFIX, validate_semantic_option, validate_table_option}; +use table::requests::{ + SEMANTIC_PREFIX, VALID_DDL_OPTION_KEYS, validate_semantic_option, validate_table_option, +}; use crate::error::{ ConvertToLogicalExpressionSnafu, InvalidSqlSnafu, InvalidTableOptionSnafu, ParseSqlValueSnafu, @@ -423,15 +425,32 @@ pub fn convert_month_day_nano_to_duration( Ok(std::time::Duration::new(adjusted_seconds, nanos_remainder)) } +/// Parses submission options, accepting only scalar `wait` and `timeout` values. +pub(crate) fn parse_ddl_options(parser: &mut Parser) -> Result { + parse_with_options_inner(parser, true) +} + pub fn parse_with_options(parser: &mut Parser) -> Result { + parse_with_options_inner(parser, false) +} + +fn parse_with_options_inner(parser: &mut Parser, ddl_only: bool) -> Result { let options = parser .parse_options(Keyword::WITH) .context(SyntaxSnafu)? .into_iter() .map(parse_option_string) .collect::>>()?; + // Validate before OptionMap can discard non-scalar secret options. for (key, value) in &options { - if key.starts_with(SEMANTIC_PREFIX) { + if ddl_only { + ensure!( + VALID_DDL_OPTION_KEYS.contains(&key.as_str()) && value.as_string().is_some(), + InvalidSqlSnafu { + msg: format!("invalid DDL option '{key}': '{value}'") + } + ); + } else if key.starts_with(SEMANTIC_PREFIX) { // Semantic keys are whitelisted and value-checked against their domain, // so a user cannot set an unknown key or an out-of-range value. let value = value.as_string().unwrap_or_default(); diff --git a/src/sql/src/statements/truncate.rs b/src/sql/src/statements/truncate.rs index 307326d9c71..9e5ae5b1814 100644 --- a/src/sql/src/statements/truncate.rs +++ b/src/sql/src/statements/truncate.rs @@ -18,11 +18,15 @@ use itertools::Itertools; use serde::Serialize; use sqlparser::ast::{ObjectName, Visit, VisitMut, Visitor, VisitorMut}; +use crate::statements::OptionMap; + /// TRUNCATE TABLE statement. #[derive(Debug, Clone, PartialEq, Eq, Serialize)] pub struct TruncateTable { table_name: ObjectName, time_ranges: Vec<(sqlparser::ast::Value, sqlparser::ast::Value)>, + /// Options for this DDL submission, never persisted as table options. + options: OptionMap, } impl Visit for TruncateTable { @@ -32,6 +36,7 @@ impl Visit for TruncateTable { start.visit(visitor)?; end.visit(visitor)?; } + Visit::visit(&self.options, visitor)?; ::std::ops::ControlFlow::Continue(()) } } @@ -43,6 +48,7 @@ impl VisitMut for TruncateTable { start.visit(visitor)?; end.visit(visitor)?; } + VisitMut::visit(&mut self.options, visitor)?; ::std::ops::ControlFlow::Continue(()) } } @@ -53,6 +59,7 @@ impl TruncateTable { Self { table_name, time_ranges: Vec::new(), + options: OptionMap::default(), } } @@ -64,9 +71,21 @@ impl TruncateTable { Self { table_name, time_ranges, + options: OptionMap::default(), } } + /// Sets the options for this DDL submission. + pub fn with_options(mut self, options: OptionMap) -> Self { + self.options = options; + self + } + + /// Returns the options for this DDL submission. + pub fn options(&self) -> &OptionMap { + &self.options + } + pub fn table_name(&self) -> &ObjectName { &self.table_name } @@ -80,19 +99,23 @@ impl Display for TruncateTable { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { let table_name = self.table_name(); write!(f, r#"TRUNCATE TABLE {table_name}"#)?; - if self.time_ranges.is_empty() { - return Ok(()); + if !self.time_ranges.is_empty() { + write!(f, " FILE RANGE ")?; + write!( + f, + "{}", + self.time_ranges + .iter() + .map(|(start, end)| format!("({}, {})", start, end)) + .join(", ") + )?; } - - write!(f, " FILE RANGE ")?; - write!( - f, - "{}", - self.time_ranges - .iter() - .map(|(start, end)| format!("({}, {})", start, end)) - .join(", ") - ) + if !self.options.is_empty() { + f.write_str(" WITH (")?; + self.options.write_sql(f, ", ", "", false)?; + f.write_str(")")?; + } + Ok(()) } } diff --git a/tests-integration/tests/main.rs b/tests-integration/tests/main.rs index 72d0a1850bd..714c438344d 100644 --- a/tests-integration/tests/main.rs +++ b/tests-integration/tests/main.rs @@ -32,6 +32,7 @@ mod sql; #[macro_use] mod region_migration; mod table_ddl_event; +mod truncate; #[macro_use] mod repartition; mod repartition_event; diff --git a/tests-integration/tests/truncate.rs b/tests-integration/tests/truncate.rs new file mode 100644 index 00000000000..c1e7735b633 --- /dev/null +++ b/tests-integration/tests/truncate.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. + +use std::sync::Arc; +use std::time::Duration; + +use client::OutputData; +use common_error::ext::ErrorExt; +use common_error::status_code::StatusCode; +use common_meta::lock_key::TableLock; +use common_procedure::{ + LockKey, Procedure, ProcedureId, ProcedureState, ProcedureWithId, Status, watcher, +}; +use datatypes::arrow::array::AsArray; +use tests_integration::cluster::GreptimeDbClusterBuilder; +use tests_integration::standalone::GreptimeDbStandaloneBuilder; +use tests_integration::test_util::{execute_sql, execute_sql_and_expect, try_execute_sql}; +use tokio::sync::Notify; + +/// Holds the table lock until explicitly released, keeping truncate pending. +struct HoldTruncateTableLock { + table_id: u32, + started: Arc, + release: Arc, +} + +#[async_trait::async_trait] +impl Procedure for HoldTruncateTableLock { + fn type_name(&self) -> &str { + "test::HoldTruncateTableLock" + } + + async fn execute(&mut self, _: &common_procedure::Context) -> common_procedure::Result { + self.started.notify_one(); + self.release.notified().await; + Ok(Status::done()) + } + + fn dump(&self) -> common_procedure::Result { + Ok(String::new()) + } + + fn lock_key(&self) -> LockKey { + LockKey::new(vec![TableLock::Write(self.table_id).into()]) + } +} + +#[rstest_reuse::apply(crate::both_deployment_cases)] +async fn test_truncate_wait_and_timeout(distributed: bool) { + let cluster = if distributed { + Some( + GreptimeDbClusterBuilder::new("truncate_wait_and_timeout") + .await + .with_datanodes(1) + .build(false) + .await, + ) + } else { + None + }; + let standalone = if distributed { + None + } else { + Some( + GreptimeDbStandaloneBuilder::new("truncate_wait_and_timeout") + .build() + .await, + ) + }; + let (frontend, manager) = if let Some(cluster) = &cluster { + ( + cluster.fe_instance().clone(), + cluster.metasrv.procedure_manager().clone(), + ) + } else { + let standalone = standalone.as_ref().unwrap(); + ( + standalone.fe_instance().clone(), + standalone.procedure_manager.clone(), + ) + }; + + execute_sql( + &frontend, + "CREATE TABLE truncate_wait_test (v INT, ts TIMESTAMP TIME INDEX)", + ) + .await; + execute_sql(&frontend, "INSERT INTO truncate_wait_test VALUES (1, 1)").await; + let definition = execute_sql(&frontend, "SHOW CREATE TABLE truncate_wait_test") + .await + .data + .pretty_print() + .await; + let table_id = frontend + .catalog_manager() + .table("greptime", "public", "truncate_wait_test", None) + .await + .unwrap() + .unwrap() + .table_info() + .table_id(); + let started = Arc::new(Notify::new()); + let release = Arc::new(Notify::new()); + let mut blocker = manager + .submit(ProcedureWithId::with_random_id(Box::new( + HoldTruncateTableLock { + table_id, + started: started.clone(), + release: release.clone(), + }, + ))) + .await + .unwrap(); + tokio::time::timeout(Duration::from_secs(10), started.notified()) + .await + .unwrap(); + + let output = execute_sql( + &frontend, + "TRUNCATE TABLE truncate_wait_test WITH (wait=false, timeout='10s')", + ) + .await; + let OutputData::RecordBatches(batches) = output.data else { + panic!("expected Procedure ID result"); + }; + let batch = batches.iter().next().unwrap(); + let id = batch + .column_by_name("Procedure ID") + .unwrap() + .as_string::() + .value(0); + let id = ProcedureId::parse_str(id).unwrap(); + let state = execute_sql(&frontend, &format!("SELECT procedure_state('{id}')")) + .await + .data + .pretty_print() + .await; + assert!(state.contains(r#""status":"Running""#), "{state}"); + + // The same table lock holds a synchronous truncate past its wait deadline. + // Give metadata reads and submission a generous budget. Using the same budget + // also verifies that the async procedure keeps running after it expires. + let err = try_execute_sql( + &frontend, + "TRUNCATE TABLE truncate_wait_test WITH (wait=true, timeout='10s')", + ) + .await + .unwrap_err(); + if distributed { + assert!(matches!( + err.status_code(), + StatusCode::DeadlineExceeded | StatusCode::Cancelled + )); + } else { + assert_eq!(err.status_code(), StatusCode::Cancelled); + assert!(err.output_msg().contains("DDL request timed out"), "{err}"); + } + let pending = manager + .list_procedures() + .await + .unwrap() + .into_iter() + .filter(|procedure| procedure.type_name == "metasrv-procedure::TruncateTable") + .collect::>(); + assert_eq!(pending.len(), 2); + let mut watchers = pending + .into_iter() + .map(|procedure| { + assert!(matches!(procedure.state, ProcedureState::Running)); + manager.procedure_watcher(procedure.id).unwrap() + }) + .collect::>(); + execute_sql_and_expect( + &frontend, + "SELECT * FROM truncate_wait_test", + r#" ++---+-------------------------+ +| v | ts | ++---+-------------------------+ +| 1 | 1970-01-01T00:00:00.001 | ++---+-------------------------+ +"#, + ) + .await; + + release.notify_one(); + tokio::time::timeout(Duration::from_secs(10), async { + watcher::wait(&mut blocker).await.unwrap(); + for watcher in &mut watchers { + watcher::wait(watcher).await.unwrap(); + } + }) + .await + .unwrap(); + execute_sql_and_expect( + &frontend, + "SELECT * FROM truncate_wait_test", + r#" ++---+----+ +| v | ts | ++---+----+ ++---+----+ +"#, + ) + .await; + assert_eq!( + definition, + execute_sql(&frontend, "SHOW CREATE TABLE truncate_wait_test") + .await + .data + .pretty_print() + .await + ); +} diff --git a/tests/cases/standalone/common/truncate/truncate.result b/tests/cases/standalone/common/truncate/truncate.result index 27fcbd90265..36239ed80e5 100644 --- a/tests/cases/standalone/common/truncate/truncate.result +++ b/tests/cases/standalone/common/truncate/truncate.result @@ -126,7 +126,7 @@ SELECT ts, host, cpu, memory FROM monitor ORDER BY ts; | 2023-09-20T13:47:44 | host3 | 88.8 | 4096.0 | +---------------------+-------+------+--------+ -TRUNCATE monitor; +TRUNCATE monitor WITH (wait = true, timeout = '30s'); Affected Rows: 0 @@ -137,6 +137,32 @@ SELECT ts, host, cpu, memory FROM monitor ORDER BY ts; +----+------+-----+--------+ +----+------+-----+--------+ +TRUNCATE monitor WITH (wait = 'invalid'); + +Error: 1004(InvalidArguments), Invalid SQL, error: invalid DDL option 'wait': 'invalid' + +TRUNCATE monitor WITH (timeout = 'invalid'); + +Error: 1004(InvalidArguments), Invalid SQL, error: invalid DDL option 'timeout': 'invalid': expected number at 0 + +TRUNCATE monitor WITH (ttl = '1d'); + +Error: 2000(InvalidSyntax), Invalid SQL, error: invalid DDL option 'ttl': ''1d'' + +-- SQLNESS REPLACE ([0-9a-fA-F]{8}-[0-9a-fA-F]{4}-4[0-9a-fA-F]{3}-[89abAB][0-9a-fA-F]{3}-[0-9a-fA-F]{12}) PROC_ID +TRUNCATE monitor WITH (wait = false, timeout = '30s'); + ++--------------------------------------+ +| Procedure ID | ++--------------------------------------+ +| PROC_ID | ++--------------------------------------+ + +-- Keep a synchronous truncate in the same case to check the original result. +TRUNCATE monitor; + +Affected Rows: 0 + DROP TABLE monitor; Affected Rows: 0 diff --git a/tests/cases/standalone/common/truncate/truncate.sql b/tests/cases/standalone/common/truncate/truncate.sql index a116664f401..8ed86547e56 100644 --- a/tests/cases/standalone/common/truncate/truncate.sql +++ b/tests/cases/standalone/common/truncate/truncate.sql @@ -51,10 +51,22 @@ INSERT INTO monitor(ts, host, cpu, memory) VALUES SELECT ts, host, cpu, memory FROM monitor ORDER BY ts; -TRUNCATE monitor; +TRUNCATE monitor WITH (wait = true, timeout = '30s'); SELECT ts, host, cpu, memory FROM monitor ORDER BY ts; +TRUNCATE monitor WITH (wait = 'invalid'); + +TRUNCATE monitor WITH (timeout = 'invalid'); + +TRUNCATE monitor WITH (ttl = '1d'); + +-- SQLNESS REPLACE ([0-9a-fA-F]{8}-[0-9a-fA-F]{4}-4[0-9a-fA-F]{3}-[89abAB][0-9a-fA-F]{3}-[0-9a-fA-F]{12}) PROC_ID +TRUNCATE monitor WITH (wait = false, timeout = '30s'); + +-- Keep a synchronous truncate in the same case to check the original result. +TRUNCATE monitor; + DROP TABLE monitor; From 5af1e3301e5f5de2bd90158acab0aacb7f18488f Mon Sep 17 00:00:00 2001 From: LFC <990479+MichaelScofield@users.noreply.github.com> Date: Wed, 21 Oct 2026 11:34:47 +0000 Subject: [PATCH 060/239] fix(mysql): cap timestamp output at microsecond precision (#9831) * fix(mysql): cap timestamp output at microsecond precision Truncate MySQL text timestamp results to microseconds while preserving the existing second and millisecond formatting. Keep stored nanoseconds and explicit integer/string representations unchanged. Fixes #9828 Signed-off-by: luofucong * chore: ignore local output directory Signed-off-by: luofucong --------- Signed-off-by: luofucong --- .gitignore | 2 + src/servers/src/mysql/writer.rs | 5 +- src/servers/tests/mysql/mysql_server_test.rs | 66 ++++++++++++++++++- src/servers/tests/mysql/mysql_writer_test.rs | 9 +-- .../timestamp_precision_display.result | 31 ++++++--- .../timestamp/timestamp_precision_display.sql | 9 ++- .../cases/mysql_timestamp_precision/case.toml | 8 +++ .../cases/mysql_timestamp_precision/setup.sql | 3 + .../mysql_timestamp_precision/verify.result | 17 +++++ .../mysql_timestamp_precision/verify.sql | 5 ++ 10 files changed, 134 insertions(+), 21 deletions(-) create mode 100644 tests/compatibility/cases/mysql_timestamp_precision/case.toml create mode 100644 tests/compatibility/cases/mysql_timestamp_precision/setup.sql create mode 100644 tests/compatibility/cases/mysql_timestamp_precision/verify.result create mode 100644 tests/compatibility/cases/mysql_timestamp_precision/verify.sql diff --git a/.gitignore b/.gitignore index dd60d96224e..dba299862bb 100644 --- a/.gitignore +++ b/.gitignore @@ -79,3 +79,5 @@ docs/specs/ # Claude Code personal/local overrides (keep shared .claude/ content tracked) .claude/*.local.* + +output/ diff --git a/src/servers/src/mysql/writer.rs b/src/servers/src/mysql/writer.rs index 049f925c650..c485a5189fc 100644 --- a/src/servers/src/mysql/writer.rs +++ b/src/servers/src/mysql/writer.rs @@ -23,7 +23,7 @@ use arrow::datatypes::{ UInt16Type, UInt32Type, UInt64Type, }; use arrow_schema::{DataType, IntervalUnit}; -use chrono::{Datelike, NaiveDateTime}; +use chrono::{Datelike, NaiveDateTime, SubsecRound}; use common_decimal::Decimal128; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; @@ -275,7 +275,8 @@ impl MysqlResultWriter { write!( &mut staged_timestamp.formatted, "{}", - datetime.format("%Y-%m-%d %H:%M:%S%.f") + // MySQL temporal types support at most microsecond precision. + datetime.trunc_subsecs(6).format("%Y-%m-%d %H:%M:%S%.f") ) .map_err(|_| { InternalSnafu { diff --git a/src/servers/tests/mysql/mysql_server_test.rs b/src/servers/tests/mysql/mysql_server_test.rs index cba667f22b1..fdbdd4fd8f9 100644 --- a/src/servers/tests/mysql/mysql_server_test.rs +++ b/src/servers/tests/mysql/mysql_server_test.rs @@ -33,7 +33,9 @@ use datafusion_expr::LogicalPlan; use datatypes::prelude::VectorRef; use datatypes::schema::{ColumnSchema, Schema}; use datatypes::value::Value; -use datatypes::vectors::{Int32Vector, TimestampMicrosecondVector, TimestampSecondVector}; +use datatypes::vectors::{ + Int32Vector, TimestampMicrosecondVector, TimestampNanosecondVector, TimestampSecondVector, +}; use mysql_async::prelude::*; use mysql_async::{Conn, Row, SslOpts}; use query::QueryEngineFactory; @@ -859,6 +861,68 @@ async fn test_mysql_timestamp_precision_slot_isolation() -> Result<()> { Ok(()) } +#[tokio::test] +async fn test_mysql_nanosecond_timestamps_use_microsecond_precision() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "ts", + datatypes::prelude::ConcreteDataType::timestamp_nanosecond_datatype(), + true, + )])); + let timestamps = vec![ + Some(1_792_572_935_718_870_300), + Some(123_000_789), + Some(999_999_999), + Some(999), + Some(-1), + Some(0), + None, + ]; + let batch = RecordBatch::new( + schema, + vec![Arc::new(TimestampNanosecondVector::from(timestamps)) as VectorRef], + ) + .unwrap(); + let table = MemTable::table("nanosecond_timestamps", batch); + let query = "SELECT ts FROM nanosecond_timestamps".to_string(); + let (result, health_check) = query_mysql_text_protocol(table.clone(), query.clone()).await; + let rows = result.unwrap(); + let expected = [ + Some("2026-10-21 08:55:35.718870"), + Some("1970-01-01 00:00:00.123"), + Some("1970-01-01 00:00:00.999999"), + Some("1970-01-01 00:00:00"), + Some("1969-12-31 23:59:59.999999"), + Some("1970-01-01 00:00:00"), + None, + ]; + assert_eq!(expected.len(), rows.len()); + for (row, expected) in rows.iter().zip(expected) { + assert_eq!( + vec![expected.map_or(Value::Null, |s| Value::from(s.as_bytes().to_vec()))], + row.values + ); + } + assert_eq!(Some(1), health_check.unwrap()); + + let (result, health_check) = query_mysql_binary_protocol(table, query).await; + let rows = result.unwrap(); + let expected = [ + mysql_async::Value::Date(2026, 10, 21, 8, 55, 35, 718_870), + mysql_async::Value::Date(1970, 1, 1, 0, 0, 0, 123_000), + mysql_async::Value::Date(1970, 1, 1, 0, 0, 0, 999_999), + mysql_async::Value::Date(1970, 1, 1, 0, 0, 0, 0), + mysql_async::Value::Date(1969, 12, 31, 23, 59, 59, 999_999), + mysql_async::Value::Date(1970, 1, 1, 0, 0, 0, 0), + mysql_async::Value::NULL, + ]; + assert_eq!(expected.len(), rows.len()); + for (row, expected) in rows.iter().zip(expected) { + assert_eq!(Some(&expected), row.as_ref(0)); + } + assert_eq!(Some(1), health_check.unwrap()); + Ok(()) +} + fn timestamp_table(table_name: &str, timestamp: Option) -> (TableRef, Arc) { let schema = Arc::new(Schema::new(vec![ ColumnSchema::new( diff --git a/src/servers/tests/mysql/mysql_writer_test.rs b/src/servers/tests/mysql/mysql_writer_test.rs index 473a6300263..6cfcd09e746 100644 --- a/src/servers/tests/mysql/mysql_writer_test.rs +++ b/src/servers/tests/mysql/mysql_writer_test.rs @@ -264,13 +264,8 @@ fn test_create_mysql_column_unsigned_flags() { // ── Timestamp precision tests (issue #8227) ─────────────────────────────────── -/// Timestamps displayed through the MySQL writer must respect their storage unit. -/// -/// Before the fix, `to_chrono_datetime_with_timezone` → `NaiveDateTime::Display` -/// always emitted 6 fractional-second digits, so: -/// - `TIMESTAMP(3)` showed `.195000` instead of `.195` -/// - `TIMESTAMP(9)` showed `.195123` instead of `.195123456` -/// +/// The common timestamp string formatter preserves full storage precision. +/// The MySQL writer separately caps temporal output at microseconds (issue #9828). #[test] fn test_timestamp_display_precision_utc() { // Use UTC timezone for deterministic, timezone-independent output. diff --git a/tests/cases/standalone/common/types/timestamp/timestamp_precision_display.result b/tests/cases/standalone/common/types/timestamp/timestamp_precision_display.result index 45b930336f5..2c5f66da409 100644 --- a/tests/cases/standalone/common/types/timestamp/timestamp_precision_display.result +++ b/tests/cases/standalone/common/types/timestamp/timestamp_precision_display.result @@ -1,7 +1,7 @@ -- description: Timestamp display precision over the MySQL protocol should --- respect the column's storage unit (issue #8227): +-- respect MySQL's maximum microsecond precision (issues #8227 and #9828): -- - TIMESTAMP(3) must not pad to 6 fractional digits --- - TIMESTAMP(9) must not truncate to 6 fractional digits +-- - TIMESTAMP(9) must truncate to at most 6 fractional digits -- SQLNESS PROTOCOL MYSQL CREATE TABLE ts_precision_display ( ts_s TIMESTAMP(0), @@ -54,13 +54,26 @@ affected_rows: 1 -- SQLNESS PROTOCOL MYSQL SELECT * FROM ts_precision_display ORDER BY v; -+---------------------+-------------------------+----------------------------+-------------------------------+---+ -| ts_s | ts_ms | ts_us | ts_ns | v | -+---------------------+-------------------------+----------------------------+-------------------------------+---+ -| 2026-06-02 03:50:00 | 2026-06-02 03:50:00.195 | 2026-06-02 03:50:00.195123 | 2026-06-02 03:50:00.195123456 | 1 | -| 2026-06-02 03:50:01 | 2026-06-02 03:50:01.100 | 2026-06-02 03:50:01.100 | 2026-06-02 03:50:01.100 | 2 | -| 2026-06-02 03:50:02 | 2026-06-02 03:50:02 | 2026-06-02 03:50:02 | 2026-06-02 03:50:02 | 3 | -+---------------------+-------------------------+----------------------------+-------------------------------+---+ ++---------------------+-------------------------+----------------------------+----------------------------+---+ +| ts_s | ts_ms | ts_us | ts_ns | v | ++---------------------+-------------------------+----------------------------+----------------------------+---+ +| 2026-06-02 03:50:00 | 2026-06-02 03:50:00.195 | 2026-06-02 03:50:00.195123 | 2026-06-02 03:50:00.195123 | 1 | +| 2026-06-02 03:50:01 | 2026-06-02 03:50:01.100 | 2026-06-02 03:50:01.100 | 2026-06-02 03:50:01.100 | 2 | +| 2026-06-02 03:50:02 | 2026-06-02 03:50:02 | 2026-06-02 03:50:02 | 2026-06-02 03:50:02 | 3 | ++---------------------+-------------------------+----------------------------+----------------------------+---+ + +-- Full nanosecond precision remains available as an integer or string. +-- SQLNESS PROTOCOL MYSQL +SELECT ts_ns::BIGINT AS epoch_ns, ts_ns::STRING AS timestamp_string +FROM ts_precision_display ORDER BY v; + ++---------------------+-------------------------------+ +| epoch_ns | timestamp_string | ++---------------------+-------------------------------+ +| 1780372200195123456 | 2026-06-02T03:50:00.195123456 | +| 1780372201100000000 | 2026-06-02T03:50:01.100 | +| 1780372202000000000 | 2026-06-02T03:50:02 | ++---------------------+-------------------------------+ -- SQLNESS PROTOCOL MYSQL DROP TABLE ts_precision_display; diff --git a/tests/cases/standalone/common/types/timestamp/timestamp_precision_display.sql b/tests/cases/standalone/common/types/timestamp/timestamp_precision_display.sql index 8eb82ed711d..38fd0d6aa03 100644 --- a/tests/cases/standalone/common/types/timestamp/timestamp_precision_display.sql +++ b/tests/cases/standalone/common/types/timestamp/timestamp_precision_display.sql @@ -1,7 +1,7 @@ -- description: Timestamp display precision over the MySQL protocol should --- respect the column's storage unit (issue #8227): +-- respect MySQL's maximum microsecond precision (issues #8227 and #9828): -- - TIMESTAMP(3) must not pad to 6 fractional digits --- - TIMESTAMP(9) must not truncate to 6 fractional digits +-- - TIMESTAMP(9) must truncate to at most 6 fractional digits -- SQLNESS PROTOCOL MYSQL CREATE TABLE ts_precision_display ( @@ -47,5 +47,10 @@ INSERT INTO ts_precision_display VALUES ( -- SQLNESS PROTOCOL MYSQL SELECT * FROM ts_precision_display ORDER BY v; +-- Full nanosecond precision remains available as an integer or string. +-- SQLNESS PROTOCOL MYSQL +SELECT ts_ns::BIGINT AS epoch_ns, ts_ns::STRING AS timestamp_string +FROM ts_precision_display ORDER BY v; + -- SQLNESS PROTOCOL MYSQL DROP TABLE ts_precision_display; diff --git a/tests/compatibility/cases/mysql_timestamp_precision/case.toml b/tests/compatibility/cases/mysql_timestamp_precision/case.toml new file mode 100644 index 00000000000..4ddd00e664f --- /dev/null +++ b/tests/compatibility/cases/mysql_timestamp_precision/case.toml @@ -0,0 +1,8 @@ +name = "mysql_timestamp_precision" +reason = "MySQL timestamp results must use at most microseconds while stored nanoseconds survive restart (issue #9828)." +introduced_by = "issue #9828" +topologies = ["standalone", "distributed"] +from_range = ["*"] +to_range = [">=v1.4.0-alpha.0"] +features = ["table", "mysql", "timestamp"] +owner = "servers" diff --git a/tests/compatibility/cases/mysql_timestamp_precision/setup.sql b/tests/compatibility/cases/mysql_timestamp_precision/setup.sql new file mode 100644 index 00000000000..b9a453b0bf4 --- /dev/null +++ b/tests/compatibility/cases/mysql_timestamp_precision/setup.sql @@ -0,0 +1,3 @@ +CREATE TABLE timestamps (ts TIMESTAMP(9) TIME INDEX); + +INSERT INTO timestamps VALUES ('2026-10-21 08:55:35.718870300'); diff --git a/tests/compatibility/cases/mysql_timestamp_precision/verify.result b/tests/compatibility/cases/mysql_timestamp_precision/verify.result new file mode 100644 index 00000000000..330d9c5e033 --- /dev/null +++ b/tests/compatibility/cases/mysql_timestamp_precision/verify.result @@ -0,0 +1,17 @@ +-- SQLNESS PROTOCOL MYSQL +SELECT * FROM timestamps; + ++----------------------------+ +| ts | ++----------------------------+ +| 2026-10-21 08:55:35.718870 | ++----------------------------+ + +-- SQLNESS PROTOCOL MYSQL +SELECT ts::BIGINT AS epoch_ns, ts::STRING AS timestamp_string FROM timestamps; + ++---------------------+-------------------------------+ +| epoch_ns | timestamp_string | ++---------------------+-------------------------------+ +| 1792572935718870300 | 2026-10-21T08:55:35.718870300 | ++---------------------+-------------------------------+ diff --git a/tests/compatibility/cases/mysql_timestamp_precision/verify.sql b/tests/compatibility/cases/mysql_timestamp_precision/verify.sql new file mode 100644 index 00000000000..8dc0214b50f --- /dev/null +++ b/tests/compatibility/cases/mysql_timestamp_precision/verify.sql @@ -0,0 +1,5 @@ +-- SQLNESS PROTOCOL MYSQL +SELECT * FROM timestamps; + +-- SQLNESS PROTOCOL MYSQL +SELECT ts::BIGINT AS epoch_ns, ts::STRING AS timestamp_string FROM timestamps; From 0dc649e0289fcbd7fd3d2ed502a0ca9c5dc65119 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Wed, 21 Oct 2026 11:55:44 +0000 Subject: [PATCH 061/239] fix(datanode): release region map guard before awaiting role changes (#9824) Signed-off-by: WenyXu --- src/datanode/src/region_server.rs | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index 15bf786c389..135d069fd77 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -486,7 +486,14 @@ impl RegionServer { region_id: RegionId, state: SettableRegionRoleState, ) -> Result { - match self.inner.region_map.get(®ion_id) { + // Cloning releases the DashMap guard before awaiting the engine operation. + let engine = self + .inner + .region_map + .get(®ion_id) + .map(|entry| entry.value().clone()); + + match engine { Some(engine) => Ok(engine .set_region_role_state_gracefully(region_id, state) .await From ec4ef45f4056c28f497f3d396c216dffa705945e Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Wed, 21 Oct 2026 13:37:34 +0000 Subject: [PATCH 062/239] feat: use pgwire auth mechanisms API to simplify startuphandler (#9634) * feat: use pgwire auth mechanisms API to simplify startuphandler * chore: update pgwire * chore: update pgwire revision to merged one * refactor: model based on upstream api * refactor: simplify implementation by introducing authsource * refactor: update to latest pgwire apis * refactor: adopt pgwire api changes * chore: switch to released pgwire * test: harden pg scram mixed-auth integration test and reject channel-binding clients - Fix MixedAuthProvider to return a SCRAM verifier for the denied user, so the test exercises rejection by authorization after successful SCRAM verification; record the selected mechanism from the returned variant and assert authorize ran exactly once per verified user. - Assert the denied connection fails with 28P01 before any terminal message, protecting the completion-ordering introduced by pgwire's authentication gate. - Add a raw-protocol test proving a client-first with a p= channel binding flag is rejected (RFC 5802 section 6), while n-flag clients (tokio_postgres) continue normally. Bump pgwire to 0.42.1, which enforces RFC 5802 channel binding negotiation and rejects short/overlong SCRAM proofs on the verifier path (sunng87/pgwire#476). Signed-off-by: Ning Sun --------- Signed-off-by: Ning Sun --- Cargo.lock | 12 +- Cargo.toml | 6 +- src/auth/src/common.rs | 12 + src/servers/Cargo.toml | 4 +- src/servers/src/postgres.rs | 22 +- src/servers/src/postgres/auth_handler.rs | 688 ++++++++--------------- src/servers/tests/postgres/mod.rs | 304 +++++++++- 7 files changed, 589 insertions(+), 459 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 7a8da222c03..4c5a0077193 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -663,8 +663,8 @@ dependencies = [ [[package]] name = "arrow-pg" -version = "0.15.0" -source = "git+https://github.com/datafusion-contrib/datafusion-postgres.git?rev=ffb14a52b3d7c7489559812e1ec38096f225ce93#ffb14a52b3d7c7489559812e1ec38096f225ce93" +version = "0.16.0" +source = "git+https://github.com/datafusion-contrib/datafusion-postgres.git?rev=3ae0a6cb30b1c5131687f4637247ff4f84ac11e4#3ae0a6cb30b1c5131687f4637247ff4f84ac11e4" dependencies = [ "arrow 59.2.0", "arrow-schema 59.2.0", @@ -4234,8 +4234,8 @@ dependencies = [ [[package]] name = "datafusion-pg-catalog" -version = "0.18.3" -source = "git+https://github.com/datafusion-contrib/datafusion-postgres.git?rev=ffb14a52b3d7c7489559812e1ec38096f225ce93#ffb14a52b3d7c7489559812e1ec38096f225ce93" +version = "0.19.0" +source = "git+https://github.com/datafusion-contrib/datafusion-postgres.git?rev=3ae0a6cb30b1c5131687f4637247ff4f84ac11e4#3ae0a6cb30b1c5131687f4637247ff4f84ac11e4" dependencies = [ "arrow-pg", "async-trait", @@ -10711,9 +10711,9 @@ dependencies = [ [[package]] name = "pgwire" -version = "0.41.0" +version = "0.42.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c667f8f4046889d52ed11e17aae954b1816d3df48dd23bdebcf199cef332df93" +checksum = "1452e6541bcc64e166ff9d350c25c09334c0c80e8f60fde5e36ff81768eeaa55" dependencies = [ "async-trait", "aws-lc-rs", diff --git a/Cargo.toml b/Cargo.toml index 20582d8eaa4..1674063f0c6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -141,7 +141,7 @@ datafusion-functions-aggregate-common = "=55.1.0" datafusion-functions-window-common = "=55.1.0" datafusion-optimizer = "=55.1.0" datafusion-orc = "0.10.0" -datafusion-pg-catalog = "0.18.3" +datafusion-pg-catalog = "0.19.0" datafusion-physical-expr = "=55.1.0" datafusion-physical-plan = "=55.1.0" datafusion-proto = "=55.1.0" @@ -382,8 +382,8 @@ datafusion-physical-plan = { git = "https://github.com/GreptimeTeam/datafusion.g datafusion-proto = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } datafusion-sql = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } datafusion-substrait = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } -datafusion-pg-catalog = { git = "https://github.com/datafusion-contrib/datafusion-postgres.git", rev = "ffb14a52b3d7c7489559812e1ec38096f225ce93" } -arrow-pg = { git = "https://github.com/datafusion-contrib/datafusion-postgres.git", rev = "ffb14a52b3d7c7489559812e1ec38096f225ce93" } +datafusion-pg-catalog = { git = "https://github.com/datafusion-contrib/datafusion-postgres.git", rev = "3ae0a6cb30b1c5131687f4637247ff4f84ac11e4" } +arrow-pg = { git = "https://github.com/datafusion-contrib/datafusion-postgres.git", rev = "3ae0a6cb30b1c5131687f4637247ff4f84ac11e4" } sqlparser = { git = "https://github.com/GreptimeTeam/sqlparser-rs.git", rev = "9e9019bb1c7040ed956f654e39378dd42ab17884" } # pprof 0.15.0 initializes framehop inside SIGPROF; pin the unreleased fix diff --git a/src/auth/src/common.rs b/src/auth/src/common.rs index d25390ed778..4b88cc60fc4 100644 --- a/src/auth/src/common.rs +++ b/src/auth/src/common.rs @@ -361,6 +361,18 @@ impl PgScramSha256Verifier { &self.salt } + /// `SHA256(ClientKey)`, used to verify the client proof in SCRAM. + pub fn stored_key(&self) -> &[u8] { + &self.stored_key + } + + /// `HMAC(SaltedPassword, "Server Key")`, used to compute the SCRAM + /// server signature. This is secret material and must not be exposed to + /// clients. + pub fn server_key(&self) -> &[u8] { + &self.server_key + } + pub fn verify_plain_password(&self, password: &[u8]) -> Result { let salted_password = pg_scram_sha256_salted_password(password, &self.salt, self.iterations)?; diff --git a/src/servers/Cargo.toml b/src/servers/Cargo.toml index 0be0e47eab7..802ee46f9d4 100644 --- a/src/servers/Cargo.toml +++ b/src/servers/Cargo.toml @@ -24,7 +24,7 @@ api.workspace = true arrow.workspace = true arrow-flight.workspace = true arrow-ipc.workspace = true -arrow-pg = "0.15" +arrow-pg = "0.16" arrow-schema.workspace = true async-trait.workspace = true auth.workspace = true @@ -98,7 +98,7 @@ otel-arrow-rust.workspace = true parking_lot.workspace = true partition.workspace = true pg_interval = "0.6" -pgwire = { version = "0.41", default-features = false, features = [ +pgwire = { version = "0.42.1", default-features = false, features = [ "server-api-aws-lc-rs", "pg-ext-types", ] } diff --git a/src/servers/src/postgres.rs b/src/servers/src/postgres.rs index 1cd5907157b..15aae48a494 100644 --- a/src/servers/src/postgres.rs +++ b/src/servers/src/postgres.rs @@ -29,17 +29,18 @@ pub(crate) const METADATA_SCHEMA: &str = "schema"; use std::collections::HashMap; use std::net::SocketAddr; use std::sync::Arc; +use std::sync::atomic::AtomicBool; -use ::auth::UserProviderRef; +use ::auth::{UserInfoRef, UserProviderRef}; use derive_builder::Builder; -use pgwire::api::auth::{ServerParameterProvider, StartupHandler}; +use pgwire::api::auth::{AuthDriver, ServerParameterProvider, StartupHandler}; use pgwire::api::query::{ExtendedQueryHandler, SimpleQueryHandler}; use pgwire::api::{ClientInfo, ErrorHandler, PgWireServerHandlers}; pub use server::PostgresServer; use session::Session; use session::context::Channel; -use self::auth_handler::PgLoginVerifier; +use self::auth_handler::{PgUserProviderSource, ZeroPidSecretKeyGenerator}; use self::handler::DefaultQueryParser; use crate::query_handler::sql::ServerSqlQueryHandlerRef; @@ -76,9 +77,11 @@ impl ServerParameterProvider for GreptimeDBStartupParameters { pub struct PostgresServerHandlerInner { query_handler: ServerSqlQueryHandlerRef, - login_verifier: PgLoginVerifier, + auth_source: Option>, force_tls: bool, - param_provider: Arc, + auth_driver: AuthDriver, UserInfoRef>, + /// SCRAM selected; failures get the uniform metric label. + is_scram_auth: AtomicBool, batching_enabled: bool, session: Arc, @@ -129,9 +132,14 @@ impl MakePostgresServerHandler { )); let handler = PostgresServerHandlerInner { query_handler: self.query_handler.clone(), - login_verifier: PgLoginVerifier::new(self.user_provider.clone()), + auth_source: self + .user_provider + .clone() + .map(|p| Arc::new(PgUserProviderSource::new(p))), force_tls: self.force_tls, - param_provider: self.param_provider.clone(), + auth_driver: AuthDriver::new(self.param_provider.clone()) + .with_pid_secret_key_generator(Arc::new(ZeroPidSecretKeyGenerator)), + is_scram_auth: AtomicBool::new(false), batching_enabled, session: session.clone(), diff --git a/src/servers/src/postgres/auth_handler.rs b/src/servers/src/postgres/auth_handler.rs index 1fa075a8e17..169ce594f5e 100644 --- a/src/servers/src/postgres/auth_handler.rs +++ b/src/servers/src/postgres/auth_handler.rs @@ -13,29 +13,31 @@ // limitations under the License. use std::fmt::Debug; +use std::sync::Arc; +use std::sync::atomic::Ordering; use ::auth::{ BEARER_TOKEN_USER, Identity, Password, PgAuthInfo, PgScramSha256Verifier, UserInfoRef, UserProviderRef, userinfo_by_name, }; use async_trait::async_trait; -use base64::Engine; -use base64::engine::general_purpose::STANDARD as BASE64; use common_catalog::parse_catalog_and_schema_from_db_string; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; use common_time::Timezone; use futures::{Sink, SinkExt}; use pgwire::api::auth::StartupHandler; -use pgwire::api::auth::sasl::SCRAM_SHA_256_METHOD; -use pgwire::api::{ClientInfo, PgWireConnectionState, auth}; +use pgwire::api::auth::cleartext::CleartextPasswordAuth; +use pgwire::api::auth::mechanism::AuthMechanismFactory; +use pgwire::api::auth::sasl::scram::{ScramAuth, ScramSecret}; +use pgwire::api::auth::{AuthOutcome, AuthSource, LoginInfo as WireLoginInfo}; +use pgwire::api::{ClientInfo, PidSecretKeyGenerator, auth}; use pgwire::error::{ErrorInfo, PgWireError, PgWireResult}; use pgwire::messages::response::ErrorResponse; -use pgwire::messages::startup::{Authentication, PasswordMessageFamily, SecretKey}; +use pgwire::messages::startup::SecretKey; use pgwire::messages::{PgWireBackendMessage, PgWireFrontendMessage}; use session::Session; use snafu::IntoError; -use tokio::sync::Mutex; use crate::error::{AuthSnafu, Result}; use crate::metrics::METRIC_AUTH_FAILURE; @@ -44,168 +46,183 @@ use crate::postgres::types::PgErrorCode; use crate::postgres::utils::convert_err; use crate::query_handler::sql::ServerSqlQueryHandlerRef; -pub(crate) struct PgLoginVerifier { - user_provider: Option, - state: Mutex, -} +/// Query cancellation is not supported yet, so announce pid/secret key 0 as +/// before the `AuthDriver` adoption. +pub(crate) struct ZeroPidSecretKeyGenerator; -impl PgLoginVerifier { - pub(crate) fn new(user_provider: Option) -> Self { - Self { - user_provider, - state: Mutex::new(PgAuthenticationState::Initial), - } +impl PidSecretKeyGenerator for ZeroPidSecretKeyGenerator { + fn generate(&self, _client: &dyn ClientInfo) -> (i32, SecretKey) { + (0, SecretKey::I32(0)) } } -enum PgAuthenticationState { - Initial, - Cleartext, - SaslInitial { - auth_info: PgAuthInfo, - }, - SaslFinal { - verifier: PgScramSha256Verifier, - user_info: Option, - channel_binding: String, - nonce: String, - client_first_bare: String, - server_first: String, - }, +fn record_auth_failure(status_code: StatusCode) { + METRIC_AUTH_FAILURE + .with_label_values(&[status_code.as_ref()]) + .inc(); } -#[allow(dead_code)] -struct LoginInfo { - user: Option, - catalog: Option, - schema: Option, - host: String, +/// The [`AuthSource`] shared by pgwire's cleartext and SCRAM mechanisms. +/// Stateless by design: it only holds the user provider, everything about +/// the current login arrives with each call, and the verified +/// [`UserInfoRef`] travels back to the startup handler as the source +/// identity, reported by [`AuthOutcome::Authenticated`]. +/// +/// Read-only: it resolves credentials and checks authorization with the +/// provider, but holds no policy — authorization runs in the startup +/// handler after verification, before the startup completes. +pub(crate) struct PgUserProviderSource { + user_provider: UserProviderRef, } -impl LoginInfo { - pub fn from_client_info(client: &C) -> LoginInfo - where - C: ClientInfo, - { - LoginInfo { - user: client.metadata().get(super::METADATA_USER).map(Into::into), - catalog: client - .metadata() - .get(super::METADATA_CATALOG) - .map(Into::into), - schema: client - .metadata() - .get(super::METADATA_SCHEMA) - .map(Into::into), - host: client.socket_addr().ip().to_string(), - } +impl Debug for PgUserProviderSource { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PgUserProviderSource") + .finish_non_exhaustive() } } -impl PgLoginVerifier { - async fn auth(&self, login: &LoginInfo, password: &str) -> Result> { - let user_provider = match &self.user_provider { - Some(provider) => provider, - None => return Ok(None), - }; - - let user_name = match &login.user { - Some(name) => name, - None => return Ok(None), - }; - let catalog = match &login.catalog { - Some(name) => name, - None => return Ok(None), - }; - let schema = match &login.schema { - Some(name) => name, - None => return Ok(None), - }; - - let result = if user_name == BEARER_TOKEN_USER { - user_provider - .auth_bearer_token(password, catalog, schema) - .await - } else { - user_provider - .auth( - Identity::UserId(user_name, None), - Password::PlainText(password.to_string().into()), - catalog, - schema, - ) - .await - }; - match result { - Err(e) => { - METRIC_AUTH_FAILURE - .with_label_values(&[e.status_code().as_ref()]) - .inc(); - Err(AuthSnafu.into_error(e)) - } - Ok(user_info) => Ok(Some(user_info)), - } +impl PgUserProviderSource { + pub(crate) fn new(user_provider: UserProviderRef) -> Self { + Self { user_provider } } - async fn postgres_auth_info(&self, login: &LoginInfo) -> Result { - let user_provider = match &self.user_provider { - Some(provider) => provider, - None => return Ok(PgAuthInfo::Cleartext), - }; + /// (user, catalog, schema) of this login. Catalog/schema are parsed from + /// the `database` startup parameter, which was validated at connection + /// start. + fn resolve(login: &WireLoginInfo<'_>) -> Option<(String, String, String)> { + let user = login.user()?.to_string(); + let (catalog, schema) = parse_catalog_and_schema_from_db_string(login.database()?); + Some((user, catalog, schema)) + } - let user_name = match &login.user { - Some(name) => name, - None => return Ok(PgAuthInfo::Cleartext), - }; - if user_name == BEARER_TOKEN_USER { + /// Which mechanism to use for this login. Called once at connection + /// start; the SCRAM machinery resolves the verifier itself at + /// client-first time. + async fn postgres_auth_info(&self, login: &WireLoginInfo<'_>) -> Result { + let Some((user, catalog, _)) = Self::resolve(login) else { return Ok(PgAuthInfo::Cleartext); - } - let catalog = match &login.catalog { - Some(name) => name, - None => return Ok(PgAuthInfo::Cleartext), }; - match user_provider - .postgres_auth_info(Identity::UserId(user_name, None), catalog) + match self + .user_provider + .postgres_auth_info(Identity::UserId(&user, None), &catalog) .await { Err(e) => { - METRIC_AUTH_FAILURE - .with_label_values(&[e.status_code().as_ref()]) - .inc(); + record_auth_failure(e.status_code()); Err(AuthSnafu.into_error(e)) } Ok(auth_info) => Ok(auth_info), } } - async fn authorize(&self, login: &LoginInfo, user_info: &UserInfoRef) -> Result<()> { - let user_provider = match &self.user_provider { - Some(provider) => provider, - None => return Ok(()), + /// Checks with the provider that the user may access this + /// catalog/schema. Called by the startup handler after verification, + /// while nothing terminal has been sent yet. + async fn authorize( + &self, + login: &WireLoginInfo<'_>, + user_info: &UserInfoRef, + ) -> PgWireResult<()> { + let Some((_, catalog, schema)) = Self::resolve(login) else { + return Err(PgWireError::UserNameRequired); }; - let catalog = match &login.catalog { - Some(name) => name, - None => return Ok(()), - }; - let schema = match &login.schema { - Some(name) => name, - None => return Ok(()), - }; - - match user_provider.authorize(catalog, schema, user_info).await { + match self + .user_provider + .authorize(&catalog, &schema, user_info) + .await + { Err(e) => { - METRIC_AUTH_FAILURE - .with_label_values(&[e.status_code().as_ref()]) - .inc(); - Err(AuthSnafu.into_error(e)) + record_auth_failure(e.status_code()); + Err(PgWireError::ApiError(Box::new(AuthSnafu.into_error(e)))) } Ok(()) => Ok(()), } } } +#[async_trait] +impl AuthSource for PgUserProviderSource { + async fn get_password( + &self, + _login: &WireLoginInfo, + ) -> PgWireResult { + Err(PgWireError::ApiError( + "no cleartext password available from the user provider".into(), + )) + } + + /// Cleartext and bearer-token verification; the provider verifies the + /// password itself. Authorization is left to the startup handler. + async fn verify_password( + &self, + login: &WireLoginInfo, + presented_password: &[u8], + ) -> PgWireResult> { + let Some((user, catalog, _schema)) = Self::resolve(login) else { + return Err(PgWireError::UserNameRequired); + }; + let password = String::from_utf8_lossy(presented_password).into_owned(); + + let result = if user == BEARER_TOKEN_USER { + self.user_provider + .authenticate_bearer_token(&password, &catalog) + .await + } else { + self.user_provider + .authenticate( + Identity::UserId(&user, None), + Password::PlainText(password.into()), + ) + .await + }; + match result { + Err(e) => { + record_auth_failure(e.status_code()); + Err(PgWireError::ApiError(Box::new(AuthSnafu.into_error(e)))) + } + Ok(user_info) => Ok(Some(user_info)), + } + } + + async fn get_scram_secret( + &self, + login: &WireLoginInfo, + ) -> PgWireResult<(ScramSecret, Option)> { + match self.postgres_auth_info(login).await { + // For unknown users the provider hands out a mock verifier with + // no identity, so their handshake fails at proof verification + // exactly like a wrong password. + Ok(PgAuthInfo::ScramSha256 { + verifier, + user_info, + }) => Ok((to_scram_secret(&verifier), user_info)), + // SCRAM was selected but the provider no longer backs it: fail + // closed with a mock verifier. + _ => { + let user = login.user().unwrap_or_default(); + Ok(( + to_scram_secret(&PgScramSha256Verifier::mock_for_unknown_user( + user.as_bytes(), + )), + None, + )) + } + } + } +} + +fn to_scram_secret(verifier: &PgScramSha256Verifier) -> ScramSecret { + ScramSecret::Verifier { + salt: verifier.salt().to_vec(), + iterations: verifier.iterations() as usize, + stored_key: verifier.stored_key().to_vec(), + server_key: verifier.server_key().to_vec(), + } +} + fn set_client_info(client: &mut C, session: &Session) where C: ClientInfo, @@ -216,12 +233,35 @@ where if let Some(current_schema) = client.metadata().get(super::METADATA_SCHEMA) { session.set_schema(current_schema.clone()); } +} - // pass generated process id and secret key to client, this information will - // be sent to postgres client for query cancellation. - // use all 0 before we actually supported query cancellation - client.set_pid_and_secret_key(0, SecretKey::I32(0)); - // set userinfo outside +enum MechanismSelection { + /// No user provider configured: skip authentication. + NoAuth, + Cleartext(Arc), + Scram(Arc), + /// Resolving the mechanism failed. + Failed, +} + +impl PostgresServerHandlerInner { + async fn select_mechanism(&self, client: &C) -> MechanismSelection + where + C: ClientInfo + Sync, + { + let Some(source) = self.auth_source.clone() else { + return MechanismSelection::NoAuth; + }; + let login = WireLoginInfo::from_client_info(client); + if login.user() == Some(BEARER_TOKEN_USER) { + return MechanismSelection::Cleartext(source); + } + match source.postgres_auth_info(&login).await { + Err(_) => MechanismSelection::Failed, + Ok(PgAuthInfo::Cleartext) => MechanismSelection::Cleartext(source), + Ok(PgAuthInfo::ScramSha256 { .. }) => MechanismSelection::Scram(source), + } + } } #[async_trait] @@ -232,7 +272,7 @@ impl StartupHandler for PostgresServerHandlerInner { message: PgWireFrontendMessage, ) -> PgWireResult<()> where - C: ClientInfo + Sink + Unpin + Send, + C: ClientInfo + Sink + Unpin + Send + Sync, C::Error: Debug, PgWireError: From<>::Error>, { @@ -248,6 +288,7 @@ impl StartupHandler for PostgresServerHandlerInner { return Ok(()); } + auth::protocol_negotiation(client, startup).await?; auth::save_startup_parameters_to_metadata(client, startup); // check if db is valid @@ -280,57 +321,85 @@ impl StartupHandler for PostgresServerHandlerInner { } } - if self.login_verifier.user_provider.is_some() { - let login_info = LoginInfo::from_client_info(client); - let auth_info = match self.login_verifier.postgres_auth_info(&login_info).await - { - Ok(auth_info) => auth_info, - Err(_) => { + match self.select_mechanism(client).await { + MechanismSelection::NoAuth => { + self.session.set_user_info(userinfo_by_name( + client.metadata().get(super::METADATA_USER).cloned(), + )); + set_client_info(client, &self.session); + self.auth_driver + .complete_authentication(client, None) + .await?; + } + MechanismSelection::Cleartext(source) => { + if self + .auth_driver + .begin_authentication( + client, + CleartextPasswordAuth::new(source).create(), + ) + .await + .is_err() + { return send_password_authentication_failed(client).await; } - }; - client.set_state(PgWireConnectionState::AuthenticationInProgress); - match auth_info { - PgAuthInfo::ScramSha256 { .. } => { - *self.login_verifier.state.lock().await = - PgAuthenticationState::SaslInitial { auth_info }; - client - .send(PgWireBackendMessage::Authentication(Authentication::SASL( - vec![SCRAM_SHA_256_METHOD.to_string()], - ))) - .await?; - } - PgAuthInfo::Cleartext => { - *self.login_verifier.state.lock().await = - PgAuthenticationState::Cleartext; - client - .send(PgWireBackendMessage::Authentication( - Authentication::CleartextPassword, - )) - .await?; + } + MechanismSelection::Scram(source) => { + self.is_scram_auth.store(true, Ordering::Relaxed); + if self + .auth_driver + .begin_authentication(client, ScramAuth::new(source).create()) + .await + .is_err() + { + return send_password_authentication_failed(client).await; } } - } else { - self.session.set_user_info(userinfo_by_name( - client.metadata().get(super::METADATA_USER).cloned(), - )); - set_client_info(client, &self.session); - auth::finish_authentication(client, self.param_provider.as_ref()).await?; + MechanismSelection::Failed => { + return send_password_authentication_failed(client).await; + } } } PgWireFrontendMessage::PasswordMessageFamily(pwd) => { - let login_info = LoginInfo::from_client_info(client); - match self - .authenticate_password_message(client, &login_info, pwd) - .await? - { - PgAuthenticationResult::Continue => {} - PgAuthenticationResult::Success(user_info) => { + match self.auth_driver.advance_authentication(client, pwd).await { + Ok(AuthOutcome::InProgress) => {} + Ok(AuthOutcome::Authenticated { + identity, + final_message, + }) => { + self.is_scram_auth.store(false, Ordering::Relaxed); + let Some(user_info) = identity else { + return send_password_authentication_failed(client).await; + }; + + // Nothing terminal has been sent yet: authorize + // before completing, so a rejection reaches the + // client as an authentication failure. Failures are + // recorded by `authorize` itself. + let authorized = { + let login = WireLoginInfo::from_client_info(client); + match &self.auth_source { + Some(source) => source.authorize(&login, &user_info).await.is_ok(), + None => false, + } + }; + if !authorized { + return send_password_authentication_failed(client).await; + } + self.session.set_user_info(user_info); set_client_info(client, &self.session); - auth::finish_authentication(client, self.param_provider.as_ref()).await?; + self.auth_driver + .complete_authentication(client, final_message) + .await?; } - PgAuthenticationResult::Failed => { + Err(_e) => { + // SCRAM failures carry a uniform label so the metric + // cannot distinguish wrong password, unknown and + // unauthorized users. + if self.is_scram_auth.swap(false, Ordering::Relaxed) { + record_auth_failure(StatusCode::UserPasswordMismatch); + } return send_password_authentication_failed(client).await; } } @@ -341,225 +410,6 @@ impl StartupHandler for PostgresServerHandlerInner { } } -enum PgAuthenticationResult { - Continue, - Success(UserInfoRef), - Failed, -} - -impl PostgresServerHandlerInner { - /// Records a rejected SCRAM attempt in [`METRIC_AUTH_FAILURE`]. The label is - /// intentionally uniform (never `UserNotFound`), so the counter cannot be - /// used to distinguish a wrong password from an unknown user. - fn record_scram_failure(result: PgAuthenticationResult) -> PgAuthenticationResult { - if matches!(result, PgAuthenticationResult::Failed) { - METRIC_AUTH_FAILURE - .with_label_values(&[StatusCode::UserPasswordMismatch.as_ref()]) - .inc(); - } - result - } - - async fn authenticate_password_message( - &self, - client: &mut C, - login_info: &LoginInfo, - pwd: PasswordMessageFamily, - ) -> PgWireResult - where - C: ClientInfo + Sink + Unpin + Send, - C::Error: Debug, - PgWireError: From<>::Error>, - { - let mut state = self.login_verifier.state.lock().await; - let current_state = std::mem::replace(&mut *state, PgAuthenticationState::Initial); - - match current_state { - PgAuthenticationState::Cleartext => { - let pwd = pwd.into_password()?; - drop(state); - match self.login_verifier.auth(login_info, &pwd.password).await { - Ok(Some(user_info)) => Ok(PgAuthenticationResult::Success(user_info)), - Ok(None) | Err(_) => Ok(PgAuthenticationResult::Failed), - } - } - PgAuthenticationState::SaslInitial { auth_info } => { - // A labeled block funnels every rejection through a single - // `record_scram_failure`, so failed SCRAM attempts still show up - // in `METRIC_AUTH_FAILURE` without sprinkling the metric call - // across each early return. - let result = 'sasl: { - let sasl_initial = pwd.into_sasl_initial_response()?; - if sasl_initial.auth_method != SCRAM_SHA_256_METHOD { - break 'sasl PgAuthenticationResult::Failed; - } - let Some(data) = sasl_initial.data else { - break 'sasl PgAuthenticationResult::Failed; - }; - let Some(client_first) = ScramClientFirst::parse(&data) else { - break 'sasl PgAuthenticationResult::Failed; - }; - let PgAuthInfo::ScramSha256 { - verifier, - user_info, - } = auth_info - else { - break 'sasl PgAuthenticationResult::Failed; - }; - - let server_nonce = BASE64.encode(rand::random::<[u8; 18]>()); - let nonce = format!("{}{}", client_first.nonce, server_nonce); - let server_first = format!( - "r={},s={},i={}", - nonce, - BASE64.encode(verifier.salt()), - verifier.iterations() - ); - client - .send(PgWireBackendMessage::Authentication( - Authentication::SASLContinue(server_first.clone().into()), - )) - .await?; - *state = PgAuthenticationState::SaslFinal { - verifier, - user_info, - channel_binding: client_first.channel_binding, - nonce, - client_first_bare: client_first.bare, - server_first, - }; - PgAuthenticationResult::Continue - }; - Ok(Self::record_scram_failure(result)) - } - PgAuthenticationState::SaslFinal { - verifier, - user_info, - channel_binding, - nonce, - client_first_bare, - server_first, - } => { - let result = 'sasl: { - let sasl_response = pwd.into_sasl_response()?; - let Some(client_final) = ScramClientFinal::parse(&sasl_response.data) else { - break 'sasl PgAuthenticationResult::Failed; - }; - if client_final.channel_binding != channel_binding { - break 'sasl PgAuthenticationResult::Failed; - } - if client_final.nonce != nonce { - break 'sasl PgAuthenticationResult::Failed; - } - - let auth_message = format!( - "{},{},{}", - client_first_bare, server_first, client_final.without_proof - ); - let Ok(client_proof) = BASE64.decode(client_final.proof.as_bytes()) else { - break 'sasl PgAuthenticationResult::Failed; - }; - let Ok(Some(server_signature)) = - verifier.verify_client_proof(auth_message.as_bytes(), &client_proof) - else { - break 'sasl PgAuthenticationResult::Failed; - }; - let Some(user_info) = user_info else { - break 'sasl PgAuthenticationResult::Failed; - }; - - drop(state); - if self - .login_verifier - .authorize(login_info, &user_info) - .await - .is_err() - { - // `authorize` already recorded this failure; return early - // to bypass `record_scram_failure` and avoid double-counting. - return Ok(PgAuthenticationResult::Failed); - } - - client - .send(PgWireBackendMessage::Authentication( - Authentication::SASLFinal( - format!("v={}", BASE64.encode(server_signature)).into(), - ), - )) - .await?; - PgAuthenticationResult::Success(user_info) - }; - Ok(Self::record_scram_failure(result)) - } - PgAuthenticationState::Initial => Ok(PgAuthenticationResult::Failed), - } - } -} - -struct ScramClientFirst { - channel_binding: String, - bare: String, - nonce: String, -} - -impl ScramClientFirst { - fn parse(data: &[u8]) -> Option { - let message = std::str::from_utf8(data).ok()?; - let mut parts = message.splitn(3, ','); - let cbind = parts.next()?; - if !matches!(cbind, "n" | "y") { - return None; - } - let authzid = parts.next()?; - if !authzid.is_empty() { - return None; - } - let channel_binding = BASE64.encode(format!("{cbind},,")); - let bare = parts.next()?.to_string(); - let nonce = bare - .split(',') - .find_map(|chunk| chunk.strip_prefix("r="))? - .to_string(); - Some(Self { - channel_binding, - bare, - nonce, - }) - } -} - -struct ScramClientFinal { - channel_binding: String, - nonce: String, - without_proof: String, - proof: String, -} - -impl ScramClientFinal { - fn parse(data: &[u8]) -> Option { - let message = std::str::from_utf8(data).ok()?; - let proof_pos = message.rfind(",p=")?; - let without_proof = message[..proof_pos].to_string(); - let proof = message[proof_pos + 3..].to_string(); - let channel_binding = without_proof - .split(',') - .find_map(|chunk| chunk.strip_prefix("c="))?; - let nonce = without_proof - .split(',') - .find_map(|chunk| chunk.strip_prefix("r="))?; - if nonce.is_empty() || proof.is_empty() { - return None; - } - - Some(Self { - channel_binding: channel_binding.to_string(), - nonce: nonce.to_string(), - without_proof, - proof, - }) - } -} - async fn send_error(client: &mut C, err_info: ErrorInfo) -> PgWireResult<()> where C: ClientInfo + Sink + Unpin + Send, @@ -616,43 +466,3 @@ where Ok(DbResolution::NotFound("Database not specified".to_owned())) } } - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_scram_client_first_parse() { - let message = b"n,,n=greptime,r=clientnonce"; - let client_first = ScramClientFirst::parse(message).unwrap(); - assert_eq!("biws", client_first.channel_binding); - assert_eq!("n=greptime,r=clientnonce", client_first.bare); - assert_eq!("clientnonce", client_first.nonce); - - let message = b"y,,n=greptime,r=clientnonce"; - let client_first = ScramClientFirst::parse(message).unwrap(); - assert_eq!("eSws", client_first.channel_binding); - assert_eq!("n=greptime,r=clientnonce", client_first.bare); - assert_eq!("clientnonce", client_first.nonce); - - assert!(ScramClientFirst::parse(b"p=tls-server-end-point,,n=greptime,r=nonce").is_none()); - assert!(ScramClientFirst::parse(b"n,a=authzid,n=greptime,r=nonce").is_none()); - } - - #[test] - fn test_scram_client_final_parse() { - let message = b"c=biws,r=clientnonceservernonce,p=dGVzdA=="; - let client_final = ScramClientFinal::parse(message).unwrap(); - assert_eq!("biws", client_final.channel_binding); - assert_eq!("clientnonceservernonce", client_final.nonce); - assert_eq!( - "c=biws,r=clientnonceservernonce", - client_final.without_proof - ); - assert_eq!("dGVzdA==", client_final.proof); - - assert!(ScramClientFinal::parse(b"r=nonce,p=dGVzdA==").is_none()); - assert!(ScramClientFinal::parse(b"c=biws,r=nonce").is_none()); - assert!(ScramClientFinal::parse(b"c=biws,r=,p=dGVzdA==").is_none()); - } -} diff --git a/src/servers/tests/postgres/mod.rs b/src/servers/tests/postgres/mod.rs index 6b2995283f8..23971145a86 100644 --- a/src/servers/tests/postgres/mod.rs +++ b/src/servers/tests/postgres/mod.rs @@ -12,18 +12,22 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::HashMap; use std::io; use std::net::SocketAddr; use std::sync::Arc; +use std::sync::Mutex; use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Duration; use async_trait::async_trait; use auth::tests::{DatabaseAuthInfo, MockUserProvider}; use auth::{ - BEARER_TOKEN_USER, Identity, Password, UserInfoRef, UserProvider, UserProviderRef, - format_pg_scram_sha256_password_verifier, user_provider_from_option, + BEARER_TOKEN_USER, Identity, Password, PgAuthInfo, PgScramSha256Verifier, UserInfoRef, + UserProvider, UserProviderRef, format_pg_scram_sha256_password_verifier, + user_provider_from_option, }; +use common_base::secrets::ExposeSecret; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; use common_runtime::Builder as RuntimeBuilder; use common_runtime::runtime::BuilderBuild; @@ -40,6 +44,9 @@ use servers::server::Server; use servers::tls::{ReloadableTlsServerConfig, TlsOption}; use table::TableRef; use table::test_util::MemTable; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::TcpStream; +use tokio_postgres::error::SqlState; use tokio_postgres::{Client, Error as PgError, NoTls, SimpleQueryMessage}; use crate::create_testing_instance; @@ -89,6 +96,116 @@ impl UserProvider for BearerProvider { } } +const MIXED_AUTH_SCRAM_USER: &str = "scram_user"; +const MIXED_AUTH_SCRAM_PASSWORD: &[u8] = b"scram-password"; +const MIXED_AUTH_CLEARTEXT_USER: &str = "clear_user"; +const MIXED_AUTH_CLEARTEXT_PASSWORD: &str = "clear-password"; +const MIXED_AUTH_DENIED_USER: &str = "denied_user"; + +/// Selects SCRAM-SHA-256 for [`MIXED_AUTH_SCRAM_USER`] and +/// [`MIXED_AUTH_DENIED_USER`] (whose authorization is denied) and cleartext +/// for everyone else, so one endpoint serves both mechanisms. +struct MixedAuthProvider { + scram_verifier: PgScramSha256Verifier, + selections: Mutex>, + authorize_calls: Mutex>, +} + +#[async_trait] +impl UserProvider for MixedAuthProvider { + fn name(&self) -> &str { + "mixed-auth-test" + } + + async fn authenticate( + &self, + id: Identity<'_>, + password: Password<'_>, + ) -> auth::error::Result { + let Identity::UserId(username, _) = id; + let Password::PlainText(password) = password else { + return auth::error::UnsupportedPasswordTypeSnafu { + password_type: "plain_text", + } + .fail(); + }; + + let verified = if username == MIXED_AUTH_SCRAM_USER || username == MIXED_AUTH_DENIED_USER { + self.scram_verifier + .verify_plain_password(password.expose_secret().as_bytes())? + } else if username == MIXED_AUTH_CLEARTEXT_USER { + password.expose_secret() == MIXED_AUTH_CLEARTEXT_PASSWORD + } else { + return auth::error::UserNotFoundSnafu { + username: username.to_string(), + } + .fail(); + }; + + if verified { + Ok(auth::userinfo_by_name(Some(username.to_string()))) + } else { + auth::error::UserPasswordMismatchSnafu { + username: username.to_string(), + } + .fail() + } + } + + async fn authorize( + &self, + catalog: &str, + schema: &str, + user_info: &UserInfoRef, + ) -> auth::error::Result<()> { + *self + .authorize_calls + .lock() + .unwrap() + .entry(user_info.username().to_string()) + .or_insert(0) += 1; + + if user_info.username() == MIXED_AUTH_DENIED_USER { + return auth::error::AccessDeniedSnafu { + catalog: catalog.to_string(), + schema: schema.to_string(), + username: user_info.username().to_string(), + } + .fail(); + } + Ok(()) + } + + async fn postgres_auth_info( + &self, + id: Identity<'_>, + _catalog: &str, + ) -> auth::error::Result { + let Identity::UserId(username, _) = id; + let auth_info = if username == MIXED_AUTH_SCRAM_USER || username == MIXED_AUTH_DENIED_USER { + PgAuthInfo::ScramSha256 { + verifier: self.scram_verifier.clone(), + user_info: Some(auth::userinfo_by_name(Some(username.to_string()))), + } + } else { + PgAuthInfo::Cleartext + }; + + // Record the mechanism from the returned variant, so the map can + // never claim a mechanism the provider did not hand out. + let method = match &auth_info { + PgAuthInfo::ScramSha256 { .. } => "scram", + PgAuthInfo::Cleartext => "cleartext", + }; + self.selections + .lock() + .unwrap() + .insert(username.to_string(), method); + + Ok(auth_info) + } +} + fn create_postgres_server( table: TableRef, check_pwd: bool, @@ -244,6 +361,92 @@ async fn test_pg_scram_sha256_auth() -> Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_pg_scram_rejects_channel_binding_request() -> Result<()> { + // The server supports no channel binding (no certificate configured), + // so RFC 5802 §6 requires rejecting a client-first with a "p" flag. + let verifier = + format_pg_scram_sha256_password_verifier(b"greptime", b"pg-scram-salt", 4096).unwrap(); + let user_provider = + user_provider_from_option(&format!("static_user_provider:cmd:greptime={verifier}")) + .unwrap(); + let (postgres_server, server_port) = + start_test_server_with_user_provider(user_provider, Default::default()).await?; + + // "p" flag: rejected with an error + let response = raw_scram_client_first( + server_port, + b"p=tls-server-end-point,,n=greptime,r=clientnonce", + ) + .await?; + assert_eq!( + b'E', response, + "expected an ErrorResponse for a p-flag client-first" + ); + + // "n" flag (tokio_postgres sends this): handshake continues normally + let response = raw_scram_client_first(server_port, b"n,,n=greptime,r=clientnonce").await?; + assert_eq!( + b'R', response, + "expected a SASLContinue for an n-flag client-first" + ); + + postgres_server.shutdown().await.unwrap(); + Ok(()) +} + +/// A minimal raw SCRAM client: sends the startup message and a +/// SASLInitialResponse carrying `client_first`, returns the message type +/// byte of the server's answer ('R' for Authentication::SASL*, 'E' for +/// ErrorResponse). +async fn raw_scram_client_first(port: u16, client_first: &[u8]) -> Result { + let mut stream = TcpStream::connect(("127.0.0.1", port)).await.unwrap(); + + fn put_i32(buf: &mut Vec, v: i32) { + buf.extend_from_slice(&v.to_be_bytes()); + } + + async fn read_message(stream: &mut TcpStream) -> (u8, Vec) { + let mut header = [0u8; 5]; + stream.read_exact(&mut header).await.unwrap(); + let len = i32::from_be_bytes(header[1..].try_into().unwrap()) as usize - 4; + let mut body = vec![0u8; len]; + stream.read_exact(&mut body).await.unwrap(); + (header[0], body) + } + + // startup message: i32 len | i32 196608 | user/database params + let mut body = Vec::new(); + put_i32(&mut body, 196608); + body.extend_from_slice(b"user\0greptime\0database\0public\0\0"); + let mut message = Vec::new(); + put_i32(&mut message, (body.len() + 4) as i32); + message.extend_from_slice(&body); + stream.write_all(&message).await.unwrap(); + + // Authentication::SASL advertisement: 'R' | len | i32 10 | mechanism cstrings + let (typ, body) = read_message(&mut stream).await; + assert_eq!(b'R', typ, "expected the SASL mechanism advertisement"); + assert_eq!(10, i32::from_be_bytes(body[..4].try_into().unwrap())); + assert!( + body.windows(b"SCRAM-SHA-256".len()) + .any(|w| w == b"SCRAM-SHA-256") + ); + + // SASLInitialResponse: 'p' | len | "SCRAM-SHA-256"\0 | i32 len | client-first + let mut body = Vec::new(); + body.extend_from_slice(b"SCRAM-SHA-256\0"); + put_i32(&mut body, client_first.len() as i32); + body.extend_from_slice(client_first); + let mut message = Vec::new(); + message.push(b'p'); + put_i32(&mut message, (body.len() + 4) as i32); + message.extend_from_slice(&body); + stream.write_all(&message).await.unwrap(); + + Ok(read_message(&mut stream).await.0) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn test_pg_cleartext_auth_fallback() -> Result<()> { common_telemetry::init_default_ut_logging(); @@ -269,6 +472,103 @@ async fn test_pg_cleartext_auth_fallback() -> Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_pg_mixed_auth_scram_and_cleartext() -> Result<()> { + let provider = Arc::new(MixedAuthProvider { + scram_verifier: PgScramSha256Verifier::from_password( + MIXED_AUTH_SCRAM_PASSWORD, + b"mixed-auth-salt", + 4096, + ) + .unwrap(), + selections: Mutex::new(HashMap::new()), + authorize_calls: Mutex::new(HashMap::new()), + }); + let (postgres_server, server_port) = + start_test_server_with_user_provider(provider.clone(), Default::default()).await?; + + // SCRAM-selected user + let client = create_plain_connection_with_credentials( + server_port, + MIXED_AUTH_SCRAM_USER, + "scram-password", + ) + .await + .unwrap(); + let rows = client + .simple_query("SELECT uint32s FROM numbers LIMIT 1") + .await + .unwrap(); + assert_eq!("0", unwrap_results(&rows)[0]); + + assert!( + create_plain_connection_with_credentials(server_port, MIXED_AUTH_SCRAM_USER, "wrong") + .await + .is_err() + ); + + // cleartext-selected user + let client = create_plain_connection_with_credentials( + server_port, + MIXED_AUTH_CLEARTEXT_USER, + MIXED_AUTH_CLEARTEXT_PASSWORD, + ) + .await + .unwrap(); + let rows = client + .simple_query("SELECT uint32s FROM numbers LIMIT 1") + .await + .unwrap(); + assert_eq!("0", unwrap_results(&rows)[0]); + + assert!( + create_plain_connection_with_credentials(server_port, MIXED_AUTH_CLEARTEXT_USER, "wrong",) + .await + .is_err() + ); + + // unknown user falls back to cleartext and fails authentication + assert!( + create_plain_connection_with_credentials(server_port, "nobody", "whatever") + .await + .is_err() + ); + + // SCRAM-selected user denied by authorization: the client proof + // verifies, authorize rejects, and the failure reaches the client as + // 28P01 before SASLFinal/AuthenticationOk/ReadyForQuery were sent — + // had the completion fired first, this connection would succeed. + let denied = create_plain_connection_with_credentials( + server_port, + MIXED_AUTH_DENIED_USER, + "scram-password", + ) + .await + .unwrap_err(); + assert_eq!(Some(&SqlState::INVALID_PASSWORD), denied.code()); + + // both mechanisms were served by the same endpoint + let selections = provider.selections.lock().unwrap().clone(); + assert_eq!(Some(&"scram"), selections.get(MIXED_AUTH_SCRAM_USER)); + assert_eq!(Some(&"scram"), selections.get(MIXED_AUTH_DENIED_USER)); + assert_eq!( + Some(&"cleartext"), + selections.get(MIXED_AUTH_CLEARTEXT_USER) + ); + + // authorize ran exactly once for every password-verified user — + // including the denied one, proving it was rejected by authorization + // after successful SCRAM verification, not by authentication + let authorize_calls = provider.authorize_calls.lock().unwrap().clone(); + assert_eq!(Some(&1), authorize_calls.get(MIXED_AUTH_SCRAM_USER)); + assert_eq!(Some(&1), authorize_calls.get(MIXED_AUTH_DENIED_USER)); + assert_eq!(Some(&1), authorize_calls.get(MIXED_AUTH_CLEARTEXT_USER)); + assert!(!authorize_calls.contains_key("nobody")); + + postgres_server.shutdown().await.unwrap(); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_bearer_token_auth_over_cleartext() -> Result<()> { let provider = Arc::new(BearerProvider::default()); From 27c531336a3b7dece7712c682c53e277e4626d28 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Wed, 21 Oct 2026 14:58:15 +0000 Subject: [PATCH 063/239] fix(mito2): restore source splitting for cross-range series merges (#9818) * fix: restore batched series merges across ranges Signed-off-by: WenyXu * fix(mito2): split range sources before last-row selection Signed-off-by: WenyXu * fix(mito2): avoid tiny splits for single-source cross-range merges Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- src/mito2/src/engine/scan_test.rs | 192 ++++++++++++++++++++++++++++++ src/mito2/src/read/range_cache.rs | 20 ++++ src/mito2/src/read/scan_util.rs | 2 +- src/mito2/src/read/seq_scan.rs | 50 ++++++-- src/mito2/src/read/series_scan.rs | 9 +- 5 files changed, 259 insertions(+), 14 deletions(-) diff --git a/src/mito2/src/engine/scan_test.rs b/src/mito2/src/engine/scan_test.rs index 39613efb383..81b9b016772 100644 --- a/src/mito2/src/engine/scan_test.rs +++ b/src/mito2/src/engine/scan_test.rs @@ -1489,6 +1489,198 @@ async fn test_series_scan_with_format(flat_format: bool) { ); } +/// Cross-range merges split long series but preserve larger batches for short series. +#[tokio::test] +async fn test_series_scan_single_source_range_batches() { + for rows_per_series in [16, 49, 50, 64] { + check_series_scan_single_source_range_batches(rows_per_series).await; + } +} + +async fn check_series_scan_single_source_range_batches(rows_per_series: usize) { + let mut env = TestEnv::with_prefix("test_series_scan_single_source_range_batches").await; + let engine = env + .create_engine(MitoConfig { + default_flat_format: true, + range_result_cache_size: ReadableSize::mb(64), + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new() + .insert_option("compaction.type", "twcs") + .insert_option("compaction.twcs.time_window", "1h") + .build(); + let schema = test_util::rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + let mut expected = Vec::new(); + for start in [0, 3600] { + let mut rows = Vec::new(); + for key in ["a", "b", "c", "d"] { + rows.extend(test_util::build_rows_for_key( + key, + start, + start + rows_per_series, + 0, + )); + expected.extend( + (0..rows_per_series) + .map(|i| (key.to_string(), i as f64, (start + i) as i64 * 1000)), + ); + } + test_util::put_rows( + &engine, + region_id, + Rows { + schema: schema.clone(), + rows, + }, + ) + .await; + test_util::flush_region(&engine, region_id, None).await; + } + expected.sort_by(|a, b| a.0.cmp(&b.0).then(a.2.cmp(&b.2))); + + // Exercise both newly read and cached range streams at the same merge boundary. + for last_row in [false, true] { + let expected = if last_row { + expected + .iter() + .filter(|(_, _, ts)| *ts == (3600 + rows_per_series as i64 - 1) * 1000) + .cloned() + .collect::>() + } else { + expected.clone() + }; + for with_cache in [false, true] { + let request = ScanRequest { + series_row_selector: last_row + .then_some(TimeSeriesRowSelector::LastRow { after_merge: true }), + // Range caching requires a tag predicate; this one retains every row. + filters: if with_cache { + vec![col("tag_0").gt_eq(lit("a"))] + } else { + vec![] + }, + ..Default::default() + }; + // Populate the ordinary scan cache first. Its unsplit batches must not + // replace the batches prepared for a later cross-range merge. + if with_cache { + tokio::time::timeout(Duration::from_secs(10), async { + loop { + let Scanner::Seq(scanner) = + engine.scanner(region_id, request.clone()).await.unwrap() + else { + panic!("expected sequential scan"); + }; + let ctx = store_api::region_engine::QueryScanContext { + explain_verbose: true, + ..Default::default() + }; + let stream = scanner + .scan_partition(&ctx, &ExecutionPlanMetricsSet::default(), 0) + .unwrap(); + RecordBatches::try_collect(stream).await.unwrap(); + if format!("{}", VerboseDisplay(scanner)).contains("\"range_cache_hit\":") { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + } + let mut first_scan = true; + tokio::time::timeout(Duration::from_secs(10), async { + loop { + let scanner = engine + .scanner( + region_id, + ScanRequest { + distribution: Some(TimeSeriesDistribution::PerSeries), + ..request.clone() + }, + ) + .await + .unwrap(); + let Scanner::Series(scanner) = scanner else { + panic!("expected series scan"); + }; + assert_eq!(2, scanner.properties().partitions[0].len()); + + let ctx = store_api::region_engine::QueryScanContext { + explain_verbose: true, + ..Default::default() + }; + let metrics = ExecutionPlanMetricsSet::default(); + let stream = scanner.scan_partition(&ctx, &metrics, 0).unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + let actual = collect_and_assert_partition_rows( + batches.schema(), + vec![batches.iter().cloned().collect()], + ); + assert_eq!(expected, actual); + let display = format!("{}", VerboseDisplay(scanner)); + let details: serde_json::Value = + serde_json::from_str(&display[display.find('{').unwrap()..]).unwrap(); + let partitions = details["metrics_per_partition"].as_array().unwrap(); + let rows: u64 = partitions + .iter() + .map(|p| { + p["metrics"]["merge_metrics"]["num_fetch_by_rows"] + .as_u64() + .unwrap_or(0) + }) + .sum(); + if last_row { + // LastRow coalesces the source splits before the cross-range + // merge. Do not turn its output into one batch per series. + let batches: u64 = partitions + .iter() + .map(|p| { + p["metrics"]["merge_metrics"]["num_fetch_by_batches"] + .as_u64() + .unwrap_or(0) + }) + .sum(); + assert!(batches <= 2, "{display}"); + } else if rows_per_series >= 50 { + assert_eq!(0, rows, "{display}"); + } else { + // Short series retain batches spanning multiple series instead + // of sending one tiny batch per series through the distributor. + assert!(rows > 0, "{display}"); + } + if with_cache && first_scan { + assert!( + partitions.iter().all(|p| { + p["metrics"]["range_cache_hit"].as_u64().unwrap_or(0) == 0 + }), + "ordinary scan cache was reused: {display}" + ); + } + first_scan = false; + if !with_cache + || partitions + .iter() + .any(|p| p["metrics"]["range_cache_hit"].as_u64().unwrap_or(0) > 0) + { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + } + } +} + #[tokio::test] async fn test_two_phase_series_scan() { for use_index in [false, true] { diff --git a/src/mito2/src/read/range_cache.rs b/src/mito2/src/read/range_cache.rs index 1108b074697..081efa7bec8 100644 --- a/src/mito2/src/read/range_cache.rs +++ b/src/mito2/src/read/range_cache.rs @@ -76,6 +76,9 @@ pub(crate) struct ScanRequestFingerprint { #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] enum RangeScanStage { Data, + /// Source batches may be split for a later cross-range merge. Keep their + /// cached batch boundaries separate from ordinary sequential scans. + DataForCrossRangeMerge, CandidateSeries, SeriesData(SeriesRange), } @@ -237,6 +240,12 @@ pub(crate) struct RangeScanCacheKey { } impl RangeScanCacheKey { + pub(crate) fn for_cross_range_merge(mut self) -> Self { + debug_assert_eq!(self.scan.stage, RangeScanStage::Data); + self.scan.stage = RangeScanStage::DataForCrossRangeMerge; + self + } + pub(crate) fn estimated_size(&self) -> usize { mem::size_of::() + self.row_groups.capacity() * mem::size_of::<(FileId, i64)>() @@ -583,6 +592,11 @@ fn build_range_cache_key_inner( let scan = match stage { Some(RangeScanStage::CandidateSeries) => scan.for_candidate_series(), Some(RangeScanStage::SeriesData(assignment)) => scan.for_series_data(assignment), + Some(RangeScanStage::DataForCrossRangeMerge) => { + let mut scan = scan; + scan.stage = RangeScanStage::DataForCrossRangeMerge; + scan + } Some(RangeScanStage::Data) | None => scan, }; @@ -1208,6 +1222,12 @@ mod tests { let series_key_0 = build_series_range_cache_key(&ctx, &part_range, range_0).unwrap(); let series_key_1 = build_series_range_cache_key(&ctx, &part_range, range_1).unwrap(); + let merge_key = data_key.clone().for_cross_range_merge(); + assert_ne!(merge_key, data_key); + assert_ne!(merge_key.scan, candidate_key.scan); + assert_ne!(merge_key.scan, series_key_0.scan); + assert_eq!(merge_key.row_groups, data_key.row_groups); + assert_eq!(merge_key, data_key.clone().for_cross_range_merge()); assert_ne!(data_key.scan, candidate_key.scan); assert_ne!(data_key.scan, series_key_0.scan); assert_ne!(candidate_key.scan, series_key_0.scan); diff --git a/src/mito2/src/read/scan_util.rs b/src/mito2/src/read/scan_util.rs index 53967949387..b122704a71e 100644 --- a/src/mito2/src/read/scan_util.rs +++ b/src/mito2/src/read/scan_util.rs @@ -1203,7 +1203,7 @@ const SPLIT_ROW_THRESHOLD: u64 = DEFAULT_ROW_GROUP_SIZE as u64; const NUM_SERIES_THRESHOLD: u64 = 10240; /// Minimum batch size after splitting. The batch size is less than 60 because a series may only have /// 60 samples per hour. -const BATCH_SIZE_THRESHOLD: u64 = 50; +pub(crate) const BATCH_SIZE_THRESHOLD: u64 = 50; /// Minimum estimated rows per series for compaction to split batches. Shorter series /// would make compaction write one small batch per series. pub(crate) const MIN_COMPACTION_SPLIT_ROWS: usize = DEFAULT_READ_BATCH_SIZE / 2; diff --git a/src/mito2/src/read/seq_scan.rs b/src/mito2/src/read/seq_scan.rs index ecfa41bcf2e..8db83ed1845 100644 --- a/src/mito2/src/read/seq_scan.rs +++ b/src/mito2/src/read/seq_scan.rs @@ -46,15 +46,26 @@ use crate::read::range_cache::{ }; use crate::read::scan_region::{ScanInput, StreamContext}; use crate::read::scan_util::{ - MIN_COMPACTION_SPLIT_ROWS, PartitionMetrics, PartitionMetricsList, SplitRecordBatchStream, - compute_parallel_channel_size, scan_flat_file_ranges, scan_flat_mem_ranges, - should_split_flat_batches_for_merge, + BATCH_SIZE_THRESHOLD, MIN_COMPACTION_SPLIT_ROWS, PartitionMetrics, PartitionMetricsList, + SplitRecordBatchStream, compute_parallel_channel_size, scan_flat_file_ranges, + scan_flat_mem_ranges, should_split_flat_batches_for_merge, }; use crate::read::stream::{ConvertBatchStream, ScanBatch, ScanBatchStream}; use crate::read::{BoxedRecordBatchStream, ScannerMetrics, scan_util}; use crate::region::options::MergeMode; use crate::sst::parquet::DEFAULT_READ_BATCH_SIZE; +/// Describes where flat source batches will be merged. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum FlatReadMode { + /// Each range is consumed independently. + Sequential, + /// Range outputs feed a later SeriesScan merge. + CrossRangeMerge, + /// Sources from all ranges feed the compaction merge. + Compaction, +} + /// Scans a region and returns rows in a sorted sequence. /// /// The output order is always `order by primary keys, time index` inside every @@ -168,7 +179,7 @@ impl SeqScan { build_flat_sources( stream_ctx, part_range, - true, + FlatReadMode::Compaction, part_metrics, partition_pruner.clone(), &mut sources, @@ -277,13 +288,19 @@ impl SeqScan { pub(crate) async fn build_flat_partition_range_read( stream_ctx: &Arc, part_range: &PartitionRange, - compaction: bool, + read_mode: FlatReadMode, part_metrics: &PartitionMetrics, partition_pruner: Arc, file_scan_semaphore: Option>, merge_semaphore: Option>, ) -> Result<(BoxedRecordBatchStream, usize)> { - let cache_key = build_range_cache_key(stream_ctx, part_range); + let cache_key = build_range_cache_key(stream_ctx, part_range).map(|key| { + if read_mode == FlatReadMode::CrossRangeMerge { + key.for_cross_range_merge() + } else { + key + } + }); if let Some(key) = cache_key.as_ref() { if let Some(value) = stream_ctx.input.cache_strategy.get_range_result(key) { @@ -297,7 +314,7 @@ impl SeqScan { let split_batch_size = build_flat_sources( stream_ctx, part_range, - compaction, + read_mode, part_metrics, partition_pruner, &mut sources, @@ -410,7 +427,11 @@ impl SeqScan { let (mut reader, _) = Self::build_flat_partition_range_read( &stream_ctx, &part_range, - compaction, + if compaction { + FlatReadMode::Compaction + } else { + FlatReadMode::Sequential + }, &part_metrics, partition_pruner.clone(), file_scan_semaphore.clone(), @@ -633,12 +654,13 @@ impl fmt::Debug for SeqScan { pub(crate) async fn build_flat_sources( stream_ctx: &Arc, part_range: &PartitionRange, - compaction: bool, + read_mode: FlatReadMode, part_metrics: &PartitionMetrics, partition_pruner: Arc, sources: &mut Vec, semaphore: Option>, ) -> Result> { + let compaction = read_mode == FlatReadMode::Compaction; // Gets range meta. let range_meta = &stream_ctx.ranges[part_range.identifier]; #[cfg(debug_assertions)] @@ -738,11 +760,17 @@ pub(crate) async fn build_flat_sources( } } - // Queries merge one range at a time; compaction merges sources from all ranges. + // SeriesScan also merges range outputs. Split at the source, before deduplication + // and LastRow selection, so their output batching is preserved by the range cache. let needs_merge = ordered_sources.iter().flatten().count() > 1 || (compaction && stream_ctx.ranges.len() > 1); - let split_batch_size = if needs_merge { + let split_batch_size = if needs_merge || read_mode == FlatReadMode::CrossRangeMerge { should_split_flat_batches_for_merge(stream_ctx, range_meta).filter(|rows_per_series| { + // A single-source range only benefits at the downstream merge. Tiny + // splits add cache replay, channel and distribution work per batch. + if !needs_merge && *rows_per_series < BATCH_SIZE_THRESHOLD as usize { + return false; + } // Compaction writes every merged batch out, so splitting short series // would hand the SST writer and indexer one tiny batch per series. !compaction || *rows_per_series >= MIN_COMPACTION_SPLIT_ROWS diff --git a/src/mito2/src/read/series_scan.rs b/src/mito2/src/read/series_scan.rs index a77cbebb557..64d5ceb07f8 100644 --- a/src/mito2/src/read/series_scan.rs +++ b/src/mito2/src/read/series_scan.rs @@ -53,7 +53,7 @@ use crate::read::scan_util::{ PartitionMetrics, PartitionMetricsList, SeriesDistributorMetrics, compute_average_batch_size, compute_parallel_channel_size, }; -use crate::read::seq_scan::SeqScan; +use crate::read::seq_scan::{FlatReadMode, SeqScan}; use crate::read::series_candidate::{SeriesCandidateScanner, is_sparse_metric_metadata}; use crate::read::series_reader::{AssignedSeriesBatch, SeriesBatchCollector, SeriesReader}; use crate::read::stream::{ConvertBatchStream, ScanBatch, ScanBatchStream}; @@ -719,6 +719,7 @@ impl SeriesDistributor { // Create PartitionPruner covering all partitions let all_partition_ranges: Vec<_> = self.partitions.iter().flatten().cloned().collect(); + let needs_merge = all_partition_ranges.len() > 1; let partition_pruner = Arc::new(PartitionPruner::new( self.pruner.clone(), &all_partition_ranges, @@ -751,7 +752,11 @@ impl SeriesDistributor { SeqScan::build_flat_partition_range_read( &stream_ctx, &part_range, - false, + if needs_merge { + FlatReadMode::CrossRangeMerge + } else { + FlatReadMode::Sequential + }, &part_metrics, partition_pruner, file_scan_semaphore, From 436d07bf2dcf69cc3e5c5736f1bfe24a0ce4dff1 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Wed, 21 Oct 2026 19:03:56 +0000 Subject: [PATCH 064/239] fix(mito): return from manual index builds only after the index is published (#9643) * fix(mito): return from manual index builds only after the index is published A manual build_index could report success while the index was missing: the request was rejected or coalesced behind another build, the manifest commit failed, the indexer dropped an error, or the region was staging. Manual builds now wait for active builds of the same SST, keep their waiters across coalescing and schema-change retries, are notified after the worker applies the edit, and verify the files against the manifest before returning. The consistency check also honors inverted_index.ignore_column_ids and a zero segment row count, so such SSTs are no longer rebuilt on every call. Signed-off-by: Dennis Zhuang * test(mito): cover manual index build waiting for a schema-change retry Signed-off-by: Dennis Zhuang * fix(mito): select manual index build files by the applied version Signed-off-by: Dennis Zhuang * chore(mito): log why index builds are skipped or aborted Signed-off-by: Dennis Zhuang * refactor(mito): split IndexBuildTask::index_build into phases Signed-off-by: Dennis Zhuang * fix(mito): stop scheduling index builds when the region state changes Signed-off-by: Dennis Zhuang * chore(mito): document region file id and silent aborts in index builds Signed-off-by: Dennis Zhuang * fix(mito): log why a manual index build is incomplete Signed-off-by: Dennis Zhuang * fix(mito): expose incomplete index builds and treat compacted files as done Signed-off-by: Dennis Zhuang * test(mito): pin index build concurrency in the compaction test Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/mito2/src/engine/index_build_test.rs | 424 +++++++++-- src/mito2/src/error.rs | 13 + src/mito2/src/request.rs | 14 +- src/mito2/src/sst/file.rs | 133 ++-- src/mito2/src/sst/index.rs | 672 +++++++++--------- src/mito2/src/worker.rs | 3 - src/mito2/src/worker/handle_compaction.rs | 2 + src/mito2/src/worker/handle_flush.rs | 8 +- src/mito2/src/worker/handle_manifest.rs | 1 + src/mito2/src/worker/handle_rebuild_index.rs | 154 +++- .../admin/build_index_table_restart.result | 7 +- .../admin/build_index_table_restart.sql | 7 +- 12 files changed, 943 insertions(+), 495 deletions(-) diff --git a/src/mito2/src/engine/index_build_test.rs b/src/mito2/src/engine/index_build_test.rs index a00b4798a03..672c5e472d5 100644 --- a/src/mito2/src/engine/index_build_test.rs +++ b/src/mito2/src/engine/index_build_test.rs @@ -1147,7 +1147,10 @@ async fn test_consecutive_schema_changes_publish_latest_index_generation() { let files = current_file_metas(&engine, region_id).await; assert_eq!(files.len(), 1); assert!( - files[0].is_index_consistent_with_region(&version.metadata.column_metadatas), + files[0].is_index_consistent_with_region( + &version.metadata.column_metadatas, + &version.options.index_options, + ), "the published index must match the latest schema generation, file: {:?}, metadata: {:?}", files[0], version.metadata @@ -1166,10 +1169,10 @@ async fn test_consecutive_schema_changes_publish_latest_index_generation() { let reopened = engine.get_region(region_id).unwrap(); let reopened_files = current_file_metas(&engine, region_id).await; assert_eq!(reopened_files.len(), 1); - assert!( - reopened_files[0] - .is_index_consistent_with_region(&reopened.version().metadata.column_metadatas) - ); + assert!(reopened_files[0].is_index_consistent_with_region( + &reopened.version().metadata.column_metadatas, + &reopened.version().options.index_options, + )); } #[tokio::test] @@ -1250,7 +1253,184 @@ async fn test_unrelated_schema_change_retries_stale_flush_index_build() { let region = engine.get_region(region_id).unwrap(); let files = current_file_metas(&engine, region_id).await; assert_eq!(files.len(), 1); - assert!(files[0].is_index_consistent_with_region(®ion.version().metadata.column_metadatas)); + let version = region.version(); + assert!(files[0].is_index_consistent_with_region( + &version.metadata.column_metadatas, + &version.options.index_options, + )); + let scanner = engine + .scanner(region_id, ScanRequest::default()) + .await + .unwrap(); + assert_eq!(num_of_index_files(&engine, &scanner, region_id).await, 1); +} + +#[tokio::test] +async fn test_manual_build_waits_for_committed_publication_to_apply() { + let mut env = + TestEnv::with_prefix("test_manual_build_waits_for_committed_publication_to_apply_").await; + let gate = Arc::new(IndexPublicationGate::new( + IndexPublicationPhase::AfterManifestCommit, + )); + let engine = Arc::new( + env.create_engine_with( + async_build_mode_config(true), + None, + Some(gate.clone()), + None, + ) + .await, + ); + + let region_id = RegionId::new(1, 1); + env.get_schema_metadata_manager() + .register_region_table_info( + region_id.table_id(), + "test_table", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + + let request = CreateRequestBuilder::new().build_with_index(); + let column_schemas = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + put_and_flush(&engine, region_id, &column_schemas, 0..20).await; + // The flush build has committed its index to the manifest, but the worker + // has not applied it to the version. + tokio::time::timeout(Duration::from_secs(10), gate.wait_entered(1)) + .await + .expect("flush index build did not commit the manifest"); + + let engine_for_build = engine.clone(); + let mut manual_build = tokio::spawn(async move { + engine_for_build + .handle_request( + region_id, + RegionRequest::BuildIndex(RegionBuildIndexRequest::default()), + ) + .await + }); + assert!( + tokio::time::timeout(Duration::from_millis(300), &mut manual_build) + .await + .is_err(), + "manual build returned before the committed index was applied" + ); + + gate.release(1); + tokio::time::timeout(Duration::from_secs(10), manual_build) + .await + .expect("manual index build request did not finish") + .expect("manual index build task panicked") + .unwrap(); + + let version = engine.get_region(region_id).unwrap().version(); + let files = version + .ssts + .levels() + .iter() + .flat_map(|level| level.files.values()) + .collect::>(); + assert_eq!(files.len(), 1); + assert!(files[0].meta_ref().is_index_consistent_with_region( + &version.metadata.column_metadatas, + &version.options.index_options, + )); +} + +#[tokio::test] +async fn test_manual_build_waits_for_schema_change_retry() { + let mut env = TestEnv::with_prefix("test_manual_build_waits_for_schema_change_retry_").await; + let gate = Arc::new(IndexPublicationGate::new( + IndexPublicationPhase::BeforeManifestCommit, + )); + let engine = Arc::new( + env.create_engine_with( + async_build_mode_config(false), + None, + Some(gate.clone()), + None, + ) + .await, + ); + + let region_id = RegionId::new(1, 1); + env.get_schema_metadata_manager() + .register_region_table_info( + region_id.table_id(), + "test_table", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + + let request = CreateRequestBuilder::new().build_with_index(); + let column_schemas = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + put_and_flush(&engine, region_id, &column_schemas, 0..20).await; + + let engine_for_build = engine.clone(); + let manual_build = tokio::spawn(async move { + engine_for_build + .handle_request( + region_id, + RegionRequest::BuildIndex(RegionBuildIndexRequest::default()), + ) + .await + }); + tokio::time::timeout(Duration::from_secs(10), gate.wait_entered(1)) + .await + .expect("manual index build did not reach manifest publication"); + + // An unrelated schema change makes the manual publication stale. + engine + .handle_request( + region_id, + RegionRequest::Alter(RegionAlterRequest { + kind: AlterKind::AddColumns { + columns: vec![AddColumn { + column_metadata: ColumnMetadata { + column_schema: ColumnSchema::new( + "field_1", + ConcreteDataType::float64_datatype(), + true, + ), + semantic_type: SemanticType::Field, + column_id: 3, + }, + location: None, + }], + }, + }), + ) + .await + .unwrap(); + + gate.release(1); + tokio::time::timeout(Duration::from_secs(10), gate.wait_entered(2)) + .await + .expect("stale manual index build was not retried"); + // The caller waits for the retry rather than the aborted attempt. + assert!(!manual_build.is_finished()); + + gate.release(1); + tokio::time::timeout(Duration::from_secs(10), manual_build) + .await + .expect("manual index build request did not finish") + .expect("manual index build task panicked") + .unwrap(); + let scanner = engine .scanner(region_id, ScanRequest::default()) .await @@ -1339,7 +1519,10 @@ async fn test_schema_change_uses_manifest_file_generation() { let files = current_file_metas(&engine, region_id).await; assert_eq!(files.len(), 1); assert!( - files[0].is_index_consistent_with_region(&version.metadata.column_metadatas), + files[0].is_index_consistent_with_region( + &version.metadata.column_metadatas, + &version.options.index_options, + ), "the latest index definition was not published" ); } @@ -1593,36 +1776,31 @@ async fn test_index_build_type_manual_duplicate_in_flight() { assert_eq!(gate.finish_count(), 0); assert_eq!(gate.abort_count(), 0); - // Issue the second manual BuildIndex for the same region/file. - // Since the file is already in building_files (from the first manual build), - // schedule_build detects the duplicate and calls on_index_build_abort. - let request = RegionRequest::BuildIndex(RegionBuildIndexRequest::default()); - engine.handle_request(region_id, request).await.unwrap(); - - // The second request should have been aborted as duplicate. - assert_eq!(gate.abort_count(), 1, "duplicate request should be aborted"); - assert_eq!(gate.begin_count(), 2, "no new begin for duplicate"); - assert_eq!(gate.finish_count(), 0, "first build hasn't finished yet"); - - // Release the gate to let the first manual build proceed. - gate.release_begin(); - tokio::time::timeout(std::time::Duration::from_secs(5), gate.wait_finish(1)) + // The second manual BuildIndex for the same file waits for the active build + // instead of returning before the index exists. + let received = gate.recv_count(); + let engine_clone = engine.clone(); + let second_handle = tokio::spawn(async move { + let request = RegionRequest::BuildIndex(RegionBuildIndexRequest::default()); + engine_clone.handle_request(region_id, request).await + }); + tokio::time::timeout(Duration::from_secs(5), gate.wait_recv(received + 1)) .await - .unwrap(); // first manual build completes - - // Final counts: only one successful build, one aborted duplicate. - assert_eq!(gate.begin_count(), 2); // flush + first manual - assert_eq!(gate.finish_count(), 1); // first manual only - assert_eq!(gate.abort_count(), 1); // second manual duplicate abort - - // Await the first build's handle_request to ensure it completed cleanly. - tokio::time::timeout(std::time::Duration::from_secs(5), first_handle) - .await - .unwrap() - .unwrap() .unwrap(); - // Verify exactly one SST and one index file. + // Release the first manual build and the second one queued behind it. + gate.release_begin(); + gate.release_begin(); + for handle in [first_handle, second_handle] { + tokio::time::timeout(Duration::from_secs(5), handle) + .await + .unwrap() + .unwrap() + .unwrap(); + } + + // The queued build starts from the published index and skips the rebuild. + assert_eq!(gate.finish_count(), 1); let scanner = engine .scanner(region_id, ScanRequest::default()) .await @@ -1631,6 +1809,162 @@ async fn test_index_build_type_manual_duplicate_in_flight() { assert_eq!(num_of_index_files(&engine, &scanner, region_id).await, 1); } +#[tokio::test] +async fn test_index_build_type_manual_waits_for_in_flight_flush_build() { + let mut env = + TestEnv::with_prefix("test_index_build_type_manual_waits_for_in_flight_flush_build_").await; + let gate = Arc::new(GateIndexBuildListener::default()); + let engine = Arc::new( + env.create_engine_with( + async_build_mode_config(false), // Flush-triggered builds create no index. + None, + Some(gate.clone()), + None, + ) + .await, + ); + + let region_id = RegionId::new(1, 1); + env.get_schema_metadata_manager() + .register_region_table_info( + region_id.table_id(), + "test_table", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + + let request = CreateRequestBuilder::new().build_with_index(); + let column_schemas = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + put_and_flush(&engine, region_id, &column_schemas, 10..20).await; + tokio::time::timeout(Duration::from_secs(5), gate.wait_begin(1)) + .await + .unwrap(); + + // The manual build arrives while the no-op flush build of the same file is + // active. Returning at this point would leave the file without an index. + let received = gate.recv_count(); + let engine_clone = engine.clone(); + let manual_handle = tokio::spawn(async move { + let request = RegionRequest::BuildIndex(RegionBuildIndexRequest::default()); + engine_clone.handle_request(region_id, request).await + }); + tokio::time::timeout(Duration::from_secs(5), gate.wait_recv(received + 1)) + .await + .unwrap(); + + gate.release_begin(); + gate.release_begin(); + tokio::time::timeout(Duration::from_secs(5), manual_handle) + .await + .unwrap() + .unwrap() + .unwrap(); + + // The index is visible to scans as soon as the manual build returns. + let scanner = engine + .scanner(region_id, ScanRequest::default()) + .await + .unwrap(); + assert_eq!(scanner.num_files(), 1); + assert_eq!(num_of_index_files(&engine, &scanner, region_id).await, 1); +} + +#[tokio::test] +async fn test_manual_build_treats_compacted_files_as_done() { + let mut env = TestEnv::with_prefix("test_manual_build_treats_compacted_files_as_done_").await; + let gate = Arc::new(GateIndexBuildListener::default()); + let mut config = async_build_mode_config(false); + config.inverted_index.create_on_compaction = Mode::Disable; + config.fulltext_index.create_on_compaction = Mode::Disable; + config.bloom_filter_index.create_on_compaction = Mode::Disable; + // The builds of both SSTs must block at the gate together; the default + // limit depends on the CPU count and can be 1. + config.max_background_index_builds = 2; + let engine = Arc::new( + env.create_engine_with(config, None, Some(gate.clone()), None) + .await, + ); + + let region_id = RegionId::new(1, 1); + env.get_schema_metadata_manager() + .register_region_table_info( + region_id.table_id(), + "test_table", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + + let request = CreateRequestBuilder::new().build_with_index(); + let column_schemas = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + put_and_flush(&engine, region_id, &column_schemas, 10..20).await; + put_and_flush(&engine, region_id, &column_schemas, 20..30).await; + // The flush builds create no index. + tokio::time::timeout(Duration::from_secs(5), gate.wait_begin(2)) + .await + .unwrap(); + gate.release_begin(); + gate.release_begin(); + + let engine_for_build = engine.clone(); + let manual_build = tokio::spawn(async move { + engine_for_build + .handle_request( + region_id, + RegionRequest::BuildIndex(RegionBuildIndexRequest::default()), + ) + .await + }); + tokio::time::timeout(Duration::from_secs(5), gate.wait_begin(4)) + .await + .unwrap(); + + // Compaction replaces both SSTs while their manual builds are blocked. + engine + .handle_request( + region_id, + RegionRequest::Compact(RegionCompactRequest { + options: compact_request::Options::StrictWindow(StrictWindow { + window_seconds: 3600, + }), + parallelism: None, + time_range: None, + }), + ) + .await + .unwrap(); + for _ in 0..3 { + gate.release_begin(); + } + + // The build covers the SSTs that existed when it was requested, so the + // unindexed compaction output does not fail it. + tokio::time::timeout(Duration::from_secs(5), manual_build) + .await + .unwrap() + .unwrap() + .unwrap(); + let scanner = engine + .scanner(region_id, ScanRequest::default()) + .await + .unwrap(); + assert_eq!(scanner.num_files(), 1); + assert_eq!(num_of_index_files(&engine, &scanner, region_id).await, 0); +} + #[tokio::test] async fn test_reopen_waits_for_active_index_build_of_previous_incarnation() { let mut env = @@ -1723,12 +2057,18 @@ async fn test_reopen_waits_for_active_index_build_of_previous_incarnation() { assert_eq!(gate.begin_count(), 2); assert_eq!(gate.abort_count(), aborts_before_rebuild); + // The closed incarnation cannot publish its build, so its caller must not + // be told the index exists. gate.release_begin(); - tokio::time::timeout(Duration::from_secs(10), old_build) + let err = tokio::time::timeout(Duration::from_secs(10), old_build) .await .expect("old index build request did not finish") .expect("old index build task panicked") - .unwrap(); + .unwrap_err(); + assert!( + err.to_string().contains("Index build did not complete"), + "{err:?}" + ); tokio::time::timeout(Duration::from_secs(10), gate.wait_begin(3)) .await .expect("reopened index build did not start after the old build stopped"); @@ -2058,14 +2398,16 @@ async fn test_sync_flush_indexes_encoded_bulk_parts( return; } - let metadata = engine.get_region(region_id).unwrap().metadata(); + let version = engine.get_region(region_id).unwrap().version(); let indexed = async { loop { let files = current_file_metas(&engine, region_id).await; - if files - .iter() - .all(|f| f.is_index_consistent_with_region(&metadata.column_metadatas)) - { + if files.iter().all(|f| { + f.is_index_consistent_with_region( + &version.metadata.column_metadatas, + &version.options.index_options, + ) + }) { return files; } tokio::time::sleep(Duration::from_millis(20)).await; diff --git a/src/mito2/src/error.rs b/src/mito2/src/error.rs index ab781f790f4..d7ecdd88b1b 100644 --- a/src/mito2/src/error.rs +++ b/src/mito2/src/error.rs @@ -775,6 +775,18 @@ pub enum Error { location: Location, }, + #[snafu(display( + "Index build did not complete for files {:?} in region {}", + file_ids, + region_id + ))] + IndexBuildIncomplete { + region_id: RegionId, + file_ids: Vec, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Failed to convert value"))] ConvertValue { source: datatypes::error::Error, @@ -1561,6 +1573,7 @@ impl ErrorExt for Error { StatusCode::StorageUnavailable } FlushRegion { source, .. } | BuildIndexAsync { source, .. } => source.status_code(), + IndexBuildIncomplete { .. } => StatusCode::IllegalState, RegionDropped { .. } => StatusCode::Cancelled, RegionClosed { .. } => StatusCode::Cancelled, StaleCompactionExecution { .. } => StatusCode::Cancelled, diff --git a/src/mito2/src/request.rs b/src/mito2/src/request.rs index 5b9aff4f82d..2e81fbc4396 100644 --- a/src/mito2/src/request.rs +++ b/src/mito2/src/request.rs @@ -65,7 +65,7 @@ use crate::metrics::COMPACTION_ELAPSED_TOTAL; use crate::region::MitoRegionRef; use crate::region::options::RegionOptions; use crate::sst::file::FileMeta; -use crate::sst::index::IndexBuildType; +use crate::sst::index::{IndexBuildType, ResultMpscSender}; use crate::wal::EntryId; use crate::wal::entry_distributor::WalEntryReceiver; @@ -1113,8 +1113,6 @@ pub(crate) enum BackgroundNotify { IndexBuildFinished(IndexBuildFinished), /// Index build has been stopped (aborted or succeeded). IndexBuildStopped(IndexBuildStopped), - /// Index build has failed. - IndexBuildFailed(IndexBuildFailed), /// An index build must be retried against the latest schema generation. IndexBuildRetry(BuildIndexRequest), /// Compaction has finished. @@ -1196,6 +1194,8 @@ impl FlushFailed { pub(crate) struct IndexBuildFinished { pub(crate) manifest_version: ManifestVersion, pub(crate) file_meta: FileMeta, + /// Waiters of the build, notified after the worker applies the edit. + pub(crate) result_senders: Vec, } /// Notifies an index build job has been stopped. @@ -1204,12 +1204,6 @@ pub(crate) struct IndexBuildStopped { pub(crate) file_id: FileId, } -/// Notifies an index build job has failed. -#[derive(Debug)] -pub(crate) struct IndexBuildFailed { - pub(crate) err: Arc, -} - /// A local unit event carries only identity, never an input reservation lease. #[derive(Debug)] pub(crate) enum CompactionUnitNotification { @@ -1419,6 +1413,8 @@ pub(crate) struct BuildIndexRequest { pub(crate) build_type: IndexBuildType, /// files need to build index, empty means all. pub(crate) file_metas: Vec, + /// Waiters carried over from a previous build of the same files. + pub(crate) result_senders: Vec, } #[derive(Debug)] diff --git a/src/mito2/src/sst/file.rs b/src/mito2/src/sst/file.rs index 5f4dcb28752..28b6218cd7b 100644 --- a/src/mito2/src/sst/file.rs +++ b/src/mito2/src/sst/file.rs @@ -37,6 +37,7 @@ use store_api::storage::{ColumnId, FileId, IndexVersion, RegionId}; use crate::access_layer::AccessLayerRef; use crate::cache::CacheManagerRef; use crate::cache::file_cache::{FileType, IndexKey}; +use crate::region::options::IndexOptions; use crate::sst::file_purger::FilePurgerRef; use crate::sst::index::fulltext_index::FULLTEXT_BLOOM_VERSION; use crate::sst::location; @@ -456,56 +457,73 @@ impl FileMeta { self.index_file_size } - /// Check whether the file index is consistent with the given region metadata. - pub fn is_index_consistent_with_region(&self, metadata: &[ColumnMetadata]) -> bool { + /// Check whether the file index is consistent with the given region metadata + /// and index options. + pub fn is_index_consistent_with_region( + &self, + metadata: &[ColumnMetadata], + index_options: &IndexOptions, + ) -> bool { let id_to_indexes = self .indexes .iter() .map(|index| (index.column_id, index)) .collect::>(); + let inverted_options = &index_options.inverted_index; for column in metadata { - if !column.column_schema.is_indexed() { + let schema = &column.column_schema; + // The index builder never creates these inverted indexes, so expecting + // them would make every manual build rebuild the file again. + let expects_inverted = schema.is_inverted_indexed() + && inverted_options.segment_row_count > 0 + && !inverted_options + .ignore_column_ids + .contains(&column.column_id); + let expects_fulltext = schema.is_fulltext_indexed(); + let expects_skipping = schema.is_skipping_indexed(); + if !expects_inverted && !expects_fulltext && !expects_skipping { continue; } - if let Some(index) = id_to_indexes.get(&column.column_id) { - let indexes = &index.created_indexes; - if column.column_schema.is_inverted_indexed() - && !indexes.contains(&IndexType::InvertedIndex) - { - return false; - } - if column.column_schema.is_fulltext_indexed() - && !indexes.contains(&IndexType::FulltextIndex) - { - return false; - } - if column - .column_schema - .fulltext_options() - .ok() - .flatten() - .is_some_and(|options| { - options.enable && options.backend == FulltextBackend::Bloom - }) - && index.fulltext_bloom_version < FULLTEXT_BLOOM_VERSION - { - debug!( - "Fulltext bloom index is outdated, region_id: {}, file_id: {}, column_id: {}, version: {}, current: {}", - self.region_id, - self.file_id, - column.column_id, - index.fulltext_bloom_version, - FULLTEXT_BLOOM_VERSION - ); - return false; - } - if column.column_schema.is_skipping_indexed() - && !indexes.contains(&IndexType::BloomFilterIndex) - { - return false; - } - } else { + let missing_index = |index_type: IndexType| { + debug!( + "File index is missing, region_id: {}, file_id: {}, column_id: {}, index_type: {:?}", + self.region_id, self.file_id, column.column_id, index_type + ); + false + }; + let Some(index) = id_to_indexes.get(&column.column_id) else { + debug!( + "File has no index for column, region_id: {}, file_id: {}, column_id: {}", + self.region_id, self.file_id, column.column_id + ); return false; + }; + let indexes = &index.created_indexes; + if expects_inverted && !indexes.contains(&IndexType::InvertedIndex) { + return missing_index(IndexType::InvertedIndex); + } + if expects_fulltext && !indexes.contains(&IndexType::FulltextIndex) { + return missing_index(IndexType::FulltextIndex); + } + if schema + .fulltext_options() + .ok() + .flatten() + .is_some_and(|options| options.enable && options.backend == FulltextBackend::Bloom) + && index.fulltext_bloom_version < FULLTEXT_BLOOM_VERSION + { + debug!( + "Fulltext bloom index is outdated, region_id: {}, file_id: {}, column_id: {}, version: {}, current: {}", + self.region_id, + self.file_id, + column.column_id, + index.fulltext_bloom_version, + FULLTEXT_BLOOM_VERSION + ); + return false; + } + if expects_skipping && !indexes.contains(&IndexType::BloomFilterIndex) { + return missing_index(IndexType::BloomFilterIndex); } } true @@ -1379,7 +1397,7 @@ mod tests { ..Default::default() }; let region_meta = vec![new_column_meta(1, "tag1", true, false, false)]; - assert!(file_meta.is_index_consistent_with_region(®ion_meta)); + assert!(file_meta.is_index_consistent_with_region(®ion_meta, &IndexOptions::default())); // Case 2: Superset match. File has more indexes than required. file_meta.indexes = vec![ColumnIndexMetadata { @@ -1391,7 +1409,7 @@ mod tests { ..Default::default() }]; let region_meta = vec![new_column_meta(1, "tag1", true, false, false)]; - assert!(file_meta.is_index_consistent_with_region(®ion_meta)); + assert!(file_meta.is_index_consistent_with_region(®ion_meta, &IndexOptions::default())); // Case 3: Missing index type. File has the column but lacks the required index type. file_meta.indexes = vec![ColumnIndexMetadata { @@ -1400,7 +1418,7 @@ mod tests { ..Default::default() }]; let region_meta = vec![new_column_meta(1, "tag1", true, true, false)]; // Requires fulltext too - assert!(!file_meta.is_index_consistent_with_region(®ion_meta)); + assert!(!file_meta.is_index_consistent_with_region(®ion_meta, &IndexOptions::default())); // Case 4: Missing column. Region requires an index on a column not in the file's index list. file_meta.indexes = vec![ColumnIndexMetadata { @@ -1409,7 +1427,7 @@ mod tests { ..Default::default() }]; let region_meta = vec![new_column_meta(1, "tag1", true, false, false)]; // Requires index on column 1 - assert!(!file_meta.is_index_consistent_with_region(®ion_meta)); + assert!(!file_meta.is_index_consistent_with_region(®ion_meta, &IndexOptions::default())); // Case 5: No indexes required by region. Should always be consistent. file_meta.indexes = vec![ColumnIndexMetadata { @@ -1418,12 +1436,12 @@ mod tests { ..Default::default() }]; let region_meta = vec![new_column_meta(1, "tag1", false, false, false)]; // No index required - assert!(file_meta.is_index_consistent_with_region(®ion_meta)); + assert!(file_meta.is_index_consistent_with_region(®ion_meta, &IndexOptions::default())); // Case 6: Empty file indexes. Region requires an index. file_meta.indexes = vec![]; let region_meta = vec![new_column_meta(1, "tag1", true, false, false)]; - assert!(!file_meta.is_index_consistent_with_region(®ion_meta)); + assert!(!file_meta.is_index_consistent_with_region(®ion_meta, &IndexOptions::default())); // Case 7: Multiple columns, one is inconsistent. file_meta.indexes = vec![ @@ -1442,7 +1460,17 @@ mod tests { new_column_meta(1, "tag1", true, false, false), new_column_meta(2, "tag2", false, true, true), // Requires Fulltext and BloomFilter ]; - assert!(!file_meta.is_index_consistent_with_region(®ion_meta)); + assert!(!file_meta.is_index_consistent_with_region(®ion_meta, &IndexOptions::default())); + + // Case 8: Inverted indexes the builder skips by index options are not required. + file_meta.indexes = vec![]; + let region_meta = vec![new_column_meta(1, "tag1", true, false, false)]; + let mut index_options = IndexOptions::default(); + index_options.inverted_index.ignore_column_ids = vec![1]; + assert!(file_meta.is_index_consistent_with_region(®ion_meta, &index_options)); + let mut index_options = IndexOptions::default(); + index_options.inverted_index.segment_row_count = 0; + assert!(file_meta.is_index_consistent_with_region(®ion_meta, &index_options)); } #[test] @@ -1475,10 +1503,13 @@ mod tests { let bloom = fulltext_column(FulltextBackend::Bloom); // Written before the version was recorded, i.e. with the old bloom blob. - assert!(!file_meta(0).is_index_consistent_with_region(&bloom)); - assert!(file_meta(FULLTEXT_BLOOM_VERSION).is_index_consistent_with_region(&bloom)); + assert!(!file_meta(0).is_index_consistent_with_region(&bloom, &IndexOptions::default())); + assert!( + file_meta(FULLTEXT_BLOOM_VERSION) + .is_index_consistent_with_region(&bloom, &IndexOptions::default()) + ); // Tantivy indexes carry no bloom version. let tantivy = fulltext_column(FulltextBackend::Tantivy); - assert!(file_meta(0).is_index_consistent_with_region(&tantivy)); + assert!(file_meta(0).is_index_consistent_with_region(&tantivy, &IndexOptions::default())); } } diff --git a/src/mito2/src/sst/index.rs b/src/mito2/src/sst/index.rs index 6b90d337202..ff734e3241b 100644 --- a/src/mito2/src/sst/index.rs +++ b/src/mito2/src/sst/index.rs @@ -30,6 +30,7 @@ pub(crate) mod store; use std::cmp::Ordering; use std::collections::{HashMap, HashSet}; use std::num::NonZeroUsize; +use std::ops::ControlFlow; use std::sync::Arc; use bloom_filter::creator::BloomFilterIndexer; @@ -66,8 +67,8 @@ use crate::region::{ IndexBuildSource, IndexPublication, IndexPublicationStale, ManifestContextRef, }; use crate::request::{ - BackgroundNotify, BuildIndexRequest, IndexBuildFailed, IndexBuildFinished, IndexBuildStopped, - WorkerRequest, WorkerRequestWithTime, + BackgroundNotify, BuildIndexRequest, IndexBuildFinished, IndexBuildStopped, WorkerRequest, + WorkerRequestWithTime, }; use crate::schedule::scheduler::{Job, SchedulerRef}; use crate::sst::file::{ @@ -649,6 +650,15 @@ pub enum IndexBuildOutcome { Aborted(String), } +/// How an index build task reports its result. +enum IndexBuildReport { + /// The task sends the outcome to its waiters. + Done(IndexBuildOutcome), + /// The worker or a retry task owns the waiters, so the task must not + /// report to them. + Deferred, +} + /// Mpsc output result sender. pub type ResultMpscSender = Sender>; @@ -668,6 +678,8 @@ pub struct IndexBuildTask { pub(crate) target_region_metadata: RegionMetadataRef, /// The manifest state this build is based on. pub(crate) source: IndexBuildSource, + /// Index options of the target region when the task was created. + pub(crate) index_options: IndexOptions, pub reason: IndexBuildType, pub access_layer: AccessLayerRef, pub(crate) listener: WorkerListener, @@ -680,8 +692,9 @@ pub struct IndexBuildTask { pub indexer_builder: Arc, /// Request sender to notify the region worker. pub(crate) request_sender: Sender, - /// Index build result sender. - pub(crate) result_sender: ResultMpscSender, + /// Senders of callers waiting for this build. Coalesced tasks and + /// schema-change retries hand their senders over to the surviving task. + pub(crate) result_senders: Vec, } impl std::fmt::Debug for IndexBuildTask { @@ -699,17 +712,20 @@ impl std::fmt::Debug for IndexBuildTask { impl IndexBuildTask { /// Notify the caller the job is success. pub async fn on_success(&self, outcome: IndexBuildOutcome) { - let _ = self.result_sender.send(Ok(outcome)).await; + for sender in &self.result_senders { + let _ = sender.send(Ok(outcome.clone())).await; + } } /// Send index build error to waiter. pub async fn on_failure(&self, err: Arc) { - let _ = self - .result_sender - .send(Err(err.clone()).context(BuildIndexAsyncSnafu { - region_id: self.region_id, - })) - .await; + for sender in &self.result_senders { + let _ = sender + .send(Err(err.clone()).context(BuildIndexAsyncSnafu { + region_id: self.region_id, + })) + .await; + } } fn into_index_build_job(mut self, version_control: VersionControlRef) -> Job { @@ -720,10 +736,7 @@ impl IndexBuildTask { async fn do_index_build(&mut self, version_control: VersionControlRef) { self.listener - .on_index_build_begin(RegionFileId::new( - self.source.file_meta.region_id, - self.source.file_meta.file_id, - )) + .on_index_build_begin(self.region_file_id()) .await; let result = if self.reason == IndexBuildType::Compact { let mut task = self.clone(); @@ -740,7 +753,13 @@ impl IndexBuildTask { self.index_build(version_control).await }; match result { - Ok(outcome) => self.on_success(outcome).await, + Ok(IndexBuildReport::Done(outcome)) => { + if let IndexBuildOutcome::Aborted(reason) = &outcome { + debug!("Index build aborted: {}", reason); + } + self.on_success(outcome).await + } + Ok(IndexBuildReport::Deferred) => {} Err(e) => { warn!( e; "Index build task failed, region: {}, file_id: {}", @@ -749,11 +768,23 @@ impl IndexBuildTask { self.on_failure(e.into()).await } } + self.notify_worker(BackgroundNotify::IndexBuildStopped(IndexBuildStopped { + file_id: self.source.file_meta.file_id, + })) + .await; + } + + fn region_file_id(&self) -> RegionFileId { + RegionFileId::new( + self.source.file_meta.region_id, + self.source.file_meta.file_id, + ) + } + + async fn notify_worker(&self, notify: BackgroundNotify) { let worker_request = WorkerRequest::Background { region_id: self.region_id, - notify: BackgroundNotify::IndexBuildStopped(IndexBuildStopped { - file_id: self.source.file_meta.file_id, - }), + notify, }; let _ = self .request_sender @@ -761,6 +792,25 @@ impl IndexBuildTask { .await; } + /// Builds an aborted outcome without a listener event. Builds that create + /// no index stop here, and the listener observes them as neither finished + /// nor aborted. + fn aborted(&self, reason: &str) -> IndexBuildReport { + IndexBuildReport::Done(IndexBuildOutcome::Aborted(format!( + "{reason}, region: {}, file_id: {}", + self.region_id, self.source.file_meta.file_id + ))) + } + + async fn abort_and_notify( + &self, + region_file_id: RegionFileId, + reason: &str, + ) -> IndexBuildReport { + self.listener.on_index_build_abort(region_file_id).await; + self.aborted(reason) + } + // Checks if the SST file still exists in object store and version to avoid conflict with compaction. async fn check_sst_file_exists(&self, version_control: &VersionControlRef) -> bool { let file_id = self.source.file_meta.file_id; @@ -796,27 +846,93 @@ impl IndexBuildTask { async fn index_build( &mut self, version_control: VersionControlRef, - ) -> Result { + ) -> Result { + // Taken before `prepare_source` replaces the source, and used for the + // index files and listener events of this build. + let region_file_id = self.region_file_id(); + if let ControlFlow::Break(report) = + self.prepare_source(region_file_id, &version_control).await + { + return Ok(report); + } let new_index_version = self .source .file_meta .index_version() .map_or(0, |version| version + 1); + let (mut indexer, index_output) = + match self.create_index(region_file_id, new_index_version).await? { + ControlFlow::Continue(created) => created, + ControlFlow::Break(report) => return Ok(report), + }; - // Check SST file existence before building index to avoid failure of parquet reader. + // Check SST file existence again after building index. if !self.check_sst_file_exists(&version_control).await { - self.listener - .on_index_build_abort(RegionFileId::new( - self.source.file_meta.region_id, - self.source.file_meta.file_id, - )) - .await; - return Ok(IndexBuildOutcome::Aborted(format!( - "SST file not found during index build, region: {}, file_id: {}", - self.region_id, self.source.file_meta.file_id - ))); + // Calls abort to clean up index files. + indexer.abort().await; + return Ok(self + .abort_and_notify(region_file_id, "SST file not found during index build") + .await); } + self.publish(region_file_id, index_output, new_index_version) + .await + } + + /// Checks that the SST still exists and moves the build source to the + /// file's current manifest state. + async fn prepare_source( + &mut self, + region_file_id: RegionFileId, + version_control: &VersionControlRef, + ) -> ControlFlow { + // Check SST file existence before building index to avoid failure of parquet reader. + if !self.check_sst_file_exists(version_control).await { + return ControlFlow::Break( + self.abort_and_notify(region_file_id, "SST file not found during index build") + .await, + ); + } + + // A task may have waited behind another build of the same SST. Build on + // the index that build published, otherwise this publication is stale. + let current_file_meta = self + .manifest_ctx + .manifest() + .await + .files + .get(&self.source.file_meta.file_id) + .cloned(); + let Some(current_file_meta) = current_file_meta else { + return ControlFlow::Break( + self.abort_and_notify( + region_file_id, + "SST file not found in manifest during index build", + ) + .await, + ); + }; + if self.reason == IndexBuildType::Manual + && current_file_meta.is_index_consistent_with_region( + &self.target_region_metadata.column_metadatas, + &self.index_options, + ) + { + return ControlFlow::Break( + self.abort_and_notify(region_file_id, "Index is already up to date") + .await, + ); + } + self.source.file_meta = current_file_meta; + ControlFlow::Continue(()) + } + + /// Reads the SST and creates the index files. + async fn create_index( + &self, + region_file_id: RegionFileId, + new_index_version: u64, + ) -> Result> { let mut parquet_reader = self .access_layer .read_sst(self.file.clone()) // use the latest file handle instead of creating a new one @@ -833,12 +949,6 @@ impl IndexBuildTask { .filter(|size| *size > 0) }); - // Use the same file_id but with new version for index file. - let region_file_id = RegionFileId::new( - self.source.file_meta.region_id, - self.source.file_meta.file_id, - ); - let index_file_id = region_file_id.file_id(); let mut indexer = self .indexer_builder .build(region_file_id, new_index_version, row_group_size) @@ -846,7 +956,7 @@ impl IndexBuildTask { // The build type or index options may leave nothing to build. Return before reading // any row; an empty indexer finishes with no output, which publishes nothing. if indexer.is_empty() { - return Ok(IndexBuildOutcome::Finished); + return Ok(ControlFlow::Break(self.aborted("No index to build"))); } if let Some(mut parquet_reader) = parquet_reader.take() { @@ -866,103 +976,100 @@ impl IndexBuildTask { } let index_output = indexer.finish().await; - if index_output.file_size > 0 { - // Check SST file existence again after building index. - if !self.check_sst_file_exists(&version_control).await { - // Calls abort to clean up index files. - indexer.abort().await; + if index_output.file_size == 0 { + // The indexer logs and drops index creation failures. + return Ok(ControlFlow::Break(self.aborted("No index was created"))); + } + Ok(ControlFlow::Continue((indexer, index_output))) + } + + /// Uploads the index files and commits them to the manifest. + async fn publish( + &self, + region_file_id: RegionFileId, + index_output: IndexOutput, + new_index_version: u64, + ) -> Result { + // Use the same file_id but with new version for index file. + self.maybe_upload_index_file( + index_output.clone(), + region_file_id.file_id(), + new_index_version, + ) + .await?; + + self.listener + .on_index_build_before_manifest_commit(region_file_id) + .await; + + // The commit may have reached storage even on error, so the uploaded + // index file is left for GC instead of being deleted here. + match self + .update_manifest(index_output, new_index_version) + .await? + { + IndexPublication::Committed { + manifest_version, + file_meta, + } => { self.listener - .on_index_build_abort(RegionFileId::new( - self.source.file_meta.region_id, - self.source.file_meta.file_id, - )) + .on_index_build_manifest_committed(region_file_id) .await; - return Ok(IndexBuildOutcome::Aborted(format!( - "SST file not found during index build, region: {}, file_id: {}", - self.region_id, self.source.file_meta.file_id - ))); - } - - // Upload index file if write cache is enabled. - self.maybe_upload_index_file(index_output.clone(), index_file_id, new_index_version) - .await?; - - self.listener - .on_index_build_before_manifest_commit(region_file_id) - .await; - - let worker_request = match self.update_manifest(index_output, new_index_version).await { - Ok(IndexPublication::Committed { + self.notify_worker(BackgroundNotify::IndexBuildFinished(IndexBuildFinished { manifest_version, file_meta, - }) => { - self.listener - .on_index_build_manifest_committed(region_file_id) - .await; - let index_build_finished = IndexBuildFinished { - manifest_version, - file_meta, - }; - WorkerRequest::Background { - region_id: self.region_id, - notify: BackgroundNotify::IndexBuildFinished(index_build_finished), - } - } - Ok(IndexPublication::Stale(stale)) => { - INDEX_PUBLICATION_STALE_TOTAL - .with_label_values(&["manifest_commit"]) - .inc(); - let index_id = RegionIndexId::new(region_file_id, new_index_version); - // If no successor publishes the same version, GC collects the - // remote artifact. Repartition requires GC, while local mode - // accepts this narrow orphan window when GC is disabled. - cleanup_stale_index_caches( - index_id, - &self.access_layer, - self.cache_manager.as_ref(), - self.write_cache.as_ref(), - ) - .await; - self.listener.on_index_build_abort(region_file_id).await; - if stale == IndexPublicationStale::SchemaChanged { - // An index-unrelated schema change also invalidates the - // generation fence. Retry to avoid leaving the SST - // unindexed; per-SST coalescing limits this to one active - // and one pending task, so repeated changes cannot storm. - let retry = BuildIndexRequest { - region_id: self.region_id, - build_type: IndexBuildType::SchemaChange, - file_metas: vec![self.source.file_meta.clone()], - }; - let worker_request = WorkerRequest::Background { - region_id: self.region_id, - notify: BackgroundNotify::IndexBuildRetry(retry), - }; - let _ = self - .request_sender - .send(WorkerRequestWithTime::new(worker_request)) - .await; - } - return Ok(IndexBuildOutcome::Aborted(format!( - "Index build source changed before publication, region: {}, file_id: {}", - self.region_id, self.source.file_meta.file_id - ))); - } - Err(e) => { - let err = Arc::new(e); - WorkerRequest::Background { - region_id: self.region_id, - notify: BackgroundNotify::IndexBuildFailed(IndexBuildFailed { err }), - } - } - }; + result_senders: self.result_senders.clone(), + })) + .await; + Ok(IndexBuildReport::Deferred) + } + IndexPublication::Stale(stale) => Ok(self + .handle_stale_publication(region_file_id, stale, new_index_version) + .await), + } + } - let _ = self - .request_sender - .send(WorkerRequestWithTime::new(worker_request)) + async fn handle_stale_publication( + &self, + region_file_id: RegionFileId, + stale: IndexPublicationStale, + new_index_version: u64, + ) -> IndexBuildReport { + INDEX_PUBLICATION_STALE_TOTAL + .with_label_values(&["manifest_commit"]) + .inc(); + // If no successor publishes the same version, GC collects the + // remote artifact. Repartition requires GC, while local mode + // accepts this narrow orphan window when GC is disabled. + cleanup_stale_index_caches( + RegionIndexId::new(region_file_id, new_index_version), + &self.access_layer, + self.cache_manager.as_ref(), + self.write_cache.as_ref(), + ) + .await; + if stale != IndexPublicationStale::SchemaChanged { + return self + .abort_and_notify( + region_file_id, + "Index build source changed before publication", + ) .await; } - Ok(IndexBuildOutcome::Finished) + + self.listener.on_index_build_abort(region_file_id).await; + // An index-unrelated schema change also invalidates the generation + // fence. Retry to avoid leaving the SST unindexed; per-SST coalescing + // limits this to one active and one pending task, so repeated changes + // cannot storm. + self.notify_worker(BackgroundNotify::IndexBuildRetry(BuildIndexRequest { + region_id: self.region_id, + build_type: IndexBuildType::SchemaChange, + file_metas: vec![self.source.file_meta.clone()], + result_senders: self.result_senders.clone(), + })) + .await; + IndexBuildReport::Deferred } async fn maybe_upload_index_file( @@ -1131,19 +1238,26 @@ struct PendingIndexBuilds { } impl PendingIndexBuilds { - fn insert(&mut self, pending: PendingIndexBuild) -> Option { + /// Inserts a pending task and returns the task it replaces or that is dropped. + /// + /// The surviving task takes over the result senders of the returned task, so + /// callers waiting on either task wait for the build that actually runs. + fn insert(&mut self, mut pending: PendingIndexBuild) -> Option { let file_id = pending.task.source.file_meta.file_id; match self.tasks.entry(file_id) { std::collections::hash_map::Entry::Vacant(entry) => { entry.insert(pending); None } - std::collections::hash_map::Entry::Occupied(mut entry) - if pending.supersedes(entry.get()) => - { - Some(entry.insert(pending).task) + std::collections::hash_map::Entry::Occupied(mut entry) => { + let survivor = entry.get_mut(); + if pending.supersedes(survivor) { + std::mem::swap(survivor, &mut pending); + } + let senders = std::mem::take(&mut pending.task.result_senders); + survivor.task.result_senders.extend(senders); + Some(pending.task) } - std::collections::hash_map::Entry::Occupied(_) => Some(pending.task), } } @@ -1227,7 +1341,7 @@ impl IndexBuildScheduler { &mut self, version_control: &VersionControlRef, task: IndexBuildTask, - ) -> Result<()> { + ) { let status = self .region_status .entry(task.region_id) @@ -1238,7 +1352,13 @@ impl IndexBuildScheduler { task.source.file_meta.region_id, task.source.file_meta.file_id, ); - let can_wait_for_active = status.retiring || task.reason == IndexBuildType::SchemaChange; + // A manual build must not report success while an active build that may + // not create the requested indexes is still running, so it waits as well. + let can_wait_for_active = status.retiring + || matches!( + task.reason, + IndexBuildType::SchemaChange | IndexBuildType::Manual + ); let (rejected, coalesced) = if status.building_files.contains(&file_id) && !can_wait_for_active { (Some(task), None) @@ -1285,7 +1405,6 @@ impl IndexBuildScheduler { if should_schedule { self.schedule_next_build_batch(); } - Ok(()) } /// Schedule tasks until reaching the files limit or no more tasks. @@ -1370,25 +1489,6 @@ impl IndexBuildScheduler { self.schedule_next_build_batch(); } - pub(crate) async fn on_failure(&mut self, region_id: RegionId, err: Arc) { - let Some(status) = self.region_status.get_mut(®ion_id) else { - error!(err; "Index build failed after scheduler state was removed, region: {}", region_id); - return; - }; - if status.retiring { - error!(err; "Index build from a retiring region incarnation failed, region: {}", region_id); - return; - } - error!( - err; "Index build scheduler encountered failure for region {}, failing pending tasks.", - region_id - ); - status.fail_pending(err).await; - if status.building_files.is_empty() { - self.region_status.remove(®ion_id); - } - } - /// Notifies the scheduler that the region is dropped. pub(crate) async fn on_region_dropped(&mut self, region_id: RegionId) { self.retire_region( @@ -1470,6 +1570,19 @@ mod tests { with_skipping_bloom: bool, } + /// Receives the published index that the worker applies before notifying waiters. + async fn recv_index_build_finished( + rx: &mut mpsc::Receiver, + ) -> IndexBuildFinished { + match rx.recv().await.unwrap().request { + WorkerRequest::Background { + notify: BackgroundNotify::IndexBuildFinished(finished), + .. + } => finished, + other => panic!("Expect index build finished, got: {other:?}"), + } + } + async fn seed_manifest_file(manifest_ctx: &ManifestContextRef, file_meta: &FileMeta) { manifest_ctx .update_manifest( @@ -1918,14 +2031,12 @@ mod tests { file_purger, indexer_builder, request_sender: tx, - result_sender: result_tx, + index_options: IndexOptions::default(), + result_senders: vec![result_tx], }; // Schedule the build task and check result. - scheduler - .schedule_build(&version_control, task) - .await - .unwrap(); + scheduler.schedule_build(&version_control, task).await; match result_rx.recv().await.unwrap() { Ok(outcome) => { if outcome == IndexBuildOutcome::Finished { @@ -1986,7 +2097,7 @@ mod tests { // Create mock task. let (tx, mut rx) = mpsc::channel(4); - let (result_tx, mut result_rx) = mpsc::channel::>(4); + let (result_tx, _result_rx) = mpsc::channel::>(4); let task = IndexBuildTask { region_id, file, @@ -2004,21 +2115,11 @@ mod tests { file_purger, indexer_builder, request_sender: tx, - result_sender: result_tx, + index_options: IndexOptions::default(), + result_senders: vec![result_tx], }; - scheduler - .schedule_build(&version_control, task) - .await - .unwrap(); - - // The task should finish successfully. - match result_rx.recv().await.unwrap() { - Ok(outcome) => { - assert_eq!(outcome, IndexBuildOutcome::Finished); - } - _ => panic!("Expect finished result"), - } + scheduler.schedule_build(&version_control, task).await; // A notification should be sent to the worker to update the manifest. let worker_req = rx.recv().await.unwrap().request; @@ -2098,8 +2199,8 @@ mod tests { let file = FileHandle::new(file_meta.clone(), file_purger.clone()); // Create mock task. - let (tx, _rx) = mpsc::channel(4); - let (result_tx, mut result_rx) = mpsc::channel::>(4); + let (tx, mut rx) = mpsc::channel(4); + let (result_tx, _result_rx) = mpsc::channel::>(4); let task = IndexBuildTask { region_id, file, @@ -2117,13 +2218,11 @@ mod tests { file_purger, indexer_builder, request_sender: tx, - result_sender: result_tx, + index_options: IndexOptions::default(), + result_senders: vec![result_tx], }; - scheduler - .schedule_build(&version_control, task) - .await - .unwrap(); + scheduler.schedule_build(&version_control, task).await; let puffin_path = location::index_file_path( env.access_layer.table_dir(), @@ -2151,13 +2250,8 @@ mod tests { ); } - // The task should finish successfully. - match result_rx.recv().await.unwrap() { - Ok(outcome) => { - assert_eq!(outcome, IndexBuildOutcome::Finished); - } - _ => panic!("Expect finished result"), - } + // The worker is notified to apply the published index. + recv_index_build_finished(&mut rx).await; // The index file should exist after the task finishes. assert!( @@ -2228,24 +2322,26 @@ mod tests { file_purger, indexer_builder, request_sender: tx, - result_sender: result_tx, + index_options: IndexOptions::default(), + result_senders: vec![result_tx], }; - scheduler - .schedule_build(&version_control, task) - .await - .unwrap(); + scheduler.schedule_build(&version_control, task).await; - // The task should finish successfully. - match result_rx.recv().await.unwrap() { - Ok(outcome) => { - assert_eq!(outcome, IndexBuildOutcome::Finished); + // No index needs to be built, so the SST is not read. + assert!(matches!( + result_rx.recv().await.unwrap(), + Ok(IndexBuildOutcome::Aborted(reason)) if reason.contains("No index to build") + )); + + // Only the stop notification reaches the worker. + assert!(matches!( + rx.recv().await.unwrap().request, + WorkerRequest::Background { + notify: BackgroundNotify::IndexBuildStopped(_), + .. } - _ => panic!("Expect finished result"), - } - - // No index is built, so no notification should be sent to the worker. - let _ = rx.recv().await.is_none(); + )); } #[tokio::test] @@ -2306,8 +2402,8 @@ mod tests { let file = FileHandle::new(file_meta.clone(), file_purger.clone()); // Create mock task. - let (tx, mut _rx) = mpsc::channel(4); - let (result_tx, mut result_rx) = mpsc::channel::>(4); + let (tx, mut rx) = mpsc::channel(4); + let (result_tx, _result_rx) = mpsc::channel::>(4); let task = IndexBuildTask { region_id, file, @@ -2325,21 +2421,13 @@ mod tests { file_purger, indexer_builder, request_sender: tx, - result_sender: result_tx, + index_options: IndexOptions::default(), + result_senders: vec![result_tx], }; - scheduler - .schedule_build(&version_control, task) - .await - .unwrap(); + scheduler.schedule_build(&version_control, task).await; - // The task should finish successfully. - match result_rx.recv().await.unwrap() { - Ok(outcome) => { - assert_eq!(outcome, IndexBuildOutcome::Finished); - } - _ => panic!("Expect finished result"), - } + recv_index_build_finished(&mut rx).await; // The write cache should contain the uploaded index file. let index_key = IndexKey::new( @@ -2400,7 +2488,8 @@ mod tests { file_purger, indexer_builder, request_sender: tx, - result_sender: result_tx, + index_options: IndexOptions::default(), + result_senders: vec![result_tx], }; (task, result_rx) } @@ -2432,10 +2521,7 @@ mod tests { IndexBuildType::Manual, ) .await; - scheduler - .schedule_build(&version_control, active) - .await - .unwrap(); + scheduler.schedule_build(&version_control, active).await; assert_eq!(job_scheduler.num_jobs(), 1); let (mut older, mut older_rx) = create_mock_task_for_schedule_with_result( @@ -2446,10 +2532,7 @@ mod tests { ) .await; older.source.schema_version = 1; - scheduler - .schedule_build(&version_control, older) - .await - .unwrap(); + scheduler.schedule_build(&version_control, older).await; let (mut latest, mut latest_rx) = create_mock_task_for_schedule_with_result( &env, @@ -2459,27 +2542,23 @@ mod tests { ) .await; latest.source.schema_version = 2; - scheduler - .schedule_build(&version_control, latest) - .await - .unwrap(); + scheduler.schedule_build(&version_control, latest).await; - let replaced = tokio::time::timeout(std::time::Duration::from_secs(5), older_rx.recv()) - .await - .expect("replaced pending task result sender was not completed") - .expect("replaced pending task result channel closed"); - assert!(matches!( - replaced, - Ok(IndexBuildOutcome::Aborted(reason)) if reason.contains("coalesced") - )); - assert!(matches!( - latest_rx.try_recv(), - Err(mpsc::error::TryRecvError::Empty) - )); + // The replaced task's caller waits for the surviving build instead of + // being told the replaced build has completed. + for rx in [&mut older_rx, &mut latest_rx] { + assert!(matches!( + rx.try_recv(), + Err(mpsc::error::TryRecvError::Empty) + )); + } let status = &scheduler.region_status[®ion_id]; assert_eq!(status.building_files.len(), 1); assert_eq!(status.pending_tasks.len(), 1); + let survivor = &status.pending_tasks.tasks[&file_id].task; + assert_eq!(survivor.source.schema_version, 2); + assert_eq!(survivor.result_senders.len(), 2); assert_eq!(job_scheduler.num_jobs(), 1); scheduler.on_task_stopped(region_id, file_id); @@ -2517,10 +2596,7 @@ mod tests { ) .await; - scheduler - .schedule_build(&version_control, task) - .await - .unwrap(); + scheduler.schedule_build(&version_control, task).await; let result = tokio::time::timeout(std::time::Duration::from_secs(5), result_rx.recv()) .await @@ -2563,12 +2639,7 @@ mod tests { // Test 1: Basic scheduling let task1 = create_mock_task_for_schedule(&env, file_id1, region_id, IndexBuildType::Flush).await; - assert!( - scheduler - .schedule_build(&version_control, task1) - .await - .is_ok() - ); + scheduler.schedule_build(&version_control, task1).await; assert!(scheduler.region_status.contains_key(®ion_id)); let status = scheduler.region_status.get(®ion_id).unwrap(); assert_eq!(status.building_files.len(), 1); @@ -2577,20 +2648,14 @@ mod tests { // Test 2: Duplicate file scheduling (should be skipped) let task1_dup = create_mock_task_for_schedule(&env, file_id1, region_id, IndexBuildType::Flush).await; - scheduler - .schedule_build(&version_control, task1_dup) - .await - .unwrap(); + scheduler.schedule_build(&version_control, task1_dup).await; let status = scheduler.region_status.get(®ion_id).unwrap(); assert_eq!(status.building_files.len(), 1); // Still only one // Test 3: Fill up to limit (2 building tasks) let task2 = create_mock_task_for_schedule(&env, file_id2, region_id, IndexBuildType::Flush).await; - scheduler - .schedule_build(&version_control, task2) - .await - .unwrap(); + scheduler.schedule_build(&version_control, task2).await; let status = scheduler.region_status.get(®ion_id).unwrap(); assert_eq!(status.building_files.len(), 2); // Reached limit assert_eq!(status.pending_tasks.len(), 0); @@ -2605,18 +2670,9 @@ mod tests { let task5 = create_mock_task_for_schedule(&env, file_id5, region_id, IndexBuildType::Manual).await; - scheduler - .schedule_build(&version_control, task3) - .await - .unwrap(); - scheduler - .schedule_build(&version_control, task4) - .await - .unwrap(); - scheduler - .schedule_build(&version_control, task5) - .await - .unwrap(); + scheduler.schedule_build(&version_control, task3).await; + scheduler.schedule_build(&version_control, task4).await; + scheduler.schedule_build(&version_control, task5).await; let status = scheduler.region_status.get(®ion_id).unwrap(); assert_eq!(status.building_files.len(), 2); // Still at limit @@ -2651,53 +2707,6 @@ mod tests { // Region should be removed when all tasks complete assert!(!scheduler.region_status.contains_key(®ion_id)); - - // Test 8: A build failure keeps active leases without retiring the region - let task6 = - create_mock_task_for_schedule(&env, file_id1, region_id, IndexBuildType::Flush).await; - let task7 = - create_mock_task_for_schedule(&env, file_id2, region_id, IndexBuildType::Flush).await; - let task8 = - create_mock_task_for_schedule(&env, file_id3, region_id, IndexBuildType::Manual).await; - - scheduler - .schedule_build(&version_control, task6) - .await - .unwrap(); - scheduler - .schedule_build(&version_control, task7) - .await - .unwrap(); - scheduler - .schedule_build(&version_control, task8) - .await - .unwrap(); - - assert!(scheduler.region_status.contains_key(®ion_id)); - let status = scheduler.region_status.get(®ion_id).unwrap(); - assert_eq!(status.building_files.len(), 2); - assert_eq!(status.pending_tasks.len(), 1); - - scheduler - .on_failure( - region_id, - Arc::new( - crate::error::UnexpectedSnafu { - reason: "index build failed".to_string(), - } - .build(), - ), - ) - .await; - let status = scheduler.region_status.get(®ion_id).unwrap(); - assert!(!status.retiring); - assert_eq!(status.building_files.len(), 2); - assert!(status.pending_tasks.is_empty()); - - scheduler.on_task_stopped(region_id, file_id1); - assert!(scheduler.region_status.contains_key(®ion_id)); - scheduler.on_task_stopped(region_id, file_id2); - assert!(!scheduler.region_status.contains_key(®ion_id)); } /// Helper to set up a scheduler with files_limit=1 and 3 scheduled tasks, @@ -2780,18 +2789,9 @@ mod tests { ) .await; - scheduler - .schedule_build(&version_control, task1) - .await - .unwrap(); - scheduler - .schedule_build(&version_control, task2) - .await - .unwrap(); - scheduler - .schedule_build(&version_control, task3) - .await - .unwrap(); + scheduler.schedule_build(&version_control, task1).await; + scheduler.schedule_build(&version_control, task2).await; + scheduler.schedule_build(&version_control, task3).await; // Verify: 1 building + 2 pending. assert!(scheduler.region_status.contains_key(®ion_id)); @@ -2900,20 +2900,12 @@ mod tests { .await; scheduler .schedule_build(&version_control, reopened_task) - .await - .unwrap(); + .await; let status = scheduler.region_status.get(®ion_id).unwrap(); assert!(status.retiring); assert_eq!(status.building_files.len(), 1); assert_eq!(status.pending_tasks.len(), 1); - // A manifest error from the old incarnation must not discard the - // reopened task before the old build reports stopped. - scheduler - .on_failure(region_id, Arc::new(RegionClosedSnafu { region_id }.build())) - .await; - assert_eq!(scheduler.region_status[®ion_id].pending_tasks.len(), 1); - scheduler.on_task_stopped(region_id, building_file_id); let status = scheduler.region_status.get(®ion_id).unwrap(); assert!(!status.retiring); diff --git a/src/mito2/src/worker.rs b/src/mito2/src/worker.rs index 556f4a0757f..c92070ac4e7 100644 --- a/src/mito2/src/worker.rs +++ b/src/mito2/src/worker.rs @@ -1435,9 +1435,6 @@ impl RegionWorkerLoop { BackgroundNotify::IndexBuildStopped(req) => { self.handle_index_build_stopped(region_id, req).await } - BackgroundNotify::IndexBuildFailed(req) => { - self.handle_index_build_failed(region_id, req).await - } BackgroundNotify::IndexBuildRetry(req) => { self.handle_rebuild_index(req, OptionOutputTx::new(None)) .await diff --git a/src/mito2/src/worker/handle_compaction.rs b/src/mito2/src/worker/handle_compaction.rs index 50147584bcd..269cf258340 100644 --- a/src/mito2/src/worker/handle_compaction.rs +++ b/src/mito2/src/worker/handle_compaction.rs @@ -79,6 +79,7 @@ impl RegionWorkerLoop { region_id, build_type: IndexBuildType::Compact, file_metas: edit.files_to_add, + result_senders: Vec::new(), }, OptionOutputTx::new(None), ) @@ -245,6 +246,7 @@ impl RegionWorkerLoop { region_id, build_type: IndexBuildType::Compact, file_metas: index_build_file_metas, + result_senders: Vec::new(), }, OptionOutputTx::new(None), ) diff --git a/src/mito2/src/worker/handle_flush.rs b/src/mito2/src/worker/handle_flush.rs index 88d911dccc3..f09efe81e47 100644 --- a/src/mito2/src/worker/handle_flush.rs +++ b/src/mito2/src/worker/handle_flush.rs @@ -443,10 +443,13 @@ impl RegionWorkerLoop { if self.config.index.build_mode == IndexBuildMode::Sync { // Sync flush builds indexes while writing SSTs, except for encoded bulk parts, // which are uploaded as-is. Build the missing indexes for those files here. - let metadata = region.metadata(); + let version = region.version(); index_build_file_metas.retain(|meta| { request.unindexed_files.contains(&meta.file_id) - && !meta.is_index_consistent_with_region(&metadata.column_metadatas) + && !meta.is_index_consistent_with_region( + &version.metadata.column_metadatas, + &version.options.index_options, + ) }); } @@ -456,6 +459,7 @@ impl RegionWorkerLoop { region_id, build_type: IndexBuildType::Flush, file_metas: index_build_file_metas, + result_senders: Vec::new(), }, OptionOutputTx::new(None), ) diff --git a/src/mito2/src/worker/handle_manifest.rs b/src/mito2/src/worker/handle_manifest.rs index 8aac67dc96f..373080aee90 100644 --- a/src/mito2/src/worker/handle_manifest.rs +++ b/src/mito2/src/worker/handle_manifest.rs @@ -197,6 +197,7 @@ impl RegionWorkerLoop { region_id: region.region_id, build_type: IndexBuildType::SchemaChange, file_metas: Vec::new(), + result_senders: Vec::new(), }, OptionOutputTx::new(None), ) diff --git a/src/mito2/src/worker/handle_rebuild_index.rs b/src/mito2/src/worker/handle_rebuild_index.rs index 17ba1ece1c8..927a5866a80 100644 --- a/src/mito2/src/worker/handle_rebuild_index.rs +++ b/src/mito2/src/worker/handle_rebuild_index.rs @@ -18,21 +18,21 @@ use std::collections::HashMap; use std::sync::Arc; use api::v1::region::build_index_request; -use common_telemetry::{debug, error, warn}; +use common_telemetry::{debug, warn}; use store_api::metadata::RegionMetadataRef; use store_api::region_request::RegionBuildIndexRequest; use store_api::storage::{FileId, RegionId}; use tokio::sync::mpsc; use crate::cache::CacheStrategy; -use crate::error::{InvalidRequestSnafu, Result, WorkerStoppedSnafu}; +use crate::error::{ + IndexBuildIncompleteSnafu, InvalidRequestSnafu, RegionStateSnafu, Result, WorkerStoppedSnafu, +}; use crate::manifest::action::RegionEdit; use crate::metrics::INDEX_PUBLICATION_STALE_TOTAL; use crate::region::options::IndexOptions; -use crate::region::{IndexBuildSource, MitoRegionRef}; -use crate::request::{ - BuildIndexRequest, IndexBuildFailed, IndexBuildFinished, IndexBuildStopped, OptionOutputTx, -}; +use crate::region::{IndexBuildSource, MitoRegionRef, RegionLeaderState, RegionRoleState}; +use crate::request::{BuildIndexRequest, IndexBuildFinished, IndexBuildStopped, OptionOutputTx}; use crate::sst::file::{FileHandle, FileMeta, RegionFileId, RegionIndexId}; use crate::sst::index::{ IndexBuildOutcome, IndexBuildTask, IndexBuildType, IndexerBuilderImpl, ResultMpscSender, @@ -57,7 +57,7 @@ impl RegionWorkerLoop { file: FileHandle, file_meta: FileMeta, build_type: IndexBuildType, - result_sender: ResultMpscSender, + result_senders: Vec, ) -> IndexBuildTask { let access_layer = region.access_layer.clone(); @@ -90,6 +90,7 @@ impl RegionWorkerLoop { file: file.clone(), target_region_metadata: snapshot.metadata.clone(), source: IndexBuildSource::new(file_meta, snapshot.metadata.schema_version), + index_options: snapshot.index_options.clone(), reason: build_type, access_layer: access_layer.clone(), listener: self.listener.clone(), @@ -99,7 +100,7 @@ impl RegionWorkerLoop { file_purger: file.file_purger(), request_sender: self.sender.clone(), indexer_builder: indexer_builder_ref.clone(), - result_sender, + result_senders, } } @@ -120,6 +121,7 @@ impl RegionWorkerLoop { region_id, build_type: IndexBuildType::Manual, file_metas: Vec::new(), + result_senders: Vec::new(), }, sender, ) @@ -168,6 +170,18 @@ impl RegionWorkerLoop { let Some(region) = self.regions.writable_region_or(region_id, &mut sender) else { return; }; + // Staging regions are writable but must not publish indexes. + if region.should_abort_index() { + sender.send( + RegionStateSnafu { + region_id, + state: region.state(), + expect: RegionRoleState::Leader(RegionLeaderState::Writable), + } + .fail(), + ); + return; + } // Only keep the parts of the version that index builds need, so the // memtables and SST version are not pinned across the awaits below. @@ -204,13 +218,20 @@ impl RegionWorkerLoop { let build_tasks = if request.file_metas.is_empty() { // If no specific files are provided, find files whose index is inconsistent with the region metadata. + // Check the version rather than the manifest: a publication committed + // to the manifest but not yet applied is still invisible to queries. + // Its build still holds the file, so the manual task waits for it. all_files .values() + .filter(|file| { + !file.meta_ref().is_index_consistent_with_region( + &snapshot.metadata.column_metadatas, + &snapshot.index_options, + ) + }) .filter_map(|file| { let file_meta = current_file_meta(file.meta_ref().file_id)?; - (!file_meta - .is_index_consistent_with_region(&snapshot.metadata.column_metadatas)) - .then(|| (file.clone(), file_meta)) + Some((file.clone(), file_meta)) }) .collect::>() } else { @@ -234,9 +255,13 @@ impl RegionWorkerLoop { return; } - let num_tasks = build_tasks.len(); - let (tx, mut rx) = mpsc::channel::>(num_tasks); - + let file_ids = build_tasks + .iter() + .map(|(file, _)| file.meta_ref().file_id) + .collect::>(); + let build_type = request.build_type.clone(); + // Each task holds a sender and reports at most once, so sends never block. + let (tx, mut rx) = mpsc::channel::>(build_tasks.len()); for (file_handle, file_meta) in build_tasks { debug!( "Scheduling index build for region {}, file_id {}", @@ -244,6 +269,7 @@ impl RegionWorkerLoop { file_handle.meta_ref().file_id ); + // `set_region_role` can change the state while tasks are scheduled. if region.should_abort_index() { warn!( "Region {} is in state {:?}, abort index rebuild process for file_id {}", @@ -254,29 +280,79 @@ impl RegionWorkerLoop { break; } + let mut result_senders = request.result_senders.clone(); + result_senders.push(tx.clone()); let task = self.new_index_build_task( ®ion, &snapshot, - file_handle.clone(), + file_handle, file_meta, request.build_type.clone(), - tx.clone(), + result_senders, ); - let _ = self - .index_build_scheduler + self.index_build_scheduler .schedule_build(®ion.version_control, task) .await; } + drop(tx); + // Wait for all index build tasks to finish and notify the caller. common_runtime::spawn_global(async move { - for _ in 0..num_tasks { - if let Some(Err(e)) = rx.recv().await { - warn!(e; "Index build task failed for region: {}", region_id); - sender.send(Err(e)); - return; + let mut abort_reasons = Vec::new(); + // The channel closes once every task has reported or been dropped. + while let Some(result) = rx.recv().await { + match result { + Ok(IndexBuildOutcome::Finished) => {} + Ok(IndexBuildOutcome::Aborted(reason)) => abort_reasons.push(reason), + Err(e) => { + warn!(e; "Index build task failed for region: {}", region_id); + sender.send(Err(e)); + return; + } } } - sender.send(Ok(0)); + + // The build covers the SSTs that existed when it was requested. A file + // removed meanwhile, e.g. replaced by compaction, counts as done: its + // replacement is indexed according to the compaction index options, + // as any later compaction would be. + let manifest = region.manifest_ctx.manifest().await; + let (present, removed): (Vec<_>, Vec<_>) = file_ids + .into_iter() + .partition(|file_id| manifest.files.contains_key(file_id)); + if !removed.is_empty() && build_type == IndexBuildType::Manual { + warn!( + "Files were removed during the manual index build, their replacements follow the compaction index options, region: {}, files: {:?}", + region_id, removed + ); + } + let incomplete = present + .into_iter() + .filter(|file_id| { + !manifest.files[file_id].is_index_consistent_with_region( + &snapshot.metadata.column_metadatas, + &snapshot.index_options, + ) + }) + .collect::>(); + let result = if incomplete.is_empty() { + Ok(0) + } else { + // Background builds have no caller, and files they leave + // unindexed by configuration are expected. + if build_type == IndexBuildType::Manual { + warn!( + "Manual index build did not complete, region: {}, files: {:?}, aborted: {:?}", + region_id, incomplete, abort_reasons + ); + } + IndexBuildIncompleteSnafu { + region_id, + file_ids: incomplete, + } + .fail() + }; + sender.send(result); }); } @@ -324,7 +400,10 @@ impl RegionWorkerLoop { } drop(manager); - if !is_current { + let outcome = if is_current { + self.listener.on_index_build_finish(region_file_id).await; + IndexBuildOutcome::Finished + } else { INDEX_PUBLICATION_STALE_TOTAL .with_label_values(&["worker_apply"]) .inc(); @@ -340,21 +419,18 @@ impl RegionWorkerLoop { ) .await; self.listener.on_index_build_abort(region_file_id).await; - return; + IndexBuildOutcome::Aborted(format!( + "Index publication was superseded before worker apply, region: {}, file_id: {}", + region_id, file_meta.file_id + )) + }; + // Notify waiters after the version is updated, so queries issued after a + // manual build returns can use the new index. `try_send` keeps the worker + // loop from waiting on a caller; a missing result is caught by the + // waiter's manifest check. + for sender in &request.result_senders { + let _ = sender.try_send(Ok(outcome.clone())); } - - self.listener.on_index_build_finish(region_file_id).await; - } - - pub(crate) async fn handle_index_build_failed( - &mut self, - region_id: RegionId, - request: IndexBuildFailed, - ) { - error!(request.err; "Index build failed for region: {}", region_id); - self.index_build_scheduler - .on_failure(region_id, request.err.clone()) - .await; } pub(crate) async fn handle_index_build_stopped( diff --git a/tests/cases/standalone/common/function/admin/build_index_table_restart.result b/tests/cases/standalone/common/function/admin/build_index_table_restart.result index 62b0a1c084a..0004d00ed1b 100644 --- a/tests/cases/standalone/common/function/admin/build_index_table_restart.result +++ b/tests/cases/standalone/common/function/admin/build_index_table_restart.result @@ -91,10 +91,6 @@ ADMIN BUILD_INDEX('build_index_restart_test'); | 0 | +-----------------------------------------------+ --- ALTER may already schedule an async schema-change rebuild. If ADMIN BUILD_INDEX --- races with it, the manual task can be duplicate-aborted immediately while the --- schema-change task is still finishing, so wait before checking query visibility. --- SQLNESS SLEEP 1s -- Fulltext index built, verify via fulltext query SELECT msg FROM build_index_restart_test WHERE MATCHES(msg, 'fox') ORDER BY ts; @@ -135,6 +131,8 @@ AND index_type LIKE 'fulltext%'; | 1 | +---------------------------+ +-- Distributed region statistics come from datanode heartbeats. +-- SQLNESS SLEEP 1s -- Fulltext index size is reported after the build SELECT index_size FROM INFORMATION_SCHEMA.REGION_STATISTICS WHERE table_id = ( @@ -200,7 +198,6 @@ ADMIN BUILD_INDEX('build_index_restart_test'); | 0 | +-----------------------------------------------+ --- SQLNESS SLEEP 1s -- Fulltext query still works after idempotent rebuild SELECT msg FROM build_index_restart_test WHERE MATCHES(msg, 'fox') ORDER BY ts; diff --git a/tests/cases/standalone/common/function/admin/build_index_table_restart.sql b/tests/cases/standalone/common/function/admin/build_index_table_restart.sql index 745e98eee2c..f13ede93be0 100644 --- a/tests/cases/standalone/common/function/admin/build_index_table_restart.sql +++ b/tests/cases/standalone/common/function/admin/build_index_table_restart.sql @@ -45,10 +45,6 @@ ALTER TABLE build_index_restart_test MODIFY COLUMN msg SET FULLTEXT INDEX; ADMIN BUILD_INDEX('build_index_restart_test'); --- ALTER may already schedule an async schema-change rebuild. If ADMIN BUILD_INDEX --- races with it, the manual task can be duplicate-aborted immediately while the --- schema-change task is still finishing, so wait before checking query visibility. --- SQLNESS SLEEP 1s -- Fulltext index built, verify via fulltext query SELECT msg FROM build_index_restart_test WHERE MATCHES(msg, 'fox') ORDER BY ts; @@ -66,6 +62,8 @@ WHERE table_id = ( ) AND index_type LIKE 'fulltext%'; +-- Distributed region statistics come from datanode heartbeats. +-- SQLNESS SLEEP 1s -- Fulltext index size is reported after the build SELECT index_size FROM INFORMATION_SCHEMA.REGION_STATISTICS WHERE table_id = ( @@ -96,7 +94,6 @@ AND index_type LIKE 'fulltext%'; -- Run ADMIN BUILD_INDEX again (idempotency) ADMIN BUILD_INDEX('build_index_restart_test'); --- SQLNESS SLEEP 1s -- Fulltext query still works after idempotent rebuild SELECT msg FROM build_index_restart_test WHERE MATCHES(msg, 'fox') ORDER BY ts; From 9e9346d2ff2ba754fb4d6740ec1be22b7c2e1800 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Thu, 22 Oct 2026 03:36:02 +0000 Subject: [PATCH 065/239] fix: preserve datanode write batches when retrying (#9720) * fix: preserve datanode write batches when retrying Signed-off-by: WenyXu * fix: isolate invalid entries in metric write batches Signed-off-by: WenyXu * refactor: adapt batch writes to the unified metric API Signed-off-by: WenyXu * fix: explicitly accept partial results for batched retries Signed-off-by: WenyXu * test: assert terminal retry hint instead of masked error text Signed-off-by: WenyXu * test: preserve application status in region failure fixtures Signed-off-by: WenyXu * fix: refresh routes after region close Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- src/common/meta/src/error.rs | 5 + src/mito2/src/error.rs | 16 ++ src/mito2/src/flush.rs | 29 +- src/operator/src/write_dispatcher.rs | 109 ++++++-- src/operator/src/write_dispatcher/tests.rs | 302 ++++++++++++++------- 5 files changed, 329 insertions(+), 132 deletions(-) diff --git a/src/common/meta/src/error.rs b/src/common/meta/src/error.rs index c048f499dff..7dd7b9b0a62 100644 --- a/src/common/meta/src/error.rs +++ b/src/common/meta/src/error.rs @@ -41,6 +41,9 @@ pub use retry_hint::{retry_hint_from_postgres_error, retry_hint_from_postgres_po #[snafu(visibility(pub))] #[stack_trace_debug] pub enum Error { + #[snafu(display("Failed to execute batch request"))] + BatchRequest { source: Arc }, + #[snafu(display( "Region request has {} failed entries; {}", response.failures.len(), @@ -1216,6 +1219,7 @@ impl ErrorExt for Error { Error::PartialRegion { response } => response .first_failure() .map_or(StatusCode::Internal, |(_, error)| error.status_code()), + Error::BatchRequest { source } => source.status_code(), IllegalServerState { .. } | EtcdTxnOpResponse { .. } | EtcdFailed { .. } @@ -1386,6 +1390,7 @@ impl ErrorExt for Error { match self { Error::PartialRegion { .. } => RetryHint::Retryable, + Error::BatchRequest { source } => source.retry_hint(), RetryLater { source, .. } => match source.retry_hint() { RetryHint::RetryAfterRouteRefresh => RetryHint::RetryAfterRouteRefresh, RetryHint::Retryable | RetryHint::NonRetryable => RetryHint::Retryable, diff --git a/src/mito2/src/error.rs b/src/mito2/src/error.rs index d7ecdd88b1b..71253e78bea 100644 --- a/src/mito2/src/error.rs +++ b/src/mito2/src/error.rs @@ -1679,6 +1679,7 @@ impl ErrorExt for Error { match self { PartitionExprVersionMismatch { .. } | RegionNotFound { .. } + | RegionClosed { .. } | RejectWriteInStaging { .. } | RegionState { state: RegionRoleState::Follower, @@ -1819,6 +1820,21 @@ mod tests { StatusCode::RegionNotFound, RetryHint::RetryAfterRouteRefresh, ), + ( + RegionClosedSnafu { region_id }.build(), + StatusCode::Cancelled, + RetryHint::RetryAfterRouteRefresh, + ), + ( + RegionDroppedSnafu { region_id }.build(), + StatusCode::Cancelled, + RetryHint::NonRetryable, + ), + ( + RegionTruncatedSnafu { region_id }.build(), + StatusCode::Cancelled, + RetryHint::NonRetryable, + ), ] { assert_eq!(error.status_code(), status); assert_eq!(error.retry_hint(), hint); diff --git a/src/mito2/src/flush.rs b/src/mito2/src/flush.rs index 52eb311de73..fba0607235b 100644 --- a/src/mito2/src/flush.rs +++ b/src/mito2/src/flush.rs @@ -1679,7 +1679,7 @@ impl FlushStatus { #[cfg(test)] mod tests { use api::v1::{OpType, Rows}; - use common_error::ext::ErrorExt; + use common_error::ext::{ErrorExt, RetryHint}; use common_error::status_code::StatusCode; use datatypes::arrow::datatypes::Schema; use datatypes::arrow::record_batch::RecordBatch; @@ -2122,28 +2122,39 @@ mod tests { } #[tokio::test] - async fn test_region_closed_notifies_pending_bulk_writes() { + async fn test_region_closed_notifies_pending_requests() { let region_id = RegionId::new(1, 1); let version_control = Arc::new(VersionControlBuilder::new().build()); - let (bulk_req, output_rx) = new_test_bulk_request(region_id); + let (write_req, write_rx) = new_test_write_request(region_id); + let (bulk_req, bulk_rx) = new_test_bulk_request(region_id); + let (ddl_tx, ddl_rx) = oneshot::channel(); let status = FlushStatus { region_id, version_control, state: CancellableTaskState::new(), pending_task: None, closing: false, - pending_ddls: Vec::new(), - pending_writes: Vec::new(), + pending_ddls: vec![SenderDdlRequest { + region_id, + sender: OptionOutputTx::from(ddl_tx), + request: DdlRequest::Truncate( + store_api::region_request::RegionTruncateRequest::All, + ), + }], + pending_writes: vec![write_req], pending_bulk_writes: vec![bulk_req], }; status.on_region_closed(Arc::new(RegionClosedSnafu { region_id }.build())); - let err = output_rx - .await - .expect("pending bulk write must receive explicit error") - .unwrap_err(); + for rx in [write_rx, bulk_rx] { + let err = rx.await.unwrap().unwrap_err(); + assert_eq!(err.status_code(), StatusCode::Cancelled); + assert_eq!(err.retry_hint(), RetryHint::RetryAfterRouteRefresh); + } + let err = ddl_rx.await.unwrap().unwrap_err(); assert_eq!(err.status_code(), StatusCode::Cancelled); + assert_eq!(err.retry_hint(), RetryHint::RetryAfterRouteRefresh); } #[tokio::test] diff --git a/src/operator/src/write_dispatcher.rs b/src/operator/src/write_dispatcher.rs index d8a6ac9905e..2d12feb74c5 100644 --- a/src/operator/src/write_dispatcher.rs +++ b/src/operator/src/write_dispatcher.rs @@ -25,7 +25,9 @@ use backon::{BackoffBuilder, ExponentialBuilder}; use catalog::CatalogManagerRef; use common_error::ext::{BoxedError, ErrorExt, RetryHint}; use common_meta::cache_invalidator::{CacheInvalidatorRef, Context}; -use common_meta::error::{Error, ExternalSnafu, Result, TableInfoNotFoundSnafu, UnexpectedSnafu}; +use common_meta::error::{ + BatchRequestSnafu, Error, ExternalSnafu, Result, TableInfoNotFoundSnafu, UnexpectedSnafu, +}; use common_meta::instruction::CacheIdent; use common_meta::node_manager::NodeManagerRef; use common_meta::peer::{Peer, PeerRef}; @@ -33,7 +35,7 @@ use common_meta::rpc::router; use futures::future; use partition::manager::{PartitionRuleManager, PartitionRuleManagerRef}; use prometheus::Histogram; -use snafu::{OptionExt, ResultExt, ensure}; +use snafu::{IntoError, OptionExt, ResultExt, ensure}; use store_api::storage::RegionId; use table::table_name::TableName; use tokio_util::sync::CancellationToken; @@ -178,24 +180,17 @@ impl WriteDispatcher { cancellation: Option<&CancellationToken>, ) -> Result { let histogram = histogram.as_ref(); - let writes = if self.retry.is_some() { - writes - .into_iter() - .flat_map(|(peer, request)| { - split_request(request) - .into_iter() - .map(move |request| (Arc::clone(&peer), request)) - }) - .collect() - } else { - writes - }; - let attempts = future::join_all( - writes - .into_iter() - .map(|(peer, request)| self.send(peer, request, histogram)), - ) - .await; + let attempts = future::join_all(writes.into_iter().map(|(peer, request)| async move { + if self.retry.is_some() { + self.send_batch(peer, request, histogram).await + } else { + vec![self.send(peer, request, histogram).await] + } + })) + .await + .into_iter() + .flatten() + .collect(); self.finish(attempts, histogram, cancellation).await } @@ -267,6 +262,80 @@ impl WriteDispatcher { resolve_peers(&self.partition_manager, regions).await } + async fn send_batch( + &self, + peer: PeerRef, + mut request: RegionRequest, + histogram: Option<&Histogram>, + ) -> Vec { + if !matches!(&request.body, Some(Body::Inserts(_) | Body::Deletes(_))) { + return vec![self.send(peer, request, histogram).await]; + } + // Upgrade datanodes first and roll back frontends first: old datanodes + // cannot report which entries succeeded, so retrying a batch may replay writes. + api::region::set_accept_partial_result(request.header.get_or_insert_default()); + let request = Arc::new(request); + let requester = self.node_manager.datanode(&peer).await; + let result = { + let _timer = histogram.map(Histogram::start_timer); + requester.handle(Arc::clone(&request)).await + }; + let error = match result { + Ok(response) => return vec![WriteAttempt::Success(response)], + Err(error) => error, + }; + let request = match Arc::try_unwrap(request) { + Ok(request) => request, + Err(_) => { + return vec![WriteAttempt::Failed { + error: UnexpectedSnafu { + err_msg: "Region request is still shared after its RPC completed", + } + .build(), + request: None, + }]; + } + }; + let requests = split_request(request); + match error { + Error::PartialRegion { response } => { + let mut failures: HashMap<_, _> = response.failures.into_iter().collect(); + let mut attempts = vec![WriteAttempt::Success(RegionResponse { + affected_rows: response.affected_rows, + extensions: response.extensions, + metadata: Vec::new(), + })]; + for (index, request) in requests.into_iter().enumerate() { + if let Some(error) = failures.remove(&index) { + let request = error.retry_hint().is_retryable().then(|| RoutedWrite { + peer: Arc::clone(&peer), + request: Arc::new(request), + }); + attempts.push(WriteAttempt::Failed { error, request }); + } + } + attempts + } + error => { + // A whole-RPC failure cannot confirm success for any entry. + let error = Arc::new(error); + requests + .into_iter() + .map(|request| { + let request = error.retry_hint().is_retryable().then(|| RoutedWrite { + peer: Arc::clone(&peer), + request: Arc::new(request), + }); + WriteAttempt::Failed { + error: BatchRequestSnafu.into_error(Arc::clone(&error)), + request, + } + }) + .collect() + } + } + } + async fn dispatch_inner( &self, requests: Vec, diff --git a/src/operator/src/write_dispatcher/tests.rs b/src/operator/src/write_dispatcher/tests.rs index a249a8dfe8e..2607a3f194d 100644 --- a/src/operator/src/write_dispatcher/tests.rs +++ b/src/operator/src/write_dispatcher/tests.rs @@ -21,8 +21,8 @@ use std::time::{Duration, Instant}; use api::v1::region::region_server::{Region, RegionServer}; use api::v1::region::{ - AlignedSchemaVersion, BulkInsertRequest, DeleteRequest, InsertRequest, RegionRequestHeader, - RegionResponse as PbRegionResponse, bulk_insert_request, + AlignedSchemaVersion, BulkInsertRequest, DeleteRequest, DeleteRequests, InsertRequest, + InsertRequests, RegionRequestHeader, RegionResponse as PbRegionResponse, bulk_insert_request, }; use api::v1::value::ValueData; use api::v1::{ @@ -78,6 +78,8 @@ use crate::test_util; #[derive(Clone, Default)] struct RecordingRegionService { + request_sizes: Arc>>, + rpc_rejections: Arc>>, requests: Arc>>, rejections: Arc>>>, release: Arc>>>, @@ -90,6 +92,70 @@ impl Region for RecordingRegionService { request: Request, ) -> std::result::Result, Status> { let request = request.into_inner(); + if !matches!(&request.body, Some(Body::Inserts(_) | Body::Deletes(_))) { + return self.handle_entry(request).await; + } + let accepts_partial = request + .header + .as_ref() + .and_then(|header| header.query_context.as_ref()) + .and_then(|context| context.extensions.get(api::region::ACCEPT_PARTIAL_RESULT)) + .is_some_and(|value| value == "true"); + let entries = split_request(request); + let return_tonic_error = !accepts_partial || entries.len() == 1; + self.request_sizes.lock().unwrap().push(entries.len()); + if let Some(status) = self.rpc_rejections.lock().unwrap().pop_front() { + self.requests.lock().unwrap().extend(entries); + return Err(status); + } + let mut response = PbRegionResponse { + header: Some(ResponseHeader { + status: Some(PbStatus::default()), + }), + ..Default::default() + }; + let results = + future::join_all(entries.into_iter().map(|entry| self.handle_entry(entry))).await; + for (index, result) in results.into_iter().enumerate() { + match result { + Ok(result) => response.affected_rows += result.into_inner().affected_rows, + Err(status) => { + if return_tonic_error { + return Err(status); + } + let error = client::Error::from(status); + response + .failures + .push(api::v1::region::RegionRequestFailure { + request_index: index as u32, + status: Some(PbStatus { + status_code: error.status_code() as u32, + err_msg: error.output_msg(), + retry_hint: error.retry_hint().as_str().to_string(), + }), + }); + } + } + } + if !response.failures.is_empty() { + response + .header + .as_mut() + .unwrap() + .status + .as_mut() + .unwrap() + .status_code = StatusCode::PartialFailure as u32; + } + Ok(Response::new(response)) + } +} + +impl RecordingRegionService { + async fn handle_entry( + &self, + request: RegionRequest, + ) -> std::result::Result, Status> { let region_id = write_region_id(&request).unwrap().as_u64(); self.requests.lock().unwrap().push(request); let release = self.release.lock().unwrap().clone(); @@ -237,9 +303,19 @@ impl RegionTestServer { hints .into_iter() .map(|hint| { - let mut status = Status::failed_precondition(format!( - "region {region_id} rejected the write" - )); + let code = match hint { + RetryHint::Retryable => StatusCode::RegionBusy, + RetryHint::RetryAfterRouteRefresh => StatusCode::Cancelled, + RetryHint::NonRetryable => StatusCode::InvalidArguments, + }; + let mut status = Status::new( + common_error::status_code::status_to_tonic_code(code), + format!("region {region_id} rejected the write"), + ); + status.metadata_mut().insert( + common_error::GREPTIME_DB_HEADER_ERROR_CODE, + (code as u32).to_string().parse().unwrap(), + ); status.metadata_mut().insert( common_error::GREPTIME_DB_HEADER_ERROR_RETRY_HINT, hint.as_str().parse().unwrap(), @@ -343,18 +419,26 @@ fn build_insert_request(regions: &[u32]) -> RegionRequest { #[tokio::test] async fn test_write_retry_disabled_and_enabled_ablation() { common_telemetry::init_default_ut_logging(); - for run in 0..3 { + for regions in [vec![1], vec![1, 2]] { for max_retries in [0, 1] { let server = RegionTestServer::start().await; server.fail_region(1, [RetryHint::Retryable]); let dispatcher = server.create_dispatcher(max_retries).await; - let result = dispatcher.dispatch(vec![build_insert_request(&[1])]).await; + let result = dispatcher + .dispatch(vec![build_insert_request(®ions)]) + .await; let attempts = server.attempt_counts()[&RegionId::new(1, 1).as_u64()]; common_telemetry::info!( - "write_retry_ablation run={run} max_retries={max_retries} success={} attempts={attempts}", + "write_retry_ablation regions={regions:?} max_retries={max_retries} success={} attempts={attempts}", result.is_ok() ); - assert_eq!(result.is_ok(), max_retries > 0); + if max_retries == 0 { + let error = result.unwrap_err(); + assert_eq!(error.status_code(), StatusCode::RegionBusy); + assert_eq!(error.retry_hint(), RetryHint::Retryable); + } else { + assert_eq!(result.unwrap().affected_rows, regions.len()); + } assert_eq!(attempts, max_retries + 1); for request in server.service.requests.lock().unwrap().iter() { let accepts_partial = request @@ -372,22 +456,29 @@ async fn test_write_retry_disabled_and_enabled_ablation() { #[tokio::test] async fn test_write_retry_does_not_replay_successful_regions() { - let server = RegionTestServer::start().await; - server.fail_region(2, [RetryHint::Retryable]); - let dispatcher = server.create_dispatcher(2).await; - let response = dispatcher - .dispatch(vec![build_insert_request(&[1, 2])]) - .await - .unwrap(); - assert_eq!(response.affected_rows, 2); - assert_eq!( - server.attempt_counts(), - HashMap::from([ - (RegionId::new(1, 1).as_u64(), 1), - (RegionId::new(1, 2).as_u64(), 2), - ]) - ); - assert!(server.invalidator.calls.lock().unwrap().is_empty()); + for regions in [vec![1, 2], vec![1, 1, 1]] { + let server = RegionTestServer::start().await; + server.fail_region(1, [RetryHint::Retryable]); + let dispatcher = server.create_dispatcher(2).await; + let response = dispatcher + .dispatch(vec![build_insert_request(®ions)]) + .await + .unwrap(); + assert_eq!(response.affected_rows, regions.len()); + let mut expected = HashMap::new(); + for region in ®ions { + *expected + .entry(RegionId::new(1, *region).as_u64()) + .or_insert(0) += 1; + } + *expected.get_mut(&RegionId::new(1, 1).as_u64()).unwrap() += 1; + assert_eq!(server.attempt_counts(), expected); + assert_eq!( + *server.service.request_sizes.lock().unwrap(), + vec![regions.len(), 1] + ); + assert!(server.invalidator.calls.lock().unwrap().is_empty()); + } } #[tokio::test] @@ -429,10 +520,10 @@ async fn test_write_retry_transport_errors() { } server .service - .rejections + .rpc_rejections .lock() .unwrap() - .insert(region_id, VecDeque::from([status])); + .push_back(status); let dispatcher = server.create_dispatcher(2).await; let attempt = dispatcher .send( @@ -462,12 +553,12 @@ async fn test_write_retry_transport_attempt_budget() { for failures in [1, 3] { let server = RegionTestServer::start().await; let region_id = RegionId::new(1, 1).as_u64(); - server.service.rejections.lock().unwrap().insert( - region_id, - (0..failures) - .map(|_| Status::new(code, "write failed")) - .collect(), - ); + server + .service + .rpc_rejections + .lock() + .unwrap() + .extend((0..failures).map(|_| Status::new(code, "write failed"))); let dispatcher = server.create_dispatcher(2).await; let response = dispatcher.dispatch(vec![build_insert_request(&[1])]).await; let expected_attempts = if failures == 1 { @@ -500,8 +591,11 @@ async fn test_write_retry_stops_when_a_sibling_is_nonretryable() { .await .unwrap_err(); let terminal_region = RegionId::new(1, 2).as_u64(); + assert_eq!(error.retry_hint(), RetryHint::NonRetryable); + assert_eq!(error.status_code(), StatusCode::InvalidArguments); assert!( - format!("{error:?}") + error + .output_msg() .contains(&format!("region {terminal_region} rejected the write")), "{error:?}" ); @@ -736,81 +830,83 @@ async fn test_finish_retries_initial_failures_in_one_round() { #[tokio::test] async fn test_write_retry_refreshes_stale_route_and_sends_only_failed_region_to_new_leader() { - let old_server = RegionTestServer::start().await; - let new_server = RegionTestServer::start().await; - old_server.fail_region(2, [RetryHint::RetryAfterRouteRefresh]); - let backend = test_util::prepare_mocked_backend().await; - // Reuse the operator fixture to populate table metadata and partition rules. - test_util::create_partition_rule_manager(backend.clone()).await; - update_table_leader(&backend, &old_server.peer).await; + for regions in [vec![2], vec![1, 2]] { + let old_server = RegionTestServer::start().await; + let new_server = RegionTestServer::start().await; + old_server.fail_region(2, [RetryHint::RetryAfterRouteRefresh]); + let backend = test_util::prepare_mocked_backend().await; + // Reuse the operator fixture to populate table metadata and partition rules. + test_util::create_partition_rule_manager(backend.clone()).await; + update_table_leader(&backend, &old_server.peer).await; - let (partition_manager, table_info_cache, registry, catalog_manager) = - create_routing_caches(backend.clone()); - let failed_region = RegionId::new(1, 2); - assert_eq!( - partition_manager - .find_region_leader(failed_region) + let (partition_manager, table_info_cache, registry, catalog_manager) = + create_routing_caches(backend.clone()); + let failed_region = RegionId::new(1, 2); + assert_eq!( + partition_manager + .find_region_leader(failed_region) + .await + .unwrap(), + old_server.peer + ); + let table_info = table_info_cache.get(1).await.unwrap().unwrap(); + let (_, versions) = partition_manager + .find_table_partition_rule(&table_info) .await - .unwrap(), - old_server.peer - ); - let table_info = table_info_cache.get(1).await.unwrap().unwrap(); - let (_, versions) = partition_manager - .find_table_partition_rule(&table_info) - .await - .unwrap(); + .unwrap(); - // Metadata has migrated while the frontend still holds the old route. - update_table_leader(&backend, &new_server.peer).await; - assert_eq!( - partition_manager - .find_region_leader(failed_region) - .await - .unwrap(), - old_server.peer - ); + // Metadata has migrated while the frontend still holds the old route. + update_table_leader(&backend, &new_server.peer).await; + assert_eq!( + partition_manager + .find_region_leader(failed_region) + .await + .unwrap(), + old_server.peer + ); - let invalidator = Arc::new(RecordingCacheInvalidator { - calls: Mutex::default(), - inner: Some(registry), - ..Default::default() - }); - let clients = Arc::new(NodeClients::new(ChannelConfig::default())); - let dispatcher = WriteDispatcher::new(clients, partition_manager).with_retry( - build_retry_options(2), - catalog_manager, - invalidator.clone(), - ); - let mut request = build_insert_request(&[1, 2]); - if let Some(Body::Inserts(inserts)) = &mut request.body { - for insert in &mut inserts.requests { - insert.partition_expr_version = versions - [&RegionId::from_u64(insert.region_id).region_number()] - .map(|value| PartitionExprVersion { value }); + let invalidator = Arc::new(RecordingCacheInvalidator { + calls: Mutex::default(), + inner: Some(registry), + ..Default::default() + }); + let clients = Arc::new(NodeClients::new(ChannelConfig::default())); + let dispatcher = WriteDispatcher::new(clients, partition_manager).with_retry( + build_retry_options(2), + catalog_manager, + invalidator.clone(), + ); + let mut request = build_insert_request(®ions); + if let Some(Body::Inserts(inserts)) = &mut request.body { + for insert in &mut inserts.requests { + insert.partition_expr_version = versions + [&RegionId::from_u64(insert.region_id).region_number()] + .map(|value| PartitionExprVersion { value }); + } } + let response = dispatcher.dispatch(vec![request]).await.unwrap(); + assert_eq!(response.affected_rows, regions.len()); + assert_eq!( + old_server.attempt_counts(), + regions + .iter() + .map(|region| (RegionId::new(1, *region).as_u64(), 1)) + .collect() + ); + assert_eq!( + new_server.attempt_counts(), + HashMap::from([(failed_region.as_u64(), 1)]) + ); + assert_eq!( + *invalidator.calls.lock().unwrap(), + vec![ + vec![CacheIdent::TableId(1)], + vec![CacheIdent::TableName(TableName::new( + "greptime", "public", "table_1" + ))], + ] + ); } - let response = dispatcher.dispatch(vec![request]).await.unwrap(); - assert_eq!(response.affected_rows, 2); - assert_eq!( - old_server.attempt_counts(), - HashMap::from([ - (RegionId::new(1, 1).as_u64(), 1), - (failed_region.as_u64(), 1), - ]) - ); - assert_eq!( - new_server.attempt_counts(), - HashMap::from([(failed_region.as_u64(), 1)]) - ); - assert_eq!( - *invalidator.calls.lock().unwrap(), - vec![ - vec![CacheIdent::TableId(1)], - vec![CacheIdent::TableName(TableName::new( - "greptime", "public", "table_1" - ))], - ] - ); } #[derive(Clone, Copy, Debug)] From 67bd27958656349400f9af9ae7afec726c6cad5e Mon Sep 17 00:00:00 2001 From: LFC <990479+MichaelScofield@users.noreply.github.com> Date: Thu, 22 Oct 2026 04:24:04 +0000 Subject: [PATCH 066/239] fix(query): preserve Arrow types when decoding dictionary output (#9858) Signed-off-by: luofucong --- src/common/recordbatch/src/lib.rs | 55 ++++++++++++++++--- .../types/timestamp/timestamp_tz.result | 47 ++++++++++++++++ .../common/types/timestamp/timestamp_tz.sql | 20 +++++++ 3 files changed, 115 insertions(+), 7 deletions(-) diff --git a/src/common/recordbatch/src/lib.rs b/src/common/recordbatch/src/lib.rs index a3455b366dd..bb4af76fec3 100644 --- a/src/common/recordbatch/src/lib.rs +++ b/src/common/recordbatch/src/lib.rs @@ -221,14 +221,19 @@ pub fn map_dictionary_to_values_schema(schema: SchemaRef) -> (SchemaRef, bool) { .arrow_schema() .fields() .iter() + .zip(schema.column_schemas()) .zip(&columns) - .map(|(field, column)| { - Arc::new( - field - .as_ref() - .clone() - .with_data_type(column.data_type.as_arrow_type()), - ) + .map(|((field, original), column)| { + if original.data_type == column.data_type { + field.clone() + } else { + Arc::new( + field + .as_ref() + .clone() + .with_data_type(column.data_type.as_arrow_type()), + ) + } }) .collect::>(); let arrow_schema = @@ -1160,6 +1165,42 @@ mod tests { ); } + #[test] + fn test_map_dictionary_to_values_preserves_timestamp_timezone() { + let dictionary = DictionaryArray::::new( + Int32Array::from(vec![Some(0), None]), + Arc::new(StringArray::from(vec!["a"])), + ); + let timestamps = datatypes::arrow::array::TimestampSecondArray::from(vec![Some(42), None]) + .with_timezone("Europe/Paris"); + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + Field::new("label", dictionary.data_type().clone(), true), + Field::new("instant", timestamps.data_type().clone(), true), + ])); + let schema = Arc::new(Schema::try_from(arrow_schema.clone()).unwrap()); + let batch = DfRecordBatch::try_new( + arrow_schema.clone(), + vec![Arc::new(dictionary), Arc::new(timestamps)], + ) + .unwrap(); + let batch = RecordBatch::from_df_record_batch(schema.clone(), batch); + + let (mapped_schema, apply_mapper) = map_dictionary_to_values_schema(schema.clone()); + assert!(apply_mapper); + assert_eq!(arrow_schema.field(1), mapped_schema.arrow_schema().field(1)); + let mapped = map_dictionary_to_values(batch.clone(), &schema, &mapped_schema).unwrap(); + assert_eq!(&ArrowDataType::Utf8, mapped.column(0).data_type()); + assert_eq!( + &StringArray::from(vec![Some("a"), None]), + mapped + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + ); + assert!(Arc::ptr_eq(batch.column(1), mapped.column(1))); + } + #[test] fn test_map_dictionary_to_values_schema_with_duplicate_columns() { let dictionary_type = ArrowDataType::Dictionary( diff --git a/tests/cases/standalone/common/types/timestamp/timestamp_tz.result b/tests/cases/standalone/common/types/timestamp/timestamp_tz.result index 8b4833066b7..73a279a4fec 100644 --- a/tests/cases/standalone/common/types/timestamp/timestamp_tz.result +++ b/tests/cases/standalone/common/types/timestamp/timestamp_tz.result @@ -30,3 +30,50 @@ SELECT '2021-04-29 10:50:09-05'::TIMESTAMPTZ::TIME; | 15:50:09 | +------------------------------------------------------------------+ +CREATE TABLE dictionary_output ( + label STRING PRIMARY KEY, + epoch_seconds BIGINT, + ts TIMESTAMP TIME INDEX +); + +Affected Rows: 0 + +INSERT INTO dictionary_output VALUES ('a', 42, 0), ('b', NULL, 1000); + +Affected Rows: 2 + +SELECT from_unixtime(epoch_seconds, 'Europe/Paris') AS instant +FROM dictionary_output ORDER BY ts; + ++---------------------------+ +| instant | ++---------------------------+ +| 1970-01-01T01:00:42+01:00 | +| | ++---------------------------+ + +SELECT label, from_unixtime(epoch_seconds, 'Europe/Paris') AS instant +FROM dictionary_output ORDER BY ts; + ++-------+---------------------------+ +| label | instant | ++-------+---------------------------+ +| a | 1970-01-01T01:00:42+01:00 | +| b | | ++-------+---------------------------+ + +-- SQLNESS PROTOCOL MYSQL +SELECT label, from_unixtime(epoch_seconds, 'Europe/Paris') AS instant +FROM dictionary_output ORDER BY ts; + ++-------+---------------------+ +| label | instant | ++-------+---------------------+ +| a | 1970-01-01 00:00:42 | +| b | | ++-------+---------------------+ + +DROP TABLE dictionary_output; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/types/timestamp/timestamp_tz.sql b/tests/cases/standalone/common/types/timestamp/timestamp_tz.sql index cfb757f02f3..3783363fa63 100644 --- a/tests/cases/standalone/common/types/timestamp/timestamp_tz.sql +++ b/tests/cases/standalone/common/types/timestamp/timestamp_tz.sql @@ -5,3 +5,23 @@ select '2021-11-15 02:30:00'::TIMESTAMP::TIMESTAMPTZ; SELECT '2021-04-29 10:50:09-05'::TIMESTAMPTZ::DATE; SELECT '2021-04-29 10:50:09-05'::TIMESTAMPTZ::TIME; + +CREATE TABLE dictionary_output ( + label STRING PRIMARY KEY, + epoch_seconds BIGINT, + ts TIMESTAMP TIME INDEX +); + +INSERT INTO dictionary_output VALUES ('a', 42, 0), ('b', NULL, 1000); + +SELECT from_unixtime(epoch_seconds, 'Europe/Paris') AS instant +FROM dictionary_output ORDER BY ts; + +SELECT label, from_unixtime(epoch_seconds, 'Europe/Paris') AS instant +FROM dictionary_output ORDER BY ts; + +-- SQLNESS PROTOCOL MYSQL +SELECT label, from_unixtime(epoch_seconds, 'Europe/Paris') AS instant +FROM dictionary_output ORDER BY ts; + +DROP TABLE dictionary_output; From 34210f5022a09f25cadfcd64b9bc96e663228b8e Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Thu, 22 Oct 2026 06:19:03 +0000 Subject: [PATCH 067/239] feat(mito): add an environment switch to disable compaction (#9863) Signed-off-by: WenyXu --- src/mito2/src/compaction/scheduler.rs | 55 +++++++++++- src/mito2/src/compaction/scheduler_test.rs | 98 +++++++++++++++++++++- 2 files changed, 149 insertions(+), 4 deletions(-) diff --git a/src/mito2/src/compaction/scheduler.rs b/src/mito2/src/compaction/scheduler.rs index 76459abf3bc..7406a639305 100644 --- a/src/mito2/src/compaction/scheduler.rs +++ b/src/mito2/src/compaction/scheduler.rs @@ -17,13 +17,13 @@ mod planning; mod state; use std::collections::HashMap; -use std::sync::Arc; +use std::sync::{Arc, LazyLock}; use api::v1::region::compact_request; use common_base::Plugins; use common_memory_manager::OnExhaustedPolicy; use common_meta::key::SchemaMetadataManagerRef; -use common_telemetry::{debug, error, info}; +use common_telemetry::{debug, error, info, warn}; use common_time::range::TimestampRange; pub(crate) use planning::CompactionPickFinished; pub use planning::CompactionRequest; @@ -39,7 +39,7 @@ use crate::compaction::task::MAX_PARALLEL_COMPACTION; use crate::config::MitoConfig; use crate::error::{ CompactionCancelledSnafu, Error, ManualCompactionAlreadyRunningSnafu, RegionClosedSnafu, - RegionDroppedSnafu, RegionTruncatedSnafu, Result, + RegionDroppedSnafu, RegionTruncatedSnafu, Result, UnsupportedOperationSnafu, }; use crate::region::version::VersionControlRef; use crate::region::{ManifestContextRef, RegionLeaderState, RegionRoleState}; @@ -49,8 +49,42 @@ use crate::schedule::RequestCancelResult; use crate::schedule::scheduler::SchedulerRef; use crate::worker::WorkerListener; +/// Internal testing switch to prevent SST compaction from changing file layouts +/// during query regression tests. This is not a public configuration option. +const DISABLE_COMPACTION_ENV: &str = "GREPTIME_DISABLE_COMPACTION"; + +/// Internal query-regression switch, read once when the first scheduler starts. +/// Applies to SST compaction only; ingestion and flushing remain enabled. +static DISABLE_COMPACTION: LazyLock = LazyLock::new(|| { + let value = std::env::var(DISABLE_COMPACTION_ENV).ok(); + let disabled = parse_disable_compaction(value.as_deref()); + if disabled { + warn!( + "SST compaction is disabled by {DISABLE_COMPACTION_ENV}; restart without it to resume" + ); + } + disabled +}); + +fn parse_disable_compaction(value: Option<&str>) -> bool { + match value { + None | Some("false") => false, + Some("true") => true, + Some(_) => { + warn!( + "Invalid {DISABLE_COMPACTION_ENV}: expected true or false; compaction remains enabled" + ); + false + } + } +} + /// Compaction scheduler tracks and manages compaction tasks. pub(crate) struct CompactionScheduler { + /// Disables SST compaction for internal query regression tests. + /// Initialized from [`DISABLE_COMPACTION`] and fixed for the scheduler's lifetime + /// in production; unit tests can override it without changing the process environment. + disable_compaction: bool, scheduler: SchedulerRef, /// All loaded regions, including idle regions and regions executing DDLs. region_status: HashMap, @@ -127,6 +161,7 @@ impl CompactionScheduler { memory_policy: OnExhaustedPolicy, ) -> Self { Self { + disable_compaction: *DISABLE_COMPACTION, scheduler, region_status: HashMap::new(), request_sender, @@ -462,6 +497,17 @@ impl CompactionScheduler { max_parallelism: usize, time_range: Option, ) -> Result { + if self.disable_compaction { + if matches!(trigger, CompactionTrigger::Manual) { + waiter.send( + UnsupportedOperationSnafu { + err_msg: format!("SST compaction is disabled by {DISABLE_COMPACTION_ENV}"), + } + .fail(), + ); + } + return Ok(false); + } let region_id = version_control.region_id(); let status = self.region_status.entry(region_id).or_insert_with(|| { CompactionStatus::new(region_id, version_control.clone(), access_layer.clone()) @@ -631,6 +677,9 @@ impl CompactionScheduler { manifest_ctx: &ManifestContextRef, schema_metadata_manager: SchemaMetadataManagerRef, ) -> bool { + if self.disable_compaction { + return false; + } let Some(status) = self.region_status.get_mut(®ion_id) else { return false; }; diff --git a/src/mito2/src/compaction/scheduler_test.rs b/src/mito2/src/compaction/scheduler_test.rs index d8c7b3944a5..2cd556ad7c3 100644 --- a/src/mito2/src/compaction/scheduler_test.rs +++ b/src/mito2/src/compaction/scheduler_test.rs @@ -2342,12 +2342,16 @@ async fn test_pending_ddl_request_failed_on_region_truncated() { assert_matches!(result, Err(_)); } +#[rstest::rstest] #[tokio::test] -async fn test_on_compaction_finished_returns_pending_ddl_requests() { +async fn test_on_compaction_finished_returns_pending_ddl_requests( + #[values(false, true)] disable_compaction: bool, +) { let job_scheduler = Arc::new(VecScheduler::default()); let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone()); let (tx, _rx) = mpsc::channel(4); let mut scheduler = env.mock_compaction_scheduler(tx); + scheduler.disable_compaction = disable_compaction; let builder = VersionControlBuilder::new(); let version_control = Arc::new(builder.build()); let region_id = builder.region_id(); @@ -2612,3 +2616,95 @@ async fn test_concurrent_memory_competition() { drop(results); assert_eq!(manager.used_bytes(), 0); } + +#[test] +fn test_parse_disable_compaction() { + for (value, expected) in [ + (None, false), + (Some("true"), true), + (Some("false"), false), + (Some(""), false), + (Some("1"), false), + (Some("TRUE"), false), + (Some(" true "), false), + ] { + assert_eq!(parse_disable_compaction(value), expected, "{value:?}"); + } +} + +#[rstest::rstest] +#[case::regular(Options::Regular(Default::default()))] +#[case::strict_window(Options::StrictWindow(StrictWindow::default()))] +#[tokio::test] +async fn test_disable_compaction_scheduling(#[case] options: Options) { + let job_scheduler = Arc::new(VecScheduler::default()); + let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone()); + let (tx, mut rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + scheduler.disable_compaction = true; + let version_control = compactable_version(); + let region_id = version_control.region_id(); + scheduler.register_region(&version_control, &env.access_layer); + let manifest_ctx = env + .mock_manifest_context(version_control.current().version.metadata.clone()) + .await; + let (schemas, _kv_backend) = mock_schema_metadata_manager(); + + assert!( + !scheduler + .schedule_automatic_compaction( + Options::Regular(Default::default()), + &version_control, + &env.access_layer, + &manifest_ctx, + schemas.clone(), + ) + .unwrap() + ); + let (output_tx, mut output_rx) = oneshot::channel(); + assert!( + !scheduler + .schedule_manual_compaction( + options, + &version_control, + &env.access_layer, + OptionOutputTx::from(output_tx), + &manifest_ctx, + schemas.clone(), + 1, + None, + ) + .unwrap() + ); + let err = output_rx.try_recv().unwrap().unwrap_err(); + assert_matches!(err, Error::UnsupportedOperation { .. }); + assert!(err.to_string().contains(DISABLE_COMPACTION_ENV)); + assert!(!scheduler.schedule_automatic_followup(region_id, &manifest_ctx, schemas.clone())); + assert!(!scheduler.is_compacting(region_id)); + assert_eq!(job_scheduler.num_jobs(), 0); + assert!(rx.try_recv().is_err()); + + // The same compactable fixture dispatches planning with the switch off. + scheduler.disable_compaction = false; + assert!( + scheduler + .schedule_automatic_compaction( + Options::Regular(Default::default()), + &version_control, + &env.access_layer, + &manifest_ctx, + schemas, + ) + .unwrap() + ); + let finished = tokio::time::timeout( + Duration::from_secs(10), + recv_compaction_pick_finished(&mut rx), + ) + .await + .expect("enabled compaction should finish planning"); + assert!(matches!( + finished.result, + CompactionPlanningResult::Prepared(_) + )); +} From f595e9bbf08700d7c6bc1ec414954af6412c2b0e Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Thu, 22 Oct 2026 06:19:13 +0000 Subject: [PATCH 068/239] ci: allow independent Linux release architecture selection (#9865) * ci: allow independent Linux release architecture selection Signed-off-by: WenyXu * ci: use Linux architecture checkboxes and preserve release defaults Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- .../release-linux-architectures-test.py | 87 +++++++++++++++++++ .../scripts/release-linux-architectures.sh | 23 +++++ .github/workflows/checks.yml | 2 + .github/workflows/release.yml | 39 +++++++-- 4 files changed, 143 insertions(+), 8 deletions(-) create mode 100644 .github/scripts/release-linux-architectures-test.py create mode 100644 .github/scripts/release-linux-architectures.sh diff --git a/.github/scripts/release-linux-architectures-test.py b/.github/scripts/release-linux-architectures-test.py new file mode 100644 index 00000000000..95fcb6bc19f --- /dev/null +++ b/.github/scripts/release-linux-architectures-test.py @@ -0,0 +1,87 @@ +#!/usr/bin/env python3 +# 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. + +import itertools +import os +import subprocess +import unittest +from pathlib import Path + + +SCRIPT = Path(__file__).with_name("release-linux-architectures.sh") +ARCHITECTURES = ("amd64", "arm64", "riscv64") + + +class ReleaseLinuxArchitecturesTest(unittest.TestCase): + def resolve(self, selected, event="workflow_dispatch", publish=False, images=False): + env = dict( + os.environ, + GITHUB_EVENT_NAME=event, + PUBLISH_GITHUB_RELEASE=str(publish).lower(), + RELEASE_IMAGES=str(images).lower(), + ) + for arch in ARCHITECTURES: + key = f"BUILD_LINUX_{arch.upper()}_ARTIFACTS" + env.pop(key, None) + if selected is not None: + env[key] = str(arch in selected).lower() + return subprocess.run( + ["bash", str(SCRIPT)], env=env, text=True, capture_output=True, check=False + ) + + def assert_targets(self, result, targets): + self.assertEqual(result.returncode, 0, result.stderr) + self.assertEqual(result.stderr, "") + self.assertEqual( + result.stdout, + "".join(f"{arch}={str(arch in targets).lower()}\n" for arch in ARCHITECTURES), + ) + + def test_checkbox_combinations(self): + for flags in itertools.product((False, True), repeat=3): + selected = {arch for arch, enabled in zip(ARCHITECTURES, flags) if enabled} + for event, publish, images in itertools.product( + ("workflow_dispatch", "push", "schedule"), (False, True), (False, True) + ): + with self.subTest(selected=selected, event=event, publish=publish, images=images): + expected = ( + set(ARCHITECTURES) + if event != "workflow_dispatch" or publish or images + else selected + ) + self.assert_targets(self.resolve(selected, event, publish, images), expected) + + def test_default_and_omitted_inputs_preserve_previous_behavior(self): + for selected in (None, set()): + for event, publish, images in itertools.product( + ("workflow_dispatch", "push", "schedule"), (False, True), (False, True) + ): + with self.subTest(selected=selected, event=event, publish=publish, images=images): + # The old build_linux_artifacts switch defaulted to false. + old_build_linux = publish or images or event in ("push", "schedule") + self.assert_targets( + self.resolve(selected, event, publish, images), + set(ARCHITECTURES) if old_build_linux else set(), + ) + + def test_unsupported_event(self): + result = self.resolve(set(), event="pull_request") + self.assertNotEqual(result.returncode, 0) + self.assertIn("Unsupported release event", result.stderr) + self.assertEqual(result.stdout, "") + + +if __name__ == "__main__": + unittest.main() diff --git a/.github/scripts/release-linux-architectures.sh b/.github/scripts/release-linux-architectures.sh new file mode 100644 index 00000000000..cf6f819e84f --- /dev/null +++ b/.github/scripts/release-linux-architectures.sh @@ -0,0 +1,23 @@ +#!/usr/bin/env bash + +set -euo pipefail + +# Publishing keeps the existing full-build behavior, regardless of checkboxes. +build_all=false +if [[ "${PUBLISH_GITHUB_RELEASE:-false}" == true || "${RELEASE_IMAGES:-false}" == true ]]; then + build_all=true +fi +case "${GITHUB_EVENT_NAME}" in + push|schedule) build_all=true ;; + workflow_dispatch) ;; + *) echo "Unsupported release event: ${GITHUB_EVENT_NAME}" >&2; exit 1 ;; +esac + +for arch in amd64 arm64 riscv64; do + input="BUILD_LINUX_${arch^^}_ARTIFACTS" + selected=false + if [[ "$build_all" == true || "${!input:-false}" == true ]]; then + selected=true + fi + printf '%s=%s\n' "$arch" "$selected" +done diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml index 030a8e2fcbb..55dcc34ccf3 100644 --- a/.github/workflows/checks.yml +++ b/.github/workflows/checks.yml @@ -82,6 +82,8 @@ jobs: version: "latest" - name: Run check-version script tests run: uv run --no-project python .github/scripts/check-version-test.py + - name: Run release Linux architecture tests + run: python3 .github/scripts/release-linux-architectures-test.py - name: Run previous-release-tag script tests run: python3 .github/scripts/previous-release-tag-test.py - name: Run enterprise license check tests diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index a309c224568..bc5d045099c 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -19,7 +19,6 @@ on: # NIGHTLY_RELEASE_MAX_AGE_DAYS days, see the 'Check nightly release necessity' step. - cron: '0 0 * * 1-5' workflow_dispatch: # Allows you to run this workflow manually. - # Notes: The GitHub Actions ONLY support 10 inputs; this final input is the tenth. inputs: linux_amd64_runner: type: choice @@ -57,9 +56,19 @@ on: description: Do not run integration tests during the build type: boolean default: true - build_linux_artifacts: + build_linux_amd64_artifacts: type: boolean - description: Build linux artifacts (amd64, arm64, and riscv64) + description: Build linux-amd64 artifacts (publishing always builds all Linux architectures) + required: false + default: false + build_linux_arm64_artifacts: + type: boolean + description: Build linux-arm64 artifacts (publishing always builds all Linux architectures) + required: false + default: false + build_linux_riscv64_artifacts: + type: boolean + description: Build linux-riscv64 artifacts (publishing always builds all Linux architectures) required: false default: false build_macos_artifacts: @@ -115,6 +124,9 @@ jobs: if: ${{ github.repository == 'GreptimeTeam/greptimedb' }} runs-on: ubuntu-latest outputs: + build-linux-amd64: ${{ steps.linux-architectures.outputs.amd64 }} + build-linux-arm64: ${{ steps.linux-architectures.outputs.arm64 }} + build-linux-riscv64: ${{ steps.linux-architectures.outputs.riscv64 }} linux-amd64-runner: ${{ steps.start-linux-amd64-runner.outputs.label }} linux-arm64-runner: ${{ steps.start-linux-arm64-runner.outputs.label }} macos-runner: ${{ inputs.macos_runner || vars.DEFAULT_MACOS_RUNNER }} @@ -143,6 +155,17 @@ jobs: fetch-depth: 0 persist-credentials: false + # Publishing and automatic releases retain the full Linux build. + - name: Resolve Linux architectures + id: linux-architectures + env: + BUILD_LINUX_AMD64_ARTIFACTS: ${{ inputs.build_linux_amd64_artifacts }} + BUILD_LINUX_ARM64_ARTIFACTS: ${{ inputs.build_linux_arm64_artifacts }} + BUILD_LINUX_RISCV64_ARTIFACTS: ${{ inputs.build_linux_riscv64_artifacts }} + PUBLISH_GITHUB_RELEASE: ${{ inputs.publish_github_release }} + RELEASE_IMAGES: ${{ inputs.release_images }} + run: bash .github/scripts/release-linux-architectures.sh >> "$GITHUB_OUTPUT" + - name: Check Rust toolchain version shell: bash run: | @@ -182,7 +205,7 @@ jobs: ./.github/scripts/check-version.sh "${{ steps.create-version.outputs.version }}" - name: Allocate linux-amd64 runner - if: ${{ (inputs.build_linux_artifacts || inputs.publish_github_release || inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule') && steps.check-nightly.outputs.nightly-required == 'true' }} + if: ${{ (steps.linux-architectures.outputs.amd64 == 'true' || steps.linux-architectures.outputs.riscv64 == 'true') && steps.check-nightly.outputs.nightly-required == 'true' }} uses: ./.github/actions/start-runner id: start-linux-amd64-runner with: @@ -196,7 +219,7 @@ jobs: subnet-id: ${{ vars.EC2_RUNNER_SUBNET_ID }} - name: Allocate linux-arm64 runner - if: ${{ (inputs.build_linux_artifacts || inputs.publish_github_release || inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule') && steps.check-nightly.outputs.nightly-required == 'true' }} + if: ${{ steps.linux-architectures.outputs.arm64 == 'true' && steps.check-nightly.outputs.nightly-required == 'true' }} uses: ./.github/actions/start-runner id: start-linux-arm64-runner with: @@ -211,7 +234,7 @@ jobs: build-linux-amd64-artifacts: name: Build linux-amd64 artifacts - if: ${{ (inputs.build_linux_artifacts || inputs.publish_github_release || inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule') && (github.event_name != 'schedule' || needs.allocate-runners.outputs.nightly-required == 'true') }} + if: ${{ needs.allocate-runners.outputs.build-linux-amd64 == 'true' && (github.event_name != 'schedule' || needs.allocate-runners.outputs.nightly-required == 'true') }} needs: [ allocate-runners, ] @@ -237,7 +260,7 @@ jobs: build-linux-arm64-artifacts: name: Build linux-arm64 artifacts - if: ${{ (inputs.build_linux_artifacts || inputs.publish_github_release || inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule') && (github.event_name != 'schedule' || needs.allocate-runners.outputs.nightly-required == 'true') }} + if: ${{ needs.allocate-runners.outputs.build-linux-arm64 == 'true' && (github.event_name != 'schedule' || needs.allocate-runners.outputs.nightly-required == 'true') }} needs: [ allocate-runners, ] @@ -264,7 +287,7 @@ jobs: build-linux-riscv64-artifacts: name: Build linux-riscv64 artifacts - if: ${{ (inputs.build_linux_artifacts || inputs.publish_github_release || inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule') && (github.event_name != 'schedule' || needs.allocate-runners.outputs.nightly-required == 'true') }} + if: ${{ needs.allocate-runners.outputs.build-linux-riscv64 == 'true' && (github.event_name != 'schedule' || needs.allocate-runners.outputs.nightly-required == 'true') }} needs: [ allocate-runners, ] From b4aae2f73b4d5787d38896be1a576b5df18079bc Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Thu, 22 Oct 2026 06:20:23 +0000 Subject: [PATCH 069/239] fix(meta): release all PostgreSQL election advisory locks on step-down (#9864) Signed-off-by: WenyXu --- src/common/meta/src/election/rds/postgres.rs | 203 ++++++++++++++++++- 1 file changed, 201 insertions(+), 2 deletions(-) diff --git a/src/common/meta/src/election/rds/postgres.rs b/src/common/meta/src/election/rds/postgres.rs index 220b33bb601..666481c45af 100644 --- a/src/common/meta/src/election/rds/postgres.rs +++ b/src/common/meta/src/election/rds/postgres.rs @@ -122,7 +122,9 @@ impl<'a> ElectionSqlFactory<'a> { } fn step_down_sql(&self) -> String { - format!("SELECT pg_advisory_unlock({})", self.lock_id) + // Campaigning repeatedly acquires the same session-level lock. Release + // every acquisition when stepping down, not just the most recent one. + "SELECT pg_advisory_unlock_all()".to_string() } fn put_value_with_lease_sql(&self) -> String { @@ -179,6 +181,9 @@ impl<'a> ElectionSqlFactory<'a> { } /// PgClient for election. +/// +/// The session is dedicated to election and candidate registration. It must not +/// hold unrelated advisory locks, as stepping down releases all session locks. pub struct ElectionPgClient { current: Option>, pool: Pool, @@ -1280,6 +1285,24 @@ mod tests { _ => panic!("Expected LeaderChangeMessage::Elected"), } + // Lease timestamps have millisecond precision. Wait for the database + // clock to advance before checking that renewal extends the lease. + tokio::time::timeout(Duration::from_secs(1), async { + loop { + let current = leader_pg_election + .get_value_with_lease(&leader_pg_election.election_key()) + .await + .unwrap() + .unwrap(); + if current.current > lease.current { + break; + } + tokio::time::sleep(Duration::from_millis(1)).await; + } + }) + .await + .unwrap(); + // Step 2: As a leader, renew the lease. let res = leader_pg_election .pg_client @@ -1481,6 +1504,145 @@ mod tests { drop_table(&leader_pg_election, table_name).await; } + #[tokio::test] + async fn test_step_down_releases_reentrant_lock() { + maybe_skip_postgres_integration_test!(); + let table_name = "test_step_down_reentrant_lock_greptime_metakv"; + let store_key_prefix = uuid::Uuid::new_v4().to_string(); + let lock_id = 28323; + let other_lock_id = 28324; + let timeout = Duration::from_secs(10); + let mut elections = Vec::new(); + for leader_value in ["leader", "follower"] { + // Keep both sessions alive so idle timeout cannot hide a leaked lock. + let client = create_postgres_client(Some(table_name), timeout, Duration::ZERO, timeout) + .await + .unwrap(); + let (tx, _) = broadcast::channel(100); + elections.push(PgElection { + leader_value: leader_value.to_string(), + pg_client: RwLock::new(client), + is_leader: AtomicBool::new(false), + leader_infancy: AtomicBool::new(false), + leader_watcher: tx, + store_key_prefix: store_key_prefix.clone(), + candidate_lease_ttl: timeout, + meta_lease_ttl: timeout, + sql_set: ElectionSqlFactory::new(lock_id, None, table_name).build(), + }); + } + let leader = &elections[0]; + let follower = &elections[1]; + let mut leader_changes = leader.subscribe_leader_change(); + let other_lock_sql = format!("SELECT pg_try_advisory_lock({other_lock_id})"); + let other_unlock_sql = format!("SELECT pg_advisory_unlock({other_lock_id})"); + let leader_pid: i32 = leader + .pg_client + .read() + .await + .query("SELECT pg_backend_pid()", &[]) + .await + .unwrap()[0] + .get(0); + + for acquisitions in [1, 3, 1000] { + // Repeat the same acquisition performed on every campaign iteration. + for _ in 0..acquisitions { + let rows = leader + .pg_client + .read() + .await + .query(&leader.sql_set.campaign, &[]) + .await + .unwrap(); + assert!(rows[0].get::<_, bool>(0)); + } + leader.leader_action().await.unwrap(); + assert_matches!( + leader_changes.try_recv().unwrap(), + LeaderChangeMessage::Elected(_) + ); + let rows = follower + .pg_client + .read() + .await + .query(&follower.sql_set.campaign, &[]) + .await + .unwrap(); + assert!(!rows[0].get::<_, bool>(0)); + follower.follower_action().await.unwrap(); + + // An unrelated lock in another session must survive the demotion. + let rows = follower + .pg_client + .read() + .await + .query(&other_lock_sql, &[]) + .await + .unwrap(); + assert!(rows[0].get::<_, bool>(0)); + + // Expire the lease explicitly instead of relying on scheduling delays. + let key = leader.election_key(); + let lease = leader.get_value_with_lease(&key).await.unwrap().unwrap(); + leader + .update_value_with_lease(&key, &lease.origin, &leader.leader_value, Duration::ZERO) + .await + .unwrap(); + leader.leader_action().await.unwrap(); + assert!(!leader.is_leader()); + assert!(leader.get_value_with_lease(&key).await.unwrap().is_none()); + assert_matches!( + leader_changes.try_recv().unwrap(), + LeaderChangeMessage::StepDown(_) + ); + assert!(leader_changes.try_recv().is_err()); + let current_pid: i32 = leader + .pg_client + .read() + .await + .query("SELECT pg_backend_pid()", &[]) + .await + .unwrap()[0] + .get(0); + assert_eq!(leader_pid, current_pid); + + let rows = leader + .pg_client + .read() + .await + .query(&other_lock_sql, &[]) + .await + .unwrap(); + assert!(!rows[0].get::<_, bool>(0)); + let rows = follower + .pg_client + .read() + .await + .query(&follower.sql_set.campaign, &[]) + .await + .unwrap(); + assert!( + rows[0].get::<_, bool>(0), + "lock still held after {acquisitions} acquisitions" + ); + let rows = follower + .pg_client + .read() + .await + .query(&other_unlock_sql, &[]) + .await + .unwrap(); + assert!(rows[0].get::<_, bool>(0)); + follower.leader_action().await.unwrap(); + assert!(follower.is_leader()); + let lease = follower.get_value_with_lease(&key).await.unwrap().unwrap(); + assert_eq!(follower.leader_value, lease.leader_value); + follower.step_down().await.unwrap(); + } + drop_table(leader, table_name).await; + } + #[tokio::test] async fn test_follower_action() { maybe_skip_postgres_integration_test!(); @@ -1619,11 +1781,48 @@ mod tests { store_key_prefix: uuid, candidate_lease_ttl, meta_lease_ttl, - sql_set: ElectionSqlFactory::new(28321, None, table_name).build(), + sql_set: ElectionSqlFactory::new(28325, None, table_name).build(), }; + let follower_client = create_postgres_client( + None, + execution_timeout, + idle_session_timeout, + statement_timeout, + ) + .await + .unwrap(); + for _ in 0..3 { + let rows = leader_pg_election + .pg_client + .read() + .await + .query(&leader_pg_election.sql_set.campaign, &[]) + .await + .unwrap(); + assert!(rows[0].get::<_, bool>(0)); + } leader_pg_election.is_leader.store(true, Ordering::Relaxed); leader_pg_election.reset_campaign().await; assert!(!leader_pg_election.is_leader()); + // Closing the old session is asynchronous. Wait until its locks are released. + tokio::time::timeout(Duration::from_secs(5), async { + loop { + let rows = follower_client + .query(&leader_pg_election.sql_set.campaign, &[]) + .await + .unwrap(); + if rows[0].get::<_, bool>(0) { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .unwrap(); + follower_client + .query(&leader_pg_election.sql_set.step_down, &[]) + .await + .unwrap(); drop_table(&leader_pg_election, table_name).await; } From 8b1d68369be339ed111a6aa88d4b960d4a55a2fc Mon Sep 17 00:00:00 2001 From: Yingwen Date: Thu, 22 Oct 2026 06:45:25 +0000 Subject: [PATCH 070/239] fix(ci): align nightly Windows and macOS artifact test stack sizes (#9860) * fix(ci): increase Windows nightly test stack size Signed-off-by: evenyag * fix(ci): align optional macOS artifact test stack size Signed-off-by: evenyag --------- Signed-off-by: evenyag --- .github/actions/build-macos-artifacts/action.yml | 1 + .github/workflows/nightly-ci.yml | 1 + 2 files changed, 2 insertions(+) diff --git a/.github/actions/build-macos-artifacts/action.yml b/.github/actions/build-macos-artifacts/action.yml index 2caefe1c382..2affb5b94f7 100644 --- a/.github/actions/build-macos-artifacts/action.yml +++ b/.github/actions/build-macos-artifacts/action.yml @@ -64,6 +64,7 @@ runs: shell: bash env: CARGO_BUILD_RUSTFLAGS: "-Clink-arg=-Wl,-ld_classic" + RUST_MIN_STACK: 8388608 # 8 MiB, aligned with Rust CI and release tests. SQLNESS_OPTS: "--preserve-state" run: | make test sqlness-test diff --git a/.github/workflows/nightly-ci.yml b/.github/workflows/nightly-ci.yml index f7f955233a1..15f8f505bea 100644 --- a/.github/workflows/nightly-ci.yml +++ b/.github/workflows/nightly-ci.yml @@ -157,6 +157,7 @@ jobs: env: CARGO_BUILD_RUSTFLAGS: "-C linker=lld-link" RUST_BACKTRACE: 1 + RUST_MIN_STACK: 8388608 # 8 MiB, aligned with Rust CI and release tests. CARGO_INCREMENTAL: 0 GT_S3_BUCKET: ${{ vars.AWS_CI_TEST_BUCKET }} GT_S3_ACCESS_KEY_ID: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} From af52dfafd8d0187c511c76b726de798d59835ef8 Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Thu, 22 Oct 2026 07:27:19 +0000 Subject: [PATCH 071/239] chore(deps): upgrade sqlx to 0.9 to pick up webpki-roots 1.0.10 (#9862) sqlx 0.8 pins the unmaintained webpki-roots 0.26 line for its rustls trust store, which still contains the recently compromised e-Tugra CA root. Upgrading sqlx to 0.9 switches its dependency to webpki-roots 1.x, and the lockfile now resolves to 1.0.10 where the distrusted root has been removed from the compiled-in Mozilla trust store. sqlx 0.9 requires dynamic SQL strings to be explicitly audited via the new SqlSafeStr trait: only &'static str literals are accepted as-is. All affected call sites (MySQL KV backend, MySQL election, integration tests and fuzz targets) now wrap their internally generated SQL templates in AssertSqlSafe; the wrapped SQL is built from operator config or test-local identifiers only, and dynamic data keeps flowing through bind parameters. The mysql-rsa feature is enabled alongside mysql to preserve 0.8 authentication behavior: sqlx 0.9 gates non-TLS RSA password auth (caching_sha2_password / sha256_password) behind it. Signed-off-by: Ning Sun --- Cargo.lock | 573 ++++++++++++++---- Cargo.toml | 2 +- src/common/meta/Cargo.toml | 1 + src/common/meta/src/election/rds/mysql.rs | 41 +- src/common/meta/src/kv_backend/rds/mysql.rs | 77 ++- src/meta-srv/Cargo.toml | 1 + tests-fuzz/Cargo.toml | 1 + tests-fuzz/src/utils.rs | 4 +- tests-fuzz/src/utils/migration.rs | 2 +- tests-fuzz/src/utils/procedure.rs | 2 +- tests-fuzz/src/validator/column.rs | 2 +- tests-fuzz/src/validator/row.rs | 9 +- tests-fuzz/src/validator/table.rs | 3 +- .../targets/ddl/fuzz_alter_logical_table.rs | 10 +- tests-fuzz/targets/ddl/fuzz_alter_table.rs | 10 +- .../targets/ddl/fuzz_create_database.rs | 4 +- .../targets/ddl/fuzz_create_logical_table.rs | 8 +- tests-fuzz/targets/ddl/fuzz_create_table.rs | 4 +- .../ddl/fuzz_repartition_metric_table.rs | 12 +- .../targets/ddl/fuzz_repartition_table.rs | 10 +- .../ddl/fuzz_repartition_table_chaos.rs | 8 +- .../failover/fuzz_failover_metric_regions.rs | 10 +- .../failover/fuzz_failover_mito_regions.rs | 6 +- tests-fuzz/targets/fuzz_insert.rs | 6 +- .../targets/fuzz_insert_logical_table.rs | 10 +- .../migration/fuzz_migrate_metric_regions.rs | 10 +- .../migration/fuzz_migrate_mito_regions.rs | 18 +- .../unstable/fuzz_create_table_standalone.rs | 14 +- .../wal/fuzz_remote_wal_logical_prune.rs | 18 +- tests-integration/Cargo.toml | 1 + tests-integration/tests/json2.rs | 53 +- tests-integration/tests/sql.rs | 12 +- 32 files changed, 672 insertions(+), 270 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 4c5a0077193..0f2a17d104f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -920,7 +920,7 @@ checksum = "37672978ae0febce7516ae0a85b53e6185159a9a28787391eb63fc44ec36037d" dependencies = [ "async-fs", "futures-lite", - "thiserror 2.0.17", + "thiserror 2.0.21", ] [[package]] @@ -1202,6 +1202,12 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d27c3610c36aee21ce8ac510e6224498de4228ad772a171ed65643a24693a5a8" +[[package]] +name = "base16ct" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fd307490d624467aa6f74b0eabb77633d1f758a7b25f12bceb0b22e08d9726f6" + [[package]] name = "base62" version = "2.2.1" @@ -1674,7 +1680,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a1d8af896b707212cd0e99c112a78c9497dd32994192a463ed2f7419d29bd8c6" dependencies = [ "serde", - "thiserror 2.0.17", + "thiserror 2.0.21", "toml 0.8.23", ] @@ -2701,7 +2707,7 @@ dependencies = [ "serde_json", "serde_with", "snafu 0.8.6", - "sqlx", + "sqlx 0.9.0", "store-api", "strum 0.27.1", "table", @@ -3457,6 +3463,20 @@ version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" +[[package]] +name = "crypto-bigint" +version = "0.7.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a52aa3fcda4e6302a9f48734f234d35d4721b96f8fe07d073f07ce9df4f0271" +dependencies = [ + "cpubits", + "ctutils", + "num-traits", + "rand_core 0.10.0", + "serdect", + "zeroize", +] + [[package]] name = "crypto-common" version = "0.1.6" @@ -3477,6 +3497,16 @@ dependencies = [ "hybrid-array", ] +[[package]] +name = "crypto-primes" +version = "0.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3633a51a39c69ebbaa4feaa694bd83d241e4093901c84a0963b19d9bb3f0cf8f" +dependencies = [ + "crypto-bigint", + "rand_core 0.10.0", +] + [[package]] name = "crypto_secretbox" version = "0.1.1" @@ -4637,7 +4667,18 @@ dependencies = [ "const-oid 0.9.6", "der_derive", "flagset", - "pem-rfc7468", + "pem-rfc7468 0.7.0", + "zeroize", +] + +[[package]] +name = "der" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a878c850e9e421b20262e9b41f9c860e4785fa07541c266b62ff9d1ef998a80a" +dependencies = [ + "const-oid 0.10.2", + "pem-rfc7468 1.0.0", "zeroize", ] @@ -4975,7 +5016,7 @@ dependencies = [ "chrono", "rust_decimal", "serde", - "thiserror 2.0.17", + "thiserror 2.0.21", "time", "winnow 0.6.26", ] @@ -5159,6 +5200,16 @@ dependencies = [ "windows-sys 0.48.0", ] +[[package]] +name = "etcetera" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "de48cc4d1c1d97a20fd819def54b890cadde72ed3ad0c614822a0a433361be96" +dependencies = [ + "cfg-if", + "windows-sys 0.61.2", +] + [[package]] name = "ethnum" version = "1.5.3" @@ -5501,6 +5552,17 @@ dependencies = [ "spin 0.9.8", ] +[[package]] +name = "flume" +version = "0.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e139bc46ca777eb5efaf62df0ab8cc5fd400866427e56c68b22e414e53bd3be" +dependencies = [ + "futures-core", + "futures-sink", + "spin 0.9.8", +] + [[package]] name = "fnv" version = "1.0.7" @@ -5767,9 +5829,9 @@ dependencies = [ [[package]] name = "futures-channel" -version = "0.3.31" +version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2dff15bf788c671c1934e366d07e30c1814a8ef514e1af724a602e8a2fbe1b10" +checksum = "b1f9e3d69d39e4862ffed03ed071a76f9a13ba1d9109d355b0f0aa6b15e393c4" dependencies = [ "futures-core", "futures-sink", @@ -5777,15 +5839,15 @@ dependencies = [ [[package]] name = "futures-core" -version = "0.3.31" +version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "05f29059c0c2090612e8d742178b0580d2dc940c837851ad723096f87af6663e" +checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e" [[package]] name = "futures-executor" -version = "0.3.31" +version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e28d1d997f585e54aebc3f97d39e72338912123a67330d723fdbb564d646c9f" +checksum = "031b47cf1a3c6cc8bc2fc76cd437f521619387907d469316e7c0bc278f1f5432" dependencies = [ "futures-core", "futures-task", @@ -5805,9 +5867,9 @@ dependencies = [ [[package]] name = "futures-io" -version = "0.3.31" +version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9e5c1b78ca4aae1ac06c48a526a655760685149f0d465d21f37abfe57ce075c6" +checksum = "53c0fa8157de1303bfffdaa1cc2a673bfffb60102f76b0ef4441659124373fed" [[package]] name = "futures-lite" @@ -5824,26 +5886,26 @@ dependencies = [ [[package]] name = "futures-macro" -version = "0.3.31" +version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "162ee34ebcb7c64a8abebc059ce0fee27c2262618d7b60ed8faf72fef13c3650" +checksum = "9fb9654ba8355388abeb8dcb4fc62f511300867002afc858860463bdd9fe0c44" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 3.0.5", ] [[package]] name = "futures-sink" -version = "0.3.31" +version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e575fab7d1e0dcb8d0c7bcf9a63ee213816ab51902e6d244a95819acacf1d4f7" +checksum = "1944426bf7d03f1d14f708785e4b33efd750b36d48a157b836b3efc15ede8e1d" [[package]] name = "futures-task" -version = "0.3.31" +version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f90f7dce0722e95104fcb095585910c0977252f286e354b5e3bd38902cd99988" +checksum = "cd417de3d1d015fc3bfd2b1ea46dfc7bab72ef86f1cc7cc9c78e728b34a6d1fd" [[package]] name = "futures-timer" @@ -5853,9 +5915,9 @@ checksum = "f288b0a4f20f9a56b5d1da57e2227c661b7b16168e2f72365f57b63326e29b24" [[package]] name = "futures-util" -version = "0.3.31" +version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9fa08315bb612088cc391249efdc3bc77536f16c91f6cf495e6fbe85b20a4a81" +checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc" dependencies = [ "futures-channel", "futures-core", @@ -5865,7 +5927,6 @@ dependencies = [ "futures-task", "memchr", "pin-project-lite", - "pin-utils", "slab", ] @@ -6254,6 +6315,15 @@ dependencies = [ "hashbrown 0.15.4", ] +[[package]] +name = "hashlink" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "824e001ac4f3012dd16a264bec811403a67ca9deb6c102fc5049b32c4574b35f" +dependencies = [ + "hashbrown 0.16.1", +] + [[package]] name = "hdfs-native" version = "0.14.6" @@ -6288,7 +6358,7 @@ dependencies = [ "regex", "roxmltree", "socket2 0.6.4", - "thiserror 2.0.17", + "thiserror 2.0.21", "tokio", "url", "uuid", @@ -6381,6 +6451,15 @@ dependencies = [ "hmac 0.12.1", ] +[[package]] +name = "hkdf" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4aaa26c720c68b866f2c96ef5c1264b3e6f473fe5d4ce61cd44bbe913e553018" +dependencies = [ + "hmac 0.13.0", +] + [[package]] name = "hmac" version = "0.12.1" @@ -6631,7 +6710,7 @@ dependencies = [ "tokio", "tokio-rustls", "tower-service", - "webpki-roots 1.0.9", + "webpki-roots 1.0.10", ] [[package]] @@ -7405,7 +7484,7 @@ dependencies = [ "jni-sys", "log", "simd_cesu8", - "thiserror 2.0.17", + "thiserror 2.0.21", "walkdir", "windows-link 0.2.1", ] @@ -7532,7 +7611,7 @@ dependencies = [ "pest_derive", "regex", "serde_json", - "thiserror 2.0.17", + "thiserror 2.0.21", ] [[package]] @@ -8060,7 +8139,7 @@ checksum = "656b3b27f8893f7bbf9485148ff9a65f019e3f33bd5cdc87c83cab16b3fd9ec8" dependencies = [ "libc", "neli", - "thiserror 2.0.17", + "thiserror 2.0.21", "windows-sys 0.59.0", ] @@ -8296,7 +8375,7 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bb4bdc8b0ce69932332cf76d24af69c3a155242af95c226b2ab6c2e371ed1149" dependencies = [ - "thiserror 2.0.17", + "thiserror 2.0.21", "zerocopy", "zerocopy-derive", ] @@ -8528,7 +8607,7 @@ dependencies = [ "serde_json", "servers", "snafu 0.8.6", - "sqlx", + "sqlx 0.9.0", "store-api", "strum 0.27.1", "table", @@ -8880,7 +8959,7 @@ dependencies = [ "twox-hash", "url", "webpki", - "webpki-roots 1.0.9", + "webpki-roots 1.0.10", ] [[package]] @@ -8916,7 +8995,7 @@ dependencies = [ "quote", "syn 2.0.117", "termcolor", - "thiserror 2.0.17", + "thiserror 2.0.21", ] [[package]] @@ -8940,13 +9019,13 @@ dependencies = [ "rustls", "serde", "socket2 0.6.4", - "thiserror 2.0.17", + "thiserror 2.0.21", "tokio", "tokio-rustls", "tokio-util", "twox-hash", "url", - "webpki-roots 1.0.9", + "webpki-roots 1.0.10", ] [[package]] @@ -9012,7 +9091,7 @@ dependencies = [ "serde_json", "sha1 0.10.6", "sha2 0.10.9", - "thiserror 2.0.17", + "thiserror 2.0.21", "time", "uuid", ] @@ -9041,7 +9120,7 @@ dependencies = [ "serde_json", "sha1 0.11.0", "sha2 0.11.0", - "thiserror 2.0.17", + "thiserror 2.0.21", "uuid", ] @@ -9583,7 +9662,7 @@ dependencies = [ "itertools 0.14.0", "parking_lot 0.12.4", "percent-encoding", - "thiserror 2.0.17", + "thiserror 2.0.21", "tokio", "tracing", "url", @@ -9924,7 +10003,7 @@ dependencies = [ "asyncband", "opendal-core", "serde", - "sqlx", + "sqlx 0.8.6", ] [[package]] @@ -9997,7 +10076,7 @@ dependencies = [ "futures-sink", "js-sys", "pin-project-lite", - "thiserror 2.0.17", + "thiserror 2.0.21", "tracing", ] @@ -10027,7 +10106,7 @@ dependencies = [ "opentelemetry_sdk", "prost 0.14.1", "reqwest 0.13.4", - "thiserror 2.0.17", + "thiserror 2.0.21", "tokio", "tonic 0.14.2", "tonic-types", @@ -10074,7 +10153,7 @@ dependencies = [ "percent-encoding", "portable-atomic", "rand 0.9.4", - "thiserror 2.0.17", + "thiserror 2.0.21", "tokio", "tokio-stream", ] @@ -10580,7 +10659,7 @@ checksum = "500fa4cdeacd98997c5865e3d0d1cb8fe7e9d7d75ecc775e07989a433a9a9a59" dependencies = [ "arrayvec", "bitflags 2.12.1", - "thiserror 2.0.17", + "thiserror 2.0.21", "zerocopy", "zerocopy-derive", ] @@ -10610,6 +10689,15 @@ dependencies = [ "base64ct", ] +[[package]] +name = "pem-rfc7468" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6305423e0e7738146434843d1694d621cce767262b2a86910beab705e4493d9" +dependencies = [ + "base64ct", +] + [[package]] name = "percent-encoding" version = "2.3.2" @@ -10629,7 +10717,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1db05f56d34358a8b1066f67cbb203ee3e7ed2ba674a6263a1d5ec6db2204323" dependencies = [ "memchr", - "thiserror 2.0.17", + "thiserror 2.0.21", "ucd-trie", ] @@ -10735,7 +10823,7 @@ dependencies = [ "serde_json", "smol_str", "stringprep", - "thiserror 2.0.17", + "thiserror 2.0.21", "tokio", "tokio-rustls", "tokio-util", @@ -10929,9 +11017,19 @@ version = "0.7.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c8ffb9f10fa047879315e6625af03c164b16962a5368d724ed16323b68ace47f" dependencies = [ - "der", - "pkcs8", - "spki", + "der 0.7.10", + "pkcs8 0.10.2", + "spki 0.7.3", +] + +[[package]] +name = "pkcs1" +version = "0.8.0-rc.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f485dbafe293ea6ccaa7d82bc591651a796eb8a9aec3df040b057b706073dbc3" +dependencies = [ + "der 0.8.2", + "spki 0.8.1", ] [[package]] @@ -10942,11 +11040,11 @@ checksum = "e847e2c91a18bfa887dd028ec33f2fe6f25db77db3619024764914affe8b69a6" dependencies = [ "aes 0.8.4", "cbc 0.1.2", - "der", + "der 0.7.10", "pbkdf2", "scrypt", "sha2 0.10.9", - "spki", + "spki 0.7.3", ] [[package]] @@ -10955,10 +11053,20 @@ version = "0.10.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f950b2377845cebe5cf8b5165cb3cc1a5e0fa5cfa3e1f7f55707d8fd82e0a7b7" dependencies = [ - "der", + "der 0.7.10", "pkcs5", "rand_core 0.6.4", - "spki", + "spki 0.7.3", +] + +[[package]] +name = "pkcs8" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "451913da69c775a56034ea8d9003d27ee8948e12443eae7c038ba100a4f21cb7" +dependencies = [ + "der 0.8.2", + "spki 0.8.1", ] [[package]] @@ -11138,7 +11246,7 @@ dependencies = [ "spin 0.10.1", "symbolic-demangle", "tempfile", - "thiserror 2.0.17", + "thiserror 2.0.21", ] [[package]] @@ -11341,7 +11449,7 @@ dependencies = [ "parking_lot 0.12.4", "procfs", "protobuf 3.7.2", - "thiserror 2.0.17", + "thiserror 2.0.21", ] [[package]] @@ -11944,7 +12052,7 @@ dependencies = [ "rustc-hash", "rustls", "socket2 0.5.10", - "thiserror 2.0.17", + "thiserror 2.0.21", "tokio", "tracing", "web-time", @@ -11967,7 +12075,7 @@ dependencies = [ "rustls", "rustls-pki-types", "slab", - "thiserror 2.0.17", + "thiserror 2.0.21", "tinyvec", "tracing", "web-time", @@ -12436,7 +12544,7 @@ dependencies = [ "pem", "percent-encoding", "reqsign-core", - "rsa", + "rsa 0.9.10", "serde", "serde_json", "sha1 0.11.0", @@ -12459,7 +12567,7 @@ dependencies = [ "log", "mea", "percent-encoding", - "rsa", + "rsa 0.9.10", "serde", "serde_json", "sha1 0.11.0", @@ -12490,7 +12598,7 @@ dependencies = [ "percent-encoding", "reqsign-aws-v4", "reqsign-core", - "rsa", + "rsa 0.9.10", "serde", "serde_json", "tokio", @@ -12536,7 +12644,7 @@ dependencies = [ "wasm-bindgen-futures", "wasm-streams 0.4.2", "web-sys", - "webpki-roots 1.0.9", + "webpki-roots 1.0.10", ] [[package]] @@ -12612,7 +12720,7 @@ dependencies = [ "reqwest 0.12.24", "reqwest-middleware", "retry-policies", - "thiserror 2.0.17", + "thiserror 2.0.21", "tokio", "wasmtimer", ] @@ -12731,16 +12839,34 @@ dependencies = [ "num-bigint-dig", "num-integer", "num-traits", - "pkcs1", - "pkcs8", + "pkcs1 0.7.5", + "pkcs8 0.10.2", "rand_core 0.6.4", "sha2 0.10.9", - "signature", - "spki", + "signature 2.2.0", + "spki 0.7.3", "subtle", "zeroize", ] +[[package]] +name = "rsa" +version = "0.10.0-rc.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33405e321e16ef941ca66e254f6904c3aafd3c7ff001fb88bf11c73eb75c2b85" +dependencies = [ + "const-oid 0.10.2", + "crypto-bigint", + "crypto-primes", + "digest 0.11.3", + "pkcs1 0.8.0-rc.5", + "pkcs8 0.11.0", + "rand_core 0.10.0", + "signature 3.0.0", + "spki 0.8.1", + "zeroize", +] + [[package]] name = "rsasl" version = "2.2.1" @@ -12757,7 +12883,7 @@ dependencies = [ "serde_json", "sha2 0.10.9", "stringprep", - "thiserror 2.0.17", + "thiserror 2.0.21", ] [[package]] @@ -12777,7 +12903,7 @@ dependencies = [ "rsasl", "rustls", "snap", - "thiserror 2.0.17", + "thiserror 2.0.21", "tokio", "tokio-rustls", "tracing", @@ -13254,7 +13380,7 @@ dependencies = [ "proc-macro2", "quote", "syn 2.0.117", - "thiserror 2.0.17", + "thiserror 2.0.21", ] [[package]] @@ -13507,6 +13633,16 @@ dependencies = [ "unsafe-libyaml", ] +[[package]] +name = "serdect" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "66cf8fedced2fcf12406bcb34223dffb92eaf34908ede12fed414c82b7f00b3e" +dependencies = [ + "base16ct", + "serde", +] + [[package]] name = "servers" version = "1.4.0-alpha.0" @@ -13781,6 +13917,16 @@ dependencies = [ "rand_core 0.6.4", ] +[[package]] +name = "signature" +version = "3.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "28d567dcbaf0049cb8ac2608a76cd95ff9e4412e1899d389ee400918ca7537f5" +dependencies = [ + "digest 0.11.3", + "rand_core 0.10.0", +] + [[package]] name = "simba" version = "0.9.0" @@ -14009,7 +14155,17 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d91ed6c858b01f942cd56b37a94b3e0a1798290327d1236e4d9cf4eaca44d29d" dependencies = [ "base64ct", - "der", + "der 0.7.10", +] + +[[package]] +name = "spki" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ef958a98b9d5da290cfc78946e9f3e61e1e62a18db0d92cac0b83cc161491a9" +dependencies = [ + "base64ct", + "der 0.8.2", ] [[package]] @@ -14136,11 +14292,24 @@ version = "0.8.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1fefb893899429669dcdd979aff487bd78f4064e5e7907e4269081e0ef7d97dc" dependencies = [ - "sqlx-core", - "sqlx-macros", - "sqlx-mysql", - "sqlx-postgres", - "sqlx-sqlite", + "sqlx-core 0.8.6", + "sqlx-macros 0.8.6", + "sqlx-mysql 0.8.6", + "sqlx-postgres 0.8.6", + "sqlx-sqlite 0.8.6", +] + +[[package]] +name = "sqlx" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "378620ccc25c62c89d8be1c819e76a88d59bdcc3304733330788948e619bfd71" +dependencies = [ + "sqlx-core 0.9.0", + "sqlx-macros 0.9.0", + "sqlx-mysql 0.9.0", + "sqlx-postgres 0.9.0", + "sqlx-sqlite 0.9.0", ] [[package]] @@ -14151,7 +14320,6 @@ checksum = "ee6798b1838b6a0f69c007c133b8df5866302197e404e8b6ee8ed3e3a5e68dc6" dependencies = [ "base64 0.22.1", "bytes", - "chrono", "crc", "crossbeam-queue", "either", @@ -14161,24 +14329,61 @@ dependencies = [ "futures-io", "futures-util", "hashbrown 0.15.4", - "hashlink", + "hashlink 0.10.0", "indexmap 2.14.2", "log", "memchr", "once_cell", "percent-encoding", + "rustls", + "serde", + "serde_json", + "sha2 0.10.9", + "smallvec", + "thiserror 2.0.21", + "tokio", + "tokio-stream", + "tracing", + "url", + "webpki-roots 0.26.11", +] + +[[package]] +name = "sqlx-core" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05b44e85bf579a8eeb4ceaa77a3a523baf2bf0e9bac7e40f405d537b5d2d5ccb" +dependencies = [ + "base64 0.22.1", + "bytes", + "cfg-if", + "chrono", + "crc", + "crossbeam-queue", + "either", + "event-listener 5.4.0", + "futures-core", + "futures-intrusive", + "futures-io", + "futures-util", + "hashbrown 0.16.1", + "hashlink 0.11.1", + "indexmap 2.14.2", + "log", + "memchr", + "percent-encoding", "rust_decimal", "rustls", "serde", "serde_json", "sha2 0.10.9", "smallvec", - "thiserror 2.0.17", + "thiserror 2.0.21", "tokio", "tokio-stream", "tracing", "url", - "webpki-roots 0.26.11", + "webpki-roots 1.0.10", ] [[package]] @@ -14189,8 +14394,21 @@ checksum = "a2d452988ccaacfbf5e0bdbc348fb91d7c8af5bee192173ac3636b5fb6e6715d" dependencies = [ "proc-macro2", "quote", - "sqlx-core", - "sqlx-macros-core", + "sqlx-core 0.8.6", + "sqlx-macros-core 0.8.6", + "syn 2.0.117", +] + +[[package]] +name = "sqlx-macros" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bd2b84f2bc39a5705ef27ec785a11c934a41bbd4a24941e257927cddc26b60bf" +dependencies = [ + "proc-macro2", + "quote", + "sqlx-core 0.9.0", + "sqlx-macros-core 0.9.0", "syn 2.0.117", ] @@ -14210,15 +14428,41 @@ dependencies = [ "serde", "serde_json", "sha2 0.10.9", - "sqlx-core", - "sqlx-mysql", - "sqlx-postgres", - "sqlx-sqlite", + "sqlx-core 0.8.6", + "sqlx-mysql 0.8.6", + "sqlx-postgres 0.8.6", + "sqlx-sqlite 0.8.6", "syn 2.0.117", "tokio", "url", ] +[[package]] +name = "sqlx-macros-core" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb8d96de5fdc85a5c4ec813432b523ec637e80ba98f046555f75f7908ddac7c3" +dependencies = [ + "cfg-if", + "dotenvy", + "either", + "heck 0.5.0", + "hex", + "proc-macro2", + "quote", + "serde", + "serde_json", + "sha2 0.10.9", + "sqlx-core 0.9.0", + "sqlx-mysql 0.9.0", + "sqlx-postgres 0.9.0", + "sqlx-sqlite 0.9.0", + "syn 2.0.117", + "thiserror 2.0.21", + "tokio", + "url", +] + [[package]] name = "sqlx-mysql" version = "0.8.6" @@ -14230,7 +14474,6 @@ dependencies = [ "bitflags 2.12.1", "byteorder", "bytes", - "chrono", "crc", "digest 0.10.7", "dotenvy", @@ -14241,7 +14484,7 @@ dependencies = [ "futures-util", "generic-array", "hex", - "hkdf", + "hkdf 0.12.4", "hmac 0.12.1", "itoa", "log", @@ -14250,19 +14493,48 @@ dependencies = [ "once_cell", "percent-encoding", "rand 0.8.5", - "rsa", - "rust_decimal", + "rsa 0.9.10", "serde", "sha1 0.10.6", "sha2 0.10.9", "smallvec", - "sqlx-core", + "sqlx-core 0.8.6", "stringprep", - "thiserror 2.0.17", + "thiserror 2.0.21", "tracing", "whoami 1.6.0", ] +[[package]] +name = "sqlx-mysql" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90b8020fe17c5f2c245bfa2505d7ef59c5604839527c740266ad2214acebea27" +dependencies = [ + "bitflags 2.12.1", + "byteorder", + "bytes", + "chrono", + "crc", + "digest 0.11.3", + "dotenvy", + "either", + "futures-core", + "futures-util", + "generic-array", + "log", + "percent-encoding", + "rand 0.10.1", + "rsa 0.10.0-rc.19", + "rust_decimal", + "serde", + "sha1 0.11.0", + "sha2 0.11.0", + "sqlx-core 0.9.0", + "thiserror 2.0.21", + "tracing", +] + [[package]] name = "sqlx-postgres" version = "0.8.6" @@ -14273,15 +14545,14 @@ dependencies = [ "base64 0.22.1", "bitflags 2.12.1", "byteorder", - "chrono", "crc", "dotenvy", - "etcetera", + "etcetera 0.8.0", "futures-channel", "futures-core", "futures-util", "hex", - "hkdf", + "hkdf 0.12.4", "hmac 0.12.1", "home", "itoa", @@ -14290,18 +14561,54 @@ dependencies = [ "memchr", "once_cell", "rand 0.8.5", - "rust_decimal", "serde", "serde_json", "sha2 0.10.9", "smallvec", - "sqlx-core", + "sqlx-core 0.8.6", "stringprep", - "thiserror 2.0.17", + "thiserror 2.0.21", "tracing", "whoami 1.6.0", ] +[[package]] +name = "sqlx-postgres" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87a2bdd6e83f6b3ea525ca9fee568030508b58355a43d0b2c1674d5f79dcd65e" +dependencies = [ + "atoi", + "base64 0.22.1", + "bitflags 2.12.1", + "byteorder", + "chrono", + "crc", + "dotenvy", + "etcetera 0.11.0", + "futures-channel", + "futures-core", + "futures-util", + "hex", + "hkdf 0.13.0", + "hmac 0.13.0", + "itoa", + "log", + "md-5 0.11.0", + "memchr", + "rand 0.10.1", + "rust_decimal", + "serde", + "serde_json", + "sha2 0.11.0", + "smallvec", + "sqlx-core 0.9.0", + "stringprep", + "thiserror 2.0.21", + "tracing", + "whoami 2.1.3", +] + [[package]] name = "sqlx-sqlite" version = "0.8.6" @@ -14309,8 +14616,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2d12fe70b2c1b4401038055f90f151b78208de1f9f89a7dbfd41587a10c3eea" dependencies = [ "atoi", - "chrono", - "flume", + "flume 0.11.1", "futures-channel", "futures-core", "futures-executor", @@ -14321,8 +14627,33 @@ dependencies = [ "percent-encoding", "serde", "serde_urlencoded", - "sqlx-core", - "thiserror 2.0.17", + "sqlx-core 0.8.6", + "thiserror 2.0.21", + "tracing", + "url", +] + +[[package]] +name = "sqlx-sqlite" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "488e99c397a62007e4229aec669a179816339afc6d2620ca6fa420dbee2e982c" +dependencies = [ + "atoi", + "chrono", + "flume 0.12.0", + "form_urlencoded", + "futures-channel", + "futures-core", + "futures-executor", + "futures-intrusive", + "futures-util", + "libsqlite3-sys", + "log", + "percent-encoding", + "serde", + "sqlx-core 0.9.0", + "thiserror 2.0.21", "tracing", "url", ] @@ -14816,7 +15147,7 @@ dependencies = [ "tantivy-stacker", "tantivy-tokenizer-api", "tempfile", - "thiserror 2.0.17", + "thiserror 2.0.21", "time", "typetag", "uuid", @@ -15046,7 +15377,7 @@ dependencies = [ "snafu 0.8.6", "sql", "sqlparser", - "sqlx", + "sqlx 0.9.0", "store-api", "strum 0.27.1", "tinytemplate", @@ -15131,7 +15462,7 @@ dependencies = [ "session", "similar-asserts", "sql", - "sqlx", + "sqlx 0.9.0", "standalone", "store-api", "substrait 1.4.0-alpha.0", @@ -15163,11 +15494,11 @@ dependencies = [ [[package]] name = "thiserror" -version = "2.0.17" +version = "2.0.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f63587ca0f12b72a0600bcba1d40081f830876000bb46dd2337a3051618f4fc8" +checksum = "09e52cb86a36cede5cb101bf8908837b3e4c6e5e59fe7fd85c23fb56200d189e" dependencies = [ - "thiserror-impl 2.0.17", + "thiserror-impl 2.0.21", ] [[package]] @@ -15183,13 +15514,13 @@ dependencies = [ [[package]] name = "thiserror-impl" -version = "2.0.17" +version = "2.0.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3ff15c8ecd7de3849db632e14d18d2571fa09dfc5ed93479bc4485c7a517c913" +checksum = "fe5197923287db20a58125f0bc85c062f7f2c892de97b18c356f9efb14b28524" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 3.0.5", ] [[package]] @@ -15821,7 +16152,7 @@ checksum = "050686193eb999b4bb3bc2acfa891a13da00f79734704c4b8b4ef1a10b368a3c" dependencies = [ "crossbeam-channel", "symlink", - "thiserror 2.0.17", + "thiserror 2.0.21", "time", "tracing-subscriber", ] @@ -15991,7 +16322,7 @@ dependencies = [ "serde", "serde_json", "syn 2.0.117", - "thiserror 2.0.17", + "thiserror 2.0.21", "unicode-ident", ] @@ -16035,7 +16366,7 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f0495d3e02c0e9b074e5631de7af419b9f4fdbfec4b1ca1b3180c8b3a926cb5" dependencies = [ - "thiserror 2.0.17", + "thiserror 2.0.21", "vu128", ] @@ -16363,7 +16694,7 @@ dependencies = [ "strum_macros 0.26.4", "syslog_loose", "termcolor", - "thiserror 2.0.17", + "thiserror 2.0.21", "tokio", "tracing", "ua-parser", @@ -16656,14 +16987,14 @@ version = "0.26.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "521bc38abb08001b01866da9f51eb7c5d647a19260e00054a8c7fd5f9e57f7a9" dependencies = [ - "webpki-roots 1.0.9", + "webpki-roots 1.0.10", ] [[package]] name = "webpki-roots" -version = "1.0.9" +version = "1.0.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7dcd9d09a39985f5344844e66b0c530a33843579125f23e21e9f0f220850f22a" +checksum = "1c1c8e20c8a38fc7d65abfe74eb7280e4aa667e8eec8bc84996d2874810748d9" dependencies = [ "rustls-pki-types", ] @@ -17320,8 +17651,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1301e935010a701ae5f8655edc0ad17c44bad3ac5ce8c39185f75453b720ae94" dependencies = [ "const-oid 0.9.6", - "der", - "spki", + "der 0.7.10", + "spki 0.7.3", "tls_codec", ] @@ -17334,13 +17665,13 @@ dependencies = [ "bcder", "bytes", "chrono", - "der", + "der 0.7.10", "hex", "pem", "ring", - "signature", - "spki", - "thiserror 2.0.17", + "signature 2.2.0", + "spki 0.7.3", + "thiserror 2.0.21", "zeroize", ] diff --git a/Cargo.toml b/Cargo.toml index 1674063f0c6..f275cb8aada 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -250,7 +250,7 @@ similar-asserts = "1.6.0" smallvec = { version = "1", features = ["serde"] } snafu = "0.8" sqlparser = { version = "0.62.0", default-features = false, features = ["std", "visitor", "serde"] } -sqlx = { version = "0.8", default-features = false, features = [ +sqlx = { version = "0.9", default-features = false, features = [ "any", "macros", "json", diff --git a/src/common/meta/Cargo.toml b/src/common/meta/Cargo.toml index 64303b070c8..d235f6e9ad1 100644 --- a/src/common/meta/Cargo.toml +++ b/src/common/meta/Cargo.toml @@ -81,6 +81,7 @@ serde_with.workspace = true snafu.workspace = true sqlx = { workspace = true, features = [ "mysql", + "mysql-rsa", "chrono", ], optional = true } store-api.workspace = true diff --git a/src/common/meta/src/election/rds/mysql.rs b/src/common/meta/src/election/rds/mysql.rs index bd694e4ae19..5b350d11957 100644 --- a/src/common/meta/src/election/rds/mysql.rs +++ b/src/common/meta/src/election/rds/mysql.rs @@ -22,7 +22,7 @@ use snafu::{OptionExt, ResultExt, ensure}; use sqlx::mysql::{MySqlArguments, MySqlRow}; use sqlx::pool::PoolConnection; use sqlx::query::Query; -use sqlx::{MySql, MySqlPool, MySqlTransaction, Row}; +use sqlx::{AssertSqlSafe, MySql, MySqlPool, MySqlTransaction, Row}; use tokio::sync::{Mutex, MutexGuard, broadcast}; use tokio::time::MissedTickBehavior; @@ -38,6 +38,17 @@ use crate::error::{ }; use crate::key::{CANDIDATES_ROOT, ELECTION_KEY}; +/// Builds a [`Query`] from a dynamic SQL template. +/// +/// sqlx 0.9 requires `&'static str` or [`AssertSqlSafe`] to reject unaudited +/// dynamic SQL strings. The templates here are generated by +/// [`ElectionSqlFactory`] from the configured table name (operator config, +/// not user input), and all dynamic data is passed via bind parameters, +/// so it is safe to mark them with [`AssertSqlSafe`]. +fn template_query(sql: &str) -> Query<'_, MySql, MySqlArguments> { + sqlx::query(AssertSqlSafe(sql)) +} + struct ElectionSqlFactory<'a> { table_name: &'a str, } @@ -268,11 +279,11 @@ impl ElectionMysqlClient { async fn ensure_table_exists(&mut self) -> Result<()> { let create_table_sql = self.create_table_sql(); - let query = sqlx::query(&create_table_sql); + let query = template_query(&create_table_sql); self.execute(query, &create_table_sql).await?; // Insert at least one row for `SELECT * FOR UPDATE` to work. let insert_once_sql = self.insert_once_sql(); - let query = sqlx::query(&insert_once_sql); + let query = template_query(&insert_once_sql); self.execute(query, &insert_once_sql).await?; Ok(()) } @@ -324,7 +335,7 @@ impl ElectionMysqlClient { /// if `current` is `None`. async fn check_version(&mut self, sql: &str) -> Result<()> { // Check if the MySQL version is supported. - let query = sqlx::query(sql); + let query = template_query(sql); // Safety: `maybe_init_client` ensures `current` is not `None`. let client = self.current.as_mut().unwrap(); let row = tokio::time::timeout(self.execution_timeout, query.fetch_one(&mut **client)) @@ -692,7 +703,7 @@ impl MySqlElection { executor: &mut Executor<'_>, ) -> Result> { let key = key.as_bytes(); - let query = sqlx::query(&self.sql_set.get_value_with_lease).bind(key); + let query = template_query(&self.sql_set.get_value_with_lease).bind(key); let res = executor .query(query, &self.sql_set.get_value_with_lease) .await?; @@ -728,7 +739,7 @@ impl MySqlElection { executor: &mut Executor<'_>, ) -> Result<(Vec<(String, Timestamp)>, Timestamp)> { let key_prefix = format!("{}%", key_prefix).as_bytes().to_vec(); - let query = sqlx::query(&self.sql_set.get_value_with_lease_by_prefix).bind(key_prefix); + let query = template_query(&self.sql_set.get_value_with_lease_by_prefix).bind(key_prefix); let res = executor .query(query, &self.sql_set.get_value_with_lease_by_prefix) .await?; @@ -765,7 +776,7 @@ impl MySqlElection { let prev = prev.as_bytes(); let updated = updated.as_bytes(); - let query = sqlx::query(&self.sql_set.update_value_with_lease) + let query = template_query(&self.sql_set.update_value_with_lease) .bind(updated) .bind(lease_ttl as f64) .bind(key) @@ -794,7 +805,7 @@ impl MySqlElection { ) -> Result { let key = key.as_bytes(); let lease_ttl_secs = lease_ttl_secs as f64; - let query = sqlx::query(&self.sql_set.put_value_with_lease) + let query = template_query(&self.sql_set.put_value_with_lease) .bind(key) .bind(value) .bind(lease_ttl_secs); @@ -808,7 +819,7 @@ impl MySqlElection { /// Caution: Should only delete the key if the lease is expired. async fn delete_value(&self, key: &str, executor: &mut Executor<'_>) -> Result { let key = key.as_bytes(); - let query = sqlx::query(&self.sql_set.delete_value).bind(key); + let query = template_query(&self.sql_set.delete_value).bind(key); let res = executor.execute(query, &self.sql_set.delete_value).await?; Ok(res == 1) @@ -837,7 +848,7 @@ impl MySqlElection { let mut client = self.client.lock().await; let txn = client.transaction().await?; let mut executor = Executor::Txn(txn); - let query = sqlx::query(&self.sql_set.campaign); + let query = template_query(&self.sql_set.campaign); executor.query(query, &self.sql_set.campaign).await?; // Safety: Checked if lease is not None above. self.renew_lease(executor, lease.unwrap()).await?; @@ -858,7 +869,7 @@ impl MySqlElection { let mut client = self.client.lock().await; let txn = client.transaction().await?; let mut executor = Executor::Txn(txn); - let query = sqlx::query(&self.sql_set.campaign); + let query = template_query(&self.sql_set.campaign); executor.query(query, &self.sql_set.campaign).await?; self.elected(executor, lease).await?; } @@ -869,7 +880,7 @@ impl MySqlElection { let mut client = self.client.lock().await; let txn = client.transaction().await?; let mut executor = Executor::Txn(txn); - let query = sqlx::query(&self.sql_set.campaign); + let query = template_query(&self.sql_set.campaign); executor.query(query, &self.sql_set.campaign).await?; // Safety: Checked if lease is not None above. self.renew_lease(executor, lease.unwrap()).await?; @@ -1027,7 +1038,7 @@ mod tests { async fn drop_table(client: &Mutex, table_name: &str) { let mut client = client.lock().await; let sql = format!("DROP TABLE IF EXISTS `{}`;", table_name); - client.execute(sqlx::query(&sql), &sql).await.unwrap(); + client.execute(template_query(&sql), &sql).await.unwrap(); } #[tokio::test] @@ -1052,7 +1063,7 @@ mod tests { let txn = a.transaction().await.unwrap(); let mut executor = Executor::Txn(txn); let raw_query = format!("SELECT * FROM `{}` FOR UPDATE;", table_name); - let query = sqlx::query(&raw_query); + let query = template_query(&raw_query); let _ = executor.query(query, &raw_query).await.unwrap(); } @@ -1258,7 +1269,7 @@ mod tests { let txn = client.transaction().await.unwrap(); let mut executor = Executor::Txn(txn); let raw_query = format!("SELECT * FROM `{}` FOR UPDATE;", table_name); - let query = sqlx::query(&raw_query); + let query = template_query(&raw_query); let _ = executor.query(query, &raw_query).await.unwrap(); election.elected(executor, expected_lease).await } diff --git a/src/common/meta/src/kv_backend/rds/mysql.rs b/src/common/meta/src/kv_backend/rds/mysql.rs index 6791d58f933..c933bb345f6 100644 --- a/src/common/meta/src/kv_backend/rds/mysql.rs +++ b/src/common/meta/src/kv_backend/rds/mysql.rs @@ -21,7 +21,7 @@ use regex::Regex; use snafu::{OptionExt, ResultExt}; use sqlx::mysql::MySqlRow; use sqlx::pool::Pool; -use sqlx::{MySql, MySqlPool, Row, Transaction as MySqlTransaction}; +use sqlx::{AssertSqlSafe, MySql, MySqlPool, Row, Transaction as MySqlTransaction}; use strum::AsRefStr; use crate::error::{ @@ -264,7 +264,9 @@ impl Executor for MySqlClient { } async fn query(&mut self, raw_query: &str, params: &[&Vec]) -> Result> { - let query = sqlx::query(raw_query); + // Safety: `raw_query` is generated by `MySqlTemplateSet`, and all dynamic data + // is passed via bind parameters. + let query = sqlx::query(AssertSqlSafe(raw_query)); let query = params.iter().fold(query, |query, param| query.bind(param)); let rows = query .fetch_all(&**self) @@ -274,7 +276,9 @@ impl Executor for MySqlClient { } async fn execute(&mut self, raw_query: &str, params: &[&Vec]) -> Result<()> { - let query = sqlx::query(raw_query); + // Safety: `raw_query` is generated by `MySqlTemplateSet`, and all dynamic data + // is passed via bind parameters. + let query = sqlx::query(AssertSqlSafe(raw_query)); let query = params.iter().fold(query, |query, param| query.bind(param)); query .execute(&**self) @@ -303,7 +307,9 @@ impl Executor for MySqlClient { #[async_trait::async_trait] impl Transaction<'_> for MySqlTxnClient { async fn query(&mut self, raw_query: &str, params: &[&Vec]) -> Result> { - let query = sqlx::query(raw_query); + // Safety: `raw_query` is generated by `MySqlTemplateSet`, and all dynamic data + // is passed via bind parameters. + let query = sqlx::query(AssertSqlSafe(raw_query)); let query = params.iter().fold(query, |query, param| query.bind(param)); // As said in https://docs.rs/sqlx/latest/sqlx/trait.Executor.html, we need a `&mut *transaction`. Weird. let rows = query @@ -314,7 +320,9 @@ impl Transaction<'_> for MySqlTxnClient { } async fn execute(&mut self, raw_query: &str, params: &[&Vec]) -> Result<()> { - let query = sqlx::query(raw_query); + // Safety: `raw_query` is generated by `MySqlTemplateSet`, and all dynamic data + // is passed via bind parameters. + let query = sqlx::query(AssertSqlSafe(raw_query)); let query = params.iter().fold(query, |query, param| query.bind(param)); // As said in https://docs.rs/sqlx/latest/sqlx/trait.Executor.html, we need a `&mut *transaction`. Weird. query @@ -562,12 +570,14 @@ impl MySqlStore { pool: &Pool, sql_template_set: &MySqlTemplateSet, ) -> Result> { - let row = sqlx::query(&sql_template_set.show_create_table_statement) - .fetch_optional(pool) - .await - .with_context(|_| MySqlExecutionSnafu { - sql: sql_template_set.show_create_table_statement.clone(), - })?; + let row = sqlx::query(AssertSqlSafe( + sql_template_set.show_create_table_statement.as_str(), + )) + .fetch_optional(pool) + .await + .with_context(|_| MySqlExecutionSnafu { + sql: sql_template_set.show_create_table_statement.clone(), + })?; Ok(row.map(|row| row.get(1))) } @@ -598,12 +608,14 @@ impl MySqlStore { match Self::parse_value_column_blob_type(&create_table_sql) { Some(ValueBlobType::Blob) => { - sqlx::query(&sql_template_set.alter_value_column_statement) - .execute(pool) - .await - .with_context(|_| MySqlExecutionSnafu { - sql: sql_template_set.alter_value_column_statement.clone(), - })?; + sqlx::query(AssertSqlSafe( + sql_template_set.alter_value_column_statement.as_str(), + )) + .execute(pool) + .await + .with_context(|_| MySqlExecutionSnafu { + sql: sql_template_set.alter_value_column_statement.clone(), + })?; info!("Upgraded MySQL metadata value column to MEDIUMBLOB for `{table_name}`"); } Some(ValueBlobType::MediumBlob | ValueBlobType::LongBlob) => { @@ -638,12 +650,14 @@ impl MySqlStore { // We check if greptime_metakv table exists, and we will create a new table // if it does not exist. let sql_template_set = MySqlTemplateFactory::new(table_name).build(); - sqlx::query(&sql_template_set.create_table_statement) - .execute(&pool) - .await - .with_context(|_| MySqlExecutionSnafu { - sql: sql_template_set.create_table_statement.clone(), - })?; + sqlx::query(AssertSqlSafe( + sql_template_set.create_table_statement.as_str(), + )) + .execute(&pool) + .await + .with_context(|_| MySqlExecutionSnafu { + sql: sql_template_set.create_table_statement.clone(), + })?; Self::maybe_upgrade_value_column_to_mediumblob(&pool, &sql_template_set).await?; Ok(Arc::new(MySqlStore { max_txn_ops, @@ -695,7 +709,10 @@ mod tests { async fn show_create_table(pool: &MySqlPool, table_name: &str) -> String { let sql = format!("SHOW CREATE TABLE `{table_name}`"); - let row = sqlx::query(&sql).fetch_one(pool).await.unwrap(); + let row = sqlx::query(AssertSqlSafe(sql.as_str())) + .fetch_one(pool) + .await + .unwrap(); row.get::(1) } @@ -703,18 +720,24 @@ mod tests { let sql = format!( "CREATE TABLE IF NOT EXISTS `{table_name}`(k VARBINARY(3072) PRIMARY KEY, v BLOB);" ); - sqlx::query(&sql).execute(pool).await.unwrap(); + sqlx::query(AssertSqlSafe(sql.as_str())) + .execute(pool) + .await + .unwrap(); } async fn drop_table(pool: &MySqlPool, table_name: &str) { let sql = format!("DROP TABLE IF EXISTS `{table_name}`;"); - sqlx::query(&sql).execute(pool).await.unwrap(); + sqlx::query(AssertSqlSafe(sql.as_str())) + .execute(pool) + .await + .unwrap(); } async fn build_mysql_kv_backend(table_name: &str) -> Option { let pool = mysql_pool().await?; let sql_templates = MySqlTemplateFactory::new(table_name).build(); - sqlx::query(&sql_templates.create_table_statement) + sqlx::query(AssertSqlSafe(sql_templates.create_table_statement.as_str())) .execute(&pool) .await .unwrap(); diff --git a/src/meta-srv/Cargo.toml b/src/meta-srv/Cargo.toml index 6baca0934e1..fa339d76774 100644 --- a/src/meta-srv/Cargo.toml +++ b/src/meta-srv/Cargo.toml @@ -67,6 +67,7 @@ servers.workspace = true snafu.workspace = true sqlx = { workspace = true, features = [ "mysql", + "mysql-rsa", "chrono", ], optional = true } store-api.workspace = true diff --git a/tests-fuzz/Cargo.toml b/tests-fuzz/Cargo.toml index 1189f53c42f..9e0431eedcf 100644 --- a/tests-fuzz/Cargo.toml +++ b/tests-fuzz/Cargo.toml @@ -59,6 +59,7 @@ sql = { workspace = true } sqlparser.workspace = true sqlx = { workspace = true, features = [ "mysql", + "mysql-rsa", "chrono", ] } store-api = { workspace = true } diff --git a/tests-fuzz/src/utils.rs b/tests-fuzz/src/utils.rs index 4b057d1dd65..0983d51352c 100644 --- a/tests-fuzz/src/utils.rs +++ b/tests-fuzz/src/utils.rs @@ -110,7 +110,7 @@ pub const GT_FUZZ_CLUSTER_NAME: &str = "GT_FUZZ_CLUSTER_NAME"; /// Flushes memtable to SST file. pub async fn flush_memtable(e: &Pool, table_name: &Ident) -> Result<()> { let sql = format!("admin flush_table(\"{}\")", table_name); - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(e) .await .context(error::ExecuteQuerySnafu { sql })?; @@ -122,7 +122,7 @@ pub async fn flush_memtable(e: &Pool, table_name: &Ident) -> Result<()> { /// Triggers a compaction for table pub async fn compact_table(e: &Pool, table_name: &Ident) -> Result<()> { let sql = format!("admin compact_table(\"{}\")", table_name); - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(e) .await .context(error::ExecuteQuerySnafu { sql })?; diff --git a/tests-fuzz/src/utils/migration.rs b/tests-fuzz/src/utils/migration.rs index 043e2271ae9..91863171163 100644 --- a/tests-fuzz/src/utils/migration.rs +++ b/tests-fuzz/src/utils/migration.rs @@ -37,7 +37,7 @@ pub async fn migrate_region( ) -> String { let sql = format!("admin migrate_region({region_id}, {from_peer_id}, {to_peer_id}, {timeout_secs});"); - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .fetch_one(e) .await .context(error::ExecuteQuerySnafu { sql }) diff --git a/tests-fuzz/src/utils/procedure.rs b/tests-fuzz/src/utils/procedure.rs index 72259272120..e89dbda183c 100644 --- a/tests-fuzz/src/utils/procedure.rs +++ b/tests-fuzz/src/utils/procedure.rs @@ -24,7 +24,7 @@ use crate::error; /// Fetches the state of a procedure. pub async fn procedure_state(e: &Pool, procedure_id: &str) -> String { let sql = format!("admin procedure_state(\"{procedure_id}\");"); - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .fetch_one(e) .await .context(error::ExecuteQuerySnafu { sql }) diff --git a/tests-fuzz/src/validator/column.rs b/tests-fuzz/src/validator/column.rs index d092d892e9c..19cd2849b96 100644 --- a/tests-fuzz/src/validator/column.rs +++ b/tests-fuzz/src/validator/column.rs @@ -221,7 +221,7 @@ pub async fn valid_marker_value( let sql = format!("SELECT * FROM {table_name}"); // cache is useless and buggy anyway since alter happens all the time // TODO(discord9): invalid prepared sql after alter - let res = sqlx::query(&sql) + let res = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .persistent(false) .fetch_all(db) .await diff --git a/tests-fuzz/src/validator/row.rs b/tests-fuzz/src/validator/row.rs index 5bd83d80d6b..d3b371cb686 100644 --- a/tests-fuzz/src/validator/row.rs +++ b/tests-fuzz/src/validator/row.rs @@ -155,14 +155,16 @@ pub struct ValueCount { } pub async fn count_values(db: &MySqlPool, sql: &str) -> Result { - sqlx::query_as::<_, ValueCount>(sql) + // Safety: the SQL is built by fuzz validators from generated identifiers only. + sqlx::query_as::<_, ValueCount>(sqlx::AssertSqlSafe(sql)) .fetch_one(db) .await .context(error::ExecuteQuerySnafu { sql }) } pub async fn count_values_all(db: &MySqlPool, sql: &str) -> Result> { - sqlx::query_as::<_, ValueCount>(sql) + // Safety: the SQL is built by fuzz validators from generated identifiers only. + sqlx::query_as::<_, ValueCount>(sqlx::AssertSqlSafe(sql)) .fetch_all(db) .await .context(error::ExecuteQuerySnafu { sql }) @@ -170,7 +172,8 @@ pub async fn count_values_all(db: &MySqlPool, sql: &str) -> Result Result> { - sqlx::query(sql) + // Safety: the SQL is built by fuzz validators from generated identifiers only. + sqlx::query(sqlx::AssertSqlSafe(sql)) .fetch_all(db) .await .context(error::ExecuteQuerySnafu { sql }) diff --git a/tests-fuzz/src/validator/table.rs b/tests-fuzz/src/validator/table.rs index f271101051e..08f6e28d9b7 100644 --- a/tests-fuzz/src/validator/table.rs +++ b/tests-fuzz/src/validator/table.rs @@ -47,7 +47,8 @@ fn parse_show_create(show_create: &str) -> Result> { /// Fetches table options from the context pub async fn fetch_table_options(db: &MySqlPool, sql: &str) -> Result> { - let fetched_rows = sqlx::query(sql) + // Safety: the SQL is built by fuzz validators from generated identifiers only. + let fetched_rows = sqlx::query(sqlx::AssertSqlSafe(sql)) .fetch_all(db) .await .context(error::ExecuteQuerySnafu { sql })?; diff --git a/tests-fuzz/targets/ddl/fuzz_alter_logical_table.rs b/tests-fuzz/targets/ddl/fuzz_alter_logical_table.rs index c7542a46237..11dd7cbae4c 100644 --- a/tests-fuzz/targets/ddl/fuzz_alter_logical_table.rs +++ b/tests-fuzz/targets/ddl/fuzz_alter_logical_table.rs @@ -123,7 +123,7 @@ async fn execute_alter_table(ctx: FuzzContext, input: FuzzInput) -> Result<()> { let create_physical_table_expr = generate_create_physical_table_expr(&mut rng).unwrap(); let translator = CreateTableExprTranslator; let sql = translator.translate(&create_physical_table_expr)?; - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql: &sql })?; @@ -154,7 +154,7 @@ async fn execute_alter_table(ctx: FuzzContext, input: FuzzInput) -> Result<()> { let create_logical_table_expr = generate_create_logical_table_expr(physical_table_ctx, &mut rng).unwrap(); let sql = translator.translate(&create_logical_table_expr)?; - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql: &sql })?; @@ -178,7 +178,7 @@ async fn execute_alter_table(ctx: FuzzContext, input: FuzzInput) -> Result<()> { .unwrap(); let translator = AlterTableExprTranslator; let sql = translator.translate(&expr)?; - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql: &sql })?; @@ -213,7 +213,7 @@ async fn execute_alter_table(ctx: FuzzContext, input: FuzzInput) -> Result<()> { // Clean up logical table let sql = format!("DROP TABLE {}", create_logical_table_expr.table_name); - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql: &sql })?; @@ -224,7 +224,7 @@ async fn execute_alter_table(ctx: FuzzContext, input: FuzzInput) -> Result<()> { // Clean up physical table let sql = format!("DROP TABLE {}", create_physical_table_expr.table_name); - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql })?; diff --git a/tests-fuzz/targets/ddl/fuzz_alter_table.rs b/tests-fuzz/targets/ddl/fuzz_alter_table.rs index 773e75c2b89..37002943128 100644 --- a/tests-fuzz/targets/ddl/fuzz_alter_table.rs +++ b/tests-fuzz/targets/ddl/fuzz_alter_table.rs @@ -229,14 +229,14 @@ async fn execute_alter_table(ctx: FuzzContext, input: FuzzInput) -> Result<()> { let expr = generate_create_table_expr(&mut rng).unwrap(); let translator = CreateTableExprTranslator; let sql = translator.translate(&expr)?; - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql: &sql })?; info!("Create table: {sql}, result: {result:?}"); let insert = generate_insert_table_sql(&mut rng, &expr)?; - let result = sqlx::query(&insert) + let result = sqlx::query(sqlx::AssertSqlSafe(insert.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql: &insert })?; @@ -248,7 +248,7 @@ async fn execute_alter_table(ctx: FuzzContext, input: FuzzInput) -> Result<()> { let expr = generate_alter_table_expr(table_ctx.clone(), &mut rng).unwrap(); let translator = AlterTableExprTranslator; let sql = translator.translate(&expr)?; - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql: &sql })?; @@ -303,7 +303,7 @@ async fn execute_alter_table(ctx: FuzzContext, input: FuzzInput) -> Result<()> { // select from table to make sure the table is still ok let sql = format!("SELECT * FROM {}", table_ctx.name); - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .persistent(false) .execute(&ctx.greptime) .await @@ -314,7 +314,7 @@ async fn execute_alter_table(ctx: FuzzContext, input: FuzzInput) -> Result<()> { // Cleans up let table_name = table_ctx.name.clone(); let sql = format!("DROP TABLE {}", table_name); - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql })?; diff --git a/tests-fuzz/targets/ddl/fuzz_create_database.rs b/tests-fuzz/targets/ddl/fuzz_create_database.rs index a6b2071cac4..60e46e1a89c 100644 --- a/tests-fuzz/targets/ddl/fuzz_create_database.rs +++ b/tests-fuzz/targets/ddl/fuzz_create_database.rs @@ -74,7 +74,7 @@ async fn execute_create_database(ctx: FuzzContext, input: FuzzInput) -> Result<( let expr = generate_expr(input)?; let translator = CreateDatabaseExprTranslator; let sql = translator.translate(&expr)?; - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql: &sql })?; @@ -82,7 +82,7 @@ async fn execute_create_database(ctx: FuzzContext, input: FuzzInput) -> Result<( // Cleans up let sql = format!("DROP DATABASE {}", expr.database_name); - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql })?; diff --git a/tests-fuzz/targets/ddl/fuzz_create_logical_table.rs b/tests-fuzz/targets/ddl/fuzz_create_logical_table.rs index e36b5782099..1500c609f1f 100644 --- a/tests-fuzz/targets/ddl/fuzz_create_logical_table.rs +++ b/tests-fuzz/targets/ddl/fuzz_create_logical_table.rs @@ -85,7 +85,7 @@ async fn execute_create_logic_table(ctx: FuzzContext, input: FuzzInput) -> Resul .generate(&mut rng)?; let translator = CreateTableExprTranslator; let sql = translator.translate(&create_physical_table_expr)?; - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql: &sql })?; @@ -128,7 +128,7 @@ async fn execute_create_logic_table(ctx: FuzzContext, input: FuzzInput) -> Resul .unwrap() .generate(&mut rng)?; let sql = translator.translate(&create_logical_table_expr)?; - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql: &sql })?; @@ -162,7 +162,7 @@ async fn execute_create_logic_table(ctx: FuzzContext, input: FuzzInput) -> Resul // Clean up logical table let sql = format!("DROP TABLE {}", create_logical_table_expr.table_name); - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql: &sql })?; @@ -173,7 +173,7 @@ async fn execute_create_logic_table(ctx: FuzzContext, input: FuzzInput) -> Resul // Clean up physical table let sql = format!("DROP TABLE {}", create_physical_table_expr.table_name); - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql })?; diff --git a/tests-fuzz/targets/ddl/fuzz_create_table.rs b/tests-fuzz/targets/ddl/fuzz_create_table.rs index 71171cb076d..a1597256bc5 100644 --- a/tests-fuzz/targets/ddl/fuzz_create_table.rs +++ b/tests-fuzz/targets/ddl/fuzz_create_table.rs @@ -93,7 +93,7 @@ async fn execute_create_table(ctx: FuzzContext, input: FuzzInput) -> Result<()> let expr = generate_expr(input)?; let translator = CreateTableExprTranslator; let sql = translator.translate(&expr)?; - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql: &sql })?; @@ -110,7 +110,7 @@ async fn execute_create_table(ctx: FuzzContext, input: FuzzInput) -> Result<()> // Cleans up let sql = format!("DROP TABLE {}", expr.table_name); - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql })?; diff --git a/tests-fuzz/targets/ddl/fuzz_repartition_metric_table.rs b/tests-fuzz/targets/ddl/fuzz_repartition_metric_table.rs index 8a6bd81fa80..461bfe51daa 100644 --- a/tests-fuzz/targets/ddl/fuzz_repartition_metric_table.rs +++ b/tests-fuzz/targets/ddl/fuzz_repartition_metric_table.rs @@ -152,7 +152,7 @@ async fn create_metric_tables( let create_physical_expr = generate_create_physical_table_expr(partitions, rng)?; let translator = CreateTableExprTranslator; let create_physical_sql = translator.translate(&create_physical_expr)?; - let result = sqlx::query(&create_physical_sql) + let result = sqlx::query(sqlx::AssertSqlSafe(create_physical_sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { @@ -180,7 +180,7 @@ async fn create_metric_tables( } let create_logical_sql = translator.translate(&create_logical_expr)?; - let result = sqlx::query(&create_logical_sql) + let result = sqlx::query(sqlx::AssertSqlSafe(create_logical_sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { @@ -212,7 +212,7 @@ async fn execute_insert_with_retry(ctx: &FuzzContext, sql: &str) -> Result<()> { let mut attempt = 0; let max_attempts = 10; loop { - match sqlx::query(sql) + match sqlx::query(sqlx::AssertSqlSafe(sql)) .persistent(false) .execute(&ctx.greptime) .await @@ -408,7 +408,7 @@ async fn cleanup_tables( ) -> Result<()> { for table_ctx in logical_tables.values() { let drop_logical_sql = format!("DROP TABLE {}", table_ctx.name); - let result = sqlx::query(&drop_logical_sql) + let result = sqlx::query(sqlx::AssertSqlSafe(drop_logical_sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { @@ -418,7 +418,7 @@ async fn cleanup_tables( } let drop_physical_sql = format!("DROP TABLE {}", physical_table_ctx.name); - let result = sqlx::query(&drop_physical_sql) + let result = sqlx::query(sqlx::AssertSqlSafe(drop_physical_sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { @@ -572,7 +572,7 @@ async fn execute_repartition_metric_table(ctx: FuzzContext, input: FuzzInput) -> info!("Repartition sql: {sql}"); let started_at_ms = current_time_millis(); let now = Instant::now(); - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql: &sql })?; diff --git a/tests-fuzz/targets/ddl/fuzz_repartition_table.rs b/tests-fuzz/targets/ddl/fuzz_repartition_table.rs index 4f6a014f2e7..d1f633014b0 100644 --- a/tests-fuzz/targets/ddl/fuzz_repartition_table.rs +++ b/tests-fuzz/targets/ddl/fuzz_repartition_table.rs @@ -187,7 +187,7 @@ async fn execute_insert_with_retry(ctx: &FuzzContext, sql: &str) -> Result<()> { match ctx .greptime // unprepared query, see - .execute(sql) + .execute(sqlx::AssertSqlSafe(sql)) .await { Ok(_) => { @@ -250,7 +250,7 @@ async fn execute_repartition_table(ctx: FuzzContext, input: FuzzInput) -> Result let expr = generate_create_expr(&input, &mut rng).unwrap(); let translator = CreateTableExprTranslator; let sql = translator.translate(&expr)?; - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql: &sql })?; @@ -271,7 +271,7 @@ async fn execute_repartition_table(ctx: FuzzContext, input: FuzzInput) -> Result let translator = RepartitionExprTranslator; let sql = translator.translate(&expr)?; info!("Initial partition sql: {sql}"); - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql: &sql })?; @@ -306,7 +306,7 @@ async fn execute_repartition_table(ctx: FuzzContext, input: FuzzInput) -> Result let translator = RepartitionExprTranslator; let sql = translator.translate(&expr)?; info!("Repartition sql: {sql}"); - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql: &sql })?; @@ -378,7 +378,7 @@ async fn execute_repartition_table(ctx: FuzzContext, input: FuzzInput) -> Result // Cleans up let table_name = table_ctx.name.clone(); let sql = format!("DROP TABLE {}", table_name); - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql })?; diff --git a/tests-fuzz/targets/ddl/fuzz_repartition_table_chaos.rs b/tests-fuzz/targets/ddl/fuzz_repartition_table_chaos.rs index c1e570f6981..ed0b963908b 100644 --- a/tests-fuzz/targets/ddl/fuzz_repartition_table_chaos.rs +++ b/tests-fuzz/targets/ddl/fuzz_repartition_table_chaos.rs @@ -184,7 +184,7 @@ async fn execute_insert_with_retry(ctx: &FuzzContext, sql: &str) -> Result<()> { let mut attempt = 0; let max_attempts = 10; loop { - match ctx.greptime.execute(sql).await { + match ctx.greptime.execute(sqlx::AssertSqlSafe(sql)).await { Ok(_) => return Ok(()), Err(err) => { tokio::time::sleep(delay).await; @@ -201,7 +201,7 @@ async fn execute_insert_with_retry(ctx: &FuzzContext, sql: &str) -> Result<()> { async fn create_table(ctx: &FuzzContext, expr: &CreateTableExpr) -> Result { let translator = CreateTableExprTranslator; let sql = translator.translate(expr)?; - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql: &sql })?; @@ -333,7 +333,7 @@ fn alter_table_partitions_expr_from_table_ctx( async fn submit_repartition_procedure(ctx: &FuzzContext, expr: &RepartitionExpr) -> Result { let translator = RepartitionExprTranslator; let sql = translator.translate(expr)?; - let (procedure_id,) = sqlx::query_as::<_, (String,)>(&sql) + let (procedure_id,) = sqlx::query_as::<_, (String,)>(sqlx::AssertSqlSafe(sql.as_str())) .fetch_one(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql: &sql })?; @@ -418,7 +418,7 @@ async fn validate_terminal_metadata( async fn drop_table(ctx: &FuzzContext, table_name: &tests_fuzz::ir::Ident) -> Result<()> { let sql = format!("DROP TABLE {}", table_name); - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql: &sql })?; diff --git a/tests-fuzz/targets/failover/fuzz_failover_metric_regions.rs b/tests-fuzz/targets/failover/fuzz_failover_metric_regions.rs index 06623c42350..1f8cc67d775 100644 --- a/tests-fuzz/targets/failover/fuzz_failover_metric_regions.rs +++ b/tests-fuzz/targets/failover/fuzz_failover_metric_regions.rs @@ -110,7 +110,7 @@ async fn create_physical_table( let create_physical_table_expr = generate_create_physical_table_expr(rng).unwrap(); let translator = CreateTableExprTranslator; let sql = translator.translate(&create_physical_table_expr)?; - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql: &sql })?; @@ -167,7 +167,7 @@ async fn insert_values( let result = ctx .greptime // unprepared query, see - .execute(sql.as_str()) + .execute(sqlx::AssertSqlSafe(sql.as_str())) .await .context(error::ExecuteQuerySnafu { sql: &sql })?; @@ -201,7 +201,7 @@ async fn execute_failover(ctx: FuzzContext, input: FuzzInput) -> Result<()> { continue; } let sql = translator.translate(&create_logical_table_expr)?; - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql: &sql })?; @@ -272,7 +272,7 @@ async fn execute_failover(ctx: FuzzContext, input: FuzzInput) -> Result<()> { // Clean up for (table_ctx, _) in tables.values() { let sql = format!("DROP TABLE {}", table_ctx.name); - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql })?; @@ -280,7 +280,7 @@ async fn execute_failover(ctx: FuzzContext, input: FuzzInput) -> Result<()> { } let sql = format!("DROP TABLE {}", physical_table_ctx.name); - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql })?; diff --git a/tests-fuzz/targets/failover/fuzz_failover_mito_regions.rs b/tests-fuzz/targets/failover/fuzz_failover_mito_regions.rs index da3b7dbcf3c..cfb2ae13420 100644 --- a/tests-fuzz/targets/failover/fuzz_failover_mito_regions.rs +++ b/tests-fuzz/targets/failover/fuzz_failover_mito_regions.rs @@ -158,7 +158,7 @@ async fn execute_create_exprs( let _permit = semaphore.acquire().await.unwrap(); let translator = CreateTableExprTranslator; let sql = translator.translate(expr).unwrap(); - sqlx::query(&sql) + sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&greptime) .await .context(error::ExecuteQuerySnafu { sql: &sql }) @@ -225,7 +225,7 @@ async fn execute_insert_exprs( let sql = translator.translate(&insert_expr)?; let result = greptime // unprepared query, see - .execute(sql.as_str()) + .execute(sqlx::AssertSqlSafe(sql.as_str())) .await .context(error::ExecuteQuerySnafu { sql: &sql })?; ensure!( @@ -337,7 +337,7 @@ async fn execute_failover(ctx: FuzzContext, input: FuzzInput) -> Result<()> { for table_ctx in table_ctxs { let sql = format!("DROP TABLE {}", table_ctx.name); - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql })?; diff --git a/tests-fuzz/targets/fuzz_insert.rs b/tests-fuzz/targets/fuzz_insert.rs index bc690a48776..8a5192d48eb 100644 --- a/tests-fuzz/targets/fuzz_insert.rs +++ b/tests-fuzz/targets/fuzz_insert.rs @@ -130,7 +130,7 @@ async fn execute_insert(ctx: FuzzContext, input: FuzzInput) -> Result<()> { let create_expr = generate_create_expr(input, &mut rng)?; let translator = CreateTableExprTranslator; let sql = translator.translate(&create_expr)?; - let _result = sqlx::query(&sql) + let _result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql: &sql })?; @@ -142,7 +142,7 @@ async fn execute_insert(ctx: FuzzContext, input: FuzzInput) -> Result<()> { let result = ctx .greptime // unprepared query, see - .execute(sql.as_str()) + .execute(sqlx::AssertSqlSafe(sql.as_str())) .await .context(error::ExecuteQuerySnafu { sql: &sql })?; @@ -202,7 +202,7 @@ async fn execute_insert(ctx: FuzzContext, input: FuzzInput) -> Result<()> { // Cleans up let sql = format!("DROP TABLE {}", create_expr.table_name); - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql })?; diff --git a/tests-fuzz/targets/fuzz_insert_logical_table.rs b/tests-fuzz/targets/fuzz_insert_logical_table.rs index 289856f771c..dcaada5edfd 100644 --- a/tests-fuzz/targets/fuzz_insert_logical_table.rs +++ b/tests-fuzz/targets/fuzz_insert_logical_table.rs @@ -141,7 +141,7 @@ async fn create_physical_table( let create_physical_table_expr = generate_create_physical_table_expr(rng).unwrap(); let translator = CreateTableExprTranslator; let sql = translator.translate(&create_physical_table_expr)?; - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql: &sql })?; @@ -211,7 +211,7 @@ async fn insert_values( let result = ctx .greptime // unprepared query, see - .execute(sql.as_str()) + .execute(sqlx::AssertSqlSafe(sql.as_str())) .await .context(error::ExecuteQuerySnafu { sql: &sql })?; @@ -248,7 +248,7 @@ async fn execute_insert(ctx: FuzzContext, input: FuzzInput) -> Result<()> { continue; } let sql = translator.translate(&create_logical_table_expr)?; - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql: &sql })?; @@ -272,7 +272,7 @@ async fn execute_insert(ctx: FuzzContext, input: FuzzInput) -> Result<()> { // Clean up logical table for (table_name, _) in tables { let sql = format!("DROP TABLE {}", table_name); - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql: &sql })?; @@ -281,7 +281,7 @@ async fn execute_insert(ctx: FuzzContext, input: FuzzInput) -> Result<()> { // Clean up physical table let sql = format!("DROP TABLE {}", physical_table_ctx.name); - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql })?; diff --git a/tests-fuzz/targets/migration/fuzz_migrate_metric_regions.rs b/tests-fuzz/targets/migration/fuzz_migrate_metric_regions.rs index d727eefa305..38bd6269966 100644 --- a/tests-fuzz/targets/migration/fuzz_migrate_metric_regions.rs +++ b/tests-fuzz/targets/migration/fuzz_migrate_metric_regions.rs @@ -109,7 +109,7 @@ async fn create_physical_table( let create_physical_table_expr = generate_create_physical_table_expr(rng).unwrap(); let translator = CreateTableExprTranslator; let sql = translator.translate(&create_physical_table_expr)?; - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql: &sql })?; @@ -166,7 +166,7 @@ async fn insert_values( let result = ctx .greptime // unprepared query, see - .execute(sql.as_str()) + .execute(sqlx::AssertSqlSafe(sql.as_str())) .await .context(error::ExecuteQuerySnafu { sql: &sql })?; @@ -208,7 +208,7 @@ async fn create_logical_table_and_insert_values( continue; } let sql = translator.translate(&create_logical_table_expr)?; - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql: &sql })?; @@ -406,7 +406,7 @@ async fn execute_migration(ctx: FuzzContext, input: FuzzInput) -> Result<()> { // Clean up for (table_ctx, _) in tables.values() { let sql = format!("DROP TABLE {}", table_ctx.name); - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql })?; @@ -414,7 +414,7 @@ async fn execute_migration(ctx: FuzzContext, input: FuzzInput) -> Result<()> { } let sql = format!("DROP TABLE {}", physical_table_ctx.name); - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql })?; diff --git a/tests-fuzz/targets/migration/fuzz_migrate_mito_regions.rs b/tests-fuzz/targets/migration/fuzz_migrate_mito_regions.rs index 17cbfb9251d..7da074f202d 100644 --- a/tests-fuzz/targets/migration/fuzz_migrate_mito_regions.rs +++ b/tests-fuzz/targets/migration/fuzz_migrate_mito_regions.rs @@ -158,7 +158,7 @@ async fn insert_values( let result = ctx .greptime // unprepared query, see - .execute(sql.as_str()) + .execute(sqlx::AssertSqlSafe(sql.as_str())) .await .context(error::ExecuteQuerySnafu { sql: &sql })?; ensure!( @@ -208,7 +208,11 @@ async fn validate_insert_exprs( primary_keys_column_list ); info!("Executing sql: {select_sql}"); - let fetched_rows = ctx.greptime.fetch_all(select_sql.as_str()).await.unwrap(); + let fetched_rows = ctx + .greptime + .fetch_all(sqlx::AssertSqlSafe(select_sql.as_str())) + .await + .unwrap(); let mut expected_rows = replace_default(&insert_expr.values_list, table_ctx, insert_expr); sort_by_primary_keys(&mut expected_rows, primary_keys_idx); validator::row::assert_eq::(&insert_expr.columns, &fetched_rows, &expected_rows)?; @@ -224,7 +228,11 @@ async fn validate_insert_exprs( column_list, table_ctx.name, ts_column.name, ts_value, primary_keys_column_list ); info!("Executing sql: {select_sql}"); - let fetched_rows = ctx.greptime.fetch_all(select_sql.as_str()).await.unwrap(); + let fetched_rows = ctx + .greptime + .fetch_all(sqlx::AssertSqlSafe(select_sql.as_str())) + .await + .unwrap(); let mut expected_rows = replace_default(&insert_expr.values_list, table_ctx, insert_expr); sort_by_primary_keys(&mut expected_rows, primary_keys_idx); validator::row::assert_eq::(&insert_expr.columns, &fetched_rows, &expected_rows)?; @@ -291,7 +299,7 @@ async fn execute_region_migration(ctx: FuzzContext, input: FuzzInput) -> Result< let create_expr = generate_create_expr(input, &mut rng)?; let translator = CreateTableExprTranslator; let sql = translator.translate(&create_expr)?; - let _result = sqlx::query(&sql) + let _result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql: &sql })?; @@ -364,7 +372,7 @@ async fn execute_region_migration(ctx: FuzzContext, input: FuzzInput) -> Result< // Cleans up let sql = format!("DROP TABLE {}", table_ctx.name); - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql })?; diff --git a/tests-fuzz/targets/unstable/fuzz_create_table_standalone.rs b/tests-fuzz/targets/unstable/fuzz_create_table_standalone.rs index 76f978e3133..88c59f78f4c 100644 --- a/tests-fuzz/targets/unstable/fuzz_create_table_standalone.rs +++ b/tests-fuzz/targets/unstable/fuzz_create_table_standalone.rs @@ -125,7 +125,10 @@ async fn execute_unstable_create_table( let mysql_public = connect_mysql(DEFAULT_MYSQL_URL, "public").await; loop { let sql = format!("CREATE DATABASE IF NOT EXISTS {FUZZ_TESTS_DATABASE}"); - match sqlx::query(&sql).execute(&mysql_public).await { + match sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) + .execute(&mysql_public) + .await + { Ok(result) => { info!("Create database: {}, result: {result:?}", sql); break; @@ -150,7 +153,9 @@ async fn execute_unstable_create_table( let translator = CreateTableExprTranslator; let sql = translator.translate(&expr).unwrap(); - let result = sqlx::query(&sql).execute(&ctx.greptime).await; + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) + .execute(&ctx.greptime) + .await; match result { Ok(result) => { let state = *rx.borrow(); @@ -179,7 +184,10 @@ async fn execute_unstable_create_table( loop { let sql = format!("DROP DATABASE IF EXISTS {FUZZ_TESTS_DATABASE}"); - match sqlx::query(&sql).execute(&mysql_public).await { + match sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) + .execute(&mysql_public) + .await + { Ok(result) => { info!("Drop database: {}, result: {result:?}", sql); break; diff --git a/tests-fuzz/targets/wal/fuzz_remote_wal_logical_prune.rs b/tests-fuzz/targets/wal/fuzz_remote_wal_logical_prune.rs index 7545f4a4a41..a79efcce7e1 100644 --- a/tests-fuzz/targets/wal/fuzz_remote_wal_logical_prune.rs +++ b/tests-fuzz/targets/wal/fuzz_remote_wal_logical_prune.rs @@ -255,7 +255,7 @@ async fn create_tables( let create_expr = generate_create_expr(input, idx, rng)?; let translator = CreateTableExprTranslator; let sql = translator.translate(&create_expr)?; - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql: &sql })?; @@ -288,7 +288,7 @@ async fn insert_table( let result = ctx .greptime // unprepared query, see - .execute(sql.as_str()) + .execute(sqlx::AssertSqlSafe(sql.as_str())) .await .context(error::ExecuteQuerySnafu { sql: &sql })?; ensure!( @@ -350,7 +350,11 @@ async fn validate_insert_exprs( primary_keys_column_list ); info!("Executing sql: {select_sql}"); - let fetched_rows = ctx.greptime.fetch_all(select_sql.as_str()).await.unwrap(); + let fetched_rows = ctx + .greptime + .fetch_all(sqlx::AssertSqlSafe(select_sql.as_str())) + .await + .unwrap(); let mut expected_rows = replace_default(&insert_expr.values_list, table_ctx, insert_expr); sort_by_primary_keys(&mut expected_rows, primary_keys_idx); validator::row::assert_eq::(&insert_expr.columns, &fetched_rows, &expected_rows)?; @@ -367,7 +371,11 @@ async fn validate_insert_exprs( column_list, table_ctx.name, ts_column.name, ts_value, primary_keys_column_list ); info!("Executing sql: {select_sql}"); - let fetched_rows = ctx.greptime.fetch_all(select_sql.as_str()).await.unwrap(); + let fetched_rows = ctx + .greptime + .fetch_all(sqlx::AssertSqlSafe(select_sql.as_str())) + .await + .unwrap(); let mut expected_rows = replace_default(&insert_expr.values_list, table_ctx, insert_expr); sort_by_primary_keys(&mut expected_rows, primary_keys_idx); validator::row::assert_eq::(&insert_expr.columns, &fetched_rows, &expected_rows)?; @@ -620,7 +628,7 @@ async fn execute_logical_prune_fuzz( for table in tables { let sql = format!("DROP TABLE {}", table.ctx.name); - let result = sqlx::query(&sql) + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) .execute(&ctx.greptime) .await .context(error::ExecuteQuerySnafu { sql })?; diff --git a/tests-integration/Cargo.toml b/tests-integration/Cargo.toml index 31a238f7bd5..4ddd9ee5a46 100644 --- a/tests-integration/Cargo.toml +++ b/tests-integration/Cargo.toml @@ -119,6 +119,7 @@ servers = { workspace = true, features = ["testing"] } similar-asserts.workspace = true sqlx = { workspace = true, features = [ "mysql", + "mysql-rsa", "postgres", "chrono", ] } diff --git a/tests-integration/tests/json2.rs b/tests-integration/tests/json2.rs index 22bcc262661..fc39a5961ff 100644 --- a/tests-integration/tests/json2.rs +++ b/tests-integration/tests/json2.rs @@ -220,7 +220,10 @@ async fn test_json2_multi_row_insert() { i, json ); - conn.execute(sql.as_str()).await.unwrap(); + // Safety: the SQL is built from test-local identifiers and payloads only. + conn.execute(sqlx::AssertSqlSafe(sql.as_str())) + .await + .unwrap(); } assert_eq!( @@ -237,9 +240,8 @@ async fn create_json2_table( table_name: &str, merge_threshold: usize, ) -> sqlx::Result<()> { - conn.execute( - format!( - r#" + let sql = format!( + r#" CREATE TABLE {table_name} ( ts TIMESTAMP TIME INDEX, host STRING PRIMARY KEY, @@ -252,10 +254,9 @@ CREATE TABLE {table_name} ( 'memtable.bulk.encode_row_threshold' = '1000000' ) "# - ) - .as_str(), - ) - .await?; + ); + // Safety: the SQL is built from test-local identifiers only. + conn.execute(sqlx::AssertSqlSafe(sql.as_str())).await?; Ok(()) } @@ -263,9 +264,8 @@ async fn create_json2_compaction_table( conn: &mut MySqlConnection, table_name: &str, ) -> sqlx::Result<()> { - conn.execute( - format!( - r#" + let sql = format!( + r#" CREATE TABLE {table_name} ( ts TIMESTAMP TIME INDEX, host STRING PRIMARY KEY, @@ -279,17 +279,15 @@ CREATE TABLE {table_name} ( 'memtable.bulk.encode_bytes_threshold' = '100000000' ) "# - ) - .as_str(), - ) - .await?; + ); + // Safety: the SQL is built from test-local identifiers only. + conn.execute(sqlx::AssertSqlSafe(sql.as_str())).await?; Ok(()) } async fn query_json2_rows(conn: &mut MySqlConnection, table_name: &str) -> sqlx::Result { - let rows = sqlx::query( - format!( - r#" + let sql = format!( + r#" SELECT host, j.payload.id::STRING AS payload_id, @@ -299,11 +297,11 @@ SELECT FROM {table_name} ORDER BY ts "# - ) - .as_str(), - ) - .fetch_all(conn) - .await?; + ); + // Safety: the SQL is built from test-local identifiers only. + let rows = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) + .fetch_all(conn) + .await?; Ok(rows .into_iter() @@ -320,9 +318,12 @@ ORDER BY ts } async fn count_table_rows(conn: &mut MySqlConnection, table_name: &str) -> sqlx::Result { - let row = sqlx::query(format!("SELECT COUNT(*) AS count FROM {table_name}").as_str()) - .fetch_one(conn) - .await?; + // Safety: the SQL is built from test-local identifiers only. + let row = sqlx::query(sqlx::AssertSqlSafe( + format!("SELECT COUNT(*) AS count FROM {table_name}").as_str(), + )) + .fetch_one(conn) + .await?; Ok(row.get("count")) } diff --git a/tests-integration/tests/sql.rs b/tests-integration/tests/sql.rs index d88fcecf1cd..4a01b0209f6 100644 --- a/tests-integration/tests/sql.rs +++ b/tests-integration/tests/sql.rs @@ -198,7 +198,8 @@ pub async fn test_mysql_auth(store_type: StorageType) { "test", "writeonly_user", |query| async move { - sqlx::query_scalar::<_, bool>(&query) + // Safety: the SQL is built from test-local identifiers only. + sqlx::query_scalar::<_, bool>(sqlx::AssertSqlSafe(query.as_str())) .fetch_one(actor_query_pool) .await .unwrap_or(false) @@ -563,7 +564,8 @@ pub async fn test_postgres_auth(store_type: StorageType) { "postgres_auth_actor", "writeonly_user", |query| async move { - sqlx::query_scalar::<_, bool>(&query) + // Safety: the SQL is built from test-local identifiers only. + sqlx::query_scalar::<_, bool>(sqlx::AssertSqlSafe(query.as_str())) .fetch_one(actor_query_pool) .await .unwrap_or(false) @@ -784,7 +786,8 @@ pub async fn test_mysql_slow_query(store_type: StorageType) { let row = tokio::time::timeout(Duration::from_secs(10), async { loop { - if let Ok(Some(row)) = sqlx::query(&query) + // Safety: the SQL is built from test-local identifiers only. + if let Ok(Some(row)) = sqlx::query(sqlx::AssertSqlSafe(query.as_str())) .bind(slow_query) .fetch_optional(&pool) .await @@ -1142,7 +1145,8 @@ pub async fn test_postgres_slow_query(store_type: StorageType) { let row = tokio::time::timeout(Duration::from_secs(10), async { loop { - if let Ok(Some(row)) = sqlx::query(&query) + // Safety: the SQL is built from test-local identifiers only. + if let Ok(Some(row)) = sqlx::query(sqlx::AssertSqlSafe(query.as_str())) .bind(slow_query) .fetch_optional(&pool) .await From 4aa4aa63feeda53029c4f59e1276b74166a9bd16 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Thu, 22 Oct 2026 10:33:18 +0000 Subject: [PATCH 072/239] fix(test): avoid MinIO connection reuse after conditional writes (#9877) Signed-off-by: WenyXu --- src/object-store/tests/object_store_test.rs | 35 ++++++++++++++++----- 1 file changed, 27 insertions(+), 8 deletions(-) diff --git a/src/object-store/tests/object_store_test.rs b/src/object-store/tests/object_store_test.rs index dc2482d1722..de526834bad 100644 --- a/src/object-store/tests/object_store_test.rs +++ b/src/object-store/tests/object_store_test.rs @@ -21,12 +21,13 @@ use common_telemetry::info; use datafusion_object_store::path::Path; use datafusion_object_store::{ObjectStore as ArrowObjectStore, ObjectStoreExt}; use futures::TryStreamExt; -use object_store::ObjectStore; use object_store::services::{Fs, S3}; use object_store::test_util::TempFolder; +use object_store::{HttpTransporter, ObjectStore, OperationContext}; use object_store_opendal::OpendalStore; use opendal::EntryMode; use opendal::services::{Azblob, Gcs, Memory, Oss}; +use opendal_http_transport_reqwest::ReqwestTransport; use prometheus::{Encoder, TextEncoder}; use tempfile::TempDir; @@ -309,7 +310,17 @@ async fn test_s3_backend() -> Result<()> { test_object_crud(&store).await?; test_object_list(&store).await?; test_object_list_start_after(&store).await?; - test_conditional_creation(&store).await?; + // MinIO can close the socket after a 412 response without Connection: close. + // A pooled read or multipart abort can race that close. Use fresh connections + // for this test rather than retrying conditional commits with ambiguous results. + let client = reqwest::Client::builder() + .pool_max_idle_per_host(0) + .build()?; + let transport = HttpTransporter::new(ReqwestTransport::new(client)); + let conditional_store = store + .clone() + .with_context(OperationContext::new().with_http_transport(transport)); + test_conditional_creation(&conditional_store).await?; assert_opendal_metrics(); guard.remove_all().await?; } @@ -345,16 +356,24 @@ async fn test_conditional_creation(store: &ObjectStore) -> Result<()> { writer.close().await } .await; - if result.is_err() { - writer.abort().await.unwrap(); - } - result + let abort_result = if result.is_err() { + writer.abort().await + } else { + Ok(()) + }; + (result, abort_result) } }; - let (a, b) = tokio::join!(attempt(payload.clone()), attempt(payload.clone())); + let ((a, a_abort), (b, b_abort)) = + tokio::join!(attempt(payload.clone()), attempt(payload.clone())); + assert!( + a_abort.is_ok() && b_abort.is_ok(), + "{path}: writes: {a:?}, {b:?}; aborts: {a_abort:?}, {b_abort:?}" + ); assert_eq!( usize::from(a.is_ok()) + usize::from(b.is_ok()), - usize::from(!preexisting) + usize::from(!preexisting), + "{path}: writes: {a:?}, {b:?}" ); for error in [a.err(), b.err()].into_iter().flatten() { assert!( From f35057d9c20f0e724df43bd1266b5911e393d350 Mon Sep 17 00:00:00 2001 From: fys <40801205+fengys1996@users.noreply.github.com> Date: Thu, 22 Oct 2026 11:42:42 +0000 Subject: [PATCH 073/239] feat(json2): support inverted index (#9324) feat: build inverted indexes for JSON2 type hints Signed-off-by: fys --- Cargo.lock | 27 +- Cargo.toml | 3 +- .../function/src/scalars/json/json_get.rs | 4 +- src/common/grpc-expr/src/alter.rs | 64 +- src/datatypes/src/json.rs | 61 +- src/index/Cargo.toml | 1 + src/index/src/target.rs | 295 ++++++- src/mito2/benches/bench_dense_index_update.rs | 7 +- src/mito2/benches/bench_index_update.rs | 7 +- src/mito2/src/cache/index/result_cache.rs | 65 +- src/mito2/src/engine/puffin_index.rs | 7 + src/mito2/src/sst/index.rs | 57 +- .../src/sst/index/bloom_filter/applier.rs | 3 +- .../src/sst/index/bloom_filter/creator.rs | 2 +- src/mito2/src/sst/index/column.rs | 111 ++- .../src/sst/index/fulltext_index/applier.rs | 12 +- .../src/sst/index/fulltext_index/creator.rs | 11 +- src/mito2/src/sst/index/indexer/finish.rs | 2 +- src/mito2/src/sst/index/intermediate.rs | 39 +- .../src/sst/index/inverted_index/applier.rs | 108 ++- .../index/inverted_index/applier/builder.rs | 473 ++++++++++-- .../inverted_index/applier/builder/between.rs | 17 +- .../applier/builder/comparison.rs | 36 +- .../inverted_index/applier/builder/eq_list.rs | 56 +- .../inverted_index/applier/builder/in_list.rs | 17 +- .../applier/builder/regex_match.rs | 25 +- .../src/sst/index/inverted_index/creator.rs | 411 +++++++--- src/operator/src/expr_helper.rs | 7 +- src/sql/src/parsers/create_parser/json.rs | 10 + src/store-api/src/region_request.rs | 42 +- .../types/json/json2_inverted_index.result | 153 ++++ .../types/json/json2_inverted_index.sql | 78 ++ .../json/json2_inverted_index_query.result | 435 +++++++++++ .../types/json/json2_inverted_index_query.sql | 195 +++++ .../json/json2_inverted_index_schema.result | 723 ++++++++++++++++++ .../json/json2_inverted_index_schema.sql | 290 +++++++ .../json/json2_inverted_index_write.result | 153 ++++ .../types/json/json2_inverted_index_write.sql | 69 ++ 38 files changed, 3700 insertions(+), 376 deletions(-) create mode 100644 tests/cases/standalone/common/types/json/json2_inverted_index.result create mode 100644 tests/cases/standalone/common/types/json/json2_inverted_index.sql create mode 100644 tests/cases/standalone/common/types/json/json2_inverted_index_query.result create mode 100644 tests/cases/standalone/common/types/json/json2_inverted_index_query.sql create mode 100644 tests/cases/standalone/common/types/json/json2_inverted_index_schema.result create mode 100644 tests/cases/standalone/common/types/json/json2_inverted_index_schema.sql create mode 100644 tests/cases/standalone/common/types/json/json2_inverted_index_write.result create mode 100644 tests/cases/standalone/common/types/json/json2_inverted_index_write.sql diff --git a/Cargo.lock b/Cargo.lock index 0f2a17d104f..b87ac7bc9f9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1657,7 +1657,7 @@ dependencies = [ "maybe-owned", "rustix 1.0.7", "rustix-linux-procfs", - "windows-sys 0.61.2", + "windows-sys 0.60.2", "winx", ] @@ -2285,7 +2285,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "117725a109d387c937a1533ce01b450cbde6b88abceea8473c4d7a85853cda3c" dependencies = [ "lazy_static", - "windows-sys 0.59.0", + "windows-sys 0.48.0", ] [[package]] @@ -6144,7 +6144,7 @@ dependencies = [ [[package]] name = "greptime-proto" version = "0.1.0" -source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=acb247415b8f4c2f0750783dab4e1f38c71f5860#acb247415b8f4c2f0750783dab4e1f38c71f5860" +source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=3b5ba5e7422e1b96500cf052bb94e978f4660c05#3b5ba5e7422e1b96500cf052bb94e978f4660c05" dependencies = [ "prost 0.14.1", "prost-types 0.14.1", @@ -6756,7 +6756,7 @@ dependencies = [ "libc", "percent-encoding", "pin-project-lite", - "socket2 0.6.4", + "socket2 0.5.10", "tokio", "tower-service", "tracing", @@ -7032,6 +7032,7 @@ dependencies = [ "common-telemetry", "common-test-util", "criterion 0.7.0", + "datatypes", "fastbloom", "fst", "futures", @@ -9317,7 +9318,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -11589,7 +11590,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac6c3320f9abac597dcbc668774ef006702672474aad53c6d596b62e487b40b1" dependencies = [ "heck 0.5.0", - "itertools 0.14.0", + "itertools 0.10.5", "log", "multimap", "once_cell", @@ -11637,7 +11638,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a56d757972c98b346a9b766e3f02746cde6dd1cd1d1d563472929fdd74bec4d" dependencies = [ "anyhow", - "itertools 0.14.0", + "itertools 0.10.5", "proc-macro2", "quote", "syn 2.0.117", @@ -11650,7 +11651,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9120690fafc389a67ba3803df527d0ec9cbbc9cc45e4cc20b332996dfb672425" dependencies = [ "anyhow", - "itertools 0.14.0", + "itertools 0.10.5", "proc-macro2", "quote", "syn 2.0.117", @@ -13180,7 +13181,7 @@ dependencies = [ "security-framework 3.7.0", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -14103,7 +14104,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "52d1cfed4120b4d927bf7c0f86d2087a4a7d6027c906d9f9d525a80573b9be51" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -14674,7 +14675,7 @@ dependencies = [ "cfg-if", "libc", "psm", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -15297,7 +15298,7 @@ dependencies = [ "getrandom 0.3.4", "once_cell", "rustix 1.0.7", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -17066,7 +17067,7 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf221c93e13a30d793f7645a0e7762c55d169dbb0a49671918a2319d289b10bb" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.48.0", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index f275cb8aada..0e75ddbd912 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -162,7 +162,7 @@ fs2 = "0.4" fst = "0.4.7" futures = "0.3" futures-util = "0.3" -greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "acb247415b8f4c2f0750783dab4e1f38c71f5860" } +greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "3b5ba5e7422e1b96500cf052bb94e978f4660c05" } hdfs-native = "0.14.6" hex = "0.4" hostname = "0.4.0" @@ -245,6 +245,7 @@ serde = { version = "1.0", features = ["derive"] } serde_json = { version = "1.0", features = ["float_roundtrip"] } serde_with = "3" serde_yaml_ng = "0.10" +sha2 = "0.10" simd-json = "0.15" similar-asserts = "1.6.0" smallvec = { version = "1", features = ["serde"] } diff --git a/src/common/function/src/scalars/json/json_get.rs b/src/common/function/src/scalars/json/json_get.rs index 9a3a1ec8616..340084245eb 100644 --- a/src/common/function/src/scalars/json/json_get.rs +++ b/src/common/function/src/scalars/json/json_get.rs @@ -42,7 +42,9 @@ pub fn parse_json_get_path(path: &str) -> std::result::Result, json parse_json_path(path.as_bytes()) } -fn json_object_path(path: &str) -> Result>> { +/// Returns object keys for an exact JSON path (e.g., `$.a.b` → `["a", "b"]`), +/// or `None` for paths containing array access or other selectors. +pub fn json_object_path(path: &str) -> Result>> { Ok(parse_json_get_path(path) .map_err(|e| exec_datafusion_err!("Invalid JSONPath {path:?}: {e}"))? .paths diff --git a/src/common/grpc-expr/src/alter.rs b/src/common/grpc-expr/src/alter.rs index ffa148c7d0b..ad3a9937887 100644 --- a/src/common/grpc-expr/src/alter.rs +++ b/src/common/grpc-expr/src/alter.rs @@ -74,9 +74,7 @@ fn json_settings_from_proto(settings: api::v1::JsonSettings) -> Result>>()?; @@ -420,6 +418,66 @@ mod tests { use super::*; + #[test] + fn test_json_settings_preserve_inverted_index() { + for enabled in [false, true] { + let settings = json_settings_from_proto(api::v1::JsonSettings { + type_hints: vec![api::v1::JsonTypeHint { + path: vec!["service".into()], + data_type: ColumnDataType::String as i32, + inverted_index: enabled, + ..Default::default() + }], + max_auto_expanded_paths: Some(0), + }) + .unwrap(); + assert_eq!(enabled, settings.type_hints()[0].inverted_index); + } + } + + #[test] + fn test_json_settings_path_validation_from_proto() { + for (path, valid_path) in [ + (vec!["nested", "leaf"], true), + (vec![""], false), + (vec!["nested", ""], false), + ( + vec!["nested", datatypes::json::JSON2_REMAINDER_FIELD_NAME], + false, + ), + ] { + for inverted_index in [false, true] { + let settings = api::v1::SetJsonSettings { + column_name: "payload".into(), + settings: Some(api::v1::JsonSettings { + type_hints: vec![api::v1::JsonTypeHint { + path: path.iter().map(|part| (*part).to_string()).collect(), + data_type: ColumnDataType::String as i32, + inverted_index, + ..Default::default() + }], + max_auto_expanded_paths: Some(0), + }), + }; + let result = alter_expr_to_request( + 1, + AlterTableExpr { + catalog_name: "greptime".into(), + schema_name: "public".into(), + table_name: "logs".into(), + kind: Some(Kind::SetJsonSettings(settings)), + }, + None, + ); + assert_eq!( + result.is_ok(), + valid_path, + "path: {path:?}, inverted_index: {inverted_index}" + ); + } + } + } + #[test] fn test_alter_expr_to_request() { let expr = AlterTableExpr { diff --git a/src/datatypes/src/json.rs b/src/datatypes/src/json.rs index e3c408698eb..78af6b240bb 100644 --- a/src/datatypes/src/json.rs +++ b/src/datatypes/src/json.rs @@ -223,17 +223,21 @@ fn validate_type_hints(type_hints: &[JsonTypeHint]) -> Result<()> { } .fail(); } - if hint - .path - .first() - .is_some_and(|x| x == JSON2_REMAINDER_FIELD_NAME) - { - return InvalidJson2SettingsSnafu { - reason: format!( - "JSON2 type hint path cannot be rooted at reserved field '{JSON2_REMAINDER_FIELD_NAME}'" - ), + for part in &hint.path { + if part.is_empty() { + return InvalidJson2SettingsSnafu { + reason: "JSON2 type hint path segments must not be empty", + } + .fail(); + } + if part == JSON2_REMAINDER_FIELD_NAME { + return InvalidJson2SettingsSnafu { + reason: format!( + "JSON2 type hint path cannot contain reserved field '{JSON2_REMAINDER_FIELD_NAME}'" + ), + } + .fail(); } - .fail(); } let data_type = match &hint.data_type { ConcreteDataType::Boolean(_) @@ -862,6 +866,41 @@ mod tests { Ok(()) } + #[test] + fn test_json_settings_path_validation() { + for (path, expected_valid) in [ + (vec![""], false), + (vec!["", "nested"], false), + (vec!["nested", ""], false), + (vec!["nested", "", "leaf"], false), + (vec!["nested", JSON2_REMAINDER_FIELD_NAME], false), + (vec!["nested", JSON2_REMAINDER_FIELD_NAME, "leaf"], false), + (vec!["nested", "leaf"], true), + (vec![], false), + (vec![JSON2_REMAINDER_FIELD_NAME], false), + (vec![JSON2_REMAINDER_FIELD_NAME, "leaf"], false), + ] { + for inverted_index in [false, true] { + let hint = JsonTypeHint { + path: path.iter().map(|part| (*part).to_string()).collect(), + data_type: ConcreteDataType::string_datatype(), + inverted_index, + }; + let serialized = json!({"type_hints": [hint.clone()]}); + assert_eq!( + JsonSettings::try_new(vec![hint], None).is_ok(), + expected_valid, + "path: {path:?}, inverted_index: {inverted_index}" + ); + assert_eq!( + serde_json::from_value::(serialized).is_ok(), + expected_valid, + "path: {path:?}, inverted_index: {inverted_index}" + ); + } + } + } + #[test] fn test_json_settings_reject_invalid_type_hints() { for type_hints in [ @@ -889,7 +928,7 @@ mod tests { )], None, ) - .is_ok() + .is_err() ); assert!( JsonSettings::try_new( diff --git a/src/index/Cargo.toml b/src/index/Cargo.toml index 4574cac18ae..9ef12e0799b 100644 --- a/src/index/Cargo.toml +++ b/src/index/Cargo.toml @@ -18,6 +18,7 @@ common-error.workspace = true common-macro.workspace = true common-runtime.workspace = true common-telemetry.workspace = true +datatypes.workspace = true fastbloom = "0.8" fst.workspace = true futures.workspace = true diff --git a/src/index/src/target.rs b/src/index/src/target.rs index cca0a281923..7a5b3811b44 100644 --- a/src/index/src/target.rs +++ b/src/index/src/target.rs @@ -13,37 +13,117 @@ // limitations under the License. use std::any::Any; -use std::fmt::{self, Display}; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; use common_macro::stack_trace_debug; +use datatypes::json::JSON2_REMAINDER_FIELD_NAME; use serde::{Deserialize, Serialize}; -use snafu::{Snafu, ensure}; +use snafu::{OptionExt, Snafu, ensure}; use store_api::storage::ColumnId; -/// Describes an index target. Column ids are the only supported variant for now. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +/// Identifies a column or a JSON object path within a column. +/// +/// # Encoding +/// +/// - Columns are encoded as decimal IDs, e.g. `1`. +/// - JSON path targets use `json-{column_id}:` followed by path segments encoded +/// as `{UTF-8 byte length}:{segment}`, e.g. `json-1:7:service4:name`. +/// IDs and lengths are unsigned decimal integers without leading zeros. +/// Segments are read by byte length, so their contents need no escaping. +/// +/// This persisted encoding is independent of the serde representation of these types. +/// +/// Note: Columns retain decimal ID encoding because existing index files use +/// these IDs as lookup keys. Changing the encoding would prevent queries from +/// finding those indexes. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] pub enum IndexTarget { ColumnId(ColumnId), + JsonPath(JsonPathTarget), } -impl Display for IndexTarget { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - IndexTarget::ColumnId(id) => write!(f, "{}", id), - } +/// An index target for a JSON object path. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] +pub struct JsonPathTarget { + column_id: ColumnId, + #[serde(deserialize_with = "deserialize_json_path")] + path: Vec, +} + +impl JsonPathTarget { + /// Returns the ID of the root column containing this JSON path. + pub fn column_id(&self) -> ColumnId { + self.column_id + } + + /// Returns the JSON object path. + pub fn path(&self) -> &[String] { + &self.path } } impl IndexTarget { - /// Parse a target key string back into an index target description. + /// Creates an index target for a JSON object path. + pub fn new_json_path(column_id: ColumnId, path: Vec) -> Result { + validate_json_path(&path)?; + let target = JsonPathTarget { column_id, path }; + Ok(Self::JsonPath(target)) + } + + /// Encodes this target using the format described in [`IndexTarget`]. + pub fn encode(&self) -> String { + let target = match self { + Self::ColumnId(id) => return id.to_string(), + Self::JsonPath(target) => target, + }; + let mut key = format!("json-{}:", target.column_id); + for part in &target.path { + key.push_str(&part.len().to_string()); + key.push(':'); + key.push_str(part); + } + key + } + + /// Decodes an index name using the format described in [`IndexTarget`]. pub fn decode(key: &str) -> Result { - validate_column_key(key)?; - let id = key - .parse::() - .map_err(|_| InvalidColumnIdSnafu { value: key }.build())?; - Ok(IndexTarget::ColumnId(id)) + ensure!(!key.is_empty(), EmptySnafu); + if key.bytes().all(|byte| byte.is_ascii_digit()) { + let id = key + .parse::() + .map_err(|_| InvalidColumnIdSnafu { value: key }.build())?; + return Ok(Self::ColumnId(id)); + } + + let invalid_key = || InvalidJsonTargetSnafu { target_key: key }; + // JSON target IDs and lengths use decimal digits without signs or leading zeros. + let is_canonical_decimal = |value: &str| { + !value.is_empty() + && value.bytes().all(|byte| byte.is_ascii_digit()) + && (value.len() == 1 || !value.starts_with('0')) + }; + + // The prefix contains the root column ID; the remainder contains path segments. + let (column_id, mut rest) = key + .strip_prefix("json-") + .and_then(|rest| rest.split_once(':')) + .with_context(invalid_key)?; + ensure!(is_canonical_decimal(column_id), invalid_key()); + let column_id = column_id.parse().ok().with_context(invalid_key)?; + + let mut path = Vec::new(); + while !rest.is_empty() { + let (length, content) = rest.split_once(':').with_context(invalid_key)?; + ensure!(is_canonical_decimal(length), invalid_key()); + let length = length.parse::().ok().with_context(invalid_key)?; + // Read exactly this many bytes, including any delimiters in the segment. + // Check bounds and the UTF-8 boundary before allocating the segment string. + let (part, remaining) = content.split_at_checked(length).with_context(invalid_key)?; + path.push(part.to_string()); + rest = remaining; + } + Self::new_json_path(column_id, path) } } @@ -59,6 +139,45 @@ pub enum TargetKeyError { #[snafu(display("failed to parse column id from '{value}'"))] InvalidColumnId { value: String }, + + #[snafu(display("invalid JSON index target key: {target_key}"))] + InvalidJsonTarget { target_key: String }, + + #[snafu(display("invalid JSON index path: {detail}"))] + InvalidJsonPath { detail: String }, +} + +fn deserialize_json_path<'de, D>(deserializer: D) -> Result, D::Error> +where + D: serde::Deserializer<'de>, +{ + let path = Vec::::deserialize(deserializer)?; + validate_json_path(&path).map_err(serde::de::Error::custom)?; + Ok(path) +} + +fn validate_json_path(path: &[String]) -> Result<(), TargetKeyError> { + ensure!( + !path.is_empty(), + InvalidJsonPathSnafu { + detail: "path must not be empty", + } + ); + for (index, part) in path.iter().enumerate() { + ensure!( + !part.is_empty(), + InvalidJsonPathSnafu { + detail: format!("path segment at index {index} must not be empty"), + } + ); + ensure!( + part != JSON2_REMAINDER_FIELD_NAME, + InvalidJsonPathSnafu { + detail: format!("path segment at index {index} uses reserved name '{part}'"), + } + ); + } + Ok(()) } impl ErrorExt for TargetKeyError { @@ -71,15 +190,6 @@ impl ErrorExt for TargetKeyError { } } -fn validate_column_key(key: &str) -> Result<(), TargetKeyError> { - ensure!(!key.is_empty(), EmptySnafu); - ensure!( - key.chars().all(|ch| ch.is_ascii_digit()), - InvalidCharactersSnafu { key } - ); - Ok(()) -} - #[cfg(test)] mod tests { use super::*; @@ -87,12 +197,26 @@ mod tests { #[test] fn encode_decode_column() { let target = IndexTarget::ColumnId(42); - let key = format!("{}", target); + let key = target.encode(); assert_eq!(key, "42"); + assert_eq!(IndexTarget::decode("42").unwrap(), target); let decoded = IndexTarget::decode(&key).unwrap(); assert_eq!(decoded, target); } + #[test] + fn encode_decode_column_id_bounds() { + for id in [0, 1, ColumnId::MAX] { + let target = IndexTarget::ColumnId(id); + let key = id.to_string(); + assert_eq!(target.encode(), key); + assert_eq!(IndexTarget::decode(&key).unwrap(), target); + } + for key in ["-1", "4294967296", "1.5", "\"1\"", "null"] { + assert!(IndexTarget::decode(key).is_err(), "{key}"); + } + } + #[test] fn decode_rejects_empty() { let err = IndexTarget::decode("").unwrap_err(); @@ -102,6 +226,125 @@ mod tests { #[test] fn decode_rejects_invalid_digits() { let err = IndexTarget::decode("1a2").unwrap_err(); - assert!(matches!(err, TargetKeyError::InvalidCharacters { .. })); + assert!(matches!(err, TargetKeyError::InvalidJsonTarget { .. })); + } + + #[test] + fn json_target_encoding_is_stable() { + let target = + IndexTarget::new_json_path(7, vec!["resource".into(), "service.name".into()]).unwrap(); + let key = "json-7:8:resource12:service.name"; + assert_eq!(target.encode(), key); + assert_eq!(IndexTarget::decode(key).unwrap(), target); + let json = r#"{"JsonPath":{"column_id":7,"path":["resource","service.name"]}}"#; + assert_eq!(serde_json::to_string(&target).unwrap(), json); + assert_eq!(serde_json::from_str::(json).unwrap(), target); + let target = IndexTarget::new_json_path(42, vec!["引号\".:[]".into()]).unwrap(); + assert_eq!(IndexTarget::decode(&target.encode()).unwrap(), target); + assert_eq!( + serde_json::to_string(&IndexTarget::ColumnId(42)).unwrap(), + r#"{"ColumnId":42}"# + ); + } + + #[test] + fn json_target_encode_decode_segments() { + let cases: &[(ColumnId, &[&str], &str)] = &[ + (42, &["service", "name"], "json-42:7:service4:name"), + (0, &["a"], "json-0:1:a"), + (ColumnId::MAX, &["a"], "json-4294967295:1:a"), + (42, &["a.b"], "json-42:3:a.b"), + (42, &["a", "b"], "json-42:1:a1:b"), + (42, &["a:b", "12:xyz"], "json-42:3:a:b6:12:xyz"), + (42, &["中文", "😀"], "json-42:6:中文4:😀"), + (42, &["\"\\\n\0"], "json-42:4:\"\\\n\0"), + (42, &["abcdefghijkl"], "json-42:12:abcdefghijkl"), + ]; + for &(column_id, path, key) in cases { + let target = IndexTarget::new_json_path( + column_id, + path.iter().map(|part| (*part).to_string()).collect(), + ) + .unwrap(); + assert_eq!(target.encode(), key); + assert_eq!(IndexTarget::decode(key).unwrap(), target); + } + } + + #[test] + fn json_target_rejects_malformed_framing() { + for key in [ + "json-", + "json-42", + "json-:1:a", + "json-+42:1:a", + "json-042:1:a", + "json-4294967296:1:a", + "json-42:", + "json-42:0:", + "json-42:1:a0:", + "json-42::a", + "json-42:+1:a", + "json-42:-1:a", + "json-42:01:a", + "json-42:1", + "json-42:2:a", + "json-42:1:中", + "json-42:3:中x", + "json-42:1:a1:", + "json-42:999999999999999999999999999999999999:a", + ] { + assert!(IndexTarget::decode(key).is_err(), "{key}"); + } + for path in [ + vec![JSON2_REMAINDER_FIELD_NAME], + vec!["nested", JSON2_REMAINDER_FIELD_NAME], + ] { + let mut key = "json-42:".to_string(); + for part in path { + key.push_str(&format!("{}:{part}", part.len())); + } + assert!(matches!( + IndexTarget::decode(&key), + Err(TargetKeyError::InvalidJsonPath { .. }) + )); + } + } + + #[test] + fn json_target_rejects_invalid_paths_and_keys() { + for path in [ + vec![], + vec![String::new()], + vec!["a".into(), String::new()], + vec![JSON2_REMAINDER_FIELD_NAME.into()], + vec!["a".into(), JSON2_REMAINDER_FIELD_NAME.into()], + vec!["a".into(), JSON2_REMAINDER_FIELD_NAME.into(), "b".into()], + ] { + assert!(matches!( + IndexTarget::new_json_path(7, path.clone()), + Err(TargetKeyError::InvalidJsonPath { .. }) + )); + let json = serde_json::json!({ + "column_id": 7, + "path": path, + }); + assert!( + serde_json::from_value::(serde_json::json!({"JsonPath": json})) + .is_err() + ); + let key = serde_json::json!({"JsonPath": {"column_id": 7, "path": path}}).to_string(); + assert!(IndexTarget::decode(&key).is_err()); + } + assert!( + serde_json::from_value::(serde_json::json!({ + "column_id": 7, + "path": [], + })) + .is_err() + ); + for key in ["j.2.7.e30", "j.1.7", "j.1.7.!", "j.1.x.e30", "j.1.7.e30"] { + assert!(IndexTarget::decode(key).is_err(), "{key}"); + } } } diff --git a/src/mito2/benches/bench_dense_index_update.rs b/src/mito2/benches/bench_dense_index_update.rs index 6f5ccfc080b..9f827c5ab80 100644 --- a/src/mito2/benches/bench_dense_index_update.rs +++ b/src/mito2/benches/bench_dense_index_update.rs @@ -38,6 +38,7 @@ use datatypes::arrow::record_batch::RecordBatch; use datatypes::data_type::ConcreteDataType; use datatypes::schema::{ColumnSchema, SkippingIndexOptions}; use datatypes::value::Value; +use index::target::IndexTarget; use mito_codec::row_converter::{DensePrimaryKeyCodec, PrimaryKeyCodecExt}; use mito2::read::{Batch, BatchBuilder}; use mito2::sst::index::Indexer; @@ -358,7 +359,11 @@ fn bench_dense_index_update(c: &mut Criterion) { intermediate.clone(), None, NonZeroUsize::new(SEGMENT_ROWS).unwrap(), - inverted_columns.clone(), + inverted_columns + .iter() + .copied() + .map(IndexTarget::ColumnId) + .collect(), ) }); let bloom = if shape.bloom { diff --git a/src/mito2/benches/bench_index_update.rs b/src/mito2/benches/bench_index_update.rs index eba89cf60ff..a76241fa53f 100644 --- a/src/mito2/benches/bench_index_update.rs +++ b/src/mito2/benches/bench_index_update.rs @@ -33,6 +33,7 @@ use datatypes::arrow::datatypes::UInt32Type; use datatypes::arrow::record_batch::RecordBatch; use datatypes::data_type::ConcreteDataType; use datatypes::schema::{ColumnSchema, SkippingIndexOptions}; +use index::target::IndexTarget; use mito_codec::row_converter::SparsePrimaryKeyCodec; use mito2::sst::index::intermediate::IntermediateManager; use mito2::sst::{FlatSchemaOptions, to_flat_sst_arrow_schema}; @@ -240,7 +241,11 @@ fn bench_index_update(c: &mut Criterion) { intermediate.clone(), None, NonZeroUsize::new(SEGMENT_ROWS).unwrap(), - inverted_columns.clone(), + inverted_columns + .iter() + .copied() + .map(IndexTarget::ColumnId) + .collect(), ) }); let mut bloom = if shape.bloom { diff --git a/src/mito2/src/cache/index/result_cache.rs b/src/mito2/src/cache/index/result_cache.rs index 5d22c739831..4448d9e0b08 100644 --- a/src/mito2/src/cache/index/result_cache.rs +++ b/src/mito2/src/cache/index/result_cache.rs @@ -17,6 +17,7 @@ use std::sync::Arc; use index::bloom_filter::applier::InListPredicate; use index::inverted_index::search::predicate::{Predicate, RangePredicate}; +use index::target::IndexTarget; use moka::notification::RemovalCause; use moka::sync::Cache; use store_api::storage::{ColumnId, FileId}; @@ -132,7 +133,7 @@ impl PredicateKey { } /// Creates a new inverted index key. - pub fn new_inverted(predicates: Arc>>) -> Self { + pub fn new_inverted(predicates: Arc>>) -> Self { Self::Inverted(InvertedIndexKey::new(predicates)) } @@ -221,27 +222,37 @@ impl BloomFilterKey { /// Key for inverted index queries. #[derive(Debug, Clone, PartialEq, Eq, Hash, Default)] pub struct InvertedIndexKey { - predicates: Arc>>, + predicates: Arc>>, mem_usage: usize, } impl InvertedIndexKey { /// Creates a new inverted index key with the given predicates. /// Calculates memory usage based on the type and size of predicates. - pub fn new(predicates: Arc>>) -> Self { + pub fn new(predicates: Arc>>) -> Self { let mem_usage = predicates - .values() - .map(|predicates| { - predicates - .iter() - .map(|predicate| match predicate { - Predicate::InList(predicate) => { - predicate.list.iter().map(|list| list.len()).sum::() - } - Predicate::Range(_) => size_of::(), - Predicate::RegexMatch(predicate) => predicate.pattern.len(), - }) - .sum::() + .iter() + .map(|(target, predicates)| { + let target_size = size_of::() + + match target { + IndexTarget::ColumnId(_) => 0, + IndexTarget::JsonPath(target) => target + .path() + .iter() + .map(|part| size_of::() + part.len()) + .sum(), + }; + target_size + + predicates + .iter() + .map(|predicate| match predicate { + Predicate::InList(predicate) => { + predicate.list.iter().map(|list| list.len()).sum::() + } + Predicate::Range(_) => size_of::(), + Predicate::RegexMatch(predicate) => predicate.pattern.len(), + }) + .sum::() }) .sum(); @@ -291,6 +302,28 @@ mod tests { }; use crate::sst::parquet::row_selection::RowGroupSelection; + #[test] + fn test_inverted_cache_separates_json_paths() { + let key = |path: &str| { + PredicateKey::new_inverted(Arc::new(BTreeMap::from([( + IndexTarget::new_json_path(1, vec![path.into()]).unwrap(), + vec![Predicate::Range(RangePredicate { + range: Range { + lower: None, + upper: None, + }, + })], + )]))) + }; + let a = key("a"); + let b = key("b"); + let cache = IndexResultCache::new(10000); + let file = FileId::random(); + cache.put(a.clone(), file, Arc::new(RowGroupSelection::new(10, 10))); + assert!(cache.get(&a, file).is_some()); + assert!(cache.get(&b, file).is_none()); + } + #[test] fn test_cache_basic_operations() { let cache = IndexResultCache::new(1000); @@ -475,7 +508,7 @@ mod tests { upper: None, }, }); - predicates3.insert(1, vec![predicate3]); + predicates3.insert(IndexTarget::ColumnId(1), vec![predicate3]); let key3 = PredicateKey::new_inverted(Arc::new(predicates3)); let selection3 = Arc::new(RowGroupSelection::from_row_ranges( vec![(0, vec![5..15])], diff --git a/src/mito2/src/engine/puffin_index.rs b/src/mito2/src/engine/puffin_index.rs index 48833f237d1..cef4f3a449b 100644 --- a/src/mito2/src/engine/puffin_index.rs +++ b/src/mito2/src/engine/puffin_index.rs @@ -50,6 +50,8 @@ const TARGET_TYPE_UNKNOWN: &str = "unknown"; const TARGET_TYPE_COLUMN: &str = "column"; +const TARGET_TYPE_JSON_PATH: &str = "json_path"; + pub(crate) struct IndexEntryContext<'a> { pub(crate) table_dir: &'a str, pub(crate) index_file_path: &'a str, @@ -370,6 +372,11 @@ fn decode_target_info(target_key: &str) -> (String, String) { TARGET_TYPE_COLUMN.to_string(), json!({ "column": id }).to_string(), ), + Ok(IndexTarget::JsonPath(target)) => ( + TARGET_TYPE_JSON_PATH.to_string(), + json!({ "json_path": { "column": target.column_id(), "path": target.path() } }) + .to_string(), + ), _ => ( TARGET_TYPE_UNKNOWN.to_string(), json!({ "error": "failed_to_decode" }).to_string(), diff --git a/src/mito2/src/sst/index.rs b/src/mito2/src/sst/index.rs index ff734e3241b..4dc7002f147 100644 --- a/src/mito2/src/sst/index.rs +++ b/src/mito2/src/sst/index.rs @@ -33,9 +33,12 @@ use std::num::NonZeroUsize; use std::ops::ControlFlow; use std::sync::Arc; +use arrow_schema::extension::ExtensionType; use bloom_filter::creator::BloomFilterIndexer; use common_telemetry::{debug, error, info, warn}; use datatypes::arrow::record_batch::RecordBatch; +use datatypes::extension::json::Json2ExtensionType; +use index::target::IndexTarget; use mito_codec::row_converter::DensePrimaryKeyCodec; use object_store::ObjectStore; use puffin_manager::SstPuffinManager; @@ -302,7 +305,7 @@ impl Indexer { let requested: HashSet<_> = self .inverted_indexer .iter() - .flat_map(|indexer| indexer.column_ids()) + .flat_map(|indexer| indexer.indexed_column_ids()) .chain( self.bloom_filter_indexer .iter() @@ -462,7 +465,13 @@ impl IndexerBuilderImpl { let indexed_column_ids = self.metadata.inverted_indexed_column_ids( self.index_options.inverted_index.ignore_column_ids.iter(), ); - if indexed_column_ids.is_empty() { + let mut index_targets: Vec<_> = indexed_column_ids + .into_iter() + .map(IndexTarget::ColumnId) + .collect(); + index_targets.extend(self.collect_json_inverted_index_targets()); + + if index_targets.is_empty() { debug!( "No columns to be indexed, skip creating inverted index, region_id: {}, file_id: {}", self.metadata.region_id, file_id, @@ -493,12 +502,54 @@ impl IndexerBuilderImpl { self.intermediate_manager.clone(), self.inverted_index_config.mem_threshold_on_create(), segment_row_count, - indexed_column_ids, + index_targets, ); Some(indexer) } + /// Collects inverted index targets for JSON type hints with indexing enabled. + /// + /// Only flush and compaction builds collect these targets; manual and schema-change + /// rebuilds return no JSON targets. + fn collect_json_inverted_index_targets(&self) -> Vec { + if !matches!( + self.build_type, + IndexBuildType::Flush | IndexBuildType::Compact + ) { + return Vec::new(); + } + + let mut index_targets = Vec::new(); + for col in &self.metadata.column_metadatas { + if !col.column_schema.data_type.is_json2() { + continue; + } + let Ok(Some(extension)) = col.column_schema.extension_type::() + else { + continue; + }; + for hint in extension.metadata().json_settings().type_hints() { + if !hint.inverted_index { + continue; + } + // JSON hints may also describe containers through non-SQL ingestion. + if !(hint.data_type.is_numeric() + || hint.data_type.is_boolean() + || hint.data_type.is_string()) + { + continue; + } + match IndexTarget::new_json_path(col.column_id, hint.path.clone()) { + Ok(target) => index_targets.push(target), + Err(err) => warn!(err; "Skipping invalid JSON inverted index target"), + } + } + } + + index_targets + } + async fn build_fulltext_indexer(&self, file_id: FileId) -> Option { let create = match self.build_type { IndexBuildType::Flush => self.fulltext_index_config.create_on_flush.auto(), diff --git a/src/mito2/src/sst/index/bloom_filter/applier.rs b/src/mito2/src/sst/index/bloom_filter/applier.rs index 1faae466c89..decf50470b7 100644 --- a/src/mito2/src/sst/index/bloom_filter/applier.rs +++ b/src/mito2/src/sst/index/bloom_filter/applier.rs @@ -376,7 +376,8 @@ impl BloomFilterIndexApplier { // TODO(ruihang): use the same util with the code in creator fn column_blob_name(column_id: ColumnId) -> String { - format!("{INDEX_BLOB_TYPE}-{}", IndexTarget::ColumnId(column_id)) + let target_key = IndexTarget::ColumnId(column_id).encode(); + format!("{INDEX_BLOB_TYPE}-{target_key}") } /// Creates a blob reader from the remote index file diff --git a/src/mito2/src/sst/index/bloom_filter/creator.rs b/src/mito2/src/sst/index/bloom_filter/creator.rs index 2776ddacc0d..ea727f1b7b4 100644 --- a/src/mito2/src/sst/index/bloom_filter/creator.rs +++ b/src/mito2/src/sst/index/bloom_filter/creator.rs @@ -408,7 +408,7 @@ impl BloomFilterIndexer { ) -> Result { let (tx, rx) = tokio::io::duplex(PIPE_BUFFER_SIZE_FOR_SENDING_BLOB); - let target_key = IndexTarget::ColumnId(*col_id); + let target_key = IndexTarget::ColumnId(*col_id).encode(); let blob_name = format!("{INDEX_BLOB_TYPE}-{target_key}"); let (index_finish, puffin_add_blob) = futures::join!( creator.finish(tx.compat_write()), diff --git a/src/mito2/src/sst/index/column.rs b/src/mito2/src/sst/index/column.rs index e2f9aaf60d1..2818efbeeda 100644 --- a/src/mito2/src/sst/index/column.rs +++ b/src/mito2/src/sst/index/column.rs @@ -13,12 +13,45 @@ // limitations under the License. use api::v1::SemanticType; -use datatypes::arrow::array::Array; +use datatypes::arrow::array::{Array, ArrayRef, StructArray}; use datatypes::arrow::record_batch::RecordBatch; +use datatypes::vectors::{Helper, VectorRef}; +use snafu::{OptionExt, ResultExt}; use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME; +use crate::error::{ConvertVectorSnafu, InvalidRecordBatchSnafu, Result}; use crate::sst::parquet::format::PrimaryKeyArray; +/// Extracts leaf values for a JSON type hint at the given path. +/// +/// Returns an error if the path does not exist, traverses a non-object, +/// or the leaf array cannot be converted to a vector. +/// +/// # Details +/// +/// The caller guarantees that the materialized leaf type matches the type hint. +/// +/// JSON2 builders propagate ancestor nulls to child arrays. SST reads and schema +/// adaptation preserve this invariant, so no ancestor null bitmap merging is needed. +pub(crate) fn extract_json_type_hint_values(root: &ArrayRef, path: &[String]) -> Result { + let mut array = root.clone(); + for segment in path { + let object = array + .as_any() + .downcast_ref::() + .with_context(|| InvalidRecordBatchSnafu { + reason: format!("JSON index path {path:?} traverses a non-object"), + })?; + array = object + .column_by_name(segment) + .with_context(|| InvalidRecordBatchSnafu { + reason: format!("JSON index path {path:?} has no materialized field {segment}"), + })? + .clone(); + } + Helper::try_into_vector(array).context(ConvertVectorSnafu) +} + /// Yields (first row, row count) for indexing a materialized column. /// Tags are constant within consecutive equal PK dictionary keys. Fields and /// timestamps must still be visited row by row. Distinct dictionary entries may @@ -49,3 +82,79 @@ pub(crate) fn column_index_rows( Some((start, row - start)) }) } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use datatypes::arrow::array::{Int64Array, NullArray}; + use datatypes::arrow::buffer::NullBuffer; + use datatypes::arrow::datatypes::Field; + use datatypes::value::Value; + + use super::*; + + #[test] + fn test_extract_json_type_hint_values_null_type() { + let nulls: ArrayRef = Arc::new(NullArray::new(2)); + let root: ArrayRef = Arc::new(StructArray::new( + vec![Field::new("a", nulls.data_type().clone(), true)].into(), + vec![nulls], + None, + )); + // A null leaf can be extracted, but cannot be traversed as an object. + assert!(extract_json_type_hint_values(&root, &["a".into(), "b".into()]).is_err()); + let leaf = extract_json_type_hint_values(&root, &["a".into()]).unwrap(); + assert_eq!(leaf.len(), 2); + assert_eq!(leaf.get(0), Value::Null); + assert_eq!(leaf.get(1), Value::Null); + + let values: ArrayRef = Arc::new(Int64Array::from(vec![None, None])); + let root: ArrayRef = Arc::new(StructArray::new( + vec![Field::new("a", values.data_type().clone(), true)].into(), + vec![values], + None, + )); + let leaf = extract_json_type_hint_values(&root, &["a".into()]).unwrap(); + assert_eq!(leaf.len(), 2); + assert_eq!(leaf.get(0), Value::Null); + assert_eq!(leaf.get(1), Value::Null); + } + + #[test] + fn test_extract_json_type_hint_values_nulls_and_slices() { + let child = Arc::new(Int64Array::from(vec![Some(11), None, None, None, Some(55)])); + let nested = StructArray::new( + vec![Field::new("a.b", child.data_type().clone(), true)].into(), + vec![child], + Some(NullBuffer::from(vec![true, false, false, true, true])), + ); + let root: ArrayRef = Arc::new(StructArray::new( + vec![Field::new("nested", nested.data_type().clone(), true)].into(), + vec![Arc::new(nested)], + Some(NullBuffer::from(vec![true, true, false, true, true])), + )); + let path = vec!["nested".into(), "a.b".into()]; + let leaf = extract_json_type_hint_values(&root, &path).unwrap(); + assert_eq!( + (0..leaf.len()).map(|i| leaf.get(i)).collect::>(), + vec![ + Value::Int64(11), + Value::Null, + Value::Null, + Value::Null, + Value::Int64(55) + ] + ); + let slice = extract_json_type_hint_values(&root.slice(1, 4), &path).unwrap(); + assert_eq!( + (0..slice.len()).map(|i| slice.get(i)).collect::>(), + vec![Value::Null, Value::Null, Value::Null, Value::Int64(55)] + ); + assert!(extract_json_type_hint_values(&root, &["missing".into()]).is_err()); + assert!( + extract_json_type_hint_values(&root, &["nested".into(), "a.b".into(), "child".into()]) + .is_err() + ); + } +} diff --git a/src/mito2/src/sst/index/fulltext_index/applier.rs b/src/mito2/src/sst/index/fulltext_index/applier.rs index 55df0dd1a06..4a18835dab4 100644 --- a/src/mito2/src/sst/index/fulltext_index/applier.rs +++ b/src/mito2/src/sst/index/fulltext_index/applier.rs @@ -287,10 +287,8 @@ impl FulltextIndexApplier { request: &FulltextRequest, metrics: Option<&mut FulltextIndexApplyMetrics>, ) -> Result>> { - let blob_key = format!( - "{INDEX_BLOB_TYPE_TANTIVY}-{}", - IndexTarget::ColumnId(column_id) - ); + let target_key = IndexTarget::ColumnId(column_id).encode(); + let blob_key = format!("{INDEX_BLOB_TYPE_TANTIVY}-{target_key}"); let dir = self .index_source .dir(file_id, &blob_key, file_size_hint, metrics) @@ -407,10 +405,8 @@ impl FulltextIndexApplier { output: &mut [(usize, Vec>)], mut metrics: Option<&mut FulltextIndexApplyMetrics>, ) -> Result { - let blob_key = format!( - "{INDEX_BLOB_TYPE_BLOOM}-{}", - IndexTarget::ColumnId(column_id) - ); + let target_key = IndexTarget::ColumnId(column_id).encode(); + let blob_key = format!("{INDEX_BLOB_TYPE_BLOOM}-{target_key}"); let Some(reader) = self .index_source .blob(file_id, &blob_key, file_size_hint, metrics.as_deref_mut()) diff --git a/src/mito2/src/sst/index/fulltext_index/creator.rs b/src/mito2/src/sst/index/fulltext_index/creator.rs index 7348d81b7f3..df2b314b5d1 100644 --- a/src/mito2/src/sst/index/fulltext_index/creator.rs +++ b/src/mito2/src/sst/index/fulltext_index/creator.rs @@ -426,22 +426,17 @@ impl AltFulltextCreator { column_id: &ColumnId, put_options: PutOptions, ) -> Result { + let target_key = IndexTarget::ColumnId(*column_id).encode(); match self { Self::Tantivy(creator) => { - let blob_key = format!( - "{INDEX_BLOB_TYPE_TANTIVY}-{}", - IndexTarget::ColumnId(*column_id) - ); + let blob_key = format!("{INDEX_BLOB_TYPE_TANTIVY}-{target_key}"); creator .finish(puffin_writer, &blob_key, put_options) .await .context(FulltextFinishSnafu) } Self::Bloom(creator) => { - let blob_key = format!( - "{INDEX_BLOB_TYPE_BLOOM}-{}", - IndexTarget::ColumnId(*column_id) - ); + let blob_key = format!("{INDEX_BLOB_TYPE_BLOOM}-{target_key}"); creator .finish(puffin_writer, &blob_key, put_options) .await diff --git a/src/mito2/src/sst/index/indexer/finish.rs b/src/mito2/src/sst/index/indexer/finish.rs index 0c1f200d599..4e46f886ebd 100644 --- a/src/mito2/src/sst/index/indexer/finish.rs +++ b/src/mito2/src/sst/index/indexer/finish.rs @@ -121,7 +121,7 @@ impl Indexer { return true; }; - let column_ids = indexer.column_ids().collect(); + let column_ids = indexer.indexed_column_ids().collect(); let err = match indexer.finish(puffin_writer).await { Ok((row_count, byte_count)) => { self.fill_inverted_index_output( diff --git a/src/mito2/src/sst/index/intermediate.rs b/src/mito2/src/sst/index/intermediate.rs index 24103179c41..e5124a60935 100644 --- a/src/mito2/src/sst/index/intermediate.rs +++ b/src/mito2/src/sst/index/intermediate.rs @@ -12,10 +12,12 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::hash::BuildHasher; use std::io::ErrorKind; use std::path::PathBuf; use async_trait::async_trait; +use common_base::hash::FixedRandomState; use common_error::ext::BoxedError; use common_telemetry::{debug, warn}; use futures::{AsyncRead, AsyncWrite}; @@ -153,14 +155,19 @@ impl IntermediateLocation { &self.files_dir } - /// Returns the path of the directory for intermediate files associated with the `file_group`: - /// `__intm/{region_id}/{sst_file_id}/{uuid}/{file_group}/` + /// Returns the directory path for intermediate files associated with `file_group`: + /// `__intm/{region_id}/{sst_file_id}/{uuid}/{hash(file_group)}/`. + /// + /// Index names may contain JSON paths with filesystem-unsafe characters or exceed + /// directory name length limits. Fixed-seed aHash maps them to 16-character hexadecimal + /// directory names. pub fn file_group_path(&self, file_group: &str) -> String { - util::join_path(&self.files_dir, &format!("{file_group}/")) + let hash = FixedRandomState::new().hash_one(file_group); + util::join_path(&self.files_dir, &format!("{hash:016x}/")) } /// Returns the path of the intermediate file with the given `file_group` and `im_file_id`: - /// `__intm/{region_id}/{sst_file_id}/{uuid}/{file_group}/{im_file_id}.im` + /// `__intm/{region_id}/{sst_file_id}/{uuid}/{hash(file_group)}/{im_file_id}.im` pub fn file_path(&self, file_group: &str, im_file_id: &str) -> String { util::join_path( &self.file_group_path(file_group), @@ -370,16 +377,21 @@ mod tests { let uuid = location.files_dir.split('/').nth(3).unwrap(); let file_group = "1"; + let group_path = location.file_group_path(file_group); + assert_eq!(group_path, location.file_group_path(file_group)); + let group_dir = group_path.trim_end_matches('/').rsplit('/').next().unwrap(); + assert_eq!(group_dir.len(), 16); + assert!(group_dir.bytes().all(|byte| byte.is_ascii_hexdigit())); assert_eq!( location.file_group_path(file_group), - format!("{INTERMEDIATE_DIR}/0/{sst_file_id}/{uuid}/{file_group}/") + format!("{INTERMEDIATE_DIR}/0/{sst_file_id}/{uuid}/{group_dir}/") ); let im_file_id = "000000000010"; let file_path = location.file_path(file_group, im_file_id); assert_eq!( file_path, - format!("{INTERMEDIATE_DIR}/0/{sst_file_id}/{uuid}/{file_group}/{im_file_id}.im") + format!("{INTERMEDIATE_DIR}/0/{sst_file_id}/{uuid}/{group_dir}/{im_file_id}.im") ); assert_eq!(location.im_file_id_from_path(&file_path), im_file_id); @@ -421,7 +433,14 @@ mod tests { let store = IntermediateManager::init_fs(path).await.unwrap(); let provider = TempFileProvider::new(location.clone(), store); - let file_group = "tag0"; + let group0 = index::target::IndexTarget::new_json_path(1, vec!["a/\":\\".repeat(100)]) + .unwrap() + .encode(); + let group1 = index::target::IndexTarget::new_json_path(1, vec!["b/\":\\".repeat(100)]) + .unwrap() + .encode(); + assert!(group0.len() > 255 && group1.len() > 255); + let file_group = group0.as_str(); let file_id = "0000000010"; let mut writer = provider.create(file_group, file_id).await.unwrap(); writer.write_all(b"hello").await.unwrap(); @@ -434,21 +453,21 @@ mod tests { writer.flush().await.unwrap(); writer.close().await.unwrap(); - let file_group = "tag1"; + let file_group = group1.as_str(); let file_id = "0000000010"; let mut writer = provider.create(file_group, file_id).await.unwrap(); writer.write_all(b"foo").await.unwrap(); writer.flush().await.unwrap(); writer.close().await.unwrap(); - let readers = provider.read_all("tag0").await.unwrap(); + let readers = provider.read_all(&group0).await.unwrap(); assert_eq!(readers.len(), 2); for (_, mut reader) in readers { let mut buf = Vec::new(); reader.read_to_end(&mut buf).await.unwrap(); assert!(matches!(buf.as_slice(), b"hello" | b"world")); } - let readers = provider.read_all("tag1").await.unwrap(); + let readers = provider.read_all(&group1).await.unwrap(); assert_eq!(readers.len(), 1); let mut reader = readers.into_iter().map(|x| x.1).next().unwrap(); let mut buf = Vec::new(); diff --git a/src/mito2/src/sst/index/inverted_index/applier.rs b/src/mito2/src/sst/index/inverted_index/applier.rs index 8f0c309a003..15b97b68348 100644 --- a/src/mito2/src/sst/index/inverted_index/applier.rs +++ b/src/mito2/src/sst/index/inverted_index/applier.rs @@ -18,9 +18,11 @@ use std::collections::BTreeMap; use std::sync::Arc; use std::time::Instant; +use arrow_schema::extension::ExtensionType; use common_base::range_read::RangeReader; use common_telemetry::warn; use datatypes::data_type::ConcreteDataType; +use datatypes::extension::json::Json2ExtensionType; use index::inverted_index::format::reader::{InvertedIndexBlobReader, InvertedIndexReadMetrics}; use index::inverted_index::search::index_apply::{ ApplyOutput, IndexApplier, IndexNotFoundStrategy, PredicatesIndexApplier, SearchContext, @@ -31,9 +33,8 @@ use object_store::ObjectStore; use puffin::puffin_manager::cache::PuffinMetadataCacheRef; use puffin::puffin_manager::{PuffinManager, PuffinReader}; use snafu::ResultExt; -use store_api::metadata::RegionMetadataRef; +use store_api::metadata::{RegionMetadata, RegionMetadataRef}; use store_api::region_request::PathType; -use store_api::storage::ColumnId; use crate::access_layer::{RegionFilePathFactory, WriteCachePathProvider}; use crate::cache::file_cache::{FileCacheRef, FileType, IndexKey}; @@ -135,13 +136,13 @@ pub(crate) struct InvertedIndexApplier { puffin_metadata_cache: Option, /// All collected predicates. - predicates: BTreeMap>, + predicates: BTreeMap>, /// Default apply plan built from all collected predicates. default_plan: SstApplyPlan, - /// Expected predicate column types from the latest region metadata. - expected_predicate_col_types: BTreeMap, + /// Expected predicate target types from the latest region metadata. + expected_predicate_target_types: BTreeMap, } pub(crate) type InvertedIndexApplierRef = Arc; @@ -159,8 +160,8 @@ impl InvertedIndexApplier { path_type: PathType, store: ObjectStore, puffin_manager_factory: PuffinManagerFactory, - predicates: BTreeMap>, - expected_predicate_col_types: BTreeMap, + predicates: BTreeMap>, + expected_predicate_target_types: BTreeMap, ) -> Result { let default_plan = Self::build_apply_plan(&predicates)?; INDEX_APPLY_MEMORY_USAGE.add(default_plan.index_applier.memory_usage() as i64); @@ -175,7 +176,7 @@ impl InvertedIndexApplier { puffin_metadata_cache: None, predicates, default_plan, - expected_predicate_col_types, + expected_predicate_target_types, }) } @@ -368,17 +369,24 @@ impl InvertedIndexApplier { let mut compatible_predicates = BTreeMap::new(); let mut has_type_mismatch = false; - for (col_id, expected) in &self.expected_predicate_col_types { - if let Some(sst_col) = sst_metadata.column_by_id(*col_id) - && sst_col.column_schema.data_type != *expected - { + for (target, predicates) in &self.predicates { + let Some(expected) = self.expected_predicate_target_types.get(target) else { + has_type_mismatch = true; + continue; + }; + let actual = index_target_data_type(sst_metadata, target); + // JSON paths require a matching type hint in this SST. Ordinary missing + // columns retain the missing-index fallback for added columns with defaults. + let compatible = match target { + IndexTarget::ColumnId(_) => actual.as_ref().is_none_or(|ty| ty == expected), + IndexTarget::JsonPath(_) => actual.as_ref() == Some(expected), + }; + if !compatible { has_type_mismatch = true; continue; } - if let Some(predicates) = self.predicates.get(col_id) { - compatible_predicates.insert(*col_id, predicates.clone()); - } + compatible_predicates.insert(target.clone(), predicates.clone()); } if compatible_predicates.is_empty() { @@ -394,17 +402,17 @@ impl InvertedIndexApplier { } fn build_apply_plan( - predicates_by_col: &BTreeMap>, + predicates_by_target: &BTreeMap>, ) -> Result { - let predicates = predicates_by_col - .iter() - .map(|(col_id, preds)| (format!("{}", IndexTarget::ColumnId(*col_id)), preds.clone())) - .collect(); - + let mut predicates = Vec::with_capacity(predicates_by_target.len()); + for (target, preds) in predicates_by_target { + let name = target.encode(); + predicates.push((name, preds.clone())); + } let index_applier = PredicatesIndexApplier::try_from(predicates).context(BuildIndexApplierSnafu)?; - let predicate_key = PredicateKey::new_inverted(Arc::new(predicates_by_col.clone())); + let predicate_key = PredicateKey::new_inverted(Arc::new(predicates_by_target.clone())); Ok(SstApplyPlan { predicate_key, index_applier: Arc::new(index_applier), @@ -418,6 +426,34 @@ impl Drop for InvertedIndexApplier { } } +/// Resolves the value type used to encode an index predicate. +/// JSON paths must have an enabled inverted-index type hint. +fn index_target_data_type( + metadata: &RegionMetadata, + target: &IndexTarget, +) -> Option { + let target = match target { + IndexTarget::ColumnId(id) => { + return metadata + .column_by_id(*id) + .map(|col| col.column_schema.data_type.clone()); + } + IndexTarget::JsonPath(target) => target, + }; + let col = metadata.column_by_id(target.column_id())?; + let extension = col + .column_schema + .extension_type::() + .ok()??; + extension + .metadata() + .json_settings() + .type_hints() + .iter() + .find(|hint| hint.path == target.path() && hint.inverted_index) + .map(|hint| hint.data_type.clone()) +} + #[cfg(test)] mod tests { use api::v1::SemanticType; @@ -442,13 +478,15 @@ mod tests { let mut predicates = BTreeMap::new(); predicates.insert( - 1, + IndexTarget::ColumnId(1), vec![Predicate::RegexMatch(RegexMatchPredicate { pattern: "foo".to_string(), })], ); - let expected_predicate_col_types = - BTreeMap::from_iter([(1, ConcreteDataType::string_datatype())]); + let expected_predicate_target_types = BTreeMap::from_iter([( + IndexTarget::ColumnId(1), + ConcreteDataType::string_datatype(), + )]); let sst_index_applier = InvertedIndexApplier::new( table_dir, @@ -456,7 +494,7 @@ mod tests { object_store, puffin_manager_factory, predicates, - expected_predicate_col_types, + expected_predicate_target_types, ) .unwrap(); let plan = sst_index_applier @@ -474,14 +512,14 @@ mod tests { let mut predicates = BTreeMap::new(); predicates.insert( - 1, + IndexTarget::ColumnId(1), vec![Predicate::RegexMatch(RegexMatchPredicate { pattern: "foo".to_string(), })], ); // Column id 1 is String in `mock_region_metadata`, set expected type to Int64. - let expected_predicate_col_types = - BTreeMap::from_iter([(1, ConcreteDataType::int64_datatype())]); + let expected_predicate_target_types = + BTreeMap::from_iter([(IndexTarget::ColumnId(1), ConcreteDataType::int64_datatype())]); let sst_index_applier = InvertedIndexApplier::new( table_dir, @@ -489,7 +527,7 @@ mod tests { object_store, puffin_manager_factory, predicates, - expected_predicate_col_types, + expected_predicate_target_types, ) .unwrap(); let plan = sst_index_applier @@ -526,20 +564,22 @@ mod tests { let mut predicates = BTreeMap::new(); predicates.insert( - 1, + IndexTarget::ColumnId(1), vec![Predicate::RegexMatch(RegexMatchPredicate { pattern: "foo".to_string(), })], ); - let expected_predicate_col_types = - BTreeMap::from_iter([(1, ConcreteDataType::string_datatype())]); + let expected_predicate_target_types = BTreeMap::from_iter([( + IndexTarget::ColumnId(1), + ConcreteDataType::string_datatype(), + )]); let sst_index_applier = InvertedIndexApplier::new( table_dir.clone(), PathType::Bare, object_store, puffin_manager_factory, predicates, - expected_predicate_col_types, + expected_predicate_target_types, ) .unwrap(); let plan = sst_index_applier diff --git a/src/mito2/src/sst/index/inverted_index/applier/builder.rs b/src/mito2/src/sst/index/inverted_index/applier/builder.rs index d2d4c544bca..62d6f69e5b4 100644 --- a/src/mito2/src/sst/index/inverted_index/applier/builder.rs +++ b/src/mito2/src/sst/index/inverted_index/applier/builder.rs @@ -20,12 +20,16 @@ mod regex_match; use std::collections::{BTreeMap, HashSet}; +use arrow_schema::extension::ExtensionType; +use common_function::scalars::json::json_get::{JsonGetWithType, json_object_path}; use common_telemetry::warn; use datafusion_common::ScalarValue; use datafusion_expr::{BinaryExpr, Expr, Operator}; -use datatypes::data_type::ConcreteDataType; +use datatypes::data_type::{ConcreteDataType, DataType}; +use datatypes::extension::json::Json2ExtensionType; use datatypes::value::Value; use index::inverted_index::search::predicate::Predicate; +use index::target::IndexTarget; use mito_codec::index::IndexValueCodec; use mito_codec::row_converter::SortField; use object_store::ObjectStore; @@ -38,7 +42,7 @@ use store_api::storage::ColumnId; use crate::cache::file_cache::FileCacheRef; use crate::cache::index::inverted_index::InvertedIndexCacheRef; use crate::error::{ColumnNotFoundSnafu, ConvertValueSnafu, EncodeSnafu, Result}; -use crate::sst::index::inverted_index::applier::InvertedIndexApplier; +use crate::sst::index::inverted_index::applier::{InvertedIndexApplier, index_target_data_type}; use crate::sst::index::puffin_manager::PuffinManagerFactory; /// Constructs an [`InvertedIndexApplier`] which applies predicates to SST files during scan. @@ -58,11 +62,11 @@ pub(crate) struct InvertedIndexApplierBuilder<'a> { /// Metadata of the region, used to get metadata like column type. metadata: &'a RegionMetadata, - /// Column ids of the columns that are indexed. + /// Column ids with column-level inverted indexes, excluding JSON path indexes. indexed_column_ids: HashSet, /// Stores predicates during traversal on the Expr tree. - output: BTreeMap>, + output: BTreeMap>, /// The puffin manager factory. puffin_manager_factory: PuffinManagerFactory, @@ -133,7 +137,7 @@ impl<'a> InvertedIndexApplierBuilder<'a> { return Ok(None); } - let expected_predicate_column_types = self.expected_predicate_column_types(); + let expected_predicate_target_types = self.expected_predicate_target_types(); Ok(Some( InvertedIndexApplier::new( @@ -142,7 +146,7 @@ impl<'a> InvertedIndexApplierBuilder<'a> { self.object_store, self.puffin_manager_factory, self.output, - expected_predicate_column_types, + expected_predicate_target_types, )? .with_file_cache(self.file_cache) .with_puffin_metadata_cache(self.puffin_metadata_cache) @@ -150,17 +154,17 @@ impl<'a> InvertedIndexApplierBuilder<'a> { )) } - /// Returns `(column_id, data_type)` pairs for predicate columns + /// Returns `(target, data_type)` pairs for predicate targets /// collected in `self.output`. /// /// The data types are resolved from the latest region manifest. Columns /// that no longer exist in the latest metadata are skipped. - fn expected_predicate_column_types(&self) -> BTreeMap { + fn expected_predicate_target_types(&self) -> BTreeMap { self.output .keys() - .filter_map(|col_id| { - let col = self.metadata.column_by_id(*col_id)?; - Some((*col_id, col.column_schema.data_type.clone())) + .filter_map(|target| { + let data_type = index_target_data_type(self.metadata, target)?; + Some((target.clone(), data_type)) }) .collect() } @@ -197,8 +201,8 @@ impl<'a> InvertedIndexApplierBuilder<'a> { } /// Helper function to add a predicate to the output. - fn add_predicate(&mut self, column_id: ColumnId, predicate: Predicate) { - self.output.entry(column_id).or_default().push(predicate); + fn add_predicate(&mut self, target: IndexTarget, predicate: Predicate) { + self.output.entry(target).or_default().push(predicate); } /// Helper function to get the column id and the column type of a column. @@ -224,6 +228,93 @@ impl<'a> InvertedIndexApplierBuilder<'a> { ))) } + /// Resolves an expression to an inverted-index target and its value type. + /// + /// Plain column references resolve only if the column is indexed. + /// JSON accesses require a matching path type hint with inverted indexing enabled. + /// Two-argument `json_get` uses the hint's type. With three arguments, the + /// third argument must be a typed NULL matching the hint's type; different + /// string representations are considered compatible. + fn expr_to_index_target(&self, expr: &Expr) -> Result> { + if let Some(name) = Self::column_name(expr) { + return Ok(self + .column_id_and_type(name)? + .map(|(id, ty)| (IndexTarget::ColumnId(id), ty))); + } + + let Expr::ScalarFunction(func) = expr else { + return Ok(None); + }; + + if func.func.name() != JsonGetWithType::NAME || !(2..=3).contains(&func.args.len()) { + return Ok(None); + } + + let Some(name) = Self::column_name(&func.args[0]) else { + return Ok(None); + }; + + let Some(col) = self.metadata.column_by_name(name) else { + return Ok(None); + }; + + if !col.column_schema.data_type.is_json2() { + return Ok(None); + } + + let Some(path_literal) = func.args[1] + .as_literal() + .and_then(|v| v.try_as_str()) + .flatten() + else { + return Ok(None); + }; + + let Some(path) = json_object_path(path_literal).ok().flatten() else { + return Ok(None); + }; + + let Ok(Some(extension)) = col.column_schema.extension_type::() else { + return Ok(None); + }; + + let Some(hint) = extension + .metadata() + .json_settings() + .type_hints() + .iter() + .find(|hint| hint.path == path && hint.inverted_index) + else { + return Ok(None); + }; + + if let Some(type_arg) = func.args.get(2) { + let Expr::Literal(value, _) = type_arg else { + return Ok(None); + }; + if !value.is_null() { + return Ok(None); + } + // A type conversion may change comparison semantics and invalidate index pruning. + let result_type = value.data_type(); + if result_type != hint.data_type.as_arrow_type() + && !(result_type.is_string() && hint.data_type.is_string()) + { + return Ok(None); + } + } + + if !(hint.data_type.is_numeric() + || hint.data_type.is_boolean() + || hint.data_type.is_string()) + { + return Ok(None); + } + Ok(IndexTarget::new_json_path(col.column_id, path) + .ok() + .map(|target| (target, hint.data_type.clone()))) + } + /// Helper function to get a non-null literal. fn nonnull_lit(expr: &Expr) -> Option<&ScalarValue> { match expr { @@ -253,6 +344,8 @@ impl<'a> InvertedIndexApplierBuilder<'a> { #[cfg(test)] mod tests { + use std::sync::Arc; + use api::v1::SemanticType; use datafusion_common::Column; use datafusion_expr::{Between, Literal}; @@ -268,6 +361,250 @@ mod tests { use super::*; + #[test] + fn test_json_target_resolution() { + use datafusion_expr::col; + let (_dir, factory) = PuffinManagerFactory::new_for_test_block("json_target_resolution"); + let metadata = json_metadata(); + let builder = InvertedIndexApplierBuilder::new( + "test".into(), + PathType::Bare, + test_object_store(), + &metadata, + HashSet::new(), + factory, + ); + for (path, expected) in [ + ("$.a", vec!["a"]), + ("$.\"a.b\"", vec!["a.b"]), + ("$.nested.value", vec!["nested", "value"]), + ] { + let (target, _) = builder + .expr_to_index_target(&json_expr(path, None)) + .unwrap() + .unwrap(); + assert_eq!( + target, + IndexTarget::new_json_path(1, expected.into_iter().map(String::from).collect()) + .unwrap() + ); + } + for path in ["$.missing", "$.unindexed", "$.a[0]", "$.*", "$"] { + assert!( + builder + .expr_to_index_target(&json_expr(path, None)) + .unwrap() + .is_none(), + "{path}" + ); + } + let data_type = ConcreteDataType::int64_datatype(); + let expected = Some(( + IndexTarget::new_json_path(1, vec!["n".into()]).unwrap(), + data_type, + )); + for ty in [None, Some(ScalarValue::Int64(None))] { + let resolved = builder.expr_to_index_target(&json_expr("$.n", ty)).unwrap(); + assert_eq!(resolved, expected); + } + assert!( + builder + .expr_to_index_target(&json_expr("$.n", Some(ScalarValue::Float64(None)))) + .unwrap() + .is_none() + ); + assert!( + builder + .expr_to_index_target(&col("attrs")) + .unwrap() + .is_none() + ); + } + + #[test] + fn test_json_plan_checks_sst_type_hints() { + use datafusion_expr::{col, lit}; + use datatypes::extension::json::JsonMetadata; + use datatypes::json::{JsonSettings, JsonTypeHint}; + + let (_dir, factory) = PuffinManagerFactory::new_for_test_block("json_plan_type_hints"); + let metadata = json_metadata(); + for include_column in [false, true] { + let builder = InvertedIndexApplierBuilder::new( + "test".into(), + PathType::Bare, + test_object_store(), + &metadata, + HashSet::from([2]), + factory.clone(), + ); + let mut exprs = vec![json_expr("$.n", None).eq(lit(42_i64))]; + if include_column { + exprs.push(col("ts").gt(lit(ScalarValue::TimestampMillisecond(Some(0), None)))); + } + let applier = builder.build(&exprs).unwrap().unwrap(); + let plan = applier + .plan_for_sst(&Arc::new(metadata.clone())) + .unwrap() + .unwrap(); + assert_eq!(plan.predicate_key, applier.default_plan.predicate_key); + + // Changed types, missing paths, and disabled hints cannot use the JSON index. + // Compatible ordinary predicates must still be applied. + for hints in [ + vec![JsonTypeHint { + path: vec!["n".into()], + data_type: ConcreteDataType::float64_datatype(), + inverted_index: true, + }], + vec![], + vec![JsonTypeHint { + path: vec!["n".into()], + data_type: ConcreteDataType::int64_datatype(), + inverted_index: false, + }], + ] { + let mut sst = metadata.clone(); + sst.column_metadatas[0].column_schema.with_extension_type( + &Json2ExtensionType::new(Arc::new(JsonMetadata::new( + JsonSettings::try_new(hints, Some(0)).unwrap(), + ))), + ); + let plan = applier.plan_for_sst(&Arc::new(sst)).unwrap(); + if !include_column { + assert!(plan.is_none()); + continue; + } + let expected = BTreeMap::from([( + IndexTarget::ColumnId(2), + applier.predicates[&IndexTarget::ColumnId(2)].clone(), + )]); + assert_eq!( + plan.unwrap().predicate_key, + InvertedIndexApplier::build_apply_plan(&expected) + .unwrap() + .predicate_key + ); + } + } + } + + #[test] + fn test_json_predicate_collection() { + use datafusion_expr::lit; + + let (_dir, factory) = PuffinManagerFactory::new_for_test_block("json_predicate_collection"); + let metadata = json_metadata(); + let mut builder = InvertedIndexApplierBuilder::new( + "test".into(), + PathType::Bare, + test_object_store(), + &metadata, + HashSet::new(), + factory, + ); + let numeric_target = IndexTarget::new_json_path(1, vec!["n".into()]).unwrap(); + let n = json_expr("$.n", Some(ScalarValue::Int64(None))); + for expr in [ + n.clone().eq(lit(42_i64)), + lit(40_i64).lt(n.clone()), + n.clone().between(lit(1_i64), lit(99_i64)), + n.clone().in_list(vec![lit(1_i64), lit(2_i64)], false), + n.clone().eq(lit(1_i64)).or(n.eq(lit(2_i64))), + ] { + builder.output.clear(); + builder.traverse_and_collect(&expr); + assert_eq!(builder.output.len(), 1, "{expr}"); + assert_eq!( + builder.output.get(&numeric_target).unwrap().len(), + 1, + "{expr}" + ); + } + builder.output.clear(); + let expr = json_expr("$.a", None) + .eq(lit("x")) + .or(json_expr("$.b", None).eq(lit("y"))); + builder.traverse_and_collect(&expr); + assert!(builder.output.is_empty()); + let string_target = IndexTarget::new_json_path(1, vec!["a".into()]).unwrap(); + for pattern in [ + ScalarValue::Utf8(Some("a.*[ip]".into())), + ScalarValue::Utf8View(Some("a.*[ip]".into())), + ScalarValue::LargeUtf8(Some("a.*[ip]".into())), + ] { + builder.output.clear(); + builder + .collect_regex_match(&json_expr("$.a", None), &Expr::Literal(pattern, None)) + .unwrap(); + assert_eq!(builder.output.len(), 1); + assert_eq!(builder.output.get(&string_target).unwrap().len(), 1); + } + } + + #[test] + fn test_collect_and_basic() { + let (_d, factory) = PuffinManagerFactory::new_for_test_block("test_collect_and_basic_"); + + let metadata = test_region_metadata(); + let mut builder = InvertedIndexApplierBuilder::new( + "test".to_string(), + PathType::Bare, + test_object_store(), + &metadata, + HashSet::from_iter([1, 2, 3]), + factory, + ); + + let expr = Expr::BinaryExpr(BinaryExpr { + left: Box::new(Expr::BinaryExpr(BinaryExpr { + left: Box::new(tag_column()), + op: Operator::RegexMatch, + right: Box::new(string_lit("bar")), + })), + op: Operator::And, + right: Box::new(Expr::Between(Between { + expr: Box::new(tag_column2()), + negated: false, + low: Box::new(int64_lit(123)), + high: Box::new(int64_lit(456)), + })), + }); + + builder.traverse_and_collect(&expr); + let predicates = builder + .output + .get(&index::target::IndexTarget::ColumnId(1)) + .unwrap(); + assert_eq!(predicates.len(), 1); + assert_eq!( + predicates[0], + Predicate::RegexMatch(RegexMatchPredicate { + pattern: "bar".to_string() + }) + ); + let predicates = builder + .output + .get(&index::target::IndexTarget::ColumnId(2)) + .unwrap(); + assert_eq!(predicates.len(), 1); + assert_eq!( + predicates[0], + Predicate::Range(RangePredicate { + range: Range { + lower: Some(Bound { + inclusive: true, + value: encoded_int64(123), + }), + upper: Some(Bound { + inclusive: true, + value: encoded_int64(456), + }), + } + }) + ); + } + pub(crate) fn test_region_metadata() -> RegionMetadata { let mut builder = RegionMetadataBuilder::new(RegionId::new(1234, 5678)); builder @@ -349,60 +686,66 @@ mod tests { bytes } - #[test] - fn test_collect_and_basic() { - let (_d, factory) = PuffinManagerFactory::new_for_test_block("test_collect_and_basic_"); + fn json_expr(path: &str, ty: Option) -> Expr { + use common_function::scalars::udf::create_udf; + use datafusion_expr::expr::ScalarFunction; + let mut args = vec![datafusion_expr::col("attrs"), string_lit(path)]; + if let Some(ty) = ty { + args.push(Expr::Literal(ty, None)); + } + Expr::ScalarFunction(ScalarFunction::new_udf( + Arc::new(create_udf(Arc::new(JsonGetWithType::default()))), + args, + )) + } - let metadata = test_region_metadata(); - let mut builder = InvertedIndexApplierBuilder::new( - "test".to_string(), - PathType::Bare, - test_object_store(), - &metadata, - HashSet::from_iter([1, 2, 3]), - factory, + fn json_metadata() -> RegionMetadata { + use datatypes::extension::json::JsonMetadata; + use datatypes::json::{JsonSettings, JsonTypeHint}; + use datatypes::types::json_type::{JsonNativeType, JsonObjectType}; + let mut column = ColumnSchema::new( + "attrs", + ConcreteDataType::json2(JsonNativeType::Object(JsonObjectType::new())), + true, ); - - let expr = Expr::BinaryExpr(BinaryExpr { - left: Box::new(Expr::BinaryExpr(BinaryExpr { - left: Box::new(tag_column()), - op: Operator::RegexMatch, - right: Box::new(string_lit("bar")), - })), - op: Operator::And, - right: Box::new(Expr::Between(Between { - expr: Box::new(tag_column2()), - negated: false, - low: Box::new(int64_lit(123)), - high: Box::new(int64_lit(456)), - })), - }); - - builder.traverse_and_collect(&expr); - let predicates = builder.output.get(&1).unwrap(); - assert_eq!(predicates.len(), 1); - assert_eq!( - predicates[0], - Predicate::RegexMatch(RegexMatchPredicate { - pattern: "bar".to_string() + let hints = [ + (vec!["a"], true), + (vec!["b"], true), + (vec!["a.b"], true), + (vec!["nested", "value"], true), + (vec!["unindexed"], false), + ] + .into_iter() + .map(|(path, inverted_index)| JsonTypeHint { + path: path.into_iter().map(String::from).collect(), + data_type: ConcreteDataType::string_datatype(), + inverted_index, + }) + .chain(std::iter::once(JsonTypeHint { + path: vec!["n".into()], + data_type: ConcreteDataType::int64_datatype(), + inverted_index: true, + })) + .collect(); + column.with_extension_type(&Json2ExtensionType::new(Arc::new(JsonMetadata::new( + JsonSettings::try_new(hints, Some(0)).unwrap(), + )))); + let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 1)); + builder + .push_column_metadata(ColumnMetadata { + column_schema: column, + column_id: 1, + semantic_type: SemanticType::Field, }) - ); - let predicates = builder.output.get(&2).unwrap(); - assert_eq!(predicates.len(), 1); - assert_eq!( - predicates[0], - Predicate::Range(RangePredicate { - range: Range { - lower: Some(Bound { - inclusive: true, - value: encoded_int64(123), - }), - upper: Some(Bound { - inclusive: true, - value: encoded_int64(456), - }), - } - }) - ); + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + column_id: 2, + semantic_type: SemanticType::Timestamp, + }); + builder.build().unwrap() } } diff --git a/src/mito2/src/sst/index/inverted_index/applier/builder/between.rs b/src/mito2/src/sst/index/inverted_index/applier/builder/between.rs index f9e9f22ba8f..7937b088124 100644 --- a/src/mito2/src/sst/index/inverted_index/applier/builder/between.rs +++ b/src/mito2/src/sst/index/inverted_index/applier/builder/between.rs @@ -25,10 +25,7 @@ impl InvertedIndexApplierBuilder<'_> { return Ok(()); } - let Some(column_name) = Self::column_name(&between.expr) else { - return Ok(()); - }; - let Some((column_id, data_type)) = self.column_id_and_type(column_name)? else { + let Some((target, data_type)) = self.expr_to_index_target(&between.expr)? else { return Ok(()); }; let Some(low) = Self::nonnull_lit(&between.low) else { @@ -51,7 +48,7 @@ impl InvertedIndexApplierBuilder<'_> { }, }); - self.add_predicate(column_id, predicate); + self.add_predicate(target, predicate); Ok(()) } } @@ -92,7 +89,10 @@ mod tests { builder.collect_between(&between).unwrap(); - let predicates = builder.output.get(&1).unwrap(); + let predicates = builder + .output + .get(&index::target::IndexTarget::ColumnId(1)) + .unwrap(); assert_eq!(predicates.len(), 1); assert_eq!( predicates[0], @@ -159,7 +159,10 @@ mod tests { builder.collect_between(&between).unwrap(); - let predicates = builder.output.get(&3).unwrap(); + let predicates = builder + .output + .get(&index::target::IndexTarget::ColumnId(3)) + .unwrap(); assert_eq!(predicates.len(), 1); assert_eq!( predicates[0], diff --git a/src/mito2/src/sst/index/inverted_index/applier/builder/comparison.rs b/src/mito2/src/sst/index/inverted_index/applier/builder/comparison.rs index 1a02862996d..b1873f148d9 100644 --- a/src/mito2/src/sst/index/inverted_index/applier/builder/comparison.rs +++ b/src/mito2/src/sst/index/inverted_index/applier/builder/comparison.rs @@ -20,9 +20,14 @@ use crate::error::Result; use crate::sst::index::inverted_index::applier::builder::InvertedIndexApplierBuilder; impl InvertedIndexApplierBuilder<'_> { - /// Collects a comparison expression in the form of - /// `column < lit`, `column > lit`, `column <= lit`, `column >= lit`, - /// `lit < column`, `lit > column`, `lit <= column`, `lit >= column`. + /// Collects `<`, `<=`, `>`, or `>=` comparisons between an indexed column or + /// JSON path expression and a non-null literal. + /// + /// The literal may appear on either side; comparisons with the literal on + /// the left are reversed before building the index predicate. + /// + /// For example, `42 < col` becomes `col > 42`, and + /// `42 < json_get(attrs, '$.value')` becomes `json_get(attrs, '$.value') > 42`. pub(crate) fn collect_comparison_expr( &mut self, left: &DfExpr, @@ -31,28 +36,28 @@ impl InvertedIndexApplierBuilder<'_> { ) -> Result<()> { match op { Operator::Lt => { - if matches!(right, DfExpr::Column(_)) { + if Self::nonnull_lit(left).is_some() { self.collect_column_gt_lit(right, left) } else { self.collect_column_lt_lit(left, right) } } Operator::LtEq => { - if matches!(right, DfExpr::Column(_)) { + if Self::nonnull_lit(left).is_some() { self.collect_column_ge_lit(right, left) } else { self.collect_column_le_lit(left, right) } } Operator::Gt => { - if matches!(right, DfExpr::Column(_)) { + if Self::nonnull_lit(left).is_some() { self.collect_column_lt_lit(right, left) } else { self.collect_column_gt_lit(left, right) } } Operator::GtEq => { - if matches!(right, DfExpr::Column(_)) { + if Self::nonnull_lit(left).is_some() { self.collect_column_le_lit(right, left) } else { self.collect_column_ge_lit(left, right) @@ -108,13 +113,10 @@ impl InvertedIndexApplierBuilder<'_> { literal: &DfExpr, range_builder: impl FnOnce(Bytes) -> Range, ) -> Result<()> { - let Some(column_name) = Self::column_name(column) else { - return Ok(()); - }; let Some(lit) = Self::nonnull_lit(literal) else { return Ok(()); }; - let Some((column_id, data_type)) = self.column_id_and_type(column_name)? else { + let Some((target, data_type)) = self.expr_to_index_target(column)? else { return Ok(()); }; @@ -122,7 +124,7 @@ impl InvertedIndexApplierBuilder<'_> { range: range_builder(Self::encode_lit(lit, data_type)?), }); - self.add_predicate(column_id, predicate); + self.add_predicate(target, predicate); Ok(()) } } @@ -243,7 +245,10 @@ mod tests { builder.collect_comparison_expr(left, op, right).unwrap(); } - let predicates = builder.output.get(&1).unwrap(); + let predicates = builder + .output + .get(&index::target::IndexTarget::ColumnId(1)) + .unwrap(); assert_eq!(predicates.len(), cases.len()); for ((_, expected), actual) in cases.into_iter().zip(predicates) { assert_eq!( @@ -290,7 +295,10 @@ mod tests { .collect_comparison_expr(&field_column(), &Operator::Lt, &string_lit("abc")) .unwrap(); - let predicates = builder.output.get(&3).unwrap(); + let predicates = builder + .output + .get(&index::target::IndexTarget::ColumnId(3)) + .unwrap(); assert_eq!(predicates.len(), 1); assert_eq!( predicates[0], diff --git a/src/mito2/src/sst/index/inverted_index/applier/builder/eq_list.rs b/src/mito2/src/sst/index/inverted_index/applier/builder/eq_list.rs index 310ea3786ca..e6cb6ad087a 100644 --- a/src/mito2/src/sst/index/inverted_index/applier/builder/eq_list.rs +++ b/src/mito2/src/sst/index/inverted_index/applier/builder/eq_list.rs @@ -18,6 +18,7 @@ use datafusion_expr::{BinaryExpr, Expr as DfExpr, Operator}; use datatypes::data_type::ConcreteDataType; use index::Bytes; use index::inverted_index::search::predicate::{InListPredicate, Predicate}; +use index::target::IndexTarget; use crate::error::Result; use crate::sst::index::inverted_index::applier::builder::InvertedIndexApplierBuilder; @@ -25,20 +26,21 @@ use crate::sst::index::inverted_index::applier::builder::InvertedIndexApplierBui impl InvertedIndexApplierBuilder<'_> { /// Collects an eq expression in the form of `column = lit`. pub(crate) fn collect_eq(&mut self, left: &DfExpr, right: &DfExpr) -> Result<()> { - let Some(column_name) = Self::column_name(left).or_else(|| Self::column_name(right)) else { + let resolved = match self.expr_to_index_target(left)? { + Some(target) => Some(target), + None => self.expr_to_index_target(right)?, + }; + let Some((target, data_type)) = resolved else { return Ok(()); }; let Some(lit) = Self::nonnull_lit(right).or_else(|| Self::nonnull_lit(left)) else { return Ok(()); }; - let Some((column_id, data_type)) = self.column_id_and_type(column_name)? else { - return Ok(()); - }; let predicate = Predicate::InList(InListPredicate { list: BTreeSet::from_iter([Self::encode_lit(lit, data_type)?]), }); - self.add_predicate(column_id, predicate); + self.add_predicate(target, predicate); Ok(()) } @@ -53,22 +55,23 @@ impl InvertedIndexApplierBuilder<'_> { return Ok(()); }; - let Some(column_name) = Self::column_name(left).or_else(|| Self::column_name(right)) else { + let resolved = match self.expr_to_index_target(left)? { + Some(target) => Some(target), + None => self.expr_to_index_target(right)?, + }; + let Some((target, data_type)) = resolved else { return Ok(()); }; let Some(lit) = Self::nonnull_lit(right).or_else(|| Self::nonnull_lit(left)) else { return Ok(()); }; - let Some((column_id, data_type)) = self.column_id_and_type(column_name)? else { - return Ok(()); - }; let bytes = Self::encode_lit(lit, data_type.clone())?; let mut inlist = BTreeSet::from_iter([bytes]); - if Self::collect_eq_list_inner(column_name, &data_type, or_list, &mut inlist)? { + if self.collect_eq_list_inner(&target, &data_type, or_list, &mut inlist)? { let predicate = Predicate::InList(InListPredicate { list: inlist }); - self.add_predicate(column_id, predicate); + self.add_predicate(target, predicate); } Ok(()) @@ -79,7 +82,8 @@ impl InvertedIndexApplierBuilder<'_> { /// Returns false if the expression doesn't match the form then /// caller can safely ignore the expression. fn collect_eq_list_inner( - column_name: &str, + &self, + target: &IndexTarget, data_type: &ConcreteDataType, expr: &DfExpr, inlist: &mut BTreeSet, @@ -94,18 +98,23 @@ impl InvertedIndexApplierBuilder<'_> { }; if op == &Operator::Or { - let r = Self::collect_eq_list_inner(column_name, data_type, left, inlist)? - .then(|| Self::collect_eq_list_inner(column_name, data_type, right, inlist)) + let r = self + .collect_eq_list_inner(target, data_type, left, inlist)? + .then(|| self.collect_eq_list_inner(target, data_type, right, inlist)) .transpose()? .unwrap_or(false); return Ok(r); } if op == &Operator::Eq { - let Some(name) = Self::column_name(left).or_else(|| Self::column_name(right)) else { + let resolved = match self.expr_to_index_target(left)? { + Some(target) => Some(target), + None => self.expr_to_index_target(right)?, + }; + let Some((candidate, _)) = resolved else { return Ok(false); }; - if column_name != name { + if target != &candidate { return Ok(false); } let Some(lit) = Self::nonnull_lit(right).or_else(|| Self::nonnull_lit(left)) else { @@ -154,7 +163,10 @@ mod tests { .collect_eq(&string_lit("bar"), &tag_column()) .unwrap(); - let predicates = builder.output.get(&1).unwrap(); + let predicates = builder + .output + .get(&index::target::IndexTarget::ColumnId(1)) + .unwrap(); assert_eq!(predicates.len(), 2); assert_eq!( predicates[0], @@ -188,7 +200,10 @@ mod tests { .collect_eq(&field_column(), &string_lit("abc")) .unwrap(); - let predicates = builder.output.get(&3).unwrap(); + let predicates = builder + .output + .get(&index::target::IndexTarget::ColumnId(3)) + .unwrap(); assert_eq!(predicates.len(), 1); assert_eq!( predicates[0], @@ -279,7 +294,10 @@ mod tests { builder.collect_or_eq_list(&eq_expr, &or_eq_list).unwrap(); - let predicates = builder.output.get(&1).unwrap(); + let predicates = builder + .output + .get(&index::target::IndexTarget::ColumnId(1)) + .unwrap(); assert_eq!(predicates.len(), 1); assert_eq!( predicates[0], diff --git a/src/mito2/src/sst/index/inverted_index/applier/builder/in_list.rs b/src/mito2/src/sst/index/inverted_index/applier/builder/in_list.rs index 297becd788f..da75297194e 100644 --- a/src/mito2/src/sst/index/inverted_index/applier/builder/in_list.rs +++ b/src/mito2/src/sst/index/inverted_index/applier/builder/in_list.rs @@ -26,10 +26,7 @@ impl InvertedIndexApplierBuilder<'_> { if inlist.negated { return Ok(()); } - let Some(column_name) = Self::column_name(&inlist.expr) else { - return Ok(()); - }; - let Some((column_id, data_type)) = self.column_id_and_type(column_name)? else { + let Some((target, data_type)) = self.expr_to_index_target(&inlist.expr)? else { return Ok(()); }; @@ -46,7 +43,7 @@ impl InvertedIndexApplierBuilder<'_> { .insert(Self::encode_lit(lit, data_type.clone())?); } - self.add_predicate(column_id, Predicate::InList(predicate)); + self.add_predicate(target, Predicate::InList(predicate)); Ok(()) } } @@ -86,7 +83,10 @@ mod tests { builder.collect_inlist(&in_list).unwrap(); - let predicates = builder.output.get(&1).unwrap(); + let predicates = builder + .output + .get(&index::target::IndexTarget::ColumnId(1)) + .unwrap(); assert_eq!(predicates.len(), 1); assert_eq!( predicates[0], @@ -142,7 +142,10 @@ mod tests { builder.collect_inlist(&in_list).unwrap(); - let predicates = builder.output.get(&3).unwrap(); + let predicates = builder + .output + .get(&index::target::IndexTarget::ColumnId(3)) + .unwrap(); assert_eq!(predicates.len(), 1); assert_eq!( predicates[0], diff --git a/src/mito2/src/sst/index/inverted_index/applier/builder/regex_match.rs b/src/mito2/src/sst/index/inverted_index/applier/builder/regex_match.rs index 2d2e9fe4ca8..b8a7b987edc 100644 --- a/src/mito2/src/sst/index/inverted_index/applier/builder/regex_match.rs +++ b/src/mito2/src/sst/index/inverted_index/applier/builder/regex_match.rs @@ -12,7 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -use datafusion_common::ScalarValue; use datafusion_expr::Expr as DfExpr; use index::inverted_index::search::predicate::{Predicate, RegexMatchPredicate}; @@ -22,23 +21,23 @@ use crate::sst::index::inverted_index::applier::builder::InvertedIndexApplierBui impl InvertedIndexApplierBuilder<'_> { /// Collects a regex match expression in the form of `column ~ pattern`. pub(crate) fn collect_regex_match(&mut self, column: &DfExpr, pattern: &DfExpr) -> Result<()> { - let Some(column_name) = Self::column_name(column) else { - return Ok(()); - }; - let Some((column_id, data_type)) = self.column_id_and_type(column_name)? else { + let Some((target, data_type)) = self.expr_to_index_target(column)? else { return Ok(()); }; if !data_type.is_string() { return Ok(()); } - let DfExpr::Literal(ScalarValue::Utf8(Some(pattern)), _) = pattern else { + let Some(pattern) = Self::nonnull_lit(pattern) + .and_then(|value| value.try_as_str()) + .flatten() + else { return Ok(()); }; let predicate = Predicate::RegexMatch(RegexMatchPredicate { - pattern: pattern.clone(), + pattern: pattern.to_owned(), }); - self.add_predicate(column_id, predicate); + self.add_predicate(target, predicate); Ok(()) } } @@ -74,7 +73,10 @@ mod tests { .collect_regex_match(&tag_column(), &string_lit("abc")) .unwrap(); - let predicates = builder.output.get(&1).unwrap(); + let predicates = builder + .output + .get(&index::target::IndexTarget::ColumnId(1)) + .unwrap(); assert_eq!(predicates.len(), 1); assert_eq!( predicates[0], @@ -102,7 +104,10 @@ mod tests { .collect_regex_match(&field_column(), &string_lit("abc")) .unwrap(); - let predicates = builder.output.get(&3).unwrap(); + let predicates = builder + .output + .get(&index::target::IndexTarget::ColumnId(3)) + .unwrap(); assert_eq!(predicates.len(), 1); assert_eq!( predicates[0], diff --git a/src/mito2/src/sst/index/inverted_index/creator.rs b/src/mito2/src/sst/index/inverted_index/creator.rs index 9c1ba5cbca0..19c55253b58 100644 --- a/src/mito2/src/sst/index/inverted_index/creator.rs +++ b/src/mito2/src/sst/index/inverted_index/creator.rs @@ -12,7 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::HashSet; use std::num::NonZeroUsize; use std::sync::Arc; use std::sync::atomic::AtomicUsize; @@ -31,7 +30,7 @@ use mito_codec::row_converter::sparse::SparsePrimaryKeyView; use mito_codec::row_converter::{SortField, SparseOffsetsCache}; use puffin::puffin_manager::{PuffinWriter, PutOptions}; use smallvec::SmallVec; -use snafu::{ResultExt, ensure}; +use snafu::{OptionExt, ResultExt, ensure}; use store_api::codec::PrimaryKeyEncoding; use store_api::metadata::RegionMetadataRef; use store_api::storage::{ColumnId, FileId}; @@ -39,12 +38,13 @@ use tokio::io::duplex; use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; use crate::error::{ - BiErrorsSnafu, DecodeSnafu, EncodeSnafu, IndexFinishSnafu, OperateAbortedIndexSnafu, - PuffinAddBlobSnafu, PushIndexValueSnafu, Result, + BiErrorsSnafu, ConvertVectorSnafu, DecodeSnafu, EncodeSnafu, IndexFinishSnafu, + InvalidRecordBatchSnafu, OperateAbortedIndexSnafu, PuffinAddBlobSnafu, PushIndexValueSnafu, + Result, }; use crate::read::Batch; use crate::sst::index::TYPE_INVERTED_INDEX; -use crate::sst::index::column::column_index_rows; +use crate::sst::index::column::{column_index_rows, extract_json_type_hint_values}; use crate::sst::index::intermediate::{ IntermediateLocation, IntermediateManager, TempFileProvider, }; @@ -81,8 +81,8 @@ pub struct InvertedIndexer { /// The memory usage of the index creator. memory_usage: Arc, - /// Ids of indexed columns and their encoded target keys. - indexed_column_ids: Vec<(ColumnId, String)>, + /// Index targets and their encoded keys. + index_targets: Vec<(IndexTarget, String)>, /// Region metadata for column lookups. metadata: RegionMetadataRef, @@ -90,14 +90,13 @@ pub struct InvertedIndexer { impl InvertedIndexer { /// Creates a new `InvertedIndexer`. - /// Should ensure that the number of tag columns is greater than 0. pub fn new( sst_file_id: FileId, metadata: &RegionMetadataRef, intermediate_manager: IntermediateManager, memory_usage_threshold: Option, segment_row_count: NonZeroUsize, - indexed_column_ids: HashSet, + index_targets: Vec, ) -> Self { let temp_file_provider = Arc::new(TempFileProvider::new( IntermediateLocation::new(&metadata.region_id, &sst_file_id), @@ -118,11 +117,11 @@ impl InvertedIndexer { metadata.primary_key_encoding, metadata.primary_key_columns(), ); - let indexed_column_ids = indexed_column_ids + let index_targets = index_targets .into_iter() - .map(|col_id| { - let target_key = format!("{}", IndexTarget::ColumnId(col_id)); - (col_id, target_key) + .map(|target| { + let key = target.encode(); + (target, key) }) .collect(); Self { @@ -134,7 +133,7 @@ impl InvertedIndexer { stats: Statistics::new(TYPE_INVERTED_INDEX), aborted: false, memory_usage, - indexed_column_ids, + index_targets, metadata: metadata.clone(), } } @@ -182,44 +181,69 @@ impl InvertedIndexer { let is_sparse = self.metadata.primary_key_encoding == PrimaryKeyEncoding::Sparse; let mut sparse_columns: SmallVec<[(ColumnId, &str); 8]> = SmallVec::new(); - for (col_id, target_key) in &self.indexed_column_ids { - let Some(column_meta) = self.metadata.column_by_id(*col_id) else { - debug!( - "Column {} not found in the metadata during building inverted index", - col_id - ); - continue; + for (target, target_key) in &self.index_targets { + let (vector, sort_field, semantic_type) = match target { + IndexTarget::ColumnId(col_id) => { + let Some(col_meta) = self.metadata.column_by_id(*col_id) else { + debug!( + "Column {} not found in the metadata during building inverted index", + col_id + ); + continue; + }; + let col_name = &col_meta.column_schema.name; + let Some(col_array) = batch.column_by_name(col_name) else { + if is_sparse && col_meta.semantic_type == SemanticType::Tag { + if self.codec.pk_col_info(*col_id).is_some() { + sparse_columns.push((*col_id, target_key)); + } + } else { + debug!( + "Column {} not found in the batch during building inverted index", + col_id + ); + } + continue; + }; + let vector = + Helper::try_into_vector(col_array.clone()).context(ConvertVectorSnafu)?; + let sort_field = SortField::new(vector.data_type()); + (vector, sort_field, col_meta.semantic_type) + } + IndexTarget::JsonPath(target) => { + let col_id = target.column_id(); + let col = self.metadata.column_by_id(col_id).with_context(|| { + InvalidRecordBatchSnafu { + reason: format!("Missing JSON index column {col_id}"), + } + })?; + let root = + batch + .column_by_name(&col.column_schema.name) + .with_context(|| InvalidRecordBatchSnafu { + reason: format!("Missing JSON index column {col_id} in batch"), + })?; + let leaf = extract_json_type_hint_values(root, target.path())?; + let field = SortField::new(leaf.data_type()); + (leaf, field, SemanticType::Field) + } }; - let column_name = &column_meta.column_schema.name; - if let Some(column_array) = batch.column_by_name(column_name) { - // Convert Arrow array to VectorRef using Helper - let vector = Helper::try_into_vector(column_array.clone()) - .context(crate::error::ConvertVectorSnafu)?; - let sort_field = SortField::new(vector.data_type()); - - for (row, count) in column_index_rows(batch, column_meta.semantic_type) { - let elem = IndexValueCodec::encode_value( - vector.get_ref(row), - &sort_field, - &mut self.value_buf, - ) - .context(EncodeSnafu)?; - if self.index_creator.push_with_name_n(target_key, elem, count) { - self.index_creator - .spill() - .await - .context(PushIndexValueSnafu)?; - } + for (row, count) in column_index_rows(batch, semantic_type) { + let value = IndexValueCodec::encode_value( + vector.get_ref(row), + &sort_field, + &mut self.value_buf, + ) + .context(EncodeSnafu)?; + if self + .index_creator + .push_with_name_n(target_key, value, count) + { + self.index_creator + .spill() + .await + .context(PushIndexValueSnafu)?; } - } else if is_sparse && column_meta.semantic_type == SemanticType::Tag { - if self.codec.pk_col_info(*col_id).is_some() { - sparse_columns.push((*col_id, target_key)); - } - } else { - debug!( - "Column {} not found in the batch during building inverted index", - col_id - ); } } @@ -293,36 +317,35 @@ impl InvertedIndexer { let n = batch.num_rows(); guard.inc_row_count(n); - for (col_id, target_key) in &self.indexed_column_ids { - match self.codec.pk_col_info(*col_id) { - // pk - Some(col_info) => { - let pk_idx = col_info.idx; - let field = &col_info.field; - let value = batch - .pk_col_value(self.codec.decoder(), pk_idx, *col_id)? - .filter(|v| !v.is_null()) - .map(|v| { - self.value_buf.clear(); - IndexValueCodec::encode_nonnull_value( - v.as_value_ref(), - field, - &mut self.value_buf, - ) - .context(EncodeSnafu)?; - Ok(self.value_buf.as_slice()) - }) - .transpose()?; + for (target, target_key) in &self.index_targets { + let (vector, sort_field) = match target { + IndexTarget::ColumnId(col_id) => { + if let Some(col_info) = self.codec.pk_col_info(*col_id) { + let pk_idx = col_info.idx; + let field = &col_info.field; + let value = batch + .pk_col_value(self.codec.decoder(), pk_idx, *col_id)? + .filter(|v| !v.is_null()) + .map(|v| { + self.value_buf.clear(); + IndexValueCodec::encode_nonnull_value( + v.as_value_ref(), + field, + &mut self.value_buf, + ) + .context(EncodeSnafu)?; + Ok(self.value_buf.as_slice()) + }) + .transpose()?; - if self.index_creator.push_with_name_n(target_key, value, n) { - self.index_creator - .spill() - .await - .context(PushIndexValueSnafu)?; + if self.index_creator.push_with_name_n(target_key, value, n) { + self.index_creator + .spill() + .await + .context(PushIndexValueSnafu)?; + } + continue; } - } - // fields - None => { let Some(values) = batch.field_col_value(*col_id) else { debug!( "Column {} not found in the batch during building inverted index", @@ -331,34 +354,34 @@ impl InvertedIndexer { continue; }; let sort_field = SortField::new(values.data.data_type()); - for i in 0..n { - self.value_buf.clear(); - let value = values.data.get_ref(i); - if value.is_null() { - if self.index_creator.push_with_name(target_key, None) { - self.index_creator - .spill() - .await - .context(PushIndexValueSnafu)?; - } - } else { - IndexValueCodec::encode_nonnull_value( - value, - &sort_field, - &mut self.value_buf, - ) - .context(EncodeSnafu)?; - if self - .index_creator - .push_with_name(target_key, Some(&self.value_buf)) - { - self.index_creator - .spill() - .await - .context(PushIndexValueSnafu)?; - } - } - } + (values.data.clone(), sort_field) + } + IndexTarget::JsonPath(target) => { + let col_id = target.column_id(); + let root = batch + .field_col_value(col_id) + .with_context(|| InvalidRecordBatchSnafu { + reason: format!("Missing JSON index column {col_id} in batch"), + })? + .data + .to_arrow_array(); + let leaf = extract_json_type_hint_values(&root, target.path())?; + let field = SortField::new(leaf.data_type()); + (leaf, field) + } + }; + for row in 0..n { + let value = IndexValueCodec::encode_value( + vector.get_ref(row), + &sort_field, + &mut self.value_buf, + ) + .context(EncodeSnafu)?; + if self.index_creator.push_with_name(target_key, value) { + self.index_creator + .spill() + .await + .context(PushIndexValueSnafu)?; } } } @@ -428,8 +451,14 @@ impl InvertedIndexer { self.temp_file_provider.cleanup().await } - pub fn column_ids(&self) -> impl Iterator + '_ { - self.indexed_column_ids.iter().map(|(col_id, _)| *col_id) + /// Returns column-level index IDs, excluding JSON path targets. + pub fn indexed_column_ids(&self) -> impl Iterator + '_ { + self.index_targets + .iter() + .filter_map(|(target, _)| match target { + IndexTarget::ColumnId(col_id) => Some(*col_id), + IndexTarget::JsonPath(_) => None, + }) } pub fn memory_usage(&self) -> usize { @@ -439,20 +468,25 @@ impl InvertedIndexer { #[cfg(test)] mod tests { - use std::collections::BTreeSet; + use std::collections::{BTreeMap, BTreeSet, HashSet}; use api::v1::SemanticType; use datafusion_expr::{Expr as DfExpr, Operator, binary_expr, col, lit}; + use datatypes::arrow::array::{Array, ArrayRef, Int64Array, StructArray}; + use datatypes::arrow::buffer::NullBuffer; + use datatypes::arrow::datatypes::{Field, Schema}; use datatypes::data_type::ConcreteDataType; use datatypes::schema::ColumnSchema; use datatypes::value::ValueRef; use datatypes::vectors::{UInt8Vector, UInt64Vector}; use futures::future::BoxFuture; + use index::inverted_index::format::reader::{InvertedIndexBlobReader, InvertedIndexReader}; + use index::inverted_index::search::predicate::{InListPredicate, Predicate}; use mito_codec::row_converter::{DensePrimaryKeyCodec, PrimaryKeyCodecExt}; use object_store::ObjectStore; use object_store::services::Memory; - use puffin::puffin_manager::PuffinManager; use puffin::puffin_manager::cache::PuffinMetadataCache; + use puffin::puffin_manager::{PuffinManager, PuffinReader}; use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; use store_api::region_request::PathType; use store_api::storage::RegionId; @@ -463,6 +497,7 @@ mod tests { use crate::metrics::CACHE_BYTES; use crate::read::BatchColumn; use crate::sst::file::{RegionFileId, RegionIndexId}; + use crate::sst::index::inverted_index::applier::InvertedIndexApplier; use crate::sst::index::inverted_index::applier::builder::InvertedIndexApplierBuilder; use crate::sst::index::puffin_manager::PuffinManagerFactory; @@ -574,7 +609,11 @@ mod tests { intm_mgr, memory_threshold, NonZeroUsize::new(segment_row_count).unwrap(), - sst_indexed_column_ids, + sst_indexed_column_ids + .iter() + .copied() + .map(IndexTarget::ColumnId) + .collect(), ); for (str_tag, i32_tag, u64_field) in &rows { @@ -661,6 +700,156 @@ mod tests { } } + #[tokio::test] + async fn test_create_json_hint_index_both_formats() { + for flat in [false, true] { + let (dir, factory) = PuffinManagerFactory::new_for_test_async("json_hint_index").await; + let mut metadata = (*mock_region_metadata()).clone(); + let column = &mut metadata.column_metadatas[3].column_schema; + column.data_type = + ConcreteDataType::json2(datatypes::types::json_type::JsonNativeType::Object( + datatypes::types::json_type::JsonObjectType::new(), + )); + column.with_extension_type(&datatypes::extension::json::Json2ExtensionType::new( + Arc::new(datatypes::extension::json::JsonMetadata::new( + datatypes::json::JsonSettings::try_new( + vec![datatypes::json::JsonTypeHint { + path: vec!["value".into()], + data_type: ConcreteDataType::int64_datatype(), + inverted_index: true, + }], + Some(0), + ) + .unwrap(), + )), + )); + let metadata = Arc::new(metadata); + let file_id = FileId::random(); + let target = IndexTarget::new_json_path(4, vec!["value".into()]).unwrap(); + let mut creator = InvertedIndexer::new( + file_id, + &metadata, + new_intm_mgr(dir.path().to_string_lossy()).await, + None, + NonZeroUsize::new(2).unwrap(), + vec![target.clone()], + ); + let values: ArrayRef = Arc::new(Int64Array::from(vec![Some(1), None, Some(1), None])); + let root: ArrayRef = Arc::new(StructArray::new( + vec![Field::new("value", values.data_type().clone(), true)].into(), + vec![values], + Some(NullBuffer::from(vec![true, false, true, true])), + )); + for start in [0, 2] { + let array = root.slice(start, 2); + if flat { + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "field_u64", + array.data_type().clone(), + true, + )])), + vec![array], + ) + .unwrap(); + creator.update_flat(&batch).await.unwrap(); + } else { + let mut batch = Batch::new( + vec![], + Arc::new(UInt64Vector::from_vec(vec![0, 1])), + Arc::new(UInt64Vector::from_vec(vec![0, 0])), + Arc::new(UInt8Vector::from_vec(vec![1, 1])), + vec![BatchColumn { + column_id: 4, + data: Helper::try_into_vector(array).unwrap(), + }], + ) + .unwrap(); + creator.update(&mut batch).await.unwrap(); + } + } + // JSON paths must not be advertised as ordinary column indexes. + assert_eq!(creator.indexed_column_ids().count(), 0); + let store = mock_object_store(); + let manager = factory.build( + store.clone(), + RegionFilePathFactory::new("table0".into(), PathType::Bare), + ); + let index_id = RegionIndexId::new(RegionFileId::new(metadata.region_id, file_id), 0); + let mut writer = manager.writer(&index_id).await.unwrap(); + let (rows, bytes) = creator.finish(&mut writer).await.unwrap(); + assert_eq!(rows, 4); + assert!(bytes > 0); + writer.finish().await.unwrap(); + let reader = manager.reader(&index_id).await.unwrap(); + let blob = reader + .blob(INDEX_BLOB_TYPE) + .await + .unwrap() + .reader() + .await + .unwrap(); + let reader = InvertedIndexBlobReader::new(blob); + let metas = reader.metadata(None).await.unwrap(); + assert_eq!(metas.total_row_count, 4); + assert_eq!(metas.metas.len(), 1); + let meta = &metas.metas[&target.encode()]; + let fst = reader + .fst( + meta.base_offset + u64::from(meta.relative_fst_offset), + meta.fst_size, + None, + ) + .await + .unwrap(); + let mut encoded = Vec::new(); + IndexValueCodec::encode_nonnull_value( + ValueRef::Int64(1), + &SortField::new(ConcreteDataType::int64_datatype()), + &mut encoded, + ) + .unwrap(); + assert!(fst.get(&encoded).is_some()); + // Null parent and leaf rows must not add values to the FST. + assert_eq!(fst.len(), 1); + assert!(meta.null_bitmap_size > 0); + // Apply the generated index to both matching and non-matching values. + for (value, expected) in [(2, vec![]), (1, vec![0, 1])] { + let mut bytes = Vec::new(); + IndexValueCodec::encode_nonnull_value( + ValueRef::Int64(value), + &SortField::new(ConcreteDataType::int64_datatype()), + &mut bytes, + ) + .unwrap(); + let predicates = BTreeMap::from([( + target.clone(), + vec![Predicate::InList(InListPredicate { + list: BTreeSet::from([bytes]), + })], + )]); + let applier = InvertedIndexApplier::new( + "table0".into(), + PathType::Bare, + store.clone(), + factory.clone(), + predicates, + BTreeMap::from([(target.clone(), ConcreteDataType::int64_datatype())]), + ) + .unwrap(); + let plan = applier.plan_for_sst(&metadata).unwrap().unwrap(); + let output = applier + .apply(index_id, None, &plan.index_applier, None) + .await + .unwrap(); + assert_eq!( + output.matched_segment_ids.iter_ones().collect::>(), + expected + ); + } + } + } + #[tokio::test] async fn test_create_and_query_get_key() { let rows = BTreeSet::from_iter([ diff --git a/src/operator/src/expr_helper.rs b/src/operator/src/expr_helper.rs index 935ce149414..563eb00789e 100644 --- a/src/operator/src/expr_helper.rs +++ b/src/operator/src/expr_helper.rs @@ -805,6 +805,7 @@ fn json_settings_to_proto(settings: JsonSettings) -> Result { path: hint.path, data_type: data_type as i32, datatype_extension, + inverted_index: hint.inverted_index, }) }) .collect::>>()?; @@ -1851,7 +1852,7 @@ SELECT max(c1), min(c2) FROM schema_2.table_2;"; #[test] fn test_to_alter_set_json_settings_expr() { - let sql = "ALTER TABLE monitor MODIFY COLUMN payload JSON2 (service STRING);"; + let sql = "ALTER TABLE monitor MODIFY COLUMN payload JSON2 (service STRING INVERTED INDEX, message STRING);"; let stmt = ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) .unwrap() @@ -1870,8 +1871,10 @@ SELECT max(c1), min(c2) FROM schema_2.table_2;"; assert_eq!("payload", modify.column_name); let settings = modify.settings.as_ref().unwrap(); assert_eq!(Some(100), settings.max_auto_expanded_paths); - assert_eq!(1, settings.type_hints.len()); + assert_eq!(2, settings.type_hints.len()); assert_eq!(["service"], &settings.type_hints[0].path[..]); + assert!(settings.type_hints[0].inverted_index); + assert!(!settings.type_hints[1].inverted_index); let sql = "ALTER TABLE monitor MODIFY COLUMN payload JSON2;"; let stmt = diff --git a/src/sql/src/parsers/create_parser/json.rs b/src/sql/src/parsers/create_parser/json.rs index 0adadc1d193..204f294a82f 100644 --- a/src/sql/src/parsers/create_parser/json.rs +++ b/src/sql/src/parsers/create_parser/json.rs @@ -194,6 +194,16 @@ fn parse_json2_type_hint(parser: &mut Parser<'_>) -> Result { ) } ); + ensure!( + !path + .iter() + .any(|segment| segment == JSON2_REMAINDER_FIELD_NAME), + InvalidSqlSnafu { + msg: format!( + "JSON2 indexed type hint path cannot contain reserved field '{JSON2_REMAINDER_FIELD_NAME}'" + ) + } + ); inverted_index = true; } else if let Token::Word(word) = parser.peek_token().token && word.value.eq_ignore_ascii_case(SKIPPING) diff --git a/src/store-api/src/region_request.rs b/src/store-api/src/region_request.rs index f0650f8d7dc..937fb9d950f 100644 --- a/src/store-api/src/region_request.rs +++ b/src/store-api/src/region_request.rs @@ -1569,7 +1569,7 @@ fn json_settings_from_proto(settings: v1::JsonSettings) -> Result Ok(JsonTypeHint { path: hint.path, data_type, - inverted_index: false, + inverted_index: hint.inverted_index, }) }) .collect::>>()?; @@ -2439,6 +2439,44 @@ mod tests { ); } + #[test] + fn test_json_settings_path_validation_from_proto() { + for (path, valid_path) in [ + (vec!["nested", "leaf"], true), + (vec![""], false), + (vec!["nested", ""], false), + ( + vec!["nested", datatypes::json::JSON2_REMAINDER_FIELD_NAME], + false, + ), + ] { + for inverted_index in [false, true] { + let settings = api::v1::SetJsonSettings { + column_name: "payload".into(), + settings: Some(api::v1::JsonSettings { + type_hints: vec![api::v1::JsonTypeHint { + path: path.iter().map(|part| (*part).to_string()).collect(), + data_type: ColumnDataType::String as i32, + inverted_index, + ..Default::default() + }], + max_auto_expanded_paths: Some(0), + }), + }; + let result = RegionAlterRequest::try_from(AlterRequest { + region_id: 0, + schema_version: 1, + kind: Some(alter_request::Kind::SetJsonSettings(settings)), + }); + assert_eq!( + result.is_ok(), + valid_path, + "path: {path:?}, inverted_index: {inverted_index}" + ); + } + } + } + #[test] fn test_from_proto_alter_request() { RegionAlterRequest::try_from(AlterRequest { @@ -2504,6 +2542,7 @@ mod tests { path: vec!["service".to_string()], data_type: ColumnDataType::String as i32, datatype_extension: None, + inverted_index: true, }], max_auto_expanded_paths: Some(10), }), @@ -2521,6 +2560,7 @@ mod tests { assert_eq!("payload", column_name); assert_eq!(Some(10), settings.max_auto_expanded_paths()); assert_eq!(1, settings.type_hints().len()); + assert!(settings.type_hints()[0].inverted_index); } #[test] diff --git a/tests/cases/standalone/common/types/json/json2_inverted_index.result b/tests/cases/standalone/common/types/json/json2_inverted_index.result new file mode 100644 index 00000000000..c03155e3961 --- /dev/null +++ b/tests/cases/standalone/common/types/json/json2_inverted_index.result @@ -0,0 +1,153 @@ +-- Index declarations round-trip through table metadata and SHOW CREATE TABLE. +CREATE TABLE json2_inverted_index ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 ( + max_auto_expanded_paths = 0, + "service.name" STRING INVERTED INDEX, + service.name STRING INVERTED INDEX, + http.status_code BIGINT INVERTED INDEX, + request.bytes BIGINT UNSIGNED INVERTED INDEX, + duration DOUBLE INVERTED INDEX, + success BOOLEAN INVERTED INDEX, + message STRING + ) +) WITH ('append_mode' = 'true'); + +Affected Rows: 0 + +SHOW CREATE TABLE json2_inverted_index; + ++----------------------+-------------------------------------------------------+ +| Table | Create Table | ++----------------------+-------------------------------------------------------+ +| json2_inverted_index | CREATE TABLE IF NOT EXISTS "json2_inverted_index" ( | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | "attrs" JSON2( | +| | max_auto_expanded_paths = 0, | +| | "service.name" STRING INVERTED INDEX, | +| | "service"."name" STRING INVERTED INDEX, | +| | "http"."status_code" BIGINT INVERTED INDEX, | +| | "request"."bytes" BIGINT UNSIGNED INVERTED INDEX, | +| | "duration" DOUBLE INVERTED INDEX, | +| | "success" BOOLEAN INVERTED INDEX, | +| | "message" STRING | +| | ) NULL, | +| | TIME INDEX ("ts") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | append_mode = 'true' | +| | ) | ++----------------------+-------------------------------------------------------+ + +-- A rejected ALTER must preserve the existing declarations. +ALTER TABLE json2_inverted_index MODIFY COLUMN attrs JSON2 ( + message STRING INVERTED INDEX INVERTED INDEX +); + +Error: 2000(InvalidSyntax), Invalid SQL, error: duplicated INVERTED INDEX option for JSON2 type hint 'message' + +ALTER TABLE json2_inverted_index MODIFY COLUMN attrs JSON2 ( + nested."!__remainder__!" STRING INVERTED INDEX +); + +Error: 2000(InvalidSyntax), Invalid SQL, error: JSON2 indexed type hint path cannot contain reserved field '!__remainder__!' + +SHOW CREATE TABLE json2_inverted_index; + ++----------------------+-------------------------------------------------------+ +| Table | Create Table | ++----------------------+-------------------------------------------------------+ +| json2_inverted_index | CREATE TABLE IF NOT EXISTS "json2_inverted_index" ( | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | "attrs" JSON2( | +| | max_auto_expanded_paths = 0, | +| | "service.name" STRING INVERTED INDEX, | +| | "service"."name" STRING INVERTED INDEX, | +| | "http"."status_code" BIGINT INVERTED INDEX, | +| | "request"."bytes" BIGINT UNSIGNED INVERTED INDEX, | +| | "duration" DOUBLE INVERTED INDEX, | +| | "success" BOOLEAN INVERTED INDEX, | +| | "message" STRING | +| | ) NULL, | +| | TIME INDEX ("ts") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | append_mode = 'true' | +| | ) | ++----------------------+-------------------------------------------------------+ + +DROP TABLE json2_inverted_index; + +Affected Rows: 0 + +-- Reserved field names cannot occur anywhere in an indexed hint path. +CREATE TABLE json2_inverted_invalid ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 ("!__remainder__!" STRING INVERTED INDEX) +); + +Error: 2000(InvalidSyntax), Invalid SQL, error: JSON2 type hint path cannot be rooted at reserved field '!__remainder__!' + +CREATE TABLE json2_inverted_invalid ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 (nested."!__remainder__!" STRING INVERTED INDEX) +); + +Error: 2000(InvalidSyntax), Invalid SQL, error: JSON2 indexed type hint path cannot contain reserved field '!__remainder__!' + +CREATE TABLE json2_inverted_invalid ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 (nested."!__remainder__!".value STRING INVERTED INDEX) +); + +Error: 2000(InvalidSyntax), Invalid SQL, error: JSON2 indexed type hint path cannot contain reserved field '!__remainder__!' + +-- Duplicate and incomplete index declarations are rejected. +CREATE TABLE json2_inverted_invalid ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 (message STRING INVERTED INDEX INVERTED INDEX) +); + +Error: 2000(InvalidSyntax), Invalid SQL, error: duplicated INVERTED INDEX option for JSON2 type hint 'message' + +CREATE TABLE json2_inverted_invalid ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 (message STRING INVERTED) +); + +Error: 2000(InvalidSyntax), Invalid SQL, error: expect INDEX after INVERTED keyword for JSON2 type hint 'message' + +-- Array paths and unsupported hint types are rejected. +CREATE TABLE json2_inverted_invalid ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 (items[0].name STRING INVERTED INDEX) +); + +Error: 2000(InvalidSyntax), Invalid SQL syntax: sql parser error: Expected: a data type name, found: [ at Line: 3, Column: 23 + +CREATE TABLE json2_inverted_invalid ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 (event_time TIMESTAMP INVERTED INDEX) +); + +Error: 2000(InvalidSyntax), Invalid SQL, error: unsupported JSON2 type hint data type: TIMESTAMP; supported types: STRING, BIGINT, BIGINT UNSIGNED, DOUBLE, BOOLEAN; supported aliases: INT64, UINT64, FLOAT64 + +-- JSON hints support only inverted indexes. +CREATE TABLE json2_inverted_invalid ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 (message STRING SKIPPING INDEX) +); + +Error: 2000(InvalidSyntax), Invalid SQL, error: JSON2 type hint SKIPPING INDEX is not supported yet + +CREATE TABLE json2_inverted_invalid ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 (message STRING FULLTEXT INDEX) +); + +Error: 2000(InvalidSyntax), Invalid SQL syntax: sql parser error: Expected: JSON2 type hint option, found: FULLTEXT at Line: 3, Column: 33 + diff --git a/tests/cases/standalone/common/types/json/json2_inverted_index.sql b/tests/cases/standalone/common/types/json/json2_inverted_index.sql new file mode 100644 index 00000000000..866a99a143e --- /dev/null +++ b/tests/cases/standalone/common/types/json/json2_inverted_index.sql @@ -0,0 +1,78 @@ +-- Index declarations round-trip through table metadata and SHOW CREATE TABLE. +CREATE TABLE json2_inverted_index ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 ( + max_auto_expanded_paths = 0, + "service.name" STRING INVERTED INDEX, + service.name STRING INVERTED INDEX, + http.status_code BIGINT INVERTED INDEX, + request.bytes BIGINT UNSIGNED INVERTED INDEX, + duration DOUBLE INVERTED INDEX, + success BOOLEAN INVERTED INDEX, + message STRING + ) +) WITH ('append_mode' = 'true'); + +SHOW CREATE TABLE json2_inverted_index; + +-- A rejected ALTER must preserve the existing declarations. +ALTER TABLE json2_inverted_index MODIFY COLUMN attrs JSON2 ( + message STRING INVERTED INDEX INVERTED INDEX +); + +ALTER TABLE json2_inverted_index MODIFY COLUMN attrs JSON2 ( + nested."!__remainder__!" STRING INVERTED INDEX +); + +SHOW CREATE TABLE json2_inverted_index; + +DROP TABLE json2_inverted_index; + +-- Reserved field names cannot occur anywhere in an indexed hint path. +CREATE TABLE json2_inverted_invalid ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 ("!__remainder__!" STRING INVERTED INDEX) +); + +CREATE TABLE json2_inverted_invalid ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 (nested."!__remainder__!" STRING INVERTED INDEX) +); + +CREATE TABLE json2_inverted_invalid ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 (nested."!__remainder__!".value STRING INVERTED INDEX) +); + +-- Duplicate and incomplete index declarations are rejected. +CREATE TABLE json2_inverted_invalid ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 (message STRING INVERTED INDEX INVERTED INDEX) +); + +CREATE TABLE json2_inverted_invalid ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 (message STRING INVERTED) +); + +-- Array paths and unsupported hint types are rejected. +CREATE TABLE json2_inverted_invalid ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 (items[0].name STRING INVERTED INDEX) +); + +CREATE TABLE json2_inverted_invalid ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 (event_time TIMESTAMP INVERTED INDEX) +); + +-- JSON hints support only inverted indexes. +CREATE TABLE json2_inverted_invalid ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 (message STRING SKIPPING INDEX) +); + +CREATE TABLE json2_inverted_invalid ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 (message STRING FULLTEXT INDEX) +); diff --git a/tests/cases/standalone/common/types/json/json2_inverted_index_query.result b/tests/cases/standalone/common/types/json/json2_inverted_index_query.result new file mode 100644 index 00000000000..1eecb960186 --- /dev/null +++ b/tests/cases/standalone/common/types/json/json2_inverted_index_query.result @@ -0,0 +1,435 @@ +-- Query indexed JSON paths alongside ordinary columns in flat SSTs. +CREATE TABLE json2_index_query ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 ( + max_auto_expanded_paths = 0, + service.name STRING INVERTED INDEX, + "service.name" STRING INVERTED INDEX, + status BIGINT INVERTED INDEX, + success BOOLEAN INVERTED INDEX, + duration DOUBLE INVERTED INDEX, + bytes BIGINT UNSIGNED INVERTED INDEX, + unindexed STRING + ), + ordinary STRING INVERTED INDEX +) WITH ('append_mode' = 'true', 'sst_format' = 'flat', 'index.inverted_index.segment_row_count' = '1'); + +Affected Rows: 0 + +INSERT INTO json2_index_query VALUES + (1, '{"service":{"name":"api"},"service.name":"literal","status":200,"success":true,"duration":1.5,"bytes":18446744073709551615,"unindexed":"x"}', 'a'), + (2, '{"service":{"name":"worker"},"service.name":"api","status":500,"success":false,"duration":2.5,"bytes":42,"unindexed":"y"}', 'b'), + (3, '{}', 'c'), (4, NULL, 'd'); + +Affected Rows: 4 + +ADMIN FLUSH_TABLE('json2_index_query'); + ++----------------------------------------+ +| ADMIN FLUSH_TABLE('json2_index_query') | ++----------------------------------------+ +| 0 | ++----------------------------------------+ + +-- Nested and literal dotted paths resolve to different targets. +SELECT ts FROM json2_index_query WHERE attrs.service.name = 'api' ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | ++-------------------------+ + +SELECT ts FROM json2_index_query WHERE json_get(attrs, '$."service.name"') = 'api' ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.002 | ++-------------------------+ + +-- Integer predicates support equality, reversed ranges, BETWEEN, and IN. +SELECT ts FROM json2_index_query WHERE attrs.status::BIGINT = 200 ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | ++-------------------------+ + +SELECT ts FROM json2_index_query WHERE 400 < attrs.status::BIGINT ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.002 | ++-------------------------+ + +SELECT ts FROM json2_index_query WHERE attrs.status::BIGINT BETWEEN 100 AND 300 ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | ++-------------------------+ + +SELECT ts FROM json2_index_query WHERE attrs.status::BIGINT IN (200, 500) ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | +| 1970-01-01T00:00:00.002 | ++-------------------------+ + +-- OR predicates work within one path and across different paths. +SELECT ts FROM json2_index_query WHERE attrs.service.name = 'api' OR attrs.service.name = 'worker' ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | +| 1970-01-01T00:00:00.002 | ++-------------------------+ + +SELECT ts FROM json2_index_query WHERE attrs.service.name = 'api' OR json_get(attrs, '$."service.name"') = 'api' ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | +| 1970-01-01T00:00:00.002 | ++-------------------------+ + +-- Regex, boolean, floating-point, and unsigned predicates return matching rows. +SELECT ts FROM json2_index_query WHERE attrs.service.name ~ '^a.*[ip]$' ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | ++-------------------------+ + +SELECT ts FROM json2_index_query WHERE attrs.success::BOOLEAN = true ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | ++-------------------------+ + +SELECT ts FROM json2_index_query WHERE attrs.duration::DOUBLE > 2.0 ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.002 | ++-------------------------+ + +SELECT ts FROM json2_index_query WHERE attrs.bytes::BIGINT UNSIGNED = 18446744073709551615::BIGINT UNSIGNED ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | ++-------------------------+ + +-- Ordinary and JSON index predicates can be combined. +SELECT ts FROM json2_index_query WHERE ordinary = 'b' AND attrs.status::BIGINT > 400 ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.002 | ++-------------------------+ + +-- Missing paths and null roots remain visible to IS NULL. +SELECT ts FROM json2_index_query WHERE attrs.service.name IS NULL ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.003 | +| 1970-01-01T00:00:00.004 | ++-------------------------+ + +-- Unindexed hints remain queryable. +SELECT ts FROM json2_index_query WHERE attrs.unindexed = 'x' ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | ++-------------------------+ + +-- Implicit comparison and explicit conversion preserve query results. +SELECT ts FROM json2_index_query WHERE attrs.status = '200' ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | ++-------------------------+ + +SELECT ts FROM json2_index_query WHERE attrs.duration::BIGINT = 1 ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | ++-------------------------+ + +-- String equality prunes all four rows; keep the pruning counter in the plan. +-- SQLNESS REPLACE "metrics_per_partition":.*("rows_inverted_filtered":\d+).*metrics= "metrics_per_partition":{$1} metrics= +-- SQLNESS REPLACE "metrics_per_partition":\s*\[.*metrics= "metrics_per_partition":REDACTED metrics= +-- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED +-- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" +-- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED +-- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED +-- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE ,\s"dyn_filters":\s\[.* REDACTED +-- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED +-- SQLNESS REPLACE "max_part_range":\s*\{[^}]*\} "max_part_range":REDACTED +EXPLAIN ANALYZE VERBOSE SELECT ts FROM json2_index_query WHERE attrs.service.name = 'absent'; + ++-+-+-+ +| stage | node | plan_| ++-+-+-+ +| 0_| 0_|_CooperativeExec metrics=REDACTED_| +|_|_|_MergeScanExec: REDACTED +|_|_|_| +| 1_| 0_|_FilterExec: json_get(attrs@1, $.service.name, NULL) = absent, projection=[ts@0] metrics=REDACTED_| +|_|_|_CooperativeExec metrics=REDACTED_| +|_|_|_UnorderedScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["ts", "attrs"], "filters": ["json_get(attrs, Utf8(\"$.service.name\"), Utf8View(NULL)) = Utf8View(\"absent\")"], "time_range": "TimestampRange{{[#,#)}}", "file_summary": {"files":1,"rows":4,"row_groups":1,"size":REDACTED,"index_size":REDACTED}, "max_part_range":REDACTED, "files": [{"file_id":"REDACTED","time_range_start":"REDACTED","time_range_end":"REDACTED","rows":4,"size":REDACTED,"index_size":REDACTED}], "flat_format":REDACTED, "metrics_per_partition":{"rows_inverted_filtered":4} metrics=REDACTED | +|_|_|_| +|_|_| Total rows: REDACTED_| ++-+-+-+ + +-- A numeric range prunes all four rows. +-- SQLNESS REPLACE "metrics_per_partition":.*("rows_inverted_filtered":\d+).*metrics= "metrics_per_partition":{$1} metrics= +-- SQLNESS REPLACE "metrics_per_partition":\s*\[.*metrics= "metrics_per_partition":REDACTED metrics= +-- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED +-- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" +-- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED +-- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED +-- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE ,\s"dyn_filters":\s\[.* REDACTED +-- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED +-- SQLNESS REPLACE "max_part_range":\s*\{[^}]*\} "max_part_range":REDACTED +EXPLAIN ANALYZE VERBOSE SELECT ts FROM json2_index_query WHERE attrs.status::BIGINT > 600; + ++-+-+-+ +| stage | node | plan_| ++-+-+-+ +| 0_| 0_|_CooperativeExec metrics=REDACTED_| +|_|_|_MergeScanExec: REDACTED +|_|_|_| +| 1_| 0_|_FilterExec: json_get(attrs@1, $.status, NULL) > 600, projection=[ts@0] metrics=REDACTED_| +|_|_|_CooperativeExec metrics=REDACTED_| +|_|_|_UnorderedScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["ts", "attrs"], "filters": ["json_get(attrs, Utf8(\"$.status\"), Int64(NULL)) > Int64(600)"], "time_range": "TimestampRange{{[#,#)}}", "file_summary": {"files":1,"rows":4,"row_groups":1,"size":REDACTED,"index_size":REDACTED}, "max_part_range":REDACTED, "files": [{"file_id":"REDACTED","time_range_start":"REDACTED","time_range_end":"REDACTED","rows":4,"size":REDACTED,"index_size":REDACTED}], "flat_format":REDACTED, "metrics_per_partition":{"rows_inverted_filtered":4} metrics=REDACTED | +|_|_|_| +|_|_| Total rows: REDACTED_| ++-+-+-+ + +-- A regex predicate prunes all four rows. +-- SQLNESS REPLACE "metrics_per_partition":.*("rows_inverted_filtered":\d+).*metrics= "metrics_per_partition":{$1} metrics= +-- SQLNESS REPLACE "metrics_per_partition":\s*\[.*metrics= "metrics_per_partition":REDACTED metrics= +-- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED +-- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" +-- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED +-- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED +-- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE ,\s"dyn_filters":\s\[.* REDACTED +-- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED +-- SQLNESS REPLACE "max_part_range":\s*\{[^}]*\} "max_part_range":REDACTED +EXPLAIN ANALYZE VERBOSE SELECT ts FROM json2_index_query WHERE attrs.service.name ~ '^z.*[ip]$'; + ++-+-+-+ +| stage | node | plan_| ++-+-+-+ +| 0_| 0_|_CooperativeExec metrics=REDACTED_| +|_|_|_MergeScanExec: REDACTED +|_|_|_| +| 1_| 0_|_FilterExec: json_get(attrs@1, $.service.name, NULL) ~ ^z.*[ip]$, projection=[ts@0] metrics=REDACTED_| +|_|_|_CooperativeExec metrics=REDACTED_| +|_|_|_UnorderedScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["ts", "attrs"], "filters": ["json_get(attrs, Utf8(\"$.service.name\"), Utf8View(NULL)) ~ Utf8View(\"^z.*[ip]$\")"], "time_range": "TimestampRange{{[#,#)}}", "file_summary": {"files":1,"rows":4,"row_groups":1,"size":REDACTED,"index_size":REDACTED}, "max_part_range":REDACTED, "files": [{"file_id":"REDACTED","time_range_start":"REDACTED","time_range_end":"REDACTED","rows":4,"size":REDACTED,"index_size":REDACTED}], "flat_format":REDACTED, "metrics_per_partition":{"rows_inverted_filtered":4} metrics=REDACTED | +|_|_|_| +|_|_| Total rows: REDACTED_| ++-+-+-+ + +-- Matching explicit types use the hint index. +SELECT ts FROM json2_index_query WHERE json_get(attrs, '$.status', NULL::BIGINT) = 200 ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | ++-------------------------+ + +-- SQLNESS REPLACE "metrics_per_partition":.*("rows_inverted_filtered":\d+).*metrics= "metrics_per_partition":{$1} metrics= +-- SQLNESS REPLACE "metrics_per_partition":\s*\[.*metrics= "metrics_per_partition":REDACTED metrics= +-- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED +-- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" +-- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED +-- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED +-- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE ,\s"dyn_filters":\s\[.* REDACTED +-- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED +-- SQLNESS REPLACE "max_part_range":\s*\{[^}]*\} "max_part_range":REDACTED +EXPLAIN ANALYZE VERBOSE SELECT ts FROM json2_index_query WHERE json_get(attrs, '$.status', NULL::BIGINT) > 600; + ++-+-+-+ +| stage | node | plan_| ++-+-+-+ +| 0_| 0_|_CooperativeExec metrics=REDACTED_| +|_|_|_MergeScanExec: REDACTED +|_|_|_| +| 1_| 0_|_FilterExec: json_get(attrs@1, $.status, NULL) > 600, projection=[ts@0] metrics=REDACTED_| +|_|_|_CooperativeExec metrics=REDACTED_| +|_|_|_UnorderedScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["ts", "attrs"], "filters": ["json_get(attrs, Utf8(\"$.status\"), Int64(NULL)) > Int64(600)"], "time_range": "TimestampRange{{[#,#)}}", "file_summary": {"files":1,"rows":4,"row_groups":1,"size":REDACTED,"index_size":REDACTED}, "max_part_range":REDACTED, "files": [{"file_id":"REDACTED","time_range_start":"REDACTED","time_range_end":"REDACTED","rows":4,"size":REDACTED,"index_size":REDACTED}], "flat_format":REDACTED, "metrics_per_partition":{"rows_inverted_filtered":4} metrics=REDACTED | +|_|_|_| +|_|_| Total rows: REDACTED_| ++-+-+-+ + +-- A different explicit type scans rows without using the hint index. +SELECT ts FROM json2_index_query WHERE json_get(attrs, '$.duration', NULL::BIGINT) = 1 ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | ++-------------------------+ + +-- SQLNESS REPLACE "metrics_per_partition":.*("rows_inverted_filtered":\d+).*metrics= "metrics_per_partition":{$1} metrics= +-- SQLNESS REPLACE "metrics_per_partition":\s*\[.*metrics= "metrics_per_partition":REDACTED metrics= +-- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED +-- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" +-- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED +-- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED +-- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE ,\s"dyn_filters":\s\[.* REDACTED +-- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED +-- SQLNESS REPLACE "max_part_range":\s*\{[^}]*\} "max_part_range":REDACTED +EXPLAIN ANALYZE VERBOSE SELECT ts FROM json2_index_query WHERE json_get(attrs, '$.duration', NULL::BIGINT) > 600; + ++-+-+-+ +| stage | node | plan_| ++-+-+-+ +| 0_| 0_|_CooperativeExec metrics=REDACTED_| +|_|_|_MergeScanExec: REDACTED +|_|_|_| +| 1_| 0_|_FilterExec: json_get(attrs@1, $.duration, NULL) > 600, projection=[ts@0] metrics=REDACTED_| +|_|_|_CooperativeExec metrics=REDACTED_| +|_|_|_UnorderedScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["ts", "attrs"], "filters": ["json_get(attrs, Utf8(\"$.duration\"), Int64(NULL)) > Int64(600)"], "time_range": "TimestampRange{{[#,#)}}", "file_summary": {"files":1,"rows":4,"row_groups":1,"size":REDACTED,"index_size":REDACTED}, "max_part_range":REDACTED, "files": [{"file_id":"REDACTED","time_range_start":"REDACTED","time_range_end":"REDACTED","rows":4,"size":REDACTED,"index_size":REDACTED}], "flat_format":REDACTED, "metrics_per_partition":REDACTED metrics=REDACTED | +|_|_|_| +|_|_| Total rows: REDACTED_| ++-+-+-+ + +-- Queries retain correct results after compaction. +INSERT INTO json2_index_query VALUES (5, '{"service":{"name":"api"},"status":201}', 'a'); + +Affected Rows: 1 + +ADMIN FLUSH_TABLE('json2_index_query'); + ++----------------------------------------+ +| ADMIN FLUSH_TABLE('json2_index_query') | ++----------------------------------------+ +| 0 | ++----------------------------------------+ + +ADMIN COMPACT_TABLE('json2_index_query', 'swcs', '86400'); + ++-----------------------------------------------------------+ +| ADMIN COMPACT_TABLE('json2_index_query', 'swcs', '86400') | ++-----------------------------------------------------------+ +| 0 | ++-----------------------------------------------------------+ + +SELECT ts FROM json2_index_query WHERE attrs.service.name = 'api' AND attrs.status::BIGINT < 300 ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | +| 1970-01-01T00:00:00.005 | ++-------------------------+ + +DROP TABLE json2_index_query; + +Affected Rows: 0 + +-- JSON-only targets also work with the default SST format. +CREATE TABLE json2_index_query_only ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 (value STRING INVERTED INDEX) +) WITH ('append_mode' = 'true', 'index.inverted_index.segment_row_count' = '1'); + +Affected Rows: 0 + +INSERT INTO json2_index_query_only VALUES (1, '{"value":"yes"}'), (2, '{"value":"no"}'), (3, NULL); + +Affected Rows: 3 + +ADMIN FLUSH_TABLE('json2_index_query_only'); + ++---------------------------------------------+ +| ADMIN FLUSH_TABLE('json2_index_query_only') | ++---------------------------------------------+ +| 0 | ++---------------------------------------------+ + +SELECT ts FROM json2_index_query_only WHERE attrs.value = 'yes'; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | ++-------------------------+ + +SELECT ts FROM json2_index_query_only WHERE attrs.value = 'missing'; + ++----+ +| ts | ++----+ ++----+ + +DROP TABLE json2_index_query_only; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/types/json/json2_inverted_index_query.sql b/tests/cases/standalone/common/types/json/json2_inverted_index_query.sql new file mode 100644 index 00000000000..0a241afbf8f --- /dev/null +++ b/tests/cases/standalone/common/types/json/json2_inverted_index_query.sql @@ -0,0 +1,195 @@ +-- Query indexed JSON paths alongside ordinary columns in flat SSTs. +CREATE TABLE json2_index_query ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 ( + max_auto_expanded_paths = 0, + service.name STRING INVERTED INDEX, + "service.name" STRING INVERTED INDEX, + status BIGINT INVERTED INDEX, + success BOOLEAN INVERTED INDEX, + duration DOUBLE INVERTED INDEX, + bytes BIGINT UNSIGNED INVERTED INDEX, + unindexed STRING + ), + ordinary STRING INVERTED INDEX +) WITH ('append_mode' = 'true', 'sst_format' = 'flat', 'index.inverted_index.segment_row_count' = '1'); + +INSERT INTO json2_index_query VALUES + (1, '{"service":{"name":"api"},"service.name":"literal","status":200,"success":true,"duration":1.5,"bytes":18446744073709551615,"unindexed":"x"}', 'a'), + (2, '{"service":{"name":"worker"},"service.name":"api","status":500,"success":false,"duration":2.5,"bytes":42,"unindexed":"y"}', 'b'), + (3, '{}', 'c'), (4, NULL, 'd'); + +ADMIN FLUSH_TABLE('json2_index_query'); + +-- Nested and literal dotted paths resolve to different targets. +SELECT ts FROM json2_index_query WHERE attrs.service.name = 'api' ORDER BY ts; + +SELECT ts FROM json2_index_query WHERE json_get(attrs, '$."service.name"') = 'api' ORDER BY ts; + +-- Integer predicates support equality, reversed ranges, BETWEEN, and IN. +SELECT ts FROM json2_index_query WHERE attrs.status::BIGINT = 200 ORDER BY ts; + +SELECT ts FROM json2_index_query WHERE 400 < attrs.status::BIGINT ORDER BY ts; + +SELECT ts FROM json2_index_query WHERE attrs.status::BIGINT BETWEEN 100 AND 300 ORDER BY ts; + +SELECT ts FROM json2_index_query WHERE attrs.status::BIGINT IN (200, 500) ORDER BY ts; + +-- OR predicates work within one path and across different paths. +SELECT ts FROM json2_index_query WHERE attrs.service.name = 'api' OR attrs.service.name = 'worker' ORDER BY ts; + +SELECT ts FROM json2_index_query WHERE attrs.service.name = 'api' OR json_get(attrs, '$."service.name"') = 'api' ORDER BY ts; + +-- Regex, boolean, floating-point, and unsigned predicates return matching rows. +SELECT ts FROM json2_index_query WHERE attrs.service.name ~ '^a.*[ip]$' ORDER BY ts; + +SELECT ts FROM json2_index_query WHERE attrs.success::BOOLEAN = true ORDER BY ts; + +SELECT ts FROM json2_index_query WHERE attrs.duration::DOUBLE > 2.0 ORDER BY ts; + +SELECT ts FROM json2_index_query WHERE attrs.bytes::BIGINT UNSIGNED = 18446744073709551615::BIGINT UNSIGNED ORDER BY ts; + +-- Ordinary and JSON index predicates can be combined. +SELECT ts FROM json2_index_query WHERE ordinary = 'b' AND attrs.status::BIGINT > 400 ORDER BY ts; + +-- Missing paths and null roots remain visible to IS NULL. +SELECT ts FROM json2_index_query WHERE attrs.service.name IS NULL ORDER BY ts; + +-- Unindexed hints remain queryable. +SELECT ts FROM json2_index_query WHERE attrs.unindexed = 'x' ORDER BY ts; + +-- Implicit comparison and explicit conversion preserve query results. +SELECT ts FROM json2_index_query WHERE attrs.status = '200' ORDER BY ts; + +SELECT ts FROM json2_index_query WHERE attrs.duration::BIGINT = 1 ORDER BY ts; + +-- String equality prunes all four rows; keep the pruning counter in the plan. +-- SQLNESS REPLACE "metrics_per_partition":.*("rows_inverted_filtered":\d+).*metrics= "metrics_per_partition":{$1} metrics= +-- SQLNESS REPLACE "metrics_per_partition":\s*\[.*metrics= "metrics_per_partition":REDACTED metrics= +-- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED +-- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" +-- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED +-- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED +-- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE ,\s"dyn_filters":\s\[.* REDACTED +-- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED +-- SQLNESS REPLACE "max_part_range":\s*\{[^}]*\} "max_part_range":REDACTED +EXPLAIN ANALYZE VERBOSE SELECT ts FROM json2_index_query WHERE attrs.service.name = 'absent'; + +-- A numeric range prunes all four rows. +-- SQLNESS REPLACE "metrics_per_partition":.*("rows_inverted_filtered":\d+).*metrics= "metrics_per_partition":{$1} metrics= +-- SQLNESS REPLACE "metrics_per_partition":\s*\[.*metrics= "metrics_per_partition":REDACTED metrics= +-- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED +-- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" +-- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED +-- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED +-- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE ,\s"dyn_filters":\s\[.* REDACTED +-- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED +-- SQLNESS REPLACE "max_part_range":\s*\{[^}]*\} "max_part_range":REDACTED +EXPLAIN ANALYZE VERBOSE SELECT ts FROM json2_index_query WHERE attrs.status::BIGINT > 600; + +-- A regex predicate prunes all four rows. +-- SQLNESS REPLACE "metrics_per_partition":.*("rows_inverted_filtered":\d+).*metrics= "metrics_per_partition":{$1} metrics= +-- SQLNESS REPLACE "metrics_per_partition":\s*\[.*metrics= "metrics_per_partition":REDACTED metrics= +-- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED +-- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" +-- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED +-- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED +-- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE ,\s"dyn_filters":\s\[.* REDACTED +-- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED +-- SQLNESS REPLACE "max_part_range":\s*\{[^}]*\} "max_part_range":REDACTED +EXPLAIN ANALYZE VERBOSE SELECT ts FROM json2_index_query WHERE attrs.service.name ~ '^z.*[ip]$'; + +-- Matching explicit types use the hint index. +SELECT ts FROM json2_index_query WHERE json_get(attrs, '$.status', NULL::BIGINT) = 200 ORDER BY ts; + +-- SQLNESS REPLACE "metrics_per_partition":.*("rows_inverted_filtered":\d+).*metrics= "metrics_per_partition":{$1} metrics= +-- SQLNESS REPLACE "metrics_per_partition":\s*\[.*metrics= "metrics_per_partition":REDACTED metrics= +-- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED +-- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" +-- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED +-- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED +-- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE ,\s"dyn_filters":\s\[.* REDACTED +-- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED +-- SQLNESS REPLACE "max_part_range":\s*\{[^}]*\} "max_part_range":REDACTED +EXPLAIN ANALYZE VERBOSE SELECT ts FROM json2_index_query WHERE json_get(attrs, '$.status', NULL::BIGINT) > 600; + +-- A different explicit type scans rows without using the hint index. +SELECT ts FROM json2_index_query WHERE json_get(attrs, '$.duration', NULL::BIGINT) = 1 ORDER BY ts; + +-- SQLNESS REPLACE "metrics_per_partition":.*("rows_inverted_filtered":\d+).*metrics= "metrics_per_partition":{$1} metrics= +-- SQLNESS REPLACE "metrics_per_partition":\s*\[.*metrics= "metrics_per_partition":REDACTED metrics= +-- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED +-- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" +-- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED +-- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED +-- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE ,\s"dyn_filters":\s\[.* REDACTED +-- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED +-- SQLNESS REPLACE "max_part_range":\s*\{[^}]*\} "max_part_range":REDACTED +EXPLAIN ANALYZE VERBOSE SELECT ts FROM json2_index_query WHERE json_get(attrs, '$.duration', NULL::BIGINT) > 600; + +-- Queries retain correct results after compaction. +INSERT INTO json2_index_query VALUES (5, '{"service":{"name":"api"},"status":201}', 'a'); + +ADMIN FLUSH_TABLE('json2_index_query'); + +ADMIN COMPACT_TABLE('json2_index_query', 'swcs', '86400'); + +SELECT ts FROM json2_index_query WHERE attrs.service.name = 'api' AND attrs.status::BIGINT < 300 ORDER BY ts; + +DROP TABLE json2_index_query; + +-- JSON-only targets also work with the default SST format. +CREATE TABLE json2_index_query_only ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 (value STRING INVERTED INDEX) +) WITH ('append_mode' = 'true', 'index.inverted_index.segment_row_count' = '1'); + +INSERT INTO json2_index_query_only VALUES (1, '{"value":"yes"}'), (2, '{"value":"no"}'), (3, NULL); + +ADMIN FLUSH_TABLE('json2_index_query_only'); + +SELECT ts FROM json2_index_query_only WHERE attrs.value = 'yes'; + +SELECT ts FROM json2_index_query_only WHERE attrs.value = 'missing'; + +DROP TABLE json2_index_query_only; diff --git a/tests/cases/standalone/common/types/json/json2_inverted_index_schema.result b/tests/cases/standalone/common/types/json/json2_inverted_index_schema.result new file mode 100644 index 00000000000..f01b1db7249 --- /dev/null +++ b/tests/cases/standalone/common/types/json/json2_inverted_index_schema.result @@ -0,0 +1,723 @@ +-- Enabling a hint index must keep rows from SSTs with only an ordinary index. +CREATE TABLE json2_index_evolution ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 (max_auto_expanded_paths = 0, value BIGINT), + ordinary STRING INVERTED INDEX +) WITH ('append_mode' = 'true', 'sst_format' = 'flat', 'index.inverted_index.segment_row_count' = '1'); + +Affected Rows: 0 + +INSERT INTO json2_index_evolution VALUES (1, '{"value":42}', 'a'); + +Affected Rows: 1 + +ADMIN FLUSH_TABLE('json2_index_evolution'); + ++--------------------------------------------+ +| ADMIN FLUSH_TABLE('json2_index_evolution') | ++--------------------------------------------+ +| 0 | ++--------------------------------------------+ + +ALTER TABLE json2_index_evolution MODIFY COLUMN attrs JSON2 (max_auto_expanded_paths = 0, value BIGINT INVERTED INDEX); + +Affected Rows: 0 + +SHOW CREATE TABLE json2_index_evolution; + ++-----------------------+------------------------------------------------------+ +| Table | Create Table | ++-----------------------+------------------------------------------------------+ +| json2_index_evolution | CREATE TABLE IF NOT EXISTS "json2_index_evolution" ( | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | "attrs" JSON2( | +| | max_auto_expanded_paths = 0, | +| | "value" BIGINT INVERTED INDEX | +| | ) NULL, | +| | "ordinary" STRING NULL INVERTED INDEX, | +| | TIME INDEX ("ts") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | append_mode = 'true', | +| | 'index.inverted_index.segment_row_count' = '1', | +| | sst_format = 'flat' | +| | ) | ++-----------------------+------------------------------------------------------+ + +-- Two-argument access uses the enabled BIGINT hint for old SSTs. +SELECT ts, json_get(attrs, '$.value') AS value FROM json2_index_evolution WHERE json_get(attrs, '$.value') = 42 ORDER BY ts; + ++-------------------------+-------+ +| ts | value | ++-------------------------+-------+ +| 1970-01-01T00:00:00.001 | 42 | ++-------------------------+-------+ + +SELECT ts FROM json2_index_evolution WHERE json_get(attrs, '$.value', NULL::BIGINT) = 42 ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | ++-------------------------+ + +INSERT INTO json2_index_evolution VALUES (2, '{"value":42}', 'b'); + +Affected Rows: 1 + +ADMIN FLUSH_TABLE('json2_index_evolution'); + ++--------------------------------------------+ +| ADMIN FLUSH_TABLE('json2_index_evolution') | ++--------------------------------------------+ +| 0 | ++--------------------------------------------+ + +-- Two-argument access returns matching rows from both old and new SSTs. +SELECT ts, json_get(attrs, '$.value') AS value FROM json2_index_evolution WHERE json_get(attrs, '$.value') = 42 ORDER BY ts; + ++-------------------------+-------+ +| ts | value | ++-------------------------+-------+ +| 1970-01-01T00:00:00.001 | 42 | +| 1970-01-01T00:00:00.002 | 42 | ++-------------------------+-------+ + +SELECT ts FROM json2_index_evolution WHERE json_get(attrs, '$.value', NULL::BIGINT) = 42 ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | +| 1970-01-01T00:00:00.002 | ++-------------------------+ + +-- Only the new SST has a usable BIGINT index; it filters one row. +-- SQLNESS REPLACE "metrics_per_partition":.*("rows_inverted_filtered":\d+).*metrics= "metrics_per_partition":{$1} metrics= +-- SQLNESS REPLACE "metrics_per_partition":\s*\[.*metrics= "metrics_per_partition":REDACTED metrics= +-- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED +-- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" +-- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED +-- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED +-- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE ,\s"dyn_filters":\s\[.* REDACTED +-- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED +-- SQLNESS REPLACE "max_part_range":\s*\{[^}]*\} "max_part_range":REDACTED +-- SQLNESS REPLACE "files":\s*\[.*\],\s*"flat_format" "files":REDACTED,"flat_format" +EXPLAIN ANALYZE VERBOSE SELECT ts FROM json2_index_evolution WHERE json_get(attrs, '$.value') > 600; + ++-+-+-+ +| stage | node | plan_| ++-+-+-+ +| 0_| 0_|_CooperativeExec metrics=REDACTED_| +|_|_|_MergeScanExec: REDACTED +|_|_|_| +| 1_| 0_|_FilterExec: json_get(attrs@1, $.value, NULL) > 600, projection=[ts@0] metrics=REDACTED_| +|_|_|_CooperativeExec metrics=REDACTED_| +|_|_|_UnorderedScan: region=REDACTED, {"partition_count":{"count":2, "mem_ranges":0, "files":2, "file_ranges":2}, "projection": ["ts", "attrs"], "filters": ["json_get(attrs, Utf8(\"$.value\"), Int64(NULL)) > Int64(600)"], "time_range": "TimestampRange{{[#,#)}}", "file_summary": {"files":2,"rows":2,"row_groups":2,"size":REDACTED,"index_size":REDACTED}, "max_part_range":REDACTED, "files":REDACTED,"flat_format":REDACTED, "metrics_per_partition":{"rows_inverted_filtered":1} metrics=REDACTED | +|_|_|_| +|_|_| Total rows: REDACTED_| ++-+-+-+ + +-- Changing the hint type must not use old indexes with a different value encoding. +ALTER TABLE json2_index_evolution MODIFY COLUMN attrs JSON2 (max_auto_expanded_paths = 0, value DOUBLE INVERTED INDEX); + +Affected Rows: 0 + +SHOW CREATE TABLE json2_index_evolution; + ++-----------------------+------------------------------------------------------+ +| Table | Create Table | ++-----------------------+------------------------------------------------------+ +| json2_index_evolution | CREATE TABLE IF NOT EXISTS "json2_index_evolution" ( | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | "attrs" JSON2( | +| | max_auto_expanded_paths = 0, | +| | "value" DOUBLE INVERTED INDEX | +| | ) NULL, | +| | "ordinary" STRING NULL INVERTED INDEX, | +| | TIME INDEX ("ts") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | append_mode = 'true', | +| | 'index.inverted_index.segment_row_count' = '1', | +| | sst_format = 'flat' | +| | ) | ++-----------------------+------------------------------------------------------+ + +INSERT INTO json2_index_evolution VALUES (3, '{"value":42.5}', 'c'); + +Affected Rows: 1 + +ADMIN FLUSH_TABLE('json2_index_evolution'); + ++--------------------------------------------+ +| ADMIN FLUSH_TABLE('json2_index_evolution') | ++--------------------------------------------+ +| 0 | ++--------------------------------------------+ + +SELECT target_type, target_json, COUNT(*) AS index_count +FROM information_schema.ssts_index_meta +WHERE table_id = (SELECT table_id FROM information_schema.tables WHERE table_name = 'json2_index_evolution') +GROUP BY target_type, target_json ORDER BY target_type, target_json; + ++-------------+---------------------------------------------+-------------+ +| target_type | target_json | index_count | ++-------------+---------------------------------------------+-------------+ +| column | {"column":2} | 3 | +| json_path | {"json_path":{"column":1,"path":["value"]}} | 2 | ++-------------+---------------------------------------------+-------------+ + +-- The DOUBLE hint matches old integer values and new fractional values. +SELECT ts, json_get(attrs, '$.value') AS value FROM json2_index_evolution WHERE json_get(attrs, '$.value') = 42.0 ORDER BY ts; + ++-------------------------+-------+ +| ts | value | ++-------------------------+-------+ +| 1970-01-01T00:00:00.001 | 42.0 | +| 1970-01-01T00:00:00.002 | 42.0 | ++-------------------------+-------+ + +SELECT ts, json_get(attrs, '$.value') AS value FROM json2_index_evolution WHERE json_get(attrs, '$.value') = 42.5 ORDER BY ts; + ++-------------------------+-------+ +| ts | value | ++-------------------------+-------+ +| 1970-01-01T00:00:00.003 | 42.5 | ++-------------------------+-------+ + +-- Only the new DOUBLE index filters a row; older SSTs fall back to scanning. +-- SQLNESS REPLACE "metrics_per_partition":.*("rows_inverted_filtered":\d+).*metrics= "metrics_per_partition":{$1} metrics= +-- SQLNESS REPLACE "metrics_per_partition":\s*\[.*metrics= "metrics_per_partition":REDACTED metrics= +-- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED +-- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" +-- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED +-- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED +-- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE ,\s"dyn_filters":\s\[.* REDACTED +-- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED +-- SQLNESS REPLACE "max_part_range":\s*\{[^}]*\} "max_part_range":REDACTED +-- SQLNESS REPLACE "files":\s*\[.*\],\s*"flat_format" "files":REDACTED,"flat_format" +EXPLAIN ANALYZE VERBOSE SELECT ts FROM json2_index_evolution WHERE json_get(attrs, '$.value') = 42.0; + ++-+-+-+ +| stage | node | plan_| ++-+-+-+ +| 0_| 0_|_CooperativeExec metrics=REDACTED_| +|_|_|_MergeScanExec: REDACTED +|_|_|_| +| 1_| 0_|_FilterExec: json_get(attrs@1, $.value, NULL) = 42, projection=[ts@0] metrics=REDACTED_| +|_|_|_CooperativeExec metrics=REDACTED_| +|_|_|_UnorderedScan: region=REDACTED, {"partition_count":{"count":3, "mem_ranges":0, "files":3, "file_ranges":3}, "projection": ["ts", "attrs"], "filters": ["json_get(attrs, Utf8(\"$.value\"), Float64(NULL)) = Float64(42)"], "time_range": "TimestampRange{{[#,#)}}", "file_summary": {"files":3,"rows":3,"row_groups":3,"size":REDACTED,"index_size":REDACTED}, "max_part_range":REDACTED, "files":REDACTED,"flat_format":REDACTED, "metrics_per_partition":{"rows_inverted_filtered":1} metrics=REDACTED | +|_|_|_| +|_|_| Total rows: REDACTED_| ++-+-+-+ + +-- Restart preserves the altered hint, query results, and index pruning. +-- SQLNESS ARG restart=true +SHOW CREATE TABLE json2_index_evolution; + ++-----------------------+------------------------------------------------------+ +| Table | Create Table | ++-----------------------+------------------------------------------------------+ +| json2_index_evolution | CREATE TABLE IF NOT EXISTS "json2_index_evolution" ( | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | "attrs" JSON2( | +| | max_auto_expanded_paths = 0, | +| | "value" DOUBLE INVERTED INDEX | +| | ) NULL, | +| | "ordinary" STRING NULL INVERTED INDEX, | +| | TIME INDEX ("ts") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | append_mode = 'true', | +| | 'index.inverted_index.segment_row_count' = '1', | +| | sst_format = 'flat' | +| | ) | ++-----------------------+------------------------------------------------------+ + +SELECT ts, json_get(attrs, '$.value') AS value FROM json2_index_evolution WHERE json_get(attrs, '$.value') = 42.0 ORDER BY ts; + ++-------------------------+-------+ +| ts | value | ++-------------------------+-------+ +| 1970-01-01T00:00:00.001 | 42.0 | +| 1970-01-01T00:00:00.002 | 42.0 | ++-------------------------+-------+ + +SELECT ts, json_get(attrs, '$.value') AS value FROM json2_index_evolution WHERE json_get(attrs, '$.value') = 42.5 ORDER BY ts; + ++-------------------------+-------+ +| ts | value | ++-------------------------+-------+ +| 1970-01-01T00:00:00.003 | 42.5 | ++-------------------------+-------+ + +-- SQLNESS REPLACE "metrics_per_partition":.*("rows_inverted_filtered":\d+).*metrics= "metrics_per_partition":{$1} metrics= +-- SQLNESS REPLACE "metrics_per_partition":\s*\[.*metrics= "metrics_per_partition":REDACTED metrics= +-- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED +-- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" +-- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED +-- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED +-- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE ,\s"dyn_filters":\s\[.* REDACTED +-- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED +-- SQLNESS REPLACE "max_part_range":\s*\{[^}]*\} "max_part_range":REDACTED +-- SQLNESS REPLACE "files":\s*\[.*\],\s*"flat_format" "files":REDACTED,"flat_format" +EXPLAIN ANALYZE VERBOSE SELECT ts FROM json2_index_evolution WHERE json_get(attrs, '$.value') = 42.0; + ++-+-+-+ +| stage | node | plan_| ++-+-+-+ +| 0_| 0_|_CooperativeExec metrics=REDACTED_| +|_|_|_MergeScanExec: REDACTED +|_|_|_| +| 1_| 0_|_FilterExec: json_get(attrs@1, $.value, NULL) = 42, projection=[ts@0] metrics=REDACTED_| +|_|_|_CooperativeExec metrics=REDACTED_| +|_|_|_UnorderedScan: region=REDACTED, {"partition_count":{"count":3, "mem_ranges":0, "files":3, "file_ranges":3}, "projection": ["ts", "attrs"], "filters": ["json_get(attrs, Utf8(\"$.value\"), Float64(NULL)) = Float64(42)"], "time_range": "TimestampRange{{[#,#)}}", "file_summary": {"files":3,"rows":3,"row_groups":3,"size":REDACTED,"index_size":REDACTED}, "max_part_range":REDACTED, "files":REDACTED,"flat_format":REDACTED, "metrics_per_partition":{"rows_inverted_filtered":1} metrics=REDACTED | +|_|_|_| +|_|_| Total rows: REDACTED_| ++-+-+-+ + +-- Compaction replaces old indexes with DOUBLE indexes for all three rows. +ADMIN COMPACT_TABLE('json2_index_evolution', 'swcs', '86400'); + ++---------------------------------------------------------------+ +| ADMIN COMPACT_TABLE('json2_index_evolution', 'swcs', '86400') | ++---------------------------------------------------------------+ +| 0 | ++---------------------------------------------------------------+ + +SELECT ts, json_get(attrs, '$.value') AS value FROM json2_index_evolution WHERE json_get(attrs, '$.value') = 42.0 ORDER BY ts; + ++-------------------------+-------+ +| ts | value | ++-------------------------+-------+ +| 1970-01-01T00:00:00.001 | 42.0 | +| 1970-01-01T00:00:00.002 | 42.0 | ++-------------------------+-------+ + +SELECT ts, json_get(attrs, '$.value') AS value FROM json2_index_evolution WHERE json_get(attrs, '$.value') = 42.5 ORDER BY ts; + ++-------------------------+-------+ +| ts | value | ++-------------------------+-------+ +| 1970-01-01T00:00:00.003 | 42.5 | ++-------------------------+-------+ + +SELECT target_type, target_json, COUNT(*) AS index_count +FROM information_schema.ssts_index_meta +WHERE table_id = (SELECT table_id FROM information_schema.tables WHERE table_name = 'json2_index_evolution') +GROUP BY target_type, target_json ORDER BY target_type, target_json; + ++-------------+---------------------------------------------+-------------+ +| target_type | target_json | index_count | ++-------------+---------------------------------------------+-------------+ +| column | {"column":2} | 1 | +| json_path | {"json_path":{"column":1,"path":["value"]}} | 1 | ++-------------+---------------------------------------------+-------------+ + +-- SQLNESS REPLACE "metrics_per_partition":.*("rows_inverted_filtered":\d+).*metrics= "metrics_per_partition":{$1} metrics= +-- SQLNESS REPLACE "metrics_per_partition":\s*\[.*metrics= "metrics_per_partition":REDACTED metrics= +-- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED +-- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" +-- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED +-- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED +-- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE ,\s"dyn_filters":\s\[.* REDACTED +-- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED +-- SQLNESS REPLACE "max_part_range":\s*\{[^}]*\} "max_part_range":REDACTED +-- SQLNESS REPLACE "files":\s*\[.*\],\s*"flat_format" "files":REDACTED,"flat_format" +EXPLAIN ANALYZE VERBOSE SELECT ts FROM json2_index_evolution WHERE json_get(attrs, '$.value') > 600.0; + ++-+-+-+ +| stage | node | plan_| ++-+-+-+ +| 0_| 0_|_CooperativeExec metrics=REDACTED_| +|_|_|_MergeScanExec: REDACTED +|_|_|_| +| 1_| 0_|_FilterExec: json_get(attrs@1, $.value, NULL) > 600, projection=[ts@0] metrics=REDACTED_| +|_|_|_CooperativeExec metrics=REDACTED_| +|_|_|_UnorderedScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["ts", "attrs"], "filters": ["json_get(attrs, Utf8(\"$.value\"), Float64(NULL)) > Float64(600)"], "time_range": "TimestampRange{{[#,#)}}", "file_summary": {"files":1,"rows":3,"row_groups":1,"size":REDACTED,"index_size":REDACTED}, "max_part_range":REDACTED, "files":REDACTED,"flat_format":REDACTED, "metrics_per_partition":{"rows_inverted_filtered":3} metrics=REDACTED | +|_|_|_| +|_|_| Total rows: REDACTED_| ++-+-+-+ + +DROP TABLE json2_index_evolution; + +Affected Rows: 0 + +-- Removing a hint preserves old JSON values and leaves its old index unused. +CREATE TABLE json2_index_schema ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 (max_auto_expanded_paths = 0, value BIGINT INVERTED INDEX), + ordinary STRING INVERTED INDEX +) WITH ('append_mode' = 'true', 'sst_format' = 'flat', 'index.inverted_index.segment_row_count' = '1'); + +Affected Rows: 0 + +INSERT INTO json2_index_schema VALUES (1, '{"value":42}', 'a'), (2, '{"value":7}', 'b'); + +Affected Rows: 2 + +ADMIN FLUSH_TABLE('json2_index_schema'); + ++-----------------------------------------+ +| ADMIN FLUSH_TABLE('json2_index_schema') | ++-----------------------------------------+ +| 0 | ++-----------------------------------------+ + +SELECT ts FROM json2_index_schema WHERE json_get(attrs, '$.value', NULL::BIGINT) = 42 ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | ++-------------------------+ + +ALTER TABLE json2_index_schema MODIFY COLUMN attrs JSON2 (max_auto_expanded_paths = 0); + +Affected Rows: 0 + +SHOW CREATE TABLE json2_index_schema; + ++--------------------+---------------------------------------------------+ +| Table | Create Table | ++--------------------+---------------------------------------------------+ +| json2_index_schema | CREATE TABLE IF NOT EXISTS "json2_index_schema" ( | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | "attrs" JSON2( | +| | max_auto_expanded_paths = 0 | +| | ) NULL, | +| | "ordinary" STRING NULL INVERTED INDEX, | +| | TIME INDEX ("ts") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | append_mode = 'true', | +| | 'index.inverted_index.segment_row_count' = '1', | +| | sst_format = 'flat' | +| | ) | ++--------------------+---------------------------------------------------+ + +SELECT target_type, COUNT(*) AS index_count +FROM information_schema.ssts_index_meta +WHERE table_id = (SELECT table_id FROM information_schema.tables WHERE table_name = 'json2_index_schema') +GROUP BY target_type ORDER BY target_type; + ++-------------+-------------+ +| target_type | index_count | ++-------------+-------------+ +| column | 1 | +| json_path | 1 | ++-------------+-------------+ + +INSERT INTO json2_index_schema VALUES (3, '{"value":42}', 'c'); + +Affected Rows: 1 + +ADMIN FLUSH_TABLE('json2_index_schema'); + ++-----------------------------------------+ +| ADMIN FLUSH_TABLE('json2_index_schema') | ++-----------------------------------------+ +| 0 | ++-----------------------------------------+ + +-- Without a hint, two-argument access uses the default string result type. +SELECT ts, json_get(attrs, '$.value') AS value FROM json2_index_schema WHERE json_get(attrs, '$.value') = '42' ORDER BY ts; + ++-------------------------+-------+ +| ts | value | ++-------------------------+-------+ +| 1970-01-01T00:00:00.001 | 42 | +| 1970-01-01T00:00:00.003 | 42 | ++-------------------------+-------+ + +SELECT ts FROM json2_index_schema WHERE json_get(attrs, '$.value', NULL::BIGINT) = 42 ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | +| 1970-01-01T00:00:00.003 | ++-------------------------+ + +-- Removing the hint disables pruning by the old JSON index. +-- SQLNESS REPLACE "metrics_per_partition":.*("rows_inverted_filtered":\d+).*metrics= "metrics_per_partition":{$1} metrics= +-- SQLNESS REPLACE "metrics_per_partition":\s*\[.*metrics= "metrics_per_partition":REDACTED metrics= +-- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED +-- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" +-- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED +-- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED +-- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE ,\s"dyn_filters":\s\[.* REDACTED +-- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED +-- SQLNESS REPLACE "max_part_range":\s*\{[^}]*\} "max_part_range":REDACTED +-- SQLNESS REPLACE "files":\s*\[.*\],\s*"flat_format" "files":REDACTED,"flat_format" +EXPLAIN ANALYZE VERBOSE SELECT ts FROM json2_index_schema WHERE json_get(attrs, '$.value') = 'missing'; + ++-+-+-+ +| stage | node | plan_| ++-+-+-+ +| 0_| 0_|_CooperativeExec metrics=REDACTED_| +|_|_|_MergeScanExec: REDACTED +|_|_|_| +| 1_| 0_|_FilterExec: json_get(attrs@1, $.value) = missing, projection=[ts@0] metrics=REDACTED_| +|_|_|_CooperativeExec metrics=REDACTED_| +|_|_|_UnorderedScan: region=REDACTED, {"partition_count":{"count":2, "mem_ranges":0, "files":2, "file_ranges":2}, "projection": ["ts", "attrs"], "filters": ["json_get(attrs, Utf8(\"$.value\")) = Utf8View(\"missing\")"], "time_range": "TimestampRange{{[#,#)}}", "file_summary": {"files":2,"rows":3,"row_groups":2,"size":REDACTED,"index_size":REDACTED}, "max_part_range":REDACTED, "files":REDACTED,"flat_format":REDACTED, "metrics_per_partition":REDACTED metrics=REDACTED | +|_|_|_| +|_|_| Total rows: REDACTED_| ++-+-+-+ + +-- Compaction after hint removal preserves values and drops the JSON index. +SELECT ts, json_get(attrs, '$.value') AS value FROM json2_index_schema ORDER BY ts; + ++-------------------------+-------+ +| ts | value | ++-------------------------+-------+ +| 1970-01-01T00:00:00.001 | 42 | +| 1970-01-01T00:00:00.002 | 7 | +| 1970-01-01T00:00:00.003 | 42 | ++-------------------------+-------+ + +ADMIN COMPACT_TABLE('json2_index_schema', 'swcs', '86400'); + ++------------------------------------------------------------+ +| ADMIN COMPACT_TABLE('json2_index_schema', 'swcs', '86400') | ++------------------------------------------------------------+ +| 0 | ++------------------------------------------------------------+ + +SELECT ts, json_get(attrs, '$.value') AS value FROM json2_index_schema ORDER BY ts; + ++-------------------------+-------+ +| ts | value | ++-------------------------+-------+ +| 1970-01-01T00:00:00.001 | 42 | +| 1970-01-01T00:00:00.002 | 7 | +| 1970-01-01T00:00:00.003 | 42 | ++-------------------------+-------+ + +SELECT ts FROM json2_index_schema WHERE json_get(attrs, '$.value', NULL::BIGINT) = 42 ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | +| 1970-01-01T00:00:00.003 | ++-------------------------+ + +SELECT target_type, target_json, COUNT(*) AS index_count +FROM information_schema.ssts_index_meta +WHERE table_id = (SELECT table_id FROM information_schema.tables WHERE table_name = 'json2_index_schema') +GROUP BY target_type, target_json ORDER BY target_type, target_json; + ++-------------+--------------+-------------+ +| target_type | target_json | index_count | ++-------------+--------------+-------------+ +| column | {"column":2} | 1 | ++-------------+--------------+-------------+ + +-- SQLNESS REPLACE "metrics_per_partition":.*("rows_inverted_filtered":\d+).*metrics= "metrics_per_partition":{$1} metrics= +-- SQLNESS REPLACE "metrics_per_partition":\s*\[.*metrics= "metrics_per_partition":REDACTED metrics= +-- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED +-- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" +-- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED +-- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED +-- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE ,\s"dyn_filters":\s\[.* REDACTED +-- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED +-- SQLNESS REPLACE "max_part_range":\s*\{[^}]*\} "max_part_range":REDACTED +-- SQLNESS REPLACE "files":\s*\[.*\],\s*"flat_format" "files":REDACTED,"flat_format" +EXPLAIN ANALYZE VERBOSE SELECT ts FROM json2_index_schema WHERE json_get(attrs, '$.value') = 'missing'; + ++-+-+-+ +| stage | node | plan_| ++-+-+-+ +| 0_| 0_|_CooperativeExec metrics=REDACTED_| +|_|_|_MergeScanExec: REDACTED +|_|_|_| +| 1_| 0_|_FilterExec: json_get(attrs@1, $.value) = missing, projection=[ts@0] metrics=REDACTED_| +|_|_|_CooperativeExec metrics=REDACTED_| +|_|_|_UnorderedScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["ts", "attrs"], "filters": ["json_get(attrs, Utf8(\"$.value\")) = Utf8View(\"missing\")"], "time_range": "TimestampRange{{[#,#)}}", "file_summary": {"files":1,"rows":3,"row_groups":1,"size":REDACTED,"index_size":REDACTED}, "max_part_range":REDACTED, "files":REDACTED,"flat_format":REDACTED, "metrics_per_partition":REDACTED metrics=REDACTED | +|_|_|_| +|_|_| Total rows: REDACTED_| ++-+-+-+ + +DROP TABLE json2_index_schema; + +Affected Rows: 0 + +-- Disable only indexing, retaining the BIGINT hint and old physical index. +CREATE TABLE json2_index_disabled ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 (max_auto_expanded_paths = 0, value BIGINT INVERTED INDEX), + ordinary STRING INVERTED INDEX +) WITH ('append_mode' = 'true', 'sst_format' = 'flat', 'index.inverted_index.segment_row_count' = '1'); + +Affected Rows: 0 + +INSERT INTO json2_index_disabled VALUES (1, '{"value":42}', 'a'), (2, '{"value":7}', 'b'); + +Affected Rows: 2 + +ADMIN FLUSH_TABLE('json2_index_disabled'); + ++-------------------------------------------+ +| ADMIN FLUSH_TABLE('json2_index_disabled') | ++-------------------------------------------+ +| 0 | ++-------------------------------------------+ + +ALTER TABLE json2_index_disabled MODIFY COLUMN attrs JSON2 (max_auto_expanded_paths = 0, value BIGINT); + +Affected Rows: 0 + +SHOW CREATE TABLE json2_index_disabled; + ++----------------------+-----------------------------------------------------+ +| Table | Create Table | ++----------------------+-----------------------------------------------------+ +| json2_index_disabled | CREATE TABLE IF NOT EXISTS "json2_index_disabled" ( | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | "attrs" JSON2( | +| | max_auto_expanded_paths = 0, | +| | "value" BIGINT | +| | ) NULL, | +| | "ordinary" STRING NULL INVERTED INDEX, | +| | TIME INDEX ("ts") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | append_mode = 'true', | +| | 'index.inverted_index.segment_row_count' = '1', | +| | sst_format = 'flat' | +| | ) | ++----------------------+-----------------------------------------------------+ + +INSERT INTO json2_index_disabled VALUES (3, '{"value":42}', 'c'); + +Affected Rows: 1 + +ADMIN FLUSH_TABLE('json2_index_disabled'); + ++-------------------------------------------+ +| ADMIN FLUSH_TABLE('json2_index_disabled') | ++-------------------------------------------+ +| 0 | ++-------------------------------------------+ + +-- The retained hint still gives two-argument access a numeric result. +SELECT ts, json_get(attrs, '$.value') + 1 AS value FROM json2_index_disabled WHERE json_get(attrs, '$.value') = 42 ORDER BY ts; + ++-------------------------+-------+ +| ts | value | ++-------------------------+-------+ +| 1970-01-01T00:00:00.001 | 43 | +| 1970-01-01T00:00:00.003 | 43 | ++-------------------------+-------+ + +-- Only the old SST retains a physical JSON index. +SELECT target_type, target_json, COUNT(*) AS index_count +FROM information_schema.ssts_index_meta +WHERE table_id = (SELECT table_id FROM information_schema.tables WHERE table_name = 'json2_index_disabled') +GROUP BY target_type, target_json ORDER BY target_type, target_json; + ++-------------+---------------------------------------------+-------------+ +| target_type | target_json | index_count | ++-------------+---------------------------------------------+-------------+ +| column | {"column":2} | 2 | +| json_path | {"json_path":{"column":1,"path":["value"]}} | 1 | ++-------------+---------------------------------------------+-------------+ + +-- Neither SST uses the disabled JSON index. +-- SQLNESS REPLACE "metrics_per_partition":.*("rows_inverted_filtered":\d+).*metrics= "metrics_per_partition":{$1} metrics= +-- SQLNESS REPLACE "metrics_per_partition":\s*\[.*metrics= "metrics_per_partition":REDACTED metrics= +-- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED +-- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" +-- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED +-- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED +-- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE ,\s"dyn_filters":\s\[.* REDACTED +-- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED +-- SQLNESS REPLACE "max_part_range":\s*\{[^}]*\} "max_part_range":REDACTED +-- SQLNESS REPLACE "files":\s*\[.*\],\s*"flat_format" "files":REDACTED,"flat_format" +EXPLAIN ANALYZE VERBOSE SELECT ts FROM json2_index_disabled WHERE json_get(attrs, '$.value') > 600; + ++-+-+-+ +| stage | node | plan_| ++-+-+-+ +| 0_| 0_|_CooperativeExec metrics=REDACTED_| +|_|_|_MergeScanExec: REDACTED +|_|_|_| +| 1_| 0_|_FilterExec: json_get(attrs@1, $.value, NULL) > 600, projection=[ts@0] metrics=REDACTED_| +|_|_|_CooperativeExec metrics=REDACTED_| +|_|_|_UnorderedScan: region=REDACTED, {"partition_count":{"count":2, "mem_ranges":0, "files":2, "file_ranges":2}, "projection": ["ts", "attrs"], "filters": ["json_get(attrs, Utf8(\"$.value\"), Int64(NULL)) > Int64(600)"], "time_range": "TimestampRange{{[#,#)}}", "file_summary": {"files":2,"rows":3,"row_groups":2,"size":REDACTED,"index_size":REDACTED}, "max_part_range":REDACTED, "files":REDACTED,"flat_format":REDACTED, "metrics_per_partition":REDACTED metrics=REDACTED | +|_|_|_| +|_|_| Total rows: REDACTED_| ++-+-+-+ + +DROP TABLE json2_index_disabled; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/types/json/json2_inverted_index_schema.sql b/tests/cases/standalone/common/types/json/json2_inverted_index_schema.sql new file mode 100644 index 00000000000..839c2207c79 --- /dev/null +++ b/tests/cases/standalone/common/types/json/json2_inverted_index_schema.sql @@ -0,0 +1,290 @@ +-- Enabling a hint index must keep rows from SSTs with only an ordinary index. +CREATE TABLE json2_index_evolution ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 (max_auto_expanded_paths = 0, value BIGINT), + ordinary STRING INVERTED INDEX +) WITH ('append_mode' = 'true', 'sst_format' = 'flat', 'index.inverted_index.segment_row_count' = '1'); + +INSERT INTO json2_index_evolution VALUES (1, '{"value":42}', 'a'); + +ADMIN FLUSH_TABLE('json2_index_evolution'); + +ALTER TABLE json2_index_evolution MODIFY COLUMN attrs JSON2 (max_auto_expanded_paths = 0, value BIGINT INVERTED INDEX); + +SHOW CREATE TABLE json2_index_evolution; + +-- Two-argument access uses the enabled BIGINT hint for old SSTs. +SELECT ts, json_get(attrs, '$.value') AS value FROM json2_index_evolution WHERE json_get(attrs, '$.value') = 42 ORDER BY ts; + +SELECT ts FROM json2_index_evolution WHERE json_get(attrs, '$.value', NULL::BIGINT) = 42 ORDER BY ts; + +INSERT INTO json2_index_evolution VALUES (2, '{"value":42}', 'b'); + +ADMIN FLUSH_TABLE('json2_index_evolution'); + +-- Two-argument access returns matching rows from both old and new SSTs. +SELECT ts, json_get(attrs, '$.value') AS value FROM json2_index_evolution WHERE json_get(attrs, '$.value') = 42 ORDER BY ts; + +SELECT ts FROM json2_index_evolution WHERE json_get(attrs, '$.value', NULL::BIGINT) = 42 ORDER BY ts; + +-- Only the new SST has a usable BIGINT index; it filters one row. +-- SQLNESS REPLACE "metrics_per_partition":.*("rows_inverted_filtered":\d+).*metrics= "metrics_per_partition":{$1} metrics= +-- SQLNESS REPLACE "metrics_per_partition":\s*\[.*metrics= "metrics_per_partition":REDACTED metrics= +-- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED +-- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" +-- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED +-- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED +-- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE ,\s"dyn_filters":\s\[.* REDACTED +-- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED +-- SQLNESS REPLACE "max_part_range":\s*\{[^}]*\} "max_part_range":REDACTED +-- SQLNESS REPLACE "files":\s*\[.*\],\s*"flat_format" "files":REDACTED,"flat_format" +EXPLAIN ANALYZE VERBOSE SELECT ts FROM json2_index_evolution WHERE json_get(attrs, '$.value') > 600; + +-- Changing the hint type must not use old indexes with a different value encoding. +ALTER TABLE json2_index_evolution MODIFY COLUMN attrs JSON2 (max_auto_expanded_paths = 0, value DOUBLE INVERTED INDEX); + +SHOW CREATE TABLE json2_index_evolution; + +INSERT INTO json2_index_evolution VALUES (3, '{"value":42.5}', 'c'); + +ADMIN FLUSH_TABLE('json2_index_evolution'); + +SELECT target_type, target_json, COUNT(*) AS index_count +FROM information_schema.ssts_index_meta +WHERE table_id = (SELECT table_id FROM information_schema.tables WHERE table_name = 'json2_index_evolution') +GROUP BY target_type, target_json ORDER BY target_type, target_json; + +-- The DOUBLE hint matches old integer values and new fractional values. +SELECT ts, json_get(attrs, '$.value') AS value FROM json2_index_evolution WHERE json_get(attrs, '$.value') = 42.0 ORDER BY ts; + +SELECT ts, json_get(attrs, '$.value') AS value FROM json2_index_evolution WHERE json_get(attrs, '$.value') = 42.5 ORDER BY ts; + +-- Only the new DOUBLE index filters a row; older SSTs fall back to scanning. +-- SQLNESS REPLACE "metrics_per_partition":.*("rows_inverted_filtered":\d+).*metrics= "metrics_per_partition":{$1} metrics= +-- SQLNESS REPLACE "metrics_per_partition":\s*\[.*metrics= "metrics_per_partition":REDACTED metrics= +-- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED +-- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" +-- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED +-- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED +-- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE ,\s"dyn_filters":\s\[.* REDACTED +-- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED +-- SQLNESS REPLACE "max_part_range":\s*\{[^}]*\} "max_part_range":REDACTED +-- SQLNESS REPLACE "files":\s*\[.*\],\s*"flat_format" "files":REDACTED,"flat_format" +EXPLAIN ANALYZE VERBOSE SELECT ts FROM json2_index_evolution WHERE json_get(attrs, '$.value') = 42.0; + +-- Restart preserves the altered hint, query results, and index pruning. +-- SQLNESS ARG restart=true +SHOW CREATE TABLE json2_index_evolution; + +SELECT ts, json_get(attrs, '$.value') AS value FROM json2_index_evolution WHERE json_get(attrs, '$.value') = 42.0 ORDER BY ts; + +SELECT ts, json_get(attrs, '$.value') AS value FROM json2_index_evolution WHERE json_get(attrs, '$.value') = 42.5 ORDER BY ts; + +-- SQLNESS REPLACE "metrics_per_partition":.*("rows_inverted_filtered":\d+).*metrics= "metrics_per_partition":{$1} metrics= +-- SQLNESS REPLACE "metrics_per_partition":\s*\[.*metrics= "metrics_per_partition":REDACTED metrics= +-- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED +-- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" +-- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED +-- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED +-- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE ,\s"dyn_filters":\s\[.* REDACTED +-- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED +-- SQLNESS REPLACE "max_part_range":\s*\{[^}]*\} "max_part_range":REDACTED +-- SQLNESS REPLACE "files":\s*\[.*\],\s*"flat_format" "files":REDACTED,"flat_format" +EXPLAIN ANALYZE VERBOSE SELECT ts FROM json2_index_evolution WHERE json_get(attrs, '$.value') = 42.0; + +-- Compaction replaces old indexes with DOUBLE indexes for all three rows. +ADMIN COMPACT_TABLE('json2_index_evolution', 'swcs', '86400'); + +SELECT ts, json_get(attrs, '$.value') AS value FROM json2_index_evolution WHERE json_get(attrs, '$.value') = 42.0 ORDER BY ts; + +SELECT ts, json_get(attrs, '$.value') AS value FROM json2_index_evolution WHERE json_get(attrs, '$.value') = 42.5 ORDER BY ts; + +SELECT target_type, target_json, COUNT(*) AS index_count +FROM information_schema.ssts_index_meta +WHERE table_id = (SELECT table_id FROM information_schema.tables WHERE table_name = 'json2_index_evolution') +GROUP BY target_type, target_json ORDER BY target_type, target_json; + +-- SQLNESS REPLACE "metrics_per_partition":.*("rows_inverted_filtered":\d+).*metrics= "metrics_per_partition":{$1} metrics= +-- SQLNESS REPLACE "metrics_per_partition":\s*\[.*metrics= "metrics_per_partition":REDACTED metrics= +-- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED +-- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" +-- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED +-- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED +-- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE ,\s"dyn_filters":\s\[.* REDACTED +-- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED +-- SQLNESS REPLACE "max_part_range":\s*\{[^}]*\} "max_part_range":REDACTED +-- SQLNESS REPLACE "files":\s*\[.*\],\s*"flat_format" "files":REDACTED,"flat_format" +EXPLAIN ANALYZE VERBOSE SELECT ts FROM json2_index_evolution WHERE json_get(attrs, '$.value') > 600.0; + +DROP TABLE json2_index_evolution; + +-- Removing a hint preserves old JSON values and leaves its old index unused. +CREATE TABLE json2_index_schema ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 (max_auto_expanded_paths = 0, value BIGINT INVERTED INDEX), + ordinary STRING INVERTED INDEX +) WITH ('append_mode' = 'true', 'sst_format' = 'flat', 'index.inverted_index.segment_row_count' = '1'); + +INSERT INTO json2_index_schema VALUES (1, '{"value":42}', 'a'), (2, '{"value":7}', 'b'); + +ADMIN FLUSH_TABLE('json2_index_schema'); + +SELECT ts FROM json2_index_schema WHERE json_get(attrs, '$.value', NULL::BIGINT) = 42 ORDER BY ts; + +ALTER TABLE json2_index_schema MODIFY COLUMN attrs JSON2 (max_auto_expanded_paths = 0); + +SHOW CREATE TABLE json2_index_schema; + +SELECT target_type, COUNT(*) AS index_count +FROM information_schema.ssts_index_meta +WHERE table_id = (SELECT table_id FROM information_schema.tables WHERE table_name = 'json2_index_schema') +GROUP BY target_type ORDER BY target_type; + +INSERT INTO json2_index_schema VALUES (3, '{"value":42}', 'c'); + +ADMIN FLUSH_TABLE('json2_index_schema'); + +-- Without a hint, two-argument access uses the default string result type. +SELECT ts, json_get(attrs, '$.value') AS value FROM json2_index_schema WHERE json_get(attrs, '$.value') = '42' ORDER BY ts; + +SELECT ts FROM json2_index_schema WHERE json_get(attrs, '$.value', NULL::BIGINT) = 42 ORDER BY ts; + +-- Removing the hint disables pruning by the old JSON index. +-- SQLNESS REPLACE "metrics_per_partition":.*("rows_inverted_filtered":\d+).*metrics= "metrics_per_partition":{$1} metrics= +-- SQLNESS REPLACE "metrics_per_partition":\s*\[.*metrics= "metrics_per_partition":REDACTED metrics= +-- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED +-- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" +-- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED +-- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED +-- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE ,\s"dyn_filters":\s\[.* REDACTED +-- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED +-- SQLNESS REPLACE "max_part_range":\s*\{[^}]*\} "max_part_range":REDACTED +-- SQLNESS REPLACE "files":\s*\[.*\],\s*"flat_format" "files":REDACTED,"flat_format" +EXPLAIN ANALYZE VERBOSE SELECT ts FROM json2_index_schema WHERE json_get(attrs, '$.value') = 'missing'; + +-- Compaction after hint removal preserves values and drops the JSON index. +SELECT ts, json_get(attrs, '$.value') AS value FROM json2_index_schema ORDER BY ts; + +ADMIN COMPACT_TABLE('json2_index_schema', 'swcs', '86400'); + +SELECT ts, json_get(attrs, '$.value') AS value FROM json2_index_schema ORDER BY ts; + +SELECT ts FROM json2_index_schema WHERE json_get(attrs, '$.value', NULL::BIGINT) = 42 ORDER BY ts; + +SELECT target_type, target_json, COUNT(*) AS index_count +FROM information_schema.ssts_index_meta +WHERE table_id = (SELECT table_id FROM information_schema.tables WHERE table_name = 'json2_index_schema') +GROUP BY target_type, target_json ORDER BY target_type, target_json; + +-- SQLNESS REPLACE "metrics_per_partition":.*("rows_inverted_filtered":\d+).*metrics= "metrics_per_partition":{$1} metrics= +-- SQLNESS REPLACE "metrics_per_partition":\s*\[.*metrics= "metrics_per_partition":REDACTED metrics= +-- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED +-- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" +-- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED +-- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED +-- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE ,\s"dyn_filters":\s\[.* REDACTED +-- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED +-- SQLNESS REPLACE "max_part_range":\s*\{[^}]*\} "max_part_range":REDACTED +-- SQLNESS REPLACE "files":\s*\[.*\],\s*"flat_format" "files":REDACTED,"flat_format" +EXPLAIN ANALYZE VERBOSE SELECT ts FROM json2_index_schema WHERE json_get(attrs, '$.value') = 'missing'; + +DROP TABLE json2_index_schema; + +-- Disable only indexing, retaining the BIGINT hint and old physical index. +CREATE TABLE json2_index_disabled ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 (max_auto_expanded_paths = 0, value BIGINT INVERTED INDEX), + ordinary STRING INVERTED INDEX +) WITH ('append_mode' = 'true', 'sst_format' = 'flat', 'index.inverted_index.segment_row_count' = '1'); + +INSERT INTO json2_index_disabled VALUES (1, '{"value":42}', 'a'), (2, '{"value":7}', 'b'); + +ADMIN FLUSH_TABLE('json2_index_disabled'); + +ALTER TABLE json2_index_disabled MODIFY COLUMN attrs JSON2 (max_auto_expanded_paths = 0, value BIGINT); + +SHOW CREATE TABLE json2_index_disabled; + +INSERT INTO json2_index_disabled VALUES (3, '{"value":42}', 'c'); + +ADMIN FLUSH_TABLE('json2_index_disabled'); + +-- The retained hint still gives two-argument access a numeric result. +SELECT ts, json_get(attrs, '$.value') + 1 AS value FROM json2_index_disabled WHERE json_get(attrs, '$.value') = 42 ORDER BY ts; + +-- Only the old SST retains a physical JSON index. +SELECT target_type, target_json, COUNT(*) AS index_count +FROM information_schema.ssts_index_meta +WHERE table_id = (SELECT table_id FROM information_schema.tables WHERE table_name = 'json2_index_disabled') +GROUP BY target_type, target_json ORDER BY target_type, target_json; + +-- Neither SST uses the disabled JSON index. +-- SQLNESS REPLACE "metrics_per_partition":.*("rows_inverted_filtered":\d+).*metrics= "metrics_per_partition":{$1} metrics= +-- SQLNESS REPLACE "metrics_per_partition":\s*\[.*metrics= "metrics_per_partition":REDACTED metrics= +-- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED +-- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" +-- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED +-- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED +-- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE ,\s"dyn_filters":\s\[.* REDACTED +-- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED +-- SQLNESS REPLACE "max_part_range":\s*\{[^}]*\} "max_part_range":REDACTED +-- SQLNESS REPLACE "files":\s*\[.*\],\s*"flat_format" "files":REDACTED,"flat_format" +EXPLAIN ANALYZE VERBOSE SELECT ts FROM json2_index_disabled WHERE json_get(attrs, '$.value') > 600; + +DROP TABLE json2_index_disabled; diff --git a/tests/cases/standalone/common/types/json/json2_inverted_index_write.result b/tests/cases/standalone/common/types/json/json2_inverted_index_write.result new file mode 100644 index 00000000000..a0baf7861a0 --- /dev/null +++ b/tests/cases/standalone/common/types/json/json2_inverted_index_write.result @@ -0,0 +1,153 @@ +-- Flush builds typed JSON indexes, including missing and null values. +CREATE TABLE json2_index_write ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 ( + max_auto_expanded_paths = 0, + "service.name" STRING INVERTED INDEX, + service.name STRING INVERTED INDEX, + status BIGINT INVERTED INDEX, + bytes BIGINT UNSIGNED INVERTED INDEX, + duration DOUBLE INVERTED INDEX, + success BOOLEAN INVERTED INDEX, + absent STRING INVERTED INDEX, + unindexed STRING + ), + ordinary STRING INVERTED INDEX +) WITH ('append_mode' = 'true', 'sst_format' = 'flat'); + +Affected Rows: 0 + +INSERT INTO json2_index_write VALUES + (1, '{"service.name":"api","service":{"name":"nested"},"status":200,"bytes":18446744073709551615,"duration":1.5,"success":true,"unindexed":"ignored"}', 'a'), + (2, '{"service":{}}', 'b'), + (3, '{}', 'c'), + (4, NULL, 'd'); + +Affected Rows: 4 + +ADMIN FLUSH_TABLE('json2_index_write'); + ++----------------------------------------+ +| ADMIN FLUSH_TABLE('json2_index_write') | ++----------------------------------------+ +| 0 | ++----------------------------------------+ + +-- Read physical index metadata, not the table's declared index settings. +SELECT index_type, target_type, target_json, blob_size > 0 AS has_index_data +FROM information_schema.ssts_index_meta +WHERE table_id = (SELECT table_id FROM information_schema.tables WHERE table_name = 'json2_index_write') +ORDER BY target_type, target_json; + ++------------+-------------+------------------------------------------------------+----------------+ +| index_type | target_type | target_json | has_index_data | ++------------+-------------+------------------------------------------------------+----------------+ +| inverted | column | {"column":2} | true | +| inverted | json_path | {"json_path":{"column":1,"path":["absent"]}} | true | +| inverted | json_path | {"json_path":{"column":1,"path":["bytes"]}} | true | +| inverted | json_path | {"json_path":{"column":1,"path":["duration"]}} | true | +| inverted | json_path | {"json_path":{"column":1,"path":["service","name"]}} | true | +| inverted | json_path | {"json_path":{"column":1,"path":["service.name"]}} | true | +| inverted | json_path | {"json_path":{"column":1,"path":["status"]}} | true | +| inverted | json_path | {"json_path":{"column":1,"path":["success"]}} | true | ++------------+-------------+------------------------------------------------------+----------------+ + +-- Each flushed SST gets its own set of indexes. +INSERT INTO json2_index_write VALUES + (5, '{"service.name":"worker","service":{"name":"batch"},"status":500}', 'e'); + +Affected Rows: 1 + +ADMIN FLUSH_TABLE('json2_index_write'); + ++----------------------------------------+ +| ADMIN FLUSH_TABLE('json2_index_write') | ++----------------------------------------+ +| 0 | ++----------------------------------------+ + +SELECT target_type, COUNT(*) AS index_count +FROM information_schema.ssts_index_meta +WHERE table_id = (SELECT table_id FROM information_schema.tables WHERE table_name = 'json2_index_write') +GROUP BY target_type ORDER BY target_type; + ++-------------+-------------+ +| target_type | index_count | ++-------------+-------------+ +| column | 2 | +| json_path | 14 | ++-------------+-------------+ + +-- Compaction recreates the same targets and preserves all rows. +ADMIN COMPACT_TABLE('json2_index_write', 'swcs', '86400'); + ++-----------------------------------------------------------+ +| ADMIN COMPACT_TABLE('json2_index_write', 'swcs', '86400') | ++-----------------------------------------------------------+ +| 0 | ++-----------------------------------------------------------+ + +SELECT index_type, target_type, target_json, blob_size > 0 AS has_index_data +FROM information_schema.ssts_index_meta +WHERE table_id = (SELECT table_id FROM information_schema.tables WHERE table_name = 'json2_index_write') +ORDER BY target_type, target_json; + ++------------+-------------+------------------------------------------------------+----------------+ +| index_type | target_type | target_json | has_index_data | ++------------+-------------+------------------------------------------------------+----------------+ +| inverted | column | {"column":2} | true | +| inverted | json_path | {"json_path":{"column":1,"path":["absent"]}} | true | +| inverted | json_path | {"json_path":{"column":1,"path":["bytes"]}} | true | +| inverted | json_path | {"json_path":{"column":1,"path":["duration"]}} | true | +| inverted | json_path | {"json_path":{"column":1,"path":["service","name"]}} | true | +| inverted | json_path | {"json_path":{"column":1,"path":["service.name"]}} | true | +| inverted | json_path | {"json_path":{"column":1,"path":["status"]}} | true | +| inverted | json_path | {"json_path":{"column":1,"path":["success"]}} | true | ++------------+-------------+------------------------------------------------------+----------------+ + +SELECT COUNT(*) FROM json2_index_write; + ++----------+ +| count(*) | ++----------+ +| 5 | ++----------+ + +DROP TABLE json2_index_write; + +Affected Rows: 0 + +-- A JSON-only index must still produce an index file and file-level availability. +CREATE TABLE json2_index_only ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 (value BIGINT INVERTED INDEX) +) WITH ('append_mode' = 'true'); + +Affected Rows: 0 + +INSERT INTO json2_index_only VALUES (1, '{"value":42}'), (2, '{}'); + +Affected Rows: 2 + +ADMIN FLUSH_TABLE('json2_index_only'); + ++---------------------------------------+ +| ADMIN FLUSH_TABLE('json2_index_only') | ++---------------------------------------+ +| 0 | ++---------------------------------------+ + +SELECT index_type, target_type, target_json, index_file_size > 0 AS has_index_file +FROM information_schema.ssts_index_meta +WHERE table_id = (SELECT table_id FROM information_schema.tables WHERE table_name = 'json2_index_only'); + ++------------+-------------+---------------------------------------------+----------------+ +| index_type | target_type | target_json | has_index_file | ++------------+-------------+---------------------------------------------+----------------+ +| inverted | json_path | {"json_path":{"column":1,"path":["value"]}} | true | ++------------+-------------+---------------------------------------------+----------------+ + +DROP TABLE json2_index_only; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/types/json/json2_inverted_index_write.sql b/tests/cases/standalone/common/types/json/json2_inverted_index_write.sql new file mode 100644 index 00000000000..60a7a5195fd --- /dev/null +++ b/tests/cases/standalone/common/types/json/json2_inverted_index_write.sql @@ -0,0 +1,69 @@ +-- Flush builds typed JSON indexes, including missing and null values. +CREATE TABLE json2_index_write ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 ( + max_auto_expanded_paths = 0, + "service.name" STRING INVERTED INDEX, + service.name STRING INVERTED INDEX, + status BIGINT INVERTED INDEX, + bytes BIGINT UNSIGNED INVERTED INDEX, + duration DOUBLE INVERTED INDEX, + success BOOLEAN INVERTED INDEX, + absent STRING INVERTED INDEX, + unindexed STRING + ), + ordinary STRING INVERTED INDEX +) WITH ('append_mode' = 'true', 'sst_format' = 'flat'); + +INSERT INTO json2_index_write VALUES + (1, '{"service.name":"api","service":{"name":"nested"},"status":200,"bytes":18446744073709551615,"duration":1.5,"success":true,"unindexed":"ignored"}', 'a'), + (2, '{"service":{}}', 'b'), + (3, '{}', 'c'), + (4, NULL, 'd'); + +ADMIN FLUSH_TABLE('json2_index_write'); + +-- Read physical index metadata, not the table's declared index settings. +SELECT index_type, target_type, target_json, blob_size > 0 AS has_index_data +FROM information_schema.ssts_index_meta +WHERE table_id = (SELECT table_id FROM information_schema.tables WHERE table_name = 'json2_index_write') +ORDER BY target_type, target_json; + +-- Each flushed SST gets its own set of indexes. +INSERT INTO json2_index_write VALUES + (5, '{"service.name":"worker","service":{"name":"batch"},"status":500}', 'e'); + +ADMIN FLUSH_TABLE('json2_index_write'); + +SELECT target_type, COUNT(*) AS index_count +FROM information_schema.ssts_index_meta +WHERE table_id = (SELECT table_id FROM information_schema.tables WHERE table_name = 'json2_index_write') +GROUP BY target_type ORDER BY target_type; + +-- Compaction recreates the same targets and preserves all rows. +ADMIN COMPACT_TABLE('json2_index_write', 'swcs', '86400'); + +SELECT index_type, target_type, target_json, blob_size > 0 AS has_index_data +FROM information_schema.ssts_index_meta +WHERE table_id = (SELECT table_id FROM information_schema.tables WHERE table_name = 'json2_index_write') +ORDER BY target_type, target_json; + +SELECT COUNT(*) FROM json2_index_write; + +DROP TABLE json2_index_write; + +-- A JSON-only index must still produce an index file and file-level availability. +CREATE TABLE json2_index_only ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 (value BIGINT INVERTED INDEX) +) WITH ('append_mode' = 'true'); + +INSERT INTO json2_index_only VALUES (1, '{"value":42}'), (2, '{}'); + +ADMIN FLUSH_TABLE('json2_index_only'); + +SELECT index_type, target_type, target_json, index_file_size > 0 AS has_index_file +FROM information_schema.ssts_index_meta +WHERE table_id = (SELECT table_id FROM information_schema.tables WHERE table_name = 'json2_index_only'); + +DROP TABLE json2_index_only; From bb3f07b4825e074c207ae29db28de1e3354c11ae Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Thu, 22 Oct 2026 13:21:39 +0000 Subject: [PATCH 074/239] fix(mito2): pick TWCS L1 merges by closed units so compaction terminates (#9741) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * refactor(mito2): pick TWCS L1 merges by closed units TWCS could keep rewriting L1 files without new writes. Overlap-only picks could split their inputs into one more file than they read, and L1 file reduction then rewrote up to 16 files to remove that single file. The outputs, split by key, created new overlaps and the cycle repeated. Replace the three global pick phases with one pick per window, scheduled by class across windows. Every pick now decreases a termination measure when the window receives no writes. - Active windows only take merges that remove files. L1 files are grouped by conflicts and merged in key order, so size-split tails are merged together. A reduction rewrites at most 4 output files' worth of bytes per removed file. - Inactive windows absorb L0 files, then take closed merges: no other file intrudes on the inputs' bounding box, so outputs cannot add overlaps. Each merge resolves conflicts or removes files, decreasing (conflicts, files). Files that do not take part in an overlap are not rewritten. - Overlaps that only share a boundary key no longer count as conflicts, so a size-split sorted run does not look overlapping. - Closed merges may take 4x the input limit to resolve a whole window once. - Cache primary key ranges per pick instead of resolving them per comparison. - Log the pick branch of each compaction output. Signed-off-by: Lei, HUANG * docs(mito2): state the compactor assumptions behind TWCS termination Signed-off-by: Lei, HUANG * feat(mito2): resolve active window overlaps on manual compaction A window that stops receiving writes keeps being the active window, whose picks only remove files, so its overlaps were never resolved. Manual compactions now pick every window as inactive, so a closed merge can resolve them. Automatic compactions are unchanged. Signed-off-by: Lei, HUANG * fix(mito2): defer overlap resolution in the window before the active one Late data keeps landing in the window right before the active window. A closed merge that resolves its overlaps early rewrites the whole window, then the late data brings the overlaps back, and the balanced run check keeps them there. Overlap-only merges now skip that window until it is older. File-reducing merges and manual compactions are unchanged. On the 12h/1GB out-of-order bench this reaches overlap depth 1.0 instead of 2.08 at the same write amplification (3.01 vs 3.02). Signed-off-by: Lei, HUANG * refactor(mito2): only check earlier-starting files when closing units Closing units only checks files that start before a unit ends, and the overlap checks compare key bounds without cloning them. Signed-off-by: Lei, HUANG * fix(mito2): keep settling windows from rewriting L1 files The window before the active one still receives late data. Skipping only its overlap-only merges was not enough: file-reducing merges of conflicting files and mixed L0/L1 merges still rewrote its L1 files, and a mixed merge could rewrite a nearly full output for each small late file. A settling window now only merges L0 files among themselves and conflict-free files; the window resolves its overlaps once it is older. On the out-of-order bench, write amplification drops from 3.88 to 3.30 at 2GB/6h, 3.05 to 2.94 at 1GB/6h and 3.01 to 2.96 at 1GB/12h, with overlap depth 1.0 in all three. Signed-off-by: Lei, HUANG * fix(mito2): resolve late overlaps once a window is quiet Late data that lands in a window older than the settling one is merged into a small file covering the whole key range. Resolving it rewrites the window's dominant sorted run, which the balanced run check refuses, so the overlap stayed. Removing the check instead rewrites the window for every late flush. A window is quiet once the region has written at least as many rows as the window holds since its newest file. Overlap merges in a quiet window may rewrite a dominant run, so each quiet period pays for at most one such rewrite. Sequences are per row and kept by compaction, so this needs no state. On the bench with late data up to 2 hours, mean overlap depth drops from 1.67 to 1.17 at 2GB/6h and from 2.33 to 1.83 at 1GB/6h, at write amplification 3.70 to 4.25 and 3.26 to 3.74; all inactive windows of 2GB/6h realtime reach depth 1 (write amplification 2.80 to 3.30). The remaining depth is in the final active window. Signed-off-by: Lei, HUANG * test(mito2): check closed merges of quiet and settling windows terminate Signed-off-by: Lei, HUANG * fix(mito2): split TWCS closures over the input limit A file spanning the keys of a large sorted run joins every run file into one closure. Once a closure held more files than one compaction accepts, closed merges skipped it, so the window never merged those files again and late files kept piling up. Split such closures into single files, like active-window picks do, and replace the closed check with a conflict bound: outputs only conflict with files intruding on the inputs' box, so a merge is taken if `outputs · I` does not exceed the conflicts its inputs had with other files. Intervals of whole closures have no such conflicts and still need to be closed, so other windows pick as before, and `(conflicts, files)` still strictly decreases. Small run files and late files over a large run can be merged again. A late file over a run that cannot shrink stays: resolving it would rewrite the whole run. Signed-off-by: Lei, HUANG * fix(mito2): address TWCS picker review findings - Manual compactions may rewrite a dominant sorted run to resolve overlaps. The active window is never quiet, so manual compactions could not resolve its late-data overlaps before. - Single-key intervals predict one output: the writer never splits a series, so large files of one series were never merged. - The active window falls back to single files when no group of conflicting files qualifies, so a large file no longer keeps its small peers unmerged. - Compute the settling window distance with checked_sub, which overflowed for windows at the timestamp limits. Signed-off-by: Lei, HUANG * fix(mito2): count inner keys as TWCS conflicts and check balance per group - A single key strictly inside another file's key range now conflicts, so a late file of one series inside a wide file is resolved like other overlaps. Files sharing only a boundary key still do not conflict. Intrusion covers the new pairs, which a randomized test checks. - Check the dominant-run guard per group of conflicting files. Runs of separate groups never overlap, so pairing one group's run with another group's peer made an interval of two dominant runs look balanced. - A window with unknown row counts, written before they were recorded, is not quiet: its threshold would drop to a single row. Signed-off-by: Lei, HUANG * perf(mito2): index TWCS files by key to avoid quadratic grouping Grouping a window's files compared every file with all files sharing its time range, and closing groups and checking intrusions scanned the whole window. A sorted run split by key shares one time range, so a window of N such files took O(N²) checks without any conflict. The 1024-file fallback only covers the sorted-run path, not these groups. Index the files by lower primary key in a segment tree that keeps each subtree's largest upper key and its time bounds. Grouping, closing, and the conflict bound only check the files whose key and time ranges meet, so a file of a key-sliced run only meets its neighbors. Conflicts among split files are precomputed per file and counted per interval with a binary search instead of pairwise. On a run of 20,000 key-sliced files with 5 late files spanning all keys, a closed merge pick takes ~62 ms instead of ~3.9 s and an active pick ~19 ms instead of ~3.5 s. Signed-off-by: Lei, HUANG * test(object-store): retry transient S3 transport failures Signed-off-by: Lei, HUANG * fix(mito2): bound TWCS closed-merge conflict graph work Budget the candidate files visited while building conflict units and split-unit adjacency. When the budget is exceeded, treat every file as a split unit and pick a bounded file reduction instead of materializing a quadratic graph. The fallback candidate is admitted only through adds_no_conflicts, so outputs do not add conflict pairs and the (P, N) termination invariant still holds. Dense picks are logged with their own branch. Signed-off-by: Lei, HUANG * fix(mito2): charge TWCS index nodes to the conflict-graph budget The conflict-graph work budget only counted files that reached a callback, so a traversal with poor pruning could still visit a quadratic number of internal nodes without triggering the dense fallback. Charge every ExtentIndex node, including pruned nodes, and stop the traversal as soon as the shared budget is exhausted. Add a regression test where the root straddles the query time but neither leaf matches, so a callback-only budget would never fire. Signed-off-by: Lei, HUANG * fix(mito2): retry smaller TWCS dense compaction candidates If a dense fallback candidate fails admission, try one single-output subinterval within it before giving up. Reuse the reduction guards and check both candidates against the full window's external conflicts, preserving bounded work and termination.\n\nCover the chain regression through both fallback unit constructions and reject one-output retries that introduce conflicts across time gaps. Signed-off-by: Lei, HUANG * fix(mito2): compact shared primary-key boundaries Count shared primary-key endpoints as conflicts and apply the same overlap rule to bounding-box intrusions. Existing compaction writers keep series intact, so fresh output ranges remain disjoint. Align simulated output splits with the writer and cover boundary-only cleanup in inactive and manual windows, plus stable real output files after a subsequent TWCS request. Signed-off-by: Lei, HUANG * fix(mito2): compact small L1 files in inactive windows Allow bounded small-L1 reductions in settling and older inactive windows when normal picks cannot proceed. Preserve full-window conflict checks, large-file protection, and LastNonNull closure semantics.\n\nCover selection limits, termination, and real SST compaction with tombstones and reopen checks. Signed-off-by: Lei, HUANG * docs(mito2): clarify TWCS unit grouping guarantees Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- src/mito2/src/compaction/last_non_null.rs | 7 +- src/mito2/src/compaction/picker.rs | 3 + src/mito2/src/compaction/scheduler.rs | 11 +- .../src/compaction/scheduler/planning.rs | 8 +- src/mito2/src/compaction/twcs.rs | 1484 +++++++++--- src/mito2/src/compaction/twcs/units.rs | 2073 +++++++++++++++++ src/mito2/src/engine/compaction_test.rs | 206 +- src/mito2/src/gc/worker_test.rs | 14 +- src/object-store/tests/object_store_test.rs | 3 + 9 files changed, 3426 insertions(+), 383 deletions(-) create mode 100644 src/mito2/src/compaction/twcs/units.rs diff --git a/src/mito2/src/compaction/last_non_null.rs b/src/mito2/src/compaction/last_non_null.rs index afcdc87e90f..906788658d3 100644 --- a/src/mito2/src/compaction/last_non_null.rs +++ b/src/mito2/src/compaction/last_non_null.rs @@ -293,7 +293,7 @@ mod tests { .unwrap() .set_compacting(busy); let request = compact_request::Options::Regular(Default::default()); - let picker = new_picker(&request, ®ion.region_options, limit, None); + let picker = new_picker(&request, ®ion.region_options, limit, None, false); let picked = picker.pick(®ion).await.unwrap().unwrap(); assert_eq!(expected_windows.len(), picked.outputs.len()); @@ -312,7 +312,7 @@ mod tests { // The LastRow picker must retain its original seed cap and priority. region.region_options.merge_mode = Some(MergeMode::LastRow); - let picker = new_picker(&request, ®ion.region_options, limit, None); + let picker = new_picker(&request, ®ion.region_options, limit, None, false); let picked = picker.pick(®ion).await.unwrap().unwrap(); assert_eq!(limit.unwrap_or(3).min(3), picked.outputs.len()); for (output, window) in picked.outputs.iter().rev().zip(windows.iter().rev()) { @@ -487,6 +487,7 @@ mod tests { ®ion.region_options, Some(1), TimestampRange::new(Timestamp::new_second(0), Timestamp::new_second(3600)), + false, ); let picked = picker.pick(®ion).await.unwrap(); @@ -543,6 +544,7 @@ mod tests { ®ion.region_options, Some(1), TimestampRange::new(Timestamp::new_second(10), Timestamp::new_second(20)), + false, ); let picked = picker.pick(®ion).await.unwrap().unwrap(); @@ -590,6 +592,7 @@ mod tests { ®ion.region_options, Some(1), None, + false, ); let picked = picker.pick(®ion).await.unwrap().unwrap(); assert!(picked.outputs.is_empty()); diff --git a/src/mito2/src/compaction/picker.rs b/src/mito2/src/compaction/picker.rs index 3628953cc31..539638e570f 100644 --- a/src/mito2/src/compaction/picker.rs +++ b/src/mito2/src/compaction/picker.rs @@ -128,6 +128,7 @@ pub fn new_picker( region_options: &RegionOptions, max_background_tasks: Option, time_range: Option, + manual: bool, ) -> Arc { if let compact_request::Options::StrictWindow(window) = compact_request_options { let window = if window.window_seconds == 0 { @@ -148,8 +149,10 @@ pub fn new_picker( time_window_seconds: twcs_opts.time_window_seconds(), max_output_file_size: twcs_opts.max_output_file_size.map(|r| r.as_bytes()), append_mode: region_options.append_mode, + merge_mode: region_options.merge_mode(), max_background_tasks, time_range, + all_windows_inactive: manual, }, }; if region_options.merge_mode() == MergeMode::LastNonNull { diff --git a/src/mito2/src/compaction/scheduler.rs b/src/mito2/src/compaction/scheduler.rs index 7406a639305..22df8c094ab 100644 --- a/src/mito2/src/compaction/scheduler.rs +++ b/src/mito2/src/compaction/scheduler.rs @@ -571,7 +571,7 @@ impl CompactionScheduler { max_parallelism, ); status.merge_waiter(waiter); - self.dispatch_compaction_planning(plan_id, request, compact_options, time_range); + self.dispatch_compaction_planning(plan_id, request, compact_options, time_range, trigger); self.listener.on_compaction_scheduled(region_id); Ok(true) } @@ -617,7 +617,13 @@ impl CompactionScheduler { // was fetched above. let plan_id = Self::next_plan_id(&mut self.next_plan_id); status.start_picking_with_trigger(plan_id, CompactionTrigger::Manual); - self.dispatch_compaction_planning(plan_id, request, options, time_range); + self.dispatch_compaction_planning( + plan_id, + request, + options, + time_range, + CompactionTrigger::Manual, + ); debug!( "Successfully scheduled manual compaction planning for region id: {}", region_id @@ -703,6 +709,7 @@ impl CompactionScheduler { request, compact_request::Options::Regular(Default::default()), None, + CompactionTrigger::Automatic, ); debug!( "Successfully scheduled next compaction planning for region id: {}", diff --git a/src/mito2/src/compaction/scheduler/planning.rs b/src/mito2/src/compaction/scheduler/planning.rs index d427044b955..2cc620ebbd1 100644 --- a/src/mito2/src/compaction/scheduler/planning.rs +++ b/src/mito2/src/compaction/scheduler/planning.rs @@ -32,7 +32,9 @@ use crate::access_layer::AccessLayerRef; use crate::cache::CacheManagerRef; use crate::compaction::compactor::{CompactionRegion, CompactionVersion}; use crate::compaction::picker::{PickerOutput, new_picker}; -use crate::compaction::scheduler::state::{CompactingFiles, CompactionExecution, CompactionPhase}; +use crate::compaction::scheduler::state::{ + CompactingFiles, CompactionExecution, CompactionPhase, CompactionTrigger, +}; use crate::compaction::scheduler::{CompactionScheduler, CompactionTransition}; use crate::compaction::unit::CompactionUnit; use crate::compaction::{CompactionOutput, find_dynamic_options, format_panic_reason}; @@ -115,6 +117,7 @@ impl CompactionScheduler { request: CompactionRequest, options: compact_request::Options, time_range: Option, + trigger: CompactionTrigger, ) { let plugins = self.plugins.clone(); let max_background_compactions = self.engine_config.max_background_compactions; @@ -127,6 +130,7 @@ impl CompactionScheduler { plugins, max_background_compactions, time_range, + trigger, ); Self::notify_planning_result(region_id, plan_id, request_sender, planning).await; }); @@ -180,6 +184,7 @@ impl CompactionScheduler { plugins: Plugins, max_background_compactions: usize, time_range: Option, + trigger: CompactionTrigger, ) -> CompactionPlanningResult { let region_id = request.region_id(); let (dynamic_compaction_opts, ttl) = find_dynamic_options( @@ -208,6 +213,7 @@ impl CompactionScheduler { ®ion_options, Some(max_picker_outputs), time_range, + trigger == CompactionTrigger::Manual, ); let region_id = request.region_id(); let CompactionRequest { diff --git a/src/mito2/src/compaction/twcs.rs b/src/mito2/src/compaction/twcs.rs index a1f51b64091..bcbb3314525 100644 --- a/src/mito2/src/compaction/twcs.rs +++ b/src/mito2/src/compaction/twcs.rs @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +mod units; + use std::collections::hash_map::Entry; use std::collections::{BTreeMap, HashMap, HashSet}; use std::fmt::Debug; @@ -34,43 +36,36 @@ use crate::compaction::run::{ Ranged, SortedRun, files_overlap, files_overlap_inclusive, find_sorted_runs, find_sorted_runs_by_time_range, merge_primary_key_ranges, primary_key_ranges_overlap, }; +use crate::compaction::twcs::units::{ClosedMergeOptions, PickLimits}; use crate::error::{JoinSnafu, Result}; +use crate::region::options::MergeMode; use crate::sst::file::{FileHandle, Level, overlaps}; use crate::sst::primary_key::PrimaryKeyRangeMapper; use crate::sst::version::LevelMeta; const LEVEL_COMPACTED: Level = 1; -#[derive(Clone, Copy, Debug)] -enum PickPhase { - HasL0, - L1FileReduction, - L1OverlapOnly, +/// Scheduling priority of a window's pick. Earlier classes are scheduled first, +/// and newer windows first within a class. +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)] +enum PickClass { + /// Merges L0 files, which also absorbs late data of inactive windows. + L0Merge, + /// Merges files to reduce the window's file count. + FileReduction, + /// Merges files of an inactive window to resolve their overlaps. + OverlapResolution, } -impl PickPhase { - fn applies_to(self, window: &Window) -> bool { - match self { - Self::HasL0 => window.files().any(|file| file.level() == 0), - Self::L1FileReduction | Self::L1OverlapOnly => { - window.files().any(|file| file.level() != 0) - } - } - } -} - -const PICK_PHASES: [PickPhase; 3] = [ - PickPhase::HasL0, - PickPhase::L1FileReduction, - PickPhase::L1OverlapOnly, -]; - -struct WindowPickContext<'a> { - active_window: Option, - files: &'a Window, - windows: &'a BTreeMap, - phase: PickPhase, - primary_key_mapper: &'a PrimaryKeyRangeMapper, +/// Compaction inputs picked for one window. +#[derive(Debug)] +struct WindowPick { + class: PickClass, + /// Name of the strategy that picked the inputs, for logs. + branch: &'static str, + inputs: Vec, + /// Number of sorted runs or units the strategy considered, for logs. + found_runs: usize, } struct WindowOutputContext { @@ -95,6 +90,11 @@ static MAX_INPUT_FILES: LazyLock = LazyLock::new(|| { parse_max_input_files(env_value.as_deref()) }); +/// A closed merge in an inactive window may take this many times +/// [`MAX_INPUT_FILES`]. Resolving overlaps across a whole window needs all of +/// its overlapping files at once, and each byte is rewritten once per window. +const CLOSED_MERGE_INPUT_FACTOR: usize = 4; + fn parse_max_input_files(env_value: Option<&str>) -> usize { env_value .and_then(|env_value| env_value.parse().ok()) @@ -121,13 +121,22 @@ pub struct TwcsPicker { pub max_output_file_size: Option, /// Whether the target region is in append mode. pub append_mode: bool, + /// LastNonNull picks expand to an overlap closure, so cannot use partial + /// small-file reductions that must leave large peers untouched. + pub(crate) merge_mode: MergeMode, /// Max background compaction tasks. pub max_background_tasks: Option, /// Optional time range that constrains candidate compaction windows. pub(crate) time_range: Option, + /// Picks the active window like an inactive one. Manual compactions set it, + /// so a window that stopped receiving writes still gets its overlaps + /// resolved. + pub(crate) all_windows_inactive: bool, } impl TwcsPicker { + /// Picks at most one compaction input per window and schedules them across + /// windows by [`PickClass`], newer windows first within a class. async fn build_output( &self, region_id: RegionId, @@ -140,7 +149,13 @@ impl TwcsPicker { max_outputs, primary_key_mapper, } = context; - let mut output = vec![]; + let region_sequence = time_windows + .values() + .flat_map(|window| window.files()) + .filter_map(|file| file.meta_ref().sequence) + .map(|sequence| sequence.get()) + .max() + .unwrap_or(0); let windows = time_windows .values() .rev() @@ -159,278 +174,330 @@ impl TwcsPicker { .map(|window| window.time_window) .collect::>(); let time_windows = Arc::new(time_windows); + let max_outputs = max_outputs.unwrap_or(usize::MAX); // Enumerating all LastNonNull seeds must not increase the number of // concurrently planned windows along with the output limit. let chunk_size = self.max_background_tasks.unwrap_or(windows.len()).max(1); - let mut selected_windows = HashSet::new(); - 'phases: for phase in PICK_PHASES { - for chunk in windows.chunks(chunk_size) { - let mut handles = Vec::with_capacity(chunk.len()); - for window in chunk { - if selected_windows.contains(window) { - continue; - } - if !time_windows - .get(window) - .is_some_and(|files| phase.applies_to(files)) - { - continue; - } - let picker = self.clone(); - let time_windows = time_windows.clone(); - let window = *window; - let primary_key_mapper = primary_key_mapper.clone(); - handles.push(common_runtime::spawn_blocking_compact(move || { - time_windows.get(&window).map(|files| { - ( - window, - picker.find_inputs( - region_id, - WindowPickContext { - active_window, - files, - windows: &time_windows, - phase, - primary_key_mapper: &primary_key_mapper, - }, - ), - ) - }) - })); - tokio::task::yield_now().await; - } - for result in futures::future::join_all(handles).await { - let Some((window, (inputs, filter_deleted))) = result.context(JoinSnafu)? - else { - continue; - }; - if inputs.is_empty() { - continue; - } - - selected_windows.insert(window); - output.push(CompactionOutput { - output_level: LEVEL_COMPACTED, // always compact to l1 - inputs, - filter_deleted, - output_time_range: None, // we do not enforce output time range in twcs compactions. + let mut picks = Vec::new(); + for chunk in windows.chunks(chunk_size) { + let mut handles = Vec::with_capacity(chunk.len()); + for &window in chunk { + let picker = self.clone(); + let time_windows = time_windows.clone(); + let primary_key_mapper = primary_key_mapper.clone(); + handles.push(common_runtime::spawn_blocking_compact(move || { + let pick = time_windows.get(&window).and_then(|files| { + picker.pick_window( + region_id, + files, + active_window, + time_window_size, + region_sequence, + &primary_key_mapper, + ) }); - - if let Some(max_outputs) = max_outputs - && output.len() >= max_outputs - { - debug!( - "Region ({:?}) compaction output limit ({}) reached, remaining candidates discarded", - region_id, max_outputs - ); - break 'phases; - } - } + (window, pick) + })); + tokio::task::yield_now().await; + } + for result in futures::future::join_all(handles).await { + let (window, pick) = result.context(JoinSnafu)?; + picks.extend(pick.map(|pick| (window, pick))); + } + // Windows are visited newest first, so older windows cannot displace + // a full set of picks of the first class. + if picks + .iter() + .filter(|(_, pick)| pick.class == PickClass::L0Merge) + .count() + >= max_outputs + { + break; } } - // The compactor pops outputs from the end, preserving phase priority and - // newer-window-first pop order within each phase. + + // The stable sort keeps newer windows first within a class. + picks.sort_by_key(|(_, pick)| pick.class); + if picks.len() > max_outputs { + debug!( + "Region ({:?}) compaction output limit ({}) reached, remaining candidates discarded", + region_id, max_outputs + ); + picks.truncate(max_outputs); + } + let mut output = picks + .into_iter() + .map(|(window, pick)| { + let files = &time_windows[&window]; + let filter_deleted = !self.append_mode + && !window_has_overlap(files, &time_windows) + && !selected_overlaps_unselected(&pick.inputs, files, &primary_key_mapper); + if pick.inputs.len() > 1 { + log_pick_result( + region_id, + window, + active_window, + &pick, + files.files.len(), + self.max_output_file_size, + filter_deleted, + ); + } + CompactionOutput { + output_level: LEVEL_COMPACTED, // always compact to l1 + inputs: pick.inputs, + filter_deleted, + output_time_range: None, // we do not enforce output time range in twcs compactions. + } + }) + .collect::>(); + // The compactor pops outputs from the end. output.reverse(); Ok(output) } - fn find_inputs( + /// Picks the compaction inputs of one window, or `None` if it needs none. + fn pick_window( &self, region_id: RegionId, - context: WindowPickContext<'_>, - ) -> (Vec, bool) { - let WindowPickContext { - active_window, - files, - windows, - phase, - primary_key_mapper, - } = context; - let is_active_window = active_window == Some(files.time_window); - let window = &files.time_window; - let mut files_to_merge: Vec<_> = files.files().cloned().collect(); + window: &Window, + active_window: Option, + time_window_size: Option, + region_sequence: u64, + mapper: &PrimaryKeyRangeMapper, + ) -> Option { + let mut files: Vec<_> = window.files().cloned().collect(); + let quiet = is_quiet_window(&files, region_sequence); // Filter out large files in append mode - they won't benefit from compaction if self.append_mode && let Some(max_size) = self.max_output_file_size { - let (kept_files, ignored_files) = files_to_merge - .into_iter() - .partition(|file| file.size() <= max_size); - files_to_merge = kept_files; + let (kept_files, ignored_files) = + files.into_iter().partition(|file| file.size() <= max_size); + files = kept_files; if !ignored_files.is_empty() { info!( "Skipped {} large files in append mode for region {}, window {}, max_size: {}", ignored_files.len(), region_id, - window, + window.time_window, max_size ); } } - let (l0_files, l1_files): (Vec<_>, Vec<_>) = files_to_merge - .into_iter() - .partition(|file| file.level() == 0); - let num_l0_files = l0_files.len(); - let num_l1_files = l1_files.len(); - if !is_active_window - && files.files.len() < self.inactive_window_trigger_file_num - && num_l1_files < self.inactive_window_l1_merge_trigger + let (l0_files, l1_files): (Vec<_>, Vec<_>) = + files.into_iter().partition(|file| file.level() == 0); + if !self.all_windows_inactive && active_window == Some(window.time_window) { + self.pick_active_window(l0_files, l1_files, mapper) + } else if window.files.len() < self.inactive_window_trigger_file_num + && l1_files.len() < self.inactive_window_l1_merge_trigger { - return (vec![], false); - } - let (inputs, found_runs) = if is_active_window { - match phase { - PickPhase::HasL0 if num_l0_files >= self.trigger_file_num => pick_candidate_files( - l0_files, - self.max_output_file_size, - pick_count_first, - primary_key_mapper, - ), - PickPhase::L1FileReduction | PickPhase::L1OverlapOnly - if num_l1_files >= self.active_window_l1_merge_trigger => - { - pick_l1_candidate_files( - l1_files, - self.max_output_file_size, - phase, - primary_key_mapper, - ) - } - _ => (vec![], 0), - } + None } else { - pick_inactive_window_files( + let settling = !self.all_windows_inactive + && is_settling_window(window.time_window, active_window, time_window_size); + let num_l1_files = l1_files.len(); + self.pick_inactive_window(l0_files, l1_files, mapper, settling, quiet) + .or_else(|| { + if self.merge_mode == MergeMode::LastNonNull + || num_l1_files < self.inactive_window_l1_merge_trigger + { + return None; + } + // Include even append-mode files excluded from normal picks: + // a small-file merge must not add conflicts with large peers. + let pick = units::pick_small_l1(&window.files, mapper, self.pick_limits())?; + Some(WindowPick { + class: PickClass::FileReduction, + branch: "inactive:pick_small_l1", + inputs: pick.inputs, + found_runs: pick.units, + }) + }) + } + } + + /// The active window keeps receiving writes, so every pick must reduce its + /// file count: + /// + /// 1. Merge L0 files once they reach the trigger. + /// 2. Otherwise, merge L1 files once they reach the L1 safety trigger. L1 + /// files are merged separately from L0 files to avoid chained rewrites. + /// + /// Overlaps are left for the window to resolve once it becomes inactive: + /// resolving them while data keeps arriving rewrites the window repeatedly, + /// and leaves a dominant run that blocks the inactive-window merge. + fn pick_active_window( + &self, + l0_files: Vec, + l1_files: Vec, + mapper: &PrimaryKeyRangeMapper, + ) -> Option { + if l0_files.len() >= self.trigger_file_num { + let (inputs, found_runs) = pick_candidate_files( l0_files, - l1_files, - InactiveWindowPick { - l0_file_num: self.inactive_window_trigger_file_num, - l1_file_num: self.inactive_window_l1_merge_trigger, - phase, - primary_key_mapper, - }, self.max_output_file_size, - ) - }; - if inputs.is_empty() { - return (inputs, false); - } - - let filter_deleted = !self.append_mode - && !window_has_overlap(files, windows) - && !selected_overlaps_unselected(&inputs, files, primary_key_mapper); - - if inputs.len() > 1 { - // If we have more than one file to compact. - log_pick_result( - region_id, - *window, - active_window, - found_runs, - files.files.len(), - self.max_output_file_size, - filter_deleted, - &inputs, + pick_count_first, + mapper, ); + if !inputs.is_empty() { + return Some(WindowPick { + class: PickClass::L0Merge, + branch: "active:pick_count_first", + inputs, + found_runs, + }); + } } - (inputs, filter_deleted) + if l1_files.len() < self.active_window_l1_merge_trigger { + return None; + } + let pick = units::pick_active_l1(&l1_files, mapper, self.pick_limits())?; + Some(WindowPick { + class: PickClass::FileReduction, + branch: "active:pick_l1_file_reduction", + inputs: pick.inputs, + found_runs: pick.units, + }) } -} -/// Picks compaction inputs for an inactive window. -/// -/// The window no longer receives fresh writes (late arrivals aside), so it -/// should converge, but merging must stay within a bounded rewrite cost: -/// -/// 1. L0 candidates are selected before pure L1 candidates across all windows. -/// 2. Balanced L0 picks below the trigger continue converging without pulling -/// L1 files into the rewrite. -/// 3. A mixed merge may bypass the balance checks, but only when the total -/// rewrite fits in the output file budget, bounding write amplification. -/// 4. Last resort: converge L0 files among themselves regardless of balance. -/// The rewrite is bounded by the L0 bytes and leaves large compacted files -/// untouched. If nothing qualifies, the window is left as-is. -#[derive(Debug, Clone, Copy)] -struct InactiveWindowPick<'a> { - l0_file_num: usize, - l1_file_num: usize, - phase: PickPhase, - primary_key_mapper: &'a PrimaryKeyRangeMapper, -} - -fn pick_inactive_window_files( - l0_files: Vec, - l1_files: Vec, - pick: InactiveWindowPick<'_>, - max_output_file_size: Option, -) -> (Vec, usize) { - if !matches!(pick.phase, PickPhase::HasL0) { - return if l1_files.len() >= pick.l1_file_num { - pick_l1_candidate_files( - l1_files, - max_output_file_size, - pick.phase, - pick.primary_key_mapper, + /// An inactive window should converge at a bounded rewrite cost: + /// + /// 1. Absorb L0 files (late data), see [`pick_inactive_l0`]. + /// 2. Otherwise, merge a closed group of files that resolves overlaps or + /// removes files, see [`units::pick_closed_merge`]. Pure file reduction + /// waits for the L1 trigger; overlap resolution does not. + /// + /// A `settling` window is the one right before the active window. Late data + /// still lands there, and any rewrite of its overlapping files would be + /// repeated once more late data arrives. It only merges L0 files among + /// themselves and conflict-free files, and leaves overlaps for when the + /// window is older. + /// + /// Late data can also land in older windows. Resolving such an overlap + /// rewrites a dominant sorted run for a small peer, which would repeat for + /// every late flush, so it waits until the window is `quiet`: the region + /// has written a window's worth of rows elsewhere since the window last + /// received data. Each quiet period then pays for at most one such rewrite. + /// A manual compaction is a one-off request to resolve overlaps, so it does + /// not wait. + /// + /// Without new writes each pick consumes an L0 file or makes progress on + /// the closed merge's termination measure, so the window stops compacting. + fn pick_inactive_window( + &self, + l0_files: Vec, + l1_files: Vec, + mapper: &PrimaryKeyRangeMapper, + settling: bool, + quiet: bool, + ) -> Option { + let num_l1_files = l1_files.len(); + if !l0_files.is_empty() + && let Some((inputs, found_runs, branch)) = pick_inactive_l0( + l0_files.clone(), + (!settling).then_some(l1_files.as_slice()), + self.max_output_file_size, + mapper, ) + { + return Some(WindowPick { + class: PickClass::L0Merge, + branch, + inputs, + found_runs, + }); + } + + let mut files = l1_files; + files.extend(l0_files); + let pick = units::pick_closed_merge( + &files, + mapper, + ClosedMergeOptions { + limits: self.closed_merge_limits(), + allow_file_reduction: num_l1_files >= self.inactive_window_l1_merge_trigger, + allow_conflicts: !settling, + allow_unbalanced: (quiet || self.all_windows_inactive) && !settling, + }, + )?; + let (class, branch) = if pick.dense_fallback { + (PickClass::FileReduction, "inactive:dense_file_reduction") + } else if pick.reduces_files { + (PickClass::FileReduction, "inactive:pick_l1_file_reduction") } else { - (vec![], 0) + (PickClass::OverlapResolution, "inactive:pick_l1_overlap") }; + Some(WindowPick { + class, + branch, + inputs: pick.inputs, + found_runs: pick.units, + }) } - if l0_files.len() >= pick.l0_file_num { - let candidate = pick_candidate_files( - l0_files.clone(), - max_output_file_size, - pick_count_first, - pick.primary_key_mapper, - ); - if !candidate.0.is_empty() { - return candidate; + fn pick_limits(&self) -> PickLimits { + PickLimits { + max_output_file_size: self.max_output_file_size, + max_input_files: *MAX_INPUT_FILES, } } - let candidate = pick_candidate_files( + fn closed_merge_limits(&self) -> PickLimits { + PickLimits { + max_input_files: MAX_INPUT_FILES.saturating_mul(CLOSED_MERGE_INPUT_FACTOR), + ..self.pick_limits() + } + } +} + +/// Absorbs L0 files of an inactive window, in order of preference: +/// +/// 1. Balanced merges of L0 files, which leave L1 files untouched. +/// 2. Mixed merges that bypass the balance checks, but only when the total +/// rewrite fits in the output file budget, bounding write amplification. +/// Skipped without `l1_files`. +/// 3. Merges of L0 files among themselves regardless of balance. The rewrite +/// is bounded by the L0 bytes and leaves large compacted files untouched. +/// +/// Every pick consumes L0 files. +fn pick_inactive_l0( + l0_files: Vec, + l1_files: Option<&[FileHandle]>, + max_output_file_size: Option, + mapper: &PrimaryKeyRangeMapper, +) -> Option<(Vec, usize, &'static str)> { + let (inputs, found_runs) = pick_candidate_files( l0_files.clone(), max_output_file_size, pick_count_first, - pick.primary_key_mapper, + mapper, ); - if !candidate.0.is_empty() { - return candidate; + if !inputs.is_empty() { + return Some((inputs, found_runs, "inactive:pick_count_first")); } - let mut all_files = l0_files.clone(); - all_files.extend(l1_files); - let candidate = pick_candidate_files( - all_files, - max_output_file_size, - pick_mixed_within_budget, - pick.primary_key_mapper, - ); - if !candidate.0.is_empty() { - return candidate; + if let Some(l1_files) = l1_files { + let mut all_files = l0_files.clone(); + all_files.extend_from_slice(l1_files); + let (inputs, found_runs) = pick_candidate_files( + all_files, + max_output_file_size, + pick_mixed_within_budget, + mapper, + ); + if !inputs.is_empty() { + return Some((inputs, found_runs, "inactive:pick_mixed_within_budget")); + } } - - pick_candidate_files( + let (inputs, found_runs) = pick_candidate_files( l0_files, max_output_file_size, pick_unbalanced_count_first, - pick.primary_key_mapper, - ) -} - -fn pick_l1_candidate_files( - l1_files: Vec, - max_output_file_size: Option, - phase: PickPhase, - mapper: &PrimaryKeyRangeMapper, -) -> (Vec, usize) { - let picker = match phase { - PickPhase::L1FileReduction => pick_l1_file_reduction, - PickPhase::L1OverlapOnly => pick_l1_overlap_only, - PickPhase::HasL0 => return (vec![], 0), - }; - pick_candidate_files(l1_files, max_output_file_size, picker, mapper) + mapper, + ); + (!inputs.is_empty()).then_some((inputs, found_runs, "inactive:pick_unbalanced_count_first")) } type CandidatePicker = @@ -459,17 +526,29 @@ struct OrderedFile<'a> { file: &'a FileHandle, run_id: usize, position_in_run: usize, + /// Primary key range in the pick's schema, resolved once per pick. + primary_key_range: Option<(bytes::Bytes, bytes::Bytes)>, } -/// Metrics of a candidate compaction input, accumulated incrementally as the -/// candidate interval expands one file at a time. +impl OrderedFile<'_> { + /// Same as [`files_overlap_inclusive`] with the cached key ranges. + fn overlaps_inclusive(&self, other: &Self) -> bool { + self.file.overlap_inclusive(other.file) + && match (&self.primary_key_range, &other.primary_key_range) { + (Some(lhs), Some(rhs)) => primary_key_ranges_overlap(lhs, rhs), + _ => true, + } + } +} + +/// Metrics of a candidate compaction input, accumulated one file at a time. #[derive(Debug, Default)] struct Candidate { - /// Number of files in the interval. + /// Number of input files. num_files: usize, - /// Total input bytes of the interval. + /// Total input bytes. total_size: usize, - /// Size of the largest single file in the interval. + /// Size of the largest single input file. largest_file_size: usize, /// Files in the interval that overlap another interval file from a different /// sorted run. Resolving these overlaps is what reduces sorted runs. @@ -493,18 +572,12 @@ impl Candidate { file: &OrderedFile, preceding: &[OrderedFile], participations: &mut Vec, - mapper: &PrimaryKeyRangeMapper, ) { - self.num_files += 1; - let file_size = file.file.size() as usize; - self.total_size += file_size; - self.largest_file_size = self.largest_file_size.max(file_size); - self.absorb_level_rows(file.file); + self.absorb_file(file.file); let mut participates = false; for (offset, other) in preceding.iter().enumerate() { - if file.run_id != other.run_id && files_overlap_inclusive(file.file, other.file, mapper) - { + if file.run_id != other.run_id && file.overlaps_inclusive(other) { if !participations[offset] { participations[offset] = true; self.overlap_participants += 1; @@ -518,6 +591,15 @@ impl Candidate { } } + /// Accumulates file statistics independently of overlap accounting. + fn absorb_file(&mut self, file: &FileHandle) { + self.num_files += 1; + let file_size = file.size() as usize; + self.total_size += file_size; + self.largest_file_size = self.largest_file_size.max(file_size); + self.absorb_level_rows(file); + } + fn absorb_level_rows(&mut self, file: &FileHandle) { let num_rows = file.num_rows(); self.has_unknown_rows |= num_rows == 0; @@ -644,28 +726,6 @@ fn pick_count_first( ) } -fn pick_l1_file_reduction( - sorted_runs: Vec>, - max_output_file_size: Option, - mapper: &PrimaryKeyRangeMapper, -) -> Vec { - pick_count_first_where(sorted_runs, max_output_file_size, mapper, |candidate| { - is_balanced_candidate(candidate) && candidate.file_reduction(max_output_file_size) > 0 - }) -} - -fn pick_l1_overlap_only( - sorted_runs: Vec>, - max_output_file_size: Option, - mapper: &PrimaryKeyRangeMapper, -) -> Vec { - pick_count_first_where(sorted_runs, max_output_file_size, mapper, |candidate| { - is_balanced_candidate(candidate) - && candidate.file_reduction(max_output_file_size) == 0 - && candidate.overlap_participants > 0 - }) -} - #[cfg(test)] fn pick_mixed_count_first( sorted_runs: Vec>, @@ -710,7 +770,7 @@ fn pick_count_first_where( mapper: &PrimaryKeyRangeMapper, is_eligible: impl Fn(&Candidate) -> bool, ) -> Vec { - let files = ordered_files(&sorted_runs); + let files = ordered_files(&sorted_runs, mapper); let mut best = None; for left in 0..files.len() { @@ -718,12 +778,7 @@ fn pick_count_first_where( let right_bound = left.saturating_add(*MAX_INPUT_FILES).min(files.len()); let mut participations: Vec = Vec::with_capacity(right_bound - left); for right in left..right_bound { - candidate.absorb( - &files[right], - &files[left..right], - &mut participations, - mapper, - ); + candidate.absorb(&files[right], &files[left..right], &mut participations); if candidate.num_files < 2 || !is_eligible(&candidate) || !candidate.makes_progress(max_output_file_size) @@ -749,7 +804,10 @@ fn pick_count_first_where( /// Flattens sorted runs into files ordered by `(start asc, end desc)`, breaking ties /// by run and position to keep the order deterministic. -fn ordered_files(sorted_runs: &[SortedRun]) -> Vec> { +fn ordered_files<'a>( + sorted_runs: &'a [SortedRun], + mapper: &PrimaryKeyRangeMapper, +) -> Vec> { let mut files = sorted_runs .iter() .enumerate() @@ -761,6 +819,7 @@ fn ordered_files(sorted_runs: &[SortedRun]) -> Vec> file, run_id, position_in_run, + primary_key_range: file.primary_key_range(mapper), }) }) .collect::>(); @@ -808,18 +867,17 @@ fn selected_overlaps_unselected( }) } -#[allow(clippy::too_many_arguments)] fn log_pick_result( region_id: RegionId, window: i64, active_window: Option, - found_runs: usize, + pick: &WindowPick, file_num: usize, max_output_file_size: Option, filter_deleted: bool, - inputs: &[FileHandle], ) { - let input_file_str: Vec = inputs + let input_file_str: Vec = pick + .inputs .iter() .map(|f| { let range = f.range(); @@ -840,13 +898,14 @@ fn log_pick_result( let active_window_str = active_window.map(|s| Timestamp::new_second(s).to_iso8601_string()); let max_output_file_size = max_output_file_size.map(|size| ReadableSize(size).to_string()); info!( - "Region ({:?}) compaction pick result: current window: {}, active window: {:?}, \ + "Region ({:?}) compaction pick result: current window: {}, active window: {:?}, branch: {}, \ found runs: {}, file num: {}, max output file size: {:?}, filter deleted: {}, \ input files: {:?}", region_id, window_str, active_window_str, - found_runs, + pick.branch, + pick.found_runs, file_num, max_output_file_size, filter_deleted, @@ -1010,6 +1069,44 @@ impl Window { } } +/// Whether `window` is the one right before `active_window`, where late data +/// still lands. +fn is_settling_window( + window: i64, + active_window: Option, + time_window_size: Option, +) -> bool { + active_window + .zip(time_window_size) + .is_some_and(|(active, size)| { + // Windows at the timestamp limits may be too far apart to subtract. + active + .checked_sub(window) + .is_some_and(|distance| (1..=size).contains(&distance)) + }) +} + +/// Whether the region has written at least as many rows as `files` hold, and at +/// least one row, since the newest of them was written. Sequences are assigned +/// per row and compaction keeps the inputs' maximum, so this needs no state. +/// Files written before row counts were recorded have zero rows, so a window +/// holding them is not quiet. +fn is_quiet_window(files: &[FileHandle], region_sequence: u64) -> bool { + if files.iter().any(|file| file.num_rows() == 0) { + return false; + } + let Some(window_sequence) = files + .iter() + .filter_map(|file| file.meta_ref().sequence) + .map(|sequence| sequence.get()) + .max() + else { + return false; + }; + let rows: u64 = files.iter().map(|file| file.num_rows() as u64).sum(); + region_sequence.saturating_sub(window_sequence) >= rows.max(1) +} + /// Assigns files to windows with predefined window size (in seconds) by their max timestamps. fn assign_to_windows<'a>( files: impl Iterator, @@ -1330,8 +1427,10 @@ mod tests { time_window_seconds: Some(3), max_output_file_size, append_mode: false, + merge_mode: MergeMode::LastRow, max_background_tasks: None, time_range: None, + all_windows_inactive: false, }; let output = picker.pick(&compaction_region).await.unwrap().unwrap(); @@ -1353,8 +1452,10 @@ mod tests { time_window_seconds: Some(3), max_output_file_size: None, append_mode: false, + merge_mode: MergeMode::LastRow, max_background_tasks: None, time_range: None, + all_windows_inactive: false, }; let compaction_region = compaction_region_with_expired_sst().await; @@ -1409,8 +1510,10 @@ mod tests { time_window_seconds: Some(3), max_output_file_size: None, append_mode: false, + merge_mode: MergeMode::LastRow, max_background_tasks: None, time_range: None, + all_windows_inactive: false, }; let output = picker.pick(&compaction_region).await.unwrap().unwrap(); @@ -1471,8 +1574,10 @@ mod tests { time_window_seconds: Some(3600), max_output_file_size: None, append_mode: false, + merge_mode: MergeMode::LastRow, max_background_tasks: None, time_range: None, + all_windows_inactive: false, }; assert!(picker.pick(&compaction_region).await.unwrap().is_none()); @@ -1629,6 +1734,17 @@ mod tests { } } + #[test] + fn test_settling_window_at_timestamp_limits() { + assert!(is_settling_window(0, Some(60), Some(60))); + assert!(!is_settling_window(60, Some(60), Some(60))); + assert!(!is_settling_window(-60, Some(60), Some(60))); + assert!(!is_settling_window(0, None, Some(60))); + // The distance between these windows does not fit in an i64. + assert!(!is_settling_window(i64::MAX, Some(i64::MIN), Some(1))); + assert!(!is_settling_window(i64::MIN, Some(i64::MAX), Some(1))); + } + #[test] fn test_window_assignment_at_timestamp_limits() { for (seconds, expected_window) in [(i64::MIN, i64::MIN), (i64::MAX, i64::MAX - 1)] { @@ -2028,8 +2144,10 @@ mod tests { time_window_seconds: None, max_output_file_size: None, append_mode: false, + merge_mode: MergeMode::LastRow, max_background_tasks: None, time_range: None, + all_windows_inactive: false, } .build_output_with_time_range(RegionId::from_u64(0), windows, active_window, None) .await @@ -2179,8 +2297,10 @@ mod tests { time_window_seconds: None, max_output_file_size: None, append_mode: false, + merge_mode: MergeMode::LastRow, max_background_tasks: None, time_range: None, + all_windows_inactive: false, } .build_output_with_time_range(RegionId::from_u64(0), windows, active_window, None) .await @@ -2250,8 +2370,10 @@ mod tests { time_window_seconds: Some(3), max_output_file_size: None, append_mode: false, + merge_mode: MergeMode::LastRow, max_background_tasks: None, time_range: None, + all_windows_inactive: false, }; let active_window = find_latest_window_in_seconds(files.iter(), 3); @@ -2298,8 +2420,10 @@ mod tests { time_window_seconds: Some(3), max_output_file_size: Some(1000), append_mode: true, + merge_mode: MergeMode::LastRow, max_background_tasks: None, time_range: None, + all_windows_inactive: false, }; let active_window = find_latest_window_in_seconds(files.iter(), 3); @@ -2332,8 +2456,10 @@ mod tests { time_window_seconds: Some(1), max_output_file_size: Some(1_000), append_mode: true, + merge_mode: MergeMode::LastRow, max_background_tasks: None, time_range: None, + all_windows_inactive: false, }; let output = picker @@ -2377,8 +2503,10 @@ mod tests { time_window_seconds: Some(3), max_output_file_size: None, append_mode: false, + merge_mode: MergeMode::LastRow, max_background_tasks: Some(max_background_tasks), time_range: None, + all_windows_inactive: false, }; let active_window = find_latest_window_in_seconds(files.iter(), 3); @@ -2404,8 +2532,10 @@ mod tests { time_window_seconds: Some(3), max_output_file_size: None, append_mode: false, + merge_mode: MergeMode::LastRow, max_background_tasks: None, time_range: None, + all_windows_inactive: false, }; let windows_no_limit = assign_to_windows(files.iter(), 3); @@ -2451,8 +2581,10 @@ mod tests { time_window_seconds: Some(3), max_output_file_size: None, append_mode: false, + merge_mode: MergeMode::LastRow, max_background_tasks: Some(max_background_tasks), time_range: None, + all_windows_inactive: false, }; let active_window = find_latest_window_in_seconds(files.iter(), 3); @@ -2503,8 +2635,10 @@ mod tests { time_window_seconds: Some(3), max_output_file_size: None, append_mode: false, + merge_mode: MergeMode::LastRow, max_background_tasks: None, time_range: None, + all_windows_inactive: false, }; let active_window = find_latest_window_in_seconds(files.iter(), 3); @@ -2550,8 +2684,10 @@ mod tests { time_window_seconds: Some(3), max_output_file_size: None, append_mode: false, + merge_mode: MergeMode::LastRow, max_background_tasks: None, time_range: None, + all_windows_inactive: false, }; let active_window = find_latest_window_in_seconds(files.iter(), 3); @@ -2599,8 +2735,10 @@ mod tests { time_window_seconds: Some(3600), max_output_file_size: None, append_mode: false, + merge_mode: MergeMode::LastRow, max_background_tasks: None, time_range: None, + all_windows_inactive: false, }; let active_window = find_latest_window_in_seconds(files.iter(), 3600); @@ -2648,8 +2786,10 @@ mod tests { time_window_seconds: Some(3600), max_output_file_size: Some(1024 * 1024 * 1024), append_mode: false, + merge_mode: MergeMode::LastRow, max_background_tasks: None, time_range: None, + all_windows_inactive: false, }; let active_window = find_latest_window_in_seconds(files.iter(), 3600); @@ -2682,8 +2822,10 @@ mod tests { time_window_seconds: Some(3), max_output_file_size: None, append_mode: false, + merge_mode: MergeMode::LastRow, max_background_tasks: Some(1), time_range: None, + all_windows_inactive: false, }; let output = picker @@ -2779,11 +2921,13 @@ mod tests { time_window_seconds: Some(3), max_output_file_size: None, append_mode: false, + merge_mode: MergeMode::LastRow, max_background_tasks: Some(1), time_range: TimestampRange::new( Timestamp::new_millisecond(7_200), Timestamp::new_millisecond(7_800), ), + all_windows_inactive: false, }; let output = picker @@ -2817,8 +2961,10 @@ mod tests { time_window_seconds: Some(100), max_output_file_size: None, append_mode: false, + merge_mode: MergeMode::LastRow, max_background_tasks: None, time_range: None, + all_windows_inactive: false, }; let output = picker @@ -2830,79 +2976,6 @@ mod tests { assert_eq!(2, output[0].inputs.len()); } - #[test] - fn test_inactive_window_l1_trigger_is_independent_of_l0_trigger() { - let l0_files = (0..6) - .map(|idx| { - new_file_handle_with_size_and_sequence( - FileId::random(), - idx * 20, - idx * 20 + 10, - 0, - idx as u64 + 1, - 10, - ) - }) - .collect(); - let l1_files = (0..2) - .map(|idx| { - new_file_handle_with_size_and_sequence( - FileId::random(), - idx * 20, - idx * 20 + 10, - 1, - idx as u64 + 10, - 10, - ) - }) - .collect(); - - let (inputs, _) = pick_inactive_window_files( - l0_files, - l1_files, - InactiveWindowPick { - l0_file_num: 8, - l1_file_num: 2, - phase: PickPhase::L1FileReduction, - primary_key_mapper: &primary_key_mapper_for_test(), - }, - None, - ); - - assert_eq!(2, inputs.len()); - assert!(inputs.iter().all(|file| file.level() == 1)); - } - - #[test] - fn test_inactive_window_does_not_fallback_to_l1_below_trigger() { - let l1_files = (0..2) - .map(|idx| { - new_file_handle_with_size_and_sequence( - FileId::random(), - idx * 20, - idx * 20 + 10, - 1, - idx as u64 + 1, - 10, - ) - }) - .collect(); - - let (inputs, _) = pick_inactive_window_files( - vec![], - l1_files, - InactiveWindowPick { - l0_file_num: 2, - l1_file_num: 8, - phase: PickPhase::L1FileReduction, - primary_key_mapper: &primary_key_mapper_for_test(), - }, - None, - ); - - assert!(inputs.is_empty()); - } - #[tokio::test] async fn test_inactive_window_requires_a_level_trigger_before_mixed_fallback() { let files = [ @@ -2918,8 +2991,10 @@ mod tests { time_window_seconds: Some(100), max_output_file_size: Some(1_000), append_mode: false, + merge_mode: MergeMode::LastRow, max_background_tasks: None, time_range: None, + all_windows_inactive: false, }; let output = picker @@ -2947,8 +3022,10 @@ mod tests { time_window_seconds: Some(100), max_output_file_size: None, append_mode: false, + merge_mode: MergeMode::LastRow, max_background_tasks: None, time_range: None, + all_windows_inactive: false, }; let output = picker @@ -2976,8 +3053,10 @@ mod tests { time_window_seconds: Some(100), max_output_file_size: None, append_mode: false, + merge_mode: MergeMode::LastRow, max_background_tasks: None, time_range: None, + all_windows_inactive: false, }; let output = picker @@ -3004,8 +3083,10 @@ mod tests { time_window_seconds: Some(100), max_output_file_size: None, append_mode: false, + merge_mode: MergeMode::LastRow, max_background_tasks: None, time_range: None, + all_windows_inactive: false, }; let output = picker @@ -3047,8 +3128,10 @@ mod tests { time_window_seconds: Some(1000), max_output_file_size: None, append_mode: false, + merge_mode: MergeMode::LastRow, max_background_tasks: None, time_range: None, + all_windows_inactive: false, }; let output = picker @@ -3229,8 +3312,10 @@ mod tests { time_window_seconds: Some(1), max_output_file_size: None, append_mode: false, + merge_mode: MergeMode::LastRow, max_background_tasks: None, time_range: None, + all_windows_inactive: false, }; for (case, files, expected_level, expected_len) in [ @@ -3280,8 +3365,10 @@ mod tests { time_window_seconds: Some(1), max_output_file_size: None, append_mode: false, + merge_mode: MergeMode::LastRow, max_background_tasks: None, time_range: None, + all_windows_inactive: false, }; let output = picker @@ -3317,8 +3404,10 @@ mod tests { time_window_seconds: Some(1), max_output_file_size: None, append_mode: false, + merge_mode: MergeMode::LastRow, max_background_tasks: None, time_range: None, + all_windows_inactive: false, }; let output = picker @@ -3367,8 +3456,10 @@ mod tests { time_window_seconds: Some(1), max_output_file_size: Some(512), append_mode: false, + merge_mode: MergeMode::LastRow, max_background_tasks: None, time_range: None, + all_windows_inactive: false, }; let output = picker @@ -3414,8 +3505,10 @@ mod tests { time_window_seconds: Some(1), max_output_file_size: Some(100), append_mode: false, + merge_mode: MergeMode::LastRow, max_background_tasks: Some(max_background_tasks), time_range: None, + all_windows_inactive: false, } } @@ -3428,6 +3521,209 @@ mod tests { .value() } + #[tokio::test] + async fn test_manual_compaction_resolves_active_window_overlaps() { + let files = priority_test_files(5, 1, 100, true); + for (all_windows_inactive, expected) in [(false, 0), (true, 1)] { + let picker = TwcsPicker { + all_windows_inactive, + ..priority_test_picker(1) + }; + let output = picker + .build_output_with_time_range( + RegionId::from_u64(1), + assign_to_windows(files.iter(), 1), + Some(5), + None, + ) + .await + .unwrap(); + assert_eq!(expected, output.len(), "manual: {all_windows_inactive}"); + } + } + + #[tokio::test] + async fn test_shared_boundary_l1_files_resolve_in_inactive_and_manual_windows() { + let files = [(b"a", b"m"), (b"m", b"z")].map(|(min, max)| { + new_file_handle_with_size_sequence_and_primary_key_range( + FileId::random(), + 0, + 100, + 1, + 1, + 100, + pk_range(min, max), + ) + }); + for (active_window, all_windows_inactive, expected) in + [(0, false, 0), (0, true, 1), (3, false, 1)] + { + let picker = TwcsPicker { + all_windows_inactive, + ..priority_test_picker(1) + }; + let output = picker + .build_output_with_time_range( + RegionId::from_u64(1), + assign_to_windows(files.iter(), 1), + Some(active_window), + Some(1), + ) + .await + .unwrap(); + assert_eq!( + expected, + output.len(), + "active: {active_window}, manual: {all_windows_inactive}" + ); + if let Some(output) = output.first() { + assert_eq!(2, output.inputs.len()); + } + } + } + + #[tokio::test] + async fn test_settling_window_defers_overlap_resolution() { + let files = priority_test_files(3, 1, 100, true); + for (active_window, all_windows_inactive, expected) in + [(4, false, 0), (5, false, 1), (2, false, 1), (4, true, 1)] + { + let picker = TwcsPicker { + all_windows_inactive, + ..priority_test_picker(1) + }; + let output = picker + .build_output_with_time_range( + RegionId::from_u64(1), + assign_to_windows(files.iter(), 1), + Some(active_window), + Some(1), + ) + .await + .unwrap(); + assert_eq!( + expected, + output.len(), + "active: {active_window}, manual: {all_windows_inactive}" + ); + } + } + + #[tokio::test] + async fn test_inactive_windows_merge_small_l1_without_large_peer() { + let large = + new_file_handle_with_size_and_sequence(FileId::random(), 3_000, 3_999, 1, 1, 1_000); + let small = priority_test_files(3, 1, 20, true); + let files = [vec![large.clone()], small.to_vec()].concat(); + // The immediate predecessor is settling; the older window is not quiet. + for (active_window, merge_mode, trigger, expected) in [ + (4, MergeMode::LastRow, 2, 1), + (5, MergeMode::LastRow, 2, 1), + (4, MergeMode::LastNonNull, 2, 0), + (5, MergeMode::LastNonNull, 2, 0), + (4, MergeMode::LastRow, 4, 0), + ] { + let picker = TwcsPicker { + merge_mode, + inactive_window_l1_merge_trigger: trigger, + ..priority_test_picker(1) + }; + let output = picker + .build_output_with_time_range( + RegionId::from_u64(1), + assign_to_windows(files.iter(), 1), + Some(active_window), + Some(1), + ) + .await + .unwrap(); + assert_eq!( + expected, + output.len(), + "active: {active_window}, mode: {merge_mode:?}, trigger: {trigger}" + ); + if expected == 0 { + continue; + } + let inputs: HashSet<_> = output[0].inputs.iter().map(FileHandle::file_id).collect(); + assert_eq!( + small + .iter() + .map(FileHandle::file_id) + .collect::>(), + inputs + ); + assert!( + !output[0].filter_deleted, + "the unselected large file still overlaps" + ); + } + } + + #[tokio::test] + async fn test_small_l1_fallback_checks_large_append_mode_peer() { + // Two conflicting small files form a non-rectangular area. Their union + // would introduce a conflict with the large file in the empty corner. + let files = [ + (0, 20, b"a", b"m", 20), + (10, 30, b"l", b"z", 20), + (0, 5, b"y", b"z", 1_000), + ] + .map(|(start, end, min, max, size)| { + new_file_handle_with_size_sequence_and_primary_key_range( + FileId::random(), + start, + end, + 1, + 1, + size, + pk_range(min, max), + ) + }); + for append_mode in [false, true] { + let picker = TwcsPicker { + append_mode, + ..priority_test_picker(1) + }; + let output = picker + .build_output_with_time_range( + RegionId::from_u64(1), + assign_to_windows(files.iter(), 1), + Some(1), + Some(1), + ) + .await + .unwrap(); + assert!(output.is_empty(), "append mode: {append_mode}"); + } + } + + #[tokio::test] + async fn test_settling_window_does_not_rewrite_l1_for_late_data() { + // A late L0 file overlapping both L1 files, which a mixed merge absorbs. + let mut files = priority_test_files(3, 1, 40, false).to_vec(); + files.push(new_file_handle_with_size_and_sequence( + FileId::random(), + 3_000, + 3_299, + 0, + 10, + 10, + )); + for (active_window, expected) in [(4, 0), (5, 1)] { + let output = priority_test_picker(1) + .build_output_with_time_range( + RegionId::from_u64(1), + assign_to_windows(files.iter(), 1), + Some(active_window), + Some(1), + ) + .await + .unwrap(); + assert_eq!(expected, output.len(), "active: {active_window}"); + } + } + #[tokio::test] async fn test_older_l0_candidate_has_priority_over_newer_pure_l1_candidate() { let files = priority_test_files(0, 0, 40, false) @@ -3447,7 +3743,7 @@ mod tests { } #[tokio::test] - async fn test_older_l1_file_reduction_has_priority_over_newer_overlap_only_l1() { + async fn test_older_l1_file_reduction_has_priority_over_newer_overlap_resolution() { let files = priority_test_files(0, 1, 40, false) .into_iter() .chain(priority_test_files(1, 1, 100, true)) @@ -3488,7 +3784,9 @@ mod tests { .map(|output| (output.inputs[0].level(), output_window(output))) .collect::>(); - assert_eq!(vec![(0, 1), (0, 0), (1, 3), (1, 2), (1, 5)], pop_priority); + // The active window 5 only reduces files, so the overlap slot goes to + // the inactive window 4. + assert_eq!(vec![(0, 1), (0, 0), (1, 3), (1, 2), (1, 4)], pop_priority); let units = crate::compaction::unit::CompactionUnit::from_picker(PickerOutput { outputs: output, time_window_size: 1, @@ -3525,8 +3823,10 @@ mod tests { time_window_seconds: Some(100), max_output_file_size: Some(2_000), append_mode: false, + merge_mode: MergeMode::LastRow, max_background_tasks: None, time_range: None, + all_windows_inactive: false, }; let output = picker @@ -3556,8 +3856,10 @@ mod tests { time_window_seconds: Some(100), max_output_file_size: Some(100_000), append_mode: false, + merge_mode: MergeMode::LastRow, max_background_tasks: None, time_range: None, + all_windows_inactive: false, }; let output = picker @@ -3587,8 +3889,10 @@ mod tests { time_window_seconds: Some(100), max_output_file_size: Some(100_000), append_mode: false, + merge_mode: MergeMode::LastRow, max_background_tasks: None, time_range: None, + all_windows_inactive: false, }; let output = picker @@ -3618,8 +3922,10 @@ mod tests { time_window_seconds: Some(100), max_output_file_size: Some(1_000_000), append_mode: false, + merge_mode: MergeMode::LastRow, max_background_tasks: None, time_range: None, + all_windows_inactive: false, }; let output = picker @@ -3710,6 +4016,436 @@ mod tests { assert_eq!(vec![(0, 19), (10, 29)], picked_ranges(&picked)); } + #[tokio::test] + async fn test_l1_partial_pk_overlap_does_not_rewrite_unrelated_files() { + // Region 1025's 2026-09-23 R2: A1..A8 are R1's outputs, plus one L1 E. + // Keep the logged, rounded sizes (tenths of MiB) and normalize timestamps + // while preserving the observed order: A3, A7, A6, A1, A4, A5, A2, A8, E. + // The PK ranges are hypothetical: the logs only contain their lengths. + let a_files = [ + (4, 5249), // A1 + (7, 5218), // A2 + (1, 5255), // A3 + (5, 5263), // A4 + (6, 5131), // A5 + (3, 5143), // A6 + (2, 5197), // A7 + (8, 4716), // A8 + ] + .into_iter() + .enumerate() + .map(|(idx, (start, size))| { + let key = b'a' + idx as u8; + new_file_handle_with_size_sequence_and_primary_key_range( + FileId::random(), + start, + 10_000, + 1, + 100, + size * 1024 * 1024 / 10, + pk_range(&[key, b'0'], &[key, b'9']), + ) + }) + .collect::>(); + let mapper = primary_key_mapper_for_test(); + for (idx, file) in a_files.iter().enumerate() { + for other in &a_files[idx + 1..] { + assert!(file.overlap_inclusive(other)); // Their time ranges overlap. + assert!(!files_overlap_inclusive(file, other, &mapper)); + } + } + + let picker = TwcsPicker { + trigger_file_num: 4, + active_window_l1_merge_trigger: 4, + inactive_window_trigger_file_num: 2, + inactive_window_l1_merge_trigger: 8, + time_window_seconds: Some(60), + max_output_file_size: Some(512 * 1024 * 1024), + append_mode: false, + merge_mode: MergeMode::LastRow, + max_background_tasks: Some(1), + time_range: None, + all_windows_inactive: false, + }; + + for (case, key_range, extra_size, overlapping_a, expected_inputs) in [ + ( + "A3 and E skip the unrelated A files", + Some((b'c', b'c')), + 5155, + &[2][..], + &[2, 8][..], + ), + ( + "A8 and E must not pull in A2", + Some((b'h', b'h')), + 5155, + &[7][..], + &[7, 8][..], + ), + ( + "E overlaps no A: no progress", + Some((b'z', b'z')), + 5155, + &[][..], + &[] as &[usize], + ), + ( + "A3 and E may split into one extra tail file", + Some((b'c', b'c')), + 5255, + &[2][..], + &[2, 8][..], + ), + ( + "two A files exceed E's run budget", + Some((b'b', b'c')), + 5155, + &[1, 2][..], + &[][..], + ), + ( + "sufficient E bytes allow two A files separated in time order", + Some((b'b', b'c')), + 6000, + &[1, 2][..], + &[1, 2, 8][..], + ), + ( + "unknown PK bounds include the whole dominant run in the budget", + None, + 5155, + &[0, 1, 2, 3, 4, 5, 6, 7][..], + &[][..], + ), + ] { + let extra = new_file_handle_with_size_sequence_and_primary_key_range( + FileId::random(), + 1_000, + 9_000, + 1, + 100, + extra_size * 1024 * 1024 / 10, + key_range.and_then(|(min, max)| pk_range(&[min, b'0'], &[max, b'9'])), + ); + for (idx, file) in a_files.iter().enumerate() { + assert!(file.overlap_inclusive(&extra)); + assert_eq!( + overlapping_a.contains(&idx), + files_overlap_inclusive(file, &extra, &mapper), + "{case}: A{}", + idx + 1, + ); + } + + let mut files = a_files.clone(); + files.push(extra); + for active_window in [Some(0), None] { + // The active window leaves pure overlap resolution to the window + // becoming inactive. + let expected_inputs = if active_window.is_some() { + &[][..] + } else { + expected_inputs + }; + let expected = expected_inputs + .iter() + .map(|&idx| files[idx].file_id()) + .collect::>(); + let windows = assign_to_windows(files.iter(), 60); + let outputs = picker + .build_output_with_time_range( + RegionId::from_u64(0), + windows, + active_window, + Some(60), + ) + .await + .unwrap(); + + assert_eq!( + usize::from(!expected_inputs.is_empty()), + outputs.len(), + "{case}, active window: {active_window:?}" + ); + let actual = outputs + .iter() + .flat_map(|output| output.inputs.iter().map(FileHandle::file_id)) + .collect::>(); + assert_eq!(expected, actual, "{case}, active window: {active_window:?}"); + } + } + } + + #[rstest::rstest] + #[case::active_filtered(Some(0), false, 1000, &[0, 2])] + #[case::inactive_filtered(None, false, 1000, &[0, 2])] + #[case::older_reduction_before_newer_overlap(Some(3), true, 1000, &[0, 2])] + #[case::contiguous_has_more_reduction(Some(0), false, 90, &[0, 1, 2, 3])] + #[case::filtered_has_equal_reduction_with_less_io(Some(0), false, 600, &[0, 2])] + #[tokio::test] + async fn test_l1_file_reduction_considers_filtered_candidates( + #[case] active_window: Option, + #[case] newer_overlap_window: bool, + #[case] unrelated_size: u64, + #[case] expected_indices: &[usize], + ) { + // Only A/B overlap. With 1000-byte unrelated files between/after them, + // no contiguous interval reduces the file count, but A+B alone does. + let mut files = [ + (b'a', 100), + (b'b', unrelated_size), + (b'a', 100), + (b'c', unrelated_size), + ] + .into_iter() + .enumerate() + .map(|(idx, (key, size))| { + new_file_handle_with_size_sequence_and_primary_key_range( + FileId::random(), + idx as i64 * 100, + 749, + 1, + 1, + size, + pk_range(&[key], &[key]), + ) + }) + .collect::>(); + let expected = expected_indices + .iter() + .map(|&idx| files[idx].file_id()) + .collect::>(); + if newer_overlap_window { + // One slot: a filtered file-reducing pick must beat a newer pure + // rewrite, rather than being deferred behind overlap resolution. + files.extend(priority_test_files(1, 1, 512, true)); + } + let picker = TwcsPicker { + max_output_file_size: Some(512), + ..priority_test_picker(1) + }; + let outputs = picker + .build_output_with_time_range( + RegionId::from_u64(0), + assign_to_windows(files.iter(), 1), + active_window, + Some(1), + ) + .await + .unwrap(); + + assert_eq!(1, outputs.len()); + assert_eq!( + expected, + outputs[0] + .inputs + .iter() + .map(FileHandle::file_id) + .collect::>() + ); + } + + #[tokio::test] + async fn test_overlap_resolution_limits_dominant_run_rewrites() { + let file_size = 512 * 1024 * 1024; + // Peers are full-size files in one run, so a pure rewrite keeps the file count. + for (case, peer_sizes, max_output_size, should_pick) in [ + ( + "small peer cannot rewrite any slice of the large run", + vec![file_size], + file_size, + false, + ), + ( + "one byte below the run budget", + [vec![file_size; 6], vec![file_size - 1]].concat(), + file_size, + false, + ), + ( + "exact run budget boundary", + vec![file_size; 7], + file_size, + true, + ), + ("equal sized runs", vec![file_size; 8], file_size, true), + ( + "file reduction keeps its own admission", + vec![file_size], + 2 * file_size, + true, + ), + ] { + let mut files = (0..8) + .map(|idx| { + let key = b'a' + idx as u8; + new_file_handle_with_size_sequence_and_primary_key_range( + FileId::random(), + idx, + 10_000, + 1, + 100, + file_size, + pk_range(&[key, b'0'], &[key, b'9']), + ) + }) + .collect::>(); + files.extend(peer_sizes.into_iter().enumerate().map(|(idx, size)| { + let start = 1_000 * (idx as i64 + 1); + new_file_handle_with_size_sequence_and_primary_key_range( + FileId::random(), + start, + start + 999, + 1, + 100, + size, + pk_range(b"a0", b"h9"), + ) + })); + // Only the last case removes files; the active window leaves pure + // overlap resolution to the window becoming inactive. + let reduces_files = max_output_size > file_size; + let picker = TwcsPicker { + trigger_file_num: 4, + active_window_l1_merge_trigger: 4, + inactive_window_trigger_file_num: 2, + inactive_window_l1_merge_trigger: 8, + time_window_seconds: Some(60), + max_output_file_size: Some(max_output_size), + append_mode: false, + merge_mode: MergeMode::LastRow, + max_background_tasks: Some(1), + time_range: None, + all_windows_inactive: false, + }; + + for active_window in [Some(0), None] { + let expected = if should_pick && (reduces_files || active_window.is_none()) { + files + .iter() + .map(FileHandle::file_id) + .collect::>() + } else { + HashSet::new() + }; + let outputs = picker + .build_output_with_time_range( + RegionId::from_u64(0), + assign_to_windows(files.iter(), 60), + active_window, + Some(60), + ) + .await + .unwrap(); + let actual = outputs + .iter() + .flat_map(|output| output.inputs.iter().map(FileHandle::file_id)) + .collect::>(); + assert_eq!(expected, actual, "{case}, active window: {active_window:?}"); + } + } + } + + #[tokio::test] + async fn test_quiet_or_manual_window_resolves_dominant_run_overlap() { + let file_size = 512 * 1024 * 1024; + let l1_file = + |start, end, sequence, size, range: Option<(bytes::Bytes, bytes::Bytes)>, num_rows| { + let file = new_file_handle_with_size_sequence_and_primary_key_range( + FileId::random(), + start, + end, + 1, + sequence, + size, + range.clone(), + ); + FileHandle::new_with_primary_key_range( + FileMeta { + num_rows, + ..file.meta_ref().clone() + }, + crate::test_util::new_noop_file_purger(), + range, + ) + }; + + // The window holds 9 rows. It is quiet once the region writes 9 more + // rows after it, and only if its row counts are known. A manual + // compaction does not wait, even for the active window. + for (newer_sequence, num_rows, manual, resolves) in [ + (108, 1, false, false), + (109, 1, false, true), + (109, 0, false, false), + (100, 1, true, true), + ] { + // A run of eight key-sliced files and one small file overlapping all + // of them: resolving the overlap rewrites the run for its small peer. + let mut window_files = (0..8) + .map(|idx| { + let key = b'a' + idx as u8; + let range = pk_range(&[key, b'0'], &[key, b'9']); + l1_file(idx, 10_000, 100, file_size, range, num_rows) + }) + .collect::>(); + window_files.push(l1_file( + 1_000, + 1_999, + 100, + file_size / 8, + pk_range(b"a0", b"h9"), + num_rows, + )); + let expected = window_files + .iter() + .map(FileHandle::file_id) + .collect::>(); + let picker = TwcsPicker { + trigger_file_num: 4, + active_window_l1_merge_trigger: 4, + inactive_window_trigger_file_num: 2, + inactive_window_l1_merge_trigger: 16, + time_window_seconds: Some(60), + max_output_file_size: Some(file_size), + append_mode: false, + merge_mode: MergeMode::LastRow, + max_background_tasks: Some(1), + time_range: None, + all_windows_inactive: manual, + }; + let mut files = window_files; + files.push(l1_file( + 120_000, + 120_999, + newer_sequence, + file_size, + pk_range(b"a0", b"h9"), + 1, + )); + let outputs = picker + .build_output_with_time_range( + RegionId::from_u64(0), + assign_to_windows(files.iter(), 60), + manual.then_some(0), + Some(60), + ) + .await + .unwrap(); + let actual = outputs + .iter() + .flat_map(|output| output.inputs.iter().map(FileHandle::file_id)) + .collect::>(); + let case = format!("sequence {newer_sequence}, rows {num_rows}, manual {manual}"); + if resolves { + assert_eq!(expected, actual, "{case}"); + } else { + assert!(actual.is_empty(), "{case}: {actual:?}"); + } + } + } + #[test] fn test_count_first_prefers_guaranteed_reduction_over_pure_rewrite() { // A window mixing large and small files: the large triple on its own is a @@ -3820,8 +4556,10 @@ mod tests { time_window_seconds: Some(100), max_output_file_size: None, append_mode: false, + merge_mode: MergeMode::LastRow, max_background_tasks: None, time_range: None, + all_windows_inactive: false, }; let output = picker diff --git a/src/mito2/src/compaction/twcs/units.rs b/src/mito2/src/compaction/twcs/units.rs new file mode 100644 index 00000000000..eeb60153614 --- /dev/null +++ b/src/mito2/src/compaction/twcs/units.rs @@ -0,0 +1,2073 @@ +// 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. + +//! Unit-based compaction picks for one TWCS window. +//! +//! Picks select a contiguous interval of *units* in a fixed order. Grouping +//! prefers rewriting related files together, but input limits and fallback +//! paths may split groups or use single-file units. Separate admission checks +//! preserve each pick's progress guarantees under the assumptions below. +//! +//! - [`pick_active_l1`] reduces the L1 file count of the active window. +//! - [`pick_small_l1`] reduces small L1 files without rewriting large peers. +//! - [`pick_closed_merge`] only merges intervals of an inactive window whose +//! outputs cannot add conflicts. Compaction outputs stay inside the bounding +//! box of the inputs, so they can only overlap files intruding on that box. +//! +//! When a window's split-unit conflict graph is too dense to materialize, +//! [`pick_closed_merge`] falls back to a bounded file reduction and admits it +//! only when the same conflict bound holds. +//! +//! # Termination +//! +//! Let `P` be the number of conflicting file pairs ([`Extent::conflicts`]) in an +//! inactive window and `N` its file count. A merge removes every pair among its +//! inputs. Its outputs are one sorted run with no internal conflicts, +//! so they only conflict with the `I` files intruding on the inputs' box, at +//! most `outputs · I` pairs. An inactive pick only merges inputs that had at +//! least as many pairs with other files, so `P` does not grow; without +//! intruders the inputs are closed and add no pairs. Every inactive pick either +//! resolves a pair (`P` drops) or removes a file (`P` does not grow and `N` +//! drops), so `(P, N)` strictly decreases in lexicographic order until the +//! window receives new writes. Active-window picks may add conflicts, so each +//! of them must remove a file: `N` strictly decreases. +//! +//! The argument relies on how the compactor writes outputs: +//! +//! - Outputs hold a subset of the input rows, so each output's time and key +//! bounds lie inside the inputs' bounding box. +//! - Outputs are one stream in (key, time) order. For regions with a primary +//! key, the writer only splits between complete series, so output key ranges +//! are disjoint. Without a primary key, outputs have disjoint or touching +//! time ranges. Neither is an [`Extent::conflicts`] pair. +//! - A split only starts a new file once the current one reaches the size +//! threshold, and outputs are not larger than inputs in total, so +//! [`PickLimits::predicted_outputs`] bounds the output count. +//! +//! `test_closed_merges_terminate` and `test_active_l1_picks_terminate` replay +//! random windows with outputs written under these rules. + +use std::cmp::Reverse; +use std::ops::Range; + +use bytes::Bytes; +use common_time::Timestamp; + +use crate::sst::file::FileHandle; +use crate::sst::primary_key::PrimaryKeyRangeMapper; + +/// An L1 file reduction may rewrite at most this many output files' worth of +/// bytes per SST it removes. +const MAX_REWRITE_FILES_PER_REDUCED_FILE: u64 = 4; + +/// A small L1 file occupies at most one quarter of the output size budget. +const SMALL_L1_SIZE_FACTOR: u64 = 4; + +/// Maximum number of index nodes visited while building conflict structures for +/// a pick. Above this, the picker falls back to a bounded file reduction instead +/// of materializing an unbounded conflict graph. +const CONFLICT_GRAPH_BUDGET: usize = 1 << 20; + +/// Output size limits and input limits shared by both picks. +#[derive(Debug, Clone, Copy)] +pub(crate) struct PickLimits { + /// Split threshold of compaction outputs; `None` writes a single output. + pub(crate) max_output_file_size: Option, + /// Maximum number of input files of one compaction. + pub(crate) max_input_files: usize, +} + +impl PickLimits { + /// Upper bound of the outputs of compacting an interval: the writer only + /// starts a new file after the current one reaches the threshold, and never + /// splits a series, so files of a single primary key compact into one file. + fn predicted_outputs(&self, stats: &IntervalStats) -> usize { + match self.max_output_file_size { + Some(max) if max > 0 && stats.single_key.is_none() => { + stats.size.div_ceil(max).max(1) as usize + } + _ => 1, + } + } + + /// Whether rewriting `size` bytes to remove `reduction` files stays within + /// the rewrite budget. An unset output size imposes no bound. + fn within_reduction_cost(&self, size: u64, reduction: usize) -> bool { + match self.max_output_file_size { + Some(max) if max > 0 => { + size <= (reduction as u64) + .saturating_mul(MAX_REWRITE_FILES_PER_REDUCED_FILE) + .saturating_mul(max) + } + _ => true, + } + } + + /// Bytes a dominant sorted run may exceed its peers by in an overlap merge. + fn run_tolerance(&self) -> u64 { + self.max_output_file_size.unwrap_or(0) + } +} + +/// A selected compaction input. +#[derive(Debug)] +pub(crate) struct UnitPick { + pub(crate) inputs: Vec, + /// Whether the pick is predicted to reduce the window's file count. Otherwise + /// it only resolves overlaps. + pub(crate) reduces_files: bool, + /// Number of units the candidate files were grouped into. + pub(crate) units: usize, + /// Whether this pick came from the bounded dense-window fallback. + pub(crate) dense_fallback: bool, +} + +/// Picks L1 files of the active window whose compaction removes files. +/// +/// Units are groups of conflicting files, so a merge does not leave a file +/// behind while rewriting its overlapping peers. A group larger than the input +/// limit is split back into single files. If no interval of groups qualifies, +/// for example because a large file dominates its group, the groups are split +/// into single files, so the smaller files can still be merged. Intervals +/// follow [`sort_units`]; the pick removes the most files, then rewrites the +/// fewest bytes. +pub(crate) fn pick_active_l1( + l1_files: &[FileHandle], + mapper: &PrimaryKeyRangeMapper, + options: PickLimits, +) -> Option { + pick_active_l1_with_budget(l1_files, mapper, options, CONFLICT_GRAPH_BUDGET) +} + +/// [`pick_active_l1`] with an explicit work budget, so tests can exercise the +/// single-file fallback without a window-sized conflict graph. +fn pick_active_l1_with_budget( + l1_files: &[FileHandle], + mapper: &PrimaryKeyRangeMapper, + options: PickLimits, + work_budget: usize, +) -> Option { + let extents = extents_of(l1_files, mapper); + let index = ExtentIndex::new(&extents); + let mut units = match conflict_units(l1_files, &extents, &index, work_budget) { + Some(units) => units, + None => single_file_units(l1_files, &extents), + }; + split_oversized_units(&mut units, l1_files, &extents, options.max_input_files); + if let Some(pick) = pick_reducing_interval(&mut units, l1_files, options) { + return Some(pick); + } + if units.len() == l1_files.len() { + return None; + } + split_oversized_units(&mut units, l1_files, &extents, 1); + pick_reducing_interval(&mut units, l1_files, options) +} + +/// Picks the interval of `units` that removes the most files within the +/// rewrite budget, then rewrites the fewest bytes. +fn pick_reducing_interval( + units: &mut [Unit], + files: &[FileHandle], + options: PickLimits, +) -> Option { + sort_units(units); + let range = pick_reducing_interval_where(units, options, |_| true)?; + Some(UnitPick { + inputs: collect_inputs(&units[range], files), + reduces_files: true, + units: units.len(), + dense_fallback: false, + }) +} + +/// Selects a reducing interval of sorted units, with an additional filter so +/// the dense fallback can retry without duplicating the reduction guards. +fn pick_reducing_interval_where( + units: &[Unit], + options: PickLimits, + is_eligible: impl Fn(&IntervalStats) -> bool, +) -> Option> { + let mut best: Option<(usize, u64, Range)> = None; + // File reductions ignore conflicts, so split files are not counted. + for_each_interval(units, &[], options.max_input_files, |range, stats| { + let reduction = stats + .num_files + .saturating_sub(options.predicted_outputs(stats)); + if stats.num_files < 2 + || reduction == 0 + || !options.within_reduction_cost(stats.size, reduction) + || stats.largest_file_size > stats.size - stats.largest_file_size + || !is_eligible(stats) + { + return; + } + let better = best.as_ref().is_none_or(|(best_reduction, best_size, _)| { + (reduction, Reverse(stats.size)) > (*best_reduction, Reverse(*best_size)) + }); + if better { + best = Some((reduction, stats.size, range)); + } + }); + + best.map(|(_, _, range)| range) +} + +/// Reduces small L1 files left behind by a closed merge, including in settling +/// windows. The input bytes fit in one output and no input dominates its peers, +/// so each merge removes files without repeatedly rewriting an accumulator. +/// +/// Large and L0 files remain in the external conflict index. Check only the +/// best reducing interval to bound planning work; rejecting it can miss other +/// admissible intervals, but cannot increase the window's conflict count. +pub(crate) fn pick_small_l1( + files: &[FileHandle], + mapper: &PrimaryKeyRangeMapper, + limits: PickLimits, +) -> Option { + let max_output = limits.max_output_file_size.filter(|size| *size > 0)?; + let extents = extents_of(files, mapper); + let mut units = single_file_units(files, &extents); + units.retain(|unit| { + let file = &files[unit.files[0]]; + file.level() == 1 && file.size() <= max_output / SMALL_L1_SIZE_FACTOR + }); + sort_units(&mut units); + let range = pick_reducing_interval_where(&units, limits, |stats| stats.size <= max_output)?; + let index = ExtentIndex::new(&extents); + if !adds_no_conflicts(&units[range.clone()], &extents, &index, 1) { + return None; + } + Some(UnitPick { + inputs: collect_inputs(&units[range], files), + reduces_files: true, + units: units.len(), + dense_fallback: false, + }) +} + +/// Thresholds of the closed merge pick. +#[derive(Debug, Clone, Copy)] +pub(crate) struct ClosedMergeOptions { + pub(crate) limits: PickLimits, + /// Whether merges that only remove files are allowed. + pub(crate) allow_file_reduction: bool, + /// Whether merges may rewrite conflicting files. Otherwise only + /// conflict-free intervals are merged. + pub(crate) allow_conflicts: bool, + /// Whether overlap merges may rewrite a dominant sorted run for a small + /// peer. Otherwise such a merge would be repeated whenever late data adds + /// another small peer. + pub(crate) allow_unbalanced: bool, +} + +/// Picks an interval of units among all files of an inactive window. +/// +/// Units are the closures of conflicting files: no file outside a unit intrudes +/// on its bounding box. A closure larger than the input limit is split back into +/// single files, so files around a late file spanning a large sorted run can +/// still be merged. Intervals follow [`sort_units`]. An interval is eligible if +/// its outputs add no more conflicts than its files have with other files (see +/// [`adds_no_conflicts`]), and it +/// +/// - removes files within the rewrite budget, or +/// - resolves conflicts without rewriting a dominant sorted run (unless +/// [`ClosedMergeOptions::allow_unbalanced`]), and adds at most one tail file. +/// +/// The pick maximizes resolved conflicts plus removed files per rewritten byte. +/// A window whose split-unit conflict graph is too dense to materialize falls +/// back to a bounded file reduction; see [`pick_dense_closed_merge`]. +pub(crate) fn pick_closed_merge( + files: &[FileHandle], + mapper: &PrimaryKeyRangeMapper, + options: ClosedMergeOptions, +) -> Option { + pick_closed_merge_with_budget(files, mapper, options, CONFLICT_GRAPH_BUDGET) +} + +/// [`pick_closed_merge`] with an explicit work budget, so tests can exercise +/// the dense fallback without building a window-sized conflict graph. +fn pick_closed_merge_with_budget( + files: &[FileHandle], + mapper: &PrimaryKeyRangeMapper, + options: ClosedMergeOptions, + work_budget: usize, +) -> Option { + let limits = options.limits; + let extents = extents_of(files, mapper); + let index = ExtentIndex::new(&extents); + let mut dense = false; + let mut units = match conflict_units(files, &extents, &index, work_budget) { + Some(mut units) => { + close_units(&mut units, &extents, &index); + split_oversized_units(&mut units, files, &extents, limits.max_input_files); + units + } + None => { + dense = true; + single_file_units(files, &extents) + } + }; + sort_units(&mut units); + if dense { + return pick_dense_closed_merge(&units, files, &extents, &index, options); + } + let Some(pairs) = split_conflicts(&units, &extents, &index, work_budget) else { + return pick_dense_closed_merge(&units, files, &extents, &index, options); + }; + + let mut best: Option<(u64, u64, Range, bool)> = None; + for_each_interval(&units, &pairs, limits.max_input_files, |range, stats| { + if stats.num_files < 2 { + return; + } + let outputs = limits.predicted_outputs(stats); + if outputs > stats.num_files + 1 { + return; + } + let reduction = stats.num_files.saturating_sub(outputs); + let reduces_files = reduction > 0 + && (options.allow_file_reduction || stats.conflicts > 0) + && limits.within_reduction_cost(stats.size, reduction); + if (stats.conflicts > 0 && !options.allow_conflicts) + || (!reduces_files && stats.conflicts == 0) + { + return; + } + let benefit = (reduction + stats.conflicts) as u64; + let better = best.as_ref().is_none_or(|(best_benefit, best_size, _, _)| { + is_more_efficient(benefit, stats.size, *best_benefit, *best_size) + }); + if !better { + return; + } + let units = &units[range.clone()]; + // Overlap merges must not rewrite a dominant run for a small peer. A + // file-reducing merge is already bounded by its rewrite budget. + if !reduces_files + && !options.allow_unbalanced + && !has_balanced_components(units, files, &extents, limits.run_tolerance()) + { + return; + } + if units.len() > 1 && !adds_no_conflicts(units, &extents, &index, outputs) { + return; + } + best = Some((benefit, stats.size, range, reduces_files)); + }); + + best.map(|(_, _, range, reduces_files)| UnitPick { + inputs: collect_inputs(&units[range], files), + reduces_files, + units: units.len(), + dense_fallback: false, + }) +} + +/// Picks a bounded file reduction when the split-unit conflict graph is too +/// dense to materialize. +/// +/// The candidate comes from [`pick_reducing_interval_where`], which ignores +/// conflicts. If rejected, it is searched once more for a one-output interval. +/// Both candidates must pass [`adds_no_conflicts`], so the `(P, N)` termination +/// argument still holds. Searching only within the original candidate adds at +/// most O(K²) interval visits for K input files and one extra conflict check, +/// rather than checking external conflicts for every interval in the window. +/// This bounded retry can still miss other safely reducible intervals. +fn pick_dense_closed_merge( + units: &[Unit], + files: &[FileHandle], + extents: &[Extent], + index: &ExtentIndex, + options: ClosedMergeOptions, +) -> Option { + let limits = options.limits; + let range = pick_reducing_interval_where(units, limits, |_| true)?; + let mut selected = &units[range]; + if !is_admissible_dense_interval(selected, extents, index, options) { + let range = pick_reducing_interval_where(selected, limits, |stats| { + limits.predicted_outputs(stats) == 1 + })?; + selected = &selected[range]; + if !is_admissible_dense_interval(selected, extents, index, options) { + return None; + } + } + Some(UnitPick { + inputs: collect_inputs(selected, files), + reduces_files: true, + units: units.len(), + dense_fallback: true, + }) +} + +/// File reduction and rewrite limits are checked by the interval picker; +/// each dense candidate must also respect the closed-window conflict policy. +fn is_admissible_dense_interval( + units: &[Unit], + extents: &[Extent], + index: &ExtentIndex, + options: ClosedMergeOptions, +) -> bool { + let stats = interval_stats_exact(units, extents); + (options.allow_file_reduction || stats.conflicts > 0) + && (options.allow_conflicts || stats.conflicts == 0) + && adds_no_conflicts( + units, + extents, + index, + options.limits.predicted_outputs(&stats), + ) +} + +/// Totals of an interval without materializing [`split_conflicts`]. +/// +/// After [`close_units`], files of different units never conflict, so only +/// split units can conflict across units, and each split unit holds one file. +/// The dense fallback may instead treat every file as a split unit, and then +/// the same pairwise loop covers every cross-unit pair. Counting them takes +/// `O(K²)` for an interval of at most `K` files, instead of building a +/// window-wide adjacency. +fn interval_stats_exact<'a>(units: &'a [Unit], extents: &[Extent]) -> IntervalStats<'a> { + let mut stats = IntervalStats::default(); + let mut first = true; + for unit in units { + stats.num_files += unit.files.len(); + stats.size += unit.size; + stats.largest_file_size = stats.largest_file_size.max(unit.largest_file_size); + stats.conflicts += unit.conflicts; + stats.single_key = if first { + unit.single_key() + } else { + stats + .single_key + .filter(|&key| unit.single_key() == Some(key)) + }; + first = false; + } + + let split_files: Vec = units + .iter() + .filter(|unit| unit.split) + .map(|unit| unit.files[0]) + .collect(); + for (offset, &lhs) in split_files.iter().enumerate() { + for &rhs in &split_files[offset + 1..] { + if extents[lhs].conflicts(&extents[rhs]) { + stats.conflicts += 1; + } + } + } + stats +} + +/// Orders units by minimum primary key, then by time, so adjacent units are +/// key-adjacent. Compaction outputs and time-sliced flushes split a window by +/// key ranges, so merging neighbors in this order keeps outputs from covering +/// the key ranges of files left behind. +fn sort_units(units: &mut [Unit]) { + units.sort_by(|lhs, rhs| { + (lhs.first_key(), lhs.extent.start, Reverse(lhs.extent.end)) + .cmp(&(rhs.first_key(), rhs.extent.start, Reverse(rhs.extent.end))) + .then_with(|| lhs.files[0].cmp(&rhs.files[0])) + }); +} + +/// Whether `benefit` per `size` bytes beats the best so far; ties prefer fewer +/// bytes, then the earlier interval. +fn is_more_efficient(benefit: u64, size: u64, best_benefit: u64, best_size: u64) -> bool { + let lhs = u128::from(benefit) * u128::from(best_size.max(1)); + let rhs = u128::from(best_benefit) * u128::from(size.max(1)); + lhs > rhs || (lhs == rhs && size < best_size) +} + +/// Time and primary key bounds of a file or a group of files. +#[derive(Debug, Clone)] +pub(crate) struct Extent { + start: Timestamp, + end: Timestamp, + /// Inclusive primary key bounds. `None` is unknown and matches any key. + primary_key: Option<(Bytes, Bytes)>, +} + +impl Extent { + pub(crate) fn of_file(file: &FileHandle, mapper: &PrimaryKeyRangeMapper) -> Self { + let (start, end) = file.time_range(); + Self { + start, + end, + primary_key: file.primary_key_range(mapper), + } + } + + fn union(&self, other: &Self) -> Self { + Self { + start: self.start.min(other.start), + end: self.end.max(other.end), + primary_key: match (&self.primary_key, &other.primary_key) { + (Some((lhs_min, lhs_max)), Some((rhs_min, rhs_max))) => Some(( + lhs_min.clone().min(rhs_min.clone()), + lhs_max.clone().max(rhs_max.clone()), + )), + _ => None, + }, + } + } + + /// Time ranges overlap; touching boundaries do not. + fn time_overlaps(&self, other: &Self) -> bool { + self.start.max(other.start) < self.end.min(other.end) + } + + /// Whether two files overlap for the picker. + /// + /// Shared boundary keys count: independently written files can hold + /// different versions of the same row there. Compaction outputs do not + /// conflict with each other because the writer never splits a series. + pub(crate) fn conflicts(&self, other: &Self) -> bool { + self.time_overlaps(other) + && match (&self.primary_key, &other.primary_key) { + (Some(lhs), Some(rhs)) => { + let (low, high) = ((&lhs.0).max(&rhs.0), (&lhs.1).min(&rhs.1)); + low <= high + } + _ => true, + } + } + + /// Whether `self` may conflict with a file inside `area`, which is the + /// bounding box of a group of files. + /// + /// A shared boundary key may occur in an output, so the same inclusive + /// key overlap used for individual files also applies to the box. + fn intrudes(&self, area: &Self) -> bool { + self.conflicts(area) + } +} + +fn is_single_key(range: &(Bytes, Bytes)) -> bool { + range.0 == range.1 +} + +/// Upper primary key bound, ordered after every key when unknown. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +enum UpperKey<'a> { + Key(&'a Bytes), + Unbounded, +} + +/// Key and time bounds of the files in a subtree of [`ExtentIndex`]. +#[derive(Debug, Clone, Copy)] +struct SubtreeBounds<'a> { + max_upper_key: UpperKey<'a>, + min_start: Timestamp, + max_end: Timestamp, +} + +impl<'a> SubtreeBounds<'a> { + fn of(extent: &'a Extent) -> Self { + Self { + max_upper_key: match &extent.primary_key { + Some((_, max)) => UpperKey::Key(max), + None => UpperKey::Unbounded, + }, + min_start: extent.start, + max_end: extent.end, + } + } + + fn merge(self, other: Self) -> Self { + Self { + max_upper_key: self.max_upper_key.max(other.max_upper_key), + min_start: self.min_start.min(other.min_start), + max_end: self.max_end.max(other.max_end), + } + } +} + +/// Files of a window ordered by their lower primary key, in a segment tree +/// holding the bounds of each subtree. +/// +/// A query visits the files whose key and time ranges may overlap an area and +/// skips subtrees whose keys end before the area or whose times miss it. The +/// files of a sorted run split by key then only meet their key neighbors, +/// instead of every file sharing their time range. +struct ExtentIndex<'a> { + extents: &'a [Extent], + /// Indices into `extents`, sorted by lower key with unknown keys first. + /// Sorted offset `i` is leaf `leaf_base + i`. + by_lower_key: Vec, + /// Array-backed tree rooted at index 1; padding leaves are `None`. + bounds: Vec>>, + /// First leaf index, a power of two. + leaf_base: usize, +} + +impl<'a> ExtentIndex<'a> { + fn new(extents: &'a [Extent]) -> Self { + let lower_key = |i: usize| extents[i].primary_key.as_ref().map(|(min, _)| min); + let mut by_lower_key: Vec = (0..extents.len()).collect(); + by_lower_key.sort_by(|&lhs, &rhs| { + (lower_key(lhs), extents[lhs].start).cmp(&(lower_key(rhs), extents[rhs].start)) + }); + let leaf_base = extents.len().next_power_of_two(); + let mut bounds = vec![None; 2 * leaf_base]; + for (offset, &i) in by_lower_key.iter().enumerate() { + bounds[leaf_base + offset] = Some(SubtreeBounds::of(&extents[i])); + } + for node in (1..leaf_base).rev() { + bounds[node] = match (bounds[2 * node], bounds[2 * node + 1]) { + (Some(lhs), Some(rhs)) => Some(lhs.merge(rhs)), + (lhs, rhs) => lhs.or(rhs), + }; + } + Self { + extents, + by_lower_key, + bounds, + leaf_base, + } + } + + /// Visits the files whose inclusive key range meets `area`'s and whose time + /// range overlaps it, a superset of the files that conflict with or intrude + /// on `area`. Takes O(log N) per visited file. + fn for_each_near(&self, area: &Extent, visit: impl FnMut(usize)) { + let mut budget = usize::MAX; + let _ = self.for_each_near_with_budget(area, &mut budget, visit); + } + + /// [`for_each_near`] with a budget on visited tree nodes. + /// + /// Every node reached by the traversal consumes one unit, including nodes + /// pruned before any file callback. Returns `false` as soon as the budget is + /// exhausted, so a caller can stop a traversal whose pruning is poor even + /// when it finds few files. + fn for_each_near_with_budget( + &self, + area: &Extent, + budget: &mut usize, + mut visit: impl FnMut(usize), + ) -> bool { + let leaves_end = match &area.primary_key { + Some((_, max)) => self.by_lower_key.partition_point(|&i| { + self.extents[i] + .primary_key + .as_ref() + .is_none_or(|(min, _)| min <= max) + }), + None => self.by_lower_key.len(), + }; + self.visit_subtree(1, 0..self.leaf_base, leaves_end, area, budget, &mut visit) + } + + fn visit_subtree( + &self, + node: usize, + leaves: Range, + leaves_end: usize, + area: &Extent, + budget: &mut usize, + visit: &mut impl FnMut(usize), + ) -> bool { + if *budget == 0 { + return false; + } + *budget -= 1; + if leaves.start >= leaves_end { + return true; + } + let Some(bounds) = self.bounds[node] else { + return true; + }; + if bounds.min_start >= area.end + || bounds.max_end <= area.start + || area + .primary_key + .as_ref() + .is_some_and(|(min, _)| bounds.max_upper_key < UpperKey::Key(min)) + { + return true; + } + if leaves.len() == 1 { + let i = self.by_lower_key[leaves.start]; + if self.extents[i].time_overlaps(area) { + visit(i); + } + return true; + } + let mid = leaves.start + leaves.len() / 2; + if !self.visit_subtree(2 * node, leaves.start..mid, leaves_end, area, budget, visit) { + return false; + } + self.visit_subtree( + 2 * node + 1, + mid..leaves.end, + leaves_end, + area, + budget, + visit, + ) + } +} + +fn extents_of(files: &[FileHandle], mapper: &PrimaryKeyRangeMapper) -> Vec { + files + .iter() + .map(|file| Extent::of_file(file, mapper)) + .collect() +} + +/// Files that are compacted together. +#[derive(Debug)] +struct Unit { + /// Ascending indices into the candidate files. + files: Vec, + extent: Extent, + size: u64, + largest_file_size: u64, + /// Number of conflicting file pairs inside the unit. + conflicts: usize, + /// A single file split from a unit over the input limit. It may conflict + /// with other split files, see [`split_conflicts`]. + split: bool, +} + +impl Unit { + fn new( + members: Vec, + files: &[FileHandle], + extents: &[Extent], + conflicts: usize, + ) -> Self { + let extent = members[1..] + .iter() + .fold(extents[members[0]].clone(), |extent, &i| { + extent.union(&extents[i]) + }); + let sizes = members.iter().map(|&i| files[i].size()); + Self { + size: sizes.clone().sum(), + largest_file_size: sizes.max().unwrap_or(0), + files: members, + extent, + conflicts, + split: false, + } + } + + /// The minimum primary key, with unknown keys ordered first. + fn first_key(&self) -> Option<&Bytes> { + self.extent.primary_key.as_ref().map(|(min, _)| min) + } + + /// The primary key of a unit whose files all hold that single key. + fn single_key(&self) -> Option<&Bytes> { + self.extent + .primary_key + .as_ref() + .filter(|range| is_single_key(range)) + .map(|(min, _)| min) + } +} + +/// Groups files into connected components of [`Extent::conflicts`], each with +/// its number of conflicting pairs. Each file only checks the files the index +/// finds near it: O((N + M) log N) for M pairs of files whose key and time +/// ranges meet. +/// +/// Returns `None` once the index traversal has visited more than `work_budget` +/// nodes, so a dense window can fall back before building a quadratic graph. +fn conflict_units( + files: &[FileHandle], + extents: &[Extent], + index: &ExtentIndex, + work_budget: usize, +) -> Option> { + let mut sets = DisjointSet::new(extents.len()); + let mut conflicts = vec![0; extents.len()]; + let mut budget = work_budget; + for i in 0..extents.len() { + let completed = index.for_each_near_with_budget(&extents[i], &mut budget, |j| { + // Counts each pair once, from its smaller index. + if j > i && extents[i].conflicts(&extents[j]) { + sets.union(i, j); + conflicts[i] += 1; + } + }); + if !completed { + return None; + } + } + Some( + sets.groups() + .into_iter() + .map(|members| { + let pairs = members.iter().map(|&i| conflicts[i]).sum(); + Unit::new(members, files, extents, pairs) + }) + .collect(), + ) +} + +/// Treats every file as its own split unit, for the dense fallback that cannot +/// afford to build conflict components. +fn single_file_units(files: &[FileHandle], extents: &[Extent]) -> Vec { + (0..files.len()) + .map(|i| Unit { + split: true, + ..Unit::new(vec![i], files, extents, 0) + }) + .collect() +} + +/// Replaces units with more files than one compaction accepts by single files, +/// which are no longer closed. +fn split_oversized_units( + units: &mut Vec, + files: &[FileHandle], + extents: &[Extent], + max_input_files: usize, +) { + let (oversized, mut kept): (Vec<_>, Vec<_>) = std::mem::take(units) + .into_iter() + .partition(|unit| unit.files.len() > max_input_files); + kept.extend( + oversized + .into_iter() + .flat_map(|unit| unit.files) + .map(|i| Unit { + split: true, + ..Unit::new(vec![i], files, extents, 0) + }), + ); + *units = kept; +} + +/// Merges units until no file intrudes on a unit it does not belong to. +/// +/// Each check only visits the files the index finds near the unit, and each +/// merge re-checks only the merged unit. +fn close_units(units: &mut Vec, extents: &[Extent], index: &ExtentIndex) { + let mut owner = vec![0; extents.len()]; + for (u, unit) in units.iter().enumerate() { + for &i in &unit.files { + owner[i] = u; + } + } + let mut alive = vec![true; units.len()]; + let mut pending: Vec = (0..units.len()).collect(); + while let Some(u) = pending.pop() { + if !alive[u] { + continue; + } + let area = &units[u].extent; + let mut intruders = std::collections::BTreeSet::new(); + index.for_each_near(area, |i| { + if owner[i] != u && extents[i].intrudes(area) { + intruders.insert(owner[i]); + } + }); + if intruders.is_empty() { + continue; + } + for v in intruders { + alive[v] = false; + let absorbed = std::mem::take(&mut units[v].files); + for &i in &absorbed { + owner[i] = u; + } + let (extent, size, largest, conflicts) = { + let other = &units[v]; + ( + other.extent.clone(), + other.size, + other.largest_file_size, + other.conflicts, + ) + }; + let unit = &mut units[u]; + unit.files.extend(absorbed); + unit.extent = unit.extent.union(&extent); + unit.size += size; + unit.largest_file_size = unit.largest_file_size.max(largest); + unit.conflicts += conflicts; + } + pending.push(u); + } + let mut closed = Vec::with_capacity(units.len()); + for (u, mut unit) in std::mem::take(units).into_iter().enumerate() { + if alive[u] { + unit.files.sort_unstable(); + closed.push(unit); + } + } + *units = closed; +} + +/// Totals of a contiguous interval of units. +#[derive(Debug, Default)] +struct IntervalStats<'a> { + num_files: usize, + size: u64, + largest_file_size: u64, + conflicts: usize, + /// The primary key of an interval whose files all hold that single key. + single_key: Option<&'a Bytes>, +} + +/// Positions of the split units each unit conflicts with, sorted. Files of +/// different units never conflict, so only split files have any. +/// +/// Returns `None` once the index traversal has visited more than `work_budget` +/// nodes, so a dense window cannot materialize an unbounded adjacency. +fn split_conflicts( + units: &[Unit], + extents: &[Extent], + index: &ExtentIndex, + work_budget: usize, +) -> Option>> { + let mut position = vec![None; extents.len()]; + for (p, unit) in units.iter().enumerate() { + if unit.split { + position[unit.files[0]] = Some(p); + } + } + let mut budget = work_budget; + let mut result = Vec::with_capacity(units.len()); + for unit in units { + if !unit.split { + result.push(Vec::new()); + continue; + } + let i = unit.files[0]; + let mut neighbors = Vec::new(); + let completed = index.for_each_near_with_budget(&extents[i], &mut budget, |j| { + if let Some(p) = position[j] + && j != i + && extents[i].conflicts(&extents[j]) + { + neighbors.push(p); + } + }); + if !completed { + return None; + } + neighbors.sort_unstable(); + result.push(neighbors); + } + Some(result) +} + +/// Visits every contiguous interval of units with at most `max_input_files` +/// files: O(U · K) intervals for an input limit K. Each interval counts the +/// conflicts of its split files from `split_conflicts` (see +/// [`split_conflicts`]) with a binary search; an empty `split_conflicts` +/// skips them. +fn for_each_interval<'a>( + units: &'a [Unit], + split_conflicts: &[Vec], + max_input_files: usize, + mut visit: impl FnMut(Range, &IntervalStats<'a>), +) { + for left in 0..units.len() { + let mut stats = IntervalStats::default(); + for (right, unit) in units.iter().enumerate().skip(left) { + stats.num_files += unit.files.len(); + if stats.num_files > max_input_files { + break; + } + stats.size += unit.size; + stats.largest_file_size = stats.largest_file_size.max(unit.largest_file_size); + stats.conflicts += unit.conflicts; + stats.single_key = if right == left { + unit.single_key() + } else { + stats + .single_key + .filter(|&key| unit.single_key() == Some(key)) + }; + if let Some(neighbors) = split_conflicts.get(right) { + // Conflicts with the split files already in the interval. + stats.conflicts += neighbors.partition_point(|&p| p < right) + - neighbors.partition_point(|&p| p < left); + } + visit(left..right + 1, &stats); + } + } +} + +/// Whether merging `units` into `outputs` files adds no more conflicts than +/// their files have with other files. Outputs stay inside the bounding box of +/// the inputs, so each of them may conflict with every file intruding on that +/// box. An interval without intruders is closed and adds no conflicts. +fn adds_no_conflicts( + units: &[Unit], + extents: &[Extent], + index: &ExtentIndex, + outputs: usize, +) -> bool { + let area = units[1..] + .iter() + .fold(units[0].extent.clone(), |area, unit| { + area.union(&unit.extent) + }); + let selected: std::collections::HashSet = units + .iter() + .flat_map(|unit| unit.files.iter().copied()) + .collect(); + let mut intruders = Vec::new(); + index.for_each_near(&area, |i| { + if !selected.contains(&i) && extents[i].intrudes(&area) { + intruders.push(i); + } + }); + if intruders.is_empty() { + return true; + } + let existing: usize = intruders + .iter() + .map(|&other| { + units + .iter() + .flat_map(|unit| unit.files.iter()) + .filter(|&&i| extents[i].conflicts(&extents[other])) + .count() + }) + .sum(); + outputs.saturating_mul(intruders.len()) <= existing +} + +/// Whether every group of conflicting files among `units` is balanced, see +/// [`has_balanced_runs`]. Runs of separate groups never overlap, so pairing a +/// dominant run with the peer of another group would hide it. Takes O(K²) +/// conflict checks for K files. +fn has_balanced_components( + units: &[Unit], + files: &[FileHandle], + extents: &[Extent], + tolerance: u64, +) -> bool { + let members: Vec = units + .iter() + .flat_map(|unit| unit.files.iter().copied()) + .collect(); + let mut sets = DisjointSet::new(members.len()); + for (i, &lhs) in members.iter().enumerate() { + for (j, &rhs) in members.iter().enumerate().skip(i + 1) { + if extents[lhs].conflicts(&extents[rhs]) { + sets.union(i, j); + } + } + } + sets.groups() + .into_iter() + .filter(|group| group.len() > 1) + .all(|group| { + has_balanced_runs( + group.into_iter().map(|i| members[i]), + files, + extents, + tolerance, + ) + }) +} + +/// Whether the largest sorted run among `files` holds at most `tolerance` bytes +/// more than all other runs together. Runs are built greedily in time order, +/// which takes O(K²) conflict checks for K files. +fn has_balanced_runs( + members: impl Iterator, + files: &[FileHandle], + extents: &[Extent], + tolerance: u64, +) -> bool { + let mut members: Vec = members.collect(); + members.sort_by(|&lhs, &rhs| { + (extents[lhs].start, Reverse(extents[lhs].end), lhs).cmp(&( + extents[rhs].start, + Reverse(extents[rhs].end), + rhs, + )) + }); + let mut runs: Vec<(Vec, u64)> = Vec::new(); + for i in members { + let size = files[i].size(); + match runs.iter_mut().find(|(run, _)| { + run.iter() + .all(|&other| !extents[i].conflicts(&extents[other])) + }) { + Some((run, run_size)) => { + run.push(i); + *run_size += size; + } + None => runs.push((vec![i], size)), + } + } + let total: u64 = runs.iter().map(|(_, size)| size).sum(); + let largest = runs.iter().map(|(_, size)| *size).max().unwrap_or(0); + largest <= (total - largest).saturating_add(tolerance) +} + +fn collect_inputs(units: &[Unit], files: &[FileHandle]) -> Vec { + let mut indices: Vec = units + .iter() + .flat_map(|unit| unit.files.iter().copied()) + .collect(); + indices.sort_unstable(); + indices.into_iter().map(|i| files[i].clone()).collect() +} + +/// Union-find over file indices; each root is the smallest index of its set. +struct DisjointSet { + parent: Vec, +} + +impl DisjointSet { + fn new(len: usize) -> Self { + Self { + parent: (0..len).collect(), + } + } + + fn find(&mut self, mut i: usize) -> usize { + while self.parent[i] != i { + self.parent[i] = self.parent[self.parent[i]]; + i = self.parent[i]; + } + i + } + + fn union(&mut self, lhs: usize, rhs: usize) { + let (lhs, rhs) = (self.find(lhs), self.find(rhs)); + let (root, child) = (lhs.min(rhs), lhs.max(rhs)); + self.parent[child] = root; + } + + /// Sets in ascending order of their smallest index, each sorted. + fn groups(&mut self) -> Vec> { + let mut groups: Vec> = Vec::new(); + let mut slot = vec![usize::MAX; self.parent.len()]; + for i in 0..self.parent.len() { + let root = self.find(i); + if slot[root] == usize::MAX { + slot[root] = groups.len(); + groups.push(Vec::new()); + } + groups[slot[root]].push(i); + } + groups + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashSet; + + use store_api::storage::FileId; + + use super::*; + use crate::compaction::test_util::{ + new_file_handle_with_size_sequence_and_primary_key_range, pk_range, + primary_key_mapper_for_test, + }; + + const MAX_OUTPUT: u64 = 16; + + fn file(start: i64, end: i64, min: &[u8], max: &[u8], size: u64) -> FileHandle { + new_file_handle_with_size_sequence_and_primary_key_range( + FileId::random(), + start, + end, + 1, + 1, + size, + pk_range(min, max), + ) + } + + /// Files that all overlap each other in time and primary key. + fn dense_files(count: usize) -> Vec { + (0..count) + .map(|idx| file(0, 100, b"a", b"z", 1 + (idx % 3) as u64)) + .collect() + } + + /// A chain of files whose key ranges overlap their immediate neighbors. + fn chain_files(count: usize, size: u64) -> Vec { + (0..count) + .map(|idx| { + let min = format!("{idx:04}"); + let max = format!("{:04}", idx + 2); + file(0, 100, min.as_bytes(), max.as_bytes(), size) + }) + .collect() + } + + fn limits(max_output_file_size: u64) -> PickLimits { + PickLimits { + max_output_file_size: Some(max_output_file_size), + max_input_files: 16, + } + } + + fn closed_merge(files: &[FileHandle], max_output_file_size: u64) -> Option { + pick_closed_merge( + files, + &primary_key_mapper_for_test(), + closed_merge_options(max_output_file_size), + ) + } + + fn closed_merge_with_budget( + files: &[FileHandle], + max_output_file_size: u64, + work_budget: usize, + ) -> Option { + pick_closed_merge_with_budget( + files, + &primary_key_mapper_for_test(), + closed_merge_options(max_output_file_size), + work_budget, + ) + } + + fn closed_merge_options(max_output_file_size: u64) -> ClosedMergeOptions { + ClosedMergeOptions { + limits: limits(max_output_file_size), + allow_file_reduction: true, + allow_conflicts: true, + allow_unbalanced: false, + } + } + + #[test] + fn test_small_l1_reduction_limits() { + let mapper = primary_key_mapper_for_test(); + // Empty/single inputs, the small-file boundary, balance, byte and count + // caps, and an unlimited output size all affect eligibility. + for (sizes, max_output, max_inputs, expected) in [ + (vec![], Some(16), 16, 0), + (vec![4], Some(16), 16, 0), + (vec![4, 4], Some(16), 16, 2), + (vec![5, 5], Some(16), 16, 0), + (vec![4, 1], Some(16), 16, 0), + (vec![4, 2, 2], Some(16), 16, 3), + (vec![4; 8], Some(16), 16, 4), + (vec![4; 8], Some(16), 3, 3), + (vec![4, 4], Some(16), 1, 0), + (vec![4, 4], Some(0), 16, 0), + (vec![4, 4], None, 16, 0), + ] { + let files = sizes + .iter() + .map(|&size| file(0, 100, b"a", b"z", size)) + .collect::>(); + let pick = pick_small_l1( + &files, + &mapper, + PickLimits { + max_output_file_size: max_output, + max_input_files: max_inputs, + }, + ); + assert_eq!( + expected, + pick.map_or(0, |pick| pick.inputs.len()), + "sizes={sizes:?}, target={max_output:?}, cap={max_inputs}" + ); + } + } + + #[test] + fn test_small_l1_checks_unselected_large_and_l0_files() { + let mapper = primary_key_mapper_for_test(); + let small = [file(0, 10, b"a", b"b", 4), file(90, 100, b"y", b"z", 4)]; + for level in [0, 1] { + for (min, max, expected) in [(b"m", b"n", false), (b"a", b"z", true)] { + let peer = new_file_handle_with_size_sequence_and_primary_key_range( + FileId::random(), + 0, + 100, + level, + 1, + 100, + pk_range(min, max), + ); + let files = [small.to_vec(), vec![peer]].concat(); + let pick = pick_small_l1(&files, &mapper, limits(16)); + assert_eq!( + expected, + pick.is_some(), + "level={level}, range={min:?}..{max:?}" + ); + if let Some(pick) = pick { + assert_eq!(ids(&small), ids(&pick.inputs)); + } + } + } + // A small L0 must not be selected as the second member of an L1 merge. + let l0 = new_file_handle_with_size_sequence_and_primary_key_range( + FileId::random(), + 90, + 100, + 0, + 1, + 4, + pk_range(b"y", b"z"), + ); + assert!(pick_small_l1(&[small[0].clone(), l0], &mapper, limits(16)).is_none()); + } + + fn ids(files: &[FileHandle]) -> HashSet { + files.iter().map(|file| file.file_id().file_id()).collect() + } + + fn extent(start: i64, end: i64, min: &[u8], max: &[u8]) -> Extent { + Extent::of_file( + &file(start, end, min, max, 1), + &primary_key_mapper_for_test(), + ) + } + + #[test] + fn test_conflicts_include_shared_boundary_keys() { + // Independently written files may share rows at their boundary keys. + assert!(extent(0, 10, b"a", b"m").conflicts(&extent(0, 10, b"m", b"z"))); + assert!(extent(0, 10, b"m", b"m").conflicts(&extent(0, 10, b"m", b"z"))); + // A writer-produced run has strictly disjoint key ranges. + assert!(!extent(0, 10, b"a", b"l").conflicts(&extent(0, 10, b"m", b"z"))); + // A single key shared by two single-key files is a real overlap. + assert!(extent(0, 10, b"m", b"m").conflicts(&extent(5, 15, b"m", b"m"))); + // So is a single key inside another range. + assert!(extent(0, 10, b"m", b"m").conflicts(&extent(0, 10, b"a", b"z"))); + assert!(extent(0, 10, b"a", b"z").conflicts(&extent(0, 10, b"m", b"m"))); + assert!(extent(0, 10, b"a", b"n").conflicts(&extent(0, 10, b"m", b"z"))); + // Touching time ranges do not overlap. + assert!(!extent(0, 10, b"a", b"z").conflicts(&extent(10, 20, b"a", b"z"))); + } + + #[test] + fn test_intrusion_covers_every_conflict_inside_an_area() { + let area = extent(0, 10, b"m", b"z"); + // Any file sharing the area's boundary key may conflict with an output. + assert!(extent(0, 10, b"m", b"m").intrudes(&area)); + assert!(extent(0, 10, b"a", b"m").intrudes(&area)); + assert!(!extent(0, 10, b"a", b"l").intrudes(&area)); + assert!(extent(0, 10, b"a", b"n").intrudes(&area)); + assert!(!extent(10, 20, b"a", b"z").intrudes(&area)); + // A range may conflict with a single-key file inside a single-key area. + let point = extent(0, 10, b"m", b"m"); + assert!(extent(0, 10, b"a", b"z").intrudes(&point)); + assert!(extent(0, 10, b"a", b"m").intrudes(&point)); + } + + #[test] + fn test_intrusion_covers_conflicts_of_random_files() { + let mut rng = Rng(23); + let random_extent = |rng: &mut Rng| { + let start = rng.next(10) as i64; + let end = start + 1 + rng.next(10) as i64; + let min = rng.next(6) as u8; + let max = min + rng.next(3) as u8; + extent(start, end, &[b'0' + min], &[b'0' + max]) + }; + for _ in 0..2_000 { + let group: Vec<_> = (0..1 + rng.next(3)) + .map(|_| random_extent(&mut rng)) + .collect(); + let area = group[1..] + .iter() + .fold(group[0].clone(), |area, extent| area.union(extent)); + let other = random_extent(&mut rng); + if group.iter().any(|extent| other.conflicts(extent)) { + assert!(other.intrudes(&area), "{other:?} conflicts inside {area:?}"); + } + } + } + + #[test] + fn test_extent_index_budget_counts_nodes() { + // The root's time bounds straddle the query, but neither leaf overlaps + // it. A callback-only budget would never fire; the node budget must. + let files = vec![file(0, 10, b"a", b"z", 1), file(90, 100, b"a", b"z", 1)]; + let extents = extents_of(&files, &primary_key_mapper_for_test()); + let index = ExtentIndex::new(&extents); + let query = extent(50, 51, b"a", b"z"); + let mut matched = 0; + + let mut budget = 2; + assert!(!index.for_each_near_with_budget(&query, &mut budget, |_| matched += 1)); + assert_eq!(0, budget); + assert_eq!(0, matched); + + let mut budget = 3; + assert!(index.for_each_near_with_budget(&query, &mut budget, |_| matched += 1)); + assert_eq!(0, budget); + assert_eq!(0, matched); + } + + #[test] + fn test_extent_index_matches_linear_scan() { + let mut rng = Rng(29); + let mapper = primary_key_mapper_for_test(); + for count in [0, 1, 2, 7, 64, 200] { + let extents: Vec<_> = (0..count) + .map(|_| { + let start = rng.next(100) as i64; + let end = start + rng.next(30) as i64; + let min = rng.next(50) as u8; + let max = min + rng.next(10) as u8; + let range = if rng.next(10) > 0 { + pk_range(&[b'0' + min], &[b'0' + max]) + } else { + None + }; + let file = new_file_handle_with_size_sequence_and_primary_key_range( + FileId::random(), + start, + end, + 1, + 1, + 1, + range, + ); + Extent::of_file(&file, &mapper) + }) + .collect(); + let index = ExtentIndex::new(&extents); + for area in &extents { + let mut near = Vec::new(); + index.for_each_near(area, |i| near.push(i)); + near.sort_unstable(); + let expected: Vec<_> = (0..extents.len()) + .filter(|&i| { + let other = &extents[i]; + let keys_meet = match (&other.primary_key, &area.primary_key) { + (Some(lhs), Some(rhs)) => lhs.0 <= rhs.1 && rhs.0 <= lhs.1, + _ => true, + }; + keys_meet && other.time_overlaps(area) + }) + .collect(); + assert_eq!(expected, near, "{area:?}"); + for (i, other) in extents.iter().enumerate() { + if other.conflicts(area) || other.intrudes(area) { + assert!(near.binary_search(&i).is_ok(), "{other:?} near {area:?}"); + } + } + } + } + } + + #[test] + fn test_active_l1_packs_pk_disjoint_files() { + let files = vec![ + file(0, 1000, b"a", b"a", 100), + file(0, 1000, b"z", b"z", 100), + ]; + let pick = pick_active_l1(&files, &primary_key_mapper_for_test(), limits(512)).unwrap(); + assert_eq!(ids(&files), ids(&pick.inputs)); + } + + #[test] + fn test_active_l1_bounds_rewrite_cost() { + let overlapping = |sizes: &[u64]| { + sizes + .iter() + .enumerate() + .map(|(idx, &size)| file(idx as i64, 100, b"a", b"z", size)) + .collect::>() + }; + let mapper = primary_key_mapper_for_test(); + assert!(pick_active_l1(&overlapping(&[300, 300, 300]), &mapper, limits(512)).is_some()); + // Removing one file would rewrite 4600 bytes, beyond 4 * 512. + let costly = overlapping(&[500, 500, 500, 500, 500, 500, 500, 500, 500, 100]); + assert!(pick_active_l1(&costly, &mapper, limits(512)).is_none()); + } + + #[test] + fn test_active_l1_merges_small_files_beside_dominant_file() { + // All files overlap. Merging the group would rewrite the large file for + // its small peers, so only the small files are merged. + let mut files = vec![file(0, 1000, b"a", b"z", 1024)]; + files.extend((1..16).map(|start| file(start, 1000, b"a", b"z", 1))); + let pick = pick_active_l1(&files, &primary_key_mapper_for_test(), limits(512)).unwrap(); + assert_eq!(ids(&files[1..]), ids(&pick.inputs)); + } + + #[test] + fn test_merges_of_a_single_key_write_one_file() { + // The writer never splits a series, so the overlapping files of one key + // compact into a single file however large they are. + let files = vec![file(0, 100, b"m", b"m", 32), file(50, 150, b"m", b"m", 32)]; + let mapper = primary_key_mapper_for_test(); + let pick = closed_merge(&files, MAX_OUTPUT).unwrap(); + assert_eq!(ids(&files), ids(&pick.inputs)); + assert!(pick.reduces_files); + let pick = pick_active_l1(&files, &mapper, limits(MAX_OUTPUT)).unwrap(); + assert_eq!(ids(&files), ids(&pick.inputs)); + } + + #[test] + fn test_closed_merge_does_not_enclose_other_files() { + let a = file(0, 10, b"a", b"c", 10); + let b = file(0, 10, b"x", b"z", 10); + let c = file(0, 10, b"m", b"m", 1); + let files = vec![a.clone(), b, c.clone()]; + let pick = closed_merge(&files, MAX_OUTPUT).unwrap(); + // Merging A and B alone would produce a file covering C. + assert_eq!(ids(&[a, c]), ids(&pick.inputs)); + assert!(pick.reduces_files); + } + + #[test] + fn test_shared_boundary_files_are_picked_without_file_reduction() { + // Independently written files can hold different versions at m. Their + // total size predicts two outputs, so only overlap resolution helps. + let size = 512 * 1024 * 1024; + let files = [ + file(0, 100, b"a", b"m", size), + file(0, 100, b"m", b"z", size), + ]; + let pick = closed_merge(&files, size).unwrap(); + assert_eq!(ids(&files), ids(&pick.inputs)); + assert!(!pick.reduces_files); + } + + /// A sorted run of `sizes.len()` key-sliced files, and files spanning all + /// of its keys in the given time ranges, which overlap the whole run. + fn run_with_late_files( + sizes: &[u64], + late: &[(i64, i64)], + ) -> (Vec, Vec) { + let run = sizes + .iter() + .enumerate() + .map(|(idx, &size)| { + let idx = 2 * idx as u8; + file(0, 1000, &[b'a', idx], &[b'a', idx + 1], size) + }) + .collect::>(); + let max_key = 2 * sizes.len() as u8; + let late = late + .iter() + .map(|&(start, end)| file(start, end, &[b'a', 0], &[b'a', max_key], 1)) + .collect(); + (run, late) + } + + #[test] + fn test_closed_merge_splits_components_over_input_limit() { + // The late file joins all 20 run files into one component, over the + // input limit. Small run files are still merged. + let (run, late) = run_with_late_files(&[1; 20], &[(100, 900)]); + let files = [run, late].concat(); + let pick = closed_merge(&files, MAX_OUTPUT).unwrap(); + assert!(pick.reduces_files); + assert!(pick.inputs.len() <= 16); + } + + #[test] + fn test_closed_merge_leaves_late_file_over_large_run() { + // Resolving the overlap would rewrite the whole run, beyond the input + // limit, and merging part of it would only spread the overlap. + let (run, late) = run_with_late_files(&[MAX_OUTPUT; 20], &[(100, 900)]); + let files = [run, late].concat(); + assert!(closed_merge(&files, MAX_OUTPUT).is_none()); + } + + #[test] + fn test_closed_merge_compacts_late_files_over_large_run() { + let (run, late) = + run_with_late_files(&[MAX_OUTPUT; 20], &[(100, 200), (300, 400), (500, 600)]); + let files = [run, late.clone()].concat(); + let pick = closed_merge(&files, MAX_OUTPUT).unwrap(); + assert_eq!(ids(&late), ids(&pick.inputs)); + assert!(pick.reduces_files); + } + + #[test] + fn test_closed_merge_rejects_dominant_runs() { + // Each of these components has a 3 GiB run and a 512 MiB peer. + let size = 512 * 1024 * 1024; + // The peer spans the keys of its run, written before or after it. + let mut files = (0..6u8) + .map(|idx| file(idx.into(), 10_000, &[b'b' + idx], &[b'b' + idx], size)) + .collect::>(); + files.push(file(6, 10_000, b"a", b"h", size)); + files.extend((0..6u8).map(|idx| { + file( + 8 + i64::from(idx), + 10_000, + &[b'n' + idx], + &[b'n' + idx], + size, + ) + })); + files.push(file(7, 10_000, b"m", b"t", size)); + // A separate balanced pair is still resolved. + let pair = [ + file(14, 10_000, b"x", b"y", size), + file(15, 10_000, b"x", b"y", size), + ]; + files.extend(pair.clone()); + + let pick = closed_merge(&files, size).unwrap(); + assert_eq!(ids(&pair), ids(&pick.inputs)); + assert!(!pick.reduces_files); + } + + #[test] + fn test_closed_merge_ignores_touching_time_ranges() { + let first = file(0, 10, b"a", b"z", 1024); + let pair = [file(10, 20, b"a", b"z", 512), file(11, 19, b"a", b"z", 512)]; + let mut files = vec![first]; + files.extend(pair.clone()); + let pick = closed_merge(&files, 512).unwrap(); + assert_eq!(ids(&pair), ids(&pick.inputs)); + } + + #[test] + fn test_conflict_units_respects_work_budget() { + let files = dense_files(20); + let mapper = primary_key_mapper_for_test(); + let extents = extents_of(&files, &mapper); + let index = ExtentIndex::new(&extents); + + assert!(conflict_units(&files, &extents, &index, 0).is_none()); + assert!(conflict_units(&files, &extents, &index, usize::MAX).is_some()); + } + + #[test] + fn test_split_conflicts_respects_work_budget() { + let files = dense_files(20); + let mapper = primary_key_mapper_for_test(); + let extents = extents_of(&files, &mapper); + let index = ExtentIndex::new(&extents); + let mut units = conflict_units(&files, &extents, &index, usize::MAX).unwrap(); + close_units(&mut units, &extents, &index); + split_oversized_units(&mut units, &files, &extents, 16); + sort_units(&mut units); + + assert_eq!(20, units.iter().filter(|unit| unit.split).count()); + assert!(split_conflicts(&units, &extents, &index, 0).is_none()); + assert!(split_conflicts(&units, &extents, &index, usize::MAX).is_some()); + } + + #[test] + fn test_interval_stats_exact_matches_for_each_interval() { + let mapper = primary_key_mapper_for_test(); + // One closed unit with internal conflicts, and one oversized component + // split into single-file units with cross-unit conflicts. + let mut files = (0..10) + .map(|_| file(0, 100, b"a", b"c", 1)) + .collect::>(); + files.extend((0..20).map(|_| file(0, 100, b"x", b"z", 1))); + let extents = extents_of(&files, &mapper); + let index = ExtentIndex::new(&extents); + let mut units = conflict_units(&files, &extents, &index, usize::MAX).unwrap(); + close_units(&mut units, &extents, &index); + split_oversized_units(&mut units, &files, &extents, 16); + sort_units(&mut units); + + let pairs = split_conflicts(&units, &extents, &index, usize::MAX).unwrap(); + for_each_interval(&units, &pairs, 16, |range, stats| { + let exact = interval_stats_exact(&units[range], &extents); + assert_eq!(stats.num_files, exact.num_files); + assert_eq!(stats.size, exact.size); + assert_eq!(stats.largest_file_size, exact.largest_file_size); + assert_eq!(stats.conflicts, exact.conflicts); + assert_eq!(stats.single_key, exact.single_key); + }); + } + + #[test] + fn test_active_l1_falls_back_on_dense_conflict_graph() { + let files = dense_files(300); + let pick = pick_active_l1_with_budget( + &files, + &primary_key_mapper_for_test(), + limits(MAX_OUTPUT), + 1 << 10, + ) + .unwrap(); + assert!(!pick.dense_fallback); + assert!(pick.reduces_files); + assert!((2..=16).contains(&pick.inputs.len())); + } + + #[test] + fn test_dense_closed_merge_reduces_files() { + let files = dense_files(300); + let pick = closed_merge_with_budget(&files, MAX_OUTPUT, 1 << 10).unwrap(); + assert!(pick.dense_fallback); + assert!(pick.reduces_files); + assert!((2..=16).contains(&pick.inputs.len())); + } + + #[test] + fn test_dense_closed_merge_resolves_conflicts_without_file_reduction_setting() { + let files = dense_files(300); + let mapper = primary_key_mapper_for_test(); + let options = ClosedMergeOptions { + allow_file_reduction: false, + ..closed_merge_options(MAX_OUTPUT) + }; + let pick = pick_closed_merge_with_budget(&files, &mapper, options, 1 << 10).unwrap(); + assert!(pick.dense_fallback); + assert!(pick.reduces_files); + } + + #[test] + fn test_dense_closed_merge_defers_when_conflicts_are_not_allowed() { + let files = dense_files(300); + let mapper = primary_key_mapper_for_test(); + let options = ClosedMergeOptions { + allow_conflicts: false, + ..closed_merge_options(MAX_OUTPUT) + }; + assert!(pick_closed_merge_with_budget(&files, &mapper, options, 1 << 10).is_none()); + } + + #[test] + fn test_dense_closed_merge_retries_one_output_interval() { + // PR #9741: the first 16 files predict two outputs but have only one + // conflict with the rest of the chain. The first eight predict one + // output, so rejecting the larger interval must not prevent progress. + let files = chain_files(300, 2); + let mapper = primary_key_mapper_for_test(); + let extents = extents_of(&files, &mapper); + let index = ExtentIndex::new(&extents); + let mut units = conflict_units(&files, &extents, &index, usize::MAX).unwrap(); + close_units(&mut units, &extents, &index); + split_oversized_units(&mut units, &files, &extents, 16); + sort_units(&mut units); + + assert!(!adds_no_conflicts(&units[..16], &extents, &index, 2)); + assert!(adds_no_conflicts(&units[..8], &extents, &index, 1)); + // Cover both single-file units from an exhausted grouping budget and + // split units from an oversized closure. + for pick in [ + closed_merge_with_budget(&files, MAX_OUTPUT, 0).unwrap(), + pick_dense_closed_merge( + &units, + &files, + &extents, + &index, + closed_merge_options(MAX_OUTPUT), + ) + .unwrap(), + ] { + assert_eq!(ids(&files[..8]), ids(&pick.inputs)); + assert!(pick.dense_fallback); + assert!(pick.reduces_files); + assert_eq!(300, pick.units); + } + } + + #[test] + fn test_dense_closed_merge_rejects_unsafe_one_output_retry() { + // Intruders lie in the time gaps between inputs, so even a one-output + // merge can introduce conflicts that no input previously had. + let mut files: Vec<_> = (0..4) + .map(|i| file(30 * i, 30 * i + 10, b"a", b"z", 6)) + .collect(); + files.extend((0..3).map(|i| file(30 * i + 15, 30 * i + 25, b"b", b"y", 1_000))); + + assert!(closed_merge_with_budget(&files, MAX_OUTPUT, 0).is_none()); + } + + /// SplitMix64, so the simulation needs no random number dependency. + struct Rng(u64); + + impl Rng { + fn next(&mut self, bound: u64) -> u64 { + self.0 = self.0.wrapping_add(0x9e37_79b9_7f4a_7c15); + let mut z = self.0; + z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9); + z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb); + (z ^ (z >> 31)) % bound + } + } + + /// A simulated file with its key bounds as numbers. + #[derive(Clone)] + struct SimFile { + handle: FileHandle, + min: u64, + max: u64, + } + + fn sim_file(start: i64, end: i64, min: u64, max: u64, size: u64) -> SimFile { + let (min_key, max_key) = (format!("{min:04}"), format!("{max:04}")); + SimFile { + handle: file(start, end, min_key.as_bytes(), max_key.as_bytes(), size), + min, + max, + } + } + + /// Writes the outputs of merging `inputs` like the compactor: one stream in + /// (key, time) order, split by size into consecutive files. The data is + /// spread evenly over keys, and a series is never split between files. + fn merge(inputs: &[SimFile], rng: &mut Rng) -> Vec { + let min = inputs.iter().map(|file| file.min).min().unwrap(); + let max = inputs.iter().map(|file| file.max).max().unwrap(); + let start = inputs + .iter() + .map(|file| file.handle.time_range().0.value()) + .min() + .unwrap(); + let end = inputs + .iter() + .map(|file| file.handle.time_range().1.value()) + .max() + .unwrap(); + let size: u64 = inputs.iter().map(|file| file.handle.size()).sum(); + // Outputs may be fewer than predicted, never more. A single key is + // never split. + let predicted = if min == max { + 1 + } else { + size.div_ceil(MAX_OUTPUT).max(1) + }; + let keys = max - min + 1; + let outputs = (predicted - rng.next(2).min(predicted - 1)).min(keys); + (0..outputs) + .map(|i| { + let first_key = min + keys * i / outputs; + let last_key = min + keys * (i + 1) / outputs - 1; + let output_size = size * (i + 1) / outputs - size * i / outputs; + sim_file(start, end, first_key, last_key, output_size) + }) + .collect() + } + + fn conflict_pairs(files: &[FileHandle]) -> usize { + let extents = extents_of(files, &primary_key_mapper_for_test()); + (0..extents.len()) + .map(|i| { + (i + 1..extents.len()) + .filter(|&j| extents[i].conflicts(&extents[j])) + .count() + }) + .sum() + } + + fn random_window(rng: &mut Rng) -> Vec { + (0..2 + rng.next(30)) + .map(|_| { + let start = rng.next(3_000) as i64; + let end = start + 1 + rng.next(600) as i64; + let min = rng.next(200); + let span = if rng.next(4) == 0 { 200 } else { 20 }; + let max = min + rng.next(span); + sim_file(start, end, min, max, 1 + rng.next(2 * MAX_OUTPUT)) + }) + .collect() + } + + /// A sorted run of more files than one compaction accepts, and a few files + /// spanning many of its keys, which join them into oversized components. + fn oversized_window(rng: &mut Rng) -> Vec { + let run_files = 17 + rng.next(40); + let mut files: Vec<_> = (0..run_files) + .map(|idx| { + let start = rng.next(100) as i64; + let end = 2_900 + rng.next(100) as i64; + sim_file( + start, + end, + 5 * idx, + 5 * (idx + 1) - 1, + 1 + rng.next(2 * MAX_OUTPUT), + ) + }) + .collect(); + files.extend((0..1 + rng.next(6)).map(|_| { + let start = rng.next(2_800) as i64; + let end = start + 1 + rng.next(200) as i64; + let min = rng.next(5 * run_files / 2); + let max = min + 5 * run_files / 2 + rng.next(5 * run_files / 2); + sim_file(start, end, min, max, 1 + rng.next(MAX_OUTPUT)) + })); + files + } + + #[test] + fn test_dense_closed_merge_terminates() { + let mut rng = Rng(23); + let mut files: Vec<_> = (0..300) + .map(|_| sim_file(0, 100, 0, 8, 1 + rng.next(2 * MAX_OUTPUT))) + .collect(); + let handles = |files: &[SimFile]| { + files + .iter() + .map(|file| file.handle.clone()) + .collect::>() + }; + let mut current = (conflict_pairs(&handles(&files)), files.len()); + let mut steps = 0; + let mut dense_steps = 0; + while let Some(pick) = closed_merge_with_budget(&handles(&files), MAX_OUTPUT, 1 << 10) { + assert!(pick.inputs.len() >= 2); + if pick.dense_fallback { + dense_steps += 1; + } + let picked = ids(&pick.inputs); + let (inputs, kept): (Vec<_>, Vec<_>) = files + .into_iter() + .partition(|file| picked.contains(&file.handle.file_id().file_id())); + files = kept; + files.extend(merge(&inputs, &mut rng)); + let next = (conflict_pairs(&handles(&files)), files.len()); + assert!(next < current, "step {steps}: {current:?} -> {next:?}"); + current = next; + steps += 1; + assert!(steps < 1_000); + } + assert!(dense_steps > 0); + assert!(files.len() < 300); + } + + /// Applies `pick` until it picks nothing, asserting that `measure` strictly + /// decreases with every merge. Returns the number of merges. + fn run_to_fixpoint( + rng: &mut Rng, + pick: impl Fn(&[FileHandle]) -> Option, + measure: impl Fn(&[FileHandle]) -> (usize, usize), + ) -> usize { + run_windows_to_fixpoint(rng, random_window, pick, measure) + } + + fn run_windows_to_fixpoint( + rng: &mut Rng, + new_window: impl Fn(&mut Rng) -> Vec, + pick: impl Fn(&[FileHandle]) -> Option, + measure: impl Fn(&[FileHandle]) -> (usize, usize), + ) -> usize { + let mut total_steps = 0; + for window in 0..300 { + let mut files = new_window(rng); + let handles = + |files: &[SimFile]| files.iter().map(|f| f.handle.clone()).collect::>(); + let mut current = measure(&handles(&files)); + let mut steps = 0; + while let Some(pick) = pick(&handles(&files)) { + assert!(pick.inputs.len() >= 2); + let picked = ids(&pick.inputs); + let (inputs, kept): (Vec<_>, Vec<_>) = files + .into_iter() + .partition(|file| picked.contains(&file.handle.file_id().file_id())); + files = kept; + files.extend(merge(&inputs, rng)); + let next = measure(&handles(&files)); + assert!(next < current, "window {window}: {current:?} -> {next:?}"); + current = next; + steps += 1; + } + total_steps += steps; + } + total_steps + } + + #[test] + fn test_closed_merges_terminate() { + let steps = run_to_fixpoint( + &mut Rng(7), + |files| closed_merge(files, MAX_OUTPUT), + |files| (conflict_pairs(files), files.len()), + ); + assert!(steps > 300); + } + + #[test] + fn test_closed_merges_of_oversized_components_terminate() { + let steps = run_windows_to_fixpoint( + &mut Rng(19), + oversized_window, + |files| closed_merge(files, MAX_OUTPUT), + |files| (conflict_pairs(files), files.len()), + ); + assert!(steps > 300); + } + + #[test] + fn test_closed_merges_of_quiet_and_settling_windows_terminate() { + let mapper = primary_key_mapper_for_test(); + // Quiet windows take merges a default window refuses; settling windows + // skip the conflicting ones. + for (seed, options, more_merges) in [ + ( + 13, + ClosedMergeOptions { + allow_unbalanced: true, + ..closed_merge_options(MAX_OUTPUT) + }, + true, + ), + ( + 17, + ClosedMergeOptions { + allow_conflicts: false, + ..closed_merge_options(MAX_OUTPUT) + }, + false, + ), + ] { + let steps = run_to_fixpoint( + &mut Rng(seed), + |files| pick_closed_merge(files, &mapper, options), + |files| (conflict_pairs(files), files.len()), + ); + let base = run_to_fixpoint( + &mut Rng(seed), + |files| closed_merge(files, MAX_OUTPUT), + |files| (conflict_pairs(files), files.len()), + ); + assert_eq!(more_merges, steps > base, "seed {seed}: {steps} vs {base}"); + } + } + + #[test] + fn test_small_l1_merges_terminate() { + let mapper = primary_key_mapper_for_test(); + let steps = run_windows_to_fixpoint( + &mut Rng(31), + |rng| { + let mut files = random_window(rng); + // Include a reducible cluster inside a large peer as well as + // random external intruders to exercise the conflict guard. + files.push(sim_file(0, 3000, 0, 400, 64)); + files.extend((0..8).map(|_| sim_file(0, 100, 0, 8, 1 + rng.next(4)))); + files + }, + |files| pick_small_l1(files, &mapper, limits(MAX_OUTPUT)), + |files| (conflict_pairs(files), files.len()), + ); + assert!(steps >= 300); + } + + #[test] + fn test_active_l1_picks_terminate() { + let mapper = primary_key_mapper_for_test(); + let steps = run_to_fixpoint( + &mut Rng(11), + |files| pick_active_l1(files, &mapper, limits(MAX_OUTPUT)), + |files| (files.len(), 0), + ); + assert!(steps > 200); + } +} diff --git a/src/mito2/src/engine/compaction_test.rs b/src/mito2/src/engine/compaction_test.rs index 82f8bcaecc8..02f76af6f5f 100644 --- a/src/mito2/src/engine/compaction_test.rs +++ b/src/mito2/src/engine/compaction_test.rs @@ -212,6 +212,7 @@ async fn test_strict_window_output_file_size( assert!(version.ssts.levels()[0].files.is_empty()); let primary_key_mapper = version.ssts.primary_key_mapper(); let files = version.ssts.levels()[1].files().collect::>(); + let max_output_size = files.iter().map(|file| file.size()).max().unwrap(); assert_eq!( 6000, files @@ -229,6 +230,14 @@ async fn test_strict_window_output_file_size( .collect::>(); if expect_split { assert!(window_files.len() > 1, "each window must split by size"); + assert_eq!( + window_files.len() as u64, + window_files + .iter() + .map(|file| file.size()) + .sum::() + .div_ceil(max_output_size) + ); } else { assert_eq!(1, window_files.len(), "zero must disable size splitting"); } @@ -253,6 +262,35 @@ async fn test_strict_window_output_file_size( sort_batches_and_print(&before, &["tag_0", "ts"]), sort_batches_and_print(&after, &["tag_0", "ts"]) ); + + // Fresh size-split outputs have disjoint series and must not look like + // independently written files sharing a boundary key to the TWCS picker. + // Raise the tiny split threshold so an overlap merge would be eligible by + // size; otherwise the output-count guard alone would reject these files. + engine + .handle_request( + region_id, + RegionRequest::Alter(RegionAlterRequest { + kind: SetRegionOptions { + options: vec![SetRegionOption::Twsc( + "compaction.twcs.max_output_file_size".into(), + format!("{max_output_size}B"), + )], + }, + }), + ) + .await + .unwrap(); + let output_ids: HashSet<_> = files.iter().map(|file| file.file_id()).collect(); + compact(&engine, region_id).await; + let version = engine.get_region(region_id).unwrap().version(); + assert_eq!( + output_ids, + version.ssts.levels()[1] + .files() + .map(|file| file.file_id()) + .collect() + ); } /// Flush may collapse versions within one file, but compaction must not promote @@ -544,6 +582,169 @@ async fn assert_partial_compaction_preserves_delete_order(flat_format: bool) { ); } +/// Exercise the small-L1 fallback with real SST sizes, preserved tombstones, +/// and both layouts. No fabricated file metadata or hand-picked merge inputs. +#[rstest::rstest] +#[tokio::test] +async fn test_inactive_small_l1_compaction_preserves_large_peer( + #[values(false, true)] flat_format: bool, + #[values(3600, 7200)] active_ts: usize, +) { + let mut env = TestEnv::new().await; + let region_id = RegionId::new(1, 1); + let (engine, columns) = env_for_manual_compaction(&mut env, region_id, flat_format).await; + let mut rng = ::seed_from_u64(41); + let mut old_rows = build_rows_for_key("m", 10, 16, 0); + for _ in 0..30_000 { + old_rows.extend(build_rows_for_key( + &format!("a{:032x}", rand::Rng::random::(&mut rng)), + 1, + 2, + 0, + )); + } + for _ in 0..2 { + put_rows( + &engine, + region_id, + Rows { + schema: columns.clone(), + rows: old_rows.clone(), + }, + ) + .await; + flush(&engine, region_id).await; + } + compact(&engine, region_id).await; + let region = engine.get_region(region_id).unwrap(); + let table_dir = region.table_dir().to_string(); + let large = region.version().ssts.levels()[1] + .files() + .next() + .unwrap() + .clone(); + const OUTPUT_SIZE: u64 = 64 * 1024; + assert!(large.size() > 8 * OUTPUT_SIZE); + engine + .handle_request( + region_id, + RegionRequest::Alter(RegionAlterRequest { + kind: SetRegionOptions { + options: vec![ + SetRegionOption::Twsc( + "compaction.twcs.max_output_file_size".into(), + "64KiB".into(), + ), + SetRegionOption::Twsc( + "compaction.twcs.inactive_window.l1_merge_trigger".into(), + "2".into(), + ), + ], + }, + }), + ) + .await + .unwrap(); + + // Each pair of L0 files becomes one small L1. The first contains a Delete + // masking the large SST; later pairs update other rows in its key/time box. + for ts in 10..14 { + for _ in 0..2 { + let rows = Rows { + schema: columns.clone(), + rows: build_rows_for_key("m", ts, ts + 2, ts), + }; + if ts == 10 { + engine + .handle_request( + region_id, + RegionRequest::Delete(RegionDeleteRequest { + rows, + hint: None, + partition_expr_version: None, + }), + ) + .await + .unwrap(); + } else { + put_rows(&engine, region_id, rows).await; + } + flush(&engine, region_id).await; + } + compact(&engine, region_id).await; + } + put_and_flush(&engine, region_id, &columns, active_ts..active_ts + 1).await; + let version = region.version(); + let small = version.ssts.levels()[1] + .files() + .filter(|file| file.file_id() != large.file_id()) + .collect::>(); + assert_eq!(4, small.len()); + assert!(small.iter().all(|file| file.size() <= OUTPUT_SIZE / 4)); + let small_ids = small + .iter() + .map(|file| file.file_id()) + .collect::>(); + let pick = preview_regular_compaction(&engine, region_id).await; + assert_eq!(1, pick.outputs.len()); + assert_eq!( + small_ids, + pick.outputs[0] + .inputs + .iter() + .map(|file| file.file_id()) + .collect() + ); + assert!(!pick.outputs[0].filter_deleted); + let before = RecordBatches::try_collect( + engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(), + ) + .await + .unwrap() + .pretty_print() + .unwrap(); + + // Even manual cleanup cannot rewrite the dominant peer within the closed + // merge's output bound, so it executes the same small-only fallback. + compact(&engine, region_id).await; + let version = region.version(); + let l1 = version.ssts.levels()[1].files().collect::>(); + assert_eq!(2, l1.len()); + assert!(l1.iter().any(|file| file.file_id() == large.file_id())); + let output = l1 + .iter() + .find(|file| file.file_id() != large.file_id()) + .unwrap(); + assert!(!small_ids.contains(&output.file_id())); + assert_eq!(1, count_sst_deletes(&engine, region_id, output).await); + let after = RecordBatches::try_collect( + engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(), + ) + .await + .unwrap() + .pretty_print() + .unwrap(); + assert_eq!(before, after); + crate::test_util::reopen_region(&engine, region_id, table_dir, true, HashMap::new()).await; + let reopened = RecordBatches::try_collect( + engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(), + ) + .await + .unwrap() + .pretty_print() + .unwrap(); + assert_eq!(before, reopened); +} + /// Region fixture containing only identifiers, so it cannot pin an old SST snapshot. struct CompactionTestRegion { region_id: RegionId, @@ -1806,6 +2007,7 @@ async fn preview_regular_compaction( &version.options, None, None, + false, ); let compaction_region = CompactionRegion { region_id, @@ -2777,9 +2979,9 @@ async fn test_compaction_region_with_format(flat_format: bool) { // -[15.........29]- (delete) // [15.....24] // Count-first compaction consumes the first 4 SSTs as soon as the trigger is reached. - // The compacted output and final flush leave 2 SSTs. + // The manual compaction then merges the compacted output with the final flush. assert_eq!( - 2, + 1, scanner.num_files(), "unexpected files: {:?}", scanner.file_ids() diff --git a/src/mito2/src/gc/worker_test.rs b/src/mito2/src/gc/worker_test.rs index 2a66294f59b..3d389e7e1d7 100644 --- a/src/mito2/src/gc/worker_test.rs +++ b/src/mito2/src/gc/worker_test.rs @@ -399,7 +399,13 @@ async fn test_gc_worker_basic_compact() { let region = engine.get_region(region_id).unwrap(); let manifest = region.manifest_ctx.manifest().await; - let removed_file_num = manifest.removed_files.removed_files[0].files.len(); + // A manual compaction may take several rounds, each recording its removals. + let removed_file_num: usize = manifest + .removed_files + .removed_files + .iter() + .map(|removed| removed.files.len()) + .sum(); assert!(removed_file_num > 0); let version = manifest.manifest_version; @@ -488,9 +494,11 @@ async fn test_gc_worker_compact_with_ref() { let file_ref_manifest = FileRefsManifest { file_refs: HashMap::from([( region_id, - manifest.removed_files.removed_files[0] - .files + manifest + .removed_files + .removed_files .iter() + .flat_map(|removed| removed.files.iter()) .map(|removed_file| match removed_file { RemovedFile::File(file_id, v) => FileRef::new(region_id, *file_id, *v), RemovedFile::Index(file_id, v) => FileRef::new(region_id, *file_id, Some(*v)), diff --git a/src/object-store/tests/object_store_test.rs b/src/object-store/tests/object_store_test.rs index de526834bad..d38480e6302 100644 --- a/src/object-store/tests/object_store_test.rs +++ b/src/object-store/tests/object_store_test.rs @@ -304,6 +304,9 @@ async fn test_s3_backend() -> Result<()> { } let store = ObjectStore::new(builder).unwrap(); + // Match production S3 clients: retry transient transport failures without + // retrying conditional-write conflicts or weakening their assertions. + let store = object_store::util::with_retry_layers(store); let store = object_store::util::with_instrument_layers(store, false); let guard = TempFolder::new(&store, "/"); From 2eff52c7575e8adc5a720cf55bd15747181deac7 Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Thu, 22 Oct 2026 13:36:30 +0000 Subject: [PATCH 075/239] feat(otlp)!: support raw delta exponential histograms (#9754) * feat(otlp): support raw delta exponential histograms Signed-off-by: shuiyisong * fix(otlp): guard point accounting and benchmark metrics ingestion Signed-off-by: shuiyisong * fix: remove otlp_metric_compat logic Signed-off-by: shuiyisong * perf(otlp): cache physical schema lookups during metric preflight Signed-off-by: shuiyisong * chore: minor rename Signed-off-by: shuiyisong * perf(otlp): batch metric rows and reuse schema validation Signed-off-by: shuiyisong * chore: comments Signed-off-by: shuiyisong * fix: replace super in import path Signed-off-by: shuiyisong --------- Signed-off-by: shuiyisong --- Cargo.lock | 1 - config/config.md | 4 +- config/frontend.example.toml | 4 +- config/standalone.example.toml | 4 +- docs/rfcs/2026-05-28-table-semantic-layer.md | 31 +- docs/rfcs/2026-08-04-native-histograms.md | 118 +- src/common/frontend/src/metrics.rs | 9 +- src/common/query/src/prelude.rs | 2 - src/frontend/AGENTS.md | 5 +- src/frontend/Cargo.toml | 1 - src/frontend/src/instance.rs | 195 +-- src/frontend/src/instance/builder.rs | 2 - src/frontend/src/instance/otlp.rs | 64 +- src/metric-engine/src/engine.rs | 20 + src/metric-engine/src/engine/alter.rs | 8 +- .../src/engine/alter/extract_new_columns.rs | 100 +- src/metric-engine/src/engine/create.rs | 8 +- .../src/engine/create/extract_new_columns.rs | 83 +- src/operator/src/insert.rs | 76 +- src/operator/src/insert/otlp.rs | 828 ++++++++++++ src/operator/src/req_convert/insert.rs | 1 + .../src/req_convert/insert/row_to_batch.rs | 2 +- src/query/src/optimizer/windowed_sort.rs | 19 + src/servers/Cargo.toml | 8 + src/servers/benches/otlp_metrics.rs | 84 ++ src/servers/benches/otlp_metrics/README.md | 51 + src/servers/benches/otlp_metrics/fixtures.rs | 157 +++ src/servers/benches/otlp_metrics_preflight.rs | 260 ++++ src/servers/src/error.rs | 11 +- src/servers/src/http/otlp.rs | 1 - src/servers/src/otlp/metrics.rs | 841 ++++++------- src/servers/src/otlp/metrics/prepare.rs | 623 +++++++++ src/servers/src/otlp/metrics/resource_info.rs | 133 +- .../otlp/metrics/resource_info/tests/delta.rs | 41 +- src/servers/src/otlp/metrics/tests/delta.rs | 645 +++++++++- src/servers/src/otlp/metrics/translator.rs | 34 - src/servers/src/row_writer.rs | 43 +- src/servers/src/semantic.rs | 11 +- src/session/src/protocol_ctx.rs | 8 +- src/table/src/requests.rs | 1 + tests-integration/src/instance.rs | 6 + tests-integration/src/otlp.rs | 1117 ++++++++++++++++- tests-integration/tests/http.rs | 248 ++++ .../promql/delta_exponential_histogram.result | 106 ++ .../promql/delta_exponential_histogram.sql | 51 + 45 files changed, 5051 insertions(+), 1014 deletions(-) create mode 100644 src/operator/src/insert/otlp.rs create mode 100644 src/servers/benches/otlp_metrics.rs create mode 100644 src/servers/benches/otlp_metrics/README.md create mode 100644 src/servers/benches/otlp_metrics/fixtures.rs create mode 100644 src/servers/benches/otlp_metrics_preflight.rs create mode 100644 src/servers/src/otlp/metrics/prepare.rs create mode 100644 tests/cases/standalone/promql/delta_exponential_histogram.result create mode 100644 tests/cases/standalone/promql/delta_exponential_histogram.sql diff --git a/Cargo.lock b/Cargo.lock index b87ac7bc9f9..c8dfd19620a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5656,7 +5656,6 @@ dependencies = [ "common-test-util", "common-time", "common-version", - "dashmap", "datafusion", "datafusion-expr", "datatypes", diff --git a/config/config.md b/config/config.md index a549f5b78ad..d450609a926 100644 --- a/config/config.md +++ b/config/config.md @@ -82,7 +82,7 @@ | `pending_rows_batcher.worker_channel_capacity` | Integer | `65536` | Maximum queued submissions per table worker. | | `pending_rows_batcher.max_inflight_requests` | Integer | `3000` | Maximum admitted original requests awaiting completion. | | `pending_rows_batcher.flow_notification_queue_capacity` | Integer | `1024` | Maximum number of queued table Flow notifications. | -| `pending_rows_batcher.logical_table` | -- | -- | Metric-engine logical-table batching for Prom remote write and non-legacy OTLP metrics.
Requires prom_store.with_metric_engine. Logs, traces and legacy metrics are not eligible.
Enable independently with protocols and a nonzero flush interval.
Omitted fields use independent defaults, not parent settings.
Empty protocols or a zero interval disables logical batching without fallback.
Omitting this entire section preserves legacy Prom batching; it does not enable OTLP batching. | +| `pending_rows_batcher.logical_table` | -- | -- | Metric-engine logical-table batching for Prom remote write and OTLP metrics.
Requires prom_store.with_metric_engine. Logs and traces are not eligible.
Enable independently with protocols and a nonzero flush interval.
Omitted fields use independent defaults, not parent settings.
Empty protocols or a zero interval disables logical batching without fallback.
Omitting this entire section preserves legacy Prom batching; it does not enable OTLP batching. | | `pending_rows_batcher.logical_table.pending_rows_flush_interval` | String | `0s` | Flush interval measured from the first pending submission. Zero disables batching. | | `pending_rows_batcher.logical_table.max_batch_rows` | Integer | `100000` | Flush after a complete submission reaches this row threshold. | | `pending_rows_batcher.logical_table.max_concurrent_flushes` | Integer | `256` | Maximum concurrent flushes shared by Prom and OTLP metrics. | @@ -358,7 +358,7 @@ | `pending_rows_batcher.worker_channel_capacity` | Integer | `65536` | Maximum queued submissions per table worker. | | `pending_rows_batcher.max_inflight_requests` | Integer | `3000` | Maximum admitted original requests awaiting completion. | | `pending_rows_batcher.flow_notification_queue_capacity` | Integer | `1024` | Maximum number of queued table Flow notifications. | -| `pending_rows_batcher.logical_table` | -- | -- | Metric-engine logical-table batching for Prom remote write and non-legacy OTLP metrics.
Requires prom_store.with_metric_engine. Logs, traces and legacy metrics are not eligible.
Enable independently with protocols and a nonzero flush interval.
Omitted fields use independent defaults, not parent settings.
Empty protocols or a zero interval disables logical batching without fallback.
Omitting this entire section preserves legacy Prom batching; it does not enable OTLP batching. | +| `pending_rows_batcher.logical_table` | -- | -- | Metric-engine logical-table batching for Prom remote write and OTLP metrics.
Requires prom_store.with_metric_engine. Logs and traces are not eligible.
Enable independently with protocols and a nonzero flush interval.
Omitted fields use independent defaults, not parent settings.
Empty protocols or a zero interval disables logical batching without fallback.
Omitting this entire section preserves legacy Prom batching; it does not enable OTLP batching. | | `pending_rows_batcher.logical_table.pending_rows_flush_interval` | String | `0s` | Flush interval measured from the first pending submission. Zero disables batching. | | `pending_rows_batcher.logical_table.max_batch_rows` | Integer | `100000` | Flush after a complete submission reaches this row threshold. | | `pending_rows_batcher.logical_table.max_concurrent_flushes` | Integer | `256` | Maximum concurrent flushes shared by Prom and OTLP metrics. | diff --git a/config/frontend.example.toml b/config/frontend.example.toml index 6d7e74824d1..c959dc10926 100644 --- a/config/frontend.example.toml +++ b/config/frontend.example.toml @@ -274,8 +274,8 @@ max_inflight_requests = 3000 ## Maximum number of queued table Flow notifications. flow_notification_queue_capacity = 1024 -## Metric-engine logical-table batching for Prom remote write and non-legacy OTLP metrics. -## Requires prom_store.with_metric_engine. Logs, traces and legacy metrics are not eligible. +## Metric-engine logical-table batching for Prom remote write and OTLP metrics. +## Requires prom_store.with_metric_engine. Logs and traces are not eligible. ## Enable independently with protocols and a nonzero flush interval. ## Omitted fields use independent defaults, not parent settings. ## Empty protocols or a zero interval disables logical batching without fallback. diff --git a/config/standalone.example.toml b/config/standalone.example.toml index 3aedde0aec0..f70277aa599 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -246,8 +246,8 @@ max_inflight_requests = 3000 ## Maximum number of queued table Flow notifications. flow_notification_queue_capacity = 1024 -## Metric-engine logical-table batching for Prom remote write and non-legacy OTLP metrics. -## Requires prom_store.with_metric_engine. Logs, traces and legacy metrics are not eligible. +## Metric-engine logical-table batching for Prom remote write and OTLP metrics. +## Requires prom_store.with_metric_engine. Logs and traces are not eligible. ## Enable independently with protocols and a nonzero flush interval. ## Omitted fields use independent defaults, not parent settings. ## Empty protocols or a zero interval disables logical batching without fallback. diff --git a/docs/rfcs/2026-05-28-table-semantic-layer.md b/docs/rfcs/2026-05-28-table-semantic-layer.md index 79764ad80d2..37a6de3e689 100644 --- a/docs/rfcs/2026-05-28-table-semantic-layer.md +++ b/docs/rfcs/2026-05-28-table-semantic-layer.md @@ -100,24 +100,19 @@ Two design decisions worth pinning down up front, because they constrain everyth - **Conflict.** Some table-level keys (`trace.conventions` lifted from `schema_url`, `metric.temporality`, ...) cannot represent the truth when a long-lived table sees rows from multiple sources. v1 records `mixed` or `unknown` rather than a fictitious single value. Downstream consumers must treat any single-valued semantic key as best-effort, not strong evidence. - **Update.** Semantic options are stamped at table creation. v1 does not specify an update path; promoting `metadata_quality` from `inferred` to `declared`, refreshing `resource.attributes_preserved`, or revising `trace.conventions` on later writes is deferred. If real usage shows update is needed, it lands as a separate RFC. -OTLP delta sums and explicit histograms store raw interval values without -accumulating across timestamps. Delta-to-cumulative conversion and its -per-series ingestion state are intentionally out of scope. Explicit histogram -bucket counts are prefix-summed only within each point to produce the classic -histogram representation. Delta exponential histograms are currently rejected; -[planned raw-delta support](2026-08-04-native-histograms.md#planned-raw-delta-exponential-histograms) -will store their interval values in the canonical native-histogram Struct using -the same temporality tag and a gauge reset hint. The tag will select histogram -summation for `increase()` and summation divided by range duration for `rate()`, -without cumulative reset correction or extrapolation. +OTLP delta sums, explicit histograms and exponential histograms store raw interval +values without accumulation across timestamps or requests. Explicit histogram +buckets are prefix-summed within each point. Exponential histograms retain their +canonical native-histogram Struct and use a gauge reset hint for delta samples. +All use the fixed String tag `otlp_aggregation_temporality="delta"`, independent +of `default_column_prefix`, as part of series identity. + +For tagged float series, the marker selects existing delta `rate()`/`increase()` +behavior. Native histogram counter functions are unchanged: query raw delta +histograms with `sum_over_time(metric[3m])`, or divide that histogram by `180` +for a per-second result. See [raw-delta exponential histograms](2026-08-04-native-histograms.md#raw-delta-exponential-histograms). +Table options never select per-row query behavior. -OTLP delta sums and explicit histograms additionally store the query-visible -String tag `otlp_aggregation_temporality="delta"` on each generated row. Its -name is fixed and does not follow `default_column_prefix`. The tag is part of -series identity and is authoritative for per-series float `rate()` and -`increase()` behavior; the table option is never used as a row-level -discriminator. Native histograms currently retain their native algorithms; -the planned delta extension will change only tagged delta histogram series. Prometheus metadata reports `unknown` for counter, histogram, and up/down-counter tables whose catalog temporality is `delta` or `mixed`. A same-request conflict can create `mixed`; a later write does not update an existing table option, so the @@ -143,7 +138,7 @@ schema lacks it; unrelated nullable labels retain their existing matching behavior. Ignoring the marker adds no marker-related matching work. Use `ignoring(otlp_aggregation_temporality)` when temporality should not -participate in matching. Apply `rate()` or `increase()` before an aggregation +participate in matching. For float series, apply `rate()` or `increase()` before an aggregation or subquery that drops the marker. `irate()` and `resets()` are not temporality-aware in v1; `delta()`, `idelta()`, and `changes()` retain their existing raw-sample semantics. diff --git a/docs/rfcs/2026-08-04-native-histograms.md b/docs/rfcs/2026-08-04-native-histograms.md index 3aef4eb7ffb..6b0a9ebb757 100644 --- a/docs/rfcs/2026-08-04-native-histograms.md +++ b/docs/rfcs/2026-08-04-native-histograms.md @@ -126,20 +126,15 @@ explicit if sampled responses are implemented. OTLP/HTTP exponential histograms are accepted by default. OTel Arrow exponential histograms are rejected because the current Arrow wire format omits `zero_threshold`; accepting them would -silently change the distribution. Currently, cumulative temporality is required; -delta and unspecified exponential histograms are rejected before their points are -converted. Explicit OTLP histograms keep their existing `_bucket`, `_sum`, and -`_count` representation. OTLP delta sums and explicit histograms are stored as -raw interval values with `otlp_aggregation_temporality="delta"`. Explicit -histogram bucket counts are prefix-summed within each point for the classic -histogram representation, but values are never accumulated across timestamps. -For these tagged float series, PromQL `increase()` sums the interval values -in the query range, and `rate()` divides that sum by the range duration. -Delta-to-cumulative conversion is intentionally out of scope, not deferred work. -Raw-delta exponential-histogram support is planned below; delta exponential -points remain rejected until it is implemented. See -[Table Semantic Layer](2026-05-28-table-semantic-layer.md#conflict-and-update-semantics) -for the per-series temporality contract. +silently change the distribution. Cumulative and delta temporality are accepted; +unspecified and unrecognized temporality are rejected before processing points. +Explicit OTLP histograms keep their existing `_bucket`, `_sum`, and `_count` +representation. All delta metrics store raw interval values with +`otlp_aggregation_temporality="delta"`; ingestion never accumulates across points +or requests. Explicit histogram bucket counts are prefix-summed within a point. +Tagged float series retain GreptimeDB's automatic delta `rate()`/`increase()` +extension. Native histograms use the explicit queries below instead. +Delta-to-cumulative conversion is outside this feature. OTLP scales `-4` through `8` map directly to Prometheus schemas. Higher scales are downscaled to schema `8`: dense counts that collide are merged before the @@ -148,11 +143,14 @@ This preserves count mass but irreversibly loses distinctions between source buckets that merge. Lower scales are rejected. OTLP counts always populate the integer histogram family. The transmitted non-negative finite zero threshold and zero count are preserved; the Struct start timestamp is stored in -milliseconds and the reset hint is unknown. Point timestamps and attributes -retain their existing mode-specific conversion rules. Legacy mode retains its -normalized name, attribute rules, and nanosecond row timestamp. Non-legacy mode -retains Prometheus-compatible translation and its millisecond row timestamp. -Both modes write the same canonical Struct. +milliseconds. The reset hint is gauge for delta and unknown for cumulative samples. +OTLP metrics always use Prometheus-compatible name and attribute conversion; +`otlp_metric_compat` remains accepted as table metadata but does not select a +conversion mode. Metric-engine conversion retains nanoseconds until alignment +to the destination time-index unit: seconds, milliseconds, microseconds, or nanoseconds. +New logical tables inherit their backing table's unit; auto-created physical +tables default to milliseconds. Ordinary insertion uses milliseconds. +Both storage paths write the same canonical Struct. An absent sum is stored as an ordinary quiet NaN, so the sample remains selectable while `histogram_sum` is unknown. Any NaN sum on a point without @@ -172,33 +170,71 @@ produced an accepted row. Zero-count buckets are omitted after downscaling, splitting nonzero runs into sparse spans without changing integer counts. Rejected data points increment `greptime_servers_otlp_exponential_histogram_rejected_data_points_total`, labeled -by `reason`: `delta_temporality`, `unspecified_temporality`, or -`invalid_data_point`. Each rejected point is counted once, including when +by `reason`: `unspecified_temporality` or `invalid_data_point`. Attribute and +schema rejections use `invalid_data_point`; other metric families do not +increment this histogram-specific counter. Each rejected point is counted once, including when resource descriptors are enabled. Minimum, maximum, and exemplars are not persisted. -### Planned raw-delta exponential histograms +### Raw-delta exponential histograms -Accept OTLP delta exponential histograms and store each point's interval counts -and sum in the existing canonical native-histogram Struct, without accumulating -across timestamps. Reuse the integer-count encoding, scale conversion, -zero-threshold preservation, stale-marker handling, and partial-success path. -OTel Arrow support remains blocked independently by its missing `zero_threshold`. +Each stored delta sample retains its interval count, sum, zero count and buckets. +The temporality tag separates delta and cumulative series under the same table +name. Query raw deltas explicitly: -Store `otlp_aggregation_temporality="delta"` as part of series identity, as for -delta sums and explicit histograms. Use the gauge reset hint for delta histogram -samples so decreasing interval counts are not interpreted as counter resets. -Prometheus uses the same gauge hint for -[raw-delta exponential-histogram ingestion](https://github.com/prometheus/prometheus/blob/9f159f6df198280b757163ec555ffde0c8bfe8e0/storage/remote/otlptranslator/prometheusremotewrite/histograms.go#L125). -The temporality tag, rather than the gauge hint alone or table metadata, -determines raw-delta query behavior. +```promql +sum_over_time(request_duration{otlp_aggregation_temporality="delta"}[3m]) +sum_over_time(request_duration{otlp_aggregation_temporality="delta"}[3m]) / 180 +histogram_count(sum_over_time(request_duration{otlp_aggregation_temporality="delta"}[3m])) +``` -Extend the per-series `increase()` and `rate()` selection to native histogram -inputs: sum interval histograms over the query range, then divide by the range -duration for `rate()`. Do not apply cumulative reset correction, synthetic-zero -insertion, or counter extrapolation to these delta samples. Preserve cumulative -histogram behavior and mixed float/histogram range warnings. Acceptance coverage -must include multiple intervals whose counts decrease, exact integer ingestion, -stale points, partial rejection, and pure-native and mixed-range query paths. +Samples are selected by their endpoint timestamp in the existing left-open, +right-closed range. Each selected sample contributes its entire interval, +including intervals crossing the left boundary. A single sample contributes its +full value; an empty window emits no sample. Counts `10,20,15` at seconds +`60,120,180` give count `45` at second `180`, or `0.25` after division by `180`. +Missing intervals do not shorten the divisor. There is no prorating, reset +correction, extrapolation, synthetic zero insertion or inferred coverage. + +Native-histogram `rate()` and `increase()` retain counter behavior, regardless +of the tag; they are not supported queries for raw deltas. Converting a histogram +to a float can enter the existing scalar-delta extension, so the explicit +expressions above are the supported examples. Existing sum/division UDFs and +serialization are reused without a new wire format. + +Under default storage options, identical table/tags/stored timestamps overwrite +with the last point, without adding buckets. Timestamp alignment can collapse +endpoints at coarser units; interval starts do not form part of the key. Preserve +source order through default resolution and row merging. Existing storage +options continue to apply. + +Metrics recovery is point-based across gauges, sums, explicit histograms, +summaries and exponential histograms. Point-local validation precedes +existing-schema validation, which precedes simultaneous detection of fresh +logical/physical conflicts. Reject the union of contributors to all conflicts; +remove every output of a rejected point. Defaults distinguish omitted columns +from explicit nulls. Auto-DDL configuration and hints remain authoritative; +metric-engine timestamp alignment remains permitted with auto-DDL disabled. +Only surviving points contribute metadata and descriptor windows. Metadata +describes actual outputs, excluding absent sums and empty quantile outputs. + +Authentication, authorization, admission, catalog and primary-write failures +retain their failure paths. Synchronous acceptance waits for writes; eligible +asynchronous batching acknowledges enqueue, and later flush errors retain their +existing metrics. Descriptor failures after primary acceptance remain warnings, +including a partial-success warning with zero rejected points. No write is +retried in smaller subsets after an ambiguous result. + +Prometheus compatibility follows raw-delta ingestion with a gauge reset hint +and explicit summation. GreptimeDB retains its temporality label, transmitted +zero threshold, NaN for absent sums, canonical stale representation, integer +bounds, sparse compaction and existing arithmetic. Equivalent distributions +need not have identical layouts or floating-point results. Start timestamps +remain stored but explicit summation does not use them for counter correction. +Public metadata is `unknown` for freshly created delta/mixed-temporality tables; +existing creation-time options are not refreshed by later writes. + +Upgrade receiving ingestion frontends before enabling delta senders. Mixed-version +query execution is unverified; this feature introduces no new query UDF or format. # PromQL Compatibility Decisions diff --git a/src/common/frontend/src/metrics.rs b/src/common/frontend/src/metrics.rs index 7a2c152fa1f..c160e57db7b 100644 --- a/src/common/frontend/src/metrics.rs +++ b/src/common/frontend/src/metrics.rs @@ -14,12 +14,12 @@ //! Cache lookup counters shared by frontend, pipeline, and protocol handlers. //! -//! The `type` label identifies `otlp_trace_aux`, `otlp_metrics_legacy`, +//! The `type` label identifies `otlp_trace_aux`, //! `pipeline_table`, `pipeline`, `pipeline_original`, `pipeline_failover`, or //! `mysql_prepared_stmt`. Counters aggregate across catalogs and connections. //! Trace lookups count distinct service/operation keys once per request, excluding -//! rechecks; legacy metric lookups count each requested metric name. Pipeline -//! loaders count successful coalesced waiters as hits and failed loads as misses. +//! rechecks. Pipeline loaders count successful coalesced waiters as hits and +//! failed loads as misses. //! Failover counts one outcome after resolving all candidate schemas; an ambiguous //! result is a miss. Table and prepared-statement caches count each lookup. //! Both outcomes for all known caches are exported from the first metric update. @@ -27,9 +27,8 @@ use lazy_static::lazy_static; use prometheus::{IntCounterVec, register_int_counter_vec}; -const CACHE_TYPES: [&str; 7] = [ +const CACHE_TYPES: [&str; 6] = [ "otlp_trace_aux", - "otlp_metrics_legacy", "pipeline_table", "pipeline", "pipeline_original", diff --git a/src/common/query/src/prelude.rs b/src/common/query/src/prelude.rs index 62e5da34207..ca70bf5520f 100644 --- a/src/common/query/src/prelude.rs +++ b/src/common/query/src/prelude.rs @@ -91,8 +91,6 @@ pub fn greptime_native_histogram() -> &'static str { const GREPTIME_TIMESTAMP: &str = "greptime_timestamp"; /// Default value column name constant for backward compatibility. const GREPTIME_VALUE: &str = "greptime_value"; -/// Default counter column name for OTLP metrics (legacy mode). -pub const GREPTIME_COUNT: &str = "greptime_count"; /// Stored series label that opts ordinary float samples into raw-delta math. pub const OTLP_AGGREGATION_TEMPORALITY_LABEL: &str = "otlp_aggregation_temporality"; /// Authoritative raw-delta value for [`OTLP_AGGREGATION_TEMPORALITY_LABEL`]. diff --git a/src/frontend/AGENTS.md b/src/frontend/AGENTS.md index 8d858242e8c..152cbbe6bee 100644 --- a/src/frontend/AGENTS.md +++ b/src/frontend/AGENTS.md @@ -74,8 +74,11 @@ including caches owned by pipeline and protocol handlers. These tables are ingestion-managed; manual mutation requires restarting serving frontends. +- **OTLP metrics** (`instance/otlp.rs`): authorize all emitted destinations before + schema preflight can reject points or expose destination schemas. + - **Logical-table batching** (`instance/logical_batcher.rs`): `Services` initializes - one shared batcher for opted-in HTTP Prom and nonlegacy OTLP metric-engine + one shared batcher for opted-in HTTP Prom and OTLP metric-engine writes. Downstream routers can initialize it when enabling replacement endpoints. `Instance::handle_otlp_metric_row_inserts` shares OTLP eligibility and dispatch between converters and falls back for incompatible tables. diff --git a/src/frontend/Cargo.toml b/src/frontend/Cargo.toml index 44094e662a8..1fd347c68e9 100644 --- a/src/frontend/Cargo.toml +++ b/src/frontend/Cargo.toml @@ -50,7 +50,6 @@ common-stat.workspace = true common-telemetry.workspace = true common-time.workspace = true common-version.workspace = true -dashmap.workspace = true datafusion.workspace = true datafusion-expr.workspace = true datatypes.workspace = true diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index cbbf2c7babf..2e49bf0b7cc 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -53,11 +53,9 @@ use common_base::cancellation::CancellableFuture; use common_batcher::request_limiter::RequestLimiter; use common_error::ext::{BoxedError, ErrorExt}; use common_event_recorder::EventRecorderRef; -use common_frontend::metrics::{CACHE_HIT, CACHE_MISS}; use common_meta::cache::TableFlownodeSetCacheRef; use common_meta::cache_invalidator::CacheInvalidatorRef; use common_meta::key::TableMetadataManagerRef; -use common_meta::key::table_name::TableNameKey; use common_meta::node_manager::NodeManagerRef; use common_meta::procedure_executor::ProcedureExecutorRef; use common_query::Output; @@ -65,12 +63,10 @@ use common_recordbatch::RecordBatchStreamWrapper; use common_recordbatch::error::StreamTimeoutSnafu; use common_telemetry::logging::SlowQueryOptions; use common_telemetry::{debug, error, tracing}; -use dashmap::DashMap; use datafusion::dataframe::DataFrame; use datafusion::physical_plan::ExecutionPlan; use datafusion_expr::LogicalPlan; use futures::{Stream, StreamExt, future}; -use lazy_static::lazy_static; use operator::delete::DeleterRef; use operator::insert::InserterRef; use operator::statement::{StatementExecutor, StatementExecutorRef}; @@ -84,14 +80,10 @@ use query::parser::{PromQuery, QueryStatement}; use query::query_engine::DescribeResult; use query::query_engine::options::{QueryOptions, validate_catalog_and_schema}; use servers::batcher::logical_table::LogicalTablePendingRowsBatcher; -use servers::error::{ - self as server_error, AuthSnafu, CommonMetaSnafu, ExecuteQuerySnafu, - OtlpMetricModeIncompatibleSnafu, UnexpectedResultSnafu, -}; +use servers::error::{self as server_error, AuthSnafu, ExecuteQuerySnafu, UnexpectedResultSnafu}; use servers::interceptor::{ PromQueryInterceptor, PromQueryInterceptorRef, SqlQueryInterceptor, SqlQueryInterceptorRef, }; -use servers::otlp::metrics::legacy_normalize_otlp_name; use servers::prometheus_handler::{ ParsedPromQuery, PrometheusHandler, resolve_schema_from_matchers, }; @@ -108,7 +100,6 @@ use sql::statements::statement::Statement; use sql::statements::tql::Tql; use sql::util::{extract_tables_from_prom_expr_checked, extract_tables_from_statement_checked}; use sqlparser::ast::{AnalyzeFormat, ObjectName}; -use table::requests::{OTLP_METRIC_COMPAT_KEY, OTLP_METRIC_COMPAT_PROM}; use tracing::Span; use crate::error::{ @@ -119,10 +110,6 @@ use crate::error::{ use crate::service_config::InfluxdbMergeMode; use crate::stream_wrapper::CancellableStreamWrapper; -lazy_static! { - static ref OTLP_LEGACY_DEFAULT_VALUE: String = "legacy".to_string(); -} - /// The frontend instance contains necessary components, and implements many /// traits, like [`servers::query_handler::grpc::GrpcQueryHandler`], /// [`servers::query_handler::sql::SqlQueryHandler`], etc. @@ -155,12 +142,6 @@ pub struct Instance { common_grpc::flow_write::Reply, >, >, - - // cache for otlp metrics - // first layer key: db-string - // key: direct input metric name - // value: if runs in legacy mode - otlp_metrics_table_legacy_cache: Arc>>, } impl Instance { @@ -499,137 +480,6 @@ impl Instance { .await .context(TableOperationSnafu) } - - async fn check_otlp_legacy( - &self, - names: &[String], - ctx: &QueryContextRef, - ) -> server_error::Result { - let db_string = ctx.get_db_string(); - // fast cache check - let cache = self - .otlp_metrics_table_legacy_cache - .entry(db_string.clone()) - .or_default(); - if let Some(flag) = fast_legacy_check(&cache, names)? { - return Ok(flag); - } - // release cache reference to avoid lock contention - drop(cache); - - let catalog = ctx.current_catalog(); - let schema = ctx.current_schema(); - - // query legacy table names - let normalized_names = names - .iter() - .map(|n| legacy_normalize_otlp_name(n)) - .collect::>(); - let table_names = normalized_names - .iter() - .map(|n| TableNameKey::new(catalog, &schema, n)) - .collect::>(); - let table_values = self - .table_metadata_manager() - .table_name_manager() - .batch_get(table_names) - .await - .context(CommonMetaSnafu)?; - let table_ids = table_values - .into_iter() - .filter_map(|v| v.map(|vi| vi.table_id())) - .collect::>(); - - // means no existing table is found, use new mode - if table_ids.is_empty() { - return Ok(false); - } - - // has existing table, check table options - let table_infos = self - .table_metadata_manager() - .table_info_manager() - .batch_get(&table_ids) - .await - .context(CommonMetaSnafu)?; - let options = table_infos - .values() - .map(|info| { - info.table_info - .meta - .options - .extra_options - .get(OTLP_METRIC_COMPAT_KEY) - .unwrap_or(&OTLP_LEGACY_DEFAULT_VALUE) - }) - .collect::>(); - if !options.is_empty() { - // check value consistency - let has_prom = options.iter().any(|opt| *opt == OTLP_METRIC_COMPAT_PROM); - let has_legacy = options - .iter() - .any(|opt| *opt == OTLP_LEGACY_DEFAULT_VALUE.as_str()); - ensure!(!(has_prom && has_legacy), OtlpMetricModeIncompatibleSnafu); - Ok(has_legacy) - } else { - // no table info, use new mode - Ok(false) - } - } - - fn cache_otlp_legacy( - &self, - names: &[String], - ctx: &QueryContextRef, - is_legacy: bool, - ) -> server_error::Result<()> { - let cache = self - .otlp_metrics_table_legacy_cache - .entry(ctx.get_db_string()) - .or_default(); - cache_legacy_mode(&cache, names, is_legacy) - } -} - -fn fast_legacy_check( - cache: &DashMap, - names: &[String], -) -> server_error::Result> { - let hit_cache = names - .iter() - .filter_map(|name| cache.get(name)) - .collect::>(); - CACHE_HIT - .with_label_values(&["otlp_metrics_legacy"]) - .inc_by(hit_cache.len() as u64); - CACHE_MISS - .with_label_values(&["otlp_metrics_legacy"]) - .inc_by((names.len() - hit_cache.len()) as u64); - if !hit_cache.is_empty() { - let hit_legacy = hit_cache.iter().any(|en| *en.value()); - let hit_prom = hit_cache.iter().any(|en| !*en.value()); - - // hit but have true and false, means both legacy and new mode are used - // we cannot handle this case, so return error - // add doc links in err msg later - ensure!(!(hit_legacy && hit_prom), OtlpMetricModeIncompatibleSnafu); - - Ok(Some(hit_legacy)) - } else { - Ok(None) - } -} - -fn cache_legacy_mode( - cache: &DashMap, - names: &[String], - is_legacy: bool, -) -> server_error::Result<()> { - for name in names { - let cached = cache.entry(name.clone()).or_insert(is_legacy); - ensure!(*cached == is_legacy, OtlpMetricModeIncompatibleSnafu); - } - Ok(()) } /// If the relevant variables are set, the timeout is enforced for all PostgreSQL statements. @@ -1996,6 +1846,7 @@ mod tests { use common_error::ext::{BoxedError, ErrorExt, PlainError}; use common_error::status_code::StatusCode; use common_event_recorder::{Event, EventRecorder, EventTypeFilter, EventTypeFilterRef}; + use common_frontend::metrics::{CACHE_HIT, CACHE_MISS}; use common_frontend::slow_query_event::SlowQueryEvent; use common_meta::cache::LayeredCacheRegistryBuilder; use common_meta::kv_backend::memory::MemoryKvBackend; @@ -3456,21 +3307,6 @@ mod tests { Ok(()) } - #[tokio::test] - async fn test_instance_shared_state() -> TestResult<()> { - let instance = - test_instance_with_tables(test_table(1024, "source")?, test_table(1025, "target")?) - .await?; - - let _event_recorder = instance.event_recorder(); - assert!(Arc::ptr_eq( - &instance.otlp_metrics_table_legacy_cache, - &instance.clone().otlp_metrics_table_legacy_cache, - )); - - Ok(()) - } - #[tokio::test] async fn test_restricted_endpoint_handlers_check_permissions() -> TestResult<()> { let checker = Arc::new(RejectEndpointPermissionChecker::default()); @@ -3927,33 +3763,6 @@ mod tests { Ok(()) } - #[test] - fn test_fast_legacy_check_is_read_only() { - let hits = CACHE_HIT.with_label_values(&["otlp_metrics_legacy"]); - let misses = CACHE_MISS.with_label_values(&["otlp_metrics_legacy"]); - let before = (hits.get(), misses.get()); - let cache = DashMap::new(); - cache.insert("metric1".to_string(), true); - - let names = vec!["metric1".to_string(), "metric2".to_string()]; - assert_eq!(Some(true), fast_legacy_check(&cache, &names).unwrap()); - assert_eq!((hits.get(), misses.get()), (before.0 + 1, before.1 + 1)); - assert!(!cache.contains_key("metric2")); - - cache_legacy_mode(&cache, &names, true).unwrap(); - assert!(*cache.get("metric2").unwrap().value()); - assert!(cache_legacy_mode(&cache, &names, false).is_err()); - assert!(*cache.get("metric2").unwrap().value()); - - let cache_incompatible = DashMap::new(); - cache_incompatible.insert("metric1".to_string(), true); - cache_incompatible.insert("metric2".to_string(), false); - assert!(fast_legacy_check(&cache_incompatible, &names).is_err()); - assert_eq!((hits.get(), misses.get()), (before.0 + 3, before.1 + 1)); - assert_eq!(None, fast_legacy_check(&cache, &[]).unwrap()); - assert_eq!((hits.get(), misses.get()), (before.0 + 3, before.1 + 1)); - } - #[test] fn test_should_track_statement_process() { assert!(should_track_statement_process(&parse_one_sql( diff --git a/src/frontend/src/instance/builder.rs b/src/frontend/src/instance/builder.rs index 0a03163199b..732c587e07b 100644 --- a/src/frontend/src/instance/builder.rs +++ b/src/frontend/src/instance/builder.rs @@ -32,7 +32,6 @@ use common_meta::key::flow::FlowMetadataManager; use common_meta::kv_backend::KvBackendRef; use common_meta::node_manager::NodeManagerRef; use common_meta::procedure_executor::ProcedureExecutorRef; -use dashmap::DashMap; use operator::batcher::PendingRowsBatcher; use operator::delete::Deleter; use operator::flow::FlowServiceOperator; @@ -432,7 +431,6 @@ impl FrontendBuilder { event_recorder, slow_query_recorder, process_manager, - otlp_metrics_table_legacy_cache: Arc::new(DashMap::new()), slow_query_options: self.options.slow_query.clone(), influxdb_default_merge_mode: self.options.influxdb.default_merge_mode, trace_ingest_chunk_size: self.options.otlp.trace_ingest_chunk_size, diff --git a/src/frontend/src/instance/otlp.rs b/src/frontend/src/instance/otlp.rs index e4f395db457..4ed79e7fb41 100644 --- a/src/frontend/src/instance/otlp.rs +++ b/src/frontend/src/instance/otlp.rs @@ -16,6 +16,7 @@ mod trace_ingest; pub mod trace_semconv; pub mod trace_types; +use std::collections::HashSet; use std::sync::Arc; use api::v1::RowInsertRequests; @@ -108,31 +109,53 @@ impl OpenTelemetryProtocolHandler for Instance { interceptor_ref.pre_execute(ctx.clone())?; let ctx = Arc::new(ctx.fork()); - let input_names = request - .resource_metrics - .iter() - .flat_map(|r| r.scope_metrics.iter()) - .flat_map(|s| s.metrics.iter().map(|m| m.name.clone())) - .collect::>(); - - // See [`OtlpMetricCtx`] for details - let is_legacy = self.check_otlp_legacy(&input_names, &ctx).await?; - let mut metric_ctx = ctx .protocol_ctx() .get_otlp_metric_ctx() .cloned() .unwrap_or_default(); - metric_ctx.is_legacy = is_legacy; metric_ctx.resource_info = self.otlp_resource_info; + let mut prepared = otlp::metrics::prepare_metrics(request, &mut metric_ctx)?; + // Check all destinations before preflight can expose schemas or discard points. + let targets = { + let catalog = ctx.current_catalog(); + let schema = ctx.current_schema(); + let mut table_names = HashSet::new(); + prepared + .requests + .iter() + .filter(|request| table_names.insert(request.table_name.as_str())) + .map(|request| PermissionTableTarget::new(catalog, &schema, &request.table_name)) + .collect() + }; + self.check_table_permission( + &ctx, + PermissionReq::Action(OTLP_WRITE), + PermissionTableTargets::resolved(targets), + ) + .context(AuthSnafu)?; + + let preflight = self + .inserter + .validate_and_normalize_otlp_metrics( + &mut prepared.requests, + &ctx, + metric_ctx.with_metric_engine, + ctx.extension(PHYSICAL_TABLE_PARAM) + .unwrap_or(GREPTIME_PHYSICAL_TABLE), + ) + .await + .map_err(BoxedError::new) + .context(error::ExecuteGrpcQuerySnafu)?; + let rejections = prepared.point_rejections(preflight.rejections)?; let otlp::metrics::MetricsConversion { requests, rows, semantic_index, resource_info, mut outcome, - } = otlp::metrics::to_grpc_insert_requests(request, &mut metric_ctx)?; + } = prepared.finish(rejections, &preflight.physical_tables)?; if outcome.rejected_data_points > 0 { warn!( "Rejected {} OTLP metrics data points: {}", @@ -144,13 +167,10 @@ impl OpenTelemetryProtocolHandler for Instance { return Ok(outcome); } - self.check_row_insert_permission(&requests, &ctx, PermissionReq::Action(OTLP_WRITE)) - .context(AuthSnafu)?; let ctx = admit_write(rows as u64, &ctx) .await .map_err(BoxedError::new) .context(error::ExecuteGrpcQuerySnafu)?; - self.cache_otlp_legacy(&input_names, &ctx, is_legacy)?; OTLP_METRICS_ROWS.inc_by(rows as u64); let ctx = { @@ -163,9 +183,7 @@ impl OpenTelemetryProtocolHandler for Instance { if let Some(index) = semantic_index.encode(&c.current_schema()) { c.set_extension(SEMANTIC_PER_TABLE_INDEX_KEY, index); } - if !is_legacy { - c.set_extension(OTLP_METRIC_COMPAT_KEY, OTLP_METRIC_COMPAT_PROM.to_string()); - } + c.set_extension(OTLP_METRIC_COMPAT_KEY, OTLP_METRIC_COMPAT_PROM.to_string()); Arc::new(c) }; @@ -333,9 +351,9 @@ impl Instance { // The bulk path converts each request's time index unit to the // destination table's unit during batch alignment, so no pre-gate // alignment is needed here. - let batcher = self.logical_batcher().filter(|_| { - ctx.logical_batching_enabled() && !metric_ctx.is_legacy && metric_ctx.with_metric_engine - }); + let batcher = self + .logical_batcher() + .filter(|_| ctx.logical_batching_enabled() && metric_ctx.with_metric_engine); let batcher = if batcher.is_some() && self .inserter @@ -349,7 +367,7 @@ impl Instance { None }; - // OTLP tables have one sample field in both the legacy and physical paths. + // OTLP tables have one sample field in both the ordinary and physical paths. if let Some(batcher) = batcher { let (rows, cost) = batcher .submit_with(requests, ctx.clone(), |mut requests| { @@ -366,7 +384,7 @@ impl Instance { OutputData::AffectedRows(rows as usize), OutputMeta::new_with_cost(cost as _), )) - } else if metric_ctx.is_legacy || !metric_ctx.with_metric_engine { + } else if !metric_ctx.with_metric_engine { self.handle_row_inserts(requests, ctx.clone(), false, true) .await .map_err(BoxedError::new) diff --git a/src/metric-engine/src/engine.rs b/src/metric-engine/src/engine.rs index d01ca4650e6..af15ea087c9 100644 --- a/src/metric-engine/src/engine.rs +++ b/src/metric-engine/src/engine.rs @@ -579,6 +579,26 @@ impl MetricEngine { } } +/// Shared definitions must agree before CREATE/ALTER mutates either region. +fn validate_shared_column( + region_id: store_api::storage::RegionId, + existing: &store_api::metadata::ColumnMetadata, + proposed: &store_api::metadata::ColumnMetadata, +) -> crate::error::Result<()> { + snafu::ensure!( + existing.column_schema.data_type == proposed.column_schema.data_type + && existing.semantic_type == proposed.semantic_type, + crate::error::InvalidRequestSnafu { + region_id, + reason: format!( + "incompatible shared physical column `{}`", + proposed.column_schema.name + ), + } + ); + Ok(()) +} + #[cfg(test)] impl MetricEngine { pub async fn scan_to_stream( diff --git a/src/metric-engine/src/engine/alter.rs b/src/metric-engine/src/engine/alter.rs index 355e84cb322..2bf59edda0f 100644 --- a/src/metric-engine/src/engine/alter.rs +++ b/src/metric-engine/src/engine/alter.rs @@ -15,7 +15,7 @@ mod extract_new_columns; mod validate; -use std::collections::{BTreeSet, HashMap, HashSet}; +use std::collections::{BTreeSet, HashMap}; use api::v1::SemanticType; use common_query::native_histogram::is_native_histogram_value_type; @@ -123,7 +123,7 @@ impl MetricEngineInner { .await?; // Finds new columns to add - let mut new_column_names = HashSet::new(); + let mut new_column_indices = HashMap::new(); let mut new_columns_to_add = vec![]; let index_options = { @@ -139,7 +139,7 @@ impl MetricEngineInner { extract_new_columns( &requests, physical_columns, - &mut new_column_names, + &mut new_column_indices, &mut new_columns_to_add, )?; @@ -190,7 +190,7 @@ impl MetricEngineInner { ) }); - let new_add_columns = new_column_names.iter().map(|name| { + let new_add_columns = new_column_indices.keys().map(|name| { // Safety: previous steps ensure the physical region exist let column_metadata = *physical_schema_map.get(name).unwrap(); (name.to_string(), column_metadata.clone()) diff --git a/src/metric-engine/src/engine/alter/extract_new_columns.rs b/src/metric-engine/src/engine/alter/extract_new_columns.rs index 36c7dbf416a..c8dc640045f 100644 --- a/src/metric-engine/src/engine/alter/extract_new_columns.rs +++ b/src/metric-engine/src/engine/alter/extract_new_columns.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::{HashMap, HashSet}; +use std::collections::HashMap; use api::v1::SemanticType; use common_query::native_histogram::is_native_histogram_value_type; @@ -31,18 +31,22 @@ use crate::error::{AddingFieldColumnSnafu, Result}; pub fn extract_new_columns<'a>( requests: &'a [(RegionId, RegionAlterRequest)], physical_columns: &HashMap, - new_column_names: &mut HashSet<&'a str>, + new_column_indices: &mut HashMap<&'a str, usize>, new_columns: &mut Vec, ) -> Result<()> { - for (_, request) in requests { + for (region_id, request) in requests { let AlterKind::AddColumns { columns } = &request.kind else { unreachable!() }; for col in columns { let column_name = col.column_metadata.column_schema.name.as_str(); - if !physical_columns.contains_key(column_name) - && !new_column_names.contains(column_name) - { + if let Some(existing) = physical_columns.get(column_name).or_else(|| { + new_column_indices + .get(column_name) + .map(|&index| &new_columns[index]) + }) { + crate::engine::validate_shared_column(*region_id, existing, &col.column_metadata)?; + } else { ensure!( col.column_metadata.semantic_type != SemanticType::Field || is_native_histogram_value_type( @@ -52,7 +56,7 @@ pub fn extract_new_columns<'a>( name: column_name.to_string(), } ); - new_column_names.insert(column_name); + new_column_indices.insert(column_name, new_columns.len()); // TODO(weny): avoid clone new_columns.push(col.column_metadata.clone()); } @@ -64,7 +68,7 @@ pub fn extract_new_columns<'a>( #[cfg(test)] mod tests { - use std::collections::{HashMap, HashSet}; + use std::collections::HashMap; use api::v1::SemanticType; use datatypes::prelude::ConcreteDataType; @@ -76,6 +80,82 @@ mod tests { use super::*; use crate::error::Error; + #[test] + fn test_alter_rejects_conflicting_shared_definitions() { + let column = ColumnMetadata { + column_schema: ColumnSchema::new("shared", ConcreteDataType::string_datatype(), true), + semantic_type: SemanticType::Tag, + column_id: 0, + }; + let mut incompatible = column.clone(); + incompatible.column_schema = + ColumnSchema::new("shared", ConcreteDataType::int64_datatype(), true); + for existing in [false, true] { + let physical = if existing { + HashMap::from([("shared".into(), column.clone())]) + } else { + HashMap::new() + }; + let requests = [( + RegionId::new(1, 1), + RegionAlterRequest { + kind: AlterKind::AddColumns { + columns: [column.clone(), incompatible.clone()] + .into_iter() + .map(|column_metadata| AddColumn { + column_metadata, + location: None, + }) + .collect(), + }, + }, + )]; + let error = + extract_new_columns(&requests, &physical, &mut HashMap::new(), &mut Vec::new()) + .unwrap_err(); + assert!(matches!(error, Error::InvalidRequest { .. })); + } + } + + #[test] + fn test_extract_new_columns_preserves_first_seen_order() { + let requests = [["b", "a"], ["a", "b"]].map(|names| { + ( + RegionId::new(1, 1), + RegionAlterRequest { + kind: AlterKind::AddColumns { + columns: names + .into_iter() + .map(|name| AddColumn { + column_metadata: ColumnMetadata { + column_schema: ColumnSchema::new( + name, + ConcreteDataType::string_datatype(), + true, + ), + semantic_type: SemanticType::Tag, + column_id: 0, + }, + location: None, + }) + .collect(), + }, + }, + ) + }); + let mut indices = HashMap::new(); + let mut columns = Vec::new(); + extract_new_columns(&requests, &HashMap::new(), &mut indices, &mut columns).unwrap(); + assert_eq!(indices, HashMap::from([("b", 0), ("a", 1)])); + assert_eq!( + columns + .iter() + .map(|c| c.column_schema.name.as_str()) + .collect::>(), + ["b", "a"] + ); + } + #[test] fn test_extract_new_columns_with_field_type() { let requests = vec![( @@ -99,13 +179,13 @@ mod tests { )]; let physical_columns = HashMap::new(); - let mut new_column_names = HashSet::new(); + let mut new_column_indices = HashMap::new(); let mut new_columns = Vec::new(); let err = extract_new_columns( &requests, &physical_columns, - &mut new_column_names, + &mut new_column_indices, &mut new_columns, ) .unwrap_err(); diff --git a/src/metric-engine/src/engine/create.rs b/src/metric-engine/src/engine/create.rs index ec35c0653c4..2f27a405b36 100644 --- a/src/metric-engine/src/engine/create.rs +++ b/src/metric-engine/src/engine/create.rs @@ -14,7 +14,7 @@ mod extract_new_columns; -use std::collections::{HashMap, HashSet}; +use std::collections::HashMap; use api::v1::SemanticType; use common_query::native_histogram::is_native_histogram_value_type; @@ -266,7 +266,7 @@ impl MetricEngineInner { }; // Finds new columns to add to physical region - let mut new_column_names = HashSet::new(); + let mut new_column_indices = HashMap::new(); let mut new_columns = Vec::new(); let index_option = { @@ -282,7 +282,7 @@ impl MetricEngineInner { extract_new_columns( &requests, physical_columns, - &mut new_column_names, + &mut new_column_indices, &mut new_columns, )?; region_state.options().index @@ -319,7 +319,7 @@ impl MetricEngineInner { ) }); - let new_add_columns = new_column_names.iter().map(|name| { + let new_add_columns = new_column_indices.keys().map(|name| { // Safety: previous steps ensure the physical region exist let column_metadata = *physical_schema_map.get(name).unwrap(); (name.to_string(), column_metadata.clone()) diff --git a/src/metric-engine/src/engine/create/extract_new_columns.rs b/src/metric-engine/src/engine/create/extract_new_columns.rs index add29916ca0..d1bac6fc80a 100644 --- a/src/metric-engine/src/engine/create/extract_new_columns.rs +++ b/src/metric-engine/src/engine/create/extract_new_columns.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::{HashMap, HashSet}; +use std::collections::HashMap; use api::v1::SemanticType; use common_query::native_histogram::is_native_histogram_value_type; @@ -27,14 +27,18 @@ use crate::error::{AddingFieldColumnSnafu, Result}; pub fn extract_new_columns<'a>( requests: &'a [(RegionId, RegionCreateRequest)], physical_columns: &HashMap, - new_column_names: &mut HashSet<&'a str>, + new_column_indices: &mut HashMap<&'a str, usize>, new_columns: &mut Vec, ) -> Result<()> { - for (_, request) in requests { + for (region_id, request) in requests { for col in &request.column_metadatas { - if !physical_columns.contains_key(&col.column_schema.name) - && !new_column_names.contains(&col.column_schema.name.as_str()) - { + if let Some(existing) = physical_columns.get(&col.column_schema.name).or_else(|| { + new_column_indices + .get(col.column_schema.name.as_str()) + .map(|&index| &new_columns[index]) + }) { + crate::engine::validate_shared_column(*region_id, existing, col)?; + } else { ensure!( col.semantic_type != SemanticType::Field || is_native_histogram_value_type(&col.column_schema.data_type), @@ -42,7 +46,7 @@ pub fn extract_new_columns<'a>( name: col.column_schema.name.clone(), } ); - new_column_names.insert(&col.column_schema.name); + new_column_indices.insert(&col.column_schema.name, new_columns.len()); // TODO(weny): avoid clone new_columns.push(col.clone()); } @@ -55,7 +59,7 @@ pub fn extract_new_columns<'a>( #[cfg(test)] mod tests { use std::assert_matches; - use std::collections::{HashMap, HashSet}; + use std::collections::HashMap; use api::v1::SemanticType; use datatypes::prelude::ConcreteDataType; @@ -67,9 +71,51 @@ mod tests { use super::*; use crate::error::Error; + #[test] + fn test_conflicting_shared_definitions_rejected() { + let column = ColumnMetadata { + column_schema: ColumnSchema::new("shared", ConcreteDataType::string_datatype(), true), + semantic_type: SemanticType::Tag, + column_id: 0, + }; + let mut incompatible = column.clone(); + incompatible.semantic_type = SemanticType::Field; + let region = RegionId::new(1, 1); + assert!(crate::engine::validate_shared_column(region, &column, &incompatible).is_err()); + incompatible.semantic_type = SemanticType::Tag; + incompatible.column_schema = + ColumnSchema::new("shared", ConcreteDataType::int64_datatype(), true); + assert!(crate::engine::validate_shared_column(region, &column, &incompatible).is_err()); + assert!(crate::engine::validate_shared_column(region, &column, &column).is_ok()); + for existing in [false, true] { + let physical = if existing { + HashMap::from([("shared".into(), column.clone())]) + } else { + HashMap::new() + }; + let requests = [( + region, + RegionCreateRequest { + column_metadatas: vec![column.clone(), incompatible.clone()], + engine: "metric".into(), + primary_key: vec![0], + options: HashMap::new(), + table_dir: "test".into(), + path_type: PathType::Bare, + partition_expr_json: None, + requirements: Default::default(), + }, + )]; + let error = + extract_new_columns(&requests, &physical, &mut HashMap::new(), &mut Vec::new()) + .unwrap_err(); + assert!(matches!(error, Error::InvalidRequest { .. })); + } + } + #[test] fn test_extract_new_columns() { - let requests = vec![ + let mut requests = vec![ ( RegionId::new(1, 1), RegionCreateRequest { @@ -126,6 +172,11 @@ mod tests { ), ]; + let mut another = requests[0].1.column_metadatas[1].clone(); + another.column_schema.name = "another_column".into(); + requests[0].1.column_metadatas.push(another.clone()); + requests[1].1.column_metadatas.push(another); + let mut physical_columns = HashMap::new(); physical_columns.insert( "existing_column".to_string(), @@ -139,20 +190,22 @@ mod tests { column_id: 0, }, ); - let mut new_column_names = HashSet::new(); + let mut new_column_indices = HashMap::new(); let mut new_columns = Vec::new(); let result = extract_new_columns( &requests, &physical_columns, - &mut new_column_names, + &mut new_column_indices, &mut new_columns, ); assert!(result.is_ok()); - assert!(new_column_names.contains("new_column")); - assert_eq!(new_columns.len(), 1); + assert_eq!(new_column_indices["new_column"], 0); + assert_eq!(new_column_indices["another_column"], 1); + assert_eq!(new_columns.len(), 2); assert_eq!(new_columns[0].column_schema.name, "new_column"); + assert_eq!(new_columns[1].column_schema.name, "another_column"); } #[test] @@ -180,13 +233,13 @@ mod tests { )]; let physical_columns = HashMap::new(); - let mut new_column_names = HashSet::new(); + let mut new_column_indices = HashMap::new(); let mut new_columns = Vec::new(); let err = extract_new_columns( &requests, &physical_columns, - &mut new_column_names, + &mut new_column_indices, &mut new_columns, ) .unwrap_err(); diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index fbf65986c56..9a0cd64a828 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -14,6 +14,9 @@ #[cfg(feature = "enterprise")] mod flow_write; +mod otlp; + +pub use otlp::MetricInsertPreflight; use std::collections::BTreeMap; use std::sync::atomic::{AtomicU64, Ordering}; @@ -583,6 +586,9 @@ impl Inserter { } /// Handles row inserts request with metric engine. + /// + /// Resolves or creates `physical_table` only when a logical target is missing. + /// Existing targets retain their backing tables and time-index units. pub async fn handle_metric_row_inserts( &self, mut requests: RowInsertRequests, @@ -601,22 +607,30 @@ impl Inserter { }); validate_column_count_match(&requests)?; - // check and create physical table - let physical_table_ref = self - .create_physical_table_on_demand(&ctx, physical_table.clone(), statement_executor) - .await?; - - // check and create logical tables; `create_or_alter_tables_on_demand` - // aligns each request's time index unit with the unit of the table it - // targets, inside its existing table lookups: existing tables keep - // their own unit (which matches the physical table they are bound - // to), and new tables use the selected physical table's unit (from - // `physical_table_ref`). Ingestion endpoints encode timestamps in a - // fixed unit (prometheus remote write always uses millisecond; OTLP - // keeps nanosecond precision on the metric engine path), and the - // metric engine requires each logical table's requests to match its - // time index unit. Narrowing conversions truncate the sub-unit part - // (floor), following `Timestamp::convert_to`. + // The selected backing table is needed only by new logical targets. + let mut has_new_tables = false; + for request in &requests.inserts { + if self + .get_table( + ctx.current_catalog(), + &ctx.current_schema(), + &request.table_name, + ) + .await? + .is_none() + { + has_new_tables = true; + break; + } + } + let new_table_unit = if has_new_tables { + let physical = self + .create_physical_table_on_demand(&ctx, physical_table.clone(), statement_executor) + .await?; + table_time_index_unit(&physical) + } else { + None + }; let CreateAlterTableResult { instant_table_ids, table_infos, @@ -628,7 +642,7 @@ impl Inserter { statement_executor, true, true, - table_time_index_unit(&physical_table_ref), + new_table_unit, ) .await?; let name_to_info = table_infos @@ -1065,11 +1079,10 @@ impl Inserter { /// remote write. This will modify the `RowInsertRequests` in place. /// `is_single_value` indicates whether the default schema only contains single value column so we can accommodate it. /// - /// `align_time_index_unit` is the selected physical metric table's time - /// index unit; passing `Some` (metric engine path only) rewrites each - /// request's time index column, inside this function's existing table - /// lookups (no extra catalog access): existing destination tables are - /// converted to their own unit, and new tables to the given unit. + /// `align_time_index_unit` supplies the selected physical table's time-index + /// unit for new logical tables. Existing logical targets always use their + /// own unit, even when this argument is `None` because no new target needs + /// the selected backing table. #[allow(clippy::too_many_arguments)] async fn create_or_alter_tables_on_demand( &self, @@ -1123,7 +1136,7 @@ impl Inserter { // Metric path: an existing destination table keeps its own // time index unit (it may be bound to another physical // table than the one selected by this request). - if align_time_index_unit.is_some() + if matches!(auto_create_table_type, AutoCreateTableType::Logical(_)) && let Some(rows) = req.rows.as_mut() && let Some(target_unit) = table_time_index_unit(&table) { @@ -1167,7 +1180,7 @@ impl Inserter { // Metric path: an existing destination table keeps its // own time index unit (it may be bound to another // physical table than the one selected by this request). - if align_time_index_unit.is_some() + if matches!(auto_create_table_type, AutoCreateTableType::Logical(_)) && let Some(rows) = req.rows.as_mut() && let Some(target_unit) = table_time_index_unit(&table) { @@ -1582,7 +1595,8 @@ impl Inserter { /// When `accommodate_existing_schema` is true, it may modify the input `req` to /// accommodate it with existing schema. See [`create_or_alter_tables_on_demand`](Self::create_or_alter_tables_on_demand) /// for more details. - /// When `is_single_value` is true, it also rejects native-histogram/float kind changes. + /// When `is_single_value` is true, it also rejects native-histogram/float kind + /// changes even when no columns need adding. /// When both options are true, it considers fields when modifying the input `req`. fn get_alter_table_expr_on_demand( &self, @@ -1612,12 +1626,6 @@ impl Inserter { .iter() .filter(|col| col.semantic_type == SemanticType::Field as i32) .count(); - let column_exprs = ColumnExpr::from_column_schemas(request_schema); - let add_columns = expr_helper::extract_add_columns_expr(&table.schema(), column_exprs)?; - let Some(mut add_columns) = add_columns else { - return Ok(None); - }; - if is_single_value { let request_is_native_histogram = request_is_native_histogram(request_schema); let table_is_native_histogram = table_is_native_histogram(table); @@ -1631,6 +1639,12 @@ impl Inserter { ); } + let column_exprs = ColumnExpr::from_column_schemas(request_schema); + let add_columns = expr_helper::extract_add_columns_expr(&table.schema(), column_exprs)?; + let Some(mut add_columns) = add_columns else { + return Ok(None); + }; + // If accommodate_existing_schema is true, update request schema for Timestamp/Field columns if accommodate_existing_schema { let table_schema = table.schema(); diff --git a/src/operator/src/insert/otlp.rs b/src/operator/src/insert/otlp.rs new file mode 100644 index 00000000000..80ead7070c5 --- /dev/null +++ b/src/operator/src/insert/otlp.rs @@ -0,0 +1,828 @@ +// 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. + +//! Side-effect-free destination validation for prepared OTLP batches. Catalog +//! failures propagate; row errors let the converter reject complete source-point fan-outs. + +use std::collections::HashSet; + +use ahash::{HashMap as AHashMap, HashMapExt}; +use api::helper::{ColumnDataTypeWrapper, to_grpc_value}; +use api::v1::alter_table_expr::Kind; +use api::v1::column_def::try_as_column_schema; +use api::v1::{ + ColumnDataType, ColumnDef, ColumnSchema, Row, RowInsertRequest, Rows, SemanticType, Value, +}; +use common_query::native_histogram::native_histogram_value_type; +use common_query::prelude::{greptime_timestamp, greptime_value}; +use common_time::timestamp::TimeUnit; +use datatypes::schema::ColumnSchema as DataColumn; +use session::context::QueryContextRef; +use snafu::prelude::*; +use store_api::metric_engine_consts::{LOGICAL_TABLE_METADATA_KEY, PHYSICAL_TABLE_METADATA_KEY}; +use table::TableRef; + +use crate::error::{self, InvalidInsertRequestSnafu, Result}; +use crate::insert::{ + AutoCreateTableType, Inserter, convert_rows_time_unit, convert_timestamp_value_data, + request_is_native_histogram, table_is_native_histogram, table_time_index_unit, +}; +use crate::req_convert::insert::{ImpureDefaultFiller, value_matches_type}; + +pub struct MetricInsertPreflight { + /// Rejections in batch order, then row order. + pub rejections: Vec>>, + pub physical_tables: std::collections::HashMap, +} + +struct PhysicalColumns { + by_name: AHashMap, + fields: Vec, +} + +impl PhysicalColumns { + fn new(columns: Vec<(DataColumn, SemanticType)>) -> Self { + let fields = columns + .iter() + .filter(|(_, semantic)| *semantic == SemanticType::Field) + .map(|(column, _)| column.name.clone()) + .collect(); + let by_name = columns + .into_iter() + .map(|(column, semantic)| (column.name.clone(), (column, semantic))) + .collect(); + Self { by_name, fields } + } +} + +impl Inserter { + /// Validate each row independently, sharing schema planning within each batch. + /// Fresh cross-point conflicts are resolved by the converter after this pass. + pub async fn validate_and_normalize_otlp_metrics( + &self, + batches: &mut [RowInsertRequest], + ctx: &QueryContextRef, + metric_engine: bool, + physical_table: &str, + ) -> Result { + let disabled = self.auto_create_disabled_reason(ctx)?; + let mut tables = AHashMap::new(); + let mut physical_tables = std::collections::HashMap::new(); + for request in batches.iter() { + if !tables.contains_key(&request.table_name) { + let table = self + .get_table( + ctx.current_catalog(), + &ctx.current_schema(), + &request.table_name, + ) + .await?; + tables.insert(request.table_name.clone(), table); + } + let backing = match tables[&request.table_name].as_ref() { + Some(table) => table + .table_info() + .meta + .options + .extra_options + .get(LOGICAL_TABLE_METADATA_KEY) + .cloned(), + None if metric_engine => Some(physical_table.to_string()), + None => None, + }; + if let Some(backing) = backing { + physical_tables.insert(request.table_name.clone(), backing); + } + } + // Shared schemas accumulate tags across logical tables; index them once per request. + let mut physical_columns = AHashMap::new(); + for backing in physical_tables.values() { + if !tables.contains_key(backing) { + let table = self + .get_table(ctx.current_catalog(), &ctx.current_schema(), backing) + .await?; + tables.insert(backing.clone(), table); + } + physical_columns.entry(backing.as_str()).or_insert_with(|| { + PhysicalColumns::new( + tables[backing] + .as_ref() + .map(table_columns) + .unwrap_or_else(default_physical_columns), + ) + }); + } + // Match ordinary insertion: impure defaults are evaluated once per table + // and request, while supplied nulls remain supplied nulls. + let mut fillers = AHashMap::new(); + let mut build_plan = |source: &mut RowInsertRequest| { + let rows = source.rows.as_mut().context(InvalidInsertRequestSnafu { + reason: "missing metric rows", + })?; + let input_columns = rows.schema.len(); + let source_timestamp = timestamp_column(&rows.schema); + let request = &mut RowInsertRequest { + table_name: source.table_name.clone(), + rows: Some(Rows { + schema: std::mem::take(&mut rows.schema), + rows: Vec::new(), + }), + }; + let auto_create = disabled.is_none() + || Self::is_auto_create_exempt_private_table( + &ctx.current_schema(), + &request.table_name, + ); + let target = tables[&request.table_name].as_ref(); + let backing = physical_tables.get(&request.table_name); + let physical = backing.and_then(|name| tables[name].as_ref()); + if let Some(physical) = physical { + ensure!( + physical + .table_info() + .meta + .options + .extra_options + .contains_key(PHYSICAL_TABLE_METADATA_KEY), + InvalidInsertRequestSnafu { + reason: "selected backing table is not a physical metric table" + } + ); + } + if metric_engine { + ensure!( + backing.is_none() || physical.is_some() || auto_create, + InvalidInsertRequestSnafu { + reason: format!( + "missing physical table and {}", + disabled.unwrap_or_default() + ), + } + ); + let unit = target + .and_then(table_time_index_unit) + .or_else(|| physical.and_then(table_time_index_unit)) + .unwrap_or(TimeUnit::Millisecond); + if let Some(rows) = &mut request.rows { + convert_rows_time_unit(rows, unit)?; + } + } + let mut columns = if let Some(table) = target { + ensure!( + !table + .table_info() + .meta + .options + .extra_options + .contains_key(PHYSICAL_TABLE_METADATA_KEY), + InvalidInsertRequestSnafu { + reason: "direct writes to a physical metric table are forbidden", + } + ); + ensure!( + request_is_native_histogram(&request.rows.as_ref().unwrap().schema) + == table_is_native_histogram(table), + InvalidInsertRequestSnafu { + reason: format!( + "Table `{}` cannot mix native histogram and float sample fields", + request.table_name + ), + } + ); + let mut columns = table_columns(table); + if auto_create { + // Ordinary OTLP insertion does not rename columns. + if let Some(alter) = self.get_alter_table_expr_on_demand( + request, + table, + ctx, + metric_engine, + true, + true, + )? && let Some(Kind::AddColumns(add)) = alter.kind + { + for column in add.add_columns { + let def = + column + .column_def + .with_context(|| InvalidInsertRequestSnafu { + reason: "missing column definition", + })?; + columns.push(column_definition(&def)?); + } + } + } + columns + } else { + ensure!( + auto_create, + InvalidInsertRequestSnafu { + reason: format!( + "missing table `{}` and {}", + request.table_name, + disabled.unwrap_or_default() + ), + } + ); + let create_type = if metric_engine { + AutoCreateTableType::Logical(physical_table.to_string()) + } else { + AutoCreateTableType::Physical + }; + let create = + self.get_create_table_expr_on_demand(request, &create_type, ctx, None)?; + if metric_engine { + ensure!( + create + .column_defs + .iter() + .filter(|c| c.semantic_type == SemanticType::Field as i32) + .count() + == 1, + InvalidInsertRequestSnafu { + reason: "logical metrics require exactly one sample field" + } + ); + } + create + .column_defs + .iter() + .map(column_definition) + .collect::>>()? + }; + let rows = request.rows.as_mut().context(InvalidInsertRequestSnafu { + reason: "missing metric rows", + })?; + if let Some(backing) = backing { + // Logical tag nullability follows the shared physical schema. + for (column, semantic) in &mut columns { + if *semantic == SemanticType::Tag { + column.set_nullable(); + } + } + validate_physical_columns(rows, &physical_columns[backing.as_str()], auto_create)?; + } + // Only defaults need values in the planning template. + rows.rows.push(Row::default()); + if let Some(table) = target { + if !fillers.contains_key(&request.table_name) { + fillers.insert( + request.table_name.clone(), + ImpureDefaultFiller::new(table.table_info())?, + ); + } + fillers[&request.table_name].fill_rows(rows); + } + RowValidationPlan::new( + input_columns, + source_timestamp, + std::mem::take(rows), + &columns, + ) + }; + let mut rejections = Vec::with_capacity(batches.len()); + for request in batches { + let num_rows = request + .rows + .as_ref() + .context(InvalidInsertRequestSnafu { + reason: "missing metric rows", + })? + .rows + .len(); + let plan = match build_plan(request) { + Ok(plan) => plan, + Err(err) if attributable(&err) => { + rejections.push(vec![Some(err.to_string()); num_rows]); + continue; + } + Err(err) => return Err(err), + }; + let rows = request.rows.as_mut().context(InvalidInsertRequestSnafu { + reason: "missing metric rows", + })?; + let mut errors = Vec::with_capacity(num_rows); + for row in &mut rows.rows { + match plan.apply_row(row) { + Ok(()) => errors.push(None), + Err(err) if attributable(&err) => errors.push(Some(err.to_string())), + Err(err) => return Err(err), + } + } + rows.schema = plan.schema; + rejections.push(errors); + } + Ok(MetricInsertPreflight { + rejections, + physical_tables, + }) + } +} + +fn attributable(err: &error::Error) -> bool { + matches!( + err, + error::Error::InvalidInsertRequest { .. } + | error::Error::BuildCreateExprOnInsertion { .. } + | error::Error::FindNewColumnsOnInsertion { .. } + | error::Error::InvalidColumnDef { .. } + ) +} + +fn invalid(reason: impl Into) -> error::Error { + InvalidInsertRequestSnafu { + reason: reason.into(), + } + .build() +} + +fn column_definition(def: &ColumnDef) -> Result<(DataColumn, SemanticType)> { + let column = try_as_column_schema(def).map_err(|err| invalid(err.to_string()))?; + let semantic = + SemanticType::try_from(def.semantic_type).map_err(|err| invalid(err.to_string()))?; + Ok((column, semantic)) +} + +fn table_columns(table: &TableRef) -> Vec<(DataColumn, SemanticType)> { + let info = table.table_info(); + let primary_keys: HashSet<_> = info.meta.primary_key_indices.iter().copied().collect(); + info.meta + .schema + .column_schemas() + .iter() + .enumerate() + .map(|(index, column)| { + let semantic = if column.is_time_index() { + SemanticType::Timestamp + } else if primary_keys.contains(&index) { + SemanticType::Tag + } else { + SemanticType::Field + }; + (column.clone(), semantic) + }) + .collect() +} + +fn default_physical_columns() -> Vec<(DataColumn, SemanticType)> { + use datatypes::prelude::ConcreteDataType; + vec![ + ( + DataColumn::new( + greptime_timestamp(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + SemanticType::Timestamp, + ), + ( + DataColumn::new(greptime_value(), ConcreteDataType::float64_datatype(), true), + SemanticType::Field, + ), + ] +} + +fn check_definition( + source: &ColumnSchema, + target: &DataColumn, + semantic: SemanticType, +) -> Result<()> { + ensure!( + api::helper::is_column_type_value_eq( + source.datatype, + source.datatype_extension.clone(), + &target.data_type + ) && source.semantic_type == semantic as i32, + InvalidInsertRequestSnafu { + reason: format!( + "incompatible datatype or semantic type for column `{}`", + source.column_name + ), + } + ); + Ok(()) +} + +fn validate_physical_columns( + rows: &Rows, + columns: &PhysicalColumns, + auto_create: bool, +) -> Result<()> { + for source in &rows.schema { + if let Some((target, semantic)) = columns.by_name.get(&source.column_name) { + check_definition(source, target, *semantic)?; + } else { + ensure!( + auto_create, + InvalidInsertRequestSnafu { + reason: format!( + "missing physical column `{}` while auto-DDL is disabled", + source.column_name + ) + } + ); + ensure!( + source.semantic_type == SemanticType::Tag as i32 + || (source.semantic_type == SemanticType::Field as i32 + && api::helper::is_column_type_value_eq( + source.datatype, + source.datatype_extension.clone(), + native_histogram_value_type() + )), + InvalidInsertRequestSnafu { + reason: format!("cannot add physical column `{}`", source.column_name) + } + ); + } + } + for name in &columns.fields { + if !rows.schema.iter().any(|c| c.column_name == *name) { + let (column, _) = &columns.by_name[name]; + ensure!( + !column.is_default_impure(), + InvalidInsertRequestSnafu { + reason: format!("impure default for physical field `{}`", column.name) + } + ); + ensure!( + column + .create_default() + .map_err(|e| invalid(e.to_string()))? + .is_some(), + InvalidInsertRequestSnafu { + reason: format!("missing required physical field `{}`", column.name) + } + ); + } + } + Ok(()) +} + +/// Schema decisions are shared; row values are always checked independently. +struct RowValidationPlan { + input_columns: usize, + schema: Vec, + nullable: Vec, + defaults: Vec, + timestamp: Option<(usize, TimeUnit, TimeUnit)>, +} + +impl RowValidationPlan { + fn new( + input_columns: usize, + source_timestamp: Option<(usize, TimeUnit)>, + mut template: Rows, + columns: &[(DataColumn, SemanticType)], + ) -> Result { + let by_name: AHashMap<_, _> = columns + .iter() + .map(|(column, semantic)| (column.name.as_str(), (column, *semantic))) + .collect(); + let mut names = HashSet::new(); + let mut nullable = Vec::with_capacity(columns.len()); + for source in &template.schema { + ensure!( + names.insert(source.column_name.as_str()), + InvalidInsertRequestSnafu { + reason: format!("duplicate column `{}`", source.column_name), + } + ); + let (target, semantic) = by_name + .get(source.column_name.as_str()) + .ok_or_else(|| invalid(format!("unknown column `{}`", source.column_name)))?; + check_definition(source, target, *semantic)?; + nullable.push(target.is_nullable()); + } + let missing = columns + .iter() + .filter(|(column, _)| !names.contains(column.name.as_str())) + .collect::>(); + let mut row = template.rows.pop().context(InvalidInsertRequestSnafu { + reason: "missing normalization template", + })?; + // Resolve omissions before schemas are merged; an explicit null never uses a default. + for (column, semantic) in missing { + let value = column + .create_default() + .map_err(|err| invalid(err.to_string()))? + .ok_or_else(|| invalid(format!("missing required column `{}`", column.name)))?; + let (datatype, datatype_extension) = + ColumnDataTypeWrapper::try_from(column.data_type.clone()) + .map_err(|err| invalid(err.to_string()))? + .into_parts(); + template.schema.push(ColumnSchema { + column_name: column.name.clone(), + datatype: datatype as i32, + semantic_type: *semantic as i32, + datatype_extension, + options: None, + }); + row.values.push(to_grpc_value(value)); + nullable.push(column.is_nullable()); + } + ensure!( + row.values.len() + input_columns == template.schema.len(), + InvalidInsertRequestSnafu { + reason: "invalid normalization template" + } + ); + let defaults = row.values; + for ((source, nullable), value) in template.schema[input_columns..] + .iter() + .zip(&nullable[input_columns..]) + .zip(&defaults) + { + validate_value(value, source, *nullable)?; + } + let timestamp = source_timestamp.and_then(|(index, source)| { + let target = ColumnDataType::try_from(template.schema[index].datatype) + .ok() + .and_then(api::helper::timestamp_unit)?; + (source != target).then_some((index, source, target)) + }); + Ok(Self { + input_columns, + schema: template.schema, + nullable, + defaults, + timestamp, + }) + } + + fn apply_row(&self, row: &mut Row) -> Result<()> { + ensure!( + row.values.len() == self.input_columns, + InvalidInsertRequestSnafu { + reason: "column count mismatch", + } + ); + if let Some((index, source, target)) = self.timestamp + && let Some(value) = row.values[index].value_data.take() + { + row.values[index].value_data = + convert_timestamp_value_data(value, source, target, index)?; + } + for ((source, nullable), value) in self.schema.iter().zip(&self.nullable).zip(&row.values) { + validate_value(value, source, *nullable)?; + } + row.values.extend_from_slice(&self.defaults); + Ok(()) + } +} + +fn timestamp_column(schema: &[ColumnSchema]) -> Option<(usize, TimeUnit)> { + let index = schema + .iter() + .position(|column| column.semantic_type == SemanticType::Timestamp as i32)?; + let unit = ColumnDataType::try_from(schema[index].datatype) + .ok() + .and_then(api::helper::timestamp_unit)?; + Some((index, unit)) +} + +fn validate_value(value: &Value, source: &ColumnSchema, nullable: bool) -> Result<()> { + ensure!( + value_matches_type(value, source.datatype, source.datatype_extension.as_ref()), + InvalidInsertRequestSnafu { + reason: format!("value datatype differs for `{}`", source.column_name) + } + ); + ensure!( + value.value_data.is_some() || nullable, + InvalidInsertRequestSnafu { + reason: format!( + "null supplied for non-nullable column `{}`", + source.column_name + ), + } + ); + Ok(()) +} + +#[cfg(test)] +fn validate_and_fill_rows(rows: &mut Rows, columns: &[(DataColumn, SemanticType)]) -> Result<()> { + let template = Rows { + schema: rows.schema.clone(), + rows: vec![Row::default()], + }; + let plan = RowValidationPlan::new( + rows.schema.len(), + timestamp_column(&rows.schema), + template, + columns, + )?; + for row in &mut rows.rows { + plan.apply_row(row)?; + } + rows.schema = plan.schema; + Ok(()) +} + +#[cfg(test)] +mod tests { + use api::v1::value::ValueData; + use api::v1::{ColumnDataType, Row}; + use datatypes::prelude::ConcreteDataType; + use datatypes::schema::ColumnDefaultConstraint; + use datatypes::value::Value as DataValue; + + use super::*; + + #[test] + fn test_reused_plan_checks_each_value_and_timestamp() { + let schema = vec![api::v1::helper::time_index_column_schema( + greptime_timestamp(), + ColumnDataType::TimestampSecond, + )]; + let mut template = Rows { + schema: schema.clone(), + rows: Vec::new(), + }; + convert_rows_time_unit(&mut template, TimeUnit::Nanosecond).unwrap(); + let columns = vec![( + DataColumn::new( + greptime_timestamp(), + ConcreteDataType::timestamp_nanosecond_datatype(), + false, + ) + .with_time_index(true), + SemanticType::Timestamp, + )]; + template.rows.push(Row::default()); + let plan = + RowValidationPlan::new(schema.len(), timestamp_column(&schema), template, &columns) + .unwrap(); + for (value, valid) in [ + (Some(ValueData::TimestampSecondValue(1)), true), + (Some(ValueData::StringValue("invalid".into())), false), + (None, false), + (Some(ValueData::TimestampSecondValue(i64::MAX)), false), + (Some(ValueData::TimestampSecondValue(2)), true), + ] { + let mut rows = Rows { + schema: schema.clone(), + rows: vec![Row { + values: vec![Value { value_data: value }], + }], + }; + assert_eq!(plan.apply_row(&mut rows.rows[0]).is_ok(), valid); + if valid { + assert_eq!( + plan.schema[0].datatype, + ColumnDataType::TimestampNanosecond as i32 + ); + assert!(matches!( + rows.rows[0].values[0].value_data, + Some(ValueData::TimestampNanosecondValue( + 1_000_000_000 | 2_000_000_000 + )) + )); + } + } + } + + #[test] + fn test_preflight_omissions_defaults_and_explicit_nulls() { + let mut columns = default_physical_columns(); + columns.push(( + DataColumn::new("host", ConcreteDataType::string_datatype(), false) + .with_default_constraint(Some(ColumnDefaultConstraint::Value(DataValue::String( + "unknown".into(), + )))) + .unwrap(), + SemanticType::Tag, + )); + let row = || Rows { + schema: vec![ColumnSchema { + column_name: greptime_timestamp().into(), + datatype: ColumnDataType::TimestampMillisecond as i32, + semantic_type: SemanticType::Timestamp as i32, + ..Default::default() + }], + rows: vec![Row { + values: vec![Value { + value_data: Some(ValueData::TimestampMillisecondValue(1)), + }], + }], + }; + let mut omitted = row(); + validate_and_fill_rows(&mut omitted, &columns).unwrap(); + assert_eq!(omitted.schema[2].column_name, "host"); + assert_eq!( + omitted.rows[0].values[2].value_data, + Some(ValueData::StringValue("unknown".into())) + ); + omitted.rows[0].values[2] = Value::default(); + assert!(validate_and_fill_rows(&mut omitted, &columns).is_err()); + columns[2].0 = DataColumn::new("host", ConcreteDataType::string_datatype(), false); + assert!(validate_and_fill_rows(&mut row(), &columns).is_err()); + columns[2].0.set_nullable(); + let mut nullable = row(); + validate_and_fill_rows(&mut nullable, &columns).unwrap(); + assert_eq!(nullable.rows[0].values[2].value_data, None); + } + + #[test] + fn test_preflight_existing_columns_and_physical_definitions() { + let columns = default_physical_columns(); + let mut rows = Rows { + schema: vec![ColumnSchema { + column_name: greptime_timestamp().into(), + datatype: ColumnDataType::TimestampMillisecond as i32, + semantic_type: SemanticType::Timestamp as i32, + ..Default::default() + }], + rows: vec![Row { + values: vec![Value { + value_data: Some(ValueData::TimestampMillisecondValue(1)), + }], + }], + }; + validate_and_fill_rows(&mut rows, &columns).unwrap(); + for case in 0..3 { + let mut invalid = rows.clone(); + match case { + 0 => invalid.schema[0].semantic_type = SemanticType::Tag as i32, + 1 => invalid.schema[0].datatype = ColumnDataType::TimestampNanosecond as i32, + _ => { + invalid.rows[0].values[0].value_data = + Some(ValueData::StringValue("bad".into())) + } + } + assert!(validate_and_fill_rows(&mut invalid, &columns).is_err()); + } + let columns = PhysicalColumns::new(columns); + rows.schema[1].column_name = "new_field".into(); + assert!(validate_physical_columns(&rows, &columns, true).is_err()); + rows.schema[1].semantic_type = SemanticType::Tag as i32; + rows.schema[1].datatype = ColumnDataType::String as i32; + assert!(validate_physical_columns(&rows, &columns, false).is_err()); + assert!(validate_physical_columns(&rows, &columns, true).is_ok()); + } + + #[test] + fn test_preflight_shared_physical_field_defaults() { + let rows = Rows { + schema: vec![api::v1::helper::time_index_column_schema( + greptime_timestamp(), + ColumnDataType::TimestampMillisecond, + )], + ..Default::default() + }; + let required = DataColumn::new("required", ConcreteDataType::float64_datatype(), false); + let defaulted = required + .clone() + .with_default_constraint(Some(ColumnDefaultConstraint::Value(DataValue::Float64( + 1.0.into(), + )))) + .unwrap(); + let impure = DataColumn::new( + "impure", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_default_constraint(Some(ColumnDefaultConstraint::Function( + "current_timestamp()".into(), + ))) + .unwrap(); + for (field, error) in [ + (required, Some("missing required physical field `required`")), + (defaulted, None), + (impure, Some("impure default for physical field `impure`")), + ] { + let mut columns = default_physical_columns(); + columns.extend((0..1024).map(|index| { + ( + DataColumn::new( + format!("unrelated_tag_{index}"), + ConcreteDataType::string_datatype(), + true, + ), + SemanticType::Tag, + ) + })); + columns.push((field, SemanticType::Field)); + let columns = PhysicalColumns::new(columns); + for auto_create in [false, true] { + let result = validate_physical_columns(&rows, &columns, auto_create); + if let Some(error) = error { + assert!(result.unwrap_err().to_string().contains(error)); + } else { + result.unwrap(); + } + } + } + } +} diff --git a/src/operator/src/req_convert/insert.rs b/src/operator/src/req_convert/insert.rs index 079e98d4f9c..50000e8dc93 100644 --- a/src/operator/src/req_convert/insert.rs +++ b/src/operator/src/req_convert/insert.rs @@ -25,6 +25,7 @@ pub use column_to_row::ColumnToRow; pub(crate) use fill_impure_default::ImpureDefaultFiller; pub use fill_impure_default::fill_reqs_with_impure_default; pub use row_to_batch::rows_to_record_batch; +pub(crate) use row_to_batch::value_matches_type; pub use row_to_region::RowToRegion; use snafu::{OptionExt, ResultExt}; pub use stmt_to_region::StatementToRegion; diff --git a/src/operator/src/req_convert/insert/row_to_batch.rs b/src/operator/src/req_convert/insert/row_to_batch.rs index ccf313737ed..440c66c93a0 100644 --- a/src/operator/src/req_convert/insert/row_to_batch.rs +++ b/src/operator/src/req_convert/insert/row_to_batch.rs @@ -151,7 +151,7 @@ pub fn rows_to_record_batch(rows: &Rows, table_info: &TableInfo) -> Result, diff --git a/src/query/src/optimizer/windowed_sort.rs b/src/query/src/optimizer/windowed_sort.rs index 800bc103d38..62c17c40b7a 100644 --- a/src/query/src/optimizer/windowed_sort.rs +++ b/src/query/src/optimizer/windowed_sort.rs @@ -85,6 +85,10 @@ impl WindowedSortPhysicalRule { let Some(scanner_info) = fetch_partition_range(sort_input.clone())? else { return Ok(Transformed::no(plan)); }; + // Windowed sorting requires partition ranges, which file-engine scans lack. + if scanner_info.partition_ranges.is_empty() { + return Ok(Transformed::no(plan)); + } let input_schema = sort_input.schema(); let first_sort_expr = sort_exec.expr().first(); @@ -350,6 +354,21 @@ mod test { use super::*; + #[test] + fn test_keep_ordinary_sort_when_scan_has_no_partition_ranges() { + let ordering = datafusion_physical_expr::LexOrdering::new([ + datafusion_physical_expr::PhysicalSortExpr { + expr: Arc::new(PhysicalColumn::new("ts", 1)), + options: Default::default(), + }, + ]) + .unwrap(); + let sort = Arc::new(SortExec::new(ordering, new_region_scan())); + let optimized = + WindowedSortPhysicalRule::do_optimize(sort, &ConfigOptions::default()).unwrap(); + assert!(optimized.is::()); + } + #[test] fn test_is_time_index_expr_tracks_aliases_through_projection() { let scan = new_region_scan(); diff --git a/src/servers/Cargo.toml b/src/servers/Cargo.toml index 802ee46f9d4..1f433318938 100644 --- a/src/servers/Cargo.toml +++ b/src/servers/Cargo.toml @@ -207,3 +207,11 @@ harness = false [[bench]] name = "flush_batch_physical" harness = false + +[[bench]] +name = "otlp_metrics" +harness = false + +[[bench]] +name = "otlp_metrics_preflight" +harness = false diff --git a/src/servers/benches/otlp_metrics.rs b/src/servers/benches/otlp_metrics.rs new file mode 100644 index 00000000000..eccf7f4f167 --- /dev/null +++ b/src/servers/benches/otlp_metrics.rs @@ -0,0 +1,84 @@ +// 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. + +#[path = "otlp_metrics/fixtures.rs"] +mod fixtures; + +use std::hint::black_box; +use std::time::Duration; + +use criterion::{BatchSize, BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; +use servers::otlp::metrics::to_grpc_insert_requests; +use session::protocol_ctx::OtlpMetricCtx; + +fn bench_conversion(c: &mut Criterion) { + let fixtures = fixtures::fixtures(); + for (group_name, delta) in [ + ("otlp_conversion/common", false), + ("otlp_conversion/delta", true), + ] { + let mut group = c.benchmark_group(group_name); + for fixture in fixtures + .iter() + .filter(|fixture| fixture.name.starts_with("delta") == delta) + { + group.throughput(Throughput::Elements(fixture.points)); + for resource_info in [false, true] { + let ctx = OtlpMetricCtx { + resource_info, + with_metric_engine: true, + ..Default::default() + }; + group.bench_function( + BenchmarkId::new( + fixture.name, + if resource_info { + "descriptors_on" + } else { + "descriptors_off" + }, + ), + |b| { + let output = + to_grpc_insert_requests(fixture.request.clone(), &mut ctx.clone()) + .unwrap(); + assert_eq!( + output.outcome.accepted_data_points, + fixture.points as i64 - fixture.rejected + ); + assert_eq!(output.outcome.rejected_data_points, fixture.rejected); + assert_eq!(output.resource_info.is_some(), resource_info); + b.iter_batched( + || fixture.request.clone(), + |request| { + black_box( + to_grpc_insert_requests(request, &mut ctx.clone()).unwrap(), + ) + }, + BatchSize::LargeInput, + ); + }, + ); + } + } + group.finish(); + } +} + +criterion_group! { + name = benches; + config = Criterion::default().sample_size(20).warm_up_time(Duration::from_secs(1)).measurement_time(Duration::from_secs(3)); + targets = bench_conversion +} +criterion_main!(benches); diff --git a/src/servers/benches/otlp_metrics/README.md b/src/servers/benches/otlp_metrics/README.md new file mode 100644 index 00000000000..c230915e967 --- /dev/null +++ b/src/servers/benches/otlp_metrics/README.md @@ -0,0 +1,51 @@ +# OTLP metrics benchmarks + +Run from the workspace root: + +```sh +cargo bench -p servers --bench otlp_metrics +cargo bench -p servers --bench otlp_metrics_preflight +``` + +The conversion benchmark measures `to_grpc_insert_requests`, including point +preparation, encoding, survivor merging, and optional resource descriptors. The +preflight benchmark measures destination validation against existing tables in +an in-memory catalog, with a fresh per-request schema cache. The same binary also +measures `otlp_ingestion`, which runs preparation, destination validation, and +survivor finalization together. Request cloning, fixture construction, and catalog +setup are outside the timed section. Neither +benchmark measures network requests, table creation, storage writes, or the full +ingestion endpoint. Both use the server's jemalloc allocator on non-Windows hosts. + +Cases cover one point, 1,024 points with eight tags, 64 tags, classic histogram +fan-out, cumulative and delta exponential histograms, and mixed valid/invalid +points. `mixed_1024` rejects every tenth point during conversion; +`required_tag_1024` converts all points but rejects every tenth during preflight +because an ordinary table requires a missing tag. Other preflight cases use +logical tables with a shared physical schema. Each case checks accepted or +rejected counts before measuring. Conversion runs with descriptors on and off; +preflight does not validate the separate descriptor write. + +To compare with `main`, copy `otlp_metrics.rs` and `otlp_metrics/fixtures.rs` +into an isolated checkout of that revision and register the `otlp_metrics` bench +in `src/servers/Cargo.toml` with `harness = false`. Use the same toolchain, build +profile, and machine for both revisions. Run serially, with other CPU-heavy work +stopped, and store both runs in the same absolute `CRITERION_HOME`: + +```sh +# In the base checkout: +cargo bench -p servers --bench otlp_metrics -- otlp_conversion/common --save-baseline main +# In the feature checkout: +cargo bench -p servers --bench otlp_metrics -- otlp_conversion/common --baseline main +``` + +Delta conversion, preflight, and combined ingestion are measured only on the +feature branch because the base does not support them. A before/after conversion +ratio measures the whole branch's conversion changes, not the overhead of +preflight. These are local microbenchmarks; use an ingestion workload to measure +end-to-end throughput. + +For staged refactors, save a baseline for the unmodified feature head and each +step using `--save-baseline `. Run both binaries serially with the same +fixtures and settings. Compare each step against the preceding step and the +original head; only common conversion cases can also be compared with `main`. diff --git a/src/servers/benches/otlp_metrics/fixtures.rs b/src/servers/benches/otlp_metrics/fixtures.rs new file mode 100644 index 00000000000..4f10548e29e --- /dev/null +++ b/src/servers/benches/otlp_metrics/fixtures.rs @@ -0,0 +1,157 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest; +use otel_arrow_rust::proto::opentelemetry::common::v1::{AnyValue, KeyValue, any_value}; +use otel_arrow_rust::proto::opentelemetry::metrics::v1::{ + AggregationTemporality, ExponentialHistogram, ExponentialHistogramDataPoint, Gauge, Histogram, + HistogramDataPoint, Metric, NumberDataPoint, ResourceMetrics, ScopeMetrics, + exponential_histogram_data_point::Buckets, metric, number_data_point, +}; +use otel_arrow_rust::proto::opentelemetry::resource::v1::Resource; + +// Match the server allocator when measuring allocation-heavy conversion and validation. +#[cfg(not(windows))] +#[global_allocator] +static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc; + +pub struct Fixture { + pub name: &'static str, + pub request: ExportMetricsServiceRequest, + pub points: u64, + pub rejected: i64, +} + +pub fn fixtures() -> Vec { + [ + ("gauge_1", 1, 8), + ("gauge_1024", 1024, 8), + ("wide_1024", 1024, 64), + ("histogram_1024", 1024, 8), + ("cumulative_1024", 1024, 8), + ("mixed_1024", 1024, 8), + ("required_tag_1024", 1024, 8), + ("delta_1024", 1024, 8), + ] + .into_iter() + .map(|(name, points, labels)| { + let attributes = |index: u64| { + (0..labels) + .filter(|&label| { + !(name == "required_tag_1024" && label == 0 && index.is_multiple_of(10)) + }) + .map(|label| { + attribute(&format!("label_{label}"), &format!("value_{}", index % 100)) + }) + .collect() + }; + let time = |index| 60_000_000_000 + index * 1_000_000; + let data = if name.starts_with("histogram") { + metric::Data::Histogram(Histogram { + aggregation_temporality: AggregationTemporality::Cumulative as i32, + data_points: (0..points) + .map(|index| HistogramDataPoint { + attributes: attributes(index), + time_unix_nano: time(index), + count: 11, + sum: Some(55.0), + bucket_counts: vec![1; 11], + explicit_bounds: (1..=10).map(f64::from).collect(), + ..Default::default() + }) + .collect(), + }) + } else if name.starts_with("cumulative") + || name.starts_with("mixed") + || name.starts_with("delta") + { + metric::Data::ExponentialHistogram(ExponentialHistogram { + aggregation_temporality: if name.starts_with("delta") { + AggregationTemporality::Delta + } else { + AggregationTemporality::Cumulative + } as i32, + data_points: (0..points) + .map(|index| ExponentialHistogramDataPoint { + attributes: attributes(index), + time_unix_nano: time(index), + count: 16, + sum: Some(100.0), + scale: if name.starts_with("mixed") && index.is_multiple_of(10) { + -5 + } else { + 0 + }, + zero_threshold: 0.001, + positive: Some(Buckets { + offset: 0, + bucket_counts: vec![1; 16], + }), + ..Default::default() + }) + .collect(), + }) + } else { + metric::Data::Gauge(Gauge { + data_points: (0..points) + .map(|index| NumberDataPoint { + attributes: attributes(index), + time_unix_nano: time(index), + value: Some(number_data_point::Value::AsDouble(index as f64)), + ..Default::default() + }) + .collect(), + }) + }; + Fixture { + name, + points, + rejected: if name.starts_with("mixed") { + points.div_ceil(10) as i64 + } else { + 0 + }, + request: ExportMetricsServiceRequest { + resource_metrics: vec![ResourceMetrics { + resource: Some(Resource { + attributes: vec![ + attribute("service.name", "bench"), + attribute("host.id", "host"), + ], + ..Default::default() + }), + scope_metrics: vec![ScopeMetrics { + metrics: vec![Metric { + name: name.into(), + data: Some(data), + ..Default::default() + }], + ..Default::default() + }], + ..Default::default() + }], + }, + } + }) + .collect() +} + +fn attribute(key: &str, value: &str) -> KeyValue { + KeyValue { + key: key.into(), + value: Some(AnyValue { + value: Some(any_value::Value::StringValue(value.into())), + }), + } +} diff --git a/src/servers/benches/otlp_metrics_preflight.rs b/src/servers/benches/otlp_metrics_preflight.rs new file mode 100644 index 00000000000..c74ccfe272b --- /dev/null +++ b/src/servers/benches/otlp_metrics_preflight.rs @@ -0,0 +1,260 @@ +// 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. + +#[path = "otlp_metrics/fixtures.rs"] +mod fixtures; + +use std::collections::BTreeMap; +use std::hint::black_box; +use std::sync::Arc; +use std::time::Duration; + +use api::helper::ColumnDataTypeWrapper; +use api::v1::{ColumnDataType, ColumnSchema, RowInsertRequests, SemanticType}; +use catalog::RegisterTableRequest; +use catalog::memory::MemoryCatalogManager; +use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; +use common_meta::cache::new_table_flownode_set_cache; +use common_meta::ddl::test_util::datanode_handler::NaiveDatanodeHandler; +use common_meta::kv_backend::memory::MemoryKvBackend; +use common_meta::test_util::MockDatanodeManager; +use common_query::prelude::greptime_value; +use criterion::{BatchSize, Criterion, Throughput, criterion_group, criterion_main}; +use datatypes::schema::{ColumnSchema as DataColumn, Schema}; +use moka::future::Cache; +use operator::insert::Inserter; +use operator::test_util::create_partition_rule_manager; +use servers::otlp::metrics::{prepare_metrics, to_grpc_insert_requests}; +use session::context::QueryContext; +use session::protocol_ctx::OtlpMetricCtx; +use store_api::metric_engine_consts::{LOGICAL_TABLE_METADATA_KEY, PHYSICAL_TABLE_METADATA_KEY}; +use table::dist_table::DummyDataSource; +use table::metadata::FilterPushDownType; +use table::test_util::table_info::test_table_info; + +fn register_table(catalog: &MemoryCatalogManager, id: u32, name: &str, columns: &[ColumnSchema]) { + let schema = Schema::new( + columns + .iter() + .map(|column| { + let timestamp = column.semantic_type == SemanticType::Timestamp as i32; + let required = name == "required_tag_1024" && column.column_name == "label_0"; + DataColumn::new( + &column.column_name, + ColumnDataTypeWrapper::try_new( + column.datatype, + column.datatype_extension.clone(), + ) + .unwrap() + .into(), + !timestamp && !required, + ) + .with_time_index(timestamp) + }) + .collect(), + ); + let mut info = test_table_info( + id, + name, + DEFAULT_SCHEMA_NAME, + DEFAULT_CATALOG_NAME, + Arc::new(schema), + ); + info.meta.primary_key_indices = columns + .iter() + .enumerate() + .filter_map(|(i, c)| (c.semantic_type == SemanticType::Tag as i32).then_some(i)) + .collect(); + info.meta.value_indices = columns + .iter() + .enumerate() + .filter_map(|(i, c)| (c.semantic_type == SemanticType::Field as i32).then_some(i)) + .collect(); + // Ordinary tables retain required tags; logical tables inherit nullable physical tags. + if name != "required_tag_1024" { + info.meta.engine = "metric".into(); + let (key, value) = if name == "physical" { + (PHYSICAL_TABLE_METADATA_KEY, "true") + } else { + (LOGICAL_TABLE_METADATA_KEY, "physical") + }; + info.meta + .options + .extra_options + .insert(key.into(), value.into()); + } + catalog + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.into(), + schema: DEFAULT_SCHEMA_NAME.into(), + table_name: name.into(), + table_id: id, + table: Arc::new(table::Table::new( + Arc::new(info), + FilterPushDownType::Unsupported, + Arc::new(DummyDataSource), + )), + }) + .unwrap(); +} + +async fn inserter(requests: &RowInsertRequests) -> Inserter { + let catalog = MemoryCatalogManager::with_default_setup(); + let mut physical = BTreeMap::new(); + physical.insert( + greptime_value().to_string(), + api::v1::helper::field_column_schema(greptime_value(), ColumnDataType::Float64), + ); + for (index, request) in requests.inserts.iter().enumerate() { + let columns = &request.rows.as_ref().unwrap().schema; + register_table(&catalog, index as u32 + 1, &request.table_name, columns); + physical.extend( + columns + .iter() + .map(|column| (column.column_name.clone(), column.clone())), + ); + } + register_table( + &catalog, + 0, + "physical", + &physical.into_values().collect::>(), + ); + let backend = Arc::new(MemoryKvBackend::default()); + Inserter::new( + catalog, + create_partition_rule_manager(backend.clone()).await, + Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)), + Arc::new(new_table_flownode_set_cache( + String::new(), + Cache::new(100), + backend, + )), + true, + ) +} + +fn bench_preflight(c: &mut Criterion) { + let runtime = tokio::runtime::Runtime::new().unwrap(); + let ctx = QueryContext::arc(); + let mut group = c.benchmark_group("otlp_preflight"); + for fixture in fixtures::fixtures() { + let mut metric_ctx = OtlpMetricCtx { + with_metric_engine: true, + ..Default::default() + }; + let converted = to_grpc_insert_requests(fixture.request.clone(), &mut metric_ctx).unwrap(); + assert_eq!( + converted.outcome.accepted_data_points, + fixture.points as i64 - fixture.rejected + ); + let inserter = runtime.block_on(inserter(&converted.requests)); + let prepared = prepare_metrics(fixture.request, &mut metric_ctx).unwrap(); + group.throughput(Throughput::Elements(prepared.num_points() as u64)); + group.bench_function(fixture.name, |b| { + let mut requests = prepared.requests.clone(); + let result = runtime + .block_on(inserter.validate_and_normalize_otlp_metrics( + &mut requests, + &ctx, + true, + "physical", + )) + .unwrap(); + let rejections = prepared.point_rejections(result.rejections).unwrap(); + assert_eq!(rejections.len(), prepared.num_points()); + assert_eq!( + rejections.iter().filter(|r| r.is_some()).count(), + if fixture.name == "required_tag_1024" { + fixture.points.div_ceil(10) as usize + } else { + 0 + } + ); + b.iter_batched_ref( + || prepared.requests.clone(), + |requests| { + black_box( + runtime + .block_on(inserter.validate_and_normalize_otlp_metrics( + requests, &ctx, true, "physical", + )) + .unwrap(), + ) + }, + BatchSize::LargeInput, + ); + }); + } + group.finish(); +} + +fn bench_ingestion(c: &mut Criterion) { + let runtime = tokio::runtime::Runtime::new().unwrap(); + let ctx = QueryContext::arc(); + let mut group = c.benchmark_group("otlp_ingestion"); + for fixture in fixtures::fixtures() { + let metric_ctx = OtlpMetricCtx { + with_metric_engine: true, + ..Default::default() + }; + let converted = + to_grpc_insert_requests(fixture.request.clone(), &mut metric_ctx.clone()).unwrap(); + let inserter = runtime.block_on(inserter(&converted.requests)); + let convert = |request| { + let mut prepared = prepare_metrics(request, &mut metric_ctx.clone()).unwrap(); + let validated = runtime + .block_on(inserter.validate_and_normalize_otlp_metrics( + &mut prepared.requests, + &ctx, + true, + "physical", + )) + .unwrap(); + let rejections = prepared.point_rejections(validated.rejections).unwrap(); + prepared + .finish(rejections, &validated.physical_tables) + .unwrap() + }; + let output = convert(fixture.request.clone()); + let rejected = if fixture.name == "required_tag_1024" { + fixture.points.div_ceil(10) as i64 + } else { + fixture.rejected + }; + assert_eq!(output.outcome.rejected_data_points, rejected); + assert_eq!( + output.outcome.accepted_data_points, + fixture.points as i64 - rejected + ); + group.throughput(Throughput::Elements(fixture.points)); + group.bench_function(fixture.name, |b| { + b.iter_batched( + || fixture.request.clone(), + |request| { + black_box(convert(request)); + }, + BatchSize::LargeInput, + ); + }); + } + group.finish(); +} + +criterion_group! { + name = benches; + config = Criterion::default().sample_size(20).warm_up_time(Duration::from_secs(1)).measurement_time(Duration::from_secs(3)); + targets = bench_preflight, bench_ingestion +} +criterion_main!(benches); diff --git a/src/servers/src/error.rs b/src/servers/src/error.rs index 2e17179206e..7421ac7bdf3 100644 --- a/src/servers/src/error.rs +++ b/src/servers/src/error.rs @@ -282,14 +282,6 @@ pub enum Error { location: Location, }, - #[snafu(display( - "OTLP metric input have incompatible existing tables, please refer to docs for details" - ))] - OtlpMetricModeIncompatible { - #[snafu(implicit)] - location: Location, - }, - #[snafu(display("Common Meta error"))] CommonMeta { #[snafu(implicit)] @@ -767,8 +759,7 @@ impl ErrorExt for Error { AddressBind { .. } | AlreadyStarted { .. } - | InvalidPromRemoteReadQueryResult { .. } - | OtlpMetricModeIncompatible { .. } => StatusCode::IllegalState, + | InvalidPromRemoteReadQueryResult { .. } => StatusCode::IllegalState, UnsupportedDataType { .. } | UnsupportedOperation { .. } => StatusCode::Unsupported, diff --git a/src/servers/src/http/otlp.rs b/src/servers/src/http/otlp.rs index 717e6ca65da..fc2e6e6aabd 100644 --- a/src/servers/src/http/otlp.rs +++ b/src/servers/src/http/otlp.rs @@ -116,7 +116,6 @@ pub async fn metrics( promote_scope_attrs: http_opts.promote_scope_attrs, with_metric_engine, // set by the frontend from its config - is_legacy: false, resource_info: false, metric_type: MetricType::Init, metric_translation_strategy: http_opts.metric_translation_strategy, diff --git a/src/servers/src/otlp/metrics.rs b/src/servers/src/otlp/metrics.rs index 7a485545bdd..27462512e64 100644 --- a/src/servers/src/otlp/metrics.rs +++ b/src/servers/src/otlp/metrics.rs @@ -18,22 +18,22 @@ use api::greptime_proto::io::prometheus::write::v2::histogram::{ }; use api::greptime_proto::io::prometheus::write::v2::{BucketSpan, Histogram as PromHistogram}; use api::v1::value::ValueData; -use api::v1::{RowInsertRequests, SemanticType, Value}; +use api::v1::{ColumnSchema, RowInsertRequests, SemanticType, Value}; use common_grpc::precision::Precision; use common_query::native_histogram::{ encode_native_histogram, native_histogram_column_schema, native_histogram_value_type, }; use common_query::prelude::{ - GREPTIME_COUNT, GREPTIME_TEMPORALITY_DELTA, OTLP_AGGREGATION_TEMPORALITY_LABEL, - greptime_timestamp, greptime_value, + GREPTIME_TEMPORALITY_DELTA, OTLP_AGGREGATION_TEMPORALITY_LABEL, greptime_timestamp, + greptime_value, }; use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS; -use common_telemetry::warn; use lazy_static::lazy_static; use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest; use otel_arrow_rust::proto::opentelemetry::common::v1::{AnyValue, KeyValue, any_value}; use otel_arrow_rust::proto::opentelemetry::metrics::v1::{metric, number_data_point, *}; use session::protocol_ctx::{MetricType, OtlpMetricCtx}; +use snafu::OptionExt; use table::requests::{ METADATA_QUALITY_DECLARED, METRIC_TEMPORALITY_CUMULATIVE, METRIC_TEMPORALITY_DELTA, SEMANTIC_METRIC_METADATA_QUALITY, SEMANTIC_METRIC_ORIGINAL_NAME, SEMANTIC_METRIC_TEMPORALITY, @@ -51,12 +51,14 @@ use crate::semantic::{ METRIC_TYPE_UPDOWN_COUNTER, }; +mod prepare; mod resource_info; + +pub use prepare::{PreparedMetrics, prepare_metrics}; mod translator; pub use resource_info::OTEL_RESOURCE_INFO_TABLE_NAME; use resource_info::ResourceInfoData; -pub use translator::legacy_normalize_otlp_name; pub(crate) use translator::ucum_to_openmetrics_unit; use translator::{translate_label_name, translate_metric_name}; @@ -120,7 +122,40 @@ pub struct MetricsConversion { pub outcome: MetricsIngestOutcome, } -/// Convert OpenTelemetry metrics to GreptimeDB insert requests +struct MetricMetadata { + name: String, + original_name: String, + native_schema: Option, + semantics: SemanticIndex, +} + +impl MetricMetadata { + fn new(metric: &Metric, ctx: &OtlpMetricCtx) -> Result { + let name = translate_metric_name(metric, &ctx.metric_type, ctx.metric_translation_strategy); + let emitted = emitted_semantic_tables(&ctx.metric_type, &name); + let mut semantics = SemanticIndex::default(); + record_metric_semantics(&mut semantics, metric, &name, &emitted); + let native_schema = matches!(metric.data, Some(metric::Data::ExponentialHistogram(_))) + .then(native_histogram_column_schema) + .transpose() + .map_err(|error| { + error::InternalSnafu { + err_msg: error.to_string(), + } + .build() + })?; + Ok(Self { + name, + original_name: metric.name.clone(), + native_schema, + semantics, + }) + } +} + +/// Converts OpenTelemetry metrics without validating destination table schemas. +/// Ingestion callers use [`prepare_metrics`], run destination preflight, then call +/// [`PreparedMetrics::finish`] to apply rejections and merge surviving points. /// /// See /// @@ -129,132 +164,40 @@ pub fn to_grpc_insert_requests( request: ExportMetricsServiceRequest, metric_ctx: &mut OtlpMetricCtx, ) -> Result { - let mut table_writer = MultiTableData::default(); - let mut semantic_index = SemanticIndex::default(); - let mut outcome = MetricsIngestOutcome::default(); - let mut resource_info = ResourceInfoData::default(); - - for resource in &request.resource_metrics { - if metric_ctx.resource_info - && !metric_ctx.is_legacy - && let Some(r) = resource.resource.as_ref() - { - resource_info.observe(&r.attributes, resource); - } - - let resource_attrs = resource.resource.as_ref().map(|r| { - let mut attrs = r.attributes.clone(); - process_resource_attrs(&mut attrs, metric_ctx); - attrs - }); - - for scope in &resource.scope_metrics { - let scope_attrs = process_scope_attrs(scope, metric_ctx); - - for metric in &scope.metrics { - if metric.data.is_none() { - continue; - } - if let Some(t) = metric.data.as_ref().map(from_metric_type) { - metric_ctx.set_metric_type(t); - } - - encode_metrics( - &mut table_writer, - metric, - resource_attrs.as_ref(), - scope_attrs.as_ref(), - metric_ctx, - &mut semantic_index, - &mut outcome, - )?; - } - } - } - - let (requests, rows) = table_writer.into_row_insert_requests(); - - validate_sample_kinds(&requests)?; - - // The metric and the descriptor would fight over the same table across - // two engines. The user's metric wins. - let resource_info = if !metric_ctx.resource_info { - None - } else if requests - .inserts - .iter() - .any(|r| r.table_name == OTEL_RESOURCE_INFO_TABLE_NAME) - { - warn!( - "Skipping OTLP resource descriptor synthesis: the request writes \ - a metric named `{OTEL_RESOURCE_INFO_TABLE_NAME}`" - ); - None - } else { - resource_info.into_row_insert_requests()? - }; - if resource_info.is_some() { - semantic_index.record_scalar( - OTEL_RESOURCE_INFO_TABLE_NAME, - SEMANTIC_METRIC_TYPE, - crate::semantic::METRIC_TYPE_INFO, - ); - semantic_index.record_scalar( - OTEL_RESOURCE_INFO_TABLE_NAME, - SEMANTIC_METRIC_METADATA_QUALITY, - METADATA_QUALITY_DECLARED, - ); - } - - Ok(MetricsConversion { - requests, - rows, - semantic_index, - outcome, - resource_info, - }) + let prepared = prepare_metrics(request, metric_ctx)?; + let rejections = vec![None; prepared.num_points()]; + prepared.finish(rejections, &std::collections::HashMap::new()) } -fn validate_sample_kinds(requests: &RowInsertRequests) -> Result<()> { - for request in &requests.inserts { - let Some(rows) = &request.rows else { - continue; - }; - let field_count = rows - .schema - .iter() - .filter(|column| column.semantic_type == SemanticType::Field as i32) - .count(); - let has_native_histogram = rows.schema.iter().any(|column| { - column.semantic_type == SemanticType::Field as i32 - && api::helper::is_column_type_value_eq( - column.datatype, - column.datatype_extension.clone(), - native_histogram_value_type(), - ) - }); - if has_native_histogram && field_count != 1 { - return Err(error::InvalidParameterSnafu { - reason: format!( - "OTLP metric `{}` cannot mix native histogram and float sample fields", - request.table_name - ), - } - .build()); +fn validate_sample_kind(table_name: &str, schema: &[ColumnSchema]) -> Result<()> { + let field_count = schema + .iter() + .filter(|column| column.semantic_type == SemanticType::Field as i32) + .count(); + let has_native_histogram = schema.iter().any(|column| { + column.semantic_type == SemanticType::Field as i32 + && api::helper::is_column_type_value_eq( + column.datatype, + column.datatype_extension.clone(), + native_histogram_value_type(), + ) + }); + if has_native_histogram && field_count != 1 { + return Err(error::InvalidParameterSnafu { + reason: format!( + "OTLP metric `{}` cannot mix native histogram and float sample fields", + table_name + ), } + .build()); } Ok(()) } /// The tables a metric emits and their per-table `metric.type`. Histogram fans /// out into `_bucket` (the histogram) plus `_sum`/`_count` counters; summary -/// fans out into the quantile table plus `_count`/`_sum` counters (legacy -/// summary stays a single table). -fn emitted_semantic_tables( - metric_type: &MetricType, - is_legacy: bool, - base: &str, -) -> Vec<(String, &'static str)> { +/// fans out into the quantile table plus `_count`/`_sum` counters. +fn emitted_semantic_tables(metric_type: &MetricType, base: &str) -> Vec<(String, &'static str)> { match metric_type { MetricType::Gauge => vec![(base.to_string(), METRIC_TYPE_GAUGE)], MetricType::MonotonicSum => vec![(base.to_string(), METRIC_TYPE_COUNTER)], @@ -267,7 +210,6 @@ fn emitted_semantic_tables( (format!("{base}{SUM_TABLE_SUFFIX}"), METRIC_TYPE_COUNTER), (format!("{base}{COUNT_TABLE_SUFFIX}"), METRIC_TYPE_COUNTER), ], - MetricType::Summary if is_legacy => vec![(base.to_string(), METRIC_TYPE_SUMMARY)], MetricType::Summary => vec![ (base.to_string(), METRIC_TYPE_SUMMARY), (format!("{base}{COUNT_TABLE_SUFFIX}"), METRIC_TYPE_COUNTER), @@ -302,9 +244,8 @@ fn record_metric_semantics( index: &mut SemanticIndex, metric: &Metric, name: &str, - metric_ctx: &OtlpMetricCtx, + emitted: &[(String, &str)], ) { - let emitted = emitted_semantic_tables(&metric_ctx.metric_type, metric_ctx.is_legacy, name); if emitted.is_empty() { return; } @@ -314,7 +255,7 @@ fn record_metric_semantics( // `original_name` is meaningful only when translation renamed the metric. let original_name = (name != metric.name.as_str()).then_some(metric.name.as_str()); - for (table, metric_type) in &emitted { + for (table, metric_type) in emitted { index.record_scalar(table, SEMANTIC_METRIC_TYPE, metric_type); index.record_scalar( table, @@ -399,10 +340,6 @@ fn string_key_value(key: &str, value: String) -> KeyValue { } fn process_resource_attrs(attrs: &mut Vec, metric_ctx: &OtlpMetricCtx) { - if metric_ctx.is_legacy { - return; - } - // remap the service identity attributes to job and instance let ServiceIdentity { job, instance } = service_identity(attrs); @@ -425,10 +362,6 @@ fn process_resource_attrs(attrs: &mut Vec, metric_ctx: &OtlpMetricCtx) } fn process_scope_attrs(scope: &ScopeMetrics, metric_ctx: &OtlpMetricCtx) -> Option> { - if metric_ctx.is_legacy { - return scope.scope.as_ref().map(|s| s.attributes.clone()); - }; - if !metric_ctx.promote_scope_attrs { return None; } @@ -461,63 +394,52 @@ fn process_scope_attrs(scope: &ScopeMetrics, metric_ctx: &OtlpMetricCtx) -> Opti fn encode_metrics( table_writer: &mut MultiTableData, metric: &Metric, + metadata: &MetricMetadata, resource_attrs: Option<&Vec>, scope_attrs: Option<&Vec>, metric_ctx: &OtlpMetricCtx, - semantic_index: &mut SemanticIndex, outcome: &mut MetricsIngestOutcome, ) -> Result<()> { - let name = if metric_ctx.is_legacy { - legacy_normalize_otlp_name(&metric.name) - } else { - translate_metric_name( - metric, - &metric_ctx.metric_type, - metric_ctx.metric_translation_strategy, - ) - }; + let name = &metadata.name; - let emitted = if let Some(data) = &metric.data { + if let Some(data) = &metric.data { match data { metric::Data::Gauge(gauge) => { encode_gauge( table_writer, - &name, + name, gauge, resource_attrs, scope_attrs, metric_ctx, )?; add_accepted_data_points(outcome, gauge.data_points.len())?; - !gauge.data_points.is_empty() } metric::Data::Sum(sum) => { encode_sum( table_writer, - &name, + name, sum, resource_attrs, scope_attrs, metric_ctx, )?; add_accepted_data_points(outcome, sum.data_points.len())?; - !sum.data_points.is_empty() } metric::Data::Summary(summary) => { encode_summary( table_writer, - &name, + name, summary, resource_attrs, scope_attrs, metric_ctx, )?; add_accepted_data_points(outcome, summary.data_points.len())?; - !summary.data_points.is_empty() } metric::Data::Histogram(hist) => encode_histogram( table_writer, - &name, + name, hist, resource_attrs, scope_attrs, @@ -526,7 +448,7 @@ fn encode_metrics( )?, metric::Data::ExponentialHistogram(hist) => encode_exponential_histogram( table_writer, - &name, + metadata, hist, resource_attrs, scope_attrs, @@ -534,13 +456,6 @@ fn encode_metrics( outcome, )?, } - } else { - false - }; - - if emitted { - // Stamp semantic metadata only after at least one row was accepted. - record_metric_semantics(semantic_index, metric, &name, metric_ctx); } Ok(()) @@ -625,32 +540,35 @@ fn append_rejection_message(message: &mut Option, reason: impl FnOnce() fn encode_exponential_histogram( table_writer: &mut MultiTableData, - name: &str, + metadata: &MetricMetadata, histogram: &ExponentialHistogram, resource_attrs: Option<&Vec>, scope_attrs: Option<&Vec>, metric_ctx: &OtlpMetricCtx, outcome: &mut MetricsIngestOutcome, -) -> Result { - if let Err(rejection) = exponential_histogram_gate(histogram) { - reject_data_points(outcome, histogram.data_points.len(), || { - rejection.message(name) - })?; - OTLP_EXPONENTIAL_HISTOGRAM_REJECTED_DATA_POINTS - .with_label_values(&[rejection.reason_label()]) - .inc_by(histogram.data_points.len() as u64); - return Ok(false); - } - - let column_schema = native_histogram_column_schema().map_err(|error| { - error::InternalSnafu { - err_msg: error.to_string(), +) -> Result<()> { + let name = &metadata.name; + let temporality = match exponential_histogram_gate(histogram) { + Ok(temporality) => temporality, + Err(rejection) => { + reject_data_points(outcome, histogram.data_points.len(), || { + rejection.message(name) + })?; + OTLP_EXPONENTIAL_HISTOGRAM_REJECTED_DATA_POINTS + .with_label_values(&[rejection.reason_label()]) + .inc_by(histogram.data_points.len() as u64); + return Ok(()); } - .build() - })?; - let mut emitted = false; + }; + + let column_schema = metadata + .native_schema + .as_ref() + .context(error::InternalSnafu { + err_msg: "missing native histogram schema", + })?; for (index, data_point) in histogram.data_points.iter().enumerate() { - let (value, timestamp_nanos) = match exponential_histogram_value(data_point) { + let (value, timestamp_nanos) = match exponential_histogram_value(data_point, temporality) { Ok(value) => value, Err(reason) => { reject_data_points(outcome, 1, || { @@ -678,6 +596,11 @@ fn encode_exponential_histogram( timestamp_nanos, metric_ctx, )?; + write_temporality_tag( + table, + &mut row, + temporality == AggregationTemporality::Delta, + )?; row_writer::write_by_schema( table, std::iter::once((column_schema.clone(), Some(value))), @@ -685,54 +608,47 @@ fn encode_exponential_histogram( )?; table.add_row(row); add_accepted_data_points(outcome, 1)?; - emitted = true; } - Ok(emitted) + Ok(()) } +#[derive(Debug)] pub(crate) enum ExponentialHistogramRejection { - DeltaTemporality, UnspecifiedTemporality, } impl ExponentialHistogramRejection { fn reason_label(&self) -> &'static str { match self { - Self::DeltaTemporality => "delta_temporality", Self::UnspecifiedTemporality => "unspecified_temporality", } } fn message(&self, name: &str) -> String { match self { - Self::DeltaTemporality => format!( - "metric `{name}` uses delta OTLP exponential histograms; only cumulative temporality is supported" - ), Self::UnspecifiedTemporality => format!( - "metric `{name}` has unspecified OTLP exponential histogram temporality; cumulative temporality is required" + "metric `{name}` has unspecified OTLP exponential histogram temporality; delta or cumulative temporality is required" ), } } } -/// Whole-metric acceptance, decided once for the encoder and for the resource -/// descriptor, which must not describe a resource whose data was rejected. -/// Individual points can still fail [`exponential_histogram_value`]. -/// Raw-delta storage currently supports sums and explicit histograms; exponential -/// histograms still require cumulative input. Delta-to-cumulative conversion is out of scope. +/// Validate temporality before processing points; delta values remain raw intervals. pub(crate) fn exponential_histogram_gate( histogram: &ExponentialHistogram, -) -> std::result::Result<(), ExponentialHistogramRejection> { +) -> std::result::Result { match AggregationTemporality::try_from(histogram.aggregation_temporality) { - Ok(AggregationTemporality::Cumulative) => Ok(()), - Ok(AggregationTemporality::Delta) => Err(ExponentialHistogramRejection::DeltaTemporality), + Ok(temporality @ (AggregationTemporality::Cumulative | AggregationTemporality::Delta)) => { + Ok(temporality) + } _ => Err(ExponentialHistogramRejection::UnspecifiedTemporality), } } pub(crate) fn exponential_histogram_value( data_point: &ExponentialHistogramDataPoint, + temporality: AggregationTemporality, ) -> std::result::Result<(ValueData, i64), String> { if data_point.start_time_unix_nano > data_point.time_unix_nano { return Err(format!( @@ -747,13 +663,18 @@ pub(crate) fn exponential_histogram_value( let start_timestamp = timestamp_millis(data_point.start_time_unix_nano, "start_time_unix_nano")?; + let reset_hint = if temporality == AggregationTemporality::Delta { + ResetHint::Gauge + } else { + ResetHint::Unspecified + } as i32; let no_recorded_value = data_point.flags & DataPointFlags::NoRecordedValueMask as u32 != 0; let histogram = if no_recorded_value { PromHistogram { sum: f64::from_bits(PROMETHEUS_STALE_NAN_BITS), schema: 0, zero_threshold: 0.0, - reset_hint: ResetHint::Unspecified as i32, + reset_hint, timestamp, start_timestamp, count: Some(PromCount::CountInt(0)), @@ -817,7 +738,7 @@ pub(crate) fn exponential_histogram_value( negative_deltas, positive_spans, positive_deltas, - reset_hint: ResetHint::Unspecified as i32, + reset_hint, timestamp, start_timestamp, count: Some(PromCount::CountInt(data_point.count)), @@ -947,7 +868,6 @@ enum AttributeType { Resource, Scope, DataPoint, - Legacy, } fn write_attributes( @@ -977,7 +897,6 @@ fn write_attributes( translate_label_name(&attr.key, metric_ctx.metric_translation_strategy) ) } - AttributeType::Legacy => legacy_normalize_otlp_name(&attr.key), }; if key == OTLP_AGGREGATION_TEMPORALITY_LABEL { return Err(error::InvalidOtlpMetricInputSnafu { @@ -1007,7 +926,7 @@ fn write_timestamp( // micro/nanosecond). Only the non-metric prometheus-compatible path is // fixed to milliseconds, to keep auto-created mito tables on the // millisecond time index. - if metric_ctx.is_legacy || metric_ctx.with_metric_engine { + if metric_ctx.with_metric_engine { row_writer::write_ts_to_nanos( table, greptime_timestamp(), @@ -1074,40 +993,22 @@ fn write_tags_and_timestamp( timestamp_nanos: i64, metric_ctx: &OtlpMetricCtx, ) -> Result<()> { - if metric_ctx.is_legacy { - write_attributes( - table, - row, - resource_attrs, - AttributeType::Legacy, - metric_ctx, - )?; - write_attributes(table, row, scope_attrs, AttributeType::Legacy, metric_ctx)?; - write_attributes( - table, - row, - data_point_attrs, - AttributeType::Legacy, - metric_ctx, - )?; - } else { - // TODO(shuiyisong): check `__type__` and `__unit__` tags in prometheus - write_attributes( - table, - row, - resource_attrs, - AttributeType::Resource, - metric_ctx, - )?; - write_attributes(table, row, scope_attrs, AttributeType::Scope, metric_ctx)?; - write_attributes( - table, - row, - data_point_attrs, - AttributeType::DataPoint, - metric_ctx, - )?; - } + // TODO(shuiyisong): check `__type__` and `__unit__` tags in prometheus + write_attributes( + table, + row, + resource_attrs, + AttributeType::Resource, + metric_ctx, + )?; + write_attributes(table, row, scope_attrs, AttributeType::Scope, metric_ctx)?; + write_attributes( + table, + row, + data_point_attrs, + AttributeType::DataPoint, + metric_ctx, + )?; write_timestamp(table, row, timestamp_nanos, metric_ctx)?; @@ -1222,7 +1123,7 @@ fn encode_histogram( scope_attrs: Option<&Vec>, metric_ctx: &OtlpMetricCtx, outcome: &mut MetricsIngestOutcome, -) -> Result { +) -> Result<()> { let normalized_name = name; let bucket_table_name = format!("{}{}", normalized_name, BUCKET_TABLE_SUFFIX); @@ -1234,7 +1135,6 @@ fn encode_histogram( Ok(AggregationTemporality::Delta) ); let stale_value = f64::from_bits(PROMETHEUS_STALE_NAN_BITS); - let mut emitted = false; for (index, data_point) in hist.data_points.iter().enumerate() { if let Some(reason) = histogram_data_point_rejection(data_point, is_delta) { reject_data_points(outcome, 1, || { @@ -1357,10 +1257,9 @@ fn encode_histogram( )?; count_table.add_row(count_row); add_accepted_data_points(outcome, 1)?; - emitted = true; } - Ok(emitted) + Ok(()) } pub(crate) fn histogram_data_point_rejection( @@ -1423,17 +1322,47 @@ fn encode_summary( scope_attrs: Option<&Vec>, metric_ctx: &OtlpMetricCtx, ) -> Result<()> { - if metric_ctx.is_legacy { - let table = table_writer.get_or_default_table_data( - name, - APPROXIMATE_COLUMN_COUNT, - summary.data_points.len(), - ); + // 1. quantile table + // 2. count table + // 3. sum table - for data_point in &summary.data_points { - let mut row = table.alloc_one_row(); + let metric_name = name; + let count_name = format!("{}{}", metric_name, COUNT_TABLE_SUFFIX); + let sum_name = format!("{}{}", metric_name, SUM_TABLE_SUFFIX); + + for data_point in &summary.data_points { + { + let quantile_table = table_writer.get_or_default_table_data( + metric_name, + APPROXIMATE_COLUMN_COUNT, + summary.data_points.len(), + ); + + for quantile in &data_point.quantile_values { + let mut row = quantile_table.alloc_one_row(); + write_tags_and_timestamp( + quantile_table, + &mut row, + resource_attrs, + scope_attrs, + Some(data_point.attributes.as_ref()), + data_point.time_unix_nano as i64, + metric_ctx, + )?; + row_writer::write_tag(quantile_table, "quantile", quantile.quantile, &mut row)?; + row_writer::write_f64(quantile_table, greptime_value(), quantile.value, &mut row)?; + quantile_table.add_row(row); + } + } + { + let count_table = table_writer.get_or_default_table_data( + &count_name, + APPROXIMATE_COLUMN_COUNT, + summary.data_points.len(), + ); + let mut row = count_table.alloc_one_row(); write_tags_and_timestamp( - table, + count_table, &mut row, resource_attrs, scope_attrs, @@ -1442,104 +1371,36 @@ fn encode_summary( metric_ctx, )?; - for quantile in &data_point.quantile_values { - row_writer::write_f64( - table, - format!("greptime_p{:02}", quantile.quantile * 100f64), - quantile.value, - &mut row, - )?; - } + row_writer::write_f64( + count_table, + greptime_value(), + data_point.count as f64, + &mut row, + )?; - row_writer::write_f64(table, GREPTIME_COUNT, data_point.count as f64, &mut row)?; - table.add_row(row); + count_table.add_row(row); } - } else { - // 1. quantile table - // 2. count table - // 3. sum table + { + let sum_table = table_writer.get_or_default_table_data( + &sum_name, + APPROXIMATE_COLUMN_COUNT, + summary.data_points.len(), + ); - let metric_name = name; - let count_name = format!("{}{}", metric_name, COUNT_TABLE_SUFFIX); - let sum_name = format!("{}{}", metric_name, SUM_TABLE_SUFFIX); + let mut row = sum_table.alloc_one_row(); + write_tags_and_timestamp( + sum_table, + &mut row, + resource_attrs, + scope_attrs, + Some(data_point.attributes.as_ref()), + data_point.time_unix_nano as i64, + metric_ctx, + )?; - for data_point in &summary.data_points { - { - let quantile_table = table_writer.get_or_default_table_data( - metric_name, - APPROXIMATE_COLUMN_COUNT, - summary.data_points.len(), - ); + row_writer::write_f64(sum_table, greptime_value(), data_point.sum, &mut row)?; - for quantile in &data_point.quantile_values { - let mut row = quantile_table.alloc_one_row(); - write_tags_and_timestamp( - quantile_table, - &mut row, - resource_attrs, - scope_attrs, - Some(data_point.attributes.as_ref()), - data_point.time_unix_nano as i64, - metric_ctx, - )?; - row_writer::write_tag(quantile_table, "quantile", quantile.quantile, &mut row)?; - row_writer::write_f64( - quantile_table, - greptime_value(), - quantile.value, - &mut row, - )?; - quantile_table.add_row(row); - } - } - { - let count_table = table_writer.get_or_default_table_data( - &count_name, - APPROXIMATE_COLUMN_COUNT, - summary.data_points.len(), - ); - let mut row = count_table.alloc_one_row(); - write_tags_and_timestamp( - count_table, - &mut row, - resource_attrs, - scope_attrs, - Some(data_point.attributes.as_ref()), - data_point.time_unix_nano as i64, - metric_ctx, - )?; - - row_writer::write_f64( - count_table, - greptime_value(), - data_point.count as f64, - &mut row, - )?; - - count_table.add_row(row); - } - { - let sum_table = table_writer.get_or_default_table_data( - &sum_name, - APPROXIMATE_COLUMN_COUNT, - summary.data_points.len(), - ); - - let mut row = sum_table.alloc_one_row(); - write_tags_and_timestamp( - sum_table, - &mut row, - resource_attrs, - scope_attrs, - Some(data_point.attributes.as_ref()), - data_point.time_unix_nano as i64, - metric_ctx, - )?; - - row_writer::write_f64(sum_table, greptime_value(), data_point.sum, &mut row)?; - - sum_table.add_row(row); - } + sum_table.add_row(row); } } @@ -1675,43 +1536,99 @@ mod tests { } #[test] - fn test_conversion_skips_descriptor_for_legacy_mode() { - set_default_prefix(None).unwrap(); - let request = gauge_request( - vec![keyvalue("service.name", "api"), keyvalue("host.id", "h-1")], - "my_gauge", - ); - let mut ctx = OtlpMetricCtx { - is_legacy: true, - ..descriptor_ctx() + fn test_descriptors_exclude_points_rejected_by_destination_preflight() { + let mut request = gauge_request(vec![keyvalue("service.name", "api")], "accepted"); + let mut rejected = request.resource_metrics[0].clone(); + let metric = &mut rejected.scope_metrics[0].metrics[0]; + metric.name = "rejected".into(); + let Some(metric::Data::Gauge(gauge)) = metric.data.as_mut() else { + unreachable!() }; - let conversion = to_grpc_insert_requests(request, &mut ctx).unwrap(); - assert!(conversion.resource_info.is_none()); + gauge.data_points[0].time_unix_nano = + 2 * common_catalog::consts::SEMANTIC_GRAPH_WINDOW_NANOS as u64; + request.resource_metrics.push(rejected); + let prepared = prepare_metrics(request, &mut descriptor_ctx()).unwrap(); + let conversion = prepared + .finish( + vec![None, Some("incompatible schema".into())], + &Default::default(), + ) + .unwrap(); + assert_eq!( + ( + conversion.outcome.accepted_data_points, + conversion.outcome.rejected_data_points + ), + (1, 1) + ); + let descriptors = conversion.resource_info.unwrap(); + let rows = descriptors.inserts[0].rows.as_ref().unwrap(); + assert_eq!(rows.rows.len(), 1); + let timestamp = rows + .schema + .iter() + .position(|c| c.column_name == greptime_timestamp()) + .unwrap(); + assert_eq!( + rows.rows[0].values[timestamp].value_data, + Some(ValueData::TimestampMillisecondValue(1)) + ); + } - // legacy tables predate job/instance and the promote filter; adding - // either would alter the schema of tables already in use - let cols = column_names(&conversion.requests, "my_gauge"); - assert!(!cols.contains(&"job".to_string())); - assert!(cols.contains(&"host_id".to_string())); + #[test] + fn test_finish_rejects_mismatched_point_counts() { + for (requests, rejections) in [(0, 1), (2, 1), (1, 0), (1, 2)] { + let mut prepared = prepare_metrics( + gauge_request(vec![], "gauge"), + &mut OtlpMetricCtx::default(), + ) + .unwrap(); + prepared.requests.resize(requests, Default::default()); + let err = prepared + .finish(vec![None; rejections], &Default::default()) + .unwrap_err(); + assert!(matches!(err, error::Error::Internal { .. })); + } } #[test] fn test_conversion_skips_descriptor_on_metric_name_collision() { set_default_prefix(None).unwrap(); - let request = gauge_request( - vec![keyvalue("service.name", "api")], - OTEL_RESOURCE_INFO_TABLE_NAME, - ); - let conversion = to_grpc_insert_requests(request, &mut descriptor_ctx()).unwrap(); - assert!(conversion.resource_info.is_none()); - // the metric itself still goes through the main path - assert!( - conversion - .requests - .inserts - .iter() - .any(|r| r.table_name == OTEL_RESOURCE_INFO_TABLE_NAME) - ); + for (enabled, attributes, warns) in [ + (true, vec![keyvalue("service.name", "api")], true), + (false, vec![keyvalue("service.name", "api")], false), + (true, vec![keyvalue("unprojected", "value")], false), + ] { + let logs = tempfile::NamedTempFile::new().unwrap(); + let writer = logs.reopen().unwrap(); + let subscriber = common_telemetry::tracing_subscriber::fmt() + .without_time() + .with_ansi(false) + .with_writer(move || writer.try_clone().unwrap()) + .finish(); + let conversion = tracing::subscriber::with_default(subscriber, || { + let mut ctx = descriptor_ctx(); + ctx.resource_info = enabled; + to_grpc_insert_requests( + gauge_request(attributes, OTEL_RESOURCE_INFO_TABLE_NAME), + &mut ctx, + ) + .unwrap() + }); + let diagnostics = std::fs::read_to_string(logs.path()).unwrap(); + assert_eq!( + diagnostics.contains("Skipping OTLP resource descriptor"), + warns + ); + assert!(conversion.resource_info.is_none()); + assert!( + conversion + .requests + .inserts + .iter() + .any(|r| r.table_name == OTEL_RESOURCE_INFO_TABLE_NAME) + ); + } } #[test] @@ -1984,47 +1901,6 @@ mod tests { ); } - #[test] - fn test_encode_legacy_summary_keeps_legacy_column_names() { - set_default_prefix(Some("custom")).unwrap(); - let mut tables = MultiTableData::default(); - let summary = Summary { - data_points: vec![SummaryDataPoint { - attributes: vec![keyvalue("host", "testserver")], - time_unix_nano: 100, - count: 25, - quantile_values: vec![ValueAtQuantile { - quantile: 0.90, - value: 1000.0, - }], - ..Default::default() - }], - }; - - encode_summary( - &mut tables, - "datamon", - &summary, - None, - None, - &OtlpMetricCtx { - is_legacy: true, - ..Default::default() - }, - ) - .unwrap(); - - let table = tables.get_or_default_table_data("datamon", 0, 0); - assert_eq!( - table - .columns() - .iter() - .map(|column| column.column_name.as_str()) - .collect::>(), - vec!["host", "custom_timestamp", "greptime_p90", GREPTIME_COUNT,] - ); - } - #[test] fn test_encode_histogram() { let mut tables = MultiTableData::default(); @@ -2134,7 +2010,12 @@ mod tests { ..Default::default() }; let mut index = SemanticIndex::default(); - record_metric_semantics(&mut index, metric, name, &ctx); + record_metric_semantics( + &mut index, + metric, + name, + &emitted_semantic_tables(&ctx.metric_type, name), + ); index } @@ -2471,7 +2352,8 @@ mod tests { negative: Some(buckets), ..Default::default() }; - let (value, _) = exponential_histogram_value(&point).unwrap(); + let (value, _) = + exponential_histogram_value(&point, AggregationTemporality::Cumulative).unwrap(); assert_eq!( native_field(&value, COUNT_I64_FIELD), Some(ValueData::I64Value(point.count as i64)) @@ -2505,7 +2387,9 @@ mod tests { ZERO_COUNT_I64_FIELD, }; - let (value, timestamp_nanos) = exponential_histogram_value(&exponential_point()).unwrap(); + let (value, timestamp_nanos) = + exponential_histogram_value(&exponential_point(), AggregationTemporality::Cumulative) + .unwrap(); assert_eq!(timestamp_nanos, 2_000_000); assert_eq!( @@ -2559,7 +2443,8 @@ mod tests { zero_threshold: f64::NAN, ..Default::default() }; - let (value, _) = exponential_histogram_value(&point).unwrap(); + let (value, _) = + exponential_histogram_value(&point, AggregationTemporality::Cumulative).unwrap(); assert_eq!( native_field(&value, SCHEMA_FIELD), @@ -2592,7 +2477,8 @@ mod tests { sum: Some(f64::from_bits(PROMETHEUS_STALE_NAN_BITS)), ..Default::default() }; - let (value, _) = exponential_histogram_value(&point).unwrap(); + let (value, _) = + exponential_histogram_value(&point, AggregationTemporality::Cumulative).unwrap(); let Some(ValueData::F64Value(sum)) = native_field(&value, SUM_FIELD) else { panic!("expected histogram sum"); }; @@ -2636,7 +2522,8 @@ mod tests { cases.push((point, "shifted bucket index overflows i32")); for (point, expected) in cases { - let error = exponential_histogram_value(&point).unwrap_err(); + let error = exponential_histogram_value(&point, AggregationTemporality::Cumulative) + .unwrap_err(); assert!( error.contains(expected), "expected {expected:?}, got {error}" @@ -2803,9 +2690,9 @@ mod tests { ( AggregationTemporality::Delta, vec![exponential_point(); 2], - "delta_temporality", - 0, + "invalid_data_point", 2, + 0, ), ( AggregationTemporality::Unspecified, @@ -2886,20 +2773,17 @@ mod tests { .. } = to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()).unwrap(); - assert_eq!(outcome.accepted_data_points, 2); - assert_eq!(outcome.rejected_data_points, 1); - assert!( - outcome - .error_message - .as_deref() - .unwrap() - .contains("only cumulative temporality is supported") - ); - assert_eq!(requests.inserts.len(), 2); + assert_eq!(outcome.accepted_data_points, 3); + assert_eq!(outcome.rejected_data_points, 0); + assert!(outcome.error_message.is_none()); + assert_eq!(requests.inserts.len(), 3); let semantics = decode(&semantic_index); assert!(semantics.contains_key("temperature")); assert!(semantics.contains_key("latency")); - assert!(!semantics.contains_key("delta_latency")); + assert_eq!( + semantics["delta_latency"][SEMANTIC_METRIC_TEMPORALITY], + "delta" + ); let empty = metrics_request(vec![exponential_metric( "empty", @@ -2934,12 +2818,10 @@ mod tests { ]); let mut ctx = OtlpMetricCtx::default(); - let error = to_grpc_insert_requests(request, &mut ctx).unwrap_err(); - assert!( - error - .to_string() - .contains("cannot mix native histogram and float sample fields") - ); + let conversion = to_grpc_insert_requests(request, &mut ctx).unwrap(); + assert_eq!(conversion.outcome.accepted_data_points, 0); + assert_eq!(conversion.outcome.rejected_data_points, 2); + assert!(conversion.requests.inserts.is_empty()); } #[test] @@ -2954,12 +2836,10 @@ mod tests { ]); let mut ctx = OtlpMetricCtx::default(); - let error = to_grpc_insert_requests(request, &mut ctx).unwrap_err(); - assert!( - error - .to_string() - .contains("cannot mix native histogram and float sample fields") - ); + let conversion = to_grpc_insert_requests(request, &mut ctx).unwrap(); + assert_eq!(conversion.outcome.accepted_data_points, 0); + assert_eq!(conversion.outcome.rejected_data_points, 2); + assert!(conversion.requests.inserts.is_empty()); } #[test] @@ -3008,34 +2888,29 @@ mod tests { flags: DataPointFlags::NoRecordedValueMask as u32, ..Default::default() }; - for temporality in [ - AggregationTemporality::Delta, + let request = metrics_request(vec![exponential_metric( + "latency", + vec![stale], AggregationTemporality::Unspecified, - ] { - let request = metrics_request(vec![exponential_metric( - "latency", - vec![stale.clone()], - temporality, - )]); - let mut ctx = OtlpMetricCtx::default(); - let MetricsConversion { - requests, - rows, - semantic_index, - outcome, - .. - } = to_grpc_insert_requests(request, &mut ctx).unwrap(); + )]); + let mut ctx = OtlpMetricCtx::default(); + let MetricsConversion { + requests, + rows, + semantic_index, + outcome, + .. + } = to_grpc_insert_requests(request, &mut ctx).unwrap(); - assert_eq!(outcome.accepted_data_points, 0); - assert_eq!(outcome.rejected_data_points, 1); - assert_eq!(rows, 0); - assert!(requests.inserts.is_empty()); - assert!(semantic_index.is_empty()); - } + assert_eq!(outcome.accepted_data_points, 0); + assert_eq!(outcome.rejected_data_points, 1); + assert_eq!(rows, 0); + assert!(requests.inserts.is_empty()); + assert!(semantic_index.is_empty()); } #[test] - fn test_exponential_histogram_legacy_and_new_modes_share_struct() { + fn test_exponential_histogram_storage_paths_share_struct() { use common_query::prelude::greptime_native_histogram; let mut point = exponential_point(); @@ -3045,40 +2920,40 @@ mod tests { vec![point], AggregationTemporality::Cumulative, )]); - let mut new_ctx = OtlpMetricCtx::default(); - let new_requests = to_grpc_insert_requests(request.clone(), &mut new_ctx) + let mut ordinary_ctx = OtlpMetricCtx::default(); + let ordinary_requests = to_grpc_insert_requests(request.clone(), &mut ordinary_ctx) .unwrap() .requests; - let mut legacy_ctx = OtlpMetricCtx { - is_legacy: true, + let mut metric_ctx = OtlpMetricCtx { + with_metric_engine: true, ..Default::default() }; - let legacy_requests = to_grpc_insert_requests(request, &mut legacy_ctx) + let metric_requests = to_grpc_insert_requests(request, &mut metric_ctx) .unwrap() .requests; - let new_insert = &new_requests.inserts[0]; - let legacy_insert = &legacy_requests.inserts[0]; - assert_eq!(new_insert.table_name, "request_duration"); - assert_eq!(legacy_insert.table_name, "request_duration"); - let new_rows = new_insert.rows.as_ref().unwrap(); - let legacy_rows = legacy_insert.rows.as_ref().unwrap(); + let ordinary_insert = &ordinary_requests.inserts[0]; + let metric_insert = &metric_requests.inserts[0]; + assert_eq!(ordinary_insert.table_name, "request_duration"); + assert_eq!(metric_insert.table_name, "request_duration"); + let ordinary_rows = ordinary_insert.rows.as_ref().unwrap(); + let metric_rows = metric_insert.rows.as_ref().unwrap(); let field = greptime_native_histogram(); - let new_histogram = new_rows.rows[0].values[new_rows + let ordinary_histogram = ordinary_rows.rows[0].values[ordinary_rows .schema .iter() .position(|column| column.column_name == field) .unwrap()] .clone(); - let legacy_histogram = legacy_rows.rows[0].values[legacy_rows + let metric_histogram = metric_rows.rows[0].values[metric_rows .schema .iter() .position(|column| column.column_name == field) .unwrap()] .clone(); - assert_eq!(new_histogram, legacy_histogram); + assert_eq!(ordinary_histogram, metric_histogram); assert!(matches!( - new_rows.rows[0].values[new_rows + ordinary_rows.rows[0].values[ordinary_rows .schema .iter() .position(|column| column.column_name == greptime_timestamp()) @@ -3087,7 +2962,7 @@ mod tests { Some(ValueData::TimestampMillisecondValue(2)) )); assert!(matches!( - legacy_rows.rows[0].values[legacy_rows + metric_rows.rows[0].values[metric_rows .schema .iter() .position(|column| column.column_name == greptime_timestamp()) @@ -3102,7 +2977,7 @@ mod tests { let request = metrics_request(vec![exponential_metric( "x".repeat(1_000), vec![ExponentialHistogramDataPoint::default()], - AggregationTemporality::Delta, + AggregationTemporality::Unspecified, )]); let outcome = to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()) .unwrap() diff --git a/src/servers/src/otlp/metrics/prepare.rs b/src/servers/src/otlp/metrics/prepare.rs new file mode 100644 index 00000000000..6254cf0120c --- /dev/null +++ b/src/servers/src/otlp/metrics/prepare.rs @@ -0,0 +1,623 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::{HashMap, HashSet}; +use std::hash::{BuildHasher, Hash, Hasher}; +use std::ops::Range; + +use ahash::{HashMap as AHashMap, HashMapExt}; +use api::v1::{ColumnSchema, RowInsertRequest, RowInsertRequests, Rows, SemanticType}; +use common_base::regex_pattern::NAME_PATTERN_REG; +use common_query::native_histogram::native_histogram_value_type; +use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest; +use otel_arrow_rust::proto::opentelemetry::common::v1::{KeyValue, any_value}; +use otel_arrow_rust::proto::opentelemetry::metrics::v1::{Metric, NumberDataPoint, metric}; +use session::protocol_ctx::OtlpMetricCtx; +use snafu::{OptionExt, ensure}; +use table::requests::{ + METADATA_QUALITY_DECLARED, SEMANTIC_METRIC_METADATA_QUALITY, SEMANTIC_METRIC_TYPE, +}; + +use crate::error::{self, Result}; +use crate::metrics::OTLP_EXPONENTIAL_HISTOGRAM_REJECTED_DATA_POINTS; +use crate::otlp::metrics::{ + MetricMetadata, MetricsConversion, OTEL_RESOURCE_INFO_TABLE_NAME, ResourceInfoData, + add_accepted_data_points, encode_metrics, from_metric_type, has_no_recorded_value, + process_resource_attrs, process_scope_attrs, reject_data_points, validate_sample_kind, +}; +use crate::query_handler::MetricsIngestOutcome; +use crate::row_writer::{MultiTableData, TableData}; +use crate::semantic::SemanticIndex; + +/// Rows share exact input schemas while retaining their source-point boundaries. +pub struct PreparedMetrics { + /// Destination/schema batches. Normalization must preserve batch and row order/counts. + pub requests: Vec, + points: Vec, + row_locations: Vec<(usize, usize)>, + metrics: Vec, + resources: Vec>, + outcome: MetricsIngestOutcome, +} + +struct PreparedPoint { + metric: usize, + resource: usize, + timestamp: i64, + rows: Range, +} + +#[derive(Default)] +struct MetricWriter { + tables: MultiTableData, + // TableData schemas only grow. A new column invalidates the cached batch. + batches: AHashMap, +} + +pub fn prepare_metrics( + request: ExportMetricsServiceRequest, + metric_ctx: &mut OtlpMetricCtx, +) -> Result { + let mut prepared = PreparedMetrics { + requests: Vec::new(), + points: Vec::new(), + row_locations: Vec::new(), + metrics: Vec::new(), + resources: Vec::new(), + outcome: MetricsIngestOutcome::default(), + }; + let mut schema_batches = AHashMap::>::new(); + for resource in request.resource_metrics { + let raw_attrs = resource.resource.map(|r| r.attributes).unwrap_or_default(); + let mut attrs = raw_attrs.clone(); + process_resource_attrs(&mut attrs, metric_ctx); + let resource_index = prepared.resources.len(); + prepared.resources.push(if metric_ctx.resource_info { + raw_attrs + } else { + Vec::new() + }); + for scope in resource.scope_metrics { + let scope_attrs = process_scope_attrs(&scope, metric_ctx); + for mut metric in scope.metrics { + let Some(data) = metric.data.as_ref() else { + continue; + }; + metric_ctx.set_metric_type(from_metric_type(data)); + let metric_index = prepared.metrics.len(); + prepared + .metrics + .push(MetricMetadata::new(&metric, metric_ctx)?); + let mut writer = MetricWriter::default(); + let mut layout: (Vec, bool) = (Vec::new(), false); + let mut writers = AHashMap::new(); + // Move each payload once; never clone the accumulated request. + macro_rules! prepare_points { + ($data:ident, $variant:ident, $omits_value:expr) => { + for point in std::mem::take(&mut $data.data_points) { + // Omissions must not become explicit nulls in a retained schema. + let omits_value = ($omits_value)(&point); + if omits_value != layout.1 + || !scalar_attribute_names(&point.attributes) + .eq(layout.0.iter().map(String::as_str)) + { + let next_layout = ( + scalar_attribute_names(&point.attributes) + .map(str::to_owned) + .collect(), + omits_value, + ); + let next_writer = writers.remove(&next_layout).unwrap_or_default(); + let previous = std::mem::replace(&mut writer, next_writer); + let previous_layout = std::mem::replace(&mut layout, next_layout); + if !previous.batches.is_empty() { + writers.insert(previous_layout, previous); + } + } + let timestamp = point.time_unix_nano; + let start = point.start_time_unix_nano; + $data.data_points.push(point); + metric.data = Some(metric::Data::$variant($data)); + if !prepared.push_point( + &metric, + metric_index, + timestamp, + start, + &attrs, + scope_attrs.as_ref(), + resource_index, + metric_ctx, + &mut writer, + &mut schema_batches, + )? { + writer = MetricWriter::default(); + } + let Some(metric::Data::$variant(mut remaining)) = metric.data.take() + else { + unreachable!() + }; + remaining.data_points.clear(); + $data = remaining; + } + }; + } + let Some(data) = metric.data.take() else { + continue; + }; + match data { + metric::Data::Gauge(mut data) => { + prepare_points!(data, Gauge, |point: &NumberDataPoint| point + .value + .is_none()) + } + metric::Data::Sum(mut data) => { + prepare_points!(data, Sum, |point: &NumberDataPoint| point.value.is_none() + && !has_no_recorded_value(point.flags)) + } + metric::Data::Histogram(mut data) => { + prepare_points!(data, Histogram, |_| false) + } + metric::Data::Summary(mut data) => prepare_points!(data, Summary, |_| false), + metric::Data::ExponentialHistogram(mut data) => { + prepare_points!(data, ExponentialHistogram, |_| false) + } + } + } + } + } + Ok(prepared) +} + +fn scalar_attribute_names(attrs: &[KeyValue]) -> impl Iterator { + attrs.iter().filter_map(|attr| { + matches!( + attr.value.as_ref()?.value.as_ref()?, + any_value::Value::StringValue(_) + | any_value::Value::BoolValue(_) + | any_value::Value::IntValue(_) + | any_value::Value::DoubleValue(_) + ) + .then_some(attr.key.as_str()) + }) +} + +impl PreparedMetrics { + #[allow(clippy::too_many_arguments)] + fn push_point( + &mut self, + metric: &Metric, + metric_index: usize, + timestamp: u64, + start: u64, + resource_attrs: &Vec, + scope_attrs: Option<&Vec>, + resource: usize, + ctx: &OtlpMetricCtx, + writer: &mut MetricWriter, + schema_batches: &mut AHashMap>, + ) -> Result { + let metadata = &self.metrics[metric_index]; + let exponential = metadata.native_schema.is_some(); + // OTLP gauges ignore StartTimeUnixNano. + let invalid_interval = + !matches!(metric.data, Some(metric::Data::Gauge(_))) && start > timestamp; + let mut outcome = MetricsIngestOutcome::default(); + let encoded = if !exponential && (timestamp > i64::MAX as u64 || invalid_interval) { + Err(error::InvalidOtlpMetricInputSnafu { + reason: "invalid OTLP data-point timestamp or interval", + } + .build()) + } else { + encode_metrics( + &mut writer.tables, + metric, + metadata, + Some(resource_attrs), + scope_attrs, + ctx, + &mut outcome, + ) + }; + if let Err(err) = encoded { + match err { + error::Error::InvalidOtlpMetricInput { .. } + | error::Error::IncompatibleSchema { .. } + | error::Error::TimestampOverflow { .. } + | error::Error::TimePrecision { .. } + | error::Error::InvalidParameter { .. } => { + self.reject(exponential, err.to_string())?; + return Ok(false); + } + _ => return Err(err), + } + } + if outcome.rejected_data_points != 0 { + // Payload encoders already counted histogram-specific rejections. + reject_data_points(&mut self.outcome, 1, || { + outcome.error_message.unwrap_or_default() + })?; + return Ok(false); + } + if let Some((name, _)) = writer.tables.iter().find(|(name, table)| { + table.num_rows() != 0 + && !writer.batches.contains_key(*name) + && !NAME_PATTERN_REG.is_match(name) + }) { + self.reject(exponential, format!("invalid output table name `{name}`"))?; + return Ok(false); + } + let first_row = self.row_locations.len(); + for (name, schema, rows) in writer.tables.drain_rows() { + let batch = match writer.batches.get(name) { + Some(&(columns, batch)) if columns == schema.len() => batch, + _ => { + let mut hash = schema_batches.hasher().build_hasher(); + name.hash(&mut hash); + for column in schema { + (&column.column_name, column.datatype, column.semantic_type) + .hash(&mut hash); + } + let candidates = schema_batches.entry(hash.finish()).or_default(); + let batch = candidates.iter().copied().find(|&index| { + self.requests[index].table_name == name + && self.requests[index] + .rows + .as_ref() + .is_some_and(|rows| rows.schema == schema) + }); + let batch = if let Some(batch) = batch { + batch + } else { + validate_sample_kind(name, schema)?; + let batch = self.requests.len(); + self.requests.push(RowInsertRequest { + table_name: name.to_owned(), + rows: Some(Rows { + schema: schema.to_vec(), + rows: Vec::new(), + }), + }); + candidates.push(batch); + batch + }; + writer + .batches + .insert(name.to_owned(), (schema.len(), batch)); + batch + } + }; + let output = self.requests[batch] + .rows + .as_mut() + .context(error::InternalSnafu { + err_msg: "missing prepared batch", + })?; + for row in rows { + self.row_locations.push((batch, output.rows.len())); + output.rows.push(row); + } + } + self.points.push(PreparedPoint { + metric: metric_index, + resource, + timestamp: timestamp as i64, + rows: first_row..self.row_locations.len(), + }); + Ok(true) + } + + fn reject(&mut self, exponential: bool, reason: String) -> Result<()> { + reject_data_points(&mut self.outcome, 1, || reason)?; + if exponential { + OTLP_EXPONENTIAL_HISTOGRAM_REJECTED_DATA_POINTS + .with_label_values(&["invalid_data_point"]) + .inc(); + } + Ok(()) + } + + /// Number of source points that survived payload conversion. + pub fn num_points(&self) -> usize { + self.points.len() + } + + fn validate_batches(&self) -> Result<()> { + let mut counts = vec![0; self.requests.len()]; + for &(batch, row) in &self.row_locations { + let count = counts.get_mut(batch).context(error::InternalSnafu { + err_msg: "missing prepared batch", + })?; + ensure!( + row == *count, + error::InternalSnafu { + err_msg: "invalid prepared row order", + } + ); + *count += 1; + } + ensure!( + self.requests.iter().zip(counts).all(|(request, count)| { + count != 0 + && request + .rows + .as_ref() + .is_some_and(|rows| rows.rows.len() == count) + }), + error::InternalSnafu { + err_msg: "OTLP prepared batch row counts changed" + } + ); + Ok(()) + } + + /// Rejects a source point if any row in its fan-out failed destination validation. + pub fn point_rejections(&self, rows: Vec>>) -> Result>> { + self.validate_batches()?; + ensure!( + rows.len() == self.requests.len() + && rows.iter().zip(&self.requests).all(|(errors, request)| { + request + .rows + .as_ref() + .is_some_and(|rows| errors.len() == rows.rows.len()) + }), + error::InternalSnafu { + err_msg: "OTLP batch rejections must match prepared rows" + } + ); + Ok(self + .points + .iter() + .map(|point| { + self.row_locations[point.rows.clone()] + .iter() + .find_map(|&(batch, row)| rows[batch][row].clone()) + }) + .collect()) + } + + /// Complete authoritative validation before calling this method. Fresh + /// logical and physical conflicts are then collected from one survivor set. + pub fn finish( + mut self, + mut rejections: Vec>, + physical_tables: &HashMap, + ) -> Result { + ensure!( + rejections.len() == self.points.len(), + error::InternalSnafu { + err_msg: "OTLP prepared points and rejections must have equal lengths", + } + ); + self.validate_batches()?; + let conflicts = conflicting_points( + &self.requests, + &self.points, + &self.row_locations, + &rejections, + physical_tables, + ); + for index in conflicts { + rejections[index] = + Some("incompatible metric sample kinds or column definitions".to_string()); + } + let reuse_batches = rejections.iter().all(Option::is_none) && { + let mut tables = HashSet::new(); + self.requests.len() <= 1 || self.requests.iter().all(|r| tables.insert(&r.table_name)) + }; + let mut outputs: Vec<(String, TableData)> = Vec::new(); + let mut output_indices = AHashMap::new(); + let mut mappings: Vec, bool)>> = if reuse_batches { + Vec::new() + } else { + vec![None; self.requests.len()] + }; + let mut semantic_index = SemanticIndex::default(); + let mut resource_info = ResourceInfoData::default(); + let mut resource_times = vec![Vec::new(); self.resources.len()]; + for (point, rejection) in std::mem::take(&mut self.points).into_iter().zip(rejections) { + if let Some(reason) = rejection { + let metadata = &self.metrics[point.metric]; + self.reject( + metadata.native_schema.is_some(), + format!("metric `{}`: {reason}", metadata.original_name), + )?; + continue; + } + for &(batch, row) in &self.row_locations[point.rows] { + let request = &mut self.requests[batch]; + semantic_index.merge_table_from( + &mut self.metrics[point.metric].semantics, + &request.table_name, + ); + if reuse_batches { + continue; + } + let rows = request.rows.as_mut().context(error::InternalSnafu { + err_msg: "missing prepared batch", + })?; + if mappings[batch].is_none() { + let output = *output_indices + .entry(request.table_name.clone()) + .or_insert_with(|| { + let index = outputs.len(); + outputs.push(( + request.table_name.clone(), + TableData::new(rows.schema.len(), rows.rows.len()), + )); + index + }); + let columns = std::mem::take(&mut rows.schema) + .into_iter() + .map(|column| outputs[output].1.ensure_column(column)) + .collect::>>()?; + let identity = columns + .iter() + .enumerate() + .all(|(index, &column)| index == column); + mappings[batch] = Some((output, columns, identity)); + } + let (output, columns, identity) = + mappings[batch].as_ref().context(error::InternalSnafu { + err_msg: "missing output column mapping", + })?; + let table = &mut outputs[*output].1; + let values = std::mem::take(&mut rows.rows[row].values); + if *identity { + table.add_row(values); + } else { + let mut output = table.alloc_one_row(); + for (value, &column) in values.into_iter().zip(columns) { + output[column] = value; + } + table.add_row(output); + } + } + add_accepted_data_points(&mut self.outcome, 1)?; + if !self.resources[point.resource].is_empty() { + resource_times[point.resource].push(point.timestamp); + } + } + for (attrs, times) in self.resources.iter().zip(resource_times) { + if !times.is_empty() { + resource_info.observe_times(attrs, times); + } + } + let requests = RowInsertRequests { + inserts: if reuse_batches { + std::mem::take(&mut self.requests) + } else { + outputs + .into_iter() + .map(|(table_name, table)| RowInsertRequest { + table_name, + rows: Some(table.into_rows()), + }) + .collect() + }, + }; + let rows = requests + .inserts + .iter() + .filter_map(|r| r.rows.as_ref()) + .map(|rows| rows.rows.len()) + .sum(); + let resource_info = if !resource_info.is_empty() + && requests + .inserts + .iter() + .any(|r| r.table_name == OTEL_RESOURCE_INFO_TABLE_NAME) + { + common_telemetry::warn!( + "Skipping OTLP resource descriptor synthesis: a metric uses its table name" + ); + None + } else { + resource_info.into_row_insert_requests()? + }; + if resource_info.is_some() { + semantic_index.record_scalar( + OTEL_RESOURCE_INFO_TABLE_NAME, + SEMANTIC_METRIC_TYPE, + crate::semantic::METRIC_TYPE_INFO, + ); + semantic_index.record_scalar( + OTEL_RESOURCE_INFO_TABLE_NAME, + SEMANTIC_METRIC_METADATA_QUALITY, + METADATA_QUALITY_DECLARED, + ); + } + Ok(MetricsConversion { + requests, + rows, + semantic_index, + resource_info, + outcome: self.outcome, + }) + } +} + +fn conflicting_points( + requests: &[RowInsertRequest], + points: &[PreparedPoint], + row_locations: &[(usize, usize)], + rejections: &[Option], + physical_tables: &HashMap, +) -> HashSet { + let mut contributors = vec![Vec::new(); requests.len()]; + for (index, point) in points + .iter() + .enumerate() + .filter(|(index, _)| rejections[*index].is_none()) + { + for &(batch, _) in &row_locations[point.rows.clone()] { + if contributors[batch].last() != Some(&index) { + contributors[batch].push(index); + } + } + } + let mut columns = AHashMap::<_, Vec<(usize, &ColumnSchema)>>::new(); + let mut kinds: AHashMap<&str, Vec<(usize, bool)>> = AHashMap::new(); + for (batch, request) in requests + .iter() + .enumerate() + .filter(|(batch, _)| !contributors[*batch].is_empty()) + { + let Some(rows) = &request.rows else { continue }; + for column in &rows.schema { + columns + .entry((false, &request.table_name, &column.column_name)) + .or_default() + .push((batch, column)); + if let Some(physical) = physical_tables.get(&request.table_name) { + columns + .entry((true, physical, &column.column_name)) + .or_default() + .push((batch, column)); + } + if column.semantic_type == SemanticType::Field as i32 { + let native = api::helper::is_column_type_value_eq( + column.datatype, + column.datatype_extension.clone(), + native_histogram_value_type(), + ); + kinds + .entry(&request.table_name) + .or_default() + .push((batch, native)); + } + } + } + let mut rejected = HashSet::new(); + for group in columns.values() { + let first = group[0].1; + if group.iter().any(|(_, column)| { + column.datatype != first.datatype + || column.datatype_extension != first.datatype_extension + || column.semantic_type != first.semantic_type + }) { + for (batch, _) in group { + rejected.extend(contributors[*batch].iter().copied()); + } + } + } + for group in kinds.values() { + if group.iter().any(|(_, native)| *native) && group.iter().any(|(_, native)| !*native) { + for (batch, _) in group { + rejected.extend(contributors[*batch].iter().copied()); + } + } + } + rejected +} diff --git a/src/servers/src/otlp/metrics/resource_info.rs b/src/servers/src/otlp/metrics/resource_info.rs index 4b4f2afcad6..c7ad72c3e47 100644 --- a/src/servers/src/otlp/metrics/resource_info.rs +++ b/src/servers/src/otlp/metrics/resource_info.rs @@ -27,15 +27,10 @@ use common_catalog::consts::SEMANTIC_GRAPH_WINDOW_NANOS; use common_grpc::precision::Precision; use common_query::prelude::{greptime_timestamp, greptime_value}; use otel_arrow_rust::proto::opentelemetry::common::v1::KeyValue; -use otel_arrow_rust::proto::opentelemetry::metrics::v1::{ - AggregationTemporality, ResourceMetrics, metric, -}; use crate::error::Result; use crate::otlp::metrics::{ - INSTANCE_KEY, JOB_KEY, ServiceIdentity, exponential_histogram_gate, - exponential_histogram_value, histogram_data_point_rejection, scalar_value_string, - service_identity, + INSTANCE_KEY, JOB_KEY, ServiceIdentity, scalar_value_string, service_identity, }; use crate::otlp::trace::{ KEY_CONTAINER_ID, KEY_CONTAINER_NAME, KEY_HOST_ID, KEY_HOST_NAME, KEY_K8S_CONTAINER_NAME, @@ -88,8 +83,15 @@ pub struct ResourceInfoData { } impl ResourceInfoData { - /// Takes the raw attributes, before the promote filter runs on them. - pub fn observe(&mut self, raw_attrs: &[KeyValue], resource: &ResourceMetrics) { + pub(crate) fn is_empty(&self) -> bool { + self.rows.is_empty() + } + + pub(super) fn observe_times( + &mut self, + raw_attrs: &[KeyValue], + times: impl IntoIterator, + ) { let mut tags = Vec::with_capacity(MAX_PROJECTED_TAGS); let ServiceIdentity { job, instance } = service_identity(raw_attrs); if let Some(job) = job { @@ -113,13 +115,13 @@ impl ResourceInfoData { tags.sort_unstable(); let mut observed: BTreeMap = BTreeMap::new(); - for_each_encoded_time(resource, |ts| { + for ts in times { let window = ts - ts.rem_euclid(SEMANTIC_GRAPH_WINDOW_NANOS); observed .entry(window) .and_modify(|newest| *newest = (*newest).max(ts)) .or_insert(ts); - }); + } if observed.is_empty() { return; } @@ -167,58 +169,6 @@ impl ResourceInfoData { } } -/// Visits the times of the data points the encoder writes rows for, so a -/// resource is described exactly where it is measured rather than wherever -/// its request happens to reach. -fn for_each_encoded_time(resource: &ResourceMetrics, mut visit: impl FnMut(i64)) { - fn visit_all(points: impl Iterator, visit: &mut impl FnMut(i64)) { - for ts in points { - visit(ts as i64); - } - } - for scope in &resource.scope_metrics { - for m in &scope.metrics { - match &m.data { - Some(metric::Data::Gauge(g)) => { - visit_all(g.data_points.iter().map(|p| p.time_unix_nano), &mut visit) - } - Some(metric::Data::Sum(s)) => { - visit_all(s.data_points.iter().map(|p| p.time_unix_nano), &mut visit) - } - Some(metric::Data::Histogram(h)) => { - let is_delta = matches!( - AggregationTemporality::try_from(h.aggregation_temporality), - Ok(AggregationTemporality::Delta) - ); - visit_all( - h.data_points - .iter() - .filter(|point| { - histogram_data_point_rejection(point, is_delta).is_none() - }) - .map(|point| point.time_unix_nano), - &mut visit, - ) - } - Some(metric::Data::Summary(s)) => { - visit_all(s.data_points.iter().map(|p| p.time_unix_nano), &mut visit) - } - Some(metric::Data::ExponentialHistogram(h)) - if exponential_histogram_gate(h).is_ok() => - { - for point in &h.data_points { - if let Ok((_, ts)) = exponential_histogram_value(point) { - visit(ts); - } - } - } - Some(metric::Data::ExponentialHistogram(_)) => {} - None => {} - } - } - } -} - #[cfg(test)] mod tests { use api::v1::SemanticType; @@ -226,8 +176,8 @@ mod tests { use common_query::prelude::set_default_prefix; use otel_arrow_rust::proto::opentelemetry::common::v1::{AnyValue, any_value}; use otel_arrow_rust::proto::opentelemetry::metrics::v1::{ - AggregationTemporality, ExponentialHistogram, ExponentialHistogramDataPoint, Gauge, Metric, - NumberDataPoint, ScopeMetrics, + AggregationTemporality, ExponentialHistogram, ExponentialHistogramDataPoint, Metric, + ResourceMetrics, ScopeMetrics, metric, }; use super::*; @@ -243,27 +193,6 @@ mod tests { } } - fn gauge_at(times: &[i64]) -> ResourceMetrics { - ResourceMetrics { - scope_metrics: vec![ScopeMetrics { - metrics: vec![Metric { - data: Some(metric::Data::Gauge(Gauge { - data_points: times - .iter() - .map(|ts| NumberDataPoint { - time_unix_nano: *ts as u64, - ..Default::default() - }) - .collect(), - })), - ..Default::default() - }], - ..Default::default() - }], - ..Default::default() - } - } - #[test] fn observe_projects_allowlist_and_dedups_per_request() { let mut data = ResourceInfoData::default(); @@ -275,7 +204,7 @@ mod tests { kv("k8s.node.name", "node-a"), kv("os.type", "linux"), ]; - data.observe(&attrs, &gauge_at(&[100, 50])); + data.observe_times(&attrs, [100, 50]); assert_eq!(data.rows.len(), 1); let (tags, windows) = data.rows.iter().next().unwrap(); assert_eq!(windows.values().copied().collect::>(), vec![100]); @@ -288,11 +217,11 @@ mod tests { .all(|(k, _)| k != "os.type" && k != "service.instance.id") ); - data.observe(&[kv("host.id", "h-2")], &gauge_at(&[10])); + data.observe_times(&[kv("host.id", "h-2")], [10]); assert_eq!(data.rows.len(), 2); let mut empty = ResourceInfoData::default(); - empty.observe(&[kv("os.type", "linux")], &gauge_at(&[100])); + empty.observe_times(&[kv("os.type", "linux")], [100]); assert!(empty.into_row_insert_requests().unwrap().is_none()); } @@ -301,9 +230,9 @@ mod tests { fn observe_keeps_one_row_per_graph_window() { let window = SEMANTIC_GRAPH_WINDOW_NANOS; let mut data = ResourceInfoData::default(); - data.observe( + data.observe_times( &[kv("service.name", "api")], - &gauge_at(&[window + 1, window + 2, 3 * window + 7]), + [window + 1, window + 2, 3 * window + 7], ); let windows = data.rows.values().next().unwrap(); @@ -320,6 +249,7 @@ mod tests { let exponential = |temporality: AggregationTemporality| ResourceMetrics { scope_metrics: vec![ScopeMetrics { metrics: vec![Metric { + name: "latency".into(), data: Some(metric::Data::ExponentialHistogram(ExponentialHistogram { data_points: vec![ExponentialHistogramDataPoint { time_unix_nano: 100, @@ -337,10 +267,21 @@ mod tests { AggregationTemporality::Delta, AggregationTemporality::Unspecified, ] { - let resource = exponential(temporality); - let mut data = ResourceInfoData::default(); - data.observe(&[kv("service.name", "api")], &resource); - assert!(data.into_row_insert_requests().unwrap().is_none()); + let mut resource = exponential(temporality); + resource.resource = Some( + otel_arrow_rust::proto::opentelemetry::resource::v1::Resource { + attributes: vec![kv("service.name", "api")], + ..Default::default() + }, + ); + let conversion = crate::otlp::metrics::to_grpc_insert_requests( + otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest { resource_metrics: vec![resource] }, + &mut session::protocol_ctx::OtlpMetricCtx { resource_info: true, ..Default::default() }, + ).unwrap(); + assert_eq!( + conversion.resource_info.is_none(), + temporality == AggregationTemporality::Unspecified + ); } } @@ -348,9 +289,9 @@ mod tests { fn rows_carry_raw_key_tags_value_and_millis_timestamp() { set_default_prefix(None).unwrap(); let mut data = ResourceInfoData::default(); - data.observe( + data.observe_times( &[kv("service.name", "api"), kv("host.id", "h-1")], - &gauge_at(&[1_700_000_000_123_456_789]), + [1_700_000_000_123_456_789], ); let requests = data.into_row_insert_requests().unwrap().unwrap(); assert_eq!(requests.inserts.len(), 1); diff --git a/src/servers/src/otlp/metrics/resource_info/tests/delta.rs b/src/servers/src/otlp/metrics/resource_info/tests/delta.rs index eb1de565c2f..d118d744d34 100644 --- a/src/servers/src/otlp/metrics/resource_info/tests/delta.rs +++ b/src/servers/src/otlp/metrics/resource_info/tests/delta.rs @@ -22,8 +22,15 @@ use super::*; fn observe_ignores_rejected_classic_histogram_windows() { let window = SEMANTIC_GRAPH_WINDOW_NANOS; let resource = ResourceMetrics { + resource: Some( + otel_arrow_rust::proto::opentelemetry::resource::v1::Resource { + attributes: vec![kv("service.name", "api")], + ..Default::default() + }, + ), scope_metrics: vec![ScopeMetrics { metrics: vec![Metric { + name: "latency".into(), data: Some(metric::Data::Histogram(Histogram { data_points: vec![ HistogramDataPoint { @@ -55,12 +62,34 @@ fn observe_ignores_rejected_classic_histogram_windows() { }], ..Default::default() }; - let mut data = ResourceInfoData::default(); - data.observe(&[kv("service.name", "api")], &resource); - - let windows = data.rows.values().next().unwrap(); + let conversion = crate::otlp::metrics::to_grpc_insert_requests( + otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest { + resource_metrics: vec![resource], + }, + &mut session::protocol_ctx::OtlpMetricCtx { resource_info: true, ..Default::default() }, + ).unwrap(); assert_eq!( - vec![window + 1, 4 * window + 1], - windows.values().copied().collect::>() + ( + conversion.outcome.accepted_data_points, + conversion.outcome.rejected_data_points + ), + (2, 1) + ); + let descriptors = conversion.resource_info.unwrap(); + let rows = descriptors.inserts[0].rows.as_ref().unwrap(); + let timestamp = rows + .schema + .iter() + .position(|c| c.column_name == greptime_timestamp()) + .unwrap(); + assert_eq!( + vec![ + Some(ValueData::TimestampMillisecondValue(window / 1_000_000)), + Some(ValueData::TimestampMillisecondValue(4 * window / 1_000_000)) + ], + rows.rows + .iter() + .map(|r| r.values[timestamp].value_data.clone()) + .collect::>() ); } diff --git a/src/servers/src/otlp/metrics/tests/delta.rs b/src/servers/src/otlp/metrics/tests/delta.rs index cc3e5b04cf7..97a6ca991ed 100644 --- a/src/servers/src/otlp/metrics/tests/delta.rs +++ b/src/servers/src/otlp/metrics/tests/delta.rs @@ -320,15 +320,16 @@ fn test_reserved_temporality_label_uses_final_persisted_key() { })), ..Default::default() }; - let error = to_grpc_insert_requests( + let conversion = to_grpc_insert_requests( metrics_request(vec![gauge(vec![keyvalue( OTLP_AGGREGATION_TEMPORALITY_LABEL, "user", )])]), &mut OtlpMetricCtx::default(), ) - .unwrap_err(); - assert!(matches!(error, error::Error::InvalidOtlpMetricInput { .. })); + .unwrap(); + assert_eq!(conversion.outcome.rejected_data_points, 1); + assert_eq!(conversion.outcome.accepted_data_points, 0); let mut request = metrics_request(vec![gauge(vec![])]); request.resource_metrics[0].scope_metrics[0].scope = Some(InstrumentationScope { @@ -392,3 +393,641 @@ fn test_histogram_semantics_follow_only_emitted_rows() { ); } } + +#[test] +fn test_shared_metric_metadata_follows_surviving_outputs() { + let mut metric = histogram_metric("latency"); + let Some(metric::Data::Histogram(histogram)) = metric.data.as_mut() else { + unreachable!() + }; + let with_sum = histogram.data_points[0].clone(); + histogram.data_points[0].sum = None; + histogram.data_points.push(with_sum); + for (rejections, has_sum) in [ + (vec![None, Some("rejected sum".into())], false), + (vec![Some("rejected first point".into()), None], true), + (vec![None, None], true), + ] { + let prepared = prepare_metrics( + metrics_request(vec![metric.clone()]), + &mut OtlpMetricCtx::default(), + ) + .unwrap(); + let conversion = prepared.finish(rejections, &Default::default()).unwrap(); + let semantics = decode(&conversion.semantic_index); + assert_eq!(semantics.contains_key("latency_sum"), has_sum); + for table in ["latency_count", "latency_bucket"] { + assert_eq!( + semantics[table][SEMANTIC_METRIC_TEMPORALITY], + METRIC_TEMPORALITY_CUMULATIVE + ); + } + } +} + +#[test] +fn test_delta_exponential_interval_encoding() { + use common_query::native_histogram::{COUNT_I64_FIELD, RESET_HINT_FIELD, SUM_FIELD}; + use common_query::prelude::greptime_native_histogram; + + for with_metric_engine in [false, true] { + let count = (1u64 << 53) + 1; + let point = ExponentialHistogramDataPoint { + start_time_unix_nano: 1_000_000_123, + time_unix_nano: 2_000_000_987, + count, + zero_count: count, + sum: Some(12.0), + zero_threshold: 0.1, + ..Default::default() + }; + let mut stale = point.clone(); + stale.flags = DataPointFlags::NoRecordedValueMask as u32; + let conversion = to_grpc_insert_requests( + metrics_request(vec![exponential_metric( + "interval", + vec![point, stale], + AggregationTemporality::Delta, + )]), + &mut OtlpMetricCtx { + with_metric_engine, + ..Default::default() + }, + ) + .unwrap(); + assert_eq!(conversion.outcome.accepted_data_points, 2); + assert_eq!(conversion.outcome.rejected_data_points, 0); + let rows = table_rows(&conversion.requests, "interval"); + let field = column_index(rows, greptime_native_histogram()); + let tag = column_index(rows, OTLP_AGGREGATION_TEMPORALITY_LABEL); + let ts = column_index(rows, greptime_timestamp()); + assert_eq!( + rows.schema[ts].datatype, + if with_metric_engine { + ColumnDataType::TimestampNanosecond + } else { + ColumnDataType::TimestampMillisecond + } as i32 + ); + for row in &rows.rows { + assert_eq!( + row.values[tag].value_data, + Some(ValueData::StringValue("delta".into())) + ); + assert_eq!( + native_field( + row.values[field].value_data.as_ref().unwrap(), + RESET_HINT_FIELD + ), + Some(ValueData::I32Value(ResetHint::Gauge as i32)) + ); + } + assert_eq!( + native_field( + rows.rows[0].values[field].value_data.as_ref().unwrap(), + COUNT_I64_FIELD + ), + Some(ValueData::I64Value(count as i64)) + ); + let Some(ValueData::F64Value(stale)) = native_field( + rows.rows[1].values[field].value_data.as_ref().unwrap(), + SUM_FIELD, + ) else { + panic!("missing sum") + }; + assert_eq!(stale.to_bits(), PROMETHEUS_STALE_NAN_BITS); + } +} + +fn gauge(name: &str) -> Metric { + Metric { + name: name.into(), + data: Some(metric::Data::Gauge(Gauge { + data_points: vec![NumberDataPoint { + value: Some(Value::AsDouble(7.0)), + ..Default::default() + }], + })), + ..Default::default() + } +} + +#[test] +fn test_shared_schemas_preserve_source_order_and_omissions() { + let mut metric = gauge("ordered"); + let Some(metric::Data::Gauge(data)) = &mut metric.data else { + unreachable!() + }; + data.data_points = [ + keyvalue("a", "first"), + keyvalue("b", "second"), + keyvalue("a", "third"), + KeyValue { + key: "a".into(), + value: None, + }, + keyvalue("a", "fifth"), + keyvalue("a", "sixth"), + ] + .into_iter() + .enumerate() + .map(|(index, attribute)| NumberDataPoint { + attributes: vec![attribute], + time_unix_nano: (index as u64 + 1) * 1_000_000, + value: (index != 5).then_some(Value::AsDouble(index as f64)), + ..Default::default() + }) + .collect(); + for reject_first in [false, true] { + let prepared = prepare_metrics( + metrics_request(vec![metric.clone()]), + &mut OtlpMetricCtx::default(), + ) + .unwrap(); + assert_eq!(prepared.num_points(), 6); + // The fourth point omits a, even though its raw attribute key is unchanged. + assert_eq!(prepared.requests.len(), 4); + assert!( + prepared + .requests + .last() + .unwrap() + .rows + .as_ref() + .unwrap() + .schema + .iter() + .all(|column| column.column_name != greptime_value()) + ); + let mut rejections = vec![None; 6]; + if reject_first { + rejections[0] = Some("rejected first point".into()); + } + let conversion = prepared.finish(rejections, &Default::default()).unwrap(); + let rows = table_rows(&conversion.requests, "ordered"); + let a = column_index(rows, "a"); + let b = column_index(rows, "b"); + assert_eq!(b < a, reject_first); + let value = column_index(rows, greptime_value()); + let timestamp = column_index(rows, greptime_timestamp()); + for (index, row) in (usize::from(reject_first)..6).zip(&rows.rows) { + assert_eq!( + row.values[value].value_data, + (index != 5).then_some(ValueData::F64Value(index as f64)) + ); + assert_eq!( + row.values[timestamp].value_data, + Some(ValueData::TimestampMillisecondValue(index as i64 + 1)) + ); + assert_eq!( + row.values[a].value_data.is_some(), + matches!(index, 0 | 2 | 4 | 5) + ); + assert_eq!(row.values[b].value_data.is_some(), index == 1); + } + assert_eq!(rows.rows.len(), 6 - usize::from(reject_first)); + } +} + +#[test] +fn test_reused_sum_layouts_preserve_omitted_and_stale_values() { + let points = [false, true, false, true, false] + .into_iter() + .enumerate() + .map(|(index, stale)| NumberDataPoint { + attributes: vec![keyvalue("host", "a")], + time_unix_nano: (index as u64 + 1) * 1_000_000, + flags: if stale { + DataPointFlags::NoRecordedValueMask as u32 + } else { + 0 + }, + ..Default::default() + }) + .collect(); + let metric = Metric { + name: "requests".into(), + data: Some(metric::Data::Sum(Sum { + data_points: points, + aggregation_temporality: AggregationTemporality::Delta as i32, + ..Default::default() + })), + ..Default::default() + }; + let prepared = + prepare_metrics(metrics_request(vec![metric]), &mut OtlpMetricCtx::default()).unwrap(); + assert_eq!(prepared.requests.len(), 2); + for request in &prepared.requests { + let rows = request.rows.as_ref().unwrap(); + let has_value = rows + .schema + .iter() + .any(|c| c.column_name == greptime_value()); + assert_eq!(rows.rows.len(), if has_value { 2 } else { 3 }); + assert!( + rows.rows + .iter() + .all(|row| row.values.len() == rows.schema.len()) + ); + } + let conversion = prepared.finish(vec![None; 5], &Default::default()).unwrap(); + assert_eq!(conversion.outcome.accepted_data_points, 5); + let rows = table_rows(&conversion.requests, "requests"); + let value = column_index(rows, greptime_value()); + for (index, row) in rows.rows.iter().enumerate() { + if index % 2 == 0 { + assert!(row.values[value].value_data.is_none()); + } else { + let Some(ValueData::F64Value(value)) = row.values[value].value_data else { + panic!("missing stale value") + }; + assert_eq!(value.to_bits(), PROMETHEUS_STALE_NAN_BITS); + } + } +} + +#[test] +fn test_unmerged_batches_preserve_normalized_rows_and_metadata() { + let mut first = gauge("same"); + first.unit = "seconds".into(); + let mut second = first.clone(); + if let Some(metric::Data::Gauge(data)) = second.data.as_mut() { + data.data_points[0].value = Some(Value::AsDouble(8.0)); + } + let mut prepared = prepare_metrics( + metrics_request(vec![first, gauge("other"), second]), + &mut OtlpMetricCtx::default(), + ) + .unwrap(); + assert_eq!(prepared.requests.len(), 2); + for request in &mut prepared.requests { + let rows = request.rows.as_mut().unwrap(); + let timestamp = column_index(rows, greptime_timestamp()); + rows.schema[timestamp].datatype = ColumnDataType::TimestampMicrosecond as i32; + for row in &mut rows.rows { + row.values[timestamp].value_data = Some(ValueData::TimestampMicrosecondValue(123)); + } + } + let expected = prepared.requests.clone(); + let conversion = prepared.finish(vec![None; 3], &Default::default()).unwrap(); + assert_eq!(conversion.requests.inserts, expected); + assert_eq!(conversion.rows, 3); + assert_eq!(conversion.outcome.accepted_data_points, 3); + let semantics = decode(&conversion.semantic_index); + for request in &conversion.requests.inserts { + assert_eq!( + semantics[&request.table_name][SEMANTIC_METRIC_TYPE], + METRIC_TYPE_GAUGE + ); + } +} + +#[test] +fn test_shared_batch_row_error_rejects_complete_histogram_point() { + let mut metric = histogram_metric("latency"); + let Some(metric::Data::Histogram(data)) = &mut metric.data else { + unreachable!() + }; + data.data_points.push(data.data_points[0].clone()); + let prepared = + prepare_metrics(metrics_request(vec![metric]), &mut OtlpMetricCtx::default()).unwrap(); + let mut errors: Vec<_> = prepared + .requests + .iter() + .map(|request| vec![None; request.rows.as_ref().unwrap().rows.len()]) + .collect(); + let bucket = prepared + .requests + .iter() + .position(|r| r.table_name == "latency_bucket") + .unwrap(); + let second_point = errors[bucket].len() / 2; + errors[bucket][second_point] = Some("invalid bucket".into()); + assert!(prepared.point_rejections(Vec::new()).is_err()); + let rejections = prepared.point_rejections(errors).unwrap(); + assert_eq!(rejections, vec![None, Some("invalid bucket".into())]); + let conversion = prepared.finish(rejections, &Default::default()).unwrap(); + assert_eq!(conversion.outcome.accepted_data_points, 1); + assert_eq!(conversion.outcome.rejected_data_points, 1); + for table in ["latency_count", "latency_sum", "latency_bucket"] { + assert_eq!(table_rows(&conversion.requests, table).rows.len(), 1); + } +} + +#[test] +fn test_gauge_ignores_start_time_but_validates_timestamp() { + for with_metric_engine in [false, true] { + let mut ctx = OtlpMetricCtx { + with_metric_engine, + ..Default::default() + }; + let point = NumberDataPoint { + time_unix_nano: 60_000_000_000, + value: Some(Value::AsDouble(7.0)), + ..Default::default() + }; + let mut metric = gauge("gauge"); + let Some(metric::Data::Gauge(data)) = &mut metric.data else { + unreachable!() + }; + data.data_points = vec![point.clone()]; + let baseline = + to_grpc_insert_requests(metrics_request(vec![metric.clone()]), &mut ctx).unwrap(); + assert_eq!(baseline.outcome.accepted_data_points, 1); + + for (start, timestamp, accepted) in [ + (120_000_000_000, point.time_unix_nano, 1), + (u64::MAX, point.time_unix_nano, 1), + (0, i64::MAX as u64 + 1, 0), + ] { + let Some(metric::Data::Gauge(data)) = &mut metric.data else { + unreachable!() + }; + data.data_points[0].start_time_unix_nano = start; + data.data_points[0].time_unix_nano = timestamp; + let conversion = + to_grpc_insert_requests(metrics_request(vec![metric.clone()]), &mut ctx).unwrap(); + assert_eq!(conversion.outcome.accepted_data_points, accepted); + assert_eq!(conversion.outcome.rejected_data_points, 1 - accepted); + if accepted == 1 { + assert_eq!(conversion.requests, baseline.requests); + } else { + assert!(conversion.requests.inserts.is_empty()); + } + } + + let sum = Metric { + name: "sum".into(), + data: Some(metric::Data::Sum(Sum { + data_points: vec![NumberDataPoint { + start_time_unix_nano: 120_000_000_000, + ..point + }], + aggregation_temporality: AggregationTemporality::Delta as i32, + ..Default::default() + })), + ..Default::default() + }; + let conversion = to_grpc_insert_requests(metrics_request(vec![sum]), &mut ctx).unwrap(); + assert_eq!(conversion.outcome.accepted_data_points, 0); + assert_eq!(conversion.outcome.rejected_data_points, 1); + assert!(conversion.requests.inserts.is_empty()); + } +} + +#[test] +fn test_all_metric_families_isolate_invalid_attributes_and_fanout() { + let families = [ + gauge("gauge"), + Metric { + name: "sum".into(), + data: Some(metric::Data::Sum(Sum { + data_points: vec![NumberDataPoint::default()], + ..Default::default() + })), + ..Default::default() + }, + histogram_metric("histogram"), + Metric { + name: "summary".into(), + data: Some(metric::Data::Summary(Summary { + data_points: vec![SummaryDataPoint { + count: 1, + sum: 1.0, + ..Default::default() + }], + })), + ..Default::default() + }, + exponential_metric( + "native", + vec![exponential_point()], + AggregationTemporality::Delta, + ), + ]; + for family in families { + for invalid_first in [false, true] { + let mut metric = family.clone(); + macro_rules! add_invalid { + ($data:ident) => {{ + let mut invalid = $data.data_points[0].clone(); + invalid.attributes = vec![ + keyvalue("rejected_only", "bad"), + keyvalue(OTLP_AGGREGATION_TEMPORALITY_LABEL, "bad"), + ]; + if invalid_first { + $data.data_points.insert(0, invalid); + } else { + $data.data_points.push(invalid); + } + }}; + } + match metric.data.as_mut().unwrap() { + metric::Data::Gauge(data) => add_invalid!(data), + metric::Data::Sum(data) => add_invalid!(data), + metric::Data::Histogram(data) => add_invalid!(data), + metric::Data::Summary(data) => add_invalid!(data), + metric::Data::ExponentialHistogram(data) => add_invalid!(data), + } + let conversion = to_grpc_insert_requests( + metrics_request(vec![metric, gauge("independent")]), + &mut OtlpMetricCtx::default(), + ) + .unwrap(); + assert_eq!( + ( + conversion.outcome.accepted_data_points, + conversion.outcome.rejected_data_points + ), + (2, 1), + "{}", + family.name + ); + assert!(conversion.requests.inserts.iter().all(|r| { + r.rows + .as_ref() + .unwrap() + .schema + .iter() + .all(|c| c.column_name != "rejected_only") + })); + } + } +} + +#[test] +fn test_fresh_conflicts_are_simultaneous_and_authoritative_rejection_wins() { + let metrics = [ + histogram_metric("latency"), + exponential_metric( + "latency_bucket", + vec![exponential_point()], + AggregationTemporality::Delta, + ), + exponential_metric( + "latency_count", + vec![exponential_point()], + AggregationTemporality::Delta, + ), + ]; + for order in [ + [0, 1, 2], + [0, 2, 1], + [1, 0, 2], + [1, 2, 0], + [2, 0, 1], + [2, 1, 0], + ] { + let mut input = order + .iter() + .map(|&i| metrics[i].clone()) + .collect::>(); + input.push(gauge("independent")); + let conversion = to_grpc_insert_requests( + metrics_request(input.clone()), + &mut OtlpMetricCtx::default(), + ) + .unwrap(); + assert_eq!( + ( + conversion.outcome.accepted_data_points, + conversion.outcome.rejected_data_points + ), + (1, 3) + ); + assert_eq!(conversion.requests.inserts.len(), 1); + assert_eq!(conversion.requests.inserts[0].table_name, "independent"); + + let prepared = + prepare_metrics(metrics_request(input), &mut OtlpMetricCtx::default()).unwrap(); + let mut rejections = vec![None; 4]; + rejections[order.iter().position(|i| *i == 0).unwrap()] = + Some("incompatible existing latency_sum".into()); + let conversion = prepared + .finish(rejections, &std::collections::HashMap::new()) + .unwrap(); + assert_eq!( + ( + conversion.outcome.accepted_data_points, + conversion.outcome.rejected_data_points + ), + (3, 1) + ); + assert!(!decode(&conversion.semantic_index).contains_key("latency_sum")); + } +} + +#[test] +fn test_semantics_only_describe_actual_outputs() { + let mut histogram = histogram_metric("latency"); + if let Some(metric::Data::Histogram(data)) = histogram.data.as_mut() { + data.data_points[0].sum = None; + } + let summary = Metric { + name: "summary".into(), + data: Some(metric::Data::Summary(Summary { + data_points: vec![SummaryDataPoint::default()], + })), + ..Default::default() + }; + for reverse in [false, true] { + let mut input = vec![ + histogram.clone(), + gauge("latency_sum"), + summary.clone(), + gauge("summary"), + ]; + if reverse { + input.reverse(); + } + let conversion = + to_grpc_insert_requests(metrics_request(input), &mut OtlpMetricCtx::default()).unwrap(); + for name in ["latency_sum", "summary"] { + let semantics = decode(&conversion.semantic_index); + assert_eq!(semantics[name][SEMANTIC_METRIC_TYPE], METRIC_TYPE_GAUGE); + assert!(!semantics[name].contains_key(SEMANTIC_METRIC_TEMPORALITY)); + } + } +} + +#[test] +fn test_invalid_emitted_names_and_valid_empty_base_fanout() { + for reverse in [false, true] { + let mut metrics = vec![gauge("!!!"), histogram_metric("!!!"), gauge("valid")]; + if reverse { + metrics.reverse(); + } + let conversion = + to_grpc_insert_requests(metrics_request(metrics), &mut OtlpMetricCtx::default()) + .unwrap(); + assert_eq!( + ( + conversion.outcome.accepted_data_points, + conversion.outcome.rejected_data_points + ), + (2, 1) + ); + assert!( + conversion + .requests + .inserts + .iter() + .all(|r| !r.table_name.is_empty()) + ); + assert!( + conversion + .requests + .inserts + .iter() + .any(|r| r.table_name == "_bucket") + ); + } +} + +#[test] +fn test_shared_physical_conflicts_and_independent_backings() { + use common_query::prelude::greptime_native_histogram; + for reverse in [false, true] { + let mut scalar = gauge("scalar"); + if let Some(metric::Data::Gauge(data)) = scalar.data.as_mut() { + data.data_points[0] + .attributes + .push(keyvalue(greptime_native_histogram(), "tag")); + } + let mut metrics = vec![ + scalar, + exponential_metric( + "native", + vec![exponential_point()], + AggregationTemporality::Delta, + ), + gauge("independent"), + ]; + if reverse { + metrics.reverse(); + } + for shared in [false, true] { + let prepared = prepare_metrics( + metrics_request(metrics.clone()), + &mut OtlpMetricCtx::default(), + ) + .unwrap(); + let backing = std::collections::HashMap::from([ + ("scalar".into(), "physical".into()), + ( + "native".into(), + if shared { "physical" } else { "other" }.into(), + ), + ("independent".into(), "physical".into()), + ]); + let conversion = prepared.finish(vec![None; 3], &backing).unwrap(); + assert_eq!( + ( + conversion.outcome.accepted_data_points, + conversion.outcome.rejected_data_points + ), + if shared { (1, 2) } else { (3, 0) } + ); + } + } +} diff --git a/src/servers/src/otlp/metrics/translator.rs b/src/servers/src/otlp/metrics/translator.rs index 8b3b45d3274..79c00f6ca97 100644 --- a/src/servers/src/otlp/metrics/translator.rs +++ b/src/servers/src/otlp/metrics/translator.rs @@ -286,16 +286,6 @@ pub fn normalize_label_name(name: &str) -> String { n.to_string() } -/// Normalize otlp instrumentation, metric and attribute names -/// -/// -/// - since the name are case-insensitive, we transform them to lowercase for -/// better sql usability -/// - replace `.` and `-` with `_` -pub fn legacy_normalize_otlp_name(name: &str) -> String { - name.to_lowercase().replace(['.', '-'], "_") -} - #[cfg(test)] mod tests { use otel_arrow_rust::proto::opentelemetry::metrics::v1::Metric; @@ -306,30 +296,6 @@ mod tests { use super::*; - #[test] - fn test_legacy_normalize_otlp_name() { - assert_eq!( - legacy_normalize_otlp_name("jvm.memory.free"), - "jvm_memory_free" - ); - assert_eq!( - legacy_normalize_otlp_name("jvm-memory-free"), - "jvm_memory_free" - ); - assert_eq!( - legacy_normalize_otlp_name("jvm_memory_free"), - "jvm_memory_free" - ); - assert_eq!( - legacy_normalize_otlp_name("JVM_MEMORY_FREE"), - "jvm_memory_free" - ); - assert_eq!( - legacy_normalize_otlp_name("JVM_memory_FREE"), - "jvm_memory_free" - ); - } - #[test] fn test_translate_metric_name_strategies() { let metric = Metric { diff --git a/src/servers/src/row_writer.rs b/src/servers/src/row_writer.rs index 0661e733510..780a19df3f8 100644 --- a/src/servers/src/row_writer.rs +++ b/src/servers/src/row_writer.rs @@ -131,6 +131,16 @@ impl TableData { (self.schema, self.rows) } + pub(crate) fn into_rows(self) -> Rows { + let (schema, mut rows) = self.into_schema_and_rows(); + for row in &mut rows { + if row.values.len() < schema.len() { + row.values.resize(schema.len(), Value::default()); + } + } + Rows { schema, rows } + } + /// Writes a field value without enforcing that later writes use the same datatype /// as the first-seen schema entry. /// @@ -207,6 +217,29 @@ impl MultiTableData { .insert(table_name.to_string(), table_data); } + pub(crate) fn iter(&self) -> impl Iterator { + self.table_data_map + .iter() + .map(|(name, table)| (name.as_str(), table)) + } + + /// Retains schemas and row capacity for the next point with the same attributes. + pub(crate) fn drain_rows( + &mut self, + ) -> impl Iterator)> { + self.table_data_map.iter_mut().filter_map(|(name, table)| { + if table.rows.is_empty() { + return None; + } + for row in &mut table.rows { + if row.values.len() < table.schema.len() { + row.values.resize(table.schema.len(), Value::default()); + } + } + Some((name.as_str(), table.schema.as_slice(), table.rows.drain(..))) + }) + } + #[allow(dead_code)] pub fn num_tables(&self) -> usize { self.table_data_map.len() @@ -220,17 +253,9 @@ impl MultiTableData { .into_iter() .map(|(table_name, table_data)| { total_rows += table_data.num_rows(); - let num_columns = table_data.num_columns(); - let (schema, mut rows) = table_data.into_schema_and_rows(); - for row in &mut rows { - if num_columns > row.values.len() { - row.values.resize(num_columns, Value { value_data: None }); - } - } - RowInsertRequest { table_name, - rows: Some(Rows { schema, rows }), + rows: Some(table_data.into_rows()), } }) .collect::>(); diff --git a/src/servers/src/semantic.rs b/src/servers/src/semantic.rs index 63501dffc31..3d27b4d80da 100644 --- a/src/servers/src/semantic.rs +++ b/src/servers/src/semantic.rs @@ -108,7 +108,7 @@ impl SemanticIndex { serde_json::to_string(&BTreeMap::from([(schema, &self.tables)])).ok() } - fn merge_from(&mut self, other: &SemanticIndex) { + pub(crate) fn merge_from(&mut self, other: &SemanticIndex) { for (table, scalars) in &other.tables { for (key, value) in scalars { self.record_scalar(table, key, value); @@ -116,6 +116,15 @@ impl SemanticIndex { } } + /// Moves a metric's metadata on its first accepted output for this table. + pub(crate) fn merge_table_from(&mut self, other: &mut SemanticIndex, table: &str) { + if let Some(scalars) = other.tables.remove(table) { + for (key, value) in scalars { + self.record_scalar(table, key, &value); + } + } + } + #[cfg(test)] fn options_of(&self, table: &str) -> Option<&BTreeMap<&'static str, String>> { self.tables.get(table) diff --git a/src/session/src/protocol_ctx.rs b/src/session/src/protocol_ctx.rs index 5165def3e49..b9bc33aa204 100644 --- a/src/session/src/protocol_ctx.rs +++ b/src/session/src/protocol_ctx.rs @@ -42,18 +42,14 @@ impl ProtocolCtx { /// If true, all scope attributes will be promoted to the final table schema. /// Along with the scope name, scope version and scope schema URL. /// - `with_metric_engine` -/// - `is_legacy` -/// If the user uses OTLP metrics ingestion before v0.16, it uses the old path. -/// So we call this path 'legacy'. -/// After v0.16, we store the OTLP metrics using prometheus compatible format, the new path. -/// The difference is how we convert the input data into the final table schema. +/// Use Metric Engine for auto-created tables and retain timestamp precision +/// until alignment with the destination table's time index. #[derive(Debug, Clone, Default)] pub struct OtlpMetricCtx { pub promote_all_resource_attrs: bool, pub resource_attrs: HashSet, pub promote_scope_attrs: bool, pub with_metric_engine: bool, - pub is_legacy: bool, /// Set from the server's `otlp.experimental_enable_resource_info`; off /// means the resource descriptor is not synthesized at all. pub resource_info: bool, diff --git a/src/table/src/requests.rs b/src/table/src/requests.rs index 81f396ec703..2674c2d0a1e 100644 --- a/src/table/src/requests.rs +++ b/src/table/src/requests.rs @@ -73,6 +73,7 @@ pub fn is_trace_table(table_info: &crate::metadata::TableInfo) -> bool { ) } +/// Retained table metadata; OTLP metrics always use Prometheus-compatible conversion. pub const OTLP_METRIC_COMPAT_KEY: &str = "otlp_metric_compat"; pub const OTLP_METRIC_COMPAT_PROM: &str = "prom"; diff --git a/tests-integration/src/instance.rs b/tests-integration/src/instance.rs index 9f48d01a467..140151d594f 100644 --- a/tests-integration/src/instance.rs +++ b/tests-integration/src/instance.rs @@ -776,10 +776,16 @@ mod tests { }], }; + let errors = frontend::OTLP_RESOURCE_INFO_WRITE_ERRORS.get(); let outcome = OpenTelemetryProtocolHandler::metrics(&*instance, request, QueryContext::arc()) .await .unwrap(); + assert_eq!( + (outcome.accepted_data_points, outcome.rejected_data_points), + (1, 0) + ); + assert_eq!(frontend::OTLP_RESOURCE_INFO_WRITE_ERRORS.get(), errors + 1); let warning = outcome .error_message .expect("denial must surface as a warning"); diff --git a/tests-integration/src/otlp.rs b/tests-integration/src/otlp.rs index 94fac669010..4998e29e0c8 100644 --- a/tests-integration/src/otlp.rs +++ b/tests-integration/src/otlp.rs @@ -140,9 +140,15 @@ mod test { ) -> common_meta::error::Result { let gate = { let mut gate = self.gate.lock().unwrap(); - if gate.as_ref().is_some_and(|gate| { - matches!(&request.body, Some(region_request::Body::Inserts(inserts)) - if inserts.requests.iter().any(|insert| insert.region_id == gate.region_id.as_u64())) + if gate.as_ref().is_some_and(|gate| match &request.body { + Some(region_request::Body::Inserts(inserts)) => inserts + .requests + .iter() + .any(|insert| insert.region_id == gate.region_id.as_u64()), + Some(region_request::Body::BulkInsert(insert)) => { + insert.region_id == gate.region_id.as_u64() + } + _ => false, }) { gate.take() } else { @@ -1561,7 +1567,9 @@ WITH( Arc::new(fixed) ) .await - .is_err() + .unwrap() + .rejected_data_points + == 1 ); } assert_eq!(results[..3], results[3..]); @@ -1604,14 +1612,21 @@ WITH( let result = output.remove(0); assert!(result.is_ok(), "{result:?}"); - let error = instance + let outcome = instance .metrics( build_sum_request("fixed.delta", AggregationTemporality::Delta, &[(60, 10)]), ctx.clone(), ) .await - .unwrap_err(); - assert!(format!("{error:?}").contains("otlp_aggregation_temporality")); + .unwrap(); + assert_eq!(outcome.accepted_data_points, 0); + assert_eq!(outcome.rejected_data_points, 1); + assert!( + outcome + .error_message + .unwrap() + .contains("otlp_aggregation_temporality") + ); let mut output = instance .do_query("SELECT COUNT(*) FROM fixed_delta_total", ctx) @@ -2229,6 +2244,1094 @@ WITH( } } + fn metric_request(metrics: Vec) -> ExportMetricsServiceRequest { + ExportMetricsServiceRequest { + resource_metrics: vec![ResourceMetrics { + scope_metrics: vec![ScopeMetrics { + metrics, + ..Default::default() + }], + ..Default::default() + }], + } + } + + fn gauge_metric(name: &str, points: &[(u64, f64, Option<&str>)]) -> Metric { + Metric { + name: name.into(), + data: Some(metric::Data::Gauge(Gauge { + data_points: points + .iter() + .map(|&(time, value, host)| NumberDataPoint { + time_unix_nano: time, + value: Some(Value::AsDouble(value)), + attributes: host + .map(|host| vec![keyvalue("host", host)]) + .unwrap_or_default(), + ..Default::default() + }) + .collect(), + })), + ..Default::default() + } + } + + fn native_metric(name: &str, scale: i32) -> Metric { + use otel_arrow_rust::proto::opentelemetry::metrics::v1::{ + ExponentialHistogram, ExponentialHistogramDataPoint, + }; + Metric { + name: name.into(), + data: Some(metric::Data::ExponentialHistogram(ExponentialHistogram { + aggregation_temporality: AggregationTemporality::Delta as i32, + data_points: vec![ExponentialHistogramDataPoint { + time_unix_nano: 60_000_000_000, + count: 10, + zero_count: 10, + scale, + sum: Some(0.0), + ..Default::default() + }], + })), + ..Default::default() + } + } + + #[test] + fn test_otlp_primary_and_descriptor_write_failures() { + for sync in [false, true] { + temp_env::with_var("PENDING_ROWS_BATCH_SYNC", Some(sync.to_string()), || { + tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .unwrap() + .block_on(async { + tokio::time::timeout( + Duration::from_secs(60), + check_metric_write_failure(sync), + ) + .await + .expect("metric acknowledgement must make progress"); + }); + }); + } + } + + async fn check_metric_write_failure(sync: bool) { + use prost::Message; + + let standalone = GreptimeDbStandaloneBuilder::new("otlp_write_failure") + .with_logical_batcher(BatcherOptions { + protocols: vec![BatchingProtocol::Otlp], + max_batch_rows: 1, + pending_rows_flush_interval: Duration::from_secs(3600), + ..Default::default() + }) + .build() + .await; + let source = standalone.fe_instance(); + let gate = Arc::new(Mutex::new(None)); + let cache_registry = Arc::new( + cache::with_default_composite_cache_registry( + LayeredCacheRegistryBuilder::default().add_cache_registry( + cache::build_fundamental_cache_registry(standalone.kv_backend.clone()), + ), + ) + .unwrap() + .build(), + ); + let instance = Arc::new( + FrontendBuilder::new( + standalone.opts.frontend_options(), + standalone.kv_backend.clone(), + cache_registry, + source.catalog_manager().clone(), + Arc::new(GatedTraceNodeManager { + inner: source.node_manager().clone(), + gate: gate.clone(), + }), + source.procedure_executor().clone(), + Arc::new(catalog::process_manager::ProcessManager::new( + "otlp_write_failure".into(), + None, + )), + ) + .try_build() + .await + .unwrap(), + ); + instance.init_logical_batcher(&standalone.opts.frontend_options()); + let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, "public"); + ctx.set_protocol_ctx(ProtocolCtx::OtlpMetric(OtlpMetricCtx { + with_metric_engine: true, + ..Default::default() + })); + let mut warm = metric_request(vec![gauge_metric( + "ack_metric", + &[(60_000_000_000, 1.0, None)], + )]); + warm.resource_metrics[0].resource = Some(Resource { + attributes: vec![ + keyvalue("service.name", "ack-service"), + keyvalue("host.id", "ack-host"), + ], + ..Default::default() + }); + instance + .metrics(warm.clone(), Arc::new(ctx.clone())) + .await + .unwrap(); + let region = async |name: &str| { + let table = instance + .catalog_manager() + .table(DEFAULT_CATALOG_NAME, "public", name, None) + .await + .unwrap() + .unwrap(); + instance + .partition_manager() + .find_physical_table_route(table.table_info().table_id()) + .await + .unwrap() + .region_routes[0] + .region + .id + }; + let (started_tx, started_rx) = oneshot::channel(); + let (release_tx, release_rx) = oneshot::channel(); + *gate.lock().unwrap() = Some(TraceAuxWriteGate { + region_id: region("ack_metric").await, + started: started_tx, + release: release_rx, + }); + ctx.set_logical_batching_enabled(true); + let request = metric_request(vec![gauge_metric( + "ack_metric", + &[(120_000_000_000, 2.0, None)], + )]); + let failures = servers::metrics::FLUSH_FAILURES.get(); + let dropped = servers::metrics::FLUSH_DROPPED_ROWS.get(); + let write = tokio::spawn({ + let instance = instance.clone(); + let ctx = Arc::new(ctx.clone()); + async move { instance.metrics(request, ctx).await } + }); + tokio::time::timeout(Duration::from_secs(10), started_rx) + .await + .expect("primary batch must reach the write gate") + .unwrap(); + if sync { + assert!(!write.is_finished()); + release_tx.send(false).unwrap(); + assert!(write.await.unwrap().is_err()); + } else { + let accepted = tokio::time::timeout(Duration::from_secs(10), write) + .await + .expect("asynchronous acknowledgement must precede physical completion") + .unwrap() + .unwrap(); + assert_eq!( + (accepted.accepted_data_points, accepted.rejected_data_points), + (1, 0) + ); + release_tx.send(false).unwrap(); + while servers::metrics::FLUSH_FAILURES.get() == failures + || servers::metrics::FLUSH_DROPPED_ROWS.get() == dropped + { + tokio::task::yield_now().await; + } + } + assert_eq!(servers::metrics::FLUSH_FAILURES.get(), failures + 1); + assert_eq!(servers::metrics::FLUSH_DROPPED_ROWS.get(), dropped + 1); + assert_trace_v2_query( + &instance, + "SELECT COUNT(*)=1 AND MAX(greptime_value)=1 FROM ack_metric", + Arc::new(ctx), + ) + .await + .unwrap(); + + // The derived write fails after the metric was accepted, including when + // primary logical batching acknowledges before its flush completes. + let (started_tx, started_rx) = oneshot::channel(); + let (release_tx, release_rx) = oneshot::channel(); + *gate.lock().unwrap() = Some(TraceAuxWriteGate { + region_id: region("greptime_otel_resource_info").await, + started: started_tx, + release: release_rx, + }); + let options = standalone.opts.clone(); + let services = frontend::server::Services::new( + options.clone(), + instance.clone(), + common_base::Plugins::default(), + ); + let server = services + .http_server_builder( + &options.frontend_options(), + services.server_memory_limiter.clone(), + ) + .build(); + let client = + servers::http::test_helpers::TestClient::new(server.build(server.make_app()).unwrap()) + .await; + let errors = frontend::OTLP_RESOURCE_INFO_WRITE_ERRORS.get(); + let accepted_rows = frontend::OTLP_METRICS_ROWS.get(); + let response = tokio::spawn(async move { + client + .post("/v1/otlp/v1/metrics") + .header("content-type", "application/x-protobuf") + .body(warm.encode_to_vec()) + .send() + .await + }); + started_rx.await.unwrap(); + release_tx.send(false).unwrap(); + let response = response.await.unwrap(); + assert_eq!(response.status().as_u16(), 200); + let bytes = response.bytes().await; + let response = otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceResponse::decode(bytes).unwrap(); + let partial = response.partial_success.unwrap(); + assert_eq!(partial.rejected_data_points, 0); + assert!(partial.error_message.contains("descriptor")); + assert_eq!(frontend::OTLP_RESOURCE_INFO_WRITE_ERRORS.get(), errors + 1); + assert_eq!(frontend::OTLP_METRICS_ROWS.get(), accepted_rows + 1); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_otlp_metric_permission_precedes_preflight() { + use common_error::ext::ErrorExt; + use common_error::status_code::StatusCode; + + #[derive(Default)] + struct DenyRestricted { + calls: Mutex>>, + } + + impl auth::PermissionChecker for DenyRestricted { + fn check_permission( + &self, + _user_info: auth::UserInfoRef, + _req: auth::PermissionReq, + ) -> auth::error::Result { + Ok(auth::PermissionResp::Allow) + } + + fn check_permission_with_table_targets( + &self, + _user_info: auth::UserInfoRef, + req: auth::PermissionReq, + targets: auth::PermissionTableTargets, + ) -> auth::error::Result { + if matches!(req, auth::PermissionReq::Action(auth::OTLP_WRITE)) { + let auth::PermissionTableTargets::Resolved(targets) = targets else { + panic!("metric destinations must be resolved"); + }; + let denied = targets.iter().any(|t| t.table == "restricted_metric"); + self.calls + .lock() + .unwrap() + .push(targets.into_iter().map(|t| t.table).collect()); + if denied { + return Ok(auth::PermissionResp::Reject); + } + } + Ok(auth::PermissionResp::Allow) + } + } + + let checker = Arc::new(DenyRestricted::default()); + let plugins = common_base::Plugins::new(); + plugins.insert::(checker.clone()); + let standalone = GreptimeDbStandaloneBuilder::new("otlp_permission_preflight") + .with_plugin(plugins) + .build() + .await; + let instance = standalone.fe_instance(); + let ctx = QueryContext::arc(); + instance + .do_query( + "CREATE TABLE restricted_metric (greptime_timestamp TIMESTAMP(3) TIME INDEX, \ + greptime_value DOUBLE, host STRING NOT NULL PRIMARY KEY) \ + WITH ('otlp_metric_compat'='prom')", + ctx.clone(), + ) + .await + .remove(0) + .unwrap(); + + for host in [Some("present"), None] { + for with_allowed in [false, true] { + let mut metrics = vec![gauge_metric( + "restricted_metric", + &[(60_000_000_000, 1.0, host), (120_000_000_000, 2.0, host)], + )]; + let mut expected_targets = vec!["restricted_metric"]; + if with_allowed { + metrics.push(gauge_metric( + "allowed_metric", + &[(60_000_000_000, 7.0, None)], + )); + expected_targets.push("allowed_metric"); + } + let err = instance + .metrics(metric_request(metrics), ctx.clone()) + .await + .unwrap_err(); + assert_eq!(err.status_code(), StatusCode::PermissionDenied); + assert!(!err.to_string().contains("missing required column")); + assert_eq!( + std::mem::take(&mut *checker.calls.lock().unwrap()), + vec![expected_targets] + ); + } + } + assert_trace_v2_query( + instance, + "SELECT COUNT(*)=0 FROM information_schema.tables WHERE table_name='allowed_metric'", + ctx, + ) + .await + .unwrap(); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_otlp_preflight_defaults_order_and_required_columns() { + let standalone = GreptimeDbStandaloneBuilder::new("otlp_preflight_defaults") + .build() + .await; + let instance = standalone.fe_instance(); + let ctx = Arc::new(QueryContext::with(DEFAULT_CATALOG_NAME, "public")); + instance.do_query("CREATE TABLE defaulted (greptime_timestamp TIMESTAMP(3) TIME INDEX, greptime_value DOUBLE, host STRING NOT NULL DEFAULT 'unknown' PRIMARY KEY)",ctx.clone()).await.remove(0).unwrap(); + let outcome = instance + .metrics( + metric_request(vec![gauge_metric( + "defaulted", + &[ + (60_000_000_000, 10.0, None), + (60_000_000_000, 20.0, Some("unknown")), + (60_000_000_000, 30.0, None), + ], + )]), + ctx.clone(), + ) + .await + .unwrap(); + assert_eq!( + (outcome.accepted_data_points, outcome.rejected_data_points), + (3, 0) + ); + for flushed in [false, true] { + if flushed { + instance + .do_query("ADMIN FLUSH_TABLE('defaulted')", ctx.clone()) + .await + .remove(0) + .unwrap(); + } + assert_trace_v2_query(instance,"SELECT COUNT(*)=1 AND MIN(greptime_value)=30 AND MIN(host)='unknown' FROM defaulted",ctx.clone()).await.unwrap(); + } + instance.do_query("CREATE TABLE required (greptime_timestamp TIMESTAMP(3) TIME INDEX, greptime_value DOUBLE, host STRING NOT NULL PRIMARY KEY)",ctx.clone()).await.remove(0).unwrap(); + for (name, points, expected) in [ + ( + "defaulted", + [ + (120_000_000_000, 1.0, None), + (180_000_000_000, 2.0, Some("supplied")), + ], + (3, 0), + ), + ( + "required", + [ + (60_000_000_000, 10.0, None), + (120_000_000_000, 20.0, Some("present")), + ], + (2, 1), + ), + ] { + for reverse in [false, true] { + let mut points = points; + if reverse { + points.reverse(); + } + let outcome = instance + .metrics( + metric_request(vec![ + gauge_metric(name, &points), + gauge_metric("independent", &[(60_000_000_000, 7.0, None)]), + ]), + ctx.clone(), + ) + .await + .unwrap(); + assert_eq!( + (outcome.accepted_data_points, outcome.rejected_data_points), + expected, + "{name}, reverse={reverse}" + ); + } + } + for sql in [ + "SELECT COUNT(*)=1 FROM defaulted WHERE greptime_value=1 AND host='unknown'", + "SELECT COUNT(*)=1 FROM defaulted WHERE greptime_value=2 AND host='supplied'", + "SELECT COUNT(*)=1 AND MIN(host)='present' FROM required", + ] { + assert_trace_v2_query(instance, sql, ctx.clone()) + .await + .unwrap(); + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_otlp_preflight_custom_backing_and_timestamp_units() { + for policy in 0..3 { + let standalone = GreptimeDbStandaloneBuilder::new(&format!("otlp_backing_{policy}")) + .with_auto_create_table(policy != 0) + .build() + .await; + let instance = standalone.fe_instance(); + for precision in [0, 3, 6, 9] { + let db = format!("backing_{precision}"); + let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, &db); + ctx.set_protocol_ctx(ProtocolCtx::OtlpMetric(OtlpMetricCtx { + with_metric_engine: true, + ..Default::default() + })); + if policy < 2 { + ctx.set_extension( + "auto_create_table", + if policy == 0 { "true" } else { "false" }, + ); + } + let ctx = Arc::new(ctx); + for sql in [ + format!("CREATE DATABASE {db}"), + format!( + "CREATE TABLE custom_physical (greptime_timestamp TIMESTAMP({precision}) TIME INDEX, greptime_value DOUBLE) ENGINE=metric WITH ('physical_metric_table'='true')" + ), + format!( + "CREATE TABLE custom_total (greptime_timestamp TIMESTAMP({precision}) TIME INDEX, greptime_value DOUBLE, \"stream\" STRING PRIMARY KEY) ENGINE=metric WITH ('on_physical_table'='custom_physical','otlp_metric_compat'='prom')" + ), + ] { + instance + .do_query(&sql, ctx.clone()) + .await + .remove(0) + .unwrap(); + } + let mut request = + build_sum_request("custom", AggregationTemporality::Cumulative, &[(60, 7)]); + if let Some(metric::Data::Sum(data)) = request.resource_metrics[0].scope_metrics[0] + .metrics[0] + .data + .as_mut() + { + data.data_points[0].time_unix_nano += 123_456_789; + } + request.resource_metrics[0].scope_metrics[0] + .metrics + .push(native_metric("bad", -5)); + let outcome = instance.metrics(request, ctx.clone()).await.unwrap(); + assert_eq!( + (outcome.accepted_data_points, outcome.rejected_data_points), + (1, 1), + "policy={policy}, precision={precision}, {outcome:?}" + ); + let expected = 60_123_456_789i64 / 10i64.pow(9 - precision); + let output = instance + .do_query( + "SELECT greptime_timestamp,greptime_value FROM custom_total", + ctx.clone(), + ) + .await + .remove(0) + .unwrap(); + let OutputData::Stream(stream) = output.data else { + unreachable!() + }; + let batches = RecordBatches::try_collect(stream).await.unwrap().take(); + let column = batches[0].column(0); + let actual = match precision { + 0 => column.as_primitive::().value(0), + 3 => column.as_primitive::().value(0), + 6 => column.as_primitive::().value(0), + _ => column.as_primitive::().value(0), + }; + assert_eq!(actual, expected); + assert_trace_v2_query(instance,&format!("SELECT COUNT(*)=0 FROM information_schema.tables WHERE table_schema='{db}' AND table_name IN ('greptime_physical_table','bad')"),ctx.clone()).await.unwrap(); + if policy == 2 && precision == 3 { + instance.do_query("CREATE TABLE not_physical (greptime_timestamp TIMESTAMP TIME INDEX, greptime_value DOUBLE)", ctx.clone()).await.remove(0).unwrap(); + let mut selected = (*ctx).clone(); + selected.set_extension( + servers::http::prom_store::PHYSICAL_TABLE_PARAM, + "not_physical", + ); + let mut request = build_sum_request( + "custom", + AggregationTemporality::Cumulative, + &[(120, 9)], + ); + request.resource_metrics[0].scope_metrics[0] + .metrics + .push(gauge_metric( + "invalid_backing", + &[(120_000_000_000, 1.0, None)], + )); + let outcome = instance.metrics(request, Arc::new(selected)).await.unwrap(); + assert_eq!( + (outcome.accepted_data_points, outcome.rejected_data_points), + (1, 1) + ); + } + } + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_otlp_delta_exponential_distributed_queries() { + use otel_arrow_rust::proto::opentelemetry::metrics::v1::exponential_histogram_data_point::Buckets; + + let cluster = tests::create_distributed_instance("delta_exponential_dist").await; + let instance = cluster.frontend(); + let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, "public"); + ctx.set_protocol_ctx(ProtocolCtx::OtlpMetric(OtlpMetricCtx { + with_metric_engine: true, + ..Default::default() + })); + let ctx = Arc::new(ctx); + let mut metric = native_metric("delta_distribution", 0); + let Some(metric::Data::ExponentialHistogram(data)) = metric.data.as_mut() else { + unreachable!() + }; + let point = data.data_points[0].clone(); + data.data_points = [(60, 10), (120, 20), (180, 15)] + .into_iter() + .map(|(end, count)| { + let mut point = point.clone(); + point.time_unix_nano = end * 1_000_000_000; + point.start_time_unix_nano = (end - 60) * 1_000_000_000; + point.count = count; + point.sum = Some(count as f64 * 1.5); + point.zero_count = 0; + point.zero_threshold = 0.001; + point.positive = Some(Buckets { + offset: 0, + bucket_counts: vec![count], + }); + point + }) + .collect(); + let outcome = instance + .metrics(metric_request(vec![metric]), ctx.clone()) + .await + .unwrap(); + assert_eq!( + (outcome.accepted_data_points, outcome.rejected_data_points), + (3, 0) + ); + for flushed in [false, true] { + if flushed { + instance + .do_query("ADMIN FLUSH_TABLE('delta_distribution')", ctx.clone()) + .await + .remove(0) + .unwrap(); + } + for (expression, expected) in [ + ( + "histogram_count(sum_over_time(delta_distribution[3m]))", + 45.0, + ), + ("histogram_sum(sum_over_time(delta_distribution[3m]))", 67.5), + ( + "histogram_count(sum_over_time(delta_distribution[3m]) / 180)", + 0.25, + ), + ( + "histogram_fraction(1, 2, sum_over_time(delta_distribution[3m]))", + 1.0, + ), + ] { + let output = instance + .do_query( + &format!("TQL EVAL (180,180,'1m') {expression}"), + ctx.clone(), + ) + .await + .remove(0) + .unwrap(); + let OutputData::Stream(stream) = output.data else { + unreachable!() + }; + let batches = RecordBatches::try_collect(stream).await.unwrap().take(); + assert_eq!(batches[0].num_rows(), 1); + assert_eq!( + batches[0] + .columns() + .iter() + .find(|column| column.data_type() == &DataType::Float64) + .unwrap() + .as_primitive::() + .value(0), + expected, + "{expression}" + ); + } + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_otlp_preflight_logical_nullability_and_multiple_backing_units() { + let standalone = GreptimeDbStandaloneBuilder::new("otlp_logical_presence") + .with_logical_batcher(BatcherOptions { + protocols: vec![BatchingProtocol::Otlp], + max_batch_rows: 1, + ..Default::default() + }) + .build() + .await; + let instance = standalone.fe_instance(); + let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, "public"); + for (physical, logical, precision) in [ + ("greptime_physical_table", "millis", 3), + ("custom_physical", "nanos", 9), + ] { + for sql in [ + format!( + "CREATE TABLE {physical} (greptime_timestamp TIMESTAMP({precision}) TIME INDEX, greptime_value DOUBLE) ENGINE=metric WITH ('physical_metric_table'='true')" + ), + format!( + "CREATE TABLE {logical} (greptime_timestamp TIMESTAMP({precision}) TIME INDEX, greptime_value DOUBLE, host STRING NOT NULL PRIMARY KEY) ENGINE=metric WITH ('on_physical_table'='{physical}')" + ), + ] { + instance + .do_query(&sql, Arc::new(ctx.clone())) + .await + .remove(0) + .unwrap(); + } + } + for with_metric_engine in [false, true] { + ctx.set_protocol_ctx(ProtocolCtx::OtlpMetric(OtlpMetricCtx { + with_metric_engine, + ..Default::default() + })); + for batching in [false, true] { + ctx.set_logical_batching_enabled(batching); + for reverse in [false, true] { + let mut points = vec![ + (60_000_000_123, 10.0, None), + (120_000_000_456, 20.0, Some("present")), + ]; + if reverse { + points.reverse(); + } + // The default-backed target remains eligible when sent alone; + // custom-backed targets retain the established fallback. + for names in [vec!["millis"], vec!["millis", "nanos"]] { + // Ordinary ingestion emits milliseconds; test it against the + // matching existing logical table without auto-creation. + if !with_metric_engine && names.contains(&"nanos") { + continue; + } + let outcome = instance + .metrics( + metric_request( + names + .iter() + .map(|name| gauge_metric(name, &points)) + .collect(), + ), + Arc::new(ctx.clone()), + ) + .await + .unwrap(); + assert_eq!( + (outcome.accepted_data_points, outcome.rejected_data_points), + (names.len() as i64 * 2, 0) + ); + } + } + } + } + for (name, unit, expected) in [ + ("greptime_physical_table", TimeUnit::Millisecond, 60_000), + ("custom_physical", TimeUnit::Nanosecond, 60_000_000_123), + ] { + let output = instance.do_query(&format!("SELECT greptime_timestamp, host, greptime_value FROM {name} ORDER BY greptime_timestamp"), Arc::new(ctx.clone())).await.remove(0).unwrap(); + let OutputData::Stream(stream) = output.data else { + unreachable!() + }; + let batches = RecordBatches::try_collect(stream).await.unwrap().take(); + assert_eq!( + batches.iter().map(|batch| batch.num_rows()).sum::(), + 2 + ); + assert!(batches[0].column(1).is_null(0)); + assert_eq!( + batches[0] + .column(2) + .as_primitive::() + .value(0), + 10.0 + ); + let timestamp = batches[0].column(0); + assert_eq!(timestamp.data_type(), &DataType::Timestamp(unit, None)); + let actual = if unit == TimeUnit::Millisecond { + timestamp + .as_primitive::() + .value(0) + } else { + timestamp.as_primitive::().value(0) + }; + assert_eq!(actual, expected); + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_otlp_preflight_shared_physical_and_fanout_conflicts() { + let standalone = GreptimeDbStandaloneBuilder::new("otlp_physical_conflicts") + .build() + .await; + let instance = standalone.fe_instance(); + for reverse in [false, true] { + let db = format!("physical_{reverse}"); + let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, &db); + ctx.set_protocol_ctx(ProtocolCtx::OtlpMetric(OtlpMetricCtx { + with_metric_engine: true, + ..Default::default() + })); + let ctx = Arc::new(ctx); + instance + .do_query(&format!("CREATE DATABASE {db}"), ctx.clone()) + .await + .remove(0) + .unwrap(); + let mut scalar = gauge_metric("scalar", &[(60_000_000_000, 1.0, None)]); + if let Some(metric::Data::Gauge(data)) = scalar.data.as_mut() { + data.data_points[0] + .attributes + .push(keyvalue("greptime_native_histogram", "tag")); + } + let mut metrics = vec![ + scalar, + native_metric("native", 0), + gauge_metric("independent", &[(60_000_000_000, 7.0, None)]), + ]; + if reverse { + metrics.reverse(); + } + let outcome = instance + .metrics(metric_request(metrics.clone()), ctx.clone()) + .await + .unwrap(); + assert_eq!( + (outcome.accepted_data_points, outcome.rejected_data_points), + (1, 2) + ); + assert_trace_v2_query(instance,&format!("SELECT COUNT(*)=0 FROM information_schema.tables WHERE table_schema='{db}' AND table_name IN ('scalar','native')"),ctx.clone()).await.unwrap(); + assert_trace_v2_query(instance,&format!("SELECT COUNT(*)=0 FROM information_schema.columns WHERE table_schema='{db}' AND column_name='greptime_native_histogram'"),ctx.clone()).await.unwrap(); + // The existing physical Float64 definition is authoritative, including + // for existing-column writes that propose no ALTER. + instance.do_query("CREATE TABLE bad_type (greptime_timestamp TIMESTAMP(3) TIME INDEX, greptime_value BIGINT) WITH ('otlp_metric_compat'='prom')",ctx.clone()).await.remove(0).unwrap(); + let outcome = instance + .metrics( + metric_request(vec![ + gauge_metric("bad_type", &[(60_000_000_000, 3.0, None)]), + gauge_metric("independent", &[(120_000_000_000, 9.0, None)]), + ]), + ctx.clone(), + ) + .await + .unwrap(); + assert_eq!( + (outcome.accepted_data_points, outcome.rejected_data_points), + (1, 1) + ); + assert_trace_v2_query(instance, "SELECT COUNT(*)=0 FROM bad_type", ctx.clone()) + .await + .unwrap(); + let compatible_name = if reverse { "scalar" } else { "native" }; + let compatible = metrics + .iter() + .find(|m| m.name == compatible_name) + .unwrap() + .clone(); + let outcome = instance + .metrics(metric_request(vec![compatible]), ctx.clone()) + .await + .unwrap(); + assert_eq!( + (outcome.accepted_data_points, outcome.rejected_data_points), + (1, 0) + ); + if !reverse { + // An existing scalar logical table cannot ALTER a shared Struct + // field into a string tag either. + metrics + .iter_mut() + .find(|m| m.name == "scalar") + .unwrap() + .name = "independent".into(); + } + let outcome = instance + .metrics(metric_request(metrics), ctx.clone()) + .await + .unwrap(); + assert_eq!( + (outcome.accepted_data_points, outcome.rejected_data_points), + (2, 1) + ); + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_otlp_metrics_ignore_compat_option() { + let standalone = GreptimeDbStandaloneBuilder::new("otlp_compat_metadata") + .build() + .await; + let instance = standalone.fe_instance(); + let ctx = Arc::new(QueryContext::with(DEFAULT_CATALOG_NAME, "public")); + let tables = [ + ("without_option", ""), + ("with_prom_option", "WITH ('otlp_metric_compat'='prom')"), + ("with_old_option", "WITH ('otlp_metric_compat'='legacy')"), + ]; + for (name, options) in tables { + instance + .do_query( + &format!("CREATE TABLE {name} (greptime_timestamp TIMESTAMP(3) TIME INDEX, greptime_value DOUBLE) {options}"), + ctx.clone(), + ) + .await + .remove(0) + .unwrap(); + } + // Neither absent nor retained option values select a different conversion + // or make tables incompatible within the same request. + for timestamp in [60_000_000_123, 120_000_000_456] { + let outcome = instance + .metrics( + metric_request( + tables + .iter() + .map(|(name, _)| gauge_metric(name, &[(timestamp, 7.0, None)])) + .collect(), + ), + ctx.clone(), + ) + .await + .unwrap(); + assert_eq!( + (outcome.accepted_data_points, outcome.rejected_data_points), + (3, 0) + ); + } + for (name, _) in tables { + assert_trace_v2_query( + instance, + &format!("SELECT COUNT(*)=2 AND MIN(greptime_value)=7 AND MIN(greptime_timestamp)='1970-01-01 00:01:00'::TIMESTAMP FROM {name}"), + ctx.clone(), + ) + .await + .unwrap(); + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_otlp_delta_exponential_storage_precision_and_replacement() { + use datatypes::arrow::datatypes::Int64Type; + let standalone = GreptimeDbStandaloneBuilder::new("otlp_histogram_precision") + .build() + .await; + let instance = standalone.fe_instance(); + for precision in [0, 3, 6, 9] { + let db = format!("native_{precision}"); + let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, &db); + ctx.set_protocol_ctx(ProtocolCtx::OtlpMetric(OtlpMetricCtx { + with_metric_engine: true, + ..Default::default() + })); + let ctx = Arc::new(ctx); + for sql in [ + format!("CREATE DATABASE {db}"), + format!( + "CREATE TABLE greptime_physical_table (greptime_timestamp TIMESTAMP({precision}) TIME INDEX, greptime_value DOUBLE) ENGINE=metric WITH ('physical_metric_table'='true')" + ), + ] { + instance + .do_query(&sql, ctx.clone()) + .await + .remove(0) + .unwrap(); + } + let large = (1u64 << 53) + 1; + // Replace the sample after timestamp truncation, then replay it exactly. + for (round, count) in [(0, 10), (1, large), (1, large)] { + let mut metric = native_metric("interval", 0); + if let Some(metric::Data::ExponentialHistogram(data)) = metric.data.as_mut() { + let point = &mut data.data_points[0]; + point.count = count; + point.zero_count = count; + point.start_time_unix_nano = round as u64 * 1_000_000_000; + point.time_unix_nano = + 60_000_000_000 + if precision == 9 { 1 } else { 1 + round as u64 }; + } + let outcome = instance + .metrics(metric_request(vec![metric]), ctx.clone()) + .await + .unwrap(); + assert_eq!( + (outcome.accepted_data_points, outcome.rejected_data_points), + (1, 0) + ); + for flushed in [false, true] { + if flushed { + instance + .do_query("ADMIN FLUSH_TABLE('interval')", ctx.clone()) + .await + .remove(0) + .unwrap(); + } + let output = instance + .do_query( + "SELECT greptime_timestamp,greptime_native_histogram FROM interval", + ctx.clone(), + ) + .await + .remove(0) + .unwrap(); + let OutputData::Stream(stream) = output.data else { + unreachable!() + }; + let batches = RecordBatches::try_collect(stream).await.unwrap().take(); + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 1); + let counts = batches[0] + .column(1) + .as_struct() + .column_by_name("count_i64") + .unwrap() + .as_primitive::(); + assert_eq!(counts.value(0), count as i64); + assert_eq!( + batches[0].column(0).data_type(), + &DataType::Timestamp( + match precision { + 0 => TimeUnit::Second, + 3 => TimeUnit::Millisecond, + 6 => TimeUnit::Microsecond, + _ => TimeUnit::Nanosecond, + }, + None + ) + ); + } + } + if precision >= 6 { + let mut metric = native_metric("interval", 0); + if let Some(metric::Data::ExponentialHistogram(data)) = metric.data.as_mut() { + data.data_points[0].time_unix_nano += if precision == 6 { 1_000 } else { 2 }; + } + instance + .metrics(metric_request(vec![metric]), ctx.clone()) + .await + .unwrap(); + assert_trace_v2_query(instance, "SELECT COUNT(*)=2 FROM interval", ctx.clone()) + .await + .unwrap(); + } + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_otlp_preflight_disabled_delta_and_schema_precedence() { + let standalone = GreptimeDbStandaloneBuilder::new("otlp_disabled_delta") + .build() + .await; + let instance = standalone.fe_instance(); + let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, "public"); + ctx.set_protocol_ctx(ProtocolCtx::OtlpMetric(OtlpMetricCtx { + with_metric_engine: true, + ..Default::default() + })); + let ctx = Arc::new(ctx); + let mut cumulative = native_metric("histogram", 0); + if let Some(metric::Data::ExponentialHistogram(data)) = cumulative.data.as_mut() { + data.aggregation_temporality = AggregationTemporality::Cumulative as i32; + } + instance + .metrics( + metric_request(vec![ + cumulative, + gauge_metric("independent", &[(60_000_000_000, 7.0, None)]), + ]), + ctx.clone(), + ) + .await + .unwrap(); + let mut fixed = ctx.fork(); + fixed.set_extension("auto_create_table", "false"); + let outcome = instance + .metrics( + metric_request(vec![ + native_metric("histogram", 0), + gauge_metric("independent", &[(120_000_000_000, 9.0, None)]), + ]), + Arc::new(fixed), + ) + .await + .unwrap(); + assert_eq!( + (outcome.accepted_data_points, outcome.rejected_data_points), + (1, 1), + "{outcome:?}" + ); + assert_trace_v2_query(instance,"SELECT COUNT(*)=0 FROM information_schema.columns WHERE column_name='otlp_aggregation_temporality'",ctx.clone()).await.unwrap(); + for reverse in [false, true] { + let prefix = if reverse { "second" } else { "first" }; + instance.do_query(&format!("CREATE TABLE {prefix}_sum (greptime_timestamp TIMESTAMP(3) TIME INDEX, greptime_value BIGINT) WITH ('otlp_metric_compat'='prom')"),ctx.clone()).await.remove(0).unwrap(); + let mut explicit = build_histogram_request( + prefix, + vec![HistogramDataPoint { + time_unix_nano: 60_000_000_000, + count: 10, + bucket_counts: vec![10], + sum: Some(10.0), + ..Default::default() + }], + ) + .resource_metrics + .remove(0) + .scope_metrics + .remove(0) + .metrics + .remove(0); + explicit.name = prefix.into(); + let mut metrics = vec![explicit, native_metric(&format!("{prefix}_bucket"), 0)]; + if reverse { + metrics.reverse(); + } + let outcome = instance + .metrics(metric_request(metrics), ctx.clone()) + .await + .unwrap(); + assert_eq!( + (outcome.accepted_data_points, outcome.rejected_data_points), + (1, 1), + "{outcome:?}" + ); + assert_trace_v2_query( + instance, + &format!("SELECT COUNT(*)=1 FROM {prefix}_bucket"), + ctx.clone(), + ) + .await + .unwrap(); + assert_trace_v2_query(instance,&format!("SELECT COUNT(*)=0 FROM information_schema.tables WHERE table_name='{prefix}_count'"),ctx.clone()).await.unwrap(); + } + } + fn build_sum_request( name: &str, temporality: AggregationTemporality, diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index f4e1d16de49..7ea4d3df2ea 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -9174,9 +9174,257 @@ pub async fn test_otlp_exponential_histogram(store_type: StorageType) { assert!(body.contains("otlp_exponential_latency"), "{query}: {body}"); } + test_otlp_delta_exponential_queries(&client).await; + test_otlp_invalid_metric_names(&client).await; + guard.remove_all().await; } +async fn test_otlp_invalid_metric_names(client: &TestClient) { + for (name, strategy) in [ + ("!!!", "UnderscoreEscapingWithSuffixes"), + ("bad/name", "NoTranslation"), + ] { + for reverse in [false, true] { + let gauge = |name| serde_json::json!({"name":name,"gauge":{"dataPoints":[{"timeUnixNano":"60000000000","asDouble":7.0}]}}); + let mut metrics = vec![gauge(name), gauge("valid_name_sibling")]; + if reverse { + metrics.reverse(); + } + let request: ExportMetricsServiceRequest = serde_json::from_value( + serde_json::json!({"resourceMetrics":[{"scopeMetrics":[{"metrics":metrics}]}]}), + ) + .unwrap(); + let res = client + .post("/v1/otlp/v1/metrics") + .header("content-type", "application/x-protobuf") + .header("x-greptime-otlp-metric-translation-strategy", strategy) + .body(request.encode_to_vec()) + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let response = ExportMetricsServiceResponse::decode(res.bytes().await).unwrap(); + assert_eq!(response.partial_success.unwrap().rejected_data_points, 1); + } + } + validate_data( + "valid_metric_name_sibling", + client, + "SELECT greptime_value FROM valid_name_sibling", + "[[7.0]]", + ) + .await; +} + +async fn test_otlp_delta_exponential_queries(client: &TestClient) { + use opentelemetry_proto::tonic::metrics::v1::{ + AggregationTemporality, ExponentialHistogram, ExponentialHistogramDataPoint, Metric, + ResourceMetrics, ScopeMetrics, exponential_histogram_data_point, metric, + }; + let request = + |name: &str, samples: &[(u64, u64, u64)], temporality| ExportMetricsServiceRequest { + resource_metrics: vec![ResourceMetrics { + scope_metrics: vec![ScopeMetrics { + metrics: vec![Metric { + name: name.into(), + data: Some(metric::Data::ExponentialHistogram(ExponentialHistogram { + aggregation_temporality: temporality as i32, + data_points: samples + .iter() + .map(|&(start, end, count)| ExponentialHistogramDataPoint { + start_time_unix_nano: start * 1_000_000_000, + time_unix_nano: end * 1_000_000_000, + count, + sum: Some(count as f64 * 1.5), + zero_threshold: 0.001, + positive: Some(exponential_histogram_data_point::Buckets { + offset: 0, + bucket_counts: vec![count], + }), + ..Default::default() + }) + .collect(), + })), + ..Default::default() + }], + ..Default::default() + }], + ..Default::default() + }], + }; + for (name, samples) in [ + ( + "delta_interval", + vec![(0, 60, 10), (60, 120, 20), (120, 180, 15)], + ), + ("delta_crossing", vec![(90, 150, 12)]), + ("delta_left", vec![(60, 120, 10)]), + ("delta_gap", vec![(120, 180, 10), (240, 300, 20)]), + ] { + let res = client + .post("/v1/otlp/v1/metrics") + .header("content-type", "application/x-protobuf") + .body(request(name, &samples, AggregationTemporality::Delta).encode_to_vec()) + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + assert!( + ExportMetricsServiceResponse::decode(res.bytes().await) + .unwrap() + .partial_success + .is_none() + ); + } + for flushed in [false, true] { + if flushed { + for name in [ + "delta_interval", + "delta_crossing", + "delta_left", + "delta_gap", + ] { + flush_table(client, name).await; + } + } + for (expression, time, expected) in [ + ( + "histogram_count(sum_over_time(delta_interval[3m]))", + 180, + Some(45.0), + ), + ( + "histogram_sum(sum_over_time(delta_interval[3m]))", + 180, + Some(67.5), + ), + ( + "histogram_count(sum_over_time(delta_interval[3m]) / 180)", + 180, + Some(0.25), + ), + ( + "histogram_sum(sum_over_time(delta_interval[3m]) / 180)", + 180, + Some(0.375), + ), + ( + "histogram_fraction(1, 2, sum_over_time(delta_interval[3m]))", + 180, + Some(1.0), + ), + ( + "histogram_count(sum_over_time(delta_crossing[3m]))", + 300, + Some(12.0), + ), + ( + "histogram_count(sum_over_time(delta_gap[3m]))", + 300, + Some(30.0), + ), + ("sum_over_time(delta_left[3m])", 300, None), + ("sum_over_time(delta_interval[3m])", 500, None), + ] { + let res = client + .get(&format!( + "/v1/prometheus/api/v1/query?query={}&time={time}", + encode(expression) + )) + .send() + .await; + assert_eq!(res.status(), StatusCode::OK, "{expression}"); + let body: serde_json::Value = res.json().await; + let results = body["data"]["result"].as_array().unwrap(); + if let Some(expected) = expected { + assert_eq!(results.len(), 1, "{expression}: {body}"); + let value = body["data"]["result"][0]["value"][1] + .as_str() + .unwrap() + .parse::() + .unwrap(); + assert!((value - expected).abs() < 1e-12, "{expression}: {body}"); + } else { + assert!(results.is_empty(), "{expression}: {body}"); + } + } + } + let metadata: serde_json::Value = client + .get("/v1/prometheus/api/v1/metadata?metric=delta_interval") + .send() + .await + .json() + .await; + assert_eq!( + metadata["data"]["delta_interval"][0]["type"], "unknown", + "{metadata}" + ); + // Both creation orders keep cumulative and delta identities distinct. + let mut mixed = request( + "mixed_created", + &[(0, 60, 10)], + AggregationTemporality::Cumulative, + ); + mixed.resource_metrics.extend( + request( + "mixed_created", + &[(0, 60, 20)], + AggregationTemporality::Delta, + ) + .resource_metrics, + ); + let response = client + .post("/v1/otlp/v1/metrics") + .header("content-type", "application/x-protobuf") + .body(mixed.encode_to_vec()) + .send() + .await; + assert_eq!(response.status(), StatusCode::OK); + let metadata: serde_json::Value = client + .get("/v1/prometheus/api/v1/metadata?metric=mixed_created") + .send() + .await + .json() + .await; + assert_eq!( + metadata["data"]["mixed_created"][0]["type"], "unknown", + "{metadata}" + ); + for (name, first, second) in [ + ( + "delta_first", + AggregationTemporality::Delta, + AggregationTemporality::Cumulative, + ), + ( + "cumulative_first", + AggregationTemporality::Cumulative, + AggregationTemporality::Delta, + ), + ] { + for temporality in [first, second] { + let res = client + .post("/v1/otlp/v1/metrics") + .header("content-type", "application/x-protobuf") + .body(request(name, &[(0, 60, 10)], temporality).encode_to_vec()) + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + } + for flushed in [false, true] { + if flushed { + flush_table(client, name).await; + } + validate_data( + "delta_identity", + client, + &format!("SELECT COUNT(*) FROM {name}"), + "[[2]]", + ) + .await; + } + } +} + pub async fn test_otlp_metric_translation_strategies(store_type: StorageType) { common_telemetry::init_default_ut_logging(); let (app, mut guard) = diff --git a/tests/cases/standalone/promql/delta_exponential_histogram.result b/tests/cases/standalone/promql/delta_exponential_histogram.result new file mode 100644 index 00000000000..b30df4a076a --- /dev/null +++ b/tests/cases/standalone/promql/delta_exponential_histogram.result @@ -0,0 +1,106 @@ +-- Infer the canonical native-histogram Struct from a typed Parquet fixture. +-- OTLP endpoint tests cover production ingestion. +CREATE TABLE delta_histogram_source ( + greptime_timestamp TIMESTAMP TIME INDEX, + histogram_schema INT, zero_threshold DOUBLE, sample_sum DOUBLE, reset_hint INT, + start_timestamp TIMESTAMP, count_i64 BIGINT, zero_count_i64 BIGINT, + absent_float DOUBLE +); + +Affected Rows: 0 + +INSERT INTO delta_histogram_source VALUES + (60000, 0, 0.001, 15.0, 3, 0, 10, 0, NULL), + (120000, 0, 0.001, 30.0, 3, 60000, 20, 0, NULL), + (180000, 0, 0.001, 22.5, 3, 120000, 15, 0, NULL); + +Affected Rows: 3 + +COPY ( + SELECT greptime_timestamp, + named_struct( + 'schema', histogram_schema, + 'zero_threshold', zero_threshold, + 'sum', sample_sum, + 'reset_hint', reset_hint, + 'start_timestamp', start_timestamp, + 'custom_values', array_slice([sample_sum], 1, 0), + 'positive_span_offsets', [CAST(histogram_schema + 1 AS INT)], + 'positive_span_lengths', [CAST(histogram_schema + 1 AS INT)], + 'negative_span_offsets', array_slice([histogram_schema], 1, 0), + 'negative_span_lengths', array_slice([histogram_schema], 1, 0), + 'count_i64', count_i64, + 'zero_count_i64', zero_count_i64, + 'positive_buckets_i64', [count_i64], + 'negative_buckets_i64', array_slice([count_i64], 1, 0), + 'count_f64', absent_float, + 'zero_count_f64', absent_float, + 'positive_buckets_f64', array_slice([sample_sum], 1, 0), + 'negative_buckets_f64', array_slice([sample_sum], 1, 0) + ) AS greptime_native_histogram + FROM delta_histogram_source +) TO 'delta_histogram/data.parquet'; + +Affected Rows: 3 + +CREATE EXTERNAL TABLE delta_exponential_histogram +WITH (location='delta_histogram/data.parquet', format='parquet'); + +Affected Rows: 0 + +TQL EVAL (180, 180, '1m') histogram_count(sum_over_time(delta_exponential_histogram[3m])); + ++---------------------+----------------------------------------------------------------------------------------------------------------------+ +| greptime_timestamp | prom_native_histogram_count(prom_native_histogram_sum_over_time(greptime_timestamp_range,greptime_native_histogram)) | ++---------------------+----------------------------------------------------------------------------------------------------------------------+ +| 1970-01-01T00:03:00 | 45.0 | ++---------------------+----------------------------------------------------------------------------------------------------------------------+ + +TQL EVAL (180, 180, '1m') histogram_sum(sum_over_time(delta_exponential_histogram[3m])); + ++---------------------+--------------------------------------------------------------------------------------------------------------------+ +| greptime_timestamp | prom_native_histogram_sum(prom_native_histogram_sum_over_time(greptime_timestamp_range,greptime_native_histogram)) | ++---------------------+--------------------------------------------------------------------------------------------------------------------+ +| 1970-01-01T00:03:00 | 67.5 | ++---------------------+--------------------------------------------------------------------------------------------------------------------+ + +TQL EVAL (180, 180, '1m') histogram_count(sum_over_time(delta_exponential_histogram[3m]) / 180); + ++---------------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| greptime_timestamp | prom_native_histogram_count(prom_native_histogram_div_scalar(prom_native_histogram_sum_over_time(greptime_timestamp_range,greptime_native_histogram),Float64(180))) | ++---------------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| 1970-01-01T00:03:00 | 0.25 | ++---------------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------+ + +TQL EVAL (180, 180, '1m') histogram_fraction(1, 2, sum_over_time(delta_exponential_histogram[3m])); + ++---------------------+-----------------------------------------------------------------------------------------------------------------------------------------------+ +| greptime_timestamp | prom_native_histogram_fraction(prom_native_histogram_sum_over_time(greptime_timestamp_range,greptime_native_histogram),Float64(1),Float64(2)) | ++---------------------+-----------------------------------------------------------------------------------------------------------------------------------------------+ +| 1970-01-01T00:03:00 | 1.0 | ++---------------------+-----------------------------------------------------------------------------------------------------------------------------------------------+ + +-- A single endpoint contributes its whole interval; the left endpoint is excluded. +TQL EVAL (180, 180, '1m') histogram_count(sum_over_time(delta_exponential_histogram[1m])); + ++---------------------+----------------------------------------------------------------------------------------------------------------------+ +| greptime_timestamp | prom_native_histogram_count(prom_native_histogram_sum_over_time(greptime_timestamp_range,greptime_native_histogram)) | ++---------------------+----------------------------------------------------------------------------------------------------------------------+ +| 1970-01-01T00:03:00 | 15.0 | ++---------------------+----------------------------------------------------------------------------------------------------------------------+ + +TQL EVAL (360, 360, '1m') histogram_count(sum_over_time(delta_exponential_histogram[3m])); + ++--------------------+----------------------------------------------------------------------------------------------------------------------+ +| greptime_timestamp | prom_native_histogram_count(prom_native_histogram_sum_over_time(greptime_timestamp_range,greptime_native_histogram)) | ++--------------------+----------------------------------------------------------------------------------------------------------------------+ ++--------------------+----------------------------------------------------------------------------------------------------------------------+ + +DROP TABLE delta_exponential_histogram; + +Affected Rows: 0 + +DROP TABLE delta_histogram_source; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/promql/delta_exponential_histogram.sql b/tests/cases/standalone/promql/delta_exponential_histogram.sql new file mode 100644 index 00000000000..0fbe4364ce9 --- /dev/null +++ b/tests/cases/standalone/promql/delta_exponential_histogram.sql @@ -0,0 +1,51 @@ +-- Infer the canonical native-histogram Struct from a typed Parquet fixture. +-- OTLP endpoint tests cover production ingestion. +CREATE TABLE delta_histogram_source ( + greptime_timestamp TIMESTAMP TIME INDEX, + histogram_schema INT, zero_threshold DOUBLE, sample_sum DOUBLE, reset_hint INT, + start_timestamp TIMESTAMP, count_i64 BIGINT, zero_count_i64 BIGINT, + absent_float DOUBLE +); +INSERT INTO delta_histogram_source VALUES + (60000, 0, 0.001, 15.0, 3, 0, 10, 0, NULL), + (120000, 0, 0.001, 30.0, 3, 60000, 20, 0, NULL), + (180000, 0, 0.001, 22.5, 3, 120000, 15, 0, NULL); + +COPY ( + SELECT greptime_timestamp, + named_struct( + 'schema', histogram_schema, + 'zero_threshold', zero_threshold, + 'sum', sample_sum, + 'reset_hint', reset_hint, + 'start_timestamp', start_timestamp, + 'custom_values', array_slice([sample_sum], 1, 0), + 'positive_span_offsets', [CAST(histogram_schema + 1 AS INT)], + 'positive_span_lengths', [CAST(histogram_schema + 1 AS INT)], + 'negative_span_offsets', array_slice([histogram_schema], 1, 0), + 'negative_span_lengths', array_slice([histogram_schema], 1, 0), + 'count_i64', count_i64, + 'zero_count_i64', zero_count_i64, + 'positive_buckets_i64', [count_i64], + 'negative_buckets_i64', array_slice([count_i64], 1, 0), + 'count_f64', absent_float, + 'zero_count_f64', absent_float, + 'positive_buckets_f64', array_slice([sample_sum], 1, 0), + 'negative_buckets_f64', array_slice([sample_sum], 1, 0) + ) AS greptime_native_histogram + FROM delta_histogram_source +) TO 'delta_histogram/data.parquet'; + +CREATE EXTERNAL TABLE delta_exponential_histogram +WITH (location='delta_histogram/data.parquet', format='parquet'); + +TQL EVAL (180, 180, '1m') histogram_count(sum_over_time(delta_exponential_histogram[3m])); +TQL EVAL (180, 180, '1m') histogram_sum(sum_over_time(delta_exponential_histogram[3m])); +TQL EVAL (180, 180, '1m') histogram_count(sum_over_time(delta_exponential_histogram[3m]) / 180); +TQL EVAL (180, 180, '1m') histogram_fraction(1, 2, sum_over_time(delta_exponential_histogram[3m])); +-- A single endpoint contributes its whole interval; the left endpoint is excluded. +TQL EVAL (180, 180, '1m') histogram_count(sum_over_time(delta_exponential_histogram[1m])); +TQL EVAL (360, 360, '1m') histogram_count(sum_over_time(delta_exponential_histogram[3m])); + +DROP TABLE delta_exponential_histogram; +DROP TABLE delta_histogram_source; From b886ae525d2d06d083e6f19d78889caca397e235 Mon Sep 17 00:00:00 2001 From: NightG1owZ <3112428897@qq.com> Date: Fri, 23 Oct 2026 01:46:55 +0000 Subject: [PATCH 076/239] feat(config): support overriding region engine options from env (#9700) * feat(config): support overriding region engine options from env `region_engine` is a `Vec` of externally tagged enums, so plain dotted environment keys cannot address it: `config` walks the path, replaces the default sequence with a table and loading fails with `invalid type: map, expected a sequence`. Address the entries by tag name instead (`__REGION_ENGINE____`) and merge the overrides onto the default entries in a dedicated source, so the engines and fields that are not mentioned keep their configured value. An entry the default sequence does not have is appended; a numeric index, an override of the whole sequence, a field below a non-table default and two variables that set a path and a field below it are reported with a dedicated error instead of a serde type error. A configuration file that configures the sequence as a whole still replaces the merged value, which the example configs document. Closes #8620 Signed-off-by: NightG1owZ <316247938+NightG1owZ@users.noreply.github.com> * docs(config): clarify that the metric engine is always enabled `region_engine` defaults to [mito, file], but the datanode builds and registers the metric engine unconditionally. Reword the example configs so `mito` and `file` read as the default configuration entries and naming `metric` only supplies its options, then regenerate config.md. Signed-off-by: NightG1owZ <316247938+NightG1owZ@users.noreply.github.com> * docs(config): drop the env-var explanation from the config reference Environment variables are documented in the GreptimeTeam/docs repository, so the example configs and the generated config reference only describe the config file. Keep the note that the metric engine is always registered. Signed-off-by: NightG1owZ <316247938+NightG1owZ@users.noreply.github.com> * test(config): cover tagged sequence environment collection Add unit tests for TaggedSeqEnvSource::collect: merging by tag, appending unknown tags, empty and null defaults, flat key removal, the error paths and the take_entries and is_index_like helpers. The variables are injected through Environment::source, so the tests do not touch the process environment. Signed-off-by: NightG1owZ <316247938+NightG1owZ@users.noreply.github.com> --------- Signed-off-by: NightG1owZ <316247938+NightG1owZ@users.noreply.github.com> Co-authored-by: NightG1owZ <316247938+NightG1owZ@users.noreply.github.com> --- config/config.md | 4 +- config/datanode.example.toml | 2 + config/standalone.example.toml | 2 + src/cmd/src/options.rs | 4 + src/cmd/tests/load_config_test.rs | 213 +++++++++ src/common/config/src/config.rs | 401 +++++++++++++++- src/common/config/src/error.rs | 87 +++- src/common/config/src/lib.rs | 1 + src/common/config/src/tagged_seq.rs | 711 ++++++++++++++++++++++++++++ src/datanode/src/config.rs | 4 + src/standalone/src/options.rs | 4 + 11 files changed, 1426 insertions(+), 7 deletions(-) create mode 100644 src/common/config/src/tagged_seq.rs diff --git a/config/config.md b/config/config.md index d450609a926..e370bb157a0 100644 --- a/config/config.md +++ b/config/config.md @@ -174,7 +174,7 @@ | `storage.http_client.timeout` | String | `30s` | The total request timeout, applied from when the request starts connecting until the response body has finished.
Also considered a total deadline. | | `storage.http_client.pool_idle_timeout` | String | `90s` | The timeout for idle sockets being kept-alive. | | `storage.http_client.skip_ssl_validation` | Bool | `false` | To skip the ssl verification
**Security Notice**: Setting `skip_ssl_validation = true` disables certificate verification, making connections vulnerable to man-in-the-middle attacks. Only use this in development or trusted private networks. | -| `[[region_engine]]` | -- | -- | The region engine options. You can configure multiple region engines.
Each engine type (mito, file, metric) may appear only once; duplicates cause startup to fail. | +| `[[region_engine]]` | -- | -- | The region engine options. You can configure multiple region engines.
Each engine type (mito, file, metric) may appear only once; duplicates cause startup to fail.
The default configuration lists `mito` and `file`; `metric` is always registered, so naming
it only supplies its options rather than enabling the engine. | | `region_engine.mito` | -- | -- | The Mito engine options. | | `region_engine.mito.num_workers` | Integer | Auto | Number of region workers. Defaults to half of the available CPU cores, rounded up. | | `region_engine.mito.worker_channel_size` | Integer | `128` | Request channel size of each worker. | @@ -650,7 +650,7 @@ | `storage.http_client.timeout` | String | `30s` | The total request timeout, applied from when the request starts connecting until the response body has finished.
Also considered a total deadline. | | `storage.http_client.pool_idle_timeout` | String | `90s` | The timeout for idle sockets being kept-alive. | | `storage.http_client.skip_ssl_validation` | Bool | `false` | To skip the ssl verification
**Security Notice**: Setting `skip_ssl_validation = true` disables certificate verification, making connections vulnerable to man-in-the-middle attacks. Only use this in development or trusted private networks. | -| `[[region_engine]]` | -- | -- | The region engine options. You can configure multiple region engines.
Each engine type (mito, file, metric) may appear only once; duplicates cause startup to fail. | +| `[[region_engine]]` | -- | -- | The region engine options. You can configure multiple region engines.
Each engine type (mito, file, metric) may appear only once; duplicates cause startup to fail.
The default configuration lists `mito` and `file`; `metric` is always registered, so naming
it only supplies its options rather than enabling the engine. | | `region_engine.mito` | -- | -- | The Mito engine options. | | `region_engine.mito.num_workers` | Integer | Auto | Number of region workers. Defaults to half of the available CPU cores, rounded up. | | `region_engine.mito.worker_channel_size` | Integer | `128` | Request channel size of each worker. | diff --git a/config/datanode.example.toml b/config/datanode.example.toml index 4b8eb41242e..bfa3afc0bcf 100644 --- a/config/datanode.example.toml +++ b/config/datanode.example.toml @@ -509,6 +509,8 @@ skip_ssl_validation = false ## The region engine options. You can configure multiple region engines. ## Each engine type (mito, file, metric) may appear only once; duplicates cause startup to fail. +## The default configuration lists `mito` and `file`; `metric` is always registered, so naming +## it only supplies its options rather than enabling the engine. [[region_engine]] ## The Mito engine options. diff --git a/config/standalone.example.toml b/config/standalone.example.toml index f70277aa599..f73a3f8a325 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -729,6 +729,8 @@ skip_ssl_validation = false ## The region engine options. You can configure multiple region engines. ## Each engine type (mito, file, metric) may appear only once; duplicates cause startup to fail. +## The default configuration lists `mito` and `file`; `metric` is always registered, so naming +## it only supplies its options rather than enabling the engine. [[region_engine]] ## The Mito engine options. diff --git a/src/cmd/src/options.rs b/src/cmd/src/options.rs index e09b3c8f733..ed88431e4fd 100644 --- a/src/cmd/src/options.rs +++ b/src/cmd/src/options.rs @@ -65,6 +65,10 @@ where fn env_list_keys() -> Option<&'static [&'static str]> { T::env_list_keys() } + + fn env_tagged_seq_keys() -> Option<&'static [&'static str]> { + T::env_tagged_seq_keys() + } } /// A serde `deserialize_with` helper that parses the `plugins` field leniently. diff --git a/src/cmd/tests/load_config_test.rs b/src/cmd/tests/load_config_test.rs index 2781183c0c8..2a41c120b23 100644 --- a/src/cmd/tests/load_config_test.rs +++ b/src/cmd/tests/load_config_test.rs @@ -902,3 +902,216 @@ fn test_load_metric_config_with_removed_sparse_primary_key_encoding() { Duration::from_secs(30) ); } + +/// Returns the mito engine config of a loaded `region_engine` sequence. +fn find_mito(region_engine: &[RegionEngineConfig]) -> &MitoConfig { + region_engine + .iter() + .find_map(|engine| match engine { + RegionEngineConfig::Mito(config) => Some(config), + _ => None, + }) + .expect("the mito engine config") +} + +#[test] +fn test_load_region_engine_mito_field_from_env() { + let env_prefix = "REGION_ENGINE_MITO_FIELD_UT"; + let env_key = [ + env_prefix, + "REGION_ENGINE", + "MITO", + "GLOBAL_WRITE_BUFFER_REJECT_SIZE", + ] + .join(ENV_VAR_SEP); + + temp_env::with_var(env_key, Some("4GB"), || { + let options = + GreptimeOptions::::load_layered_options(None, env_prefix).unwrap(); + + // The addressed field wins, the fields next to it keep their default value. + let mito = find_mito(&options.component.region_engine); + assert_eq!(mito.global_write_buffer_reject_size, ReadableSize::gb(4)); + assert_eq!( + mito.global_write_buffer_size, + MitoConfig::default().global_write_buffer_size + ); + assert_eq!(mito.index, MitoConfig::default().index); + + // The other entries of the default sequence are intact. + assert!( + options + .component + .region_engine + .iter() + .any(|engine| matches!(engine, RegionEngineConfig::File(_))), + "the file engine config should survive an env override of the mito engine" + ); + }); +} + +#[test] +fn test_load_region_engine_mito_field_from_env_standalone() { + let env_prefix = "REGION_ENGINE_MITO_FIELD_STANDALONE_UT"; + let env_key = [ + env_prefix, + "REGION_ENGINE", + "MITO", + "GLOBAL_WRITE_BUFFER_REJECT_SIZE", + ] + .join(ENV_VAR_SEP); + + temp_env::with_var(env_key, Some("4GB"), || { + let options = + GreptimeOptions::::load_layered_options(None, env_prefix).unwrap(); + + let mito = find_mito(&options.component.region_engine); + assert_eq!(mito.global_write_buffer_reject_size, ReadableSize::gb(4)); + assert_eq!( + mito.global_write_buffer_size, + MitoConfig::default().global_write_buffer_size + ); + assert!( + options + .component + .region_engine + .iter() + .any(|engine| matches!(engine, RegionEngineConfig::File(_))), + "the file engine config should survive an env override of the mito engine" + ); + }); +} + +#[test] +fn test_load_region_engine_metric_entry_from_env() { + let env_prefix = "REGION_ENGINE_METRIC_FIELD_UT"; + let env_key = [ + env_prefix, + "REGION_ENGINE", + "METRIC", + "FLUSH_METADATA_REGION_INTERVAL", + ] + .join(ENV_VAR_SEP); + + temp_env::with_var(env_key, Some("1m"), || { + let options = + GreptimeOptions::::load_layered_options(None, env_prefix).unwrap(); + + // `metric` is not in the default sequence, so it is appended instead of replacing + // the whole sequence. + assert_eq!( + options.component.region_engine.len(), + DatanodeOptions::default().region_engine.len() + 1 + ); + let metric = options + .component + .region_engine + .iter() + .find_map(|engine| match engine { + RegionEngineConfig::Metric(config) => Some(config), + _ => None, + }) + .expect("the metric engine config"); + assert_eq!( + metric.flush_metadata_region_interval, + Duration::from_secs(60) + ); + + // The engines of the default sequence keep their default options. + let mito = find_mito(&options.component.region_engine); + assert_eq!( + mito.global_write_buffer_reject_size, + MitoConfig::default().global_write_buffer_reject_size + ); + assert!( + options + .component + .region_engine + .iter() + .any(|engine| matches!(engine, RegionEngineConfig::File(_))) + ); + }); +} + +#[test] +fn test_load_region_engine_env_patch_is_replaced_by_config_file() { + let config = tempfile::NamedTempFile::new().unwrap(); + std::fs::write( + config.path(), + "[[region_engine]]\n[region_engine.metric]\nflush_metadata_region_interval = \"2m\"\n", + ) + .unwrap(); + + let env_prefix = "REGION_ENGINE_FILE_WINS_UT"; + let env_key = [ + env_prefix, + "REGION_ENGINE", + "MITO", + "GLOBAL_WRITE_BUFFER_REJECT_SIZE", + ] + .join(ENV_VAR_SEP); + + temp_env::with_var(env_key, Some("4GB"), || { + let options = GreptimeOptions::::load_layered_options( + config.path().to_str(), + env_prefix, + ) + .unwrap(); + + // A sequence is a single value: the config file replaces the whole value produced by + // the environment, so the environment patch (which is merged onto the default + // entries) is dropped. That keeps the precedence file > env > default of every other + // option, and means a config file has to spell out every engine it wants. + assert_eq!( + options.component.region_engine, + vec![RegionEngineConfig::Metric(MetricEngineConfig { + flush_metadata_region_interval: Duration::from_secs(120), + })] + ); + }); +} + +#[test] +fn test_load_region_engine_env_patch_with_config_file_lacking_region_engine() { + // The usual deployment shape: a config file that sets other options, plus an environment + // variable patching one region engine. The file says nothing about `region_engine`, so + // there is nothing that could replace the patched sequence. + let config = tempfile::NamedTempFile::new().unwrap(); + std::fs::write(config.path(), "node_id = 7\n").unwrap(); + + let env_prefix = "REGION_ENGINE_FILE_NO_ENGINE_UT"; + let env_key = [ + env_prefix, + "REGION_ENGINE", + "MITO", + "GLOBAL_WRITE_BUFFER_REJECT_SIZE", + ] + .join(ENV_VAR_SEP); + + temp_env::with_var(env_key, Some("4GB"), || { + let options = GreptimeOptions::::load_layered_options( + config.path().to_str(), + env_prefix, + ) + .unwrap(); + + // The option of the file is applied... + assert_eq!(options.component.node_id, Some(7)); + // ...the environment patch applies on top of the default engine entries... + let mito = find_mito(&options.component.region_engine); + assert_eq!(mito.global_write_buffer_reject_size, ReadableSize::gb(4)); + assert_eq!( + mito.global_write_buffer_size, + MitoConfig::default().global_write_buffer_size + ); + // ...and the rest of the default sequence is still configured. + assert!( + options + .component + .region_engine + .iter() + .any(|engine| matches!(engine, RegionEngineConfig::File(_))), + "the file engine entry should survive" + ); + }); +} diff --git a/src/common/config/src/config.rs b/src/common/config/src/config.rs index 85ce3d206f6..8dfd195d0bc 100644 --- a/src/common/config/src/config.rs +++ b/src/common/config/src/config.rs @@ -17,7 +17,10 @@ use serde::Serialize; use serde::de::DeserializeOwned; use snafu::ResultExt; -use crate::error::{LoadLayeredConfigSnafu, Result, SerdeJsonSnafu, TomlFormatSnafu}; +use crate::error::{ + LoadLayeredConfigSnafu, Result, SerdeJsonSnafu, TaggedSeqNotTopLevelSnafu, TomlFormatSnafu, +}; +use crate::tagged_seq::TaggedSeqEnvSource; /// Separator for environment variables. For example, `DATANODE__STORAGE__MANIFEST__CHECKPOINT_MARGIN`. pub const ENV_VAR_SEP: &str = "__"; @@ -64,9 +67,27 @@ pub trait Configurable: Serialize + DeserializeOwned + Default + Sized { let default_config = File::from_str(&json_str, FileFormat::Json); // Add default values and environment variables as the sources of the configuration. - let mut layered_config = config::Config::builder() - .add_source(default_config) - .add_source(env_source); + let mut layered_config = config::Config::builder().add_source(default_config); + + // A sequence of tagged tables cannot be addressed by plain dotted env keys: the path + // walk of `config` drops the default sequence. Such keys are merged onto the default + // entries by tag name instead. + layered_config = match Self::env_tagged_seq_keys() { + Some(keys) if !keys.is_empty() => { + // The default entries are looked up by literal key, so a dotted key would find + // none of them. Reject it instead of silently dropping the defaults. + if let Some(key) = keys.iter().find(|key| key.contains('.')) { + return TaggedSeqNotTopLevelSnafu { + key: (*key).to_string(), + } + .fail(); + } + + let defaults = serde_json::to_value(&default_opts).context(SerdeJsonSnafu)?; + layered_config.add_source(TaggedSeqEnvSource::new(env_source, keys, &defaults)) + } + _ => layered_config.add_source(env_source), + }; // Add config file as the source of the configuration if it is specified. if let Some(config_file) = config_file { @@ -93,6 +114,22 @@ pub trait Configurable: Serialize + DeserializeOwned + Default + Sized { None } + /// Keys whose value is a sequence of single-key tagged tables (e.g. `region_engine`, + /// serialized from `Vec` with externally tagged variants). Environment variables + /// address them by tag name (`KEY____`) instead of numeric index, and the + /// overrides are merged onto the corresponding default entries. + /// + /// Keys must name a **top-level** field of the options struct: the default entries to + /// merge onto are looked up by literal key in the serialized defaults, so a dotted key is + /// rejected instead of silently leaving the defaults out of the merged sequence. + /// + /// A configuration file that configures such a key as a whole replaces the merged value, + /// so the environment variables for it are ineffective then and the file has to spell out + /// every entry it wants. + fn env_tagged_seq_keys() -> Option<&'static [&'static str]> { + None + } + /// Serialize the configuration to a TOML string. fn to_toml(&self) -> Result { toml::to_string(&self).context(TomlFormatSnafu) @@ -102,7 +139,9 @@ pub trait Configurable: Serialize + DeserializeOwned + Default + Sized { #[cfg(test)] mod tests { use std::io::Write; + use std::time::Duration; + use common_base::readable_size::ReadableSize; use common_telemetry::logging::LoggingOptions; use common_test_util::temp_dir::create_named_temp_file; use common_wal::config::DatanodeWalConfig; @@ -264,4 +303,358 @@ mod tests { }, ); } + + /// A sequence of single-key tagged tables, the shape of `Vec` with externally + /// tagged variants. + #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] + #[serde(default)] + struct TestIndexConfig { + aux_path: String, + cache_percent: u8, + } + + impl Default for TestIndexConfig { + fn default() -> Self { + Self { + aux_path: "/default/aux".to_string(), + cache_percent: 30, + } + } + } + + #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] + #[serde(default)] + struct TestMitoConfig { + global_write_buffer_size: ReadableSize, + global_write_buffer_reject_size: ReadableSize, + index: TestIndexConfig, + } + + impl Default for TestMitoConfig { + fn default() -> Self { + Self { + global_write_buffer_size: ReadableSize::mb(64), + global_write_buffer_reject_size: ReadableSize::gb(2), + index: TestIndexConfig::default(), + } + } + } + + #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] + #[serde(default)] + struct TestFileConfig { + data_dir: String, + } + + impl Default for TestFileConfig { + fn default() -> Self { + Self { + data_dir: "/default/data".to_string(), + } + } + } + + #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] + #[serde(default)] + struct TestMetricConfig { + #[serde(with = "humantime_serde")] + flush_metadata_region_interval: Duration, + } + + impl Default for TestMetricConfig { + fn default() -> Self { + Self { + flush_metadata_region_interval: Duration::from_secs(30), + } + } + } + + #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] + #[serde(rename_all = "snake_case")] + enum TestRegionEngine { + Mito(TestMitoConfig), + File(TestFileConfig), + Metric(TestMetricConfig), + } + + #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] + struct TestTaggedSeqConfig { + region_engine: Vec, + } + + impl Default for TestTaggedSeqConfig { + fn default() -> Self { + Self { + region_engine: vec![ + TestRegionEngine::Mito(TestMitoConfig::default()), + TestRegionEngine::File(TestFileConfig::default()), + ], + } + } + } + + impl Configurable for TestTaggedSeqConfig { + fn env_tagged_seq_keys() -> Option<&'static [&'static str]> { + Some(&["region_engine"]) + } + } + + /// A config whose tagged sequence key is not a top-level field. + #[derive(Debug, Serialize, Deserialize, Default)] + struct TestDottedKeyConfig { + region_engine: Vec, + } + + impl Configurable for TestDottedKeyConfig { + fn env_tagged_seq_keys() -> Option<&'static [&'static str]> { + Some(&["region_engine.mito"]) + } + } + + fn test_mito(config: &TestTaggedSeqConfig) -> &TestMitoConfig { + config + .region_engine + .iter() + .find_map(|engine| match engine { + TestRegionEngine::Mito(config) => Some(config), + _ => None, + }) + .expect("the mito entry") + } + + fn test_metric(config: &TestTaggedSeqConfig) -> &TestMetricConfig { + config + .region_engine + .iter() + .find_map(|engine| match engine { + TestRegionEngine::Metric(config) => Some(config), + _ => None, + }) + .expect("the metric entry") + } + + /// Loads with the given prefix and returns the message of the error, including the + /// source of the layered loading error. + fn load_error_text(env_prefix: &str) -> String { + let err = TestTaggedSeqConfig::load_layered_options(None, env_prefix).unwrap_err(); + let mut text = err.to_string(); + let mut source = std::error::Error::source(&err); + while let Some(error) = source { + text.push_str(&format!(": {error}")); + source = std::error::Error::source(error); + } + text + } + + #[test] + fn test_tagged_seq_env_overrides_a_field() { + let env_prefix = "TAGGED_SEQ_FIELD_UT"; + let env_key = [ + env_prefix, + "REGION_ENGINE", + "MITO", + "GLOBAL_WRITE_BUFFER_REJECT_SIZE", + ] + .join(ENV_VAR_SEP); + + temp_env::with_var(env_key, Some("4GB"), || { + let opts = TestTaggedSeqConfig::load_layered_options(None, env_prefix).unwrap(); + let defaults = TestTaggedSeqConfig::default(); + + // The entries of the sequence that the environment does not mention are intact. + assert_eq!(opts.region_engine.len(), defaults.region_engine.len()); + assert_eq!(opts.region_engine[1], defaults.region_engine[1]); + + // The addressed field wins, the fields next to it keep their default value. + let mito = test_mito(&opts); + assert_eq!(mito.global_write_buffer_reject_size, ReadableSize::gb(4)); + assert_eq!( + mito.global_write_buffer_size, + TestMitoConfig::default().global_write_buffer_size + ); + assert_eq!(mito.index, TestIndexConfig::default()); + }); + } + + #[test] + fn test_tagged_seq_env_appends_an_entry() { + let env_prefix = "TAGGED_SEQ_APPEND_UT"; + let env_key = [ + env_prefix, + "REGION_ENGINE", + "METRIC", + "FLUSH_METADATA_REGION_INTERVAL", + ] + .join(ENV_VAR_SEP); + + temp_env::with_var(env_key, Some("1m"), || { + let opts = TestTaggedSeqConfig::load_layered_options(None, env_prefix).unwrap(); + let defaults = TestTaggedSeqConfig::default(); + + // The entry is not in the default sequence, so it is appended. + assert_eq!(opts.region_engine.len(), defaults.region_engine.len() + 1); + assert_eq!(opts.region_engine[0], defaults.region_engine[0]); + assert_eq!(opts.region_engine[1], defaults.region_engine[1]); + assert_eq!( + test_metric(&opts).flush_metadata_region_interval, + Duration::from_secs(60) + ); + }); + } + + #[test] + fn test_tagged_seq_env_overrides_a_nested_field() { + let env_prefix = "TAGGED_SEQ_NESTED_UT"; + let env_key = [env_prefix, "REGION_ENGINE", "MITO", "INDEX", "AUX_PATH"].join(ENV_VAR_SEP); + + temp_env::with_var(env_key, Some("/tmp/aux"), || { + let opts = TestTaggedSeqConfig::load_layered_options(None, env_prefix).unwrap(); + + let mito = test_mito(&opts); + assert_eq!(mito.index.aux_path, "/tmp/aux"); + assert_eq!( + mito.index.cache_percent, + TestIndexConfig::default().cache_percent + ); + assert_eq!( + mito.global_write_buffer_size, + TestMitoConfig::default().global_write_buffer_size + ); + }); + } + + #[test] + fn test_tagged_seq_env_rejects_whole_key() { + let env_prefix = "TAGGED_SEQ_NO_TAG_UT"; + let env_key = [env_prefix, "REGION_ENGINE"].join(ENV_VAR_SEP); + + temp_env::with_var(env_key, Some("mito"), || { + let text = load_error_text(env_prefix); + assert!(text.contains("region_engine"), "{text}"); + assert!(text.contains("region_engine____"), "{text}"); + }); + } + + #[test] + fn test_tagged_seq_env_rejects_numeric_index() { + let env_prefix = "TAGGED_SEQ_NUMERIC_UT"; + let env_key = [ + env_prefix, + "REGION_ENGINE", + "0", + "MITO", + "GLOBAL_WRITE_BUFFER_SIZE", + ] + .join(ENV_VAR_SEP); + + temp_env::with_var(env_key, Some("4GB"), || { + let text = load_error_text(env_prefix); + assert!( + text.contains("region_engine.0.mito.global_write_buffer_size"), + "{text}" + ); + assert!(text.contains("addressed by tag name"), "{text}"); + }); + } + + #[test] + fn test_tagged_seq_env_rejects_scalar_and_field_under_same_key() { + let env_prefix = "TAGGED_SEQ_NOT_A_TABLE_UT"; + let whole_key = [env_prefix, "REGION_ENGINE", "MITO", "INDEX"].join(ENV_VAR_SEP); + let below_key = + [env_prefix, "REGION_ENGINE", "MITO", "INDEX", "AUX_PATH"].join(ENV_VAR_SEP); + + temp_env::with_vars( + [ + (whole_key, Some("not-a-table")), + (below_key, Some("/tmp/aux")), + ], + || { + // One variable sets `index` as a whole, another sets a field below it. The + // conflict is reported as such instead of failing later inside serde with a + // type error that does not name the variables involved. + let text = load_error_text(env_prefix); + assert!( + text.contains("another environment variable sets 'region_engine.mito.index'"), + "{text}" + ); + assert!(text.contains("region_engine.mito.index.aux_path"), "{text}"); + }, + ); + } + + #[test] + fn test_tagged_seq_env_rejects_signed_and_float_indices() { + for index in ["-1", "+1", "1e3"] { + let env_prefix = "TAGGED_SEQ_INDEX_SHAPE_UT"; + let env_key = [ + env_prefix, + "REGION_ENGINE", + index, + "MITO", + "GLOBAL_WRITE_BUFFER_SIZE", + ] + .join(ENV_VAR_SEP); + + temp_env::with_var(env_key, Some("4GB"), || { + let text = load_error_text(env_prefix); + assert!(text.contains("addressed by tag name"), "{index}: {text}"); + }); + } + } + + #[test] + fn test_tagged_seq_env_rejects_tag_without_field() { + let env_prefix = "TAGGED_SEQ_NO_FIELD_UT"; + let env_key = [env_prefix, "REGION_ENGINE", "MITO"].join(ENV_VAR_SEP); + + temp_env::with_var(env_key, Some("mito"), || { + let text = load_error_text(env_prefix); + assert!(text.contains("region_engine.mito"), "{text}"); + assert!(text.contains("region_engine____"), "{text}"); + }); + } + + #[test] + fn test_tagged_seq_env_rejects_field_below_scalar_default() { + // The `data_dir` field of the file engine defaults to a string, so nothing can be set + // below it. + let env_prefix = "TAGGED_SEQ_SCALAR_DEFAULT_UT"; + let env_key = [env_prefix, "REGION_ENGINE", "FILE", "DATA_DIR", "SUB"].join(ENV_VAR_SEP); + + temp_env::with_var(env_key, Some("/tmp/sub"), || { + let text = load_error_text(env_prefix); + assert!(text.contains("region_engine.file.data_dir"), "{text}"); + assert!(text.contains("default value is not a table"), "{text}"); + }); + } + + #[test] + fn test_tagged_seq_env_type_error_names_the_environment() { + let env_prefix = "TAGGED_SEQ_TYPE_ERROR_UT"; + let env_key = [ + env_prefix, + "REGION_ENGINE", + "MITO", + "INDEX", + "CACHE_PERCENT", + ] + .join(ENV_VAR_SEP); + + temp_env::with_var(env_key, Some("not-a-number"), || { + // The override replaces the default value as the environment collected it, so the + // type error still says where that value comes from. + let text = load_error_text(env_prefix); + assert!(text.contains("in the environment"), "{text}"); + }); + } + + #[test] + fn test_tagged_seq_key_must_be_top_level() { + let err = TestDottedKeyConfig::load_layered_options(None, "TAGGED_SEQ_DOTTED_KEY_UT") + .unwrap_err(); + let text = err.to_string(); + assert!(text.contains("region_engine.mito"), "{text}"); + assert!(text.contains("top-level"), "{text}"); + } } diff --git a/src/common/config/src/error.rs b/src/common/config/src/error.rs index 9e788d1d300..d7f81014e05 100644 --- a/src/common/config/src/error.rs +++ b/src/common/config/src/error.rs @@ -65,6 +65,84 @@ pub enum Error { #[snafu(implicit)] location: Location, }, + + #[snafu(display( + "Invalid environment variable key '{key}': '{target}' is a sequence that cannot be \ + overridden as a whole; address one of its entries with '{target}____'" + ))] + TaggedSeqMissingTag { + key: String, + target: String, + #[snafu(implicit)] + location: Location, + }, + + #[snafu(display( + "Invalid environment variable key '{key}': '{tag}' is a numeric index, but '{target}' \ + is addressed by tag name; use '{target}____' instead, e.g. \ + '{target}__MITO__'" + ))] + TaggedSeqNumericTag { + key: String, + tag: String, + target: String, + #[snafu(implicit)] + location: Location, + }, + + #[snafu(display( + "Invalid environment variable key '{key}': the entry '{tag}' of '{target}' needs a \ + field to override; use '{target}____' instead" + ))] + TaggedSeqMissingField { + key: String, + tag: String, + target: String, + #[snafu(implicit)] + location: Location, + }, + + #[snafu(display( + "Invalid environment variable key '{key}': another environment variable sets '{path}' \ + as a whole, so this key cannot set a field below it; set either the whole '{path}' or \ + its fields, not both" + ))] + TaggedSeqOverridesConflict { + key: String, + path: String, + #[snafu(implicit)] + location: Location, + }, + + #[snafu(display( + "'{path}' cannot be overridden from the environment: its default value is not a table, \ + so a field below it cannot be set" + ))] + TaggedSeqNotATable { + path: String, + #[snafu(implicit)] + location: Location, + }, + + #[snafu(display( + "'{key}' is registered as a sequence of tagged tables, but it is not a top-level field \ + of the options; the default entries to merge onto are looked up by literal key" + ))] + TaggedSeqNotTopLevel { + key: String, + #[snafu(implicit)] + location: Location, + }, + + #[snafu(display( + "The default value of '{target}' is not a sequence of single-key tables: {reason}" + ))] + TaggedSeqInvalidDefault { + target: String, + reason: String, + #[snafu(implicit)] + location: Location, + }, } impl ErrorExt for Error { @@ -73,7 +151,14 @@ impl ErrorExt for Error { Error::TomlFormat { .. } | Error::LoadLayeredConfig { .. } | Error::FileWatch { .. } - | Error::InvalidPath { .. } => StatusCode::InvalidArguments, + | Error::InvalidPath { .. } + | Error::TaggedSeqMissingTag { .. } + | Error::TaggedSeqNumericTag { .. } + | Error::TaggedSeqMissingField { .. } + | Error::TaggedSeqOverridesConflict { .. } + | Error::TaggedSeqNotATable { .. } + | Error::TaggedSeqNotTopLevel { .. } + | Error::TaggedSeqInvalidDefault { .. } => StatusCode::InvalidArguments, Error::SerdeJson { .. } => StatusCode::Unexpected, } } diff --git a/src/common/config/src/lib.rs b/src/common/config/src/lib.rs index eea3b1351d8..7dc962fe5b1 100644 --- a/src/common/config/src/lib.rs +++ b/src/common/config/src/lib.rs @@ -15,6 +15,7 @@ pub mod config; pub mod error; pub mod file_watcher; +mod tagged_seq; use std::time::Duration; diff --git a/src/common/config/src/tagged_seq.rs b/src/common/config/src/tagged_seq.rs new file mode 100644 index 00000000000..7d2faa0b9bd --- /dev/null +++ b/src/common/config/src/tagged_seq.rs @@ -0,0 +1,711 @@ +// 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. + +//! Environment overrides for configuration keys holding a sequence of tagged tables. +//! +//! Such a sequence is what a `Vec` of externally tagged enums serializes to, for +//! example the `region_engine` option: +//! +//! ```text +//! [{"mito": {...}}, {"file": {...}}] +//! ``` +//! +//! Handing the plain [`Environment`] source to `config` cannot work for these keys: +//! `config` writes every flat dotted key by walking down its path, and the walk resets +//! any intermediate value that is not a table. Walking `region_engine.mito.index.aux_path` +//! therefore destroys the default sequence before anything is set. +//! +//! This source instead addresses the entries by their tag name and merges the +//! overrides onto the matching default entries: +//! +//! ```text +//! GREPTIMEDB_STANDALONE__REGION_ENGINE__MITO__INDEX__AUX_PATH=... +//! ``` +//! +//! Entries that only the environment mentions are appended, and every field that is +//! not mentioned keeps its default value. A configuration file that configures the +//! whole sequence replaces the merged one, which makes these variables ineffective; +//! the file has to spell out every entry it wants. + +use config::{ConfigError, Environment, Map, Source, Value, ValueKind}; +use serde_json::Value as JsonValue; + +use crate::error::{ + Error, Result, TaggedSeqInvalidDefaultSnafu, TaggedSeqMissingFieldSnafu, + TaggedSeqMissingTagSnafu, TaggedSeqNotATableSnafu, TaggedSeqNumericTagSnafu, + TaggedSeqOverridesConflictSnafu, +}; + +/// A [`Source`] that merges environment overrides into sequences of tagged tables. +/// +/// See the [module documentation](self) for the addressed syntax. +#[derive(Debug)] +pub(crate) struct TaggedSeqEnvSource { + /// The wrapped environment source. Everything it collects is kept as is, apart + /// from the entries that address one of the `targets`. + inner: Environment, + /// The keys whose entries are addressed by tag name. + targets: Vec, +} + +/// One sequence of tagged tables that the environment may override. +#[derive(Debug, Clone)] +struct Target { + /// The dotted key of the sequence, e.g. `region_engine`. + key: String, + /// The default value of `key`, whose entries the overrides are merged into. + default: Option, +} + +impl TaggedSeqEnvSource { + /// Wraps `inner` so that the sequence at every key of `keys` accepts tag name + /// overrides from the environment, merged onto the value that the same key has + /// in `defaults`. + pub(crate) fn new( + inner: Environment, + keys: &'static [&'static str], + defaults: &JsonValue, + ) -> Self { + let targets = keys + .iter() + .map(|key| Target { + key: (*key).to_string(), + default: defaults.get(key).cloned(), + }) + .collect(); + + Self { inner, targets } + } +} + +impl Source for TaggedSeqEnvSource { + fn clone_into_box(&self) -> Box { + Box::new(Self { + inner: self.inner.clone(), + targets: self.targets.clone(), + }) + } + + fn collect(&self) -> std::result::Result, ConfigError> { + let mut collected = self.inner.collect()?; + + for target in &self.targets { + // Take the entries out first: `config` walks the collected map in an + // unspecified order, so a leftover `region_engine.mito.x` entry could be + // written after the sequence and destroy it again. + let entries = take_entries(&mut collected, &target.key); + let merged = merge(target, entries).map_err(|e| ConfigError::Message(e.to_string()))?; + collected.insert(target.key.clone(), merged); + } + + Ok(collected) + } +} + +/// Removes and returns the entries that address `key` itself or one of its children. +fn take_entries(collected: &mut Map, key: &str) -> Vec<(String, Value)> { + let child_prefix = format!("{key}."); + let mut keys = collected + .keys() + .filter(|existing| *existing == key || existing.starts_with(&child_prefix)) + .cloned() + .collect::>(); + // `Map` is unordered, sort so that a reported error is deterministic. + keys.sort(); + + keys.into_iter() + .filter_map(|key| collected.remove(&key).map(|value| (key, value))) + .collect() +} + +/// Rebuilds the sequence of `target` with the environment overrides `entries` applied. +fn merge(target: &Target, entries: Vec<(String, Value)>) -> Result { + let mut overrides = group_entries(target, entries)?; + + let mut merged = Vec::with_capacity(overrides.len() + 1); + for entry in default_entries(target)? { + let mut entry = json_to_value(entry); + let tag = single_tag(target, &entry)?; + + if let Some(value) = overrides.remove(&tag) { + apply_override(target, &tag, &mut entry, value)?; + } + + merged.push(entry); + } + + // Entries that the default sequence does not have are appended. + let mut appended = overrides.into_iter().collect::>(); + appended.sort_by(|a, b| a.0.cmp(&b.0)); + for (tag, value) in appended { + let mut table = Map::new(); + table.insert(tag, value); + merged.push(Value::new(None, ValueKind::Table(table))); + } + + Ok(Value::new(None, ValueKind::Array(merged))) +} + +/// Merges the override of the entry `tag` into `entry`, whose only key is `tag`. +fn apply_override(target: &Target, tag: &str, entry: &mut Value, value: Value) -> Result<()> { + let ValueKind::Table(body) = &mut entry.kind else { + return invalid_default(target, "expected an array of tables"); + }; + + let Some(field) = body.get_mut(tag) else { + return invalid_default(target, "expected each entry to hold exactly one tag"); + }; + + // Both sides are tables, this only fails on a field whose default value cannot hold + // the sub-table that the environment sets. + deep_merge(field, value).map_err(|path| not_a_table(&entry_path(target, tag, &path))) +} + +/// Returns the default entries of `target`, each one a single-key table `{tag: table}`. +fn default_entries(target: &Target) -> Result<&[JsonValue]> { + match &target.default { + None | Some(JsonValue::Null) => Ok(&[]), + Some(JsonValue::Array(entries)) => Ok(entries), + Some(_) => invalid_default(target, "expected an array of tables"), + } +} + +/// Returns the tag of `entry`, which must map one tag to a table. +fn single_tag(target: &Target, entry: &Value) -> Result { + let ValueKind::Table(table) = &entry.kind else { + return invalid_default(target, "expected an array of tables"); + }; + + let mut keys = table.keys(); + let (Some(tag), None) = (keys.next(), keys.next()) else { + return invalid_default(target, "expected each entry to hold exactly one tag"); + }; + if !matches!(&table[tag].kind, ValueKind::Table(_)) { + return invalid_default(target, "expected each entry to map its tag to a table"); + } + + Ok(tag.clone()) +} + +/// Groups the overrides by tag, nesting every field path into a table. +fn group_entries(target: &Target, entries: Vec<(String, Value)>) -> Result> { + // Every entry but `key` itself is addressed as `.....`. + let offset = target.key.len() + 1; + let mut groups: Map = Map::new(); + + for (key, value) in entries { + if key == target.key { + return TaggedSeqMissingTagSnafu { + key, + target: target.key.clone(), + } + .fail(); + } + + let rest = key[offset..].to_string(); + let Some((tag, field)) = rest.split_once('.') else { + return TaggedSeqMissingFieldSnafu { + key, + tag: rest, + target: target.key.clone(), + } + .fail(); + }; + + if is_index_like(tag) { + return TaggedSeqNumericTagSnafu { + key, + tag: tag.to_string(), + target: target.key.clone(), + } + .fail(); + } + + let field = nest(field, value); + match groups.get_mut(tag) { + Some(group) => deep_merge(group, field) + .map_err(|path| overrides_conflict(&key, &entry_path(target, tag, &path)))?, + None => { + groups.insert(tag.to_string(), field); + } + } + } + + Ok(groups) +} + +/// Whether `tag` looks like a position in the sequence rather than an engine name. +/// +/// `config` also accepts negative and floating point looking path segments (`[-1]`, +/// `[1e3]`), so anything that starts like a number is reported as a misplaced index; +/// otherwise those variables would only fail later with the generic `unknown variant` +/// error from serde. +fn is_index_like(tag: &str) -> bool { + let mut chars = tag.chars(); + match chars.next() { + Some(first) if first.is_ascii_digit() => true, + Some('-') | Some('+') => { + matches!(chars.next(), Some(next) if next.is_ascii_digit() || next == '.') + } + _ => false, + } +} + +/// Builds a table whose key path is the dotted `path`, e.g. `a.b` and `v` give +/// `{a: {b: v}}`. +fn nest(path: &str, value: Value) -> Value { + let mut table = Map::new(); + match path.split_once('.') { + Some((head, tail)) => { + table.insert(head.to_string(), nest(tail, value)); + } + None => { + table.insert(path.to_string(), value); + } + } + + Value::new(None, ValueKind::Table(table)) +} + +/// Merges `incoming` into `base`, recursively for tables. Returns the dotted path of +/// the field that cannot be merged into, when `base` holds a non-table value where +/// `incoming` needs a table. +fn deep_merge(base: &mut Value, incoming: Value) -> std::result::Result<(), String> { + merge_at(base, incoming, "") +} + +fn merge_at(base: &mut Value, incoming: Value, path: &str) -> std::result::Result<(), String> { + match incoming.kind { + ValueKind::Table(incoming) => { + let ValueKind::Table(base) = &mut base.kind else { + return Err(path.to_string()); + }; + for (field, value) in incoming { + let field_path = if path.is_empty() { + field.clone() + } else { + format!("{path}.{field}") + }; + // A non-table override replaces whatever the key already holds. It is inserted + // as is, so the value keeps the origin it was collected with. + if matches!(value.kind, ValueKind::Table(_)) { + match base.get_mut(&field) { + Some(existing) => merge_at(existing, value, &field_path)?, + None => { + base.insert(field, value); + } + } + } else { + base.insert(field, value); + } + } + Ok(()) + } + kind => { + // A non-table override replaces whatever the key already holds. + base.kind = kind; + Ok(()) + } + } +} + +fn invalid_default(target: &Target, reason: &str) -> Result { + TaggedSeqInvalidDefaultSnafu { + target: target.key.clone(), + reason: reason.to_string(), + } + .fail() +} + +/// Builds the dotted path of `path` below the entry `tag` of `target`. +fn entry_path(target: &Target, tag: &str, path: &str) -> String { + if path.is_empty() { + format!("{}.{}", target.key, tag) + } else { + format!("{}.{}.{}", target.key, tag, path) + } +} + +/// Error for a field whose default value is not a table, so nothing can be set below it. +fn not_a_table(path: &str) -> Error { + TaggedSeqNotATableSnafu { + path: path.to_string(), + } + .build() +} + +/// Error for a key that sets a field below a path another key sets as a whole. +fn overrides_conflict(key: &str, path: &str) -> Error { + TaggedSeqOverridesConflictSnafu { + key: key.to_string(), + path: path.to_string(), + } + .build() +} + +/// Converts a default value parsed from the JSON serialization of the options. +fn json_to_value(json: &JsonValue) -> Value { + let kind = match json { + JsonValue::Null => ValueKind::Nil, + JsonValue::Bool(value) => ValueKind::Boolean(*value), + JsonValue::Number(value) => match (value.as_i64(), value.as_u64()) { + (Some(value), _) => ValueKind::I64(value), + (None, Some(value)) => ValueKind::U64(value), + (None, None) => ValueKind::Float(value.as_f64().unwrap_or_default()), + }, + JsonValue::String(value) => ValueKind::String(value.clone()), + JsonValue::Array(values) => ValueKind::Array(values.iter().map(json_to_value).collect()), + JsonValue::Object(values) => ValueKind::Table( + values + .iter() + .map(|(key, value)| (key.clone(), json_to_value(value))) + .collect(), + ), + }; + + Value::new(None, kind) +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + const KEY: &str = "region_engine"; + const KEYS: &[&str] = &[KEY]; + + /// The default sequence used by the tests: two exported entries, `mito` before `file`, + /// with one nested table to override. + fn defaults() -> JsonValue { + json!({ + "region_engine": [ + { + "mito": { + "global_write_buffer_size": "64MB", + "index": { + "aux_path": "/default/aux", + "cache_percent": 30 + } + } + }, + {"file": {"data_dir": "/default/data"}} + ], + "data_home": "/default/data_home" + }) + } + + /// A fixed set of variables, so the tests never touch the process environment and can + /// run in parallel. + fn env(vars: &[(&str, &str)]) -> Environment { + let mut vars_map = Map::new(); + for (key, value) in vars { + vars_map.insert((*key).to_string(), (*value).to_string()); + } + + Environment::default() + .separator("__") + .try_parsing(true) + .source(Some(vars_map)) + } + + fn collect_with(vars: &[(&str, &str)], defaults: &JsonValue) -> Map { + TaggedSeqEnvSource::new(env(vars), KEYS, defaults) + .collect() + .unwrap() + } + + fn collect(vars: &[(&str, &str)]) -> Map { + collect_with(vars, &defaults()) + } + + fn to_json(collected: Map) -> JsonValue { + let mut map = serde_json::Map::new(); + for (key, value) in collected { + map.insert(key, value.try_deserialize().unwrap()); + } + + JsonValue::Object(map) + } + + fn collect_json(vars: &[(&str, &str)]) -> JsonValue { + to_json(collect(vars)) + } + + fn collect_json_with(vars: &[(&str, &str)], defaults: &JsonValue) -> JsonValue { + to_json(collect_with(vars, defaults)) + } + + fn collect_error_with(vars: &[(&str, &str)], defaults: &JsonValue) -> String { + match TaggedSeqEnvSource::new(env(vars), KEYS, defaults).collect() { + Ok(_) => panic!("expected an error, got {vars:?}"), + Err(ConfigError::Message(message)) => message, + Err(other) => panic!("expected a message error, got: {other}"), + } + } + + fn collect_error(vars: &[(&str, &str)]) -> String { + collect_error_with(vars, &defaults()) + } + + #[test] + fn test_collect_overrides_a_scalar_field() { + let collected = collect_json(&[("REGION_ENGINE__MITO__GLOBAL_WRITE_BUFFER_SIZE", "128MB")]); + + let engines = collected["region_engine"].as_array().unwrap(); + assert_eq!(engines.len(), 2); + assert_eq!(engines[0]["mito"]["global_write_buffer_size"], "128MB"); + // The fields next to the override keep their default value. + assert_eq!(engines[0]["mito"]["index"]["aux_path"], "/default/aux"); + assert_eq!(engines[1]["file"]["data_dir"], "/default/data"); + } + + #[test] + fn test_collect_overrides_a_nested_field() { + let collected = collect_json(&[("REGION_ENGINE__MITO__INDEX__AUX_PATH", "/tmp/aux")]); + + assert_eq!( + collected["region_engine"][0]["mito"]["index"]["aux_path"], + "/tmp/aux" + ); + assert_eq!( + collected["region_engine"][0]["mito"]["index"]["cache_percent"], + 30 + ); + assert_eq!( + collected["region_engine"][0]["mito"]["global_write_buffer_size"], + "64MB" + ); + } + + #[test] + fn test_collect_merges_several_variables_of_one_entry() { + let collected = collect_json(&[ + ("REGION_ENGINE__MITO__INDEX__AUX_PATH", "/tmp/aux"), + ("REGION_ENGINE__MITO__INDEX__CACHE_PERCENT", "50"), + ]); + + assert_eq!( + collected["region_engine"][0]["mito"]["index"]["aux_path"], + "/tmp/aux" + ); + assert_eq!( + collected["region_engine"][0]["mito"]["index"]["cache_percent"], + 50 + ); + } + + #[test] + fn test_collect_appends_unknown_tags_in_tag_order() { + let collected = collect_json(&[ + ( + "REGION_ENGINE__METRIC__FLUSH_METADATA_REGION_INTERVAL", + "1m", + ), + ("REGION_ENGINE__AAA__DATA_DIR", "/tmp/aaa"), + ]); + + let engines = collected["region_engine"].as_array().unwrap(); + assert_eq!(engines.len(), 4); + // The default entries stay first and keep their order. + assert_eq!(engines[0]["mito"]["global_write_buffer_size"], "64MB"); + assert_eq!(engines[1]["file"]["data_dir"], "/default/data"); + // Entries that only the environment mentions are appended, sorted by tag. + assert_eq!(engines[2], json!({"aaa": {"data_dir": "/tmp/aaa"}})); + assert_eq!( + engines[3], + json!({"metric": {"flush_metadata_region_interval": "1m"}}) + ); + } + + #[test] + fn test_collect_keeps_the_default_sequence_without_overrides() { + let collected = collect_json(&[("DATA_HOME", "/srv/data")]); + + assert_eq!(collected["region_engine"], defaults()["region_engine"]); + assert_eq!(collected["data_home"], "/srv/data"); + } + + #[test] + fn test_collect_removes_the_flat_keys_of_the_sequence() { + let collected = collect(&[ + ("REGION_ENGINE__MITO__INDEX__AUX_PATH", "/tmp/aux"), + ("REGION_ENGINE_EXTRA__MITO__AUX_PATH", "/tmp/other"), + ]); + + // A leftover `region_engine.x` key would be written after the sequence and destroy + // it again, so all of them must be consumed by the merge. + assert!( + !collected + .keys() + .any(|key| key.starts_with("region_engine.")), + "flat keys left in the collected map: {:?}", + collected.keys().collect::>() + ); + // A key that only shares the prefix is not part of the sequence. + assert!(collected.contains_key("region_engine_extra.mito.aux_path")); + } + + #[test] + fn test_collect_without_default_entries() { + let collected = collect_json_with( + &[ + ( + "REGION_ENGINE__METRIC__FLUSH_METADATA_REGION_INTERVAL", + "1m", + ), + ("REGION_ENGINE__FILE__DATA_DIR", "/tmp/data"), + ], + &json!({}), + ); + + assert_eq!( + collected["region_engine"], + json!([ + {"file": {"data_dir": "/tmp/data"}}, + {"metric": {"flush_metadata_region_interval": "1m"}} + ]) + ); + } + + #[test] + fn test_collect_with_null_default_entries() { + let collected = collect_json_with( + &[("REGION_ENGINE__FILE__DATA_DIR", "/tmp/data")], + &json!({"region_engine": null}), + ); + + assert_eq!( + collected["region_engine"], + json!([{"file": {"data_dir": "/tmp/data"}}]) + ); + } + + #[test] + fn test_collect_rejects_setting_the_whole_sequence() { + let message = collect_error(&[("REGION_ENGINE", "mito")]); + + assert!( + message.contains("cannot be overridden as a whole"), + "{message}" + ); + assert!( + message.contains("region_engine____"), + "{message}" + ); + } + + #[test] + fn test_collect_rejects_a_numeric_index() { + let message = collect_error(&[("REGION_ENGINE__0__MITO__INDEX__AUX_PATH", "/tmp/aux")]); + + assert!(message.contains("numeric index"), "{message}"); + assert!(message.contains("addressed by tag name"), "{message}"); + } + + #[test] + fn test_collect_rejects_a_tag_without_field() { + let message = collect_error(&[("REGION_ENGINE__MITO", "mito")]); + + assert!(message.contains("needs a field"), "{message}"); + } + + #[test] + fn test_collect_rejects_conflicting_overrides() { + let message = collect_error(&[ + ("REGION_ENGINE__MITO__INDEX", "not-a-table"), + ("REGION_ENGINE__MITO__INDEX__AUX_PATH", "/tmp/aux"), + ]); + + assert!( + message.contains("another environment variable sets 'region_engine.mito.index'"), + "{message}" + ); + } + + #[test] + fn test_collect_rejects_a_field_below_a_scalar_default() { + let message = collect_error(&[("REGION_ENGINE__FILE__DATA_DIR__SUB", "/tmp/sub")]); + + assert!(message.contains("region_engine.file.data_dir"), "{message}"); + assert!( + message.contains("default value is not a table"), + "{message}" + ); + } + + #[test] + fn test_collect_rejects_a_scalar_default_sequence() { + let message = collect_error_with( + &[("REGION_ENGINE__MITO__INDEX__AUX_PATH", "/tmp/aux")], + &json!({"region_engine": "not-an-array"}), + ); + + assert!( + message.contains("not a sequence of single-key tables"), + "{message}" + ); + } + + #[test] + fn test_collect_rejects_a_default_entry_with_two_tags() { + let message = collect_error_with( + &[("REGION_ENGINE__MITO__INDEX__AUX_PATH", "/tmp/aux")], + &json!({"region_engine": [{"mito": {}, "file": {}}]}), + ); + + assert!(message.contains("exactly one tag"), "{message}"); + } + + #[test] + fn test_take_entries_keeps_unrelated_keys() { + let mut collected = Map::new(); + for key in [ + "region_engine", + "region_engine.mito.index", + "region_engine_extra.mito", + "other", + ] { + collected.insert( + (*key).to_string(), + Value::new(None, ValueKind::String("v".to_string())), + ); + } + + let entries = take_entries(&mut collected, "region_engine"); + + assert_eq!( + entries + .iter() + .map(|(key, _)| key.as_str()) + .collect::>(), + vec!["region_engine", "region_engine.mito.index"] + ); + assert!(!collected.contains_key("region_engine")); + assert!(!collected.contains_key("region_engine.mito.index")); + assert!(collected.contains_key("region_engine_extra.mito")); + assert!(collected.contains_key("other")); + } + + #[test] + fn test_is_index_like() { + for tag in ["0", "12", "-1", "+1", "1e3", "-1.5"] { + assert!(is_index_like(tag), "{tag} should look like an index"); + } + + for tag in ["mito", "file", "metric", "", "m1", "-", "+"] { + assert!(!is_index_like(tag), "{tag} should not look like an index"); + } + } +} diff --git a/src/datanode/src/config.rs b/src/datanode/src/config.rs index e627d8dc245..34a41ce2801 100644 --- a/src/datanode/src/config.rs +++ b/src/datanode/src/config.rs @@ -173,6 +173,10 @@ impl Configurable for DatanodeOptions { "wal.broker_endpoints", ]) } + + fn env_tagged_seq_keys() -> Option<&'static [&'static str]> { + Some(&["region_engine"]) + } } #[allow(clippy::large_enum_variant)] diff --git a/src/standalone/src/options.rs b/src/standalone/src/options.rs index 657961e9298..bbbad14b93d 100644 --- a/src/standalone/src/options.rs +++ b/src/standalone/src/options.rs @@ -137,6 +137,10 @@ impl Configurable for StandaloneOptions { "pending_rows_batcher.logical_table.protocols", ]) } + + fn env_tagged_seq_keys() -> Option<&'static [&'static str]> { + Some(&["region_engine"]) + } } /// The [`StandaloneOptions`] is only defined in `standalone` crate, From 00ab09012d41d53fcf4a3c75fc030642c1c3faa3 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Fri, 23 Oct 2026 02:27:21 +0000 Subject: [PATCH 077/239] ci: fix CI setup steps that fail on Ubuntu 26.04 runners (#9901) * ci: install protoc with taiki-e/install-action Signed-off-by: Dennis Zhuang * ci: run multi-language SDK tests on Python 3.13 Signed-off-by: Dennis Zhuang * ci: drop the unused libfuzzer-14-dev package from fuzz builds Signed-off-by: Dennis Zhuang * ci: build fuzz targets with the bundled libFuzzer Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- .../build-windows-artifacts/action.yml | 4 ++- .github/workflows/apidoc.yml | 4 +-- .github/workflows/checks.yml | 4 +-- .github/workflows/integration.yml | 22 +++++++-------- .github/workflows/nightly-ci.yml | 12 ++++---- .github/workflows/release.yml | 4 +-- .github/workflows/run-multi-lang-tests.yml | 4 +-- .github/workflows/rust.yml | 28 +++++++++---------- 8 files changed, 41 insertions(+), 41 deletions(-) diff --git a/.github/actions/build-windows-artifacts/action.yml b/.github/actions/build-windows-artifacts/action.yml index e5da671bad8..d25dd27e117 100644 --- a/.github/actions/build-windows-artifacts/action.yml +++ b/.github/actions/build-windows-artifacts/action.yml @@ -22,7 +22,9 @@ inputs: runs: using: composite steps: - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 + with: + tool: protoc@3.23.4 - name: Install rust toolchain uses: actions-rust-lang/setup-rust-toolchain@v1 diff --git a/.github/workflows/apidoc.yml b/.github/workflows/apidoc.yml index eb7a720a421..fb5601fdc5b 100644 --- a/.github/workflows/apidoc.yml +++ b/.github/workflows/apidoc.yml @@ -19,9 +19,9 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: actions-rust-lang/setup-rust-toolchain@v1 - run: cargo doc --workspace --no-deps --document-private-items - run: | diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml index 55dcc34ccf3..3bc54372d06 100644 --- a/.github/workflows/checks.yml +++ b/.github/workflows/checks.yml @@ -119,9 +119,9 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: actions-rust-lang/setup-rust-toolchain@v1 - name: Rust Cache uses: Swatinem/rust-cache@v2 diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml index db4f4bef65b..afce4a8c796 100644 --- a/.github/workflows/integration.yml +++ b/.github/workflows/integration.yml @@ -54,9 +54,9 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: actions-rust-lang/setup-rust-toolchain@v1 - name: Build Binaries Cache id: build-binaries-cache @@ -193,9 +193,9 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: actions-rust-lang/setup-rust-toolchain@v1 - name: Rust Cache uses: Swatinem/rust-cache@v2 @@ -327,9 +327,9 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: actions-rust-lang/setup-rust-toolchain@v1 with: cache: false @@ -364,13 +364,11 @@ jobs: env: RUSTUP_TOOLCHAIN: ${{ env.FUZZ_RUST_TOOLCHAIN }} run: | - sudo apt-get install -y libfuzzer-14-dev cargo install cargo-fuzz --version "${CARGO_FUZZ_VERSION}" --locked - name: Build fuzz target binaries if: steps.fuzz-binaries.outputs.cache-hit != 'true' shell: bash env: - CUSTOM_LIBFUZZER_PATH: /usr/lib/llvm-14/lib/libFuzzer.a RUSTUP_TOOLCHAIN: ${{ env.FUZZ_RUST_TOOLCHAIN }} FUZZ_TARGETS: | fuzz_create_database @@ -587,9 +585,9 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: actions-rust-lang/setup-rust-toolchain@v1 - name: Build Profile-CI Cache id: build-greptime-ci-cache @@ -705,9 +703,9 @@ jobs: uses: ./.github/actions/setup-etcd-cluster # Only the Kafka WAL helper needs the Rust toolchain. - if: matrix.mode.kafkaWalHelper - uses: arduino/setup-protoc@v3 + uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - if: matrix.mode.kafkaWalHelper uses: actions-rust-lang/setup-rust-toolchain@v1 with: diff --git a/.github/workflows/nightly-ci.yml b/.github/workflows/nightly-ci.yml index 15f8f505bea..f6158a8d7ad 100644 --- a/.github/workflows/nightly-ci.yml +++ b/.github/workflows/nightly-ci.yml @@ -53,9 +53,9 @@ jobs: fetch-depth: 0 persist-credentials: false - uses: ./.github/actions/setup-cyborg - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: actions-rust-lang/setup-rust-toolchain@v1 - name: Rust Cache uses: Swatinem/rust-cache@v2 @@ -82,9 +82,9 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: rui314/setup-mold@v1 - name: Install toolchain uses: actions-rust-lang/setup-rust-toolchain@v1 @@ -134,9 +134,9 @@ jobs: fetch-depth: 0 persist-credentials: false - uses: ./.github/actions/setup-cyborg - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: KyleMayes/install-llvm-action@v1 with: version: "14.0" diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index bc5d045099c..b292b586453 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -516,9 +516,9 @@ jobs: fetch-depth: 0 persist-credentials: false - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: actions-rust-lang/setup-rust-toolchain@v1 with: diff --git a/.github/workflows/run-multi-lang-tests.yml b/.github/workflows/run-multi-lang-tests.yml index 52fdca2117c..dad4fbaa080 100644 --- a/.github/workflows/run-multi-lang-tests.yml +++ b/.github/workflows/run-multi-lang-tests.yml @@ -83,10 +83,10 @@ jobs: java-version: '17' cache: 'maven' - - name: Setup Python 3.8 + - name: Setup Python 3.13 uses: actions/setup-python@v5 with: - python-version: '3.8' + python-version: '3.13' - name: Setup uv uses: astral-sh/setup-uv@v3 diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index e9c77554e07..d1b93a60e43 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -40,9 +40,9 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: actions-rust-lang/setup-rust-toolchain@v1 with: components: rustfmt @@ -58,9 +58,9 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: actions-rust-lang/setup-rust-toolchain@v1 with: components: clippy @@ -102,9 +102,9 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - name: Install riscv64 cross toolchain run: sudo apt-get update && sudo apt-get install -y gcc-riscv64-linux-gnu g++-riscv64-linux-gnu - uses: actions-rust-lang/setup-rust-toolchain@v1 @@ -130,9 +130,9 @@ jobs: with: persist-credentials: false - uses: ./.github/actions/setup-cyborg - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: KyleMayes/install-llvm-action@v1 with: version: "14.0" @@ -166,9 +166,9 @@ jobs: with: persist-credentials: false - uses: ./.github/actions/setup-cyborg - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: KyleMayes/install-llvm-action@v1 with: version: "14.0" @@ -210,9 +210,9 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: rui314/setup-mold@v1 - name: Install toolchain uses: actions-rust-lang/setup-rust-toolchain@v1 @@ -265,9 +265,9 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: rui314/setup-mold@v1 - name: Install toolchain uses: actions-rust-lang/setup-rust-toolchain@v1 From 3483f4a2842a0b4d83a5d3654041fdeee6e8195d Mon Sep 17 00:00:00 2001 From: liyang Date: Fri, 23 Oct 2026 03:06:23 +0000 Subject: [PATCH 078/239] fix(ci): move monitoring datanodeStorage out of standalone base (#9902) The `datanodeStorage` field belongs directly under `monitoring.standalone` (GreptimeDBStandaloneSpec), not under `monitoring.standalone.base` (PodTemplateSpec). Nesting it under `base` renders an invalid GreptimeDBCluster resource and the CRD schema rejects it during server-side apply: .spec.monitoring.standalone.base.datanodeStorage: field not declared in schema --- .../setup-greptimedb-cluster/with-observability.yaml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/actions/setup-greptimedb-cluster/with-observability.yaml b/.github/actions/setup-greptimedb-cluster/with-observability.yaml index 581f60dea60..98659a89941 100644 --- a/.github/actions/setup-greptimedb-cluster/with-observability.yaml +++ b/.github/actions/setup-greptimedb-cluster/with-observability.yaml @@ -15,9 +15,9 @@ monitoring: limits: cpu: 500m memory: 256Mi - datanodeStorage: - fs: - storageSize: 5Gi + datanodeStorage: + fs: + storageSize: 5Gi vector: resources: requests: From 6a34078706da86e42d249577f5cae375ab5aeb50 Mon Sep 17 00:00:00 2001 From: LFC <990479+MichaelScofield@users.noreply.github.com> Date: Fri, 23 Oct 2026 03:40:48 +0000 Subject: [PATCH 079/239] fix(pipeline): accept auto-transformed values in existing JSON columns (#9895) * fix(pipeline): accept auto-transformed values in existing JSON columns Signed-off-by: luofucong * test(pipeline): compare decoded JSON values in regression test Signed-off-by: luofucong --------- Signed-off-by: luofucong --- .../src/etl/transform/transformer/greptime.rs | 90 ++++++++++++++++++- 1 file changed, 89 insertions(+), 1 deletion(-) diff --git a/src/pipeline/src/etl/transform/transformer/greptime.rs b/src/pipeline/src/etl/transform/transformer/greptime.rs index 16c9e3399b6..a2f2fc2101e 100644 --- a/src/pipeline/src/etl/transform/transformer/greptime.rs +++ b/src/pipeline/src/etl/transform/transformer/greptime.rs @@ -503,6 +503,12 @@ fn resolve_schema( let column_type = &mut schema_info.schema[index].column_schema.data_type; match (column_type, value_type) { (column_type, value_type) if column_type == value_type => Ok(()), + // Auto-transformed arrays and objects use JSONB-encoded binary values. + (ConcreteDataType::Json(json_type), ConcreteDataType::Binary(_)) + if !json_type.is_json2() => + { + Ok(()) + } (ConcreteDataType::Json(column_type), ConcreteDataType::Json(value_type)) if column_type.is_json2() && value_type.is_json2() => { @@ -1064,7 +1070,89 @@ mod tests { use table::test_util::MemTable; use super::*; - use crate::{PipelineDefinition, identity_pipeline}; + use crate::{Content, PipelineDefinition, identity_pipeline, parse, setup_pipeline}; + + #[test] + fn test_auto_transform_existing_json_column() + -> std::result::Result<(), Box> { + let pipeline = parse(&Content::Yaml( + r#" +version: 2 +transform: + - field: ts + type: epoch, ms + index: timestamp +"#, + ))?; + let (pipeline, mut schema_info, pipeline_definition, pipeline_params) = + setup_pipeline!(pipeline); + let pipeline_context = + PipelineContext::new(&pipeline_definition, &pipeline_params, Channel::Unknown); + schema_info.set_table(Some(MemTable::table( + "events", + RecordBatch::new_empty(Arc::new(Schema::new(vec![DatatypeColumnSchema::new( + "payload", + ConcreteDataType::Json(Default::default()), + true, + )]))), + ))); + let inputs = [ + serde_json::Value::Null, + serde_json::json!([2, 4, 6]), + serde_json::json!({"active": true}), + serde_json::json!({"a\"b": "\u{0001}"}), + serde_json::json!([]), + serde_json::json!({}), + serde_json::Value::Null, + ]; + let records = inputs + .iter() + .enumerate() + .map(|(i, payload)| serde_json::json!({"ts": i, "payload": payload})) + .collect::>(); + let output = pipeline + .exec_mut( + serde_json::Value::Array(records).into(), + &pipeline_context, + &mut schema_info, + )? + .into_transformed() + .ok_or("expected transformed rows")?; + assert_eq!(output.len(), inputs.len()); + let index = schema_info.index["payload"]; + assert_eq!( + schema_info.schema[index].column_schema.data_type, + ConcreteDataType::Json(Default::default()) + ); + for ((row, _), expected) in output.iter().zip(&inputs) { + if expected.is_null() { + assert_eq!( + row.values.get(index).and_then(|v| v.value_data.as_ref()), + None + ); + } else { + let Some(ValueData::BinaryValue(value)) = &row.values[index].value_data else { + return Err("expected JSONB-encoded value".into()); + }; + assert_eq!(datatypes::types::jsonb_to_serde_json(value)?, *expected); + } + } + let error = pipeline + .exec_mut( + serde_json::json!({"ts": 10, "payload": 7}).into(), + &pipeline_context, + &mut schema_info, + ) + .err() + .ok_or("expected a datatype mismatch")?; + assert!( + error + .to_string() + .contains("expected datatype: Json, actual datatype: Int64"), + "{error:?}" + ); + Ok(()) + } #[test] fn test_column_metadata_caches_json_settings() -> Result<()> { From c4043efc27a16e72d915422877ab4931318e9756 Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Fri, 23 Oct 2026 04:15:46 +0000 Subject: [PATCH 080/239] fix(mito2): require flat SST output for JSON2 (#9893) * fix(mito2): require flat SST output for JSON2 Signed-off-by: shuiyisong * fix(mito2): avoid contradictory JSON2 recovery warnings Signed-off-by: shuiyisong --------- Signed-off-by: shuiyisong --- src/mito2/src/engine/alter_test.rs | 190 ++++++++ src/mito2/src/engine/create_test.rs | 442 ++++++++++++++++-- src/mito2/src/memtable.rs | 31 +- src/mito2/src/region/opener.rs | 82 ++-- src/mito2/src/worker/handle_alter.rs | 7 +- .../common/types/json/json2_limit.result | 108 +++++ .../common/types/json/json2_limit.sql | 44 ++ 7 files changed, 826 insertions(+), 78 deletions(-) diff --git a/src/mito2/src/engine/alter_test.rs b/src/mito2/src/engine/alter_test.rs index 7fb9f86c122..aee05880125 100644 --- a/src/mito2/src/engine/alter_test.rs +++ b/src/mito2/src/engine/alter_test.rs @@ -17,6 +17,7 @@ use std::collections::HashMap; use std::sync::Arc; use std::time::Duration; +use api::helper::encode_json_value; use api::v1::bulk_wal_entry::Body; use api::v1::helper::{row, tag_column_schema}; use api::v1::region::{StrictWindow, compact_request}; @@ -33,10 +34,13 @@ use datafusion_common::ScalarValue; use datafusion_expr::{col, lit}; use datatypes::arrow::array::{ArrayRef, Float64Array, StringArray, TimestampMillisecondArray}; use datatypes::arrow::datatypes::{DataType, Field, Schema, TimeUnit}; +use datatypes::extension::json::Json2ExtensionType; +use datatypes::json::value::JsonValue; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ ColumnDefaultConstraint, ColumnSchema, FulltextAnalyzer, FulltextBackend, FulltextOptions, }; +use datatypes::types::json_type::{JsonNativeType, JsonObjectType}; use store_api::codec::PrimaryKeyEncoding; use store_api::logstore::provider::Provider; use store_api::metadata::{ColumnMetadata, RegionMetadata}; @@ -140,6 +144,192 @@ fn add_nullable_field1() -> RegionAlterRequest { } } +#[rstest::rstest] +#[case::implicit_time_series(false, None, false)] +#[case::explicit_time_series(false, Some("time_series"), false)] +#[case::explicit_bulk(false, Some("bulk"), true)] +#[case::default_flat(true, None, true)] +#[case::flat_overrides_time_series(true, Some("time_series"), true)] +#[tokio::test] +async fn test_engine_add_json2_format_validation( + #[case] default_flat_format: bool, + #[case] memtable: Option<&str>, + #[case] accepted: bool, + #[values(false, true)] write_before_alter: bool, +) { + let mut env = TestEnv::with_prefix("add-json2-memtable-validation").await; + let engine = env + .create_engine(MitoConfig { + default_flat_format, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + let mut request = CreateRequestBuilder::new().build(); + if let Some(memtable) = memtable { + request + .options + .insert("memtable.type".into(), memtable.into()); + } + engine + .handle_request(region_id, RegionRequest::Create(request.clone())) + .await + .unwrap(); + if write_before_alter { + put_rows( + &engine, + region_id, + Rows { + schema: rows_schema(&request), + rows: build_rows(0, 1), + }, + ) + .await; + } + let region = engine.get_region(region_id).unwrap(); + let old_metadata = region.metadata(); + let mut column_schema = ColumnSchema::new( + "json", + ConcreteDataType::json2(JsonNativeType::Object(JsonObjectType::new())), + true, + ); + column_schema.with_extension_type(&Json2ExtensionType::default()); + let result = engine + .handle_request( + region_id, + RegionRequest::Alter(RegionAlterRequest { + kind: AlterKind::AddColumns { + columns: vec![AddColumn { + column_metadata: ColumnMetadata { + column_schema, + semantic_type: SemanticType::Field, + column_id: 3, + }, + location: None, + }], + }, + }), + ) + .await; + if accepted { + result.unwrap(); + assert!(region.metadata().column_by_name("json").is_some()); + } else { + let err = result.unwrap_err(); + assert_eq!(StatusCode::InvalidArguments, err.status_code()); + assert!( + err.output_msg() + .contains("JSON2 columns require sst_format=flat") + ); + assert_eq!(old_metadata, region.metadata()); + } + let expected_metadata = region.metadata(); + assert_eq!( + expected_metadata, + region.manifest_ctx.manifest().await.metadata + ); + reopen_region(&engine, region_id, request.table_dir, true, request.options).await; + assert_eq!( + expected_metadata, + engine.get_region(region_id).unwrap().metadata() + ); +} + +#[rstest::rstest] +#[tokio::test] +async fn test_engine_rejects_json2_primary_key_format_change( + #[values("bulk", "time_series")] memtable: &str, + #[values(false, true)] write_before_alter: bool, +) { + let mut env = TestEnv::with_prefix("alter-json2-format-validation").await; + let engine = env.create_engine(MitoConfig::default()).await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new() + .field_datatype(ConcreteDataType::json2(JsonNativeType::Object( + JsonObjectType::new(), + ))) + .insert_option("memtable.type", memtable) + .build(); + engine + .handle_request(region_id, RegionRequest::Create(request.clone())) + .await + .unwrap(); + let json_rows = |n: i64| Rows { + schema: rows_schema(&request), + rows: vec![row(vec![ + ValueData::StringValue("tag".into()), + ValueData::JsonValue(encode_json_value(JsonValue::from([("a", n)]))), + ValueData::TimestampMillisecondValue(n * 1000), + ])], + }; + if write_before_alter { + put_rows(&engine, region_id, json_rows(1)).await; + } + let region = engine.get_region(region_id).unwrap(); + let manifest_version = region.manifest_ctx.manifest_version().await; + let err = engine + .handle_request( + region_id, + RegionRequest::Alter(RegionAlterRequest { + kind: AlterKind::SetRegionOptions { + options: vec![SetRegionOption::Format("primary_key".into())], + }, + }), + ) + .await + .unwrap_err(); + assert_eq!(StatusCode::InvalidArguments, err.status_code()); + assert!( + err.output_msg() + .contains("JSON2 columns require sst_format=flat") + ); + assert_eq!(Some(FormatType::Flat), region.version().options.sst_format); + assert_eq!( + FormatType::Flat, + region.manifest_ctx.manifest().await.sst_format + ); + if !write_before_alter { + assert_eq!( + manifest_version, + region.manifest_ctx.manifest_version().await + ); + } + put_rows(&engine, region_id, json_rows(2)).await; + let batches = RecordBatches::try_collect( + engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(), + ) + .await + .unwrap(); + assert_eq!( + usize::from(write_before_alter) + 1, + batches.iter().map(|b| b.num_rows()).sum::() + ); + let expected = batches.pretty_print().unwrap(); + flush_region(&engine, region_id, None).await; + reopen_region(&engine, region_id, request.table_dir, true, request.options).await; + assert_eq!( + Some(FormatType::Flat), + engine + .get_region(region_id) + .unwrap() + .version() + .options + .sst_format + ); + let batches = RecordBatches::try_collect( + engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(expected, batches.pretty_print().unwrap()); +} + fn build_row_with_added_field( metadata: &[ColumnMetadata], tag_0: &str, diff --git a/src/mito2/src/engine/create_test.rs b/src/mito2/src/engine/create_test.rs index 00c3686c764..05696ad6109 100644 --- a/src/mito2/src/engine/create_test.rs +++ b/src/mito2/src/engine/create_test.rs @@ -12,20 +12,39 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::assert_matches; +use std::sync::Arc; use std::time::Duration; -use api::v1::Rows; +use api::helper::encode_json_value; +use api::v1::helper::row; +use api::v1::region::{StrictWindow, compact_request}; +use api::v1::value::ValueData; +use api::v1::{Rows, WriteHint}; use common_error::ext::{ErrorExt, WhateverResult}; use common_error::status_code::StatusCode; use common_recordbatch::RecordBatches; +use datatypes::json::value::JsonValue; use datatypes::prelude::ConcreteDataType; use datatypes::types::json_type::{JsonNativeType, JsonObjectType}; -use store_api::region_engine::RegionEngine; -use store_api::region_request::{PathType, RegionCloseRequest, RegionOpenRequest, RegionRequest}; +use store_api::codec::PrimaryKeyEncoding; +use store_api::region_engine::{RegionEngine, RegionManifestInfo, RegionRole}; +use store_api::region_request::{ + RegionCloseRequest, RegionCompactRequest, RegionOpenRequest, RegionPutRequest, RegionRequest, +}; +use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME; use store_api::storage::{RegionId, ScanRequest}; use crate::config::MitoConfig; +use crate::engine::MitoEngine; +use crate::error::Error; +use crate::manifest::action::{RegionChange, RegionMetaAction, RegionMetaActionList}; +use crate::manifest::storage::{delta_file, manifest_dir}; use crate::region::options::MemtableOptions; +use crate::sst::FormatType; +use crate::sst::location::region_dir_from_table_dir; +use crate::test_util::batch_util::sort_batches_and_print; +use crate::test_util::sst_util::{new_sparse_primary_key, sst_region_metadata_with_encoding}; use crate::test_util::{ CreateRequestBuilder, TestEnv, build_rows, flush_region, put_rows, reopen_region, rows_schema, }; @@ -104,36 +123,89 @@ async fn test_engine_create_new_region() { test_engine_create_new_region_with_format(true).await; } +#[rstest::rstest] +#[case::implicit_time_series(false, None, None, false)] +#[case::explicit_time_series(false, Some("time_series"), None, false)] +#[case::explicit_bulk(false, Some("bulk"), None, true)] +#[case::flat_overrides_time_series(false, Some("time_series"), Some("flat"), true)] +#[case::default_flat(true, None, None, true)] +#[case::default_flat_overrides_time_series(true, Some("time_series"), None, true)] +#[case::primary_key_overrides_default_flat(true, None, Some("primary_key"), false)] #[tokio::test] -async fn test_engine_rejects_json2_with_time_series_memtable_on_create_and_open() --> WhateverResult<()> { - let mut env = TestEnv::with_prefix("json2-rejects-time-series-memtable").await; +async fn test_engine_create_json2_format_validation( + #[case] default_flat_format: bool, + #[case] memtable: Option<&str>, + #[case] sst_format: Option<&str>, + #[case] accepted: bool, +) -> WhateverResult<()> { + let mut env = TestEnv::with_prefix("create-json2-memtable-validation").await; let engine = env .create_engine(MitoConfig { - default_flat_format: false, + default_flat_format, ..Default::default() }) .await; let region_id = RegionId::new(1, 1); - let request = CreateRequestBuilder::new() + let mut request = CreateRequestBuilder::new() .field_datatype(ConcreteDataType::json2(JsonNativeType::Object( JsonObjectType::new(), ))) .insert_option("append_mode", "true") - .insert_option("memtable.type", "time_series") .build(); + if let Some(memtable) = memtable { + request + .options + .insert("memtable.type".into(), memtable.into()); + } + if let Some(sst_format) = sst_format { + request + .options + .insert("sst_format".into(), sst_format.into()); + } - let err = engine - .handle_request(region_id, RegionRequest::Create(request)) - .await - .unwrap_err(); - assert_eq!(StatusCode::InvalidArguments, err.status_code()); - assert!( - err.to_string() - .contains("JSON2 columns only support BulkMemtable") - ); + let result = engine + .handle_request(region_id, RegionRequest::Create(request.clone())) + .await; + if accepted { + result?; + } else { + let err = result.unwrap_err(); + assert_eq!(StatusCode::InvalidArguments, err.status_code()); + assert!( + err.output_msg() + .contains("JSON2 columns require sst_format=flat") + ); + assert!(engine.get_region(region_id).is_none()); + request + .options + .insert("memtable.type".into(), "bulk".into()); + engine + .handle_request(region_id, RegionRequest::Create(request.clone())) + .await?; + } + reopen_region(&engine, region_id, request.table_dir, true, request.options).await; + Ok(()) +} - let request = CreateRequestBuilder::new() +#[rstest::rstest] +#[case::implicit_primary_key(false, false)] +#[case::default_flat(true, false)] +#[case::explicit_primary_key(false, true)] +#[tokio::test] +async fn test_engine_recover_json2_region( + #[case] default_flat_format: bool, + #[case] explicit_primary_key: bool, + #[values(false, true)] use_create: bool, + #[values(FormatType::PrimaryKey, FormatType::Flat)] persisted_format: FormatType, +) -> WhateverResult<()> { + let mut env = TestEnv::with_prefix("recover-legacy-json2").await; + let config = MitoConfig { + default_flat_format, + ..Default::default() + }; + let engine = env.create_engine(config.clone()).await; + let region_id = RegionId::new(1, 1); + let mut request = CreateRequestBuilder::new() .field_datatype(ConcreteDataType::json2(JsonNativeType::Object( JsonObjectType::new(), ))) @@ -141,8 +213,208 @@ async fn test_engine_rejects_json2_with_time_series_memtable_on_create_and_open( .insert_option("memtable.type", "bulk") .build(); engine - .handle_request(region_id, RegionRequest::Create(request)) + .handle_request(region_id, RegionRequest::Create(request.clone())) .await?; + let schema = rows_schema(&request); + let json_row = |n: i64| { + row(vec![ + ValueData::StringValue("tag".into()), + ValueData::JsonValue(encode_json_value(JsonValue::from([("a", n)]))), + ValueData::TimestampMillisecondValue(n * 1000), + ]) + }; + for n in 1..=2 { + put_rows( + &engine, + region_id, + Rows { + schema: schema.clone(), + rows: vec![json_row(n)], + }, + ) + .await; + flush_region(&engine, region_id, None).await; + } + // Leave one row in the WAL to exercise replay into the recovered Bulk memtable. + put_rows( + &engine, + region_id, + Rows { + schema: schema.clone(), + rows: vec![json_row(3)], + }, + ) + .await; + let expected = scan_region_data(&engine, region_id).await; + + request.options.remove("memtable.type"); + if explicit_primary_key { + request + .options + .insert("sst_format".into(), "primary_key".into()); + request + .options + .insert("memtable.type".into(), "time_series".into()); + } + let open_request = RegionOpenRequest { + engine: request.engine.clone(), + table_dir: request.table_dir.clone(), + path_type: request.path_type, + options: request.options.clone(), + skip_wal_replay: false, + checkpoint: None, + requirements: Default::default(), + }; + let follower = env.create_follower_engine(config).await; + follower + .handle_request( + region_id, + RegionRequest::Open(RegionOpenRequest { + skip_wal_replay: true, + ..open_request.clone() + }), + ) + .await?; + + let region = engine.get_region(region_id).unwrap(); + // Reproduce the persisted state from older binaries, bypassing the new ALTER validation. + region + .manifest_ctx + .manifest_manager + .write() + .await + .update( + RegionMetaActionList::new(vec![RegionMetaAction::Change(RegionChange { + metadata: region.metadata(), + sst_format: persisted_format, + append_mode: Some(true), + })]), + false, + ) + .await?; + let manifest = region.manifest_ctx.manifest().await; + assert_eq!(2, manifest.files.len()); + assert_eq!(persisted_format, manifest.sst_format); + follower + .sync_region( + region_id, + RegionManifestInfo::mito(manifest.manifest_version, 0, 0).into(), + ) + .await?; + let follower_region = follower.get_region(region_id).unwrap(); + assert_eq!( + Some(FormatType::Flat), + follower_region.version().options.sst_format + ); + assert_eq!( + manifest.files, + follower_region.manifest_ctx.manifest().await.files + ); + follower.stop().await?; + + engine + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) + .await?; + let recovery = if use_create { + RegionRequest::Create(request.clone()) + } else { + RegionRequest::Open(open_request) + }; + engine.handle_request(region_id, recovery).await?; + let region = engine.get_region(region_id).unwrap(); + let recovered_manifest = region.manifest_ctx.manifest().await; + assert_eq!( + manifest.manifest_version, + recovered_manifest.manifest_version + ); + assert_eq!(manifest.metadata, recovered_manifest.metadata); + assert_eq!(manifest.files, recovered_manifest.files); + assert_eq!(persisted_format, recovered_manifest.sst_format); + assert_eq!(Some(FormatType::Flat), region.version().options.sst_format); + assert_eq!(expected, scan_region_data(&engine, region_id).await); + + engine + .set_region_role(region_id, RegionRole::Leader) + .unwrap(); + put_rows( + &engine, + region_id, + Rows { + schema, + rows: vec![json_row(4)], + }, + ) + .await; + let expected = scan_region_data(&engine, region_id).await; + flush_region(&engine, region_id, None).await; + engine + .handle_request( + region_id, + RegionRequest::Compact(RegionCompactRequest { + options: compact_request::Options::StrictWindow(StrictWindow { + window_seconds: 3600, + }), + ..Default::default() + }), + ) + .await?; + assert_eq!(1, region.manifest_ctx.manifest().await.files.len()); + assert_eq!(expected, scan_region_data(&engine, region_id).await); + reopen_region( + &engine, + region_id, + request.table_dir, + false, + request.options, + ) + .await; + assert_eq!(expected, scan_region_data(&engine, region_id).await); + Ok(()) +} + +async fn scan_region_data(engine: &MitoEngine, region_id: RegionId) -> String { + let stream = engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + sort_batches_and_print(&batches, &["ts"]) +} + +#[tokio::test] +async fn test_engine_create_preserves_manifest_on_open_error() -> WhateverResult<()> { + let mut env = TestEnv::with_prefix("create-open-error").await; + let engine = env + .create_engine(MitoConfig { + compress_manifest: false, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + engine + .handle_request(region_id, RegionRequest::Create(request.clone())) + .await?; + put_rows( + &engine, + region_id, + Rows { + schema: rows_schema(&request), + rows: build_rows(0, 3), + }, + ) + .await; + flush_region(&engine, region_id, None).await; + let expected = scan_region_data(&engine, region_id).await; + let manifest = engine + .get_region(region_id) + .unwrap() + .manifest_ctx + .manifest() + .await; engine .handle_request( region_id, @@ -150,32 +422,124 @@ async fn test_engine_rejects_json2_with_time_series_memtable_on_create_and_open( ) .await?; - let options = [ - ("append_mode".to_string(), "true".to_string()), - ("memtable.type".to_string(), "time_series".to_string()), - ] - .into_iter() - .collect(); + let store = env.get_object_store().unwrap(); + let dir = manifest_dir(®ion_dir_from_table_dir( + &request.table_dir, + region_id, + request.path_type, + )); + let path = format!("{dir}{}", delta_file(0)); + let original = store.read(&path).await.unwrap().to_bytes(); + store.write(&path, "invalid manifest").await.unwrap(); + let mut manifest_bytes = Vec::new(); + for entry in store.list(&dir).await.unwrap() { + if entry.metadata().is_file() { + manifest_bytes.push(( + entry.path().to_string(), + store.read(entry.path()).await.unwrap().to_bytes(), + )); + } + } let err = engine - .handle_request( - region_id, - RegionRequest::Open(RegionOpenRequest { - engine: String::new(), - table_dir: "test".to_string(), - path_type: PathType::Bare, - options, - skip_wal_replay: false, - checkpoint: None, - requirements: Default::default(), - }), - ) + .handle_request(region_id, RegionRequest::Create(request.clone())) + .await + .unwrap_err(); + assert_matches!( + err.as_any().downcast_ref::().unwrap(), + Error::SerdeJson { .. } + ); + assert!(!engine.is_region_exists(region_id)); + for (path, bytes) in manifest_bytes { + assert_eq!(bytes, store.read(&path).await.unwrap().to_bytes()); + } + + store.write(&path, original).await.unwrap(); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await?; + let recovered = engine + .get_region(region_id) + .unwrap() + .manifest_ctx + .manifest() + .await; + assert_eq!(manifest.manifest_version, recovered.manifest_version); + assert_eq!(manifest.files, recovered.files); + assert_eq!(expected, scan_region_data(&engine, region_id).await); + Ok(()) +} + +#[tokio::test] +async fn test_engine_create_json2_sparse_requires_flat() -> WhateverResult<()> { + let mut env = TestEnv::with_prefix("create-json2-sparse").await; + let engine = env.create_engine(MitoConfig::default()).await; + let region_id = RegionId::new(1, 1); + let mut request = CreateRequestBuilder::new() + .field_datatype(ConcreteDataType::json2(JsonNativeType::Object( + JsonObjectType::new(), + ))) + .insert_option("memtable.type", "time_series") + .insert_option("primary_key_encoding", "sparse") + .insert_option("sst_format", "primary_key") + .build(); + let json_column = request.column_metadatas[1].column_schema.clone(); + let metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + request.column_metadatas = metadata.column_metadatas.clone(); + request.primary_key = metadata.primary_key.clone(); + request.column_metadatas[4].column_schema = json_column; + let err = engine + .handle_request(region_id, RegionRequest::Create(request.clone())) .await .unwrap_err(); assert_eq!(StatusCode::InvalidArguments, err.status_code()); assert!( err.output_msg() - .contains("JSON2 columns only support BulkMemtable") + .contains("JSON2 columns require sst_format=flat") ); + assert!(!engine.is_region_exists(region_id)); + + request.options.insert("sst_format".into(), "flat".into()); + engine + .handle_request(region_id, RegionRequest::Create(request.clone())) + .await?; + let schema = rows_schema(&request); + let mut pk_schema = schema[0].clone(); + pk_schema.column_name = PRIMARY_KEY_COLUMN_NAME.into(); + pk_schema.datatype = api::v1::ColumnDataType::Binary.into(); + let rows = Rows { + schema: vec![pk_schema, schema[5].clone(), schema[4].clone()], + rows: vec![row(vec![ + ValueData::BinaryValue(new_sparse_primary_key(&["a", "x"], &metadata, 10, 0)), + ValueData::TimestampMillisecondValue(1000), + ValueData::JsonValue(encode_json_value(JsonValue::from([("a", 1i64)]))), + ])], + }; + engine + .handle_request( + region_id, + RegionRequest::Put(RegionPutRequest { + rows, + hint: Some(WriteHint { + primary_key_encoding: api::v1::PrimaryKeyEncoding::Sparse.into(), + }), + skip_wal: false, + partition_expr_version: None, + }), + ) + .await?; + let expected = scan_region_data(&engine, region_id).await; + flush_region(&engine, region_id, None).await; + reopen_region( + &engine, + region_id, + request.table_dir, + false, + request.options, + ) + .await; + assert_eq!(expected, scan_region_data(&engine, region_id).await); Ok(()) } diff --git a/src/mito2/src/memtable.rs b/src/mito2/src/memtable.rs index afb546340b1..691ba1c5623 100644 --- a/src/mito2/src/memtable.rs +++ b/src/mito2/src/memtable.rs @@ -437,8 +437,8 @@ pub(crate) struct MemtableBuilderProvider { compact_dispatcher: Arc, } -/// Ensures JSON2 columns are not used with [`TimeSeriesMemtable`]. -pub(crate) fn ensure_json2_not_use_time_series_memtable( +/// Ensures JSON2 columns use flat SST output after resolving the SST format. +pub(crate) fn ensure_json2_uses_flat_format( metadata: &RegionMetadata, options: &RegionOptions, ) -> Result<()> { @@ -447,10 +447,11 @@ pub(crate) fn ensure_json2_not_use_time_series_memtable( .iter() .any(|x| x.column_schema.data_type.is_json2()) { + // Bulk memtables alone do not ensure the SST writer supports JSON2. ensure!( - !matches!(&options.memtable, Some(MemtableOptions::TimeSeries)), + options.sst_format == Some(FormatType::Flat), InvalidRegionOptionsSnafu { - reason: "JSON2 columns only support BulkMemtable", + reason: "JSON2 columns require sst_format=flat", } ); } @@ -994,7 +995,7 @@ mod tests { } #[test] - fn test_json2_requires_bulk_memtable() -> WhateverResult<()> { + fn test_json2_requires_flat_format() -> WhateverResult<()> { let mut metadata = sst_region_metadata(); metadata.column_metadatas[2].column_schema.data_type = ConcreteDataType::json2(JsonNativeType::Object(JsonObjectType::new())); @@ -1005,14 +1006,28 @@ mod tests { ..Default::default() }; - let err = ensure_json2_not_use_time_series_memtable(&metadata, &options).unwrap_err(); + let err = ensure_json2_uses_flat_format(&metadata, &options).unwrap_err(); assert!( err.to_string() - .contains("JSON2 columns only support BulkMemtable") + .contains("JSON2 columns require sst_format=flat") ); + options.memtable = None; + assert!(ensure_json2_uses_flat_format(&metadata, &options).is_err()); + options.memtable = Some(MemtableOptions::Bulk(BulkMemtableConfig::default())); - ensure_json2_not_use_time_series_memtable(&metadata, &options)?; + assert!(ensure_json2_uses_flat_format(&metadata, &options).is_err()); + + options.memtable = Some(MemtableOptions::TimeSeries); + options.sst_format = Some(FormatType::Flat); + ensure_json2_uses_flat_format(&metadata, &options)?; + + options.sst_format = Some(FormatType::PrimaryKey); + options.primary_key_encoding = Some(PrimaryKeyEncoding::Sparse); + assert!(ensure_json2_uses_flat_format(&metadata, &options).is_err()); + + options.sst_format = Some(FormatType::Flat); + ensure_json2_uses_flat_format(&metadata, &options)?; Ok(()) } } diff --git a/src/mito2/src/region/opener.rs b/src/mito2/src/region/opener.rs index 35ca3129fe0..d10545198f7 100644 --- a/src/mito2/src/region/opener.rs +++ b/src/mito2/src/region/opener.rs @@ -59,7 +59,7 @@ use crate::manifest::action::RegionManifest; use crate::manifest::manager::{RegionManifestManager, RegionManifestOptions}; use crate::memtable::bulk::part::BulkPart; use crate::memtable::time_partition::TimePartitions; -use crate::memtable::{MemtableBuilderProvider, ensure_json2_not_use_time_series_memtable}; +use crate::memtable::{MemtableBuilderProvider, ensure_json2_uses_flat_format}; use crate::metrics::{CACHE_FILL_DOWNLOADED_FILES, CACHE_FILL_PENDING_FILES}; use crate::read::series_candidate::is_sparse_metric_metadata; use crate::region::options::RegionOptions; @@ -337,9 +337,9 @@ impl RegionOpener { let region_id = self.region_id; let region_dir = self.region_dir(); let metadata = self.build_metadata()?; - // Tries to open the region. - match self.maybe_open(config, wal).await { - Ok(Some(region)) => { + // Only an absent region may be created; an open error must not replace its manifest. + match self.maybe_open(config, wal).await? { + Some(region) => { let recovered = region.metadata(); // Checks the schema of the region. let expect = &metadata; @@ -354,21 +354,12 @@ impl RegionOpener { return Ok(region); } - Ok(None) => { + None => { debug!( "No data under directory {}, region_id: {}", region_dir, self.region_id ); } - // The region exists but the WAL does not hold its history: - // creating it again would replace its manifest. - Err(e @ error::Error::CheckWalHistory { .. }) => return Err(e), - Err(e) => { - warn!(e; - "Failed to open region {} before creating it, region_dir: {}", - self.region_id, region_dir - ); - } } // Safety: must be set before calling this method. let mut options = self.options.take().unwrap(); @@ -386,7 +377,7 @@ impl RegionOpener { options.sst_format = Some(FormatType::PrimaryKey); FormatType::PrimaryKey }; - ensure_json2_not_use_time_series_memtable(&metadata, &options)?; + ensure_json2_uses_flat_format(&metadata, &options)?; options.validate_column_compression(&metadata)?; // The region's WAL history starts in this log store, so reopening // the region on it passes the history check without any entry. It is @@ -549,7 +540,7 @@ impl RegionOpener { let metadata = maybe_upgrade_json2_layout(metadata)?; // Updates the region options with the manifest. sanitize_region_options(&manifest, &mut region_options); - ensure_json2_not_use_time_series_memtable(&metadata, ®ion_options)?; + ensure_json2_uses_flat_format(&metadata, ®ion_options)?; region_options.validate_column_compression(&metadata)?; let region_id = self.region_id; @@ -804,23 +795,29 @@ pub(crate) fn version_builder_from_manifest( /// Updates region options by persistent options. pub(crate) fn sanitize_region_options(manifest: &RegionManifest, options: &mut RegionOptions) { - // The caller-supplied options win when they specify an explicit `sst_format`, - // so re-parsed options (e.g. bulk memtable forcing flat format) take effect on - // the running region instead of being clobbered by the persisted manifest value. - // Fall back to the manifest only when the caller did not provide a value. - match options.sst_format { - Some(format) if format != manifest.sst_format => { - common_telemetry::warn!( + // Caller options take precedence over the manifest, except JSON2 always requires flat output. + let format = options.sst_format.unwrap_or(manifest.sst_format); + // Older versions allowed JSON2 with primary-key output; recover without rewriting its history. + if format != FormatType::Flat + && manifest + .metadata + .column_metadatas + .iter() + .any(|column| column.column_schema.data_type.is_json2()) + { + warn!( + "Overriding SST format from {:?} to flat for existing JSON2 region {}", + format, manifest.metadata.region_id, + ); + options.sst_format = Some(FormatType::Flat); + } else { + if format != manifest.sst_format { + warn!( "Overriding SST format from {:?} (manifest) to {:?} (options) for region {}", - manifest.sst_format, - format, - manifest.metadata.region_id, + manifest.sst_format, format, manifest.metadata.region_id, ); } - Some(_) => {} - None => { - options.sst_format = Some(manifest.sst_format); - } + options.sst_format = Some(format); } if let Some(manifest_append_mode) = manifest.append_mode && options.append_mode != manifest_append_mode @@ -1656,6 +1653,31 @@ mod tests { assert_eq!(options.sst_format, Some(FormatType::Flat)); } + #[rstest::rstest] + fn test_sanitize_region_options_recovers_json2( + #[values(FormatType::PrimaryKey, FormatType::Flat)] persisted: FormatType, + #[values(None, Some(FormatType::PrimaryKey), Some(FormatType::Flat))] requested: Option< + FormatType, + >, + ) { + let mut manifest = build_test_manifest(persisted); + let mut metadata = manifest.metadata.as_ref().clone(); + metadata.column_metadatas[2].column_schema.data_type = + ConcreteDataType::json2(JsonNativeType::Object(JsonObjectType::new())); + manifest.metadata = Arc::new( + RegionMetadataBuilder::from_existing(metadata) + .build() + .unwrap(), + ); + let mut options = RegionOptions { + sst_format: requested, + ..Default::default() + }; + sanitize_region_options(&manifest, &mut options); + assert_eq!(Some(FormatType::Flat), options.sst_format); + assert_eq!(persisted, manifest.sst_format); + } + fn sst_parquet_bytes(batch: &RecordBatch) -> Vec { let key_value_meta = KeyValue::new( PARQUET_METADATA_KEY.to_string(), diff --git a/src/mito2/src/worker/handle_alter.rs b/src/mito2/src/worker/handle_alter.rs index 2bb1fadc7dc..1798f9a79bf 100644 --- a/src/mito2/src/worker/handle_alter.rs +++ b/src/mito2/src/worker/handle_alter.rs @@ -36,6 +36,7 @@ use store_api::storage::RegionId; use crate::error::{InvalidMetadataSnafu, InvalidRegionRequestSnafu, Result}; use crate::flush::FlushReason; use crate::manifest::action::RegionChange; +use crate::memtable::ensure_json2_uses_flat_format; use crate::region::MitoRegionRef; use crate::region::options::CompactionOptions::Twcs; use crate::region::options::{RegionOptions, TwcsOptions}; @@ -195,12 +196,16 @@ impl RegionWorkerLoop { return; } }; - // Persist the metadata to region's manifest. let options = new_options.as_ref().unwrap_or(&version.options); + if let Err(e) = ensure_json2_uses_flat_format(&new_meta, options) { + sender.send(Err(e)); + return; + } if let Err(e) = options.validate_column_compression(&new_meta) { sender.send(Err(e)); return; } + // Persist the metadata to region's manifest. let change = RegionChange { metadata: new_meta, sst_format: options.sst_format.unwrap_or_default(), diff --git a/tests/cases/standalone/common/types/json/json2_limit.result b/tests/cases/standalone/common/types/json/json2_limit.result index 555e76f1e5e..6c362a5cc62 100644 --- a/tests/cases/standalone/common/types/json/json2_limit.result +++ b/tests/cases/standalone/common/types/json/json2_limit.result @@ -112,3 +112,111 @@ drop table json2_set_append_mode_false; Affected Rows: 0 +-- JSON2 requires flat SST output even when append mode is enabled. +create table json2_primary_key_format ( + ts timestamp time index, + j json2 +) with ( + 'append_mode' = 'true', + 'sst_format' = 'primary_key' +); + +Error: 1004(InvalidArguments), Invalid region options, JSON2 columns require sst_format=flat + +create table json2_alter_primary_key_format ( + ts timestamp time index, + j json2 +) with ( + 'append_mode' = 'true', + 'memtable.type' = 'bulk' +); + +Affected Rows: 0 + +insert into json2_alter_primary_key_format values (1, '{"a": 1}'); + +Affected Rows: 1 + +alter table json2_alter_primary_key_format set 'sst_format' = 'primary_key'; + +Error: 1004(InvalidArguments), Invalid region options, JSON2 columns require sst_format=flat + +show create table json2_alter_primary_key_format; + ++--------------------------------+---------------------------------------------------------------+ +| Table | Create Table | ++--------------------------------+---------------------------------------------------------------+ +| json2_alter_primary_key_format | CREATE TABLE IF NOT EXISTS "json2_alter_primary_key_format" ( | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | "j" JSON2( | +| | max_auto_expanded_paths = 100 | +| | ) NULL, | +| | TIME INDEX ("ts") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | append_mode = 'true', | +| | 'memtable.type' = 'bulk' | +| | ) | ++--------------------------------+---------------------------------------------------------------+ + +insert into json2_alter_primary_key_format values (2, '{"a": 2}'); + +Affected Rows: 1 + +admin flush_table('json2_alter_primary_key_format'); + ++-----------------------------------------------------+ +| ADMIN flush_table('json2_alter_primary_key_format') | ++-----------------------------------------------------+ +| 0 | ++-----------------------------------------------------+ + +select ts, json_get(j, 'a') from json2_alter_primary_key_format order by ts; + ++-------------------------+------------------------------------------------------+ +| ts | json_get(json2_alter_primary_key_format.j,Utf8("a")) | ++-------------------------+------------------------------------------------------+ +| 1970-01-01T00:00:00.001 | 1 | +| 1970-01-01T00:00:00.002 | 2 | ++-------------------------+------------------------------------------------------+ + +drop table json2_alter_primary_key_format; + +Affected Rows: 0 + +create table json2_add_to_primary_key_format ( + ts timestamp time index +) with ( + 'append_mode' = 'true', + 'sst_format' = 'primary_key' +); + +Affected Rows: 0 + +alter table json2_add_to_primary_key_format add column j json2; + +Error: 1004(InvalidArguments), Invalid region options, JSON2 columns require sst_format=flat + +show create table json2_add_to_primary_key_format; + ++---------------------------------+----------------------------------------------------------------+ +| Table | Create Table | ++---------------------------------+----------------------------------------------------------------+ +| json2_add_to_primary_key_format | CREATE TABLE IF NOT EXISTS "json2_add_to_primary_key_format" ( | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | TIME INDEX ("ts") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | append_mode = 'true', | +| | sst_format = 'primary_key' | +| | ) | ++---------------------------------+----------------------------------------------------------------+ + +drop table json2_add_to_primary_key_format; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/types/json/json2_limit.sql b/tests/cases/standalone/common/types/json/json2_limit.sql index 653fdcec12c..7b428abb5e0 100644 --- a/tests/cases/standalone/common/types/json/json2_limit.sql +++ b/tests/cases/standalone/common/types/json/json2_limit.sql @@ -68,3 +68,47 @@ create table json2_set_append_mode_false ( alter table json2_set_append_mode_false set 'append_mode' = 'false'; drop table json2_set_append_mode_false; + +-- JSON2 requires flat SST output even when append mode is enabled. +create table json2_primary_key_format ( + ts timestamp time index, + j json2 +) with ( + 'append_mode' = 'true', + 'sst_format' = 'primary_key' +); + +create table json2_alter_primary_key_format ( + ts timestamp time index, + j json2 +) with ( + 'append_mode' = 'true', + 'memtable.type' = 'bulk' +); + +insert into json2_alter_primary_key_format values (1, '{"a": 1}'); + +alter table json2_alter_primary_key_format set 'sst_format' = 'primary_key'; + +show create table json2_alter_primary_key_format; + +insert into json2_alter_primary_key_format values (2, '{"a": 2}'); + +admin flush_table('json2_alter_primary_key_format'); + +select ts, json_get(j, 'a') from json2_alter_primary_key_format order by ts; + +drop table json2_alter_primary_key_format; + +create table json2_add_to_primary_key_format ( + ts timestamp time index +) with ( + 'append_mode' = 'true', + 'sst_format' = 'primary_key' +); + +alter table json2_add_to_primary_key_format add column j json2; + +show create table json2_add_to_primary_key_format; + +drop table json2_add_to_primary_key_format; From 567849908e4de7c5c7fedcf753b96f0a27427120 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Fri, 23 Oct 2026 06:13:54 +0000 Subject: [PATCH 081/239] fix(ci): specify Kind storage class for datanodes (#9908) Signed-off-by: WenyXu --- .../actions/setup-greptimedb-cluster/with-observability.yaml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/.github/actions/setup-greptimedb-cluster/with-observability.yaml b/.github/actions/setup-greptimedb-cluster/with-observability.yaml index 98659a89941..6cc923aede3 100644 --- a/.github/actions/setup-greptimedb-cluster/with-observability.yaml +++ b/.github/actions/setup-greptimedb-cluster/with-observability.yaml @@ -1,3 +1,7 @@ +datanode: + storage: + storageClassName: standard + tracing: enabled: true endpoint: "http://my-greptimedb-monitor-standalone.my-greptimedb:4000/v1/otlp/v1/traces" From 5523cf387411c9e998bac97843d60819ed29cecd Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Fri, 23 Oct 2026 06:46:56 +0000 Subject: [PATCH 082/239] feat(meta): track per-node scheduling state and fence region placement (#9884) feat: track per-node scheduling state and fence region placement Add a cluster-wide, persistent "can this node receive new regions" flag and make every metasrv placement path respect it. This is the fence half of `ADMIN drain_node`; the drain procedure, gRPC entry and region evacuation land separately. State layout - New namespace `__node_scheduling/{datanode,flownode}/{node_id}` holding `{"schedulable": false}`. A missing key means schedulable, so undrain is a delete and no stale `schedulable: true` records are left behind. - `NodeSchedulingKey`/`NodeSchedulingValue` mirror the existing KV conventions (`NodeInfoKey`'s `{role, node_id}` shape, `LeaseValue`'s `*Value` suffix). The role is part of the prefix, so extending to flownodes needs no key change. - `NodeSchedulingManager` writes with `put_if_not_exists`/`compare_and_put`, so concurrent marks change the state exactly once, and reads a single prefix range per placement decision (DDL/failover only, not per query). Placement - `NodeInfo` gains `#[serde(skip)] schedulable: Option` plus `is_schedulable()`; `None` (unresolved) counts as schedulable, and the field never reaches heartbeat bytes, so the wire/in-memory format is unchanged. - `SchedulingAwarePeerDiscovery` stamps the flag on `active_datanodes`. It is additive (never hides nodes) because failover still needs the full active set to compute its exclusion list. - All three built-in selectors now fetch candidates through `selector::common::active_schedulable_datanodes`; `RegionSupervisor` adds unschedulable datanodes to `exclude_peer_ids` for both selector variants. Enforcing the single fence - `MetaPeerClient` no longer implements `PeerDiscovery`; its discovery methods are inherent, and `DatabaseOperator` (frontends only) takes the concrete client. `SchedulingAwarePeerDiscovery`, built once in `MetasrvBuilder`, is therefore the only way to obtain a `PeerDiscoveryRef`, so bypassing the fence is a compile error rather than a silent misplacement. Compatibility - New key namespace only; `NodeInfo` serialization is byte-identical, so old metasrv nodes and persisted heartbeats keep working. The flag itself is stored in the durable KV, so a follower metasrv sees marks written by the leader. Verification - `cargo nextest run -p meta-srv node_scheduling` (13), `selector` (12), `supervisor` (14), `-p common-meta cluster` (9), `-p meta-client` (36), `-p standalone` (compiles, no tests). - `cargo clippy -p meta-srv -p common-meta -p meta-client -p standalone --all-targets -- -D warnings`, `cargo fmt --all -- --check`. Signed-off-by: WenyXu --- src/common/meta/src/affinity.rs | 1 + src/common/meta/src/cluster.rs | 78 ++++ src/meta-client/src/client.rs | 4 + src/meta-client/src/client/util.rs | 1 + src/meta-srv/src/discovery.rs | 19 +- src/meta-srv/src/discovery/lease.rs | 5 +- src/meta-srv/src/discovery/node_scheduling.rs | 151 +++++++ src/meta-srv/src/error.rs | 8 + .../handler/collect_cluster_info_handler.rs | 1 + src/meta-srv/src/key.rs | 2 + src/meta-srv/src/key/node_scheduling.rs | 153 +++++++ src/meta-srv/src/lib.rs | 1 + src/meta-srv/src/metasrv.rs | 11 + src/meta-srv/src/metasrv/builder.rs | 81 +++- src/meta-srv/src/node_scheduling.rs | 401 ++++++++++++++++++ src/meta-srv/src/peer.rs | 4 +- src/meta-srv/src/region/supervisor.rs | 87 +++- src/meta-srv/src/selector/common.rs | 23 +- src/meta-srv/src/selector/lease_based.rs | 12 +- src/meta-srv/src/selector/load_based.rs | 12 +- src/meta-srv/src/selector/round_robin.rs | 41 +- src/meta-srv/src/selector/test_utils.rs | 10 +- src/meta-srv/src/test_util.rs | 17 + src/meta-srv/src/utils/database.rs | 1 + src/standalone/src/information_extension.rs | 1 + 25 files changed, 1059 insertions(+), 66 deletions(-) create mode 100644 src/meta-srv/src/discovery/node_scheduling.rs create mode 100644 src/meta-srv/src/key/node_scheduling.rs create mode 100644 src/meta-srv/src/node_scheduling.rs diff --git a/src/common/meta/src/affinity.rs b/src/common/meta/src/affinity.rs index 2bcb0e934f0..5c8bb73f661 100644 --- a/src/common/meta/src/affinity.rs +++ b/src/common/meta/src/affinity.rs @@ -171,6 +171,7 @@ mod tests { memory_usage_bytes: 0, hostname: String::new(), env_vars, + schedulable: None, } } diff --git a/src/common/meta/src/cluster.rs b/src/common/meta/src/cluster.rs index 4e7ffb7f920..2e0631622b8 100644 --- a/src/common/meta/src/cluster.rs +++ b/src/common/meta/src/cluster.rs @@ -138,6 +138,22 @@ pub struct NodeInfo { /// Environment variables reported by the node. #[serde(default)] pub env_vars: HashMap, + /// Whether metasrv may place new regions on this node. + /// + /// `None` means the flag was not resolved (deserialized heartbeat data, tests, + /// standalone); callers treat `None` as schedulable. Filled by the metasrv + /// discovery decorator, never by the heartbeat writer. + #[serde(skip)] + pub schedulable: Option, +} + +impl NodeInfo { + /// Returns `false` only if the node is explicitly marked unschedulable. + /// + /// `None` (unresolved) counts as schedulable. + pub fn is_schedulable(&self) -> bool { + self.schedulable != Some(false) + } } #[derive(Debug, Clone, Copy, Eq, Hash, PartialEq, Serialize, Deserialize, PartialOrd, Ord)] @@ -324,6 +340,35 @@ mod tests { use crate::cluster::{DatanodeStatus, FlownodeStatus, NodeInfo, NodeInfoKey, NodeStatus}; use crate::peer::Peer; + fn new_test_node_info() -> NodeInfo { + NodeInfo { + peer: Peer { + id: 1, + addr: "127.0.0.1".to_string(), + }, + last_activity_ts: 123, + status: NodeStatus::Datanode(DatanodeStatus { + rcus: 1, + wcus: 2, + leader_regions: 3, + follower_regions: 4, + workloads: DatanodeWorkloads { + types: vec![DatanodeWorkloadType::Hybrid.to_i32()], + }, + }), + version: "".to_string(), + git_commit: "".to_string(), + start_time_ms: 1, + total_cpu_millicores: 0, + total_memory_bytes: 0, + cpu_usage_millicores: 0, + memory_usage_bytes: 0, + hostname: "test_hostname".to_string(), + env_vars: Default::default(), + schedulable: None, + } + } + #[test] fn test_node_info_key_round_trip() { let key = NodeInfoKey { @@ -364,10 +409,16 @@ mod tests { memory_usage_bytes: 0, hostname: "test_hostname".to_string(), env_vars: Default::default(), + schedulable: None, }; let node_info_bytes: Vec = node_info.try_into().unwrap(); + assert!( + !String::from_utf8_lossy(&node_info_bytes).contains("schedulable"), + "the runtime scheduling flag must not be persisted" + ); let new_node_info: NodeInfo = node_info_bytes.try_into().unwrap(); + assert!(new_node_info.schedulable.is_none()); assert_matches!( new_node_info, @@ -387,6 +438,31 @@ mod tests { ); } + #[test] + fn test_is_schedulable() { + assert!( + NodeInfo { + schedulable: None, + ..new_test_node_info() + } + .is_schedulable() + ); + assert!( + NodeInfo { + schedulable: Some(true), + ..new_test_node_info() + } + .is_schedulable() + ); + assert!( + !NodeInfo { + schedulable: Some(false), + ..new_test_node_info() + } + .is_schedulable() + ); + } + #[test] fn test_node_info_key_prefix() { let prefix = NodeInfoKey::key_prefix(); @@ -461,6 +537,7 @@ mod tests { memory_usage_bytes: 0, hostname: "test_hostname".to_string(), env_vars: Default::default(), + schedulable: None, }; let node_info_bytes: Vec = node_info.try_into().unwrap(); @@ -523,6 +600,7 @@ mod tests { memory_usage_bytes: 0, hostname: "test".to_string(), env_vars, + schedulable: None, }; let node_info_bytes: Vec = node_info.try_into().unwrap(); diff --git a/src/meta-client/src/client.rs b/src/meta-client/src/client.rs index 44b6b21dd23..bc147c07e5e 100644 --- a/src/meta-client/src/client.rs +++ b/src/meta-client/src/client.rs @@ -509,6 +509,7 @@ impl ClusterInfo for MetaClient { memory_usage_bytes: node_info.memory_usage_bytes, hostname: node_info.hostname, env_vars: Default::default(), + schedulable: None, } } else { // TODO(zyy17): It's for backward compatibility. Remove this when the deprecated fields are removed from the proto. @@ -525,6 +526,7 @@ impl ClusterInfo for MetaClient { memory_usage_bytes: 0, hostname: "".to_string(), env_vars: Default::default(), + schedulable: None, } } }) @@ -543,6 +545,7 @@ impl ClusterInfo for MetaClient { memory_usage_bytes: node_info.memory_usage_bytes, hostname: node_info.hostname, env_vars: Default::default(), + schedulable: None, } } else { // TODO(zyy17): It's for backward compatibility. Remove this when the deprecated fields are removed from the proto. @@ -559,6 +562,7 @@ impl ClusterInfo for MetaClient { memory_usage_bytes: 0, hostname: "".to_string(), env_vars: Default::default(), + schedulable: None, } } })) diff --git a/src/meta-client/src/client/util.rs b/src/meta-client/src/client/util.rs index 62566f43783..649601741ed 100644 --- a/src/meta-client/src/client/util.rs +++ b/src/meta-client/src/client/util.rs @@ -169,6 +169,7 @@ mod tests { memory_usage_bytes: 0, hostname: String::new(), env_vars: Default::default(), + schedulable: None, } } diff --git a/src/meta-srv/src/discovery.rs b/src/meta-srv/src/discovery.rs index b0d14e4a948..0e0e476bb8f 100644 --- a/src/meta-srv/src/discovery.rs +++ b/src/meta-srv/src/discovery.rs @@ -14,6 +14,7 @@ pub mod lease; pub mod node_info; +pub mod node_scheduling; pub mod utils; use api::v1::meta::heartbeat_request::NodeWorkloads; @@ -21,7 +22,7 @@ use common_error::ext::BoxedError; use common_meta::cluster::NodeInfo; use common_meta::distributed_time_constants::default_distributed_time_constants; use common_meta::error::Result; -use common_meta::peer::{Peer, PeerDiscovery, PeerResolver}; +use common_meta::peer::{Peer, PeerResolver}; use common_meta::{DatanodeId, FlownodeId}; use common_time::util::DefaultSystemTimer; use snafu::ResultExt; @@ -29,9 +30,13 @@ use snafu::ResultExt; use crate::cluster::MetaPeerClient; use crate::discovery::lease::{LeaseValueAccessor, LeaseValueType}; -#[async_trait::async_trait] -impl PeerDiscovery for MetaPeerClient { - async fn active_frontends(&self) -> Result> { +// NOTE: these are inherent methods on purpose. `MetaPeerClient` does not implement +// `PeerDiscovery`, so the only way to obtain a `PeerDiscoveryRef` inside metasrv is +// `discovery::SchedulingAwarePeerDiscovery`, which stamps node scheduling states. Wrapping the raw +// client directly would silently drop the drain fence. +impl MetaPeerClient { + /// Returns all frontends that reported a heartbeat recently. + pub async fn active_frontends(&self) -> Result> { utils::alive_frontend_infos( &DefaultSystemTimer, self, @@ -42,7 +47,8 @@ impl PeerDiscovery for MetaPeerClient { .context(common_meta::error::ExternalSnafu) } - async fn active_datanodes( + /// Returns all datanodes that reported a heartbeat recently, without scheduling states. + pub async fn active_datanodes( &self, filter: Option fn(&'a NodeWorkloads) -> bool>, ) -> Result> { @@ -57,7 +63,8 @@ impl PeerDiscovery for MetaPeerClient { .context(common_meta::error::ExternalSnafu) } - async fn active_flownodes( + /// Returns all flownodes that reported a heartbeat recently. + pub async fn active_flownodes( &self, filter: Option fn(&'a NodeWorkloads) -> bool>, ) -> Result> { diff --git a/src/meta-srv/src/discovery/lease.rs b/src/meta-srv/src/discovery/lease.rs index 585676d41b5..b40bee2dcf2 100644 --- a/src/meta-srv/src/discovery/lease.rs +++ b/src/meta-srv/src/discovery/lease.rs @@ -103,7 +103,7 @@ mod tests { }; use common_meta::distributed_time_constants::default_distributed_time_constants; use common_meta::kv_backend::ResettableKvBackendRef; - use common_meta::peer::{Peer, PeerDiscovery}; + use common_meta::peer::Peer; use common_meta::rpc::store::PutRequest; use common_time::util::current_time_millis; use common_workload::DatanodeWorkloadType; @@ -157,6 +157,7 @@ mod tests { memory_usage_bytes: 0, hostname: String::new(), env_vars, + schedulable: None, }, ) .await; @@ -199,6 +200,7 @@ mod tests { memory_usage_bytes: 0, hostname: String::new(), env_vars: Default::default(), + schedulable: None, }, ) .await; @@ -236,6 +238,7 @@ mod tests { memory_usage_bytes: 0, hostname: "test_hostname".to_string(), env_vars: Default::default(), + schedulable: None, }; let key_prefix = NodeInfoKey::key_prefix_with_role(Role::Frontend); diff --git a/src/meta-srv/src/discovery/node_scheduling.rs b/src/meta-srv/src/discovery/node_scheduling.rs new file mode 100644 index 00000000000..bfcaa0e5965 --- /dev/null +++ b/src/meta-srv/src/discovery/node_scheduling.rs @@ -0,0 +1,151 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use api::v1::meta::heartbeat_request::NodeWorkloads; +use common_error::ext::BoxedError; +use common_meta::cluster::NodeInfo; +use common_meta::error::Result; +use common_meta::peer::PeerDiscovery; +use snafu::ResultExt; + +use crate::cluster::MetaPeerClientRef; +use crate::key::NodeSchedulingRole; +use crate::node_scheduling::NodeSchedulingManagerRef; + +/// A [`PeerDiscovery`] decorator that stamps the scheduling state of datanodes +/// onto [`NodeInfo::schedulable`]. +/// +/// It never hides nodes: callers that place regions must filter with +/// `selector::common::active_schedulable_datanodes`. Keeping it additive is what allows it to be +/// the only peer discovery in the placement path (built once by `MetasrvBuilder`), because +/// failover still needs the full active set to compute its exclusions. +/// +/// This is the only `PeerDiscovery` implementation over metasrv's local state: the raw +/// [`MetaPeerClient`](crate::cluster::MetaPeerClient) only exposes inherent discovery methods, so +/// a `PeerDiscoveryRef` cannot be built without the scheduling fence. Do not build a second +/// instance; `MetasrvBuilder` owns the single production one and exposes it as +/// `Metasrv::selector_ctx().peer_discovery`. +pub struct SchedulingAwarePeerDiscovery { + manager: NodeSchedulingManagerRef, + inner: MetaPeerClientRef, +} + +impl SchedulingAwarePeerDiscovery { + pub fn new(manager: NodeSchedulingManagerRef, inner: MetaPeerClientRef) -> Self { + Self { manager, inner } + } +} + +#[async_trait::async_trait] +impl PeerDiscovery for SchedulingAwarePeerDiscovery { + async fn active_frontends(&self) -> Result> { + self.inner.active_frontends().await + } + + async fn active_datanodes( + &self, + filter: Option fn(&'a NodeWorkloads) -> bool>, + ) -> Result> { + let mut datanodes = self.inner.active_datanodes(filter).await?; + if datanodes.is_empty() { + return Ok(datanodes); + } + + let unschedulable = self + .manager + .list_unschedulable_node_ids(NodeSchedulingRole::Datanode) + .await + .map_err(BoxedError::new) + .context(common_meta::error::ExternalSnafu)?; + for datanode in &mut datanodes { + datanode.schedulable = Some(!unschedulable.contains(&datanode.peer.id)); + } + Ok(datanodes) + } + + async fn active_flownodes( + &self, + filter: Option fn(&'a NodeWorkloads) -> bool>, + ) -> Result> { + self.inner.active_flownodes(filter).await + } +} + +#[cfg(test)] +mod tests { + use common_meta::peer::Peer; + + use super::*; + use crate::cluster::MetaPeerClientRef; + use crate::test_util::{create_meta_peer_client, new_node_scheduling_discovery, put_datanodes}; + + async fn new_discovery() -> ( + MetaPeerClientRef, + NodeSchedulingManagerRef, + SchedulingAwarePeerDiscovery, + ) { + let client = create_meta_peer_client(); + put_datanodes( + &client, + vec![Peer::empty(1), Peer::empty(2), Peer::empty(3)], + ) + .await; + let (manager, discovery) = new_node_scheduling_discovery(&client); + (client, manager, discovery) + } + + #[tokio::test] + async fn test_active_datanodes_marks_unschedulable() { + let (_, manager, discovery) = new_discovery().await; + + manager + .mark_unschedulable(NodeSchedulingRole::Datanode, 2) + .await + .unwrap(); + + let datanodes = discovery.active_datanodes(None).await.unwrap(); + assert_eq!(3, datanodes.len()); + + let schedulable = |id: u64| { + datanodes + .iter() + .find(|datanode| datanode.peer.id == id) + .unwrap() + .schedulable + }; + assert_eq!(Some(true), schedulable(1)); + assert_eq!(Some(false), schedulable(2)); + assert_eq!(Some(true), schedulable(3)); + } + + #[tokio::test] + async fn test_schedulable_marks_are_not_persisted() { + let (client, manager, discovery) = new_discovery().await; + + manager + .mark_unschedulable(NodeSchedulingRole::Datanode, 2) + .await + .unwrap(); + + discovery.active_datanodes(None).await.unwrap(); + + let datanodes = client.active_datanodes(None).await.unwrap(); + assert_eq!(3, datanodes.len()); + assert!( + datanodes + .iter() + .all(|datanode| datanode.schedulable.is_none()) + ); + } +} diff --git a/src/meta-srv/src/error.rs b/src/meta-srv/src/error.rs index 70997ebe7d1..011576f9b27 100644 --- a/src/meta-srv/src/error.rs +++ b/src/meta-srv/src/error.rs @@ -335,6 +335,13 @@ pub enum Error { location: Location, }, + #[snafu(display("Invalid node scheduling key: {}", key))] + InvalidNodeSchedulingKey { + #[snafu(implicit)] + location: Location, + key: String, + }, + #[snafu(display("Failed to parse lease key from utf8"))] LeaseKeyFromUtf8 { #[snafu(source)] @@ -1217,6 +1224,7 @@ impl ErrorExt for Error { | Error::InvalidLeaseKey { .. } | Error::InvalidStatKey { .. } | Error::InvalidInactiveRegionKey { .. } + | Error::InvalidNodeSchedulingKey { .. } | Error::ParseNum { .. } | Error::ParseBool { .. } | Error::ParseAddr { .. } diff --git a/src/meta-srv/src/handler/collect_cluster_info_handler.rs b/src/meta-srv/src/handler/collect_cluster_info_handler.rs index 12fee395136..f884171a931 100644 --- a/src/meta-srv/src/handler/collect_cluster_info_handler.rs +++ b/src/meta-srv/src/handler/collect_cluster_info_handler.rs @@ -196,6 +196,7 @@ impl NodeInfoUpdate { memory_usage_bytes: info.memory_usage_bytes, hostname: info.hostname, env_vars: self.env_vars, + schedulable: None, }; let value = value.try_into().context(InvalidClusterInfoFormatSnafu)?; ctx.in_memory diff --git a/src/meta-srv/src/key.rs b/src/meta-srv/src/key.rs index abdc92ff6d4..eea9573580a 100644 --- a/src/meta-srv/src/key.rs +++ b/src/meta-srv/src/key.rs @@ -14,12 +14,14 @@ mod datanode; mod flownode; +mod node_scheduling; use std::str::FromStr; use api::v1::meta::heartbeat_request::NodeWorkloads; pub use datanode::*; pub use flownode::*; +pub use node_scheduling::*; use serde::{Deserialize, Serialize}; use snafu::{OptionExt, ResultExt, ensure}; diff --git a/src/meta-srv/src/key/node_scheduling.rs b/src/meta-srv/src/key/node_scheduling.rs new file mode 100644 index 00000000000..1ec398871c7 --- /dev/null +++ b/src/meta-srv/src/key/node_scheduling.rs @@ -0,0 +1,153 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::str::FromStr; + +use snafu::OptionExt; + +use crate::error; +use crate::error::Result; + +/// Key prefix of datanode scheduling states. +const DATANODE_SCHEDULING_PREFIX: &str = "__node_scheduling/datanode/"; +/// Key prefix of flownode scheduling states. +const FLOWNODE_SCHEDULING_PREFIX: &str = "__node_scheduling/flownode/"; + +/// The node roles whose scheduling state is tracked. +/// +/// A subset of [`common_meta::cluster::Role`]: frontends and metasrv nodes never receive +/// regions, so they are out of scope. +#[derive(Debug, Clone, Copy, Eq, Hash, PartialEq)] +pub enum NodeSchedulingRole { + Datanode, + Flownode, +} + +impl NodeSchedulingRole { + /// Returns the key prefix of this role, e.g. `__node_scheduling/datanode/`. + fn key_prefix(&self) -> &'static str { + match self { + NodeSchedulingRole::Datanode => DATANODE_SCHEDULING_PREFIX, + NodeSchedulingRole::Flownode => FLOWNODE_SCHEDULING_PREFIX, + } + } +} + +/// Key of a node's scheduling state: `__node_scheduling/{role}/{node_id}`, +/// for example `__node_scheduling/datanode/7`. +#[derive(Debug, Clone, Copy, Eq, Hash, PartialEq)] +pub struct NodeSchedulingKey { + pub role: NodeSchedulingRole, + pub node_id: u64, +} + +impl NodeSchedulingKey { + pub fn new(role: NodeSchedulingRole, node_id: u64) -> NodeSchedulingKey { + NodeSchedulingKey { role, node_id } + } + + /// Returns the prefix covering every node of `role`. + pub fn prefix_key(role: NodeSchedulingRole) -> Vec { + role.key_prefix().as_bytes().to_vec() + } +} + +impl From for Vec { + fn from(key: NodeSchedulingKey) -> Vec { + format!("{}{}", key.role.key_prefix(), key.node_id).into_bytes() + } +} + +impl FromStr for NodeSchedulingKey { + type Err = error::Error; + + fn from_str(key: &str) -> Result { + for role in [NodeSchedulingRole::Datanode, NodeSchedulingRole::Flownode] { + let Some(node_id) = key.strip_prefix(role.key_prefix()) else { + continue; + }; + + let node_id = node_id + .parse::() + .ok() + .context(error::InvalidNodeSchedulingKeySnafu { key })?; + return Ok(NodeSchedulingKey { role, node_id }); + } + + error::InvalidNodeSchedulingKeySnafu { key }.fail() + } +} + +impl TryFrom> for NodeSchedulingKey { + type Error = error::Error; + + fn try_from(bytes: Vec) -> Result { + let key = String::from_utf8(bytes).map_err(|err| { + error::InvalidNodeSchedulingKeySnafu { + key: String::from_utf8_lossy(err.as_bytes()).into_owned(), + } + .build() + })?; + key.parse() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_key_round_trip() { + for (role, expected) in [ + (NodeSchedulingRole::Datanode, "__node_scheduling/datanode/7"), + (NodeSchedulingRole::Flownode, "__node_scheduling/flownode/7"), + ] { + let key = NodeSchedulingKey::new(role, 7); + let bytes: Vec = key.into(); + assert_eq!(expected.as_bytes(), bytes.as_slice()); + + let parsed = NodeSchedulingKey::try_from(bytes).unwrap(); + assert_eq!(key, parsed); + assert_eq!(role, parsed.role); + } + } + + #[test] + fn test_key_prefix() { + assert_eq!( + b"__node_scheduling/datanode/", + NodeSchedulingKey::prefix_key(NodeSchedulingRole::Datanode).as_slice() + ); + assert_eq!( + b"__node_scheduling/flownode/", + NodeSchedulingKey::prefix_key(NodeSchedulingRole::Flownode).as_slice() + ); + } + + #[test] + fn test_invalid_key() { + for key in [ + "__node_scheduling/frontend/7", + "__node_scheduling/datanode/", + "__node_scheduling/datanode/x", + "__node_scheduling", + ] { + let err = key.parse::().unwrap_err(); + assert!( + matches!(err, error::Error::InvalidNodeSchedulingKey { .. }), + "unexpected error for {key}: {err:?}" + ); + } + } +} diff --git a/src/meta-srv/src/lib.rs b/src/meta-srv/src/lib.rs index 948cab240d4..9ae4e6db59b 100644 --- a/src/meta-srv/src/lib.rs +++ b/src/meta-srv/src/lib.rs @@ -26,6 +26,7 @@ pub mod metasrv; pub mod metrics; #[cfg(feature = "mock")] pub mod mocks; +pub mod node_scheduling; pub mod peer; pub mod procedure; pub mod pubsub; diff --git a/src/meta-srv/src/metasrv.rs b/src/meta-srv/src/metasrv.rs index f9b7a8dd5d7..fd63deaa75a 100644 --- a/src/meta-srv/src/metasrv.rs +++ b/src/meta-srv/src/metasrv.rs @@ -74,6 +74,7 @@ use crate::error::{ use crate::failure_detector::PhiAccrualFailureDetectorOptions; use crate::gc::{GcSchedulerOptions, GcTickerRef}; use crate::handler::{HeartbeatHandlerGroupBuilder, HeartbeatHandlerGroupRef}; +use crate::node_scheduling::NodeSchedulingManagerRef; use crate::procedure::ProcedureManagerListenerAdapter; use crate::procedure::region_migration::manager::RegionMigrationManagerRef; use crate::procedure::repartition::gc_requirement::RepartitionGcRequirementManagerRef; @@ -611,6 +612,7 @@ pub struct Metasrv { resource_stat: ResourceStatRef, gc_ticker: Option, database_operator: DatabaseOperatorRef, + node_scheduling_manager: NodeSchedulingManagerRef, plugins: Plugins, } @@ -852,6 +854,10 @@ impl Metasrv { &self.kv_backend } + /// Returns the undecorated peer client. + /// + /// Its `active_datanodes` does not know about node scheduling states. Use + /// [`Metasrv::selector_ctx`] for any consumer that places regions or picks datanodes. pub fn meta_peer_client(&self) -> &MetaPeerClientRef { &self.meta_peer_client } @@ -860,6 +866,7 @@ impl Metasrv { &self.selector } + /// Returns the selector context, whose `peer_discovery` stamps node scheduling states. pub fn selector_ctx(&self) -> &SelectorContext { &self.selector_ctx } @@ -924,6 +931,10 @@ impl Metasrv { &self.database_operator } + pub fn node_scheduling_manager(&self) -> &NodeSchedulingManagerRef { + &self.node_scheduling_manager + } + pub fn plugins(&self) -> &Plugins { &self.plugins } diff --git a/src/meta-srv/src/metasrv/builder.rs b/src/meta-srv/src/metasrv/builder.rs index 93ce228a156..55fdf457044 100644 --- a/src/meta-srv/src/metasrv/builder.rs +++ b/src/meta-srv/src/metasrv/builder.rs @@ -56,6 +56,7 @@ use store_api::storage::MAX_REGION_SEQ; use crate::bootstrap::build_default_meta_peer_client; use crate::cache_invalidator::MetasrvCacheInvalidator; use crate::cluster::MetaPeerClientRef; +use crate::discovery::node_scheduling::SchedulingAwarePeerDiscovery; use crate::error::{self, BuildWalProviderSnafu, OtherSnafu, Result}; use crate::event::EventHandlerImpl; use crate::gc::{DefaultGcSchedulerCtx, GcScheduler}; @@ -70,6 +71,7 @@ use crate::metasrv::{ RegionStatAwareSelectorRef, SelectTarget, SelectorContext, SelectorFactoryRef, SelectorRef, TABLE_ID_SEQ, }; +use crate::node_scheduling::NodeSchedulingManager; use crate::peer::MetasrvPeerAllocator; use crate::procedure::region_migration::DefaultContextFactory; use crate::procedure::region_migration::manager::RegionMigrationManager; @@ -238,7 +240,31 @@ impl MetasrvBuilder { let meta_peer_client = meta_peer_client .unwrap_or_else(|| build_default_meta_peer_client(&election, &in_memory)); - let database_operator = Arc::new(DatabaseOperator::new(meta_peer_client.clone())); + + let node_scheduling_manager = Arc::new(NodeSchedulingManager::new(kv_backend.clone())); + + // The only place a node scheduling decorator is built. Every consumer that reads the + // datanode list (region placement, failover exclusions) must go through + // `Metasrv::selector_ctx().peer_discovery`, and building a `PeerDiscoveryRef` from the raw + // client is impossible by construction. Bypassing this decorator silently drops the + // drain fence, so do not build a second instance. + let selector_ctx = SelectorContext { + affinity_fallback: plugins + .as_ref() + .and_then(|plugins| plugins.get::()) + .unwrap_or_default(), + affinity: plugins + .as_ref() + .and_then(|plugins| plugins.get::()) + .unwrap_or_else(|| Arc::new(NoopAffinity)), + peer_discovery: Arc::new(SchedulingAwarePeerDiscovery::new( + node_scheduling_manager.clone(), + meta_peer_client.clone(), + )), + }; + // Frontends only; the decorator passes them through unchanged. + let database_operator = + Arc::new(DatabaseOperator::new(selector_ctx.peer_discovery.clone())); let event_inserter = Box::new(InsertForwarder::new( database_operator.clone(), @@ -275,18 +301,6 @@ impl MetasrvBuilder { leader_cached_kv_backend.clone() as _, )); - let selector_ctx = SelectorContext { - affinity_fallback: plugins - .as_ref() - .and_then(|plugins| plugins.get::()) - .unwrap_or_default(), - affinity: plugins - .as_ref() - .and_then(|plugins| plugins.get::()) - .unwrap_or_else(|| Arc::new(NoopAffinity)), - peer_discovery: meta_peer_client.clone(), - }; - let wal_provider = build_wal_provider(&options.wal, kv_backend.clone()) .await .context(BuildWalProviderSnafu)?; @@ -679,6 +693,7 @@ impl MetasrvBuilder { resource_stat: Arc::new(resource_stat), gc_ticker, database_operator, + node_scheduling_manager, }) } } @@ -917,4 +932,44 @@ mod tests { &resolver )); } + + #[tokio::test] + async fn test_builder_wires_node_scheduling_discovery() { + use common_meta::peer::Peer; + + use crate::key::NodeSchedulingRole; + use crate::test_util::put_datanodes; + + let metasrv = MetasrvBuilder::new() + .options(MetasrvOptions { + enable_telemetry: false, + ..Default::default() + }) + .build() + .await + .unwrap(); + + put_datanodes( + metasrv.meta_peer_client(), + vec![Peer::empty(1), Peer::empty(2)], + ) + .await; + metasrv + .node_scheduling_manager() + .mark_unschedulable(NodeSchedulingRole::Datanode, 2) + .await + .unwrap(); + + let datanodes = metasrv + .selector_ctx() + .peer_discovery + .active_datanodes(None) + .await + .unwrap(); + assert_eq!(2, datanodes.len()); + for datanode in datanodes { + let expected = datanode.peer.id != 2; + assert_eq!(Some(expected), datanode.schedulable); + } + } } diff --git a/src/meta-srv/src/node_scheduling.rs b/src/meta-srv/src/node_scheduling.rs new file mode 100644 index 00000000000..cb7c34875ec --- /dev/null +++ b/src/meta-srv/src/node_scheduling.rs @@ -0,0 +1,401 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::HashSet; +use std::sync::Arc; + +use common_meta::kv_backend::KvBackendRef; +use common_meta::kv_backend::txn::{Compare, CompareOp, Txn, TxnOp}; +use common_meta::rpc::store::RangeRequest; +use serde::{Deserialize, Serialize}; +use snafu::ResultExt; + +use crate::error::{self, Result}; +use crate::key::{NodeSchedulingKey, NodeSchedulingRole}; + +pub type NodeSchedulingManagerRef = Arc; + +/// The scheduling state persisted for a node. +/// +/// Only `schedulable: false` is written today; a missing key already means schedulable. +/// New metadata fields must use `#[serde(default)]`. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct NodeSchedulingValue { + /// Whether the node may receive new regions. + pub schedulable: bool, +} + +/// Persists which nodes may receive new regions. +/// +/// A missing key means the node is schedulable; a present key stores a +/// [`NodeSchedulingValue`] whose `schedulable` field decides placement. +pub struct NodeSchedulingManager { + kv_backend: KvBackendRef, +} + +impl NodeSchedulingManager { + pub fn new(kv_backend: KvBackendRef) -> Self { + Self { kv_backend } + } + + /// Marks the node as unschedulable. Idempotent. + /// + /// Returns whether the stored state changed. + pub async fn mark_unschedulable(&self, role: NodeSchedulingRole, node_id: u64) -> Result { + let key: Vec = NodeSchedulingKey::new(role, node_id).into(); + let value = serde_json::to_vec(&NodeSchedulingValue { schedulable: false }).context( + error::SerializeToJsonSnafu { + input: String::from_utf8_lossy(&key).into_owned(), + }, + )?; + + let txn = match self + .kv_backend + .get(&key) + .await + .context(error::KvBackendSnafu)? + { + None => Txn::put_if_not_exists(key, value), + Some(kv) if kv.value == value => return Ok(false), + Some(kv) => Txn::compare_and_put(key, kv.value, value), + }; + let response = self + .kv_backend + .txn(txn) + .await + .context(error::KvBackendSnafu)?; + Ok(response.succeeded) + } + + /// Marks the node as schedulable by deleting its state. Idempotent. + /// + /// Returns whether a state was removed; a node without state is already schedulable, + /// so this returns `false` and writes nothing. + pub async fn mark_schedulable(&self, role: NodeSchedulingRole, node_id: u64) -> Result { + let key: Vec = NodeSchedulingKey::new(role, node_id).into(); + // `NotEqual` against a missing target means "the key exists". + let txn = Txn::new() + .when(vec![Compare::new(key.clone(), CompareOp::NotEqual, None)]) + .and_then(vec![TxnOp::Delete(key)]); + let response = self + .kv_backend + .txn(txn) + .await + .context(error::KvBackendSnafu)?; + Ok(response.succeeded) + } + + /// Returns `false` iff the node is marked unschedulable. + pub async fn is_schedulable(&self, role: NodeSchedulingRole, node_id: u64) -> Result { + let key = NodeSchedulingKey::new(role, node_id); + Ok(self + .get_value(&key) + .await? + .is_none_or(|value| value.schedulable)) + } + + /// Returns the ids of nodes of `role` that must not receive new regions. + /// + /// Returns a set because callers test membership per node. + pub async fn list_unschedulable_node_ids( + &self, + role: NodeSchedulingRole, + ) -> Result> { + let request = RangeRequest::new().with_prefix(NodeSchedulingKey::prefix_key(role)); + let response = self + .kv_backend + .range(request) + .await + .context(error::KvBackendSnafu)?; + + let mut node_ids = HashSet::new(); + for kv in response.kvs { + let key = NodeSchedulingKey::try_from(kv.key)?; + if !parse_value(&kv.value)?.schedulable { + node_ids.insert(key.node_id); + } + } + Ok(node_ids) + } + + async fn get_value(&self, key: &NodeSchedulingKey) -> Result> { + let key: Vec = (*key).into(); + let Some(kv) = self + .kv_backend + .get(&key) + .await + .context(error::KvBackendSnafu)? + else { + return Ok(None); + }; + + parse_value(&kv.value).map(Some) + } +} + +fn parse_value(value: &[u8]) -> Result { + serde_json::from_slice(value).context(error::DeserializeFromJsonSnafu { + input_len: value.len(), + }) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use common_meta::kv_backend::KvBackend; + use common_meta::kv_backend::memory::MemoryKvBackend; + + use super::*; + use crate::metasrv::SelectorContext; + + fn new_manager() -> ( + Arc>, + NodeSchedulingManager, + ) { + let kv_backend = Arc::new(MemoryKvBackend::new()); + let manager = NodeSchedulingManager::new(kv_backend.clone()); + (kv_backend, manager) + } + + #[tokio::test] + async fn test_mark_unschedulable_is_idempotent() { + let (kv_backend, manager) = new_manager(); + + assert!( + manager + .is_schedulable(NodeSchedulingRole::Datanode, 1) + .await + .unwrap() + ); + assert!( + manager + .mark_unschedulable(NodeSchedulingRole::Datanode, 1) + .await + .unwrap() + ); + assert!( + !manager + .mark_unschedulable(NodeSchedulingRole::Datanode, 1) + .await + .unwrap() + ); + + assert!( + !manager + .is_schedulable(NodeSchedulingRole::Datanode, 1) + .await + .unwrap() + ); + assert_eq!( + HashSet::from([1]), + manager + .list_unschedulable_node_ids(NodeSchedulingRole::Datanode) + .await + .unwrap() + ); + + let key: Vec = NodeSchedulingKey::new(NodeSchedulingRole::Datanode, 1).into(); + let kv = kv_backend.get(&key).await.unwrap().unwrap(); + assert_eq!(br#"{"schedulable":false}"#, kv.value.as_slice()); + + assert!( + manager + .mark_schedulable(NodeSchedulingRole::Datanode, 1) + .await + .unwrap() + ); + assert!( + !manager + .mark_schedulable(NodeSchedulingRole::Datanode, 1) + .await + .unwrap() + ); + + assert!( + manager + .is_schedulable(NodeSchedulingRole::Datanode, 1) + .await + .unwrap() + ); + assert!( + manager + .list_unschedulable_node_ids(NodeSchedulingRole::Datanode) + .await + .unwrap() + .is_empty() + ); + } + + #[tokio::test] + async fn test_mark_schedulable_without_state_writes_nothing() { + let (kv_backend, manager) = new_manager(); + + assert!( + !manager + .mark_schedulable(NodeSchedulingRole::Datanode, 2) + .await + .unwrap() + ); + + let key: Vec = NodeSchedulingKey::new(NodeSchedulingRole::Datanode, 2).into(); + assert!(kv_backend.get(&key).await.unwrap().is_none()); + } + + #[tokio::test] + async fn test_scheduling_state_is_role_scoped() { + let (_, manager) = new_manager(); + + manager + .mark_unschedulable(NodeSchedulingRole::Datanode, 7) + .await + .unwrap(); + + assert!( + manager + .is_schedulable(NodeSchedulingRole::Flownode, 7) + .await + .unwrap() + ); + assert!( + manager + .list_unschedulable_node_ids(NodeSchedulingRole::Flownode) + .await + .unwrap() + .is_empty() + ); + } + + #[tokio::test] + async fn test_scheduling_state_is_visible_to_other_manager() { + let (kv_backend, manager_a) = new_manager(); + let manager_b = NodeSchedulingManager::new(kv_backend); + + manager_a + .mark_unschedulable(NodeSchedulingRole::Datanode, 3) + .await + .unwrap(); + + assert!( + !manager_b + .is_schedulable(NodeSchedulingRole::Datanode, 3) + .await + .unwrap() + ); + } + + #[tokio::test] + async fn test_concurrent_mark_unschedulable_changes_once() { + let (_, manager) = new_manager(); + + let (a, b) = tokio::join!( + manager.mark_unschedulable(NodeSchedulingRole::Datanode, 2), + manager.mark_unschedulable(NodeSchedulingRole::Datanode, 2) + ); + let changed = [a.unwrap(), b.unwrap()]; + assert_eq!(1, changed.iter().filter(|changed| **changed).count()); + } + + /// Returns a selector context that stamps scheduling marks of datanodes 1, 2, 3 + /// where datanode 2 is unschedulable. + async fn new_selector_context() -> SelectorContext { + use common_meta::peer::Peer; + + use crate::test_util::{ + create_meta_peer_client, new_node_scheduling_discovery, put_datanodes, + }; + + let client = create_meta_peer_client(); + put_datanodes( + &client, + vec![Peer::empty(1), Peer::empty(2), Peer::empty(3)], + ) + .await; + + let (manager, discovery) = new_node_scheduling_discovery(&client); + manager + .mark_unschedulable(NodeSchedulingRole::Datanode, 2) + .await + .unwrap(); + + SelectorContext { + affinity_fallback: Default::default(), + affinity: Arc::new(common_meta::affinity::NoopAffinity), + peer_discovery: Arc::new(discovery), + } + } + + #[tokio::test] + async fn test_lease_based_selector_skips_unschedulable() { + use crate::error; + use crate::selector::lease_based::LeaseBasedSelector; + use crate::selector::{Selector, SelectorOptions}; + + let ctx = new_selector_context().await; + let selector = LeaseBasedSelector; + + let peers = selector + .select( + &ctx, + SelectorOptions { + min_required_items: 2, + ..Default::default() + }, + ) + .await + .unwrap(); + assert_eq!(2, peers.len()); + for peer in peers { + assert!([1, 3].contains(&peer.id), "unexpected peer: {peer:?}"); + } + + let err = selector + .select( + &ctx, + SelectorOptions { + min_required_items: 3, + ..Default::default() + }, + ) + .await + .unwrap_err(); + assert!( + matches!(err, error::Error::NoEnoughAvailableNode { .. }), + "unexpected error: {err:?}" + ); + } + + #[tokio::test] + async fn test_round_robin_selector_skips_unschedulable() { + use crate::metasrv::SelectTarget; + use crate::selector::round_robin::RoundRobinSelector; + use crate::selector::{Selector, SelectorOptions}; + + let ctx = new_selector_context().await; + let selector = RoundRobinSelector::new(SelectTarget::Datanode); + + let peers = selector + .select( + &ctx, + SelectorOptions { + min_required_items: 2, + ..Default::default() + }, + ) + .await + .unwrap(); + let ids = peers.iter().map(|peer| peer.id).collect::>(); + assert_eq!(vec![1, 3], ids); + } +} diff --git a/src/meta-srv/src/peer.rs b/src/meta-srv/src/peer.rs index 8843d58ad87..e0a0589ab37 100644 --- a/src/meta-srv/src/peer.rs +++ b/src/meta-srv/src/peer.rs @@ -112,7 +112,7 @@ mod tests { use super::*; use crate::selector::Selector; use crate::selector::round_robin::RoundRobinSelector; - use crate::test_util::{create_meta_peer_client, put_datanodes}; + use crate::test_util::{create_meta_peer_client, new_node_scheduling_discovery, put_datanodes}; struct OnlySecondNode; @@ -244,7 +244,7 @@ mod tests { let allocator = MetasrvPeerAllocator::new( SelectorContext { affinity_fallback: fallback, - peer_discovery: discovery.clone(), + peer_discovery: Arc::new(new_node_scheduling_discovery(&discovery).1), affinity, }, selector.clone(), diff --git a/src/meta-srv/src/region/supervisor.rs b/src/meta-srv/src/region/supervisor.rs index 5cf3f8e7fa8..c7c5cf7dab7 100644 --- a/src/meta-srv/src/region/supervisor.rs +++ b/src/meta-srv/src/region/supervisor.rs @@ -702,7 +702,19 @@ impl RegionSupervisor { } else { None }; - let exclude_peer_ids = HashSet::from_iter(failure_datanodes.iter().copied()); + let mut exclude_peer_ids: HashSet = + HashSet::from_iter(failure_datanodes.iter().copied()); + let unschedulable_datanodes = self + .selector_context + .peer_discovery + .active_datanodes(None) + .await + .map_err(common_error::ext::BoxedError::new) + .context(error::OtherSnafu)? + .into_iter() + .filter(|datanode| !datanode.is_schedulable()) + .map(|datanode| datanode.peer.id); + exclude_peer_ids.extend(unschedulable_datanodes); match &self.selector { RegionSupervisorSelector::NaiveSelector(selector) => { let opts = SelectorOptions { @@ -944,6 +956,7 @@ pub(crate) mod tests { use std::sync::{Arc, Mutex}; use std::time::Duration; + use common_meta::affinity::NoopAffinity; use common_meta::ddl::RegionFailureDetectorController; use common_meta::ddl::test_util::{ test_create_logical_table_task, test_create_physical_table_task, @@ -964,6 +977,7 @@ pub(crate) mod tests { use tokio::time::sleep; use super::RegionSupervisorSelector; + use crate::metasrv::SelectorContext; use crate::metrics::METRIC_META_HEARTBEAT_DROPPED; use crate::procedure::region_migration::RegionMigrationTriggerReason; use crate::procedure::region_migration::manager::{ @@ -977,9 +991,19 @@ pub(crate) mod tests { use crate::selector::test_utils::{RandomNodeSelector, new_test_selector_context}; pub(crate) fn new_test_supervisor() -> (RegionSupervisor, Sender) { + new_test_supervisor_with( + new_test_selector_context(), + RegionSupervisorSelector::NaiveSelector(Arc::new(RandomNodeSelector::new(vec![ + Peer::empty(1), + ]))), + ) + } + + pub(crate) fn new_test_supervisor_with( + selector_context: SelectorContext, + selector: RegionSupervisorSelector, + ) -> (RegionSupervisor, Sender) { let env = TestingEnv::new(); - let selector_context = new_test_selector_context(); - let selector = Arc::new(RandomNodeSelector::new(vec![Peer::empty(1)])); let context_factory = env.context_factory(); let region_migration_manager = Arc::new(RegionMigrationManager::new( env.procedure_manager().clone(), @@ -996,7 +1020,7 @@ pub(crate) mod tests { rx, Default::default(), selector_context, - RegionSupervisorSelector::NaiveSelector(selector), + selector, region_migration_manager, runtime_switch_manager, peer_resolver, @@ -1204,6 +1228,61 @@ pub(crate) mod tests { } } + #[tokio::test] + async fn test_select_peers_excludes_unschedulable_datanodes() { + use std::collections::HashSet as StdHashSet; + + use crate::key::NodeSchedulingRole; + use crate::selector::{Selector, SelectorOptions}; + use crate::test_util::{ + create_meta_peer_client, new_node_scheduling_discovery, put_datanodes, + }; + + /// Selector that records the excluded peers and returns a fixed peer. + struct ExcludeAssertSelector; + + #[async_trait::async_trait] + impl Selector for ExcludeAssertSelector { + type Context = SelectorContext; + type Output = Vec; + + async fn select( + &self, + _ctx: &Self::Context, + opts: SelectorOptions<'_>, + ) -> crate::error::Result { + assert_eq!(StdHashSet::from([7, 8]), opts.exclude_peer_ids); + Ok(vec![Peer::empty(9)]) + } + } + + // Datanode 7 is marked unschedulable; datanode 8 failed. + let client = create_meta_peer_client(); + put_datanodes(&client, vec![Peer::empty(7)]).await; + let (manager, discovery) = new_node_scheduling_discovery(&client); + manager + .mark_unschedulable(NodeSchedulingRole::Datanode, 7) + .await + .unwrap(); + + let selector_context = SelectorContext { + affinity_fallback: Default::default(), + affinity: Arc::new(NoopAffinity), + peer_discovery: Arc::new(discovery), + }; + let (supervisor, _tx) = new_test_supervisor_with( + selector_context, + RegionSupervisorSelector::NaiveSelector(Arc::new(ExcludeAssertSelector)), + ); + + let region_peers = supervisor + .select_peers(7, &[RegionId::new(1, 1)], &[8]) + .await + .unwrap(); + assert_eq!(1, region_peers.len()); + assert_eq!(Peer::empty(9), region_peers[0].1); + } + #[tokio::test] async fn test_heartbeat() { let (mut supervisor, sender) = new_test_supervisor(); diff --git a/src/meta-srv/src/selector/common.rs b/src/meta-srv/src/selector/common.rs index 00e72b23096..0fec9a4ec66 100644 --- a/src/meta-srv/src/selector/common.rs +++ b/src/meta-srv/src/selector/common.rs @@ -14,15 +14,33 @@ use std::collections::HashSet; +use api::v1::meta::heartbeat_request::NodeWorkloads; +use common_meta::cluster::NodeInfo; use common_meta::peer::Peer; -use snafu::ensure; +use snafu::{ResultExt, ensure}; use crate::error; use crate::error::Result; -use crate::metasrv::SelectTarget; +use crate::metasrv::{SelectTarget, SelectorContext}; use crate::selector::SelectorOptions; use crate::selector::weighted_choose::{WeightedChoose, WeightedItem}; +/// Returns active datanodes that may receive new regions. +pub(crate) async fn active_schedulable_datanodes( + ctx: &SelectorContext, + filter: Option fn(&'a NodeWorkloads) -> bool>, +) -> Result> { + let datanodes = ctx + .peer_discovery + .active_datanodes(filter) + .await + .context(error::ListActiveDatanodesSnafu)?; + Ok(datanodes + .into_iter() + .filter(NodeInfo::is_schedulable) + .collect()) +} + /// Filter out the excluded peers from the `weight_array`. pub fn filter_out_excluded_peers( weight_array: &mut Vec>, @@ -268,6 +286,7 @@ mod tests { env_vars: pool .map(|pool| HashMap::from([("pool".into(), pool.into())])) .unwrap_or_default(), + schedulable: None, } } diff --git a/src/meta-srv/src/selector/lease_based.rs b/src/meta-srv/src/selector/lease_based.rs index b9ddd852dc4..83363929a57 100644 --- a/src/meta-srv/src/selector/lease_based.rs +++ b/src/meta-srv/src/selector/lease_based.rs @@ -14,11 +14,10 @@ use common_meta::affinity::AffinityFallback; use common_meta::peer::Peer; -use snafu::ResultExt; -use crate::error::{self, ListActiveDatanodesSnafu, Result}; +use crate::error::{self, Result}; use crate::metasrv::SelectorContext; -use crate::selector::common::{choose_items, filter_out_excluded_peers}; +use crate::selector::common::{self, choose_items, filter_out_excluded_peers}; use crate::selector::weighted_choose::{RandomWeightedChoose, WeightedItem}; use crate::selector::{Selector, SelectorOptions}; @@ -33,11 +32,8 @@ impl LeaseBasedSelector { opts: &SelectorOptions<'_>, ) -> Result> { // 1. get alive datanodes. - let alive_datanodes = ctx - .peer_discovery - .active_datanodes(opts.workload_filter) - .await - .context(ListActiveDatanodesSnafu)? + let alive_datanodes = common::active_schedulable_datanodes(ctx, opts.workload_filter) + .await? .into_iter() .filter(|node| { opts.affinity diff --git a/src/meta-srv/src/selector/load_based.rs b/src/meta-srv/src/selector/load_based.rs index 84196f3973e..175c65b4696 100644 --- a/src/meta-srv/src/selector/load_based.rs +++ b/src/meta-srv/src/selector/load_based.rs @@ -18,12 +18,11 @@ use common_meta::affinity::AffinityFallback; use common_meta::datanode::{DatanodeStatKey, DatanodeStatValue}; use common_meta::peer::Peer; use common_telemetry::debug; -use snafu::ResultExt; use crate::cluster::MetaPeerClientRef; -use crate::error::{self, ListActiveDatanodesSnafu, Result}; +use crate::error::{self, Result}; use crate::metasrv::SelectorContext; -use crate::selector::common::{choose_items, filter_out_excluded_peers}; +use crate::selector::common::{self, choose_items, filter_out_excluded_peers}; use crate::selector::weight_compute::WeightCompute; use crate::selector::weighted_choose::RandomWeightedChoose; use crate::selector::{Selector, SelectorOptions}; @@ -52,11 +51,8 @@ where opts: &SelectorOptions<'_>, ) -> Result> { // 1. get alive datanodes. - let alive_datanodes = ctx - .peer_discovery - .active_datanodes(opts.workload_filter) - .await - .context(ListActiveDatanodesSnafu)? + let alive_datanodes = common::active_schedulable_datanodes(ctx, opts.workload_filter) + .await? .into_iter() .filter(|node| { opts.affinity diff --git a/src/meta-srv/src/selector/round_robin.rs b/src/meta-srv/src/selector/round_robin.rs index 05921921145..45e2f4f296b 100644 --- a/src/meta-srv/src/selector/round_robin.rs +++ b/src/meta-srv/src/selector/round_robin.rs @@ -18,11 +18,9 @@ use common_meta::affinity::AffinityFallback; use common_meta::peer::Peer; use snafu::{ResultExt, ensure}; -use crate::error::{ - self, ListActiveDatanodesSnafu, ListActiveFlownodesSnafu, NoEnoughAvailableNodeSnafu, Result, -}; +use crate::error::{self, ListActiveFlownodesSnafu, NoEnoughAvailableNodeSnafu, Result}; use crate::metasrv::{SelectTarget, SelectorContext}; -use crate::selector::{Selector, SelectorOptions}; +use crate::selector::{Selector, SelectorOptions, common}; /// Round-robin selector that returns the next peer in the list in sequence. /// Datanodes are ordered by their node_id. @@ -59,19 +57,22 @@ impl RoundRobinSelector { ctx: &SelectorContext, ) -> Result> { let mut peers = match self.select_target { - SelectTarget::Datanode => ctx - .peer_discovery - .active_datanodes(opts.workload_filter) - .await - .context(ListActiveDatanodesSnafu)? - .into_iter() - .filter(|node| { - opts.affinity - .is_none_or(|requirements| ctx.affinity.matches(requirements, node)) - }) - .map(|node| node.peer) - .filter(|peer| !opts.exclude_peer_ids.contains(&peer.id)) - .collect::>(), + SelectTarget::Datanode => { + // 1. get alive, schedulable datanodes. + let alive_datanodes = + common::active_schedulable_datanodes(ctx, opts.workload_filter).await?; + + // 2. filter by affinity, then drop the excluded datanodes. + alive_datanodes + .into_iter() + .filter(|node| { + opts.affinity + .is_none_or(|requirements| ctx.affinity.matches(requirements, node)) + }) + .map(|node| node.peer) + .filter(|peer| !opts.exclude_peer_ids.contains(&peer.id)) + .collect::>() + } SelectTarget::Flownode => ctx .peer_discovery .active_flownodes(opts.workload_filter) @@ -157,7 +158,7 @@ mod test { use common_meta::affinity::NoopAffinity; use super::*; - use crate::test_util::{create_meta_peer_client, put_datanodes}; + use crate::test_util::{create_meta_peer_client, new_node_scheduling_discovery, put_datanodes}; #[tokio::test] async fn test_round_robin_selector() { @@ -166,7 +167,7 @@ mod test { let ctx = SelectorContext { affinity_fallback: Default::default(), affinity: Arc::new(NoopAffinity), - peer_discovery: meta_peer_client.clone(), + peer_discovery: Arc::new(new_node_scheduling_discovery(&meta_peer_client).1), }; // add three nodes let peer1 = Peer { @@ -227,7 +228,7 @@ mod test { let ctx = SelectorContext { affinity_fallback: Default::default(), affinity: Arc::new(NoopAffinity), - peer_discovery: meta_peer_client.clone(), + peer_discovery: Arc::new(new_node_scheduling_discovery(&meta_peer_client).1), }; // add three nodes let peer1 = Peer { diff --git a/src/meta-srv/src/selector/test_utils.rs b/src/meta-srv/src/selector/test_utils.rs index 9b79a435e90..8073e41344e 100644 --- a/src/meta-srv/src/selector/test_utils.rs +++ b/src/meta-srv/src/selector/test_utils.rs @@ -23,13 +23,19 @@ use crate::cluster::MetaPeerClientBuilder; use crate::error::Result; use crate::metasrv::SelectorContext; use crate::selector::{Selector, SelectorOptions}; +use crate::test_util::new_node_scheduling_discovery; /// Returns [SelectorContext] for test purpose. +/// +/// Its `peer_discovery` is decorated with +/// [`SchedulingAwarePeerDiscovery`](crate::discovery::node_scheduling::SchedulingAwarePeerDiscovery) +/// backed by an empty scheduling state, so every datanode is schedulable until a test marks it +/// otherwise in the same in-memory backend. pub fn new_test_selector_context() -> SelectorContext { let kv_backend = Arc::new(MemoryKvBackend::new()); let meta_peer_client = MetaPeerClientBuilder::default() .election(None) - .in_memory(kv_backend.clone()) + .in_memory(kv_backend) .build() .map(Arc::new) .unwrap(); @@ -37,7 +43,7 @@ pub fn new_test_selector_context() -> SelectorContext { SelectorContext { affinity_fallback: Default::default(), affinity: Arc::new(NoopAffinity), - peer_discovery: meta_peer_client, + peer_discovery: Arc::new(new_node_scheduling_discovery(&meta_peer_client).1), } } diff --git a/src/meta-srv/src/test_util.rs b/src/meta-srv/src/test_util.rs index 1169a93309d..4696db380c7 100644 --- a/src/meta-srv/src/test_util.rs +++ b/src/meta-srv/src/test_util.rs @@ -24,7 +24,9 @@ use common_time::util as time_util; use common_workload::DatanodeWorkloadType; use crate::cluster::{MetaPeerClientBuilder, MetaPeerClientRef}; +use crate::discovery::node_scheduling::SchedulingAwarePeerDiscovery; use crate::key::{DatanodeLeaseKey, LeaseValue}; +use crate::node_scheduling::{NodeSchedulingManager, NodeSchedulingManagerRef}; pub(crate) fn new_region_route(region_id: u64, peers: &[Peer], leader_node: u64) -> RegionRoute { let region = Region { @@ -55,6 +57,20 @@ pub(crate) fn create_meta_peer_client() -> MetaPeerClientRef { .unwrap() } +/// Wraps `meta_peer_client` into the node scheduling discovery. +/// +/// The manager shares the client's in-memory backend, so marks written through the returned +/// manager are visible to the discovery. +pub(crate) fn new_node_scheduling_discovery( + meta_peer_client: &MetaPeerClientRef, +) -> (NodeSchedulingManagerRef, SchedulingAwarePeerDiscovery) { + let manager = Arc::new(NodeSchedulingManager::new( + meta_peer_client.memory_backend().as_kv_backend_ref(), + )); + let discovery = SchedulingAwarePeerDiscovery::new(manager.clone(), meta_peer_client.clone()); + (manager, discovery) +} + pub(crate) async fn put_datanodes(meta_peer_client: &MetaPeerClientRef, datanodes: Vec) { let backend = meta_peer_client.memory_backend(); for datanode in datanodes { @@ -102,6 +118,7 @@ pub(crate) async fn put_datanodes(meta_peer_client: &MetaPeerClientRef, datanode memory_usage_bytes: 0, hostname: String::new(), env_vars: Default::default(), + schedulable: None, }; let node_info_key = NodeInfoKey { role: Role::Datanode, diff --git a/src/meta-srv/src/utils/database.rs b/src/meta-srv/src/utils/database.rs index 796d63009d8..a67cdb14aab 100644 --- a/src/meta-srv/src/utils/database.rs +++ b/src/meta-srv/src/utils/database.rs @@ -182,6 +182,7 @@ mod tests { memory_usage_bytes: 0, hostname: String::new(), env_vars: Default::default(), + schedulable: None, }]) } diff --git a/src/standalone/src/information_extension.rs b/src/standalone/src/information_extension.rs index 8d52eb12ccc..f42d7a84f7f 100644 --- a/src/standalone/src/information_extension.rs +++ b/src/standalone/src/information_extension.rs @@ -89,6 +89,7 @@ impl InformationExtension for StandaloneInformationExtension { .to_string_lossy() .to_string(), env_vars: Default::default(), + schedulable: None, }; Ok(vec![node_info]) } From 815da8a6ef58839de11370b50686646125f5435b Mon Sep 17 00:00:00 2001 From: Yingwen Date: Fri, 23 Oct 2026 07:45:15 +0000 Subject: [PATCH 083/239] feat(mito): expose scan concurrency metrics and dashboard (#9758) * feat(mito): expose planned scan concurrency histogram Signed-off-by: evenyag * feat(grafana): visualize planned scan concurrency Signed-off-by: evenyag * feat(mito): distinguish query and compaction scan metrics Signed-off-by: evenyag --------- Signed-off-by: evenyag --- .../dashboards/metrics/cluster/dashboard.json | 142 +++++++++++++++++- .../dashboards/metrics/cluster/dashboard.md | 1 + .../dashboards/metrics/cluster/dashboard.yaml | 20 +++ .../metrics/standalone/dashboard.json | 142 +++++++++++++++++- .../metrics/standalone/dashboard.md | 1 + .../metrics/standalone/dashboard.yaml | 20 +++ src/mito2/src/metrics.rs | 8 + src/mito2/src/read/seq_scan.rs | 18 ++- src/mito2/src/read/series_scan.rs | 13 +- 9 files changed, 343 insertions(+), 22 deletions(-) diff --git a/grafana/dashboards/metrics/cluster/dashboard.json b/grafana/dashboards/metrics/cluster/dashboard.json index 2fffd27b25f..fd1d0257594 100644 --- a/grafana/dashboards/metrics/cluster/dashboard.json +++ b/grafana/dashboards/metrics/cluster/dashboard.json @@ -9112,6 +9112,132 @@ "title": "Read Stage P99 and Avg per Instance", "type": "timeseries" }, + { + "datasource": { + "type": "prometheus", + "uid": "${metrics}" + }, + "description": "Estimated maximum concurrent files by instance and type (query or compaction). Query samples are region scan limit evaluations, including memtable sources counted by the limit, rejected scans, and scans using memory budgeting. Compaction samples count all input files when building a compaction reader. A SQL query can produce multiple region scan evaluations. P95/P99 are histogram estimates, not exact per-query maxima or live open-file counts.", + "fieldConfig": { + "defaults": { + "color": { + "mode": "palette-classic" + }, + "custom": { + "axisBorderShow": false, + "axisCenteredZero": false, + "axisColorMode": "text", + "axisLabel": "", + "axisPlacement": "auto", + "barAlignment": 0, + "barWidthFactor": 0.6, + "drawStyle": "line", + "fillOpacity": 0, + "gradientMode": "none", + "hideFrom": { + "legend": false, + "tooltip": false, + "viz": false + }, + "insertNulls": false, + "lineInterpolation": "linear", + "lineWidth": 1, + "pointSize": 5, + "scaleDistribution": { + "type": "linear" + }, + "showPoints": "auto", + "spanNulls": false, + "stacking": { + "group": "A", + "mode": "none" + }, + "thresholdsStyle": { + "mode": "off" + } + }, + "mappings": [], + "thresholds": { + "mode": "absolute", + "steps": [ + { + "color": "green" + }, + { + "color": "red", + "value": 80 + } + ] + }, + "unit": "short", + "min": 0 + }, + "overrides": [] + }, + "gridPos": { + "h": 8, + "w": 8, + "x": 8, + "y": 154 + }, + "id": 565, + "options": { + "legend": { + "calcs": [ + "lastNotNull" + ], + "displayMode": "table", + "placement": "bottom", + "showLegend": true + }, + "tooltip": { + "hideZeros": false, + "mode": "multi", + "sort": "desc" + } + }, + "pluginVersion": "11.6.0", + "targets": [ + { + "datasource": { + "type": "prometheus", + "uid": "${metrics}" + }, + "editorMode": "code", + "expr": "histogram_quantile(0.95, sum by(instance, pod, type, le) (rate(greptime_mito_scan_max_concurrent_files_bucket{instance=~\"$datanode\"}[$__rate_interval])))", + "instant": false, + "legendFormat": "[{{instance}}]-[{{pod}}]-{{type}}-p95", + "range": true, + "refId": "A" + }, + { + "datasource": { + "type": "prometheus", + "uid": "${metrics}" + }, + "editorMode": "code", + "expr": "histogram_quantile(0.99, sum by(instance, pod, type, le) (rate(greptime_mito_scan_max_concurrent_files_bucket{instance=~\"$datanode\"}[$__rate_interval])))", + "instant": false, + "legendFormat": "[{{instance}}]-[{{pod}}]-{{type}}-p99", + "range": true, + "refId": "B" + }, + { + "datasource": { + "type": "prometheus", + "uid": "${metrics}" + }, + "editorMode": "code", + "expr": "sum by(instance, pod, type) (rate(greptime_mito_scan_max_concurrent_files_sum{instance=~\"$datanode\"}[$__rate_interval])) / sum by(instance, pod, type) (rate(greptime_mito_scan_max_concurrent_files_count{instance=~\"$datanode\"}[$__rate_interval]))", + "instant": false, + "legendFormat": "[{{instance}}]-[{{pod}}]-{{type}}-avg", + "range": true, + "refId": "C" + } + ], + "title": "Scan Max Concurrent Files per Instance", + "type": "timeseries" + }, { "datasource": { "type": "prometheus", @@ -9176,7 +9302,7 @@ "gridPos": { "h": 8, "w": 8, - "x": 8, + "x": 16, "y": 154 }, "id": 490, @@ -9292,8 +9418,8 @@ "gridPos": { "h": 8, "w": 8, - "x": 16, - "y": 154 + "x": 0, + "y": 162 }, "id": 495, "options": { @@ -9391,7 +9517,7 @@ "gridPos": { "h": 8, "w": 8, - "x": 0, + "x": 8, "y": 162 }, "id": 501, @@ -9505,7 +9631,7 @@ "gridPos": { "h": 8, "w": 8, - "x": 8, + "x": 16, "y": 162 }, "id": 502, @@ -9619,8 +9745,8 @@ "gridPos": { "h": 8, "w": 8, - "x": 16, - "y": 162 + "x": 0, + "y": 170 }, "id": 503, "options": { @@ -9736,7 +9862,7 @@ "gridPos": { "h": 8, "w": 8, - "x": 0, + "x": 8, "y": 170 }, "id": 504, diff --git a/grafana/dashboards/metrics/cluster/dashboard.md b/grafana/dashboards/metrics/cluster/dashboard.md index 455268cdd43..1c167af5152 100644 --- a/grafana/dashboards/metrics/cluster/dashboard.md +++ b/grafana/dashboards/metrics/cluster/dashboard.md @@ -104,6 +104,7 @@ | Write Rows per Instance | `sum by (instance, pod) (rate(greptime_mito_write_rows_total{instance=~"$datanode"}[$__rate_interval]))` | `timeseries` | Ingestion size by row counts. | `prometheus` | `rowsps` | `[{{instance}}]-[{{pod}}]` | | Read Stage OPS per Instance | `sum by(instance, pod) (rate(greptime_mito_read_stage_elapsed_count{instance=~"$datanode", stage="total"}[$__rate_interval]))` | `timeseries` | Read Stage OPS per Instance. | `prometheus` | `ops` | `[{{instance}}]-[{{pod}}]` | | Read Stage P99 and Avg per Instance | `histogram_quantile(0.99, sum by(instance, pod, le, stage) (rate(greptime_mito_read_stage_elapsed_bucket{instance=~"$datanode"}[$__rate_interval])))`
`sum by(instance, pod, stage) (rate(greptime_mito_read_stage_elapsed_sum{instance=~"$datanode"}[$__rate_interval])) / sum by(instance, pod, stage) (rate(greptime_mito_read_stage_elapsed_count{instance=~"$datanode"}[$__rate_interval]))` | `timeseries` | Read Stage P99 and average per Instance. | `prometheus` | `s` | `[{{instance}}]-[{{pod}}]-[{{stage}}]` | +| Scan Max Concurrent Files per Instance | `histogram_quantile(0.95, sum by(instance, pod, type, le) (rate(greptime_mito_scan_max_concurrent_files_bucket{instance=~"$datanode"}[$__rate_interval])))`
`histogram_quantile(0.99, sum by(instance, pod, type, le) (rate(greptime_mito_scan_max_concurrent_files_bucket{instance=~"$datanode"}[$__rate_interval])))`
`sum by(instance, pod, type) (rate(greptime_mito_scan_max_concurrent_files_sum{instance=~"$datanode"}[$__rate_interval])) / sum by(instance, pod, type) (rate(greptime_mito_scan_max_concurrent_files_count{instance=~"$datanode"}[$__rate_interval]))` | `timeseries` | Estimated maximum concurrent files by instance and type (query or compaction). Query samples are region scan limit evaluations, including memtable sources counted by the limit, rejected scans, and scans using memory budgeting. Compaction samples count all input files when building a compaction reader. A SQL query can produce multiple region scan evaluations. P95/P99 are histogram estimates, not exact per-query maxima or live open-file counts. | `prometheus` | `short` | `[{{instance}}]-[{{pod}}]-{{type}}-p95` | | Write Stage P99 and Avg per Instance | `histogram_quantile(0.99, sum by(instance, pod, le, stage) (rate(greptime_mito_write_stage_elapsed_bucket{instance=~"$datanode"}[$__rate_interval])))`
`sum by(instance, pod, stage) (rate(greptime_mito_write_stage_elapsed_sum{instance=~"$datanode"}[$__rate_interval])) / sum by(instance, pod, stage) (rate(greptime_mito_write_stage_elapsed_count{instance=~"$datanode"}[$__rate_interval]))` | `timeseries` | Write Stage P99 and average per Instance. | `prometheus` | `s` | `[{{instance}}]-[{{pod}}]-[{{stage}}]` | | Cached Bytes per Instance | `greptime_mito_cache_bytes{instance=~"$datanode"}` | `timeseries` | Cached Bytes per Instance. | `prometheus` | `decbytes` | `[{{instance}}]-[{{pod}}]-[{{type}}]` | | Region Worker Handle Bulk Insert Requests | `histogram_quantile(0.95, sum by(le,instance, stage, pod) (rate(greptime_region_worker_handle_write_bucket[$__rate_interval])))`
`sum by(instance, stage, pod) (rate(greptime_region_worker_handle_write_sum[$__rate_interval]))/sum by(instance, stage, pod) (rate(greptime_region_worker_handle_write_count[$__rate_interval]))` | `timeseries` | Per-stage elapsed time for region worker to handle bulk insert region requests. | `prometheus` | `s` | `[{{instance}}]-[{{pod}}]-[{{stage}}]-P95` | diff --git a/grafana/dashboards/metrics/cluster/dashboard.yaml b/grafana/dashboards/metrics/cluster/dashboard.yaml index 7513c5e50e1..01fdb7bbc31 100644 --- a/grafana/dashboards/metrics/cluster/dashboard.yaml +++ b/grafana/dashboards/metrics/cluster/dashboard.yaml @@ -1279,6 +1279,26 @@ groups: type: prometheus uid: ${metrics} legendFormat: '[{{instance}}]-[{{pod}}]-[{{stage}}]-avg' + - title: Scan Max Concurrent Files per Instance + type: timeseries + description: Estimated maximum concurrent files by instance and type (query or compaction). Query samples are region scan limit evaluations, including memtable sources counted by the limit, rejected scans, and scans using memory budgeting. Compaction samples count all input files when building a compaction reader. A SQL query can produce multiple region scan evaluations. P95/P99 are histogram estimates, not exact per-query maxima or live open-file counts. + unit: short + queries: + - expr: histogram_quantile(0.95, sum by(instance, pod, type, le) (rate(greptime_mito_scan_max_concurrent_files_bucket{instance=~"$datanode"}[$__rate_interval]))) + datasource: + type: prometheus + uid: ${metrics} + legendFormat: '[{{instance}}]-[{{pod}}]-{{type}}-p95' + - expr: histogram_quantile(0.99, sum by(instance, pod, type, le) (rate(greptime_mito_scan_max_concurrent_files_bucket{instance=~"$datanode"}[$__rate_interval]))) + datasource: + type: prometheus + uid: ${metrics} + legendFormat: '[{{instance}}]-[{{pod}}]-{{type}}-p99' + - expr: sum by(instance, pod, type) (rate(greptime_mito_scan_max_concurrent_files_sum{instance=~"$datanode"}[$__rate_interval])) / sum by(instance, pod, type) (rate(greptime_mito_scan_max_concurrent_files_count{instance=~"$datanode"}[$__rate_interval])) + datasource: + type: prometheus + uid: ${metrics} + legendFormat: '[{{instance}}]-[{{pod}}]-{{type}}-avg' - title: Write Stage P99 and Avg per Instance type: timeseries description: Write Stage P99 and average per Instance. diff --git a/grafana/dashboards/metrics/standalone/dashboard.json b/grafana/dashboards/metrics/standalone/dashboard.json index f17d40b4f96..d980fb500d6 100644 --- a/grafana/dashboards/metrics/standalone/dashboard.json +++ b/grafana/dashboards/metrics/standalone/dashboard.json @@ -9112,6 +9112,132 @@ "title": "Read Stage P99 and Avg per Instance", "type": "timeseries" }, + { + "datasource": { + "type": "prometheus", + "uid": "${metrics}" + }, + "description": "Estimated maximum concurrent files by instance and type (query or compaction). Query samples are region scan limit evaluations, including memtable sources counted by the limit, rejected scans, and scans using memory budgeting. Compaction samples count all input files when building a compaction reader. A SQL query can produce multiple region scan evaluations. P95/P99 are histogram estimates, not exact per-query maxima or live open-file counts.", + "fieldConfig": { + "defaults": { + "color": { + "mode": "palette-classic" + }, + "custom": { + "axisBorderShow": false, + "axisCenteredZero": false, + "axisColorMode": "text", + "axisLabel": "", + "axisPlacement": "auto", + "barAlignment": 0, + "barWidthFactor": 0.6, + "drawStyle": "line", + "fillOpacity": 0, + "gradientMode": "none", + "hideFrom": { + "legend": false, + "tooltip": false, + "viz": false + }, + "insertNulls": false, + "lineInterpolation": "linear", + "lineWidth": 1, + "pointSize": 5, + "scaleDistribution": { + "type": "linear" + }, + "showPoints": "auto", + "spanNulls": false, + "stacking": { + "group": "A", + "mode": "none" + }, + "thresholdsStyle": { + "mode": "off" + } + }, + "mappings": [], + "thresholds": { + "mode": "absolute", + "steps": [ + { + "color": "green" + }, + { + "color": "red", + "value": 80 + } + ] + }, + "unit": "short", + "min": 0 + }, + "overrides": [] + }, + "gridPos": { + "h": 8, + "w": 8, + "x": 8, + "y": 154 + }, + "id": 565, + "options": { + "legend": { + "calcs": [ + "lastNotNull" + ], + "displayMode": "table", + "placement": "bottom", + "showLegend": true + }, + "tooltip": { + "hideZeros": false, + "mode": "multi", + "sort": "desc" + } + }, + "pluginVersion": "11.6.0", + "targets": [ + { + "datasource": { + "type": "prometheus", + "uid": "${metrics}" + }, + "editorMode": "code", + "expr": "histogram_quantile(0.95, sum by(instance, pod, type, le) (rate(greptime_mito_scan_max_concurrent_files_bucket[$__rate_interval])))", + "instant": false, + "legendFormat": "[{{instance}}]-[{{pod}}]-{{type}}-p95", + "range": true, + "refId": "A" + }, + { + "datasource": { + "type": "prometheus", + "uid": "${metrics}" + }, + "editorMode": "code", + "expr": "histogram_quantile(0.99, sum by(instance, pod, type, le) (rate(greptime_mito_scan_max_concurrent_files_bucket[$__rate_interval])))", + "instant": false, + "legendFormat": "[{{instance}}]-[{{pod}}]-{{type}}-p99", + "range": true, + "refId": "B" + }, + { + "datasource": { + "type": "prometheus", + "uid": "${metrics}" + }, + "editorMode": "code", + "expr": "sum by(instance, pod, type) (rate(greptime_mito_scan_max_concurrent_files_sum[$__rate_interval])) / sum by(instance, pod, type) (rate(greptime_mito_scan_max_concurrent_files_count[$__rate_interval]))", + "instant": false, + "legendFormat": "[{{instance}}]-[{{pod}}]-{{type}}-avg", + "range": true, + "refId": "C" + } + ], + "title": "Scan Max Concurrent Files per Instance", + "type": "timeseries" + }, { "datasource": { "type": "prometheus", @@ -9176,7 +9302,7 @@ "gridPos": { "h": 8, "w": 8, - "x": 8, + "x": 16, "y": 154 }, "id": 490, @@ -9292,8 +9418,8 @@ "gridPos": { "h": 8, "w": 8, - "x": 16, - "y": 154 + "x": 0, + "y": 162 }, "id": 495, "options": { @@ -9391,7 +9517,7 @@ "gridPos": { "h": 8, "w": 8, - "x": 0, + "x": 8, "y": 162 }, "id": 501, @@ -9505,7 +9631,7 @@ "gridPos": { "h": 8, "w": 8, - "x": 8, + "x": 16, "y": 162 }, "id": 502, @@ -9619,8 +9745,8 @@ "gridPos": { "h": 8, "w": 8, - "x": 16, - "y": 162 + "x": 0, + "y": 170 }, "id": 503, "options": { @@ -9736,7 +9862,7 @@ "gridPos": { "h": 8, "w": 8, - "x": 0, + "x": 8, "y": 170 }, "id": 504, diff --git a/grafana/dashboards/metrics/standalone/dashboard.md b/grafana/dashboards/metrics/standalone/dashboard.md index 5c7a4042444..7ed99d2ab2c 100644 --- a/grafana/dashboards/metrics/standalone/dashboard.md +++ b/grafana/dashboards/metrics/standalone/dashboard.md @@ -104,6 +104,7 @@ | Write Rows per Instance | `sum by (instance, pod) (rate(greptime_mito_write_rows_total[$__rate_interval]))` | `timeseries` | Ingestion size by row counts. | `prometheus` | `rowsps` | `[{{instance}}]-[{{pod}}]` | | Read Stage OPS per Instance | `sum by(instance, pod) (rate(greptime_mito_read_stage_elapsed_count{stage="total"}[$__rate_interval]))` | `timeseries` | Read Stage OPS per Instance. | `prometheus` | `ops` | `[{{instance}}]-[{{pod}}]` | | Read Stage P99 and Avg per Instance | `histogram_quantile(0.99, sum by(instance, pod, le, stage) (rate(greptime_mito_read_stage_elapsed_bucket[$__rate_interval])))`
`sum by(instance, pod, stage) (rate(greptime_mito_read_stage_elapsed_sum[$__rate_interval])) / sum by(instance, pod, stage) (rate(greptime_mito_read_stage_elapsed_count[$__rate_interval]))` | `timeseries` | Read Stage P99 and average per Instance. | `prometheus` | `s` | `[{{instance}}]-[{{pod}}]-[{{stage}}]` | +| Scan Max Concurrent Files per Instance | `histogram_quantile(0.95, sum by(instance, pod, type, le) (rate(greptime_mito_scan_max_concurrent_files_bucket[$__rate_interval])))`
`histogram_quantile(0.99, sum by(instance, pod, type, le) (rate(greptime_mito_scan_max_concurrent_files_bucket[$__rate_interval])))`
`sum by(instance, pod, type) (rate(greptime_mito_scan_max_concurrent_files_sum[$__rate_interval])) / sum by(instance, pod, type) (rate(greptime_mito_scan_max_concurrent_files_count[$__rate_interval]))` | `timeseries` | Estimated maximum concurrent files by instance and type (query or compaction). Query samples are region scan limit evaluations, including memtable sources counted by the limit, rejected scans, and scans using memory budgeting. Compaction samples count all input files when building a compaction reader. A SQL query can produce multiple region scan evaluations. P95/P99 are histogram estimates, not exact per-query maxima or live open-file counts. | `prometheus` | `short` | `[{{instance}}]-[{{pod}}]-{{type}}-p95` | | Write Stage P99 and Avg per Instance | `histogram_quantile(0.99, sum by(instance, pod, le, stage) (rate(greptime_mito_write_stage_elapsed_bucket[$__rate_interval])))`
`sum by(instance, pod, stage) (rate(greptime_mito_write_stage_elapsed_sum[$__rate_interval])) / sum by(instance, pod, stage) (rate(greptime_mito_write_stage_elapsed_count[$__rate_interval]))` | `timeseries` | Write Stage P99 and average per Instance. | `prometheus` | `s` | `[{{instance}}]-[{{pod}}]-[{{stage}}]` | | Cached Bytes per Instance | `greptime_mito_cache_bytes` | `timeseries` | Cached Bytes per Instance. | `prometheus` | `decbytes` | `[{{instance}}]-[{{pod}}]-[{{type}}]` | | Region Worker Handle Bulk Insert Requests | `histogram_quantile(0.95, sum by(le,instance, stage, pod) (rate(greptime_region_worker_handle_write_bucket[$__rate_interval])))`
`sum by(instance, stage, pod) (rate(greptime_region_worker_handle_write_sum[$__rate_interval]))/sum by(instance, stage, pod) (rate(greptime_region_worker_handle_write_count[$__rate_interval]))` | `timeseries` | Per-stage elapsed time for region worker to handle bulk insert region requests. | `prometheus` | `s` | `[{{instance}}]-[{{pod}}]-[{{stage}}]-P95` | diff --git a/grafana/dashboards/metrics/standalone/dashboard.yaml b/grafana/dashboards/metrics/standalone/dashboard.yaml index 627cd8b5d63..7f70dba0364 100644 --- a/grafana/dashboards/metrics/standalone/dashboard.yaml +++ b/grafana/dashboards/metrics/standalone/dashboard.yaml @@ -1279,6 +1279,26 @@ groups: type: prometheus uid: ${metrics} legendFormat: '[{{instance}}]-[{{pod}}]-[{{stage}}]-avg' + - title: Scan Max Concurrent Files per Instance + type: timeseries + description: Estimated maximum concurrent files by instance and type (query or compaction). Query samples are region scan limit evaluations, including memtable sources counted by the limit, rejected scans, and scans using memory budgeting. Compaction samples count all input files when building a compaction reader. A SQL query can produce multiple region scan evaluations. P95/P99 are histogram estimates, not exact per-query maxima or live open-file counts. + unit: short + queries: + - expr: histogram_quantile(0.95, sum by(instance, pod, type, le) (rate(greptime_mito_scan_max_concurrent_files_bucket[$__rate_interval]))) + datasource: + type: prometheus + uid: ${metrics} + legendFormat: '[{{instance}}]-[{{pod}}]-{{type}}-p95' + - expr: histogram_quantile(0.99, sum by(instance, pod, type, le) (rate(greptime_mito_scan_max_concurrent_files_bucket[$__rate_interval]))) + datasource: + type: prometheus + uid: ${metrics} + legendFormat: '[{{instance}}]-[{{pod}}]-{{type}}-p99' + - expr: sum by(instance, pod, type) (rate(greptime_mito_scan_max_concurrent_files_sum[$__rate_interval])) / sum by(instance, pod, type) (rate(greptime_mito_scan_max_concurrent_files_count[$__rate_interval])) + datasource: + type: prometheus + uid: ${metrics} + legendFormat: '[{{instance}}]-[{{pod}}]-{{type}}-avg' - title: Write Stage P99 and Avg per Instance type: timeseries description: Write Stage P99 and average per Instance. diff --git a/src/mito2/src/metrics.rs b/src/mito2/src/metrics.rs index cd1116c4fbd..cffa6de39a1 100644 --- a/src/mito2/src/metrics.rs +++ b/src/mito2/src/metrics.rs @@ -230,6 +230,14 @@ lazy_static! { "Number of SSTs to scan in a scan task", vec![1.0, 4.0, 8.0, 16.0, 32.0, 64.0, 256.0, 1024.0], ).unwrap(); + /// Planned maximum concurrent sources per query scan limit evaluation or compaction reader. + /// Includes memtables counted by the file limit, rejected scans, and memory-budgeted scans. + pub static ref SCAN_MAX_CONCURRENT_FILES: HistogramVec = register_histogram_vec!( + "greptime_mito_scan_max_concurrent_files", + "Estimated maximum concurrent files (including memtable sources) per query scan limit evaluation or compaction reader", + &[TYPE_LABEL], + vec![0.0, 1.0, 2.0, 4.0, 8.0, 16.0, 32.0, 64.0, 128.0, 256.0, 384.0, 512.0, 1024.0, 2048.0, 4096.0, 8192.0], + ).unwrap(); /// Histogram for the number of rows returned per query. pub static ref READ_ROWS_RETURN: Histogram = register_histogram!( "greptime_mito_read_rows_return", diff --git a/src/mito2/src/read/seq_scan.rs b/src/mito2/src/read/seq_scan.rs index 8db83ed1845..42119135805 100644 --- a/src/mito2/src/read/seq_scan.rs +++ b/src/mito2/src/read/seq_scan.rs @@ -36,6 +36,7 @@ use store_api::storage::TimeSeriesRowSelector; use tokio::sync::Semaphore; use crate::error::{PartitionOutOfRangeSnafu, Result, TooManyFilesToReadSnafu}; +use crate::metrics::SCAN_MAX_CONCURRENT_FILES; use crate::read::flat_dedup::{FlatDedupReader, FlatLastNonNull, FlatLastRow}; use crate::read::flat_merge::FlatMergeReader; use crate::read::last_row::FlatLastRowReader; @@ -148,6 +149,11 @@ impl SeqScan { pub(crate) async fn build_flat_reader_for_compaction(&self) -> Result { assert!(self.stream_ctx.input.compaction); + // Compaction merges all input files together instead of scanning ranges sequentially. + SCAN_MAX_CONCURRENT_FILES + .with_label_values(&["compaction"]) + .observe(self.stream_ctx.input.num_files() as f64); + let metrics_set = ExecutionPlanMetricsSet::new(); let part_metrics = self.new_partition_metrics(false, &metrics_set, 0); debug_assert_eq!(1, self.properties.partitions.len()); @@ -526,9 +532,6 @@ impl SeqScan { /// Checks resource limit for the scanner. pub(crate) fn check_scan_limit(&self) -> Result<()> { - if self.stream_ctx.input.uses_scan_memory_limit() { - return Ok(()); - } // Check max file count limit for all partitions since we scan them in parallel. let total_max_files: usize = self .properties @@ -537,6 +540,15 @@ impl SeqScan { .map(|partition| Self::max_files_in_partition(&self.stream_ctx.ranges, partition)) .sum(); + if !self.stream_ctx.input.compaction { + SCAN_MAX_CONCURRENT_FILES + .with_label_values(&["query"]) + .observe(total_max_files as f64); + } + if self.stream_ctx.input.uses_scan_memory_limit() { + return Ok(()); + } + let max_concurrent_files = self.stream_ctx.input.max_concurrent_scan_files; if total_max_files > max_concurrent_files { return TooManyFilesToReadSnafu { diff --git a/src/mito2/src/read/series_scan.rs b/src/mito2/src/read/series_scan.rs index 64d5ceb07f8..f8af9e51b91 100644 --- a/src/mito2/src/read/series_scan.rs +++ b/src/mito2/src/read/series_scan.rs @@ -46,6 +46,7 @@ use crate::error::{ Error, InvalidSenderSnafu, JoinSnafu, PartitionOutOfRangeSnafu, Result, ScanMultiTimesSnafu, ScanSeriesSnafu, TooManyFilesToReadSnafu, }; +use crate::metrics::SCAN_MAX_CONCURRENT_FILES; use crate::read::ScannerMetrics; use crate::read::pruner::{PartitionPruner, Pruner, PrunerOptions}; use crate::read::scan_region::{ScanInput, StreamContext}; @@ -480,9 +481,6 @@ impl SeriesScan { /// Checks resource limit for the scanner. pub(crate) fn check_scan_limit(&self) -> Result<()> { - if self.stream_ctx.input.uses_scan_memory_limit() { - return Ok(()); - } // Sum the total number of files across all partitions let total_files: usize = self .properties @@ -495,6 +493,15 @@ impl SeriesScan { }) .sum(); + if !self.stream_ctx.input.compaction { + SCAN_MAX_CONCURRENT_FILES + .with_label_values(&["query"]) + .observe(total_files as f64); + } + if self.stream_ctx.input.uses_scan_memory_limit() { + return Ok(()); + } + let max_concurrent_files = self.stream_ctx.input.max_concurrent_scan_files; if total_files > max_concurrent_files { return TooManyFilesToReadSnafu { From afbe77e31ff223da8a8ec52e6b7efedd559958ab Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Fri, 23 Oct 2026 08:05:38 +0000 Subject: [PATCH 084/239] perf: avoid duplicate OTLP histogram conversion (#9880) Signed-off-by: shuiyisong --- src/servers/benches/otlp_metrics.rs | 228 +++++++++++++++++- src/servers/src/otlp/metrics.rs | 89 +++++++ src/servers/src/otlp/metrics/resource_info.rs | 17 ++ 3 files changed, 333 insertions(+), 1 deletion(-) diff --git a/src/servers/benches/otlp_metrics.rs b/src/servers/benches/otlp_metrics.rs index eccf7f4f167..130537696f9 100644 --- a/src/servers/benches/otlp_metrics.rs +++ b/src/servers/benches/otlp_metrics.rs @@ -12,16 +12,242 @@ // See the License for the specific language governing permissions and // limitations under the License. +//! Measures OTLP request-to-row conversion, excluding request cloning and storage. +//! Save a baseline with `cargo bench -p servers --bench otlp_metrics -- --save-baseline before`, +//! then compare with `cargo bench -p servers --bench otlp_metrics -- --baseline before`. + #[path = "otlp_metrics/fixtures.rs"] mod fixtures; use std::hint::black_box; use std::time::Duration; +use api::v1::value::ValueData; +use common_catalog::consts::SEMANTIC_GRAPH_WINDOW_NANOS; +use common_query::prelude::{greptime_timestamp, set_default_prefix}; use criterion::{BatchSize, BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; +use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest; +use otel_arrow_rust::proto::opentelemetry::common::v1::{AnyValue, KeyValue, any_value}; +use otel_arrow_rust::proto::opentelemetry::metrics::v1::{ + AggregationTemporality, ExponentialHistogram, ExponentialHistogramDataPoint, Gauge, Histogram, + HistogramDataPoint, Metric, NumberDataPoint, ResourceMetrics, ScopeMetrics, + exponential_histogram_data_point::Buckets, metric, number_data_point, +}; +use otel_arrow_rust::proto::opentelemetry::resource::v1::Resource; use servers::otlp::metrics::to_grpc_insert_requests; use session::protocol_ctx::OtlpMetricCtx; +const POINTS: usize = 128; + +fn attribute(key: &str, value: &str) -> KeyValue { + KeyValue { + key: key.to_string(), + value: Some(AnyValue { + value: Some(any_value::Value::StringValue(value.to_string())), + }), + } +} + +fn timestamp(index: usize) -> u64 { + ((index % 4 + 1) as i64 * SEMANTIC_GRAPH_WINDOW_NANOS) as u64 + index as u64 * 1_000_000 +} + +fn point_attributes() -> Vec { + vec![ + attribute("http.request.method", "GET"), + attribute("route", "/api"), + ] +} + +fn request(data: metric::Data) -> ExportMetricsServiceRequest { + ExportMetricsServiceRequest { + resource_metrics: vec![ResourceMetrics { + resource: Some(Resource { + attributes: vec![ + attribute("service.name", "api"), + attribute("service.namespace", "shop"), + attribute("service.instance.id", "instance-1"), + attribute("host.id", "host-1"), + attribute("k8s.pod.name", "api-1"), + ], + ..Default::default() + }), + scope_metrics: vec![ScopeMetrics { + metrics: vec![Metric { + name: "request_duration".to_string(), + data: Some(data), + ..Default::default() + }], + ..Default::default() + }], + ..Default::default() + }], + } +} + +fn exponential(bucket_counts: Vec, scale: i32, reject: bool) -> metric::Data { + let count = 2 + 2 * bucket_counts.iter().sum::(); + metric::Data::ExponentialHistogram(ExponentialHistogram { + aggregation_temporality: AggregationTemporality::Cumulative as i32, + data_points: (0..POINTS) + .map(|index| ExponentialHistogramDataPoint { + attributes: point_attributes(), + start_time_unix_nano: 1, + time_unix_nano: timestamp(index), + count: count + u64::from(reject && index % 4 == 3), + sum: Some(100.0), + scale, + zero_count: 2, + positive: Some(Buckets { + offset: -128, + bucket_counts: bucket_counts.clone(), + }), + negative: Some(Buckets { + offset: -128, + bucket_counts: bucket_counts.clone(), + }), + zero_threshold: 1e-6, + ..Default::default() + }) + .collect(), + }) +} + +fn bench_otlp_metrics(c: &mut Criterion) { + set_default_prefix(None).unwrap(); + let gauge = metric::Data::Gauge(Gauge { + data_points: (0..POINTS) + .map(|index| NumberDataPoint { + attributes: point_attributes(), + time_unix_nano: timestamp(index), + value: Some(number_data_point::Value::AsDouble(1.0)), + ..Default::default() + }) + .collect(), + }); + let histogram = metric::Data::Histogram(Histogram { + aggregation_temporality: AggregationTemporality::Delta as i32, + data_points: (0..POINTS) + .map(|index| HistogramDataPoint { + attributes: point_attributes(), + start_time_unix_nano: 1, + time_unix_nano: timestamp(index), + count: 17, + sum: Some(100.0), + bucket_counts: vec![1; 17], + explicit_bounds: (0..16).map(|index| index as f64).collect(), + ..Default::default() + }) + .collect(), + }); + let workloads = [ + ("gauge", gauge, POINTS, POINTS), + ("classic_16", histogram, POINTS, POINTS * 19), + ( + "exponential_16", + exponential(vec![1; 16], 0, false), + POINTS, + POINTS, + ), + ( + "exponential_256", + exponential(vec![1; 256], 0, false), + POINTS, + POINTS, + ), + ( + "exponential_sparse_256", + exponential((0..256).map(|i| u64::from(i % 16 == 0)).collect(), 0, false), + POINTS, + POINTS, + ), + ( + "exponential_downscale_256", + exponential(vec![1; 256], 10, false), + POINTS, + POINTS, + ), + ( + "exponential_reject_quarter_256", + exponential(vec![1; 256], 0, true), + POINTS * 3 / 4, + POINTS * 3 / 4, + ), + ]; + + let mut group = c.benchmark_group("otlp_metrics"); + group.sample_size(30); + group.warm_up_time(Duration::from_secs(1)); + group.measurement_time(Duration::from_secs(3)); + group.throughput(Throughput::Elements(POINTS as u64)); + for (name, data, accepted, rows) in workloads { + let request = request(data); + for resource_info in [false, true] { + let ctx = OtlpMetricCtx { + resource_info, + ..Default::default() + }; + let conversion = to_grpc_insert_requests(request.clone(), &mut ctx.clone()).unwrap(); + assert_eq!( + conversion.outcome.accepted_data_points, accepted as i64, + "{name}" + ); + assert_eq!( + conversion.outcome.rejected_data_points, + (POINTS - accepted) as i64, + "{name}" + ); + assert_eq!(conversion.rows, rows, "{name}"); + assert_eq!(conversion.resource_info.is_some(), resource_info, "{name}"); + if let Some(descriptor) = conversion.resource_info { + let descriptor_rows = descriptor.inserts[0].rows.as_ref().unwrap(); + let ts_index = descriptor_rows + .schema + .iter() + .position(|column| column.column_name == greptime_timestamp()) + .unwrap(); + let expected = (0..4) + .filter(|window| accepted == POINTS || *window != 3) + .map(|window| { + Some(ValueData::TimestampMillisecondValue( + (timestamp(POINTS - 4 + window) / 1_000_000) as i64, + )) + }) + .collect::>(); + let actual = descriptor_rows + .rows + .iter() + .map(|row| row.values[ts_index].value_data.clone()) + .collect::>(); + assert_eq!(actual, expected, "{name}"); + } + group.bench_with_input( + BenchmarkId::new( + name, + if resource_info { + "descriptor_on" + } else { + "descriptor_off" + }, + ), + &request, + |b, request| { + b.iter_batched( + || (request.clone(), ctx.clone()), + |(request, mut ctx)| { + black_box( + to_grpc_insert_requests(black_box(request), &mut ctx).unwrap(), + ) + }, + BatchSize::LargeInput, + ); + }, + ); + } + } + group.finish(); +} + fn bench_conversion(c: &mut Criterion) { let fixtures = fixtures::fixtures(); for (group_name, delta) in [ @@ -79,6 +305,6 @@ fn bench_conversion(c: &mut Criterion) { criterion_group! { name = benches; config = Criterion::default().sample_size(20).warm_up_time(Duration::from_secs(1)).measurement_time(Duration::from_secs(3)); - targets = bench_conversion + targets = bench_conversion, bench_otlp_metrics } criterion_main!(benches); diff --git a/src/servers/src/otlp/metrics.rs b/src/servers/src/otlp/metrics.rs index 27462512e64..c1226d6ed6f 100644 --- a/src/servers/src/otlp/metrics.rs +++ b/src/servers/src/otlp/metrics.rs @@ -1591,6 +1591,95 @@ mod tests { } } + #[test] + fn test_conversion_describes_only_accepted_exponential_histogram_windows() { + set_default_prefix(None).unwrap(); + let window = common_catalog::consts::SEMANTIC_GRAPH_WINDOW_NANOS as u64; + let mut valid = exponential_point(); + valid.time_unix_nano = window + 1_000_000; + + let mut invalid_count = valid.clone(); + invalid_count.time_unix_nano += 8_000_000; + invalid_count.count += 1; + let mut invalid_scale = valid.clone(); + invalid_scale.time_unix_nano += window; + invalid_scale.scale = -5; + // OTLP bucket conversion succeeds, but the shared native encoder rejects this index. + let invalid_native = ExponentialHistogramDataPoint { + time_unix_nano: 3 * window + 1_000_000, + count: 1, + positive: Some(exponential_buckets(2000, vec![1])), + ..Default::default() + }; + assert!( + exponential_histogram_value(&invalid_native, AggregationTemporality::Cumulative) + .unwrap_err() + .contains("out of range") + ); + let stale = ExponentialHistogramDataPoint { + time_unix_nano: 4 * window + 1_000_000, + count: u64::MAX, + scale: i32::MIN, + flags: DataPointFlags::NoRecordedValueMask as u32, + ..Default::default() + }; + let invalid_time = ExponentialHistogramDataPoint { + time_unix_nano: 5 * window + 1_000_000, + start_time_unix_nano: 5 * window + 2_000_000, + ..stale.clone() + }; + let rejected = vec![invalid_count, invalid_scale, invalid_native, invalid_time]; + let mut mixed = rejected.clone(); + mixed.extend([valid, stale]); + for (points, expected) in [ + (rejected, vec![]), + ( + mixed, + vec![ + (window / 1_000_000 + 1) as i64, + (4 * window / 1_000_000 + 1) as i64, + ], + ), + ] { + let mut request = metrics_request(vec![exponential_metric( + "latency", + points, + AggregationTemporality::Cumulative, + )]); + request.resource_metrics[0].resource = Some(Resource { + attributes: vec![keyvalue("service.name", "api")], + ..Default::default() + }); + let conversion = to_grpc_insert_requests(request, &mut descriptor_ctx()).unwrap(); + assert_eq!(conversion.outcome.rejected_data_points, 4); + assert_eq!( + conversion.outcome.accepted_data_points, + expected.len() as i64 + ); + assert_eq!(conversion.rows, expected.len()); + assert_eq!(conversion.resource_info.is_some(), !expected.is_empty()); + let actual = conversion + .resource_info + .map(|descriptor| { + let rows = descriptor.inserts[0].rows.as_ref().unwrap(); + let ts_index = rows + .schema + .iter() + .position(|column| column.column_name == greptime_timestamp()) + .unwrap(); + rows.rows + .iter() + .map(|row| match row.values[ts_index].value_data { + Some(ValueData::TimestampMillisecondValue(ts)) => ts, + _ => panic!("expected descriptor timestamp"), + }) + .collect::>() + }) + .unwrap_or_default(); + assert_eq!(actual, expected); + } + } + #[test] fn test_conversion_skips_descriptor_on_metric_name_collision() { set_default_prefix(None).unwrap(); diff --git a/src/servers/src/otlp/metrics/resource_info.rs b/src/servers/src/otlp/metrics/resource_info.rs index c7ad72c3e47..e9c88cf4bc5 100644 --- a/src/servers/src/otlp/metrics/resource_info.rs +++ b/src/servers/src/otlp/metrics/resource_info.rs @@ -217,6 +217,15 @@ mod tests { .all(|(k, _)| k != "os.type" && k != "service.instance.id") ); + let mut reordered = attrs.clone(); + reordered.reverse(); + data.observe_times(&reordered, [200, 75]); + assert_eq!(data.rows.len(), 1); + assert_eq!( + data.rows.values().next().unwrap().values().next(), + Some(&200) + ); + data.observe_times(&[kv("host.id", "h-2")], [10]); assert_eq!(data.rows.len(), 2); @@ -225,6 +234,14 @@ mod tests { assert!(empty.into_row_insert_requests().unwrap().is_none()); } + #[test] + fn observe_without_accepted_points_does_not_emit_rows() { + let mut data = ResourceInfoData::default(); + data.observe_times(&[kv("service.name", "api")], []); + assert!(data.is_empty()); + assert!(data.into_row_insert_requests().unwrap().is_none()); + } + /// Earlier windows would keep their metric rows but lose their entities. #[test] fn observe_keeps_one_row_per_graph_window() { From 1cb01d0d4480d2743ed841c220f2eb6166106723 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Fri, 23 Oct 2026 09:04:03 +0000 Subject: [PATCH 085/239] fix(mito2): keep SSTs when a manifest save ends with RateLimited or PermissionDenied (#9898) * fix(mito2): keep SSTs when a manifest save ends with RateLimited The object store's retry layer retries temporary errors and returns the last attempt's error. A manifest delta PUT can be stored by the server while the client sees a timeout, reset or 5xx, and the retries that follow can be throttled, so the save ends with RateLimited even though the delta exists. may_have_persisted_manifest_update treated RateLimited as proof that nothing was written, so flush and compaction deleted the SSTs the stored delta references. A later open or follower sync replays the delta and reads fail on the missing files. Treat RateLimited like other ambiguous errors: the SSTs are kept and become orphans when the save really failed. Fixes #9897 Signed-off-by: jeremyhi * fix(mito2): keep SSTs when a manifest save ends with PermissionDenied The retry layer returns the last attempt's error, so a 403 from a revoked key, a changed bucket policy or clock skew can follow a delta PUT that the server already stored. Like RateLimited, PermissionDenied does not prove that nothing was written, so it no longer lets flush or compaction delete the SSTs; they become orphans when the save really failed. The compaction test that rejects a manifest write before persistence now injects Unsupported, which is raised before any request is sent. Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi Co-authored-by: Yingwen --- src/mito2/src/engine/compaction_test.rs | 2 +- src/mito2/src/error.rs | 23 +++++++++++++++-------- 2 files changed, 16 insertions(+), 9 deletions(-) diff --git a/src/mito2/src/engine/compaction_test.rs b/src/mito2/src/engine/compaction_test.rs index 02f76af6f5f..b6594ddc136 100644 --- a/src/mito2/src/engine/compaction_test.rs +++ b/src/mito2/src/engine/compaction_test.rs @@ -1414,7 +1414,7 @@ impl object_store::layers::mock::Write for FailingManifestWriter { ) -> object_store::layers::mock::Result<()> { if matches!(self.failure, ManifestWriteFailure::BeforeWrite) { return Err(object_store::layers::mock::Error::new( - object_store::layers::mock::ErrorKind::PermissionDenied, + object_store::layers::mock::ErrorKind::Unsupported, "manifest write rejected before persistence", )); } diff --git a/src/mito2/src/error.rs b/src/mito2/src/error.rs index 71253e78bea..9bb908720fd 100644 --- a/src/mito2/src/error.rs +++ b/src/mito2/src/error.rs @@ -1445,7 +1445,9 @@ impl Error { /// Returns whether a failed manifest update may have been persisted. /// /// Unknown errors are treated conservatively because deleting output SSTs after an ambiguous - /// manifest write could leave the manifest referencing missing files. + /// manifest write could leave the manifest referencing missing files. The retry layer returns + /// the last attempt's error, so an error that can follow a stored attempt (throttling, or a 403 + /// from a revoked key, a changed bucket policy or clock skew) never proves nothing was written. pub(crate) fn may_have_persisted_manifest_update(&self) -> bool { match self { Error::UpdateManifest { .. } @@ -1459,11 +1461,9 @@ impl Error { ErrorKind::Unsupported | ErrorKind::ConfigInvalid | ErrorKind::NotFound - | ErrorKind::PermissionDenied | ErrorKind::IsADirectory | ErrorKind::NotADirectory | ErrorKind::AlreadyExists - | ErrorKind::RateLimited | ErrorKind::IsSameFile | ErrorKind::ConditionNotMatch | ErrorKind::RangeNotSatisfied @@ -1847,11 +1847,9 @@ mod tests { ErrorKind::Unsupported, ErrorKind::ConfigInvalid, ErrorKind::NotFound, - ErrorKind::PermissionDenied, ErrorKind::IsADirectory, ErrorKind::NotADirectory, ErrorKind::AlreadyExists, - ErrorKind::RateLimited, ErrorKind::IsSameFile, ErrorKind::ConditionNotMatch, ErrorKind::RangeNotSatisfied, @@ -1864,9 +1862,18 @@ mod tests { ); } - let error = - OpenDalSnafu {}.into_error(object_store::Error::new(ErrorKind::Unexpected, "test")); - assert!(error.may_have_persisted_manifest_update()); + // The last retry can fail with these kinds even though an earlier attempt was stored. + for kind in [ + ErrorKind::Unexpected, + ErrorKind::RateLimited, + ErrorKind::PermissionDenied, + ] { + let error = OpenDalSnafu {}.into_error(object_store::Error::new(kind, "test")); + assert!( + error.may_have_persisted_manifest_update(), + "error kind {kind:?} should not prove the manifest was not persisted" + ); + } let region_id = RegionId::new(1, 1); let error = RegionTruncatedSnafu { region_id }.build(); From f2be11a69e6875d746f9f82ad40efa681bb467c1 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Fri, 23 Oct 2026 10:39:14 +0000 Subject: [PATCH 086/239] fix: upload staging manifests in chunks (#9919) Signed-off-by: WenyXu --- src/mito2/src/manifest/storage/staging.rs | 71 +++++++++++++++++++++-- 1 file changed, 65 insertions(+), 6 deletions(-) diff --git a/src/mito2/src/manifest/storage/staging.rs b/src/mito2/src/manifest/storage/staging.rs index a0b350a76cd..e4e9e7f177f 100644 --- a/src/mito2/src/manifest/storage/staging.rs +++ b/src/mito2/src/manifest/storage/staging.rs @@ -23,7 +23,7 @@ use store_api::ManifestVersion; use crate::error::{OpenDalSnafu, Result}; use crate::manifest::storage::delta::DeltaStorage; use crate::manifest::storage::size_tracker::NoopTracker; -use crate::manifest::storage::utils::sort_manifests; +use crate::manifest::storage::utils::{sort_manifests, write_and_put_cache}; use crate::manifest::storage::{file_version, is_delta_file}; /// A simple blob storage for arbitrary binary data in the staging directory. @@ -79,10 +79,7 @@ impl StagingBlobStorage { self.object_store.info().root(), bytes.len() ); - self.object_store - .write(&path, bytes) - .await - .context(OpenDalSnafu)?; + write_and_put_cache(&self.object_store, None, &path, bytes).await?; Ok(()) } @@ -222,13 +219,75 @@ impl StagingStorage { #[cfg(test)] mod tests { + use std::sync::{Arc, Mutex}; + use common_datasource::compression::CompressionType; use common_test_util::temp_dir::create_temp_dir; use futures::TryStreamExt; + use object_store::layers::mock::{MockLayerBuilder, ObservedWriter, WriterObservations, oio}; use object_store::services::Fs; use object_store::{ErrorKind, ObjectStore}; - use super::{StagingStorage, staging_blob_path, staging_manifest_path}; + use super::{StagingBlobStorage, StagingStorage, staging_blob_path, staging_manifest_path}; + + // A test-only per-write limit, not a storage provider's default limit. + const WRITE_LIMIT: usize = 64 * 1024 * 1024; + + struct LimitedWriter { + inner: oio::Writer, + } + + impl oio::Write for LimitedWriter { + async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> { + if bytes.len() > WRITE_LIMIT { + return Err(object_store::Error::new( + ErrorKind::Unexpected, + "test write size limit exceeded", + )); + } + self.inner.write(bytes).await + } + + async fn close(&mut self) -> object_store::Result { + self.inner.close().await + } + + async fn abort(&mut self) -> object_store::Result<()> { + self.inner.abort().await + } + } + + #[tokio::test] + async fn test_staging_blob_put_chunked() { + let tmp_dir = create_temp_dir("test_staging_blob_put_chunked"); + let observations = Arc::new(Mutex::new(WriterObservations::default())); + let captured = observations.clone(); + let layer = MockLayerBuilder::default() + .writer_factory(Arc::new(move |_, _, inner| { + Box::new(ObservedWriter::new( + Box::new(LimitedWriter { inner }), + captured.clone(), + )) + })) + .build() + .unwrap(); + let store = ObjectStore::new(Fs::default().root(&tmp_dir.path().to_string_lossy())) + .unwrap() + .layer(layer); + let storage = StagingBlobStorage::new("staging/blob/".to_string(), store); + + for size in [17, WRITE_LIMIT, WRITE_LIMIT + 17] { + let bytes = (0..size).map(|i| (i % 251) as u8).collect::>(); + *observations.lock().unwrap() = WriterObservations::default(); + storage.put("manifest", bytes.clone()).await.unwrap(); + assert_eq!(storage.get("manifest").await.unwrap(), bytes); + let observed = observations.lock().unwrap(); + assert_eq!(observed.write_sizes.iter().sum::(), size); + assert!(observed.write_sizes.iter().all(|size| *size <= WRITE_LIMIT)); + assert_eq!(observed.write_sizes.len(), size.div_ceil(WRITE_LIMIT)); + assert_eq!(observed.close_calls, 1); + } + } #[test] fn test_staging_path() { From 1c5d15fd365e1d78b301a2f81bd877f558b6fd71 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Fri, 23 Oct 2026 14:27:39 +0000 Subject: [PATCH 087/239] perf(log-store): read object store WAL footers while the recovery listing goes on (#9889) * perf(log-store): read object store WAL footers while the recovery listing goes on Recovery listed the whole prefix before it read any footer, so its LIST pages, about 80-120 ms each on S3 and one per 1000 objects, added to the footer reads. The listing now runs in a task of its own and feeds the footer reads as its pages arrive; a listing that fails part-way still fails recovery, and a failed footer read stops the listing. Signed-off-by: jeremyhi * perf(log-store): fail object store WAL recovery as soon as its listing fails A listing that failed part-way was reported only after every object it had already listed was read, which on S3 can be tens of thousands of reads. The stream now settles a finished listing before it yields another object, so its failure goes first. Tests cover a large listed backlog and a listing task that panics; the counting and parking test object accesses forward list_ahead. Signed-off-by: jeremyhi * test(log-store): hold the listing until the first footer reads wait in the backlog test Signed-off-by: jeremyhi * test(log-store): check a failed listing settles before buffered objects with a finished task Signed-off-by: jeremyhi * perf(log-store): bound the objects the recovery listing holds ahead of the reads The listing fed the footer reads through an unbounded channel, so a listing much faster than the reads could hold most of the prefix in memory. It now holds at most 4096 objects, four LIST pages of S3, and waits for the reads while that many are queued. Signed-off-by: jeremyhi * test(log-store): poll the recovery listing directly to show it waits on a full backlog Signed-off-by: jeremyhi * test(log-store): check the bound with the channel the recovery listing uses Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- src/log-store/src/object_store_wal/io.rs | 227 +++++++++++++++++++- src/log-store/src/object_store_wal/store.rs | 218 +++++++++++++++++-- 2 files changed, 417 insertions(+), 28 deletions(-) diff --git a/src/log-store/src/object_store_wal/io.rs b/src/log-store/src/object_store_wal/io.rs index f9ed5eee36b..43f018a7ecb 100644 --- a/src/log-store/src/object_store_wal/io.rs +++ b/src/log-store/src/object_store_wal/io.rs @@ -15,9 +15,14 @@ //! Object store access for WAL objects: a deterministic key layout and //! conditional creates that make a retry of the same object a no-op. +use async_stream::try_stream; use bytes::Bytes; -use object_store::{ErrorKind, ObjectStore}; +use futures::TryStreamExt; +use futures::stream::BoxStream; +use object_store::{Entry, ErrorKind, ObjectStore}; use snafu::{OptionExt, ResultExt, ensure}; +use tokio::sync::mpsc; +use tokio_util::task::AbortOnDropHandle; use crate::error::{ CorruptedWalObjectSnafu, InvalidWalObjectStoreSnafu, Result, WalObjectConflictSnafu, @@ -28,6 +33,10 @@ use crate::error::{ /// [`u64::MAX`] so that lexicographic and numeric order agree. const OBJECT_SEQ_WIDTH: usize = 20; const OBJECT_SUFFIX: &str = ".wal"; +/// Objects a listing may hold ahead of the reads that take them, four LIST +/// pages of S3, so the listing stays pages ahead of the reads while what it +/// holds is bounded. +const LISTING_AHEAD_OBJECTS: usize = 4096; /// Outcome of a conditional create. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -48,6 +57,7 @@ pub(crate) struct ListedObject { } /// Reads and writes the WAL objects under one prefix. +#[derive(Clone)] pub(crate) struct ObjectStoreIo { store: ObjectStore, object_prefix: String, @@ -160,20 +170,50 @@ impl ObjectStoreIo { path: self.object_prefix.clone(), })?; let mut objects = entries - .into_iter() - .filter_map(|entry| { - self.parse_object_seq(entry.path()) - .map(|object_seq| ListedObject { - object_seq, - path: entry.path().to_string(), - size: entry.metadata().content_length(), - }) - }) + .iter() + .filter_map(|entry| self.listed_object(entry)) .collect::>(); objects.sort_unstable_by_key(|object| object.object_seq); Ok(objects) } + /// Lists the objects under the prefix like [`list`](Self::list), but from + /// a task of its own and in the order the object store returns them, so + /// the listing goes on while the caller consumes the objects listed so + /// far. A listing that fails, however it fails, ends the stream with an + /// error, before any object it listed that the caller has not taken yet; + /// dropping the stream stops the listing. + pub(crate) fn list_ahead(&self) -> BoxStream<'static, Result> { + let (listing, rx) = self.clone().listing(); + let listing = AbortOnDropHandle::new(common_runtime::spawn_global(listing)); + settle_listing(rx, listing, self.object_prefix.clone()) + } + + /// Returns the listing of the objects under the prefix and the receiver it + /// sends them to; it waits while the receiver holds + /// [`LISTING_AHEAD_OBJECTS`] of them. + fn listing( + self, + ) -> ( + impl Future> + Send + 'static, + mpsc::Receiver, + ) { + let (tx, rx) = mpsc::channel(LISTING_AHEAD_OBJECTS); + let listing = async move { + let mut lister = self.store.lister(&self.object_prefix).await?; + while let Some(entry) = lister.try_next().await? { + // The receiver is gone only once the stream was dropped. + if let Some(object) = self.listed_object(&entry) + && tx.send(object).await.is_err() + { + break; + } + } + Ok::<_, object_store::Error>(()) + }; + (listing, rx) + } + pub(crate) fn object_path(&self, object_seq: u64) -> String { format!( "{}{object_seq:0OBJECT_SEQ_WIDTH$}{OBJECT_SUFFIX}", @@ -193,6 +233,15 @@ impl ObjectStoreIo { } } + fn listed_object(&self, entry: &Entry) -> Option { + self.parse_object_seq(entry.path()) + .map(|object_seq| ListedObject { + object_seq, + path: entry.path().to_string(), + size: entry.metadata().content_length(), + }) + } + fn parse_object_seq(&self, path: &str) -> Option { let value = path .strip_prefix(&self.object_prefix)? @@ -204,6 +253,38 @@ impl ObjectStoreIo { } } +/// Yields the objects `listing` sends over `rx`. Once the listing finished it +/// is settled before any further object is yielded, so its failure goes +/// before the objects it listed that the caller has not taken yet. +fn settle_listing( + mut rx: mpsc::Receiver, + listing: AbortOnDropHandle>, + path: String, +) -> BoxStream<'static, Result> { + let mut listing = Some(listing); + Box::pin(try_stream! { + loop { + let object = rx.recv().await; + if let Some(listing) = + listing.take_if(|listing| object.is_none() || listing.is_finished()) + { + listing + .await + .unwrap_or_else(|error| { + let stopped = "the listing task did not complete"; + Err(object_store::Error::new(ErrorKind::Unexpected, stopped) + .set_source(error)) + }) + .context(WalObjectStoreSnafu { operation: "list", path: &path })?; + } + let Some(object) = object else { + break; + }; + yield object; + } + }) +} + /// Returns true when a conditional create failed because the object exists. /// Only the precondition failure says so: a store may report `AlreadyExists` /// for an unrelated path, such as a parent that is a file. @@ -233,10 +314,13 @@ fn normalize_prefix(prefix: &str) -> Result { #[cfg(test)] mod tests { + use std::pin::pin; use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; use common_error::ext::{ErrorExt, RetryHint}; use common_test_util::temp_dir::create_temp_dir; + use futures::poll; use object_store::layers::mock::{self, MockLayerBuilder, oio}; use object_store::secure_fs::SecureFsRoot; use object_store::services::Memory; @@ -513,6 +597,129 @@ mod tests { ); } + #[tokio::test] + async fn test_io_list_ahead_yields_the_objects_list_returns() { + let io = memory_io(); + assert!( + io.list_ahead() + .try_collect::>() + .await + .unwrap() + .is_empty() + ); + + io.store + .write( + "datanodes/1/epochs/2/objects/not-a-sequence.wal", + "malformed", + ) + .await + .unwrap(); + for object_seq in [10, 2, 1] { + io.put_if_absent(object_seq, Bytes::from(object_seq.to_string())) + .await + .unwrap(); + } + let mut listed = io.list_ahead().try_collect::>().await.unwrap(); + listed.sort_unstable_by_key(|object| object.object_seq); + assert_eq!(io.list().await.unwrap(), listed); + } + + struct CountingLister { + inner: oio::Lister, + listed: Arc, + } + + impl oio::List for CountingLister { + async fn next(&mut self) -> mock::Result> { + self.listed.fetch_add(1, Ordering::SeqCst); + self.inner.next().await + } + } + + #[tokio::test] + async fn test_io_listing_waits_while_the_backlog_is_full() { + let store = memory_store(); + let objects = LISTING_AHEAD_OBJECTS + 100; + let io = ObjectStoreIo::new(store.clone(), "datanodes/1/epochs/2/").unwrap(); + for object_seq in 0..objects as u64 { + io.put_if_absent(object_seq, Bytes::from_static(b"wal")) + .await + .unwrap(); + } + let listed = Arc::new(AtomicUsize::new(0)); + let lister_listed = listed.clone(); + let mock = MockLayerBuilder::default() + .lister_factory(Arc::new(move |_, _, inner| { + let listed = lister_listed.clone(); + Box::new(CountingLister { inner, listed }) + })) + .build() + .unwrap(); + let io = ObjectStoreIo::new(store.layer(mock), "datanodes/1/epochs/2/").unwrap(); + let (listing, mut rx) = io.listing(); + // Out of the cooperative budget, so it waits only on a full channel. + let mut listing = pin!(tokio::task::unconstrained(listing)); + + // Nothing takes the objects, so the listing waits once the backlog is + // full, with the directory entry and the object waiting to be sent + // listed on top of it. + for _ in 0..100 { + assert!(poll!(listing.as_mut()).is_pending()); + } + let ahead = listed.load(Ordering::SeqCst); + assert!( + (LISTING_AHEAD_OBJECTS + 1..=LISTING_AHEAD_OBJECTS + 2).contains(&ahead), + "{ahead} entries listed" + ); + + let taken = async { + let mut taken = 0; + while rx.recv().await.is_some() { + taken += 1; + } + taken + }; + let (listed, taken) = futures::join!(listing, taken); + listed.unwrap(); + assert_eq!(objects, taken); + } + + #[tokio::test] + async fn test_io_failed_listing_goes_before_the_objects_it_listed() { + let io = memory_io(); + let (tx, rx) = mpsc::channel(LISTING_AHEAD_OBJECTS); + for object_seq in 0..1000 { + let path = io.object_path(object_seq); + tx.try_send(ListedObject { + object_seq, + path, + size: 0, + }) + .unwrap(); + } + let listing = AbortOnDropHandle::new(tokio::spawn(async move { + drop(tx); + Err(injected_failure("list")) + })); + while !listing.is_finished() { + tokio::task::yield_now().await; + } + + let mut objects = settle_listing(rx, listing, io.object_prefix.clone()); + let error = objects.try_next().await.unwrap_err(); + assert!( + matches!( + error, + Error::WalObjectStore { + operation: "list", + .. + } + ), + "unexpected error: {error:?}" + ); + } + #[tokio::test] async fn test_io_list_ignores_objects_outside_the_wal_layout() { let io = memory_io(); diff --git a/src/log-store/src/object_store_wal/store.rs b/src/log-store/src/object_store_wal/store.rs index d5096c073eb..40d6e75326c 100644 --- a/src/log-store/src/object_store_wal/store.rs +++ b/src/log-store/src/object_store_wal/store.rs @@ -27,8 +27,8 @@ use bytes::Bytes; use common_telemetry::{debug, error, info, warn}; use common_wal::config::object_store::{AckMode, CorruptedSegmentAction, ObjectStoreWalConfig}; use futures::future::{BoxFuture, OptionFuture}; -use futures::stream::FuturesUnordered; -use futures::{StreamExt, TryStreamExt}; +use futures::stream::{BoxStream, FuturesUnordered}; +use futures::{Stream, StreamExt, TryStreamExt}; use object_store::{ErrorKind, ObjectStore}; use snafu::{IntoError, OptionExt, ResultExt, ensure}; use store_api::logstore::entry::{Entry, NaiveEntry}; @@ -3069,12 +3069,12 @@ struct FetchedObject { /// are. Segments are not read, except those of release records, see /// [`read_releases`]; a segment checksum is verified by the read that decodes /// it. Footers are fetched for up to [`RECOVERY_CONCURRENCY`] objects -/// at a time and indexed in sequence order, so the catalog checks the entry -/// ranges of every object against its predecessors like a sequential replay. -/// Only the objects on the chain [`select_chain`] picks are indexed. +/// at a time while the listing goes on, and indexed in sequence order, so the +/// catalog checks the entry ranges of every object against its predecessors +/// like a sequential replay. Only the objects on the chain [`select_chain`] +/// picks are indexed. async fn recover(io: &dyn WalObjectIo) -> Result { - let objects = io.list().await?; - let mut fetched = fetch_footers(io, objects, RECOVERY_CONCURRENCY).await?; + let mut fetched = fetch_footers(io, io.list_ahead(), RECOVERY_CONCURRENCY).await?; let releases = fetched .objects .iter_mut() @@ -3372,18 +3372,18 @@ async fn stale_conflict( .fail() } -/// Fetches and verifies the headers and footers of `objects`, up to -/// `concurrency` objects at a time, and returns them ordered by object sequence -/// whatever the order the fetches complete in. An object the object store no -/// longer has is returned as vanished; any other failure abandons the -/// remaining fetches. +/// Fetches and verifies the headers and footers of `objects` as they are +/// listed, up to `concurrency` objects at a time, and returns them ordered by +/// object sequence whatever the order the fetches complete in. An object the +/// object store no longer has is returned as vanished; a failed listing or +/// any other failure abandons the remaining fetches. async fn fetch_footers( io: &dyn WalObjectIo, - objects: Vec, + objects: impl Stream>, concurrency: usize, ) -> Result { - let outcomes = futures::stream::iter(objects) - .map(|object| async move { + let outcomes = objects + .map_ok(|object| async move { match fetch_footer(io, &object).await { Ok((header, footer, releases)) => Ok(Ok(FetchedObject { object, @@ -3395,7 +3395,7 @@ async fn fetch_footers( Err(error) => Err(error), } }) - .buffer_unordered(concurrency) + .try_buffer_unordered(concurrency) .try_collect::>() .await?; let mut objects = Vec::with_capacity(outcomes.len()); @@ -3549,6 +3549,16 @@ pub(crate) trait WalObjectIo: Send + Sync { async fn list(&self) -> Result>; + /// See [`ObjectStoreIo::list_ahead`]; by default the objects + /// [`list`](Self::list) returns. + fn list_ahead(&self) -> BoxStream<'_, Result> { + Box::pin(try_stream! { + for object in self.list().await? { + yield object; + } + }) + } + fn object_path(&self, object_seq: u64) -> String; /// Describes the object store the objects live in, for messages. @@ -3577,6 +3587,10 @@ impl WalObjectIo for ObjectStoreIo { ObjectStoreIo::list(self).await } + fn list_ahead(&self) -> BoxStream<'_, Result> { + ObjectStoreIo::list_ahead(self) + } + fn object_path(&self, object_seq: u64) -> String { ObjectStoreIo::object_path(self, object_seq) } @@ -5174,7 +5188,7 @@ mod tests { let mut fetch = { let io = io.clone(); tokio::spawn(async move { - fetch_footers(io.as_ref(), objects, RECOVERY_CONCURRENCY) + fetch_footers(io.as_ref(), listing(objects), RECOVERY_CONCURRENCY) .await .unwrap() .objects @@ -5261,6 +5275,170 @@ mod tests { store.stop().await.unwrap(); } + fn listing(objects: Vec) -> impl Stream> { + futures::stream::iter(objects.into_iter().map(Ok)) + } + + const LISTING_PAGE: usize = 4; + + /// Serves the entries of a listing in pages of [`LISTING_PAGE`], each + /// only once `pages` allows it, fails, or panics with `panics`, after + /// `fail_after` entries, and records when it is dropped. + struct PagedLister { + inner: oio::Lister, + served: usize, + pages: watch::Receiver, + fail_after: Option, + panics: bool, + dropped: Arc, + } + + impl oio::List for PagedLister { + async fn next(&mut self) -> mock::Result> { + if self.fail_after == Some(self.served) { + assert!(!self.panics, "injected listing panic"); + return Err(mock::Error::new( + mock::ErrorKind::Unexpected, + "injected listing failure", + )); + } + let page = self.served / LISTING_PAGE; + if self.pages.wait_for(|pages| *pages > page).await.is_err() { + return Err(mock::Error::new( + mock::ErrorKind::Unexpected, + "listing abandoned", + )); + } + self.served += 1; + self.inner.next().await + } + } + + impl Drop for PagedLister { + fn drop(&mut self) { + self.dropped.store(true, Ordering::SeqCst); + } + } + + /// The prefix of `object_store` listed by [`PagedLister`]s, whose reads + /// go through `read`. + fn paged_listing( + object_store: &ObjectStore, + pages: watch::Receiver, + fail_after: Option, + panics: bool, + read: mock::MockReaderFactory, + ) -> (ObjectStoreIo, Arc) { + let dropped = Arc::new(AtomicBool::new(false)); + let lister_dropped = dropped.clone(); + let mock = MockLayerBuilder::default() + .lister_factory(Arc::new(move |_, _, inner| { + Box::new(PagedLister { + inner, + served: 0, + pages: pages.clone(), + fail_after, + panics, + dropped: lister_dropped.clone(), + }) + })) + .reader_factory(read) + .build() + .unwrap(); + let io = ObjectStoreIo::new(object_store.clone().layer(mock), PREFIX).unwrap(); + (io, dropped) + } + + #[tokio::test] + async fn test_store_recovery_reads_footers_while_the_listing_goes_on() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + populate(&object_store, 3 * LISTING_PAGE + 1, 3).await; + // Every read allows one more page, so recovery completes only if it + // reads footers before the listing ends. + let (pages, listed_pages) = watch::channel(1); + let (io, _) = paged_listing( + &object_store, + listed_pages, + None, + false, + Arc::new(move |_, _, inner| { + pages.send_modify(|pages| *pages += 1); + inner + }), + ); + + let recovered = timeout(WAIT, recover(&io)).await.unwrap().unwrap(); + let expected = recover_by_decoding(&ObjectStoreIo::new(object_store, PREFIX).unwrap()) + .await + .unwrap(); + assert_same_recovery(&expected, &recovered); + assert_eq!(3 * LISTING_PAGE as u64 + 1, recovered.recovered_objects); + } + + #[tokio::test] + async fn test_store_recovery_fails_on_a_listing_that_fails_part_way() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + populate(&object_store, 3 * LISTING_PAGE, 3).await; + for (panics, reason) in [ + (false, "injected listing failure"), + (true, "the listing task did not complete"), + ] { + let (_pages, listed_pages) = watch::channel(usize::MAX); + let (io, _) = paged_listing( + &object_store, + listed_pages, + Some(LISTING_PAGE + 1), + panics, + Arc::new(|_, _, inner| inner), + ); + + let error = timeout(WAIT, recover(&io)).await.unwrap().unwrap_err(); + assert!( + matches!(&error, Error::WalObjectStore { operation: "list", error, .. } + if error.to_string().contains(reason)), + "unexpected error: {error:?}" + ); + } + } + + #[tokio::test] + async fn test_store_recovery_stops_the_listing_when_a_footer_read_fails() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + populate(&object_store, 3 * LISTING_PAGE, 3).await; + // Only the first page is ever listed, so the listing waits for good + // unless the failed read stops it. + let (_pages, listed_pages) = watch::channel(1); + let (io, dropped) = paged_listing( + &object_store, + listed_pages, + None, + false, + Arc::new(|_, _, _| Box::new(RefusedReader(ErrorKind::PermissionDenied))), + ); + + let error = timeout(WAIT, recover(&io)).await.unwrap().unwrap_err(); + assert!( + matches!( + &error, + Error::WalObjectStore { + operation: "read", + .. + } + ), + "unexpected error: {error:?}" + ); + timeout(WAIT, async { + while !dropped.load(Ordering::SeqCst) { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .unwrap(); + } + #[tokio::test] async fn test_store_recovery_reads_only_header_and_tail_of_large_object() { let _serialized = SERIALIZED.lock().await; @@ -8880,7 +9058,7 @@ mod tests { } async fn recover_listed(io: &ObjectStoreIo, listed: Vec) -> Result { - finish_recovery(fetch_footers(io, listed, RECOVERY_CONCURRENCY).await?) + finish_recovery(fetch_footers(io, listing(listed), RECOVERY_CONCURRENCY).await?) } #[tokio::test] @@ -10125,6 +10303,10 @@ mod tests { self.inner.list().await } + fn list_ahead(&self) -> BoxStream<'_, Result> { + self.inner.list_ahead() + } + fn object_path(&self, object_seq: u64) -> String { self.inner.object_path(object_seq) } From 80df1c397b796768174d3f836fd24db20d34e736 Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Mon, 26 Oct 2026 07:08:12 +0000 Subject: [PATCH 088/239] feat: COPY FROM STDIN for postgres interface (#9297) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: postgres copy-in support * feat: add postgres copy-in framework * feat(postgres): route COPY FROM STDIN flushes through the pending-rows batcher Follow-up to #9302: begin_copy_in now uses the batching-aware new_query_context, so copy-in batches are submitted to the pending-rows batcher like SQL INSERTs when PostgreSQL batching is enabled for the connection, and keep the direct write path when it is not. The batcher acknowledges synchronously, so the CommandComplete row count stays exact in both modes. Adds test_postgres_copy_in_with_batching asserting the batcher flush metric increments only when batching is enabled. Signed-off-by: Ning Sun * refactor(postgres): open up the COPY FROM STDIN WITH clause for format-specific options The WITH (...) option list is now open instead of restricted to sqlparser's COPY grammar: - COPY ... FROM STDIN statements are parsed textually (new copy_in::parse module, tokenizer-based) so arbitrary KEY value entries reach the codecs; sqlparser would reject options it does not know before we ever see them. WITH is optional, matching PostgreSQL. - CopyOptionSet becomes an open ordered map of KEY -> String | Number | Ident | Boolean | Flag with typed take_* accessors (take_string, take_single_byte, take_flag, raw take_value) and case-insensitive keys; duplicate options are rejected. - Each format codec takes the options it understands and rejects the leftovers with one message naming them ('COPY option(s) X are not supported by format Y'), so a typo cannot be silently ignored. The text format now also rejects HEADER/QUOTE/ESCAPE instead of ignoring them, matching PostgreSQL. A future format (e.g. influxdb line protocol with WITH (FORMAT influxdb_line, precision 'ns')) only needs its codec and a registry entry; no parser changes. Signed-off-by: Ning Sun * refactor(postgres): reuse arrow-csv decoder for COPY FROM STDIN csv parsing Replace the hand-rolled CSV record parser and per-value conversion with arrow-csv's incremental Decoder — the same engine backing COPY ... FROM 'file.csv' in common-datasource — so both COPY routes accept identical data. The decoder handles record splitting (quoting, escaped quotes, custom delimiters, CRLF, headers, chunk boundaries) and typed conversion against the targeted columns' schema in one pass; decoded batches become insert rows via the existing api::helper::vectors_to_rows and datatypes Helper::try_into_vectors, and the decode schema comes from the greptime Schema's arrow projection. The codec framework now has parsers emit complete insertable rows (api Row) instead of raw byte fields: create_parser takes the targeted columns and is fallible; text converts with the shared text-representation helpers, csv with the arrow decoder. Documented consequences of sharing the engine: JSON and Binary columns and bare epoch-integer timestamps are rejected for csv, and an empty field — quoted or not — is NULL; a custom NULL string is rejected accordingly. Signed-off-by: Ning Sun * docs(postgres): make the copy-in framework explicitly format-open Reframe the codec framework docs so nothing assumes csv is the only other text format: the shared pieces (Field, convert_text_field, check_text_representable_type) are described as building blocks for any format whose records carry values as text, while formats with their own typed decoding (csv today) build rows themselves. State explicitly that formats are not limited to PostgreSQL's own — GreptimeDB-specific formats (e.g. a line protocol) register like any other and define their own options through the open WITH clause. No code changes beyond a test rename and a doc clarifier. Signed-off-by: Ning Sun * test(postgres): driver-level coverage for CSV copy-in options Add an end-to-end test through tokio-postgres exercising HEADER, custom DELIMITER, quoted fields with embedded newlines, escaped quotes and CRLF record separators, with read-back verification through the real insert path. WITH is omitted, matching PostgreSQL. Signed-off-by: Ning Sun * test(postgres): drop the raw-socket copy-in tests The simple-protocol raw client duplicated what tokio-postgres and the integration tests already cover; keep the driver-based tests only. Signed-off-by: Ning Sun * refactor(postgres): address review: fix text-format O(n^2) drain, share type checks, hex::decode - text feed: take the longest complete line prefix in one drain and iterate lines over the slice, instead of draining per line (each drain memmoved the whole buffer; ~1300 drains for a 64KB chunk of short lines) - csv check_column_type: delegate to the shared check_text_representable_type and additionally reject Binary/Json - decode_bytea: use the hex crate instead of a hand-rolled decoder - semantic_type_of: use ColumnSchema::is_time_index() like the flow adapter's column_schemas_to_proto, instead of index lookups - parse: drop the unused expect_ident parameter; document the identifier-case decision - PostgresServer::new: document that query_handler and copy_in_handler are the same frontend instance Signed-off-by: Ning Sun * chore: remove unrelated files * feat(postgres)!: drop the text codec; csv is the only COPY FROM STDIN format Remove the hand-rolled PostgreSQL text-format codec (backslash escapes, octal, \N nulls) that arrow-csv cannot express, keeping arrow-csv-based csv as the only implemented format. - COPY FROM STDIN without FORMAT now fails with guidance ('requires an explicit format in this version; use WITH (FORMAT csv)') rather than silently defaulting to the removed format - FORMAT text/txt and FORMAT binary report 'not supported yet' via registry stubs instead of unknown-format errors - codec.rs shrinks to the trait framework: the text-representation helpers (Field, convert_text_field, bytea hex, timestamp parsing) are removed with the codec; the module docs note where to find their shape should a future text-serializing format need them - csv implements check_column_type directly (no shared default) Breaking for clients relying on the PostgreSQL default text format; explicit WITH (FORMAT csv) is required. Signed-off-by: Ning Sun * perf(postgres): raise COPY FROM STDIN flush batch to 32768 rows and harden flush diagnostics - COPY_IN_FLUSH_ROWS 8192 -> 32768: without the pending-rows batcher each flush pays the full insert pipeline (catalog lookup, schema check, partition split, region write round trip), so 4x larger batches cut round trips 4x at linearly-bounded memory; the constant also sizes the arrow decoder batch. Comment documents the trade-off and points high-rate users at the pending-rows batcher, which turns flushes into enqueues. - rows buffer pre-allocates COPY_IN_FLUSH_ROWS. - flush_copy_in warns on unexpected output variants instead of silently counting 0 rows. - batching wire test feeds 70000 rows to stay above the new threshold. Verified against a real psql 18.6 client: \copy ... WITH (FORMAT csv) sends FORMAT csv and succeeds; default \copy (text) gets the 'requires an explicit format' guidance; COPY TO STDOUT round-trips. Signed-off-by: Ning Sun * fix(postgres): reset copy-in state on successful CopyDone on_copy_done left the CopyInState in place (inert but holding buffers) until the next COPY overwrote it; take it like on_copy_fail does, so buffers are freed eagerly and a stray post-done CopyData gets the protocol-violation error instead of a finished parser. Signed-off-by: Ning Sun * chore: bypass typo check * chore: format change * fix(postgres): multi-chunk COPY FROM STDIN corrupted at chunk boundaries The csv feed loop flushed the arrow-csv decoder at every CopyData chunk end — including partial batches at record boundaries — but arrow-csv's contract permits flush() only after decode() returned 0 (a full batch of complete records). Flushing outside that point corrupts subsequent decoding: fields mis-parse shifted by the size of the record prefix buffered so far ('000001,host-001,...' parses as '00001h,ost-001u,...'), either erroring (22P02/22P04) or, worse, silently writing wrong values. Which chunks trigger it depends on where record boundaries land, so it appears intermittent with random-length data (psql 8KB / pgx ~64KB chunking both hit it). Adopt the strict decoder protocol: feed decodes until input is exhausted and flushes only when decode returns 0 (batch full); trailing partial batches surface at the EOF signal in finish(). With batch_size == COPY_IN_FLUSH_ROWS the flush threshold is unchanged; sub-threshold streams now emit their rows at CopyDone instead of per chunk, which is semantically identical. Adds test_csv_chunked_feeds_are_alignment_safe: the same payload fed at chunk sizes 1..65531 (plus shifted starts) with every row's content verified — reproduces the corruption on the old loop at chunk sizes 7, 8, 13, 64 and 1000. Verified with the reporter's Go reproducer (pgx, content-checked): 10/10 clean at -chunk 8192, and 4096/1000/97/65531 plus the previously failing benchmark sizes 500/1000/2000/200/800 all round- trip exactly. Signed-off-by: Ning Sun * chore: bypass spell check * chore: bypass spell check * fix(postgres): address PR review — cross-catalog guard, identifier folding, tab delimiter Three review findings (shuiyisong, PR #9297): [P1] Enforce the catalog restriction on COPY targets. copy_in_table now applies validate_catalog_and_schema when QueryOptions:: disallow_cross_catalog_query is set, mirroring SQL INSERT's check_permission path, BEFORE begin_copy_in switches the query context to the target catalog (which previously made the restriction unenforceable downstream). Cross-catalog COPY is now rejected with 'Illegal access to catalog' like INSERT. [P2] Canonicalize unquoted identifiers like the SQL parser (ParserContext::canonicalize_identifier): unquoted table-name parts and columns fold to lowercase, quoted ones keep their case. COPY Metrics (Host, Ts) now addresses the table CREATE TABLE Metrics created, and no longer silently misroutes into a coexisting quoted "Metrics". [P2] Allow a tab as CSV delimiter: take_single_byte accepts a literal tab while still rejecting newlines and other control characters, and the option parser now also accepts E'...' escape-string constants (DELIMITER E'\t'), processing PostgreSQL's backslash escapes. Verified end-to-end with psql 18.6: unquoted/quoted mixed-case round- trip into the right tables, cross-catalog rejection asserted in a new frontend unit test (QueryOptions plugin + second catalog), and both DELIMITER '' and DELIMITER E'\t' import tab-separated data. Signed-off-by: Ning Sun * fix(postgres): accept PG boolean spellings in COPY IN; drop copy state on failure Two review findings (shuiyisong, PR #9297): [P2] Boolean columns round-trip between the COPY directions. arrow-csv only parses true/false while PostgreSQL COPY OUT emits t/f, so boolean tables failed to re-import their own export (22P02). Boolean columns now decode as Utf8 in the arrow schema and flush_batch converts them with PostgreSQL's boolean input grammar (true/false, t/f, yes/no, y/n, on/off, 1/0, case-insensitive, whitespace trimmed). Round-trip covered end-to-end: COPY ... TO STDOUT's exact bytes are fed back via COPY FROM STDIN and verified. [P2] Drop the COPY state on every server-side failure. on_copy_data and on_copy_done now take ownership of the state up front, so parse, conversion or insert errors drop the buffered data instead of retaining it on long-lived or pooled connections until the next COPY or connection close (the reviewer measured ~17 MiB retained after a failed 8 MiB row); on_copy_data success puts the state back. An e2e test drives a failed COPY followed by a successful one on the same connection. Signed-off-by: Ning Sun * fix(postgres): interpret offset-free COPY timestamps in the session timezone Review finding (shuiyisong, PR #9297): the arrow decode schema carried Timestamp fields with no timezone, so the CSV decoder parsed offset-free timestamp strings as UTC while SQL INSERT parses them in the session timezone — the same input string silently stored a different instant depending on the ingestion method (8h apart in Asia/Shanghai). CopyInCodec::create_parser now receives the session timezone and the csv decode schema attaches it to Timestamp fields, so naive strings parse like INSERT; explicit offsets in the input (Z, +08:00) still take precedence, and UTC sessions are unchanged. Covered by test_csv_timestamp_session_timezone (naive vs +08:00 vs Z) and the end-to-end test_postgres_copy_in_session_timezone_parity, which SETs TIMEZONE TO 'Asia/Shanghai', inserts and COPY-loads the same string, and asserts both CAST(ts AS BIGINT) readbacks agree on 1704038400000. Signed-off-by: Ning Sun * fix(postgres): don't unescape E-string COPY options twice Review finding (shuiyisong, PR #9297): sqlparser's tokenizer already decodes PostgreSQL escape-string constants before returning Token::EscapedStringLiteral, so the second unescape_pg_string pass altered the decoded value — DELIMITER E'\\t' (a literal backslash followed by t, two characters) was turned into a tab and wrongly accepted as a single-byte delimiter. Use the EscapedStringLiteral payload as-is, like SingleQuotedString, and delete unescape_pg_string. Regression coverage distinguishes the three cases: E'\t' -> tab, E'\\' -> single backslash, and E'\\t' -> rejected as a multi-character delimiter. Signed-off-by: Ning Sun * fix(postgres): COPY OUT timestamps carry session-tz offset; quoted empty strings survive COPY IN Two review findings (shuiyisong, PR #9297): [P2] Preserve timestamp instants across COPY OUT/IN in non-UTC sessions. COPY OUT rendered naive UTC text while COPY IN (since the session-tz fix) interprets offset-free strings as local time, so re-importing an export shifted instants by the zone offset. COPY OUT now renders timestamps with an explicit session-timezone offset (2024-01-01 08:00:00.000000+08): the copy-out arrow schema carries the session tz (arrow-pg converts to a tz-aware DateTime) and the pg field type becomes TIMESTAMPTZ so the offset-bearing text format applies. Embedded offsets win on re-import regardless of session, keeping the OUT -> IN round trip exact; UTC sessions are unchanged. [P2] Preserve quoted empty strings. The shared arrow-csv decoder nulls every empty field, so a quoted empty string ("" — an empty STRING in PostgreSQL) imported as NULL. A shadow quote-mask scanner runs over the same bytes with the decoder's own splitting rules and records per field whether it was quoted; flush_batch restores quoted empties to empty strings for string columns, while unquoted empty fields stay NULL. Masks track records across chunk boundaries and align with decoder batches. Covered end-to-end: test_postgres_copy_out_in_timestamp_round_trip (Asia/Shanghai export asserted to carry +08, exact-epoch re-import comparison) and test_postgres_copy_in_quoted_empty_string (NULL, quoted empty and non-empty in one import), plus unit tests including a chunk-boundary sweep for the quoted-empty mask alignment. Signed-off-by: Ning Sun * fix(postgres): keep COPY IN quote masks aligned after doubled quotes; allow quoted empties in NOT NULL columns Two review findings (fengjiachun, PR #9297): [P2] Doubled quotes desynchronized the shadow quote-mask scanner. In QuoteInQuoted the doubled-quote branch skipped the second quote AND the following byte, so a closing quote or delimiter directly after an escaped quote was swallowed and later fields lost their masks -- import semantics depended on CopyData chunk boundaries. The branch now consumes only the escaped quote and returns to the in-quoted state, mirroring the decoder. Regression: "a""b","" ... verified across chunk sizes 1-64. [P2] A NOT NULL string column rejected quoted empty strings before restoration could run: the decode schema preserved the target nullability, and arrow-csv's batch validation rejects the temporary NULL a quoted empty decodes to. String fields now decode as nullable (safe: quoted-vs-unquoted is exactly the restoration's job), and the target NOT NULL constraint is enforced in flush_batch after restoration, so "" imports as an empty string while an unquoted empty field fails with a not-null-constraint error. Regressions at unit and e2e level: NOT NULL column accepting "" and rejecting unquoted NULL. Signed-off-by: Ning Sun * fix(postgres): align COPY IN quote-mask scanner with the CSV decoder on blank lines and CR terminators Two review findings (fengjiachun, PR #9297): [P2] Blank lines emitted a spurious empty mask. csv-core skips blank lines, but the scanner queued one mask per blank record, so every subsequent decoded row consumed the wrong mask (a leading blank line turned a quoted empty into NULL and an unquoted empty into an empty string, and tripped NOT NULL on the wrong row). A blank line before the header also shifted the header expectation. end_record now returns no mask for a completely blank record without consuming the header expectation. [P2] Standalone CR was not a record terminator. The decoder accepts CR, CRLF and LF alike; the scanner only recognized CR before LF or at EOF, so CR-separated input desynchronized masks and flush_batch failed with 'quote masks out of sync'. is_record_end now treats CR alone as a terminator, consuming one byte (two for CRLF), while a CR at a chunk boundary still defers via carry so CRLF split across CopyData chunks stays one terminator. Regressions: blank lines before/inside records and before a header (NULL/empty preservation plus NOT NULL behavior), and mixed CR/CRLF/LF terminators across chunk sizes 1-64 with a trailing CR-only record. Signed-off-by: Ning Sun * refactor(postgres)!: drop the quote-mask scanner; COPY IN CSV behaves exactly like file COPY The shadow quote-mask scanner — kept to preserve PostgreSQL's quoted-empty-string vs unquoted-NULL distinction — mirrored the decoder's record-splitting grammar by hand and required a fix for every rule it drifted on (doubled quotes, blank lines, CR terminators, NOT NULL interplay). Per maintainers' decision, that complexity is not worth the divergence: COPY FROM STDIN now behaves EXACTLY like the file-based COPY ... FROM 'file.csv' — an empty field, quoted or not, imports as NULL (arrow-csv's engine semantics, identical in both routes). Removed: QuoteMaskScanner and helpers, the always-nullable string decode schema, the restoration + NOT NULL enforcement in flush_batch, set_string_at, and the scanner-era unit tests. The decode schema again preserves target nullability, so the decoder enforces NOT NULL itself, same as the file route. Kept (small, schema-level, no second grammar): boolean spellings (t/f) and session-timezone timestamp decoding, both addressing separate review findings. E2E: test_postgres_copy_in_empty_field_is_null pins the parity — quoted empty, unquoted empty and non-empty in one import. Signed-off-by: Ning Sun * chore: remove empty file * chore: lint issues --------- Signed-off-by: Ning Sun --- Cargo.lock | 1 + src/frontend/src/instance.rs | 66 ++ src/frontend/src/instance/copy_in.rs | 87 +++ src/frontend/src/server.rs | 1 + src/servers/src/metrics.rs | 5 + src/servers/src/postgres.rs | 14 + src/servers/src/postgres/copy_in/codec.rs | 116 +++ src/servers/src/postgres/copy_in/csv.rs | 703 ++++++++++++++++++ src/servers/src/postgres/copy_in/mod.rs | 691 +++++++++++++++++ src/servers/src/postgres/copy_in/options.rs | 281 +++++++ src/servers/src/postgres/copy_in/parse.rs | 443 +++++++++++ src/servers/src/postgres/handler.rs | 115 ++- src/servers/src/postgres/server.rs | 11 +- src/servers/src/postgres/types.rs | 21 +- src/servers/src/postgres/types/error.rs | 12 + src/servers/src/query_handler.rs | 31 + src/servers/tests/mod.rs | 63 +- src/servers/tests/postgres/mod.rs | 267 ++++++- tests-integration/Cargo.toml | 1 + tests-integration/src/test_util.rs | 7 +- tests-integration/src/tests.rs | 1 + .../src/tests/postgres_copy_in_test.rs | 521 +++++++++++++ tests-integration/tests/sql.rs | 109 +++ 23 files changed, 3549 insertions(+), 18 deletions(-) create mode 100644 src/frontend/src/instance/copy_in.rs create mode 100644 src/servers/src/postgres/copy_in/codec.rs create mode 100644 src/servers/src/postgres/copy_in/csv.rs create mode 100644 src/servers/src/postgres/copy_in/mod.rs create mode 100644 src/servers/src/postgres/copy_in/options.rs create mode 100644 src/servers/src/postgres/copy_in/parse.rs create mode 100644 tests-integration/src/tests/postgres_copy_in_test.rs diff --git a/Cargo.lock b/Cargo.lock index c8dfd19620a..16739107f16 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -15394,6 +15394,7 @@ dependencies = [ "auth", "axum 0.8.4", "base64 0.22.1", + "bytes", "cache", "catalog", "chrono", diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index 2e49bf0b7cc..67f1e6003a0 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -13,6 +13,7 @@ // limitations under the License. pub mod builder; +mod copy_in; mod dashboard; mod entity_graph; mod export_database; @@ -4699,4 +4700,69 @@ mod tests { assert!(matches!(&submitted[0], DdlTask::CreateTable(_))); Ok(()) } + + /// COPY FROM STDIN must enforce the same cross-catalog restriction as + /// SQL INSERT when `disallow_cross_catalog_query` is configured. + #[tokio::test] + async fn test_copy_in_cross_catalog_restriction() -> TestResult<()> { + use catalog::memory::MemoryCatalogManager; + use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; + use servers::query_handler::CopyInHandler; + + let catalog = MemoryCatalogManager::new_with_table(test_table(1024, "source")?); + // A table in another catalog. + catalog.register_catalog_sync("other").unwrap(); + catalog + .register_schema_sync(catalog::RegisterSchemaRequest { + catalog: "other".to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + }) + .unwrap(); + catalog + .register_table_sync(catalog::RegisterTableRequest { + catalog: "other".to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: "copy_guard".to_string(), + table_id: 2048, + table: test_table(2048, "copy_guard")?, + }) + .unwrap(); + let executor = Arc::new(MockProcedureExecutor::new(catalog.clone())); + let plugins = Plugins::new(); + plugins.insert(QueryOptions { + disallow_cross_catalog_query: true, + }); + let instance = test_instance_with_catalog_manager( + catalog, + test_table(1025, "target")?, + plugins, + None, + executor, + ) + .await?; + + let ctx = Arc::new(QueryContext::with( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + )); + + // Cross-catalog COPY is rejected before any data is streamed. + let err = instance + .copy_in_table("other", DEFAULT_SCHEMA_NAME, "copy_guard", ctx.clone()) + .await + .unwrap_err(); + assert!( + format!("{err:?}").contains("Illegal access to catalog"), + "{err}" + ); + + // Same-catalog COPY resolves normally. + let found = instance + .copy_in_table(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, "target", ctx) + .await + .unwrap(); + assert!(found.is_some()); + + Ok(()) + } } diff --git a/src/frontend/src/instance/copy_in.rs b/src/frontend/src/instance/copy_in.rs new file mode 100644 index 00000000000..6580cc44b9d --- /dev/null +++ b/src/frontend/src/instance/copy_in.rs @@ -0,0 +1,87 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use api::v1::RowInsertRequests; +use async_trait::async_trait; +use auth::{PermissionChecker, PermissionCheckerRef, PermissionReq}; +use common_error::ext::BoxedError; +use common_query::Output; +use query::query_engine::options::{QueryOptions, validate_catalog_and_schema}; +use servers::error::{AuthSnafu, ExecuteGrpcRequestSnafu, ExecuteQuerySnafu, Result}; +use servers::query_handler::CopyInHandler; +use session::context::QueryContextRef; +use snafu::ResultExt; +use table::metadata::TableInfoRef; + +use crate::instance::Instance; + +#[async_trait] +impl CopyInHandler for Instance { + async fn copy_in_table( + &self, + catalog: &str, + schema: &str, + table: &str, + query_ctx: QueryContextRef, + ) -> Result> { + // Mirror SQL INSERT's catalog restriction (see `check_permission`): + // when cross-catalog queries are disallowed, a COPY may only + // target the current catalog. This must run against the ORIGINAL + // query context — `begin_copy_in` switches the context to the + // target catalog afterwards, so nothing downstream would catch + // it. + let disallow_cross_catalog_query = self + .plugins + .get::() + .map(|opts| opts.disallow_cross_catalog_query) + .unwrap_or_default(); + if disallow_cross_catalog_query { + validate_catalog_and_schema(catalog, schema, &query_ctx) + .map_err(BoxedError::new) + .context(ExecuteQuerySnafu)?; + } + + // A COPY writes rows like a SQL INSERT; check the write permission + // before the client starts streaming data. + self.plugins + .get::() + .as_ref() + .check_permission( + query_ctx.current_user(), + PermissionReq::BulkInsert { + catalog, + schema, + table, + }, + ) + .context(AuthSnafu)?; + + self.catalog_manager() + .table(catalog, schema, table, None) + .await + .map(|table| table.map(|table| table.table_info())) + .map_err(servers::error::Error::from) + } + + async fn copy_in_insert( + &self, + requests: RowInsertRequests, + query_ctx: QueryContextRef, + ) -> Result { + self.handle_row_inserts(requests, query_ctx, false, false) + .await + .map_err(BoxedError::new) + .context(ExecuteGrpcRequestSnafu) + } +} diff --git a/src/frontend/src/server.rs b/src/frontend/src/server.rs index 8de7bfe340f..0b1e2b9713e 100644 --- a/src/frontend/src/server.rs +++ b/src/frontend/src/server.rs @@ -420,6 +420,7 @@ where let pg_server = Box::new( PostgresServer::new( + instance.clone(), instance.clone(), opts.tls.should_force_tls(), tls_server_config, diff --git a/src/servers/src/metrics.rs b/src/servers/src/metrics.rs index 888bdc42cdf..d928c60c741 100644 --- a/src/servers/src/metrics.rs +++ b/src/servers/src/metrics.rs @@ -320,6 +320,11 @@ lazy_static! { "servers postgres prepared count" ) .unwrap(); + pub static ref METRIC_POSTGRES_COPY_IN_ROWS: IntCounter = register_int_counter!( + "greptime_servers_postgres_copy_in_rows", + "servers postgres COPY FROM STDIN rows written" + ) + .unwrap(); pub static ref METRIC_SERVER_GRPC_DB_REQUEST_TIMER: HistogramVec = register_histogram_vec!( "greptime_servers_grpc_db_request_elapsed", "servers grpc db request elapsed", diff --git a/src/servers/src/postgres.rs b/src/servers/src/postgres.rs index 15aae48a494..e13ec7392d5 100644 --- a/src/servers/src/postgres.rs +++ b/src/servers/src/postgres.rs @@ -13,6 +13,7 @@ // limitations under the License. mod auth_handler; +mod copy_in; mod fixtures; mod handler; mod server; @@ -42,6 +43,7 @@ use session::context::Channel; use self::auth_handler::{PgUserProviderSource, ZeroPidSecretKeyGenerator}; use self::handler::DefaultQueryParser; +use crate::query_handler::CopyInHandlerRef; use crate::query_handler::sql::ServerSqlQueryHandlerRef; pub(crate) struct GreptimeDBStartupParameters { @@ -78,6 +80,7 @@ impl ServerParameterProvider for GreptimeDBStartupParameters { pub struct PostgresServerHandlerInner { query_handler: ServerSqlQueryHandlerRef, auth_source: Option>, + copy_in_handler: CopyInHandlerRef, force_tls: bool, auth_driver: AuthDriver, UserInfoRef>, /// SCRAM selected; failures get the uniform metric label. @@ -86,11 +89,16 @@ pub struct PostgresServerHandlerInner { batching_enabled: bool, session: Arc, query_parser: Arc, + + /// State of an in-progress `COPY ... FROM STDIN`; set when the copy + /// starts and taken when it finishes, fails or the connection drops. + copy_in_state: tokio::sync::Mutex>, } #[derive(Builder)] pub(crate) struct MakePostgresServerHandler { query_handler: ServerSqlQueryHandlerRef, + copy_in_handler: CopyInHandlerRef, user_provider: Option, #[builder(default = "Arc::new(GreptimeDBStartupParameters::new())")] param_provider: Arc, @@ -112,6 +120,10 @@ impl PgWireServerHandlers for PostgresServerHandler { self.0.clone() } + fn copy_handler(&self) -> Arc { + self.0.clone() + } + fn error_handler(&self) -> Arc { self.0.clone() } @@ -136,6 +148,7 @@ impl MakePostgresServerHandler { .user_provider .clone() .map(|p| Arc::new(PgUserProviderSource::new(p))), + copy_in_handler: self.copy_in_handler.clone(), force_tls: self.force_tls, auth_driver: AuthDriver::new(self.param_provider.clone()) .with_pid_secret_key_generator(Arc::new(ZeroPidSecretKeyGenerator)), @@ -144,6 +157,7 @@ impl MakePostgresServerHandler { batching_enabled, session: session.clone(), query_parser: Arc::new(DefaultQueryParser::new(self.query_handler.clone(), session)), + copy_in_state: tokio::sync::Mutex::new(None), }; PostgresServerHandler(Arc::new(handler)) } diff --git a/src/servers/src/postgres/copy_in/codec.rs b/src/servers/src/postgres/copy_in/codec.rs new file mode 100644 index 00000000000..37a94811309 --- /dev/null +++ b/src/servers/src/postgres/copy_in/codec.rs @@ -0,0 +1,116 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! The COPY FROM STDIN codec framework. +//! +//! A *codec* binds one COPY data format to its implementation. Codecs are +//! stateless strategies: they validate the options they accept at +//! statement-parse time, describe which column types they can ingest, and +//! create the stateful [`RecordParser`] per stream. +//! +//! Formats are not limited to PostgreSQL's own — GreptimeDB-specific +//! formats (a line protocol, a JSON-lines dialect, ...) register exactly +//! like `csv` does, and may define their own options through the open +//! `WITH (...)` clause. +//! +//! # Parsing and conversion +//! +//! A [`RecordParser`] consumes raw `CopyData` bytes and emits complete, +//! insertable rows: one [`ParsedRow`] per record, in the order of the +//! targeted columns. Parsers validate record shape (field counts) and +//! convert field values themselves — `csv` delegates both splitting and +//! typed conversion to arrow-csv's incremental decoder (the engine +//! behind `COPY ... FROM 'file.csv'`) and turns decoded batches into +//! rows with the shared `api::helper::vectors_to_rows`. A future format +//! that carries values as text (e.g. a line protocol) would reintroduce +//! shared field-conversion helpers here; see the removed `text` codec in +//! git history for the shape they had. +//! +//! # Options +//! +//! The `WITH (...)` clause is open: every `KEY value` entry the client +//! writes is collected into a [`CopyOptionSet`] (see `options`) regardless +//! of whether any format understands it — the statement is parsed +//! textually (see `parse`) precisely so options beyond sqlparser's COPY +//! grammar can reach the codecs. A codec builder takes the options it +//! understands (`take_string`, `take_single_byte`, `take_flag`, or raw +//! `take_value` for codec-specific shapes) and rejects the rest with +//! `reject_leftovers`, so a future format like +//! `COPY t FROM STDIN WITH (FORMAT influxdb_line, precision 'ns')` only +//! needs its own codec and registry entry. +//! +//! To add a format: +//! +//! 1. implement [`CopyInCodec`] for a config struct holding the format's +//! validated options (see `csv::CsvCodec` for the example), +//! 2. write a builder `fn(CopyOptionSet) -> PgWireResult>` that applies defaults and rejects options the format +//! does not accept, +//! 3. register the builder under every accepted `FORMAT` name (including +//! aliases) in `REGISTERED_FORMATS` in the parent module. + +use datatypes::prelude::ConcreteDataType; +use datatypes::schema::ColumnSchema; +use pgwire::error::PgWireResult; + +/// A parsed, insertable row: one gRPC value per targeted column, in COPY +/// field order. Every parser emits these directly, whatever mechanism it +/// used to produce them. +pub(super) type ParsedRow = api::v1::Row; + +/// Incremental parser of one COPY FROM STDIN stream. +/// +/// The wire protocol delivers the payload as an arbitrary sequence of +/// `CopyData` chunks; implementations must keep state across [`feed`] +/// calls because records (and even single fields) may span chunk +/// boundaries. [`finish`] is called at `CopyDone`; a final record without +/// a trailing record separator is accepted. +pub(super) trait RecordParser: Send { + /// Feeds a chunk of `CopyData` bytes and returns all complete rows. + fn feed(&mut self, data: &[u8]) -> PgWireResult>; + /// Parses the buffered tail at `CopyDone`; returns no row when the + /// stream ended exactly at a record boundary. + fn finish(&mut self) -> PgWireResult>; +} + +/// A COPY FROM STDIN data format implementation. +/// +/// Implementations are shared per statement (hence `Send + Sync`) and must +/// be stateless; per-stream state belongs to the [`RecordParser`] they +/// create. Every method is tailored to the format: `check_column_type` +/// declares the types the format can decode (csv is bounded by what +/// arrow-csv supports). +pub(super) trait CopyInCodec: Send + Sync + std::fmt::Debug { + /// The canonical format name, used in error messages. + fn name(&self) -> &'static str; + /// The format code advertised in the `CopyInResponse` header: 0 when + /// the client sends textual data (the default), 1 for binary. + fn format_code(&self) -> i8 { + 0 + } + /// Validates that a column of `data_type` can be ingested by this + /// format. + fn check_column_type(&self, data_type: &ConcreteDataType) -> PgWireResult<()>; + /// Creates the stateful parser for one copy stream over the targeted + /// columns (in COPY field order). + /// + /// `timezone` is the session timezone (an IANA name or a UTC offset): + /// formats that parse offset-free timestamp strings from text should + /// interpret them in this zone, matching SQL INSERT. + fn create_parser( + &self, + columns: &[ColumnSchema], + timezone: &str, + ) -> PgWireResult>; +} diff --git a/src/servers/src/postgres/copy_in/csv.rs b/src/servers/src/postgres/copy_in/csv.rs new file mode 100644 index 00000000000..b2919210e3f --- /dev/null +++ b/src/servers/src/postgres/copy_in/csv.rs @@ -0,0 +1,703 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! The `csv` format — currently the only implemented COPY FROM STDIN +//! format. +//! +//! Parsing and typed conversion both reuse arrow-csv's incremental +//! [`Decoder`](arrow::csv::reader::Decoder) — the same engine backing +//! GreptimeDB's file-based `COPY ... FROM 'file.csv'` (see +//! `common-datasource`'s `csv` file format) — so both COPY routes accept +//! identical data. The decoder splits records (quoting, escaped quotes, +//! custom delimiters, `\n`/`\r\n` separators, optional headers) and +//! converts fields against the targeted columns' schema in one pass. +//! +//! Consequences of sharing the engine: JSON and Binary columns are +//! rejected, timestamps must be date/time strings (offset-free ones +//! parse in the session timezone), and an empty field — quoted or not — +//! is NULL, exactly like file-based COPY. + +use std::sync::Arc; + +use api::helper::vectors_to_rows; +use arrow::csv::ReaderBuilder; +use arrow::csv::reader::Decoder as CsvDecoder; +use arrow_schema::{ArrowError, DataType as ArrowDataType, Field, Schema as ArrowSchema}; +use datatypes::data_type::DataType; +use datatypes::prelude::ConcreteDataType; +use datatypes::schema::{ColumnSchema, Schema as GreptimeSchema}; +use datatypes::vectors::{BooleanVector, Helper, StringVector, Vector}; +use pgwire::error::{PgWireError, PgWireResult}; + +use crate::postgres::copy_in::codec::{CopyInCodec, ParsedRow, RecordParser}; +use crate::postgres::copy_in::options::CopyOptionSet; +use crate::postgres::copy_in::{COPY_IN_FLUSH_ROWS, bad_copy_data, invalid_text_repr, unsupported}; + +/// The validated `csv` format configuration. +#[derive(Debug)] +pub(super) struct CsvCodec { + delimiter: u8, + quote: u8, + escape: u8, + header: bool, +} + +/// Builds the CSV codec from the parsed options. +pub(super) fn build_csv_codec(mut options: CopyOptionSet) -> PgWireResult> { + let delimiter = options.take_single_byte("DELIMITER")?.unwrap_or(b','); + let quote = options.take_single_byte("QUOTE")?.unwrap_or(b'"'); + // In PostgreSQL the default escape character for CSV is the quote + // character itself (doubled quotes inside quoted fields). + let escape = options.take_single_byte("ESCAPE")?.unwrap_or(quote); + // An empty string is the default NULL marker (an empty field is NULL); + // a custom NULL string is not supported because arrow-csv's null + // matching does not distinguish quoted from unquoted fields. + if let Some(null) = options.take_string("NULL")? + && !null.is_empty() + { + return Err(unsupported( + "custom NULL string is not supported by format csv; an empty field is NULL", + )); + } + let header = options.take_flag("HEADER")?.unwrap_or(false); + options.reject_leftovers("csv")?; + if delimiter == quote { + return Err(bad_copy_data("DELIMITER must not be equal to QUOTE")); + } + Ok(Arc::new(CsvCodec { + delimiter, + quote, + escape, + header, + })) +} + +impl CopyInCodec for CsvCodec { + fn name(&self) -> &'static str { + "csv" + } + + /// The column types arrow-csv can decode against. + fn check_column_type(&self, data_type: &ConcreteDataType) -> PgWireResult<()> { + let supported = matches!( + data_type, + ConcreteDataType::Boolean(_) + | ConcreteDataType::Int8(_) + | ConcreteDataType::Int16(_) + | ConcreteDataType::Int32(_) + | ConcreteDataType::Int64(_) + | ConcreteDataType::UInt8(_) + | ConcreteDataType::UInt16(_) + | ConcreteDataType::UInt32(_) + | ConcreteDataType::UInt64(_) + | ConcreteDataType::Float32(_) + | ConcreteDataType::Float64(_) + | ConcreteDataType::String(_) + | ConcreteDataType::Date(_) + | ConcreteDataType::Timestamp(_) + | ConcreteDataType::Decimal128(_) + ); + if !supported { + return Err(unsupported(format!( + "column type \"{}\" is not supported by format csv", + data_type.name() + ))); + } + Ok(()) + } + + fn create_parser( + &self, + columns: &[ColumnSchema], + timezone: &str, + ) -> PgWireResult> { + // Decode schema: the greptime schema's arrow projection, except + // Boolean columns decode as Utf8 — arrow-csv only accepts + // `true`/`false`, while PostgreSQL COPY IN accepts its full set + // of boolean spellings (`t`, `f`, `yes`, `on`, ...), which COPY + // OUT emits — and Timestamp fields carry the SESSION timezone so + // offset-free strings parse like SQL INSERT does instead of as + // UTC. The substitutions are applied in `flush_batch` and by the + // decoder's tz-aware parsing respectively; explicit offsets in + // the input always win. + let boolean_columns: Vec = columns + .iter() + .enumerate() + .filter(|(_, c)| matches!(c.data_type, ConcreteDataType::Boolean(_))) + .map(|(i, _)| i) + .collect(); + let arrow_schema = GreptimeSchema::try_new(columns.to_vec()) + .map_err(|e| unsupported(format!("invalid columns for format csv: {e}")))? + .arrow_schema() + .clone(); + let timezone = timezone.to_string(); + let fields = arrow_schema + .fields() + .iter() + .zip(columns) + .map(|(field, column)| match &column.data_type { + ConcreteDataType::Boolean(_) => { + Field::new(field.name(), ArrowDataType::Utf8, field.is_nullable()) + } + ConcreteDataType::Timestamp(_) => match field.data_type() { + ArrowDataType::Timestamp(unit, None) => Field::new( + field.name(), + ArrowDataType::Timestamp(*unit, Some(Arc::from(timezone.as_str()))), + field.is_nullable(), + ), + _ => field.as_ref().clone(), + }, + _ => field.as_ref().clone(), + }) + .collect::>(); + let mut builder = ReaderBuilder::new(Arc::new(ArrowSchema::new(fields))) + .with_header(self.header) + .with_delimiter(self.delimiter) + .with_quote(self.quote) + .with_batch_size(COPY_IN_FLUSH_ROWS); + // An escape character equal to the quote is the doubled-quote + // default; csv-core handles that through quoting alone. + if self.escape != self.quote { + builder = builder.with_escape(self.escape); + } + Ok(Box::new(CsvRecordParser { + decoder: builder.build_decoder(), + boolean_columns, + })) + } +} + +/// Incremental parser backed by arrow-csv's decoder. +struct CsvRecordParser { + decoder: CsvDecoder, + /// Indices of Boolean columns, decoded as Utf8 and converted by + /// [`flush_batch`] to accept PostgreSQL's boolean spellings. + boolean_columns: Vec, +} + +impl RecordParser for CsvRecordParser { + fn feed(&mut self, data: &[u8]) -> PgWireResult> { + let mut buf = data; + let mut rows = Vec::new(); + while !buf.is_empty() { + let consumed = self.decoder.decode(buf).map_err(decode_error)?; + if consumed == 0 { + while let Some(batch) = self.decoder.flush().map_err(decode_error)? { + rows.extend(self.flush_batch(&batch)?); + } + } else { + buf = &buf[consumed..]; + } + } + Ok(rows) + } + + fn finish(&mut self) -> PgWireResult> { + // An empty decode signals EOF, finalizing a trailing record that + // lacks its record separator. + self.decoder.decode(&[]).map_err(decode_error)?; + let mut rows = Vec::new(); + while let Some(batch) = self.decoder.flush().map_err(decode_error)? { + rows.extend(self.flush_batch(&batch)?); + } + Ok(rows) + } +} + +impl CsvRecordParser { + fn flush_batch( + &mut self, + batch: &arrow::record_batch::RecordBatch, + ) -> PgWireResult> { + let mut vectors = Helper::try_into_vectors(batch.columns()) + .map_err(|e| bad_copy_data(format!("failed to convert decoded CSV batch: {e}")))?; + for &index in &self.boolean_columns { + let Some(vector) = vectors[index].as_any().downcast_ref::() else { + return Err(bad_copy_data( + "expected boolean column to decode as a string vector", + )); + }; + let mut values = Vec::with_capacity(batch.num_rows()); + for row_index in 0..batch.num_rows() { + values.push(match vector.get(row_index) { + datatypes::value::Value::Null => None, + datatypes::value::Value::String(s) => Some(parse_pg_bool(s.as_utf8())?), + other => { + return Err(bad_copy_data(format!( + "expected string values in a boolean column, got {other:?}" + ))); + } + }); + } + vectors[index] = Arc::new(BooleanVector::from(values)); + } + Ok(vectors_to_rows(vectors.iter(), batch.num_rows())) + } +} + +/// Parses PostgreSQL's boolean spellings: `true`/`false`, `t`/`f`, +/// `yes`/`no`, `on`/`off` and `1`/`0`, case-insensitive with surrounding +/// whitespace allowed (see the BOOLEAN type's input grammar). COPY OUT +/// emits `t`/`f`, so accepting them keeps both COPY directions +/// round-trippable. +fn parse_pg_bool(input: &str) -> PgWireResult { + match input.trim().to_ascii_lowercase().as_str() { + "true" | "t" | "yes" | "y" | "on" | "1" => Ok(true), + "false" | "f" | "no" | "n" | "off" | "0" => Ok(false), + _ => Err(invalid_text_repr(format!( + "invalid input syntax for type boolean: \"{input}\"" + ))), + } +} + +fn decode_error(error: ArrowError) -> PgWireError { + match &error { + ArrowError::CsvError(message) => { + bad_copy_data(format!("failed to parse CSV data: {message}")) + } + ArrowError::ParseError(message) | ArrowError::CastError(message) => { + invalid_text_repr(format!("failed to parse CSV data: {message}")) + } + _ => bad_copy_data(format!("failed to parse CSV data: {error}")), + } +} + +#[cfg(test)] +mod tests { + use datatypes::schema::ColumnSchema; + + use super::*; + use crate::postgres::copy_in::options::CopyOptionValue; + + fn test_columns() -> Vec { + vec![ + ColumnSchema::new("host", ConcreteDataType::string_datatype(), true), + ColumnSchema::new("val", ConcreteDataType::float64_datatype(), true), + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + ] + } + + fn parser() -> Box { + build_csv_codec(CopyOptionSet::default()) + .unwrap() + .create_parser(&test_columns(), "UTC") + .unwrap() + } + + fn string_at(row: &ParsedRow, index: usize) -> Option { + row.values[index].value_data.as_ref().map(|v| match v { + api::v1::value::ValueData::StringValue(s) => s.clone(), + other => panic!("expected string, got {other:?}"), + }) + } + + fn rows_with(mut set: CopyOptionSet) -> Box { + set.take_format().unwrap(); + build_csv_codec(set) + .unwrap() + .create_parser(&test_columns(), "UTC") + .unwrap() + } + + /// Feeds chunks and then the EOF signal, returning all rows: with the + /// strict decoder protocol, feeds only emit full + /// `COPY_IN_FLUSH_ROWS` batches; a trailing partial batch surfaces at + /// `finish`. + fn feed_all(parser: &mut Box, chunks: &[&[u8]]) -> Vec { + let mut rows = Vec::new(); + for chunk in chunks { + rows.extend(parser.feed(chunk).unwrap()); + } + rows.extend(parser.finish().unwrap()); + rows + } + + #[test] + fn test_build_defaults() { + let codec = build_csv_codec(CopyOptionSet::default()).unwrap(); + assert_eq!(codec.name(), "csv"); + } + + #[test] + fn test_build_rejects_bad_options() { + let mut set = CopyOptionSet::default(); + set.insert("delimiter".into(), CopyOptionValue::String("\"".into())) + .unwrap(); + let err = build_csv_codec(set).unwrap_err(); + assert!( + err.to_string() + .contains("DELIMITER must not be equal to QUOTE") + ); + + let mut set = CopyOptionSet::default(); + set.insert("null".into(), CopyOptionValue::String("NULL".into())) + .unwrap(); + assert!( + build_csv_codec(set) + .unwrap_err() + .to_string() + .contains("custom NULL string") + ); + + let mut set = CopyOptionSet::default(); + set.take_format().unwrap(); + set.insert("precision".into(), CopyOptionValue::Ident("ns".into())) + .unwrap(); + assert!( + build_csv_codec(set) + .unwrap_err() + .to_string() + .contains("not supported by format csv") + ); + } + + #[test] + fn test_check_column_type() { + let codec = build_csv_codec(CopyOptionSet::default()).unwrap(); + assert!( + codec + .check_column_type(&ConcreteDataType::json_datatype()) + .is_err() + ); + assert!( + codec + .check_column_type(&ConcreteDataType::binary_datatype()) + .is_err() + ); + assert!( + codec + .check_column_type(&ConcreteDataType::string_datatype()) + .is_ok() + ); + } + + #[test] + fn test_csv_parse_records() { + let mut parser = parser(); + let rows = feed_all( + &mut parser, + &[b"host1,1.5,2023-11-14 22:13:20\n\ + ,2.5,2023-11-14 22:13:21\n\ + \"quoted, comma\",3.5,2023-11-14 22:13:22\n"], + ); + assert_eq!(rows.len(), 3); + assert_eq!(string_at(&rows[0], 0), Some("host1".into())); + // An empty field is NULL. + assert_eq!(string_at(&rows[1], 0), None); + assert_eq!(string_at(&rows[2], 0), Some("quoted, comma".into())); + } + + #[test] + fn test_csv_escaped_quotes() { + let mut parser = parser(); + let rows = feed_all( + &mut parser, + &[b"\"say \"\"hi\"\"\",1.0,2023-11-14 22:13:20\n"], + ); + assert_eq!(string_at(&rows[0], 0), Some("say \"hi\"".into())); + } + + #[test] + fn test_csv_custom_escape() { + let mut set = CopyOptionSet::default(); + set.insert("escape".into(), CopyOptionValue::String("\\".into())) + .unwrap(); + let mut parser = rows_with(set); + let rows = feed_all(&mut parser, &[b"\"a\\\"b\",1.0,2023-11-14 22:13:20\n"]); + assert_eq!(string_at(&rows[0], 0), Some("a\"b".into())); + } + + #[test] + fn test_csv_quoted_newline_and_chunk_boundaries() { + let mut parser = parser(); + // Split inside a quoted field containing a newline, and a CRLF + // record separator split across chunks. No batch is full, so all + // feeds return nothing and both rows surface at the EOF signal. + assert!(parser.feed(b"\"multi").unwrap().is_empty()); + assert!( + parser + .feed(b"line\nfield\",1.0,2023-11-14 22:13:20") + .unwrap() + .is_empty() + ); + assert!( + parser + .feed(b"\r\nhost2,2.0,2023-11-14 22:13:21\n") + .unwrap() + .is_empty() + ); + let rows = parser.finish().unwrap(); + assert_eq!(rows.len(), 2); + assert_eq!(string_at(&rows[0], 0), Some("multiline\nfield".into())); + assert_eq!(string_at(&rows[1], 0), Some("host2".into())); + } + + #[test] + fn test_csv_header() { + let mut set = CopyOptionSet::default(); + set.insert("header".into(), CopyOptionValue::Flag).unwrap(); + let mut parser = rows_with(set); + let rows = feed_all( + &mut parser, + &[b"host,val,ts\nhost1,1.0,2023-11-14 22:13:20\n"], + ); + assert_eq!(rows.len(), 1); + assert_eq!(string_at(&rows[0], 0), Some("host1".into())); + } + + #[test] + fn test_csv_trailing_record_without_newline() { + let mut parser = parser(); + let fed = parser + .feed(b"host1,1.0,2023-11-14 22:13:20\nhost2,2.0,2023-11-14 22:13:21") + .unwrap(); + // While the trailing record is pending, flush cannot emit the + // completed one either; both surface at the EOF signal. + assert!(fed.is_empty()); + let rows = parser.finish().unwrap(); + assert_eq!(rows.len(), 2); + assert_eq!(string_at(&rows[0], 0), Some("host1".into())); + assert_eq!(string_at(&rows[1], 0), Some("host2".into())); + } + + #[test] + fn test_csv_finish_at_record_boundary() { + let mut parser = parser(); + // The complete record stays buffered (no full batch); the feed + // emits nothing and `finish` surfaces it. + assert!( + parser + .feed(b"host1,1.0,2023-11-14 22:13:20\n") + .unwrap() + .is_empty() + ); + let rows = parser.finish().unwrap(); + assert_eq!(rows.len(), 1); + assert_eq!(string_at(&rows[0], 0), Some("host1".into())); + } + + /// Regression test for chunk-boundary corruption: CopyData chunking is + /// arbitrary (records may split mid-field), and an earlier feed loop + /// flushed partial batches at record-boundary chunk ends, which + /// arrow-csv does not permit mid-stream — fields then parse + /// ("00001h,ost-001u,..."). Feed the same payload at many chunk sizes + /// and verify EVERY cell. + #[test] + fn test_csv_chunked_feeds_are_alignment_safe() { + // (int32 id, timestamp) rows of near-equal length. + let sweep_columns = vec![ + ColumnSchema::new("i", ConcreteDataType::int32_datatype(), true), + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + ]; + // A payload wider than one decoder batch so batch-full flushes + // interleave with chunk boundaries too. + let total = 2 * COPY_IN_FLUSH_ROWS + 137; + let payload: Vec = (0..total) + .map(|i| { + format!( + "{:06},2024-01-01 00:00:{:02}.{:03}\n", + i, + (i / 1000) % 60, + i % 1000 + ) + }) + .collect::() + .into_bytes(); + + let check = |chunk: usize| { + let mut parser = build_csv_codec(CopyOptionSet::default()) + .unwrap() + .create_parser(&sweep_columns, "UTC") + .unwrap(); + let mut ids: Vec = Vec::new(); + for c in payload.chunks(chunk) { + for row in parser.feed(c).unwrap() { + ids.push(match row.values[0].value_data.as_ref().unwrap() { + api::v1::value::ValueData::I32Value(v) => *v, + other => panic!("expected i32, got {other:?}"), + }); + } + } + for row in parser.finish().unwrap() { + ids.push(match row.values[0].value_data.as_ref().unwrap() { + api::v1::value::ValueData::I32Value(v) => *v, + other => panic!("expected i32, got {other:?}"), + }); + } + assert_eq!(ids.len(), total, "chunk={chunk}"); + for (i, v) in ids.iter().enumerate() { + assert_eq!(*v, i as i32, "chunk={chunk} row {i} corrupted"); + } + }; + + // Small chunk sizes hit every possible boundary alignment, + // including multiple decode calls per record. + for chunk in [1usize, 2, 3, 7, 8, 13, 31, 64, 97, 128, 1024] { + check(chunk); + } + // Client-natural sizes with a shifted start to vary where + // boundaries land relative to records. + for chunk in [4096usize, 8192, 65531] { + check(chunk); + let mut shifted = payload.clone(); + shifted.drain(..37); + let mut parser = build_csv_codec(CopyOptionSet::default()) + .unwrap() + .create_parser(&sweep_columns, "UTC") + .unwrap(); + let mut count = 0usize; + for c in shifted.chunks(chunk) { + count += parser.feed(c).unwrap().len(); + } + count += parser.finish().unwrap().len(); + assert_eq!(count, total - 1, "chunk={chunk} shifted"); + } + } + + #[test] + fn test_csv_boolean_spellings() { + // PostgreSQL COPY OUT emits `t`/`f`; COPY IN accepts the full + // boolean input grammar. + let columns = vec![ + ColumnSchema::new("flag", ConcreteDataType::boolean_datatype(), true), + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + ]; + let mut parser = build_csv_codec(CopyOptionSet::default()) + .unwrap() + .create_parser(&columns, "UTC") + .unwrap(); + let rows = feed_all( + &mut parser, + &[b"t,2024-01-01 00:00:00\n\ + f,2024-01-01 00:00:01\n\ + TRUE,2024-01-01 00:00:02\n\ + no,2024-01-01 00:00:03\n\ + on,2024-01-01 00:00:04\n\ + 0,2024-01-01 00:00:05\n\ + ,2024-01-01 00:00:06\n"], + ); + let expected = [ + Some(true), + Some(false), + Some(true), + Some(false), + Some(true), + Some(false), + None, + ]; + for (row, want) in rows.iter().zip(expected) { + assert_eq!( + row.values[0].value_data, + want.map(api::v1::value::ValueData::BoolValue), + "{:?}", + row.values[0].value_data + ); + } + } + + #[test] + fn test_csv_boolean_rejects_garbage() { + let columns = vec![ + ColumnSchema::new("flag", ConcreteDataType::boolean_datatype(), true), + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + ]; + let mut parser = build_csv_codec(CopyOptionSet::default()) + .unwrap() + .create_parser(&columns, "UTC") + .unwrap(); + parser.feed(b"maybe,2024-01-01 00:00:00\n").unwrap(); + let err = match parser.finish() { + Ok(_) => panic!("expected parse error"), + Err(e) => e, + }; + assert!( + err.to_string() + .contains("invalid input syntax for type boolean"), + "{err}" + ); + } + + #[test] + fn test_csv_timestamp_session_timezone() { + // Offset-free timestamp strings must parse in the session + // timezone (parity with SQL INSERT); explicit offsets win. + let columns = vec![ + ColumnSchema::new("host", ConcreteDataType::string_datatype(), true), + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + ]; + let mut parser = build_csv_codec(CopyOptionSet::default()) + .unwrap() + .create_parser(&columns, "Asia/Shanghai") + .unwrap(); + let rows = feed_all( + &mut parser, + &[b"host1,2024-01-01 00:00:00\n\ + host2,2024-01-01T00:00:00+08:00\n\ + host3,2024-01-01T00:00:00Z\n"], + ); + let millis = |row: &ParsedRow| match row.values[1].value_data.as_ref().unwrap() { + api::v1::value::ValueData::TimestampMillisecondValue(v) => *v, + other => panic!("expected ts, got {other:?}"), + }; + // 2024-01-01T00:00:00+08:00 == 2023-12-31T16:00:00Z + assert_eq!(millis(&rows[0]), 1704038400000); + // Explicit +08:00 offset: same instant as the naive parse. + assert_eq!(millis(&rows[1]), 1704038400000); + // Explicit Z: parsed as UTC regardless of the session zone. + assert_eq!(millis(&rows[2]), 1704067200000); + } + + #[test] + fn test_csv_errors() { + let mut p = parser(); + let err = match p.feed(b"host1,1.0\n") { + Ok(_) => panic!("expected parse error"), + Err(e) => e, + }; + assert!( + err.to_string().contains("incorrect number of fields"), + "{err}" + ); + + // Value conversion runs at flush time, so the bad double + // surfaces at the EOF signal, not at feed time. + let mut p = parser(); + p.feed(b"host1,not_a_double,2023-11-14 22:13:20\n").unwrap(); + let err = match p.finish() { + Ok(_) => panic!("expected parse error"), + Err(e) => e, + }; + assert!(err.to_string().contains("failed to parse"), "{err}"); + } +} diff --git a/src/servers/src/postgres/copy_in/mod.rs b/src/servers/src/postgres/copy_in/mod.rs new file mode 100644 index 00000000000..a3959a23830 --- /dev/null +++ b/src/servers/src/postgres/copy_in/mod.rs @@ -0,0 +1,691 @@ +// 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. + +//! PostgreSQL `COPY ... FROM STDIN` (copy-in) support. +//! +//! The wire handshake is driven by pgwire: returning [`Response::CopyIn`] +//! from a query handler makes the server send a `CopyInResponse` and enter +//! the copy-in sub-protocol; subsequent `CopyData` messages from the client +//! are dispatched to [`CopyHandler`], whose implementation below parses +//! the data stream incrementally and flushes it into the query handler as +//! [`RowInsertRequests`] batches. +//! +//! # Format framework +//! +//! Each COPY data format is implemented as a [`codec::CopyInCodec`]: a +//! stateless strategy that validates its options at statement-parse time, +//! describes the column types it can ingest, and creates the stateful +//! record parser for one stream. Formats register a builder under every +//! accepted `FORMAT` name in [`REGISTERED_FORMATS`] — `csv` is the only +//! implemented format today; PostgreSQL's `text`/`binary` and +//! GreptimeDB-specific formats alike can be added the same way. See the +//! `codec` module docs. + +mod codec; +mod csv; +mod options; +mod parse; + +use std::fmt::Debug; +use std::sync::Arc; + +use api::helper::ColumnDataTypeWrapper; +use api::v1::column_def::options_from_column_schema; +use api::v1::{ + Row as GrpcRow, RowInsertRequest, RowInsertRequests, Rows as GrpcRows, SemanticType, +}; +use async_trait::async_trait; +use codec::{CopyInCodec, RecordParser}; +use common_query::OutputData; +use common_telemetry::{info, warn}; +use datafusion::sql::sqlparser::ast::{CopyTarget, Statement as SqlParserStatement}; +use datatypes::data_type::DataType; +use datatypes::schema::{ColumnSchema, SchemaRef}; +use futures::{Sink, SinkExt, stream}; +use pgwire::api::copy::CopyHandler; +use pgwire::api::results::{CopyResponse, Response, Tag}; +use pgwire::api::{ClientInfo, PgWireConnectionState}; +use pgwire::error::{ErrorInfo, PgWireError, PgWireResult}; +use pgwire::messages::PgWireBackendMessage; +use pgwire::messages::copy::{CopyData, CopyDone, CopyFail}; +use session::context::QueryContextRef; +use table::metadata::TableInfo; + +use crate::metrics::METRIC_POSTGRES_COPY_IN_ROWS; +use crate::postgres::PostgresServerHandlerInner; +use crate::postgres::types::PgErrorCode; +use crate::postgres::utils::convert_err; + +/// Number of parsed rows buffered before flushing a batch to the handler. +/// +/// Each flush runs the full insert pipeline (catalog lookup, on-demand +/// schema check, partition split, region write) and, without the +/// pending-rows batcher, waits for it — so larger batches mean fewer +/// round trips at the cost of linearly more memory per connection +/// (decoder batch + rows buffer). The constant doubles as the arrow +/// decoder's batch size. High-rate COPY users should enable the +/// pending-rows batcher (`global-table-batch` / `BatcherOptions.protocols` +/// containing `Postgres`), which turns flushes into enqueues and makes +/// the per-batch fixed costs negligible. +pub(crate) const COPY_IN_FLUSH_ROWS: usize = 32768; + +/// Builds a codec from the format-neutral option set. +type CodecBuilder = fn(options::CopyOptionSet) -> PgWireResult>; + +/// The registered COPY FROM STDIN formats: accepted `FORMAT` name to +/// builder. `csv` is the only implemented format; the others are +/// registered so they report a clear "not supported yet" instead of an +/// unknown-format error. +const REGISTERED_FORMATS: &[(&str, CodecBuilder)] = &[ + ("csv", csv::build_csv_codec), + ("text", build_unimplemented_text_codec), + ("txt", build_unimplemented_text_codec), + ("binary", build_unimplemented_binary_codec), +]; + +fn build_unimplemented_text_codec( + _options: options::CopyOptionSet, +) -> PgWireResult> { + Err(unsupported( + "the text format is not supported yet; use WITH (FORMAT csv)", + )) +} + +fn build_unimplemented_binary_codec( + _options: options::CopyOptionSet, +) -> PgWireResult> { + Err(unsupported( + "COPY FROM STDIN WITH (FORMAT binary) is not supported yet", + )) +} + +/// Resolves the format codec from the parsed options. +fn resolve_codec(options: &mut options::CopyOptionSet) -> PgWireResult> { + let name = options.take_format()?.ok_or_else(|| { + unsupported( + "COPY FROM STDIN requires an explicit format in this version; use WITH (FORMAT csv)", + ) + })?; + let builder = REGISTERED_FORMATS + .iter() + .find(|(registered, _)| *registered == name) + .map(|(_, builder)| *builder) + .ok_or_else(|| bad_copy_data(format!("unknown COPY format \"{}\"", name)))?; + builder(std::mem::take(options)) +} + +/// The parsed `COPY tbl [(columns)] FROM STDIN [(options)]` statement. +#[derive(Debug, Clone)] +pub(crate) struct CopyFromStdin { + /// Table name parts as written (`t`, `schema.t`, `catalog.schema.t`); + /// quoted names keep their case. + table: Vec, + /// Requested column list; empty means all columns in table order. + columns: Vec, + /// The resolved format codec, validated and ready to run. + codec: Arc, +} + +fn copy_in_error(code: PgErrorCode, message: impl Into) -> PgWireError { + PgWireError::UserError(Box::new(code.to_err_info(message.into()))) +} + +fn unsupported(message: impl Into) -> PgWireError { + copy_in_error(PgErrorCode::Ec0A000, message) +} + +fn bad_copy_data(message: impl Into) -> PgWireError { + copy_in_error(PgErrorCode::Ec22P04, message) +} + +fn invalid_text_repr(message: impl Into) -> PgWireError { + copy_in_error(PgErrorCode::Ec22P02, message) +} + +pub(crate) fn copy_in_must_run_alone_error() -> PgWireError { + PgWireError::UserError(Box::new(ErrorInfo::new( + "ERROR".to_string(), + PgErrorCode::Ec0A000.code(), + "COPY FROM STDIN must be executed alone".to_string(), + ))) +} + +/// Parses a leading `COPY ... FROM STDIN` statement from SQL text. +/// +/// The statement is parsed textually — not through sqlparser — because the +/// `WITH (...)` option list is open: formats may accept options beyond +/// sqlparser's COPY grammar, and unknown options must reach the format +/// codec instead of failing statement parsing. Returns `Ok(None)` when +/// the string does not lead with a copy-in statement. +pub(crate) fn copy_from_stdin_from_sql(sql: &str) -> PgWireResult> { + let raw = match parse::scan_copy_from_stdin(sql)? { + parse::CopyInScan::NotCopyIn => return Ok(None), + parse::CopyInScan::CopyIn(raw) => raw, + }; + let mut option_set = raw.options; + let codec = resolve_codec(&mut option_set)?; + Ok(Some(CopyFromStdin { + table: raw.table, + columns: raw.columns, + codec, + })) +} + +/// Whether an already-parsed statement is a `COPY ... FROM STDIN`; used to +/// reject copy-in statements trailing other statements, where the textual +/// scanner (which only sees statement-leading COPYs) does not see them. +pub(crate) fn is_copy_from_stdin(statement: &SqlParserStatement) -> bool { + matches!( + statement, + SqlParserStatement::Copy { + to: false, + target: CopyTarget::Stdin, + .. + } + ) +} + +/// Resolves table name parts against the query context defaults. +fn resolve_table_name( + query_ctx: &QueryContextRef, + parts: &[String], +) -> PgWireResult<(String, String, String)> { + match parts { + [table] => Ok(( + query_ctx.current_catalog().to_string(), + query_ctx.current_schema(), + table.clone(), + )), + [schema, table] => Ok(( + query_ctx.current_catalog().to_string(), + schema.clone(), + table.clone(), + )), + [catalog, schema, table] => Ok((catalog.clone(), schema.clone(), table.clone())), + _ => Err(unsupported(format!( + "invalid table name for COPY: {}", + parts.join(".") + ))), + } +} + +/// Per-connection state of an in-progress COPY FROM STDIN. +pub(crate) struct CopyInState { + query_ctx: QueryContextRef, + /// The gRPC column schemas of the targeted columns. Cloned into each + /// flushed batch: `Rows.schema` is an owned `Vec` in the proto + /// request, so one copy per request is inherent. + grpc_schema: Vec, + table_name: String, + parser: Box, + /// The resolved codec's name, for logging. + codec_name: &'static str, + /// The wire format code of the resolved codec. + codec_format_code: i8, + /// Parsed rows not yet flushed to the handler. + rows: Vec, + rows_written: u64, +} + +impl CopyInState { + fn new( + codec: Arc, + columns: Vec, + table_info: Arc, + query_ctx: QueryContextRef, + table_name: String, + ) -> PgWireResult { + let schema = table_info.meta.schema.clone(); + let timestamp_index = schema.timestamp_index(); + + let columns = resolve_columns(codec.as_ref(), columns, &schema, &table_info)?; + if let Some(ts_index) = timestamp_index { + let ts_column = &schema.column_schemas()[ts_index]; + if !columns.iter().any(|column| column.name == ts_column.name) { + return Err(unsupported(format!( + "COPY FROM STDIN requires the time index column \"{}\" in the column list", + ts_column.name + ))); + } + } + let parser = codec.create_parser(&columns, query_ctx.timezone().to_string().as_str())?; + + let grpc_schema = columns + .iter() + .map(|column| build_grpc_column_schema(column, &table_info)) + .collect::>>()?; + + Ok(Self { + query_ctx, + grpc_schema, + table_name, + parser, + codec_name: codec.name(), + codec_format_code: codec.format_code(), + rows: Vec::with_capacity(COPY_IN_FLUSH_ROWS), + rows_written: 0, + }) + } + + /// The wire format code advertised to the client. + pub(crate) fn format_code(&self) -> i8 { + self.codec_format_code + } + + /// Feeds a `CopyData` chunk into the parser and buffers parsed rows. + fn feed(&mut self, data: &[u8]) -> PgWireResult<()> { + self.rows.extend(self.parser.feed(data)?); + Ok(()) + } + + /// Parses the buffered tail at `CopyDone`; a final record without a + /// trailing newline is accepted. + fn finish(&mut self) -> PgWireResult<()> { + self.rows.extend(self.parser.finish()?); + Ok(()) + } + + fn should_flush(&self) -> bool { + self.rows.len() >= COPY_IN_FLUSH_ROWS + } + + /// Takes the buffered rows as an insert request, if any. + fn take_requests(&mut self) -> Option { + if self.rows.is_empty() { + return None; + } + let rows = GrpcRows { + schema: self.grpc_schema.clone(), + rows: std::mem::take(&mut self.rows), + }; + Some(RowInsertRequests { + inserts: vec![RowInsertRequest { + table_name: self.table_name.clone(), + rows: Some(rows), + }], + }) + } +} + +fn resolve_columns( + codec: &dyn CopyInCodec, + columns: Vec, + schema: &SchemaRef, + table_info: &Arc, +) -> PgWireResult> { + let table_name = &table_info.name; + let selected: Vec = if columns.is_empty() { + schema.column_schemas().to_vec() + } else { + let mut selected = Vec::with_capacity(columns.len()); + for name in &columns { + let column = schema + .column_schema_by_name(name) + .ok_or_else(|| { + copy_in_error( + PgErrorCode::Ec42703, + format!( + "column \"{}\" of relation \"{}\" does not exist", + name, table_name + ), + ) + })? + .clone(); + if selected.iter().any(|c: &ColumnSchema| &c.name == name) { + return Err(bad_copy_data(format!( + "column \"{}\" specified more than once", + name + ))); + } + selected.push(column); + } + selected + }; + + for column in &selected { + codec.check_column_type(&column.data_type)?; + } + + Ok(selected) +} + +/// Derives a column's semantic type the same way the flow adapter does +/// (`flow::adapter::util::column_schemas_to_proto`): the time-index flag +/// for timestamps, primary-key membership for tags. +fn semantic_type_of(table_info: &TableInfo, column: &ColumnSchema) -> PgWireResult { + let meta = &table_info.meta; + if column.is_time_index() { + Ok(SemanticType::Timestamp) + } else if meta + .primary_key_indices + .iter() + .any(|&index| meta.schema.column_schemas()[index].name == column.name) + { + Ok(SemanticType::Tag) + } else { + Ok(SemanticType::Field) + } +} + +fn build_grpc_column_schema( + column: &ColumnSchema, + table_info: &Arc, +) -> PgWireResult { + let (datatype, datatype_extension) = ColumnDataTypeWrapper::try_from(column.data_type.clone()) + .map_err(|_| { + unsupported(format!( + "column type \"{}\" is not supported for COPY FROM STDIN yet", + column.data_type.name() + )) + })? + .to_parts(); + Ok(api::v1::ColumnSchema { + column_name: column.name.clone(), + datatype: datatype as i32, + semantic_type: semantic_type_of(table_info, column)? as i32, + datatype_extension, + options: options_from_column_schema(column), + }) +} + +// --------------------------------------------------------------------------- +// pgwire handlers +// --------------------------------------------------------------------------- + +impl PostgresServerHandlerInner { + /// Prepares a COPY FROM STDIN: resolves the target table, validates + /// columns and stashes the per-connection copy state. + pub(crate) async fn begin_copy_in(&self, stmt: CopyFromStdin) -> PgWireResult { + // The batching-aware context routes copy-in flushes through the + // pending-rows batcher like SQL INSERT when it is enabled for + // this connection. + let query_ctx = self.new_query_context(); + let (catalog, schema, table) = resolve_table_name(&query_ctx, &stmt.table)?; + + let Some(table_info) = self + .copy_in_handler + .copy_in_table(&catalog, &schema, &table, query_ctx.clone()) + .await + .map_err(convert_err)? + else { + return Err(copy_in_error( + PgErrorCode::Ec42P01, + format!("relation \"{}\" does not exist", stmt.table.join(".")), + )); + }; + + // The insert path resolves the bare table name against the query + // context, so point the context at the resolved catalog and schema + // when the COPY target is qualified. + let query_ctx = adjust_query_context(query_ctx, &catalog, &schema); + + let state = CopyInState::new(stmt.codec, stmt.columns, table_info, query_ctx, table)?; + let num_columns = state.grpc_schema.len(); + let format_code = state.format_code(); + + *self.copy_in_state.lock().await = Some(state); + + Ok(Response::CopyIn(CopyResponse::new( + format_code, + num_columns, + stream::empty(), + ))) + } + + /// Flushes buffered rows of an in-progress copy. Returns the number of + /// rows written by the flushed batch. + async fn flush_copy_in(&self, state: &mut CopyInState) -> PgWireResult { + let Some(requests) = state.take_requests() else { + return Ok(0); + }; + let query_ctx = state.query_ctx.clone(); + let output = self + .copy_in_handler + .copy_in_insert(requests, query_ctx) + .await + .map_err(convert_err)?; + match output.data { + OutputData::AffectedRows(rows) => Ok(rows as u64), + OutputData::RecordBatches(_) | OutputData::Stream(_) => { + warn!("unexpected output from COPY FROM STDIN insert"); + Ok(0) + } + } + } +} + +fn adjust_query_context( + query_ctx: QueryContextRef, + catalog: &str, + schema: &str, +) -> QueryContextRef { + if query_ctx.current_catalog() == catalog && query_ctx.current_schema() == schema { + return query_ctx; + } + let mut forked = query_ctx.fork(); + forked.set_current_catalog(catalog); + forked.set_current_schema(schema); + Arc::new(forked) +} + +#[async_trait] +impl CopyHandler for PostgresServerHandlerInner { + async fn on_copy_data(&self, _client: &mut C, copy_data: CopyData) -> PgWireResult<()> + where + C: ClientInfo + Sink + Unpin + Send + Sync, + C::Error: Debug, + PgWireError: From<>::Error>, + { + let mut guard = self.copy_in_state.lock().await; + let Some(mut state) = guard.take() else { + return Err(copy_in_error( + PgErrorCode::Ec08P01, + "CopyData message received outside COPY FROM STDIN", + )); + }; + // The state is owned here: any parse, conversion or insert error + // below aborts the whole copy and drops the buffered data — + // otherwise a pooled connection would retain failed-import + // buffers until its next COPY. On success it is put back. + state.feed(©_data.data)?; + if state.should_flush() { + let rows = self.flush_copy_in(&mut state).await?; + state.rows_written += rows; + METRIC_POSTGRES_COPY_IN_ROWS.inc_by(rows); + } + *guard = Some(state); + Ok(()) + } + + async fn on_copy_done(&self, client: &mut C, _done: CopyDone) -> PgWireResult<()> + where + C: ClientInfo + Sink + Unpin + Send + Sync, + C::Error: Debug, + PgWireError: From<>::Error>, + { + let mut guard = self.copy_in_state.lock().await; + // Take ownership up front: even finalization failures (parse + // errors in the trailing record, insert errors) drop the state + // instead of retaining it on a long-lived connection. + let Some(mut state) = guard.take() else { + return Err(copy_in_error( + PgErrorCode::Ec08P01, + "CopyDone message received outside COPY FROM STDIN", + )); + }; + state.finish()?; + let rows = self.flush_copy_in(&mut state).await?; + state.rows_written += rows; + METRIC_POSTGRES_COPY_IN_ROWS.inc_by(rows); + + let rows_written = state.rows_written; + let codec_name = state.codec_name; + info!( + "PostgreSQL COPY FROM STDIN (format {}) finished: {} rows written", + codec_name, rows_written + ); + client + .send(PgWireBackendMessage::CommandComplete( + Tag::new("COPY").with_rows(rows_written as usize).into(), + )) + .await?; + + // After a successful extended-protocol copy-in, pgwire leaves the + // connection in `CopyInProgress(true)` and swallows the client's + // terminating `Sync` message, so `ReadyForQuery` would never be + // sent and the connection would deadlock. Route the state to + // `AwaitingSync` so the `Sync` is dispatched to `on_sync`, which + // completes the round trip. + if let PgWireConnectionState::CopyInProgress(true) = client.state() { + client.set_state(PgWireConnectionState::AwaitingSync); + } + Ok(()) + } + + async fn on_copy_fail(&self, _client: &mut C, fail: CopyFail) -> PgWireError + where + C: ClientInfo + Sink + Unpin + Send + Sync, + C::Error: Debug, + PgWireError: From<>::Error>, + { + // Drop the buffered state; partial data of a failed copy is not + // written. + self.copy_in_state.lock().await.take(); + copy_in_error( + PgErrorCode::Ec57014, + format!("COPY FROM STDIN mode terminated by user: {}", fail.message), + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_parse_copy_from_stdin_requires_format() { + let err = copy_from_stdin_from_sql("COPY t FROM STDIN").unwrap_err(); + assert!( + err.to_string() + .contains("requires an explicit format in this version"), + "{err}" + ); + + let copy = copy_from_stdin_from_sql("COPY t FROM STDIN WITH (FORMAT csv)") + .unwrap() + .unwrap(); + assert_eq!(copy.table, vec!["t"]); + assert!(copy.columns.is_empty()); + assert_eq!(copy.codec.name(), "csv"); + assert_eq!(copy.codec.format_code(), 0); + } + + #[test] + fn test_parse_copy_from_stdin_csv() { + let copy = copy_from_stdin_from_sql( + "COPY t (a, b) FROM STDIN WITH (FORMAT csv, HEADER, DELIMITER ';')", + ) + .unwrap() + .unwrap(); + assert_eq!(copy.columns, vec!["a", "b"]); + assert_eq!(copy.codec.name(), "csv"); + } + + #[test] + fn test_parse_copy_from_stdin_text_unsupported() { + for sql in [ + "COPY t FROM STDIN WITH (FORMAT text)", + "COPY t FROM STDIN WITH (FORMAT txt)", + ] { + let err = copy_from_stdin_from_sql(sql).unwrap_err(); + assert!( + err.to_string().contains("text format is not supported yet"), + "{err}" + ); + } + } + + #[test] + fn test_parse_copy_from_stdin_binary_unsupported() { + let err = copy_from_stdin_from_sql("COPY t FROM STDIN WITH (FORMAT binary)").unwrap_err(); + assert!(err.to_string().contains("not supported")); + } + + #[test] + fn test_parse_copy_from_stdin_unknown_format() { + let err = copy_from_stdin_from_sql("COPY t FROM STDIN WITH (FORMAT parquet)").unwrap_err(); + assert!(err.to_string().contains("unknown COPY format")); + } + + #[test] + fn test_parse_copy_from_stdin_not_copy_in() { + assert!( + copy_from_stdin_from_sql("COPY (SELECT 1) TO STDOUT") + .unwrap() + .is_none() + ); + assert!( + copy_from_stdin_from_sql("COPY t TO '/tmp/x.csv'") + .unwrap() + .is_none() + ); + // GreptimeDB's own file-based COPY FROM. + assert!( + copy_from_stdin_from_sql("COPY t FROM '/tmp/x.csv' WITH (FORMAT csv)") + .unwrap() + .is_none() + ); + } + + #[test] + fn test_redundant_options() { + let err = copy_from_stdin_from_sql("COPY t FROM STDIN WITH (FORMAT csv, FORMAT text)") + .unwrap_err(); + assert!(err.to_string().contains("redundant FORMAT")); + + let err = copy_from_stdin_from_sql("COPY t FROM STDIN WITH (DELIMITER ',', DELIMITER ';')") + .unwrap_err(); + assert!(err.to_string().contains("redundant DELIMITER")); + } + + #[test] + fn test_format_specific_options_rejected() { + // The open WITH section parses anything, but the format rejects + // the options it does not understand. + let err = + copy_from_stdin_from_sql("COPY t FROM STDIN WITH (FORMAT csv, PRECISION 'ns', FREEZE)") + .unwrap_err(); + assert!( + err.to_string().contains("not supported by format csv"), + "{err}" + ); + } + + #[test] + fn test_multi_statement_rejected() { + let err = copy_from_stdin_from_sql("COPY t FROM STDIN; SELECT 1").unwrap_err(); + assert!(err.to_string().contains("alone")); + // Trailing separators and quoted semicolons are fine. + assert!( + copy_from_stdin_from_sql("COPY t FROM STDIN WITH (FORMAT csv);") + .unwrap() + .is_some() + ); + assert!( + copy_from_stdin_from_sql("COPY t FROM STDIN WITH (FORMAT csv, DELIMITER ';')") + .unwrap() + .is_some() + ); + } +} diff --git a/src/servers/src/postgres/copy_in/options.rs b/src/servers/src/postgres/copy_in/options.rs new file mode 100644 index 00000000000..d710a0ca374 --- /dev/null +++ b/src/servers/src/postgres/copy_in/options.rs @@ -0,0 +1,281 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Format-neutral collection of COPY `WITH (...)` options. +//! +//! The option list is **open**: every `KEY value` entry the client writes +//! is collected here, regardless of whether any format understands it. +//! Each format codec then takes the options it knows ([`CopyOptionSet::take_string`] +//! and friends) and rejects everything else via +//! [`CopyOptionSet::reject_leftovers`]. This way a format can introduce +//! its own options — say a line-protocol format accepting +//! `WITH (FORMAT influxdb_line, precision 'ns')` — without touching the +//! statement parser or any other codec. + +use pgwire::error::{PgWireError, PgWireResult}; + +use crate::postgres::copy_in::bad_copy_data; +use crate::postgres::types::PgErrorCode; + +/// The value of one COPY option. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum CopyOptionValue { + /// A single-quoted string literal. + String(String), + /// A numeric literal, kept as written. + Number(String), + /// A bare identifier. + Ident(String), + /// `TRUE` / `FALSE`. + Boolean(bool), + /// The option appears without a value. + Flag, +} + +impl CopyOptionValue { + /// The value as a descriptive word for error messages. + pub(crate) fn kind(&self) -> &'static str { + match self { + CopyOptionValue::String(_) => "a string value", + CopyOptionValue::Number(_) => "a numeric value", + CopyOptionValue::Ident(_) => "an identifier value", + CopyOptionValue::Boolean(_) => "a boolean value", + CopyOptionValue::Flag => "no value", + } + } +} + +/// The options of a `COPY ... FROM STDIN WITH (...)` statement. +/// +/// Entries keep statement order; keys are lowercased for lookup and +/// displayed uppercased in error messages. +#[derive(Debug, Default)] +pub(crate) struct CopyOptionSet { + entries: Vec<(String, CopyOptionValue)>, +} + +impl CopyOptionSet { + /// Adds an option; duplicate keys are rejected like PostgreSQL does. + pub(crate) fn insert(&mut self, key: String, value: CopyOptionValue) -> PgWireResult<()> { + if self.entries.iter().any(|(k, _)| *k == key) { + return Err(bad_copy_data(format!( + "conflicting or redundant {} option", + key.to_uppercase() + ))); + } + self.entries.push((key, value)); + Ok(()) + } + + /// Takes the raw value of `key`; for codec-specific options whose + /// shape this module does not prescribe. + pub(crate) fn take_value(&mut self, key: &str) -> Option { + let key = key.to_lowercase(); + let index = self.entries.iter().position(|(k, _)| *k == key)?; + Some(self.entries.remove(index).1) + } + + /// Takes the `FORMAT` option, lowercased. Accepts identifier or string + /// values. + pub(crate) fn take_format(&mut self) -> PgWireResult> { + let Some(value) = self.take_value("format") else { + return Ok(None); + }; + match value { + CopyOptionValue::Ident(name) | CopyOptionValue::String(name) => { + Ok(Some(name.to_lowercase())) + } + other => Err(bad_copy_data(format!( + "FORMAT requires an identifier value, got {}", + other.kind() + ))), + } + } + + /// Takes a string-valued option. + pub(crate) fn take_string(&mut self, key: &str) -> PgWireResult> { + let Some(value) = self.take_value(key) else { + return Ok(None); + }; + match value { + CopyOptionValue::String(s) => Ok(Some(s)), + other => Err(bad_copy_data(format!( + "COPY option {} requires a string value, got {}", + key.to_uppercase(), + other.kind() + ))), + } + } + + /// Takes a single one-byte character option. A tab is a legal (and + /// common) CSV delimiter; newlines and other control characters are + /// not, matching PostgreSQL's rules. + pub(crate) fn take_single_byte(&mut self, key: &str) -> PgWireResult> { + let Some(value) = self.take_string(key)? else { + return Ok(None); + }; + let mut chars = value.chars(); + let (c, rest) = (chars.next(), chars.as_str()); + match c { + Some('\t') if rest.is_empty() => Ok(Some(b'\t')), + Some(c) if rest.is_empty() && c.is_ascii() && !c.is_ascii_control() => { + Ok(Some(c as u8)) + } + _ => Err(bad_copy_data(format!( + "COPY option {} must be a single one-byte character other than a newline or carriage return", + key.to_uppercase() + ))), + } + } + + /// Takes a boolean option; a valueless option means `TRUE`. + pub(crate) fn take_flag(&mut self, key: &str) -> PgWireResult> { + let Some(value) = self.take_value(key) else { + return Ok(None); + }; + match value { + CopyOptionValue::Flag => Ok(Some(true)), + CopyOptionValue::Boolean(b) => Ok(Some(b)), + other => Err(bad_copy_data(format!( + "COPY option {} requires a boolean value, got {}", + key.to_uppercase(), + other.kind() + ))), + } + } + + /// Rejects every option not consumed by the format, with one message + /// naming them all. Codecs call this after their `take_*` calls so a + /// typo cannot be silently ignored. + pub(crate) fn reject_leftovers(self, format: &str) -> PgWireResult<()> { + if self.entries.is_empty() { + return Ok(()); + } + let unknown = self + .entries + .iter() + .map(|(key, _)| key.to_uppercase()) + .collect::>() + .join(", "); + Err(PgWireError::UserError(Box::new( + PgErrorCode::Ec0A000.to_err_info(format!( + "COPY option(s) {unknown} are not supported by format {format}" + )), + ))) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn new_set(entries: &[(&str, CopyOptionValue)]) -> CopyOptionSet { + let mut set = CopyOptionSet::default(); + for (key, value) in entries { + set.insert((*key).to_string(), value.clone()).unwrap(); + } + set + } + + #[test] + fn test_insert_rejects_duplicates() { + let mut set = CopyOptionSet::default(); + set.insert("format".to_string(), CopyOptionValue::Ident("csv".into())) + .unwrap(); + let err = set + .insert("format".to_string(), CopyOptionValue::Ident("text".into())) + .unwrap_err(); + assert!(err.to_string().contains("redundant FORMAT")); + } + + #[test] + fn test_take_single_byte_tab_allowed() { + // A literal tab is a legal delimiter (PostgreSQL permits it). + let mut set = new_set(&[("delimiter", CopyOptionValue::String("\t".into()))]); + assert_eq!(set.take_single_byte("DELIMITER").unwrap(), Some(b'\t')); + + // Newlines and other control characters stay rejected. + for bad in ["\n", "\r", "\0"] { + let mut set = new_set(&[("delimiter", CopyOptionValue::String(bad.into()))]); + assert!( + set.take_single_byte("DELIMITER").is_err(), + "{bad:?} must be rejected" + ); + } + } + + #[test] + fn test_take_string_and_single_byte() { + let mut set = new_set(&[ + ("null", CopyOptionValue::String("nil".into())), + ("delimiter", CopyOptionValue::String("|".into())), + ]); + assert_eq!(set.take_string("null").unwrap().as_deref(), Some("nil")); + assert_eq!(set.take_single_byte("delimiter").unwrap(), Some(b'|')); + + let mut set = new_set(&[("delimiter", CopyOptionValue::Number("42".into()))]); + let err = set.take_single_byte("delimiter").unwrap_err(); + assert!(err.to_string().contains("requires a string value")); + + let mut set = new_set(&[("delimiter", CopyOptionValue::String("ab".into()))]); + let err = set.take_single_byte("delimiter").unwrap_err(); + assert!(err.to_string().contains("single one-byte")); + } + + #[test] + fn test_take_flag() { + let mut set = new_set(&[ + ("header", CopyOptionValue::Flag), + ("frozen", CopyOptionValue::Boolean(false)), + ]); + assert_eq!(set.take_flag("header").unwrap(), Some(true)); + assert_eq!(set.take_flag("frozen").unwrap(), Some(false)); + assert_eq!(set.take_flag("absent").unwrap(), None); + } + + #[test] + fn test_take_format_normalizes() { + let mut set = new_set(&[("format", CopyOptionValue::Ident("CSV".into()))]); + assert_eq!(set.take_format().unwrap().as_deref(), Some("csv")); + } + + #[test] + fn test_reject_leftovers_lists_all() { + let mut set = new_set(&[ + ("precision", CopyOptionValue::String("ns".into())), + ("weird", CopyOptionValue::Flag), + ]); + set.take_string("precision").unwrap(); + let err = set.reject_leftovers("text").unwrap_err(); + assert!( + err.to_string().contains("WEIRD"), + "leftover option must be listed: {err}" + ); + assert!( + err.to_string().contains("not supported by format text"), + "{err}" + ); + // The consumed option is not listed. + assert!(!err.to_string().contains("PRECISION")); + } + + #[test] + fn test_raw_value_escape_hatch() { + let mut set = new_set(&[("precision", CopyOptionValue::Ident("ns".into()))]); + assert_eq!( + set.take_value("precision"), + Some(CopyOptionValue::Ident("ns".into())) + ); + } +} diff --git a/src/servers/src/postgres/copy_in/parse.rs b/src/servers/src/postgres/copy_in/parse.rs new file mode 100644 index 00000000000..8139960bfa4 --- /dev/null +++ b/src/servers/src/postgres/copy_in/parse.rs @@ -0,0 +1,443 @@ +// 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. + +//! Textual parser for `COPY ... FROM STDIN` statements. +//! +//! COPY FROM STDIN is parsed here instead of through sqlparser because the +//! `WITH (...)` option list is open: formats may accept options beyond +//! sqlparser's COPY grammar (for example a line-protocol format taking +//! `precision`), and unknown options must reach the format codec rather +//! than fail statement parsing. +//! +//! The parser commits to error reporting only after `FROM STDIN` has been +//! seen; anything else (like `COPY ... TO STDOUT`) is reported as +//! [`CopyInScan::NotCopyIn`] so the caller falls back to the regular SQL +//! path. + +use datafusion::sql::sqlparser::dialect::GenericDialect; +use datafusion::sql::sqlparser::tokenizer::{Token, Tokenizer}; +use pgwire::error::PgWireResult; + +use crate::postgres::copy_in::options::{CopyOptionSet, CopyOptionValue}; +use crate::postgres::copy_in::{bad_copy_data, copy_in_must_run_alone_error}; + +/// A scanned `COPY tbl [(cols)] FROM STDIN [WITH (...)]` statement. +/// +/// Unquoted identifiers are folded to lowercase and quoted ones keep their +/// case (`COPY Metrics` addresses the same table as `CREATE TABLE Metrics`, +/// namely `metrics`), matching `ParserContext::canonicalize_identifier` in +/// the SQL parser and PostgreSQL's own folding rule. +#[derive(Debug)] +pub(super) struct RawCopyFromStdin { + /// Table name parts, as written (`t`, `schema.t`, + /// `catalog.schema.t`); quoted names keep their case. + pub(super) table: Vec, + /// Requested column list; empty means all columns in table order. + pub(super) columns: Vec, + /// The open option list, handed to the resolved format codec. + pub(super) options: CopyOptionSet, +} + +pub(super) enum CopyInScan { + /// Not a `COPY ... FROM STDIN` statement. + NotCopyIn, + /// A complete copy-in statement with only trailing separators left. + CopyIn(RawCopyFromStdin), +} + +/// Scans a SQL string for a leading `COPY ... FROM STDIN` statement. +pub(super) fn scan_copy_from_stdin(sql: &str) -> PgWireResult { + // Fast path: the statement must start with the COPY keyword, possibly + // behind comments; only then is a full tokenization worth its cost. + let trimmed = sql.trim_start(); + let starts_like_copy = trimmed + .as_bytes() + .get(..4) + .is_some_and(|prefix| prefix.eq_ignore_ascii_case(b"copy")); + let starts_with_comment = trimmed.starts_with("--") || trimmed.starts_with("/*"); + if !starts_like_copy && !starts_with_comment { + return Ok(CopyInScan::NotCopyIn); + } + + let tokens = match Tokenizer::new(&GenericDialect {}, sql).tokenize() { + Ok(tokens) => tokens, + // Let the regular SQL path produce the parse error. + Err(_) => return Ok(CopyInScan::NotCopyIn), + }; + // The tokenizer keeps whitespace; the scanner works on tokens only. + let tokens: Vec = tokens + .into_iter() + .filter(|token| !matches!(token, Token::Whitespace(_))) + .collect(); + let mut cursor = Tokens { + tokens: &tokens, + pos: 0, + }; + + if !cursor.take_keyword("copy") { + return Ok(CopyInScan::NotCopyIn); + } + + let table = match parse_object_name(&mut cursor) { + Some(table) => table, + None => return Ok(CopyInScan::NotCopyIn), + }; + let columns = match parse_column_list(&mut cursor) { + Some(columns) => columns, + None => return Ok(CopyInScan::NotCopyIn), + }; + + if !cursor.take_keyword("from") || !cursor.take_keyword("stdin") { + // A COPY with another direction or source; not ours to handle. + return Ok(CopyInScan::NotCopyIn); + } + + let mut options = CopyOptionSet::default(); + // PostgreSQL allows the WITH keyword to be omitted. + let has_options = if cursor.take_keyword("with") { + if !matches!(cursor.peek(), Some(Token::LParen)) { + return Err(bad_copy_data("expected '(' after WITH in COPY FROM STDIN")); + } + true + } else { + matches!(cursor.peek(), Some(Token::LParen)) + }; + if has_options { + // Consume the opening '('. + cursor.next(); + parse_option_list(&mut cursor, &mut options)?; + } + + // Only statement separators may follow. + while let Some(token) = cursor.next() { + if !matches!(token, Token::SemiColon) { + return Err(copy_in_must_run_alone_error()); + } + } + + Ok(CopyInScan::CopyIn(RawCopyFromStdin { + table, + columns, + options, + })) +} + +struct Tokens<'a> { + tokens: &'a [Token], + pos: usize, +} + +impl<'a> Tokens<'a> { + fn next(&mut self) -> Option<&'a Token> { + let token = self.tokens.get(self.pos)?; + self.pos += 1; + Some(token) + } + + fn peek(&mut self) -> Option<&'a Token> { + self.tokens.get(self.pos) + } + + /// Consumes the next token if it is the unquoted keyword `keyword`. + fn take_keyword(&mut self, keyword: &str) -> bool { + match self.peek() { + Some(Token::Word(w)) + if w.quote_style.is_none() && w.value.eq_ignore_ascii_case(keyword) => + { + self.pos += 1; + true + } + _ => false, + } + } +} + +/// Returns the identifier with the SQL parser's canonicalization: quoted +/// identifiers keep their case, unquoted ones fold to lowercase. +fn expect_ident(cursor: &mut Tokens<'_>) -> Option { + match cursor.next() { + Some(Token::Word(w)) => Some(if w.quote_style.is_none() { + w.value.to_lowercase() + } else { + w.value.clone() + }), + _ => None, + } +} + +fn parse_object_name(cursor: &mut Tokens<'_>) -> Option> { + let mut parts = vec![expect_ident(cursor)?]; + while matches!(cursor.peek(), Some(Token::Period)) { + cursor.next(); + parts.push(expect_ident(cursor)?); + } + Some(parts) +} + +/// Parses an optional `(column, ...)` list. `None` means the statement is +/// not a copy-in shape (either no list, or a malformed list better left to +/// the regular parser to report). +fn parse_column_list(cursor: &mut Tokens<'_>) -> Option> { + if !matches!(cursor.peek(), Some(Token::LParen)) { + return Some(Vec::new()); + } + cursor.next(); + let mut columns = vec![expect_ident(cursor)?]; + loop { + match cursor.next() { + Some(Token::Comma) => columns.push(expect_ident(cursor)?), + Some(Token::RParen) => break, + _ => return None, + } + } + Some(columns) +} + +/// Parses `NAME [value] (, NAME [value])*` until the closing `)`. +/// +/// Option values are open-ended: a string literal, a number, an identifier +/// or a boolean word. Whether an option makes sense is decided by the +/// format codec, not here. +fn parse_option_list(cursor: &mut Tokens<'_>, options: &mut CopyOptionSet) -> PgWireResult<()> { + loop { + let key = match cursor.next() { + Some(Token::Word(w)) if w.quote_style.is_none() => w.value.to_lowercase(), + _ => return Err(bad_copy_data("expected an option name after '(' or ','")), + }; + let value = match cursor.peek() { + Some(Token::Comma) | Some(Token::RParen) => CopyOptionValue::Flag, + Some(Token::SingleQuotedString(s)) => { + let value = CopyOptionValue::String(s.clone()); + cursor.next(); + value + } + Some(Token::EscapedStringLiteral(s)) => { + // E'...' escape-string constants, e.g. DELIMITER E'\t': + // sqlparser's tokenizer already decoded the escapes, so + // the payload is used as-is, like SingleQuotedString. + let value = CopyOptionValue::String(s.clone()); + cursor.next(); + value + } + Some(Token::Number(n, _)) => { + let value = CopyOptionValue::Number(n.clone()); + cursor.next(); + value + } + Some(Token::Word(w)) => { + let value = if w.quote_style.is_none() { + match w.value.to_ascii_lowercase().as_str() { + "true" => CopyOptionValue::Boolean(true), + "false" => CopyOptionValue::Boolean(false), + _ => CopyOptionValue::Ident(w.value.clone()), + } + } else { + CopyOptionValue::Ident(w.value.clone()) + }; + cursor.next(); + value + } + _ => return Err(bad_copy_data("expected an option value or ',' or ')'")), + }; + options.insert(key, value)?; + match cursor.next() { + Some(Token::Comma) => continue, + Some(Token::RParen) => break, + _ => { + return Err(bad_copy_data( + "expected ',' or ')' in COPY FROM STDIN options", + )); + } + } + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn scan(sql: &str) -> PgWireResult> { + Ok(match scan_copy_from_stdin(sql)? { + CopyInScan::CopyIn(raw) => Some(raw), + CopyInScan::NotCopyIn => None, + }) + } + + #[test] + fn test_scan_basic() { + let raw = scan("COPY t FROM STDIN").unwrap().unwrap(); + assert_eq!(raw.table, vec!["t"]); + assert!(raw.columns.is_empty()); + } + + #[test] + fn test_scan_qualified_table_and_columns() { + let raw = scan("COPY db.public.Metrics (ts, \"Host\") FROM STDIN") + .unwrap() + .unwrap(); + assert_eq!(raw.table, vec!["db", "public", "metrics"]); + assert_eq!(raw.columns, vec!["ts", "Host"]); + } + + #[test] + fn test_scan_not_copy_in() { + assert!(scan("COPY (SELECT 1) TO STDOUT").unwrap().is_none()); + assert!( + scan("COPY t TO '/tmp/x.csv' WITH (FORMAT csv)") + .unwrap() + .is_none() + ); + assert!(scan("SELECT 1").unwrap().is_none()); + // GreptimeDB's own file-based COPY FROM. + assert!(scan("COPY t FROM '/tmp/x.csv'").unwrap().is_none()); + } + + #[test] + fn test_scan_leading_comment() { + let raw = scan("-- a comment\nCOPY t FROM STDIN").unwrap().unwrap(); + assert_eq!(raw.table, vec!["t"]); + } + + #[test] + fn test_scan_folds_unquoted_identifiers() { + // Unquoted identifiers fold to lowercase, like the SQL parser. + let raw = scan("copy Metrics (Host, ts) from stdin").unwrap().unwrap(); + assert_eq!(raw.table, vec!["metrics"]); + assert_eq!(raw.columns, vec!["host", "ts"]); + } + + #[test] + fn test_scan_preserves_quoted_identifiers() { + let raw = scan("COPY \"Metrics\" (\"Host\") FROM STDIN") + .unwrap() + .unwrap(); + assert_eq!(raw.table, vec!["Metrics"]); + assert_eq!(raw.columns, vec!["Host"]); + } + + #[test] + fn test_scan_options_kinds() { + let raw = scan( + "COPY t FROM STDIN WITH (FORMAT csv, DELIMITER ';', HEADER, \ + skip_bad_records true, precision 3, mode strict)", + ) + .unwrap() + .unwrap(); + let mut options = raw.options; + assert_eq!(options.take_format().unwrap().as_deref(), Some("csv")); + assert_eq!(options.take_single_byte("DELIMITER").unwrap(), Some(b';')); + assert_eq!(options.take_flag("HEADER").unwrap(), Some(true)); + assert_eq!( + options.take_value("skip_bad_records"), + Some(CopyOptionValue::Boolean(true)) + ); + assert_eq!( + options.take_value("precision"), + Some(CopyOptionValue::Number("3".into())) + ); + assert_eq!( + options.take_value("mode"), + Some(CopyOptionValue::Ident("strict".into())) + ); + } + + #[test] + fn test_scan_escaped_string_constants() { + // sqlparser's tokenizer already decodes E'...' escapes, so the + // payload is used as-is: E'\t' is a tab. + let mut raw = scan("COPY t FROM STDIN WITH (FORMAT csv, DELIMITER E'\\t')") + .unwrap() + .unwrap(); + assert_eq!( + raw.options.take_single_byte("DELIMITER").unwrap(), + Some(b'\t') + ); + + // E'\\' is a single literal backslash (a legal one-byte + // delimiter), NOT an escaped backslash-t. + let mut raw = scan("COPY t FROM STDIN WITH (FORMAT csv, DELIMITER E'\\\\')") + .unwrap() + .unwrap(); + assert_eq!( + raw.options.take_single_byte("DELIMITER").unwrap(), + Some(b'\\') + ); + + // E'\\t' decodes to TWO characters — backslash followed by t — + // and must be rejected as a multi-character delimiter instead of + // being unescaped a second time into a tab. + // The scanner leaves option validation to the codec; use the + // full resolution path to assert the delimiter rejection. + let err = match crate::postgres::copy_in::copy_from_stdin_from_sql( + "COPY t FROM STDIN WITH (FORMAT csv, DELIMITER E'\\\\t')", + ) { + Ok(Some(_)) => panic!("expected the multi-character delimiter to be rejected"), + Ok(None) => panic!("expected the statement to scan as copy-in"), + Err(e) => e, + }; + assert!( + err.to_string().contains("single one-byte character"), + "{err}" + ); + } + + #[test] + fn test_scan_semicolon_inside_option_value() { + // A quoted semicolon is not a statement separator. + let mut raw = scan("COPY t FROM STDIN WITH (NULL ';')").unwrap().unwrap(); + assert_eq!( + raw.options.take_string("NULL").unwrap().as_deref(), + Some(";") + ); + } + + #[test] + fn test_scan_trailing_statements_rejected() { + let err = scan("COPY t FROM STDIN; SELECT 1").unwrap_err(); + assert!(err.to_string().contains("alone")); + + // Trailing separators alone are fine. + assert!(scan("COPY t FROM STDIN;").unwrap().is_some()); + assert!(scan("COPY t FROM STDIN ; ;").unwrap().is_some()); + } + + #[test] + fn test_scan_options_without_with_keyword() { + let raw = scan("COPY t FROM STDIN (FORMAT csv, HEADER)") + .unwrap() + .unwrap(); + let mut options = raw.options; + assert_eq!(options.take_format().unwrap().as_deref(), Some("csv")); + assert_eq!(options.take_flag("HEADER").unwrap(), Some(true)); + assert!(options.reject_leftovers("csv").is_ok()); + } + + #[test] + fn test_scan_malformed_after_stdin() { + let err = scan("COPY t FROM STDIN WITH").unwrap_err(); + assert!(err.to_string().contains("'('"), "{err}"); + let err = scan("COPY t FROM STDIN WITH (").unwrap_err(); + assert!(err.to_string().contains("option"), "{err}"); + // A malformed column list never commits to the copy-in shape and + // falls through to the regular parser. + assert!(scan("COPY t (a,").unwrap().is_none()); + } + + #[test] + fn test_scan_redundant_option() { + let err = scan("COPY t FROM STDIN WITH (FORMAT csv, FORMAT text)").unwrap_err(); + assert!(err.to_string().contains("redundant FORMAT")); + } +} diff --git a/src/servers/src/postgres/handler.rs b/src/servers/src/postgres/handler.rs index 9d169811984..d03747f1b50 100644 --- a/src/servers/src/postgres/handler.rs +++ b/src/servers/src/postgres/handler.rs @@ -56,11 +56,14 @@ use crate::SqlPlan; use crate::error::{DataFusionSnafu, InferParameterTypesSnafu, Result}; use crate::postgres::types::*; use crate::postgres::utils::convert_err; -use crate::postgres::{PostgresServerHandlerInner, fixtures}; +use crate::postgres::{PostgresServerHandlerInner, copy_in, fixtures}; use crate::query_handler::sql::ServerSqlQueryHandlerRef; impl PostgresServerHandlerInner { - fn new_query_context(&self) -> QueryContextRef { + /// Creates a query context carrying this connection's batching + /// selection, so write paths can route through the pending-rows + /// batcher when it is enabled. + pub(super) fn new_query_context(&self) -> QueryContextRef { let mut ctx = self.session.new_query_context(); Arc::make_mut(&mut ctx).set_batching_enabled(self.batching_enabled); ctx @@ -87,8 +90,25 @@ impl SimpleQueryHandler for PostgresServerHandlerInner { return Ok(vec![Response::EmptyQuery]); } + // A leading COPY FROM STDIN takes over the connection for the copy + // sub-protocol; it is parsed textually so the WITH clause may carry + // options beyond sqlparser's COPY grammar. It must run alone. + if let Some(copy_from_stdin) = copy_in::copy_from_stdin_from_sql(query)? { + let response = self.begin_copy_in(copy_from_stdin).await?; + send_warning_opt(client, query_ctx).await?; + return Ok(vec![response]); + } + let parsed_query = self.query_parser.compatibility_parser.parse(query); + // A COPY FROM STDIN trailing other statements must be rejected; + // the textual scanner above only sees statement-leading COPYs. + if let Ok(statements) = &parsed_query + && statements.iter().any(copy_in::is_copy_from_stdin) + { + return Err(copy_in::copy_in_must_run_alone_error()); + } + let query = if let Ok(statements) = &parsed_query { // Comments, whitespace and empty statements also require EmptyQueryResponse. if statements.is_empty() { @@ -282,6 +302,36 @@ where _ => 0, }; + // Render timestamps with an explicit session-timezone offset so the + // exported text carries the instant unambiguously: the timezone-aware + // COPY IN decoder honors embedded offsets, keeping COPY OUT -> COPY IN + // round trips exact in non-UTC sessions (matching SQL INSERT parity). + // This touches both halves of the encoder: the arrow schema drives + // value conversion into a tz-aware DateTime, and the pg field type + // selects the offset-bearing text format. + let arrow_schema_override = copy_out_schema_with_timezone(&schema, &query_ctx); + let pg_fields = if arrow_schema_override.is_some() { + pg_fields + .into_iter() + .map(|mut field| { + if field.datatype() == &Type::TIMESTAMP { + field = FieldInfo::new( + field.name().to_string(), + None, + None, + Type::TIMESTAMPTZ, + field.format(), + ) + .with_type_size(field.type_size()) + .with_type_modifier(field.type_modifier()); + } + field + }) + .collect::>() + } else { + pg_fields + }; + let pg_schema = Arc::new(pg_fields); let num_columns = pg_schema.len(); @@ -297,7 +347,8 @@ where schema.clone(), recordbatches_stream, copy_encoder, - ); + ) + .with_arrow_schema_override(arrow_schema_override); let copy_stream: RowStream = Box::pin( row_stream @@ -336,6 +387,7 @@ impl DefaultQueryParser { pub struct PgSqlPlan { pub(crate) plan: SqlPlan, pub(crate) copy_to_stdout_format: Option, + pub(crate) copy_from_stdin: Option, } #[async_trait] @@ -360,6 +412,19 @@ impl QueryParser for DefaultQueryParser { return Ok(Some(PgSqlPlan { plan: SqlPlan::Shortcut(sql.to_string()), copy_to_stdout_format: None, + copy_from_stdin: None, + })); + } + + // A COPY FROM STDIN is parsed textually — not through sqlparser — + // so the WITH clause may carry options beyond sqlparser's COPY + // grammar; the extended query handler starts the copy sub-protocol + // when the statement is executed. + if let Some(copy_from_stdin) = copy_in::copy_from_stdin_from_sql(sql)? { + return Ok(Some(PgSqlPlan { + plan: SqlPlan::Shortcut(sql.to_string()), + copy_to_stdout_format: None, + copy_from_stdin: Some(copy_from_stdin), })); } @@ -370,7 +435,7 @@ impl QueryParser for DefaultQueryParser { } let first_stmt = statements.remove(0); let format = check_copy_to_stdout(&first_stmt); - (statement_to_query(&first_stmt), format) + (first_stmt.to_string(), format) } else { // bypass the error: it can run into error because of different // versions of sqlparser @@ -400,11 +465,13 @@ impl QueryParser for DefaultQueryParser { Ok(Some(PgSqlPlan { plan: SqlPlan::Plan(logical_plan, stmt), copy_to_stdout_format, + copy_from_stdin: None, })) } else { Ok(Some(PgSqlPlan { plan: SqlPlan::Statement(stmt, sql), copy_to_stdout_format, + copy_from_stdin: None, })) } } @@ -460,6 +527,14 @@ impl ExtendedQueryHandler for PostgresServerHandlerInner { let pg_sql_plan = &portal.statement.statement; let sql_plan = &pg_sql_plan.plan; + // A COPY FROM STDIN takes over the connection for the copy + // sub-protocol. + if let Some(copy_from_stdin) = &pg_sql_plan.copy_from_stdin { + let response = self.begin_copy_in(copy_from_stdin.clone()).await?; + send_warning_opt(client, query_ctx).await?; + return Ok(response); + } + let output = match sql_plan { SqlPlan::Empty => { // early return if query is empty @@ -782,6 +857,38 @@ impl ErrorHandler for PostgresServerHandlerInner { } } +/// Rewrites timestamp fields of the table's arrow schema to carry the +/// session timezone, so COPY OUT renders them with an explicit offset +/// (`2024-01-01 08:00:00+08:00`). Returns `None` when there is nothing +/// to rewrite (no naive timestamps or a UTC session leaving output as-is +/// is fine either way — offsets are always unambiguous on re-import). +fn copy_out_schema_with_timezone( + schema: &SchemaRef, + query_ctx: &QueryContextRef, +) -> Option> { + let timezone = query_ctx.timezone().to_string(); + let arrow_schema = schema.arrow_schema(); + let mut changed = false; + let fields = arrow_schema + .fields() + .iter() + .map(|field| match field.data_type() { + arrow::datatypes::DataType::Timestamp(unit, None) => { + changed = true; + field + .as_ref() + .clone() + .with_data_type(arrow::datatypes::DataType::Timestamp( + *unit, + Some(Arc::from(timezone.as_str())), + )) + } + _ => field.as_ref().clone(), + }) + .collect::>(); + changed.then(|| Arc::new(arrow::datatypes::Schema::new(fields))) +} + fn check_copy_to_stdout(statement: &SqlParserStatement) -> Option { if let SqlParserStatement::Copy { target, options, .. diff --git a/src/servers/src/postgres/server.rs b/src/servers/src/postgres/server.rs index 6e4aa942f9d..7e3b3e2b7bd 100644 --- a/src/servers/src/postgres/server.rs +++ b/src/servers/src/postgres/server.rs @@ -29,6 +29,7 @@ use tokio_rustls::TlsAcceptor; use crate::error::Result; use crate::postgres::{MakePostgresServerHandler, MakePostgresServerHandlerBuilder}; +use crate::query_handler::CopyInHandlerRef; use crate::query_handler::sql::ServerSqlQueryHandlerRef; use crate::server::{AbortableStream, BaseTcpServer, Server}; use crate::tls::ReloadableTlsServerConfig; @@ -44,9 +45,16 @@ pub struct PostgresServer { } impl PostgresServer { - /// Creates a new Postgres server with provided query_handler and async runtime + /// Creates a new Postgres server with provided query_handler and async runtime. + /// + /// `query_handler` and `copy_in_handler` are expected to be handles to + /// the same frontend instance (the SQL and COPY FROM STDIN facets of + /// one object); they are separate parameters only because they are + /// separate traits. + #[allow(clippy::too_many_arguments)] pub fn new( query_handler: ServerSqlQueryHandlerRef, + copy_in_handler: CopyInHandlerRef, force_tls: bool, tls_server_config: Arc, keep_alive_secs: u64, @@ -57,6 +65,7 @@ impl PostgresServer { let make_handler = Arc::new( MakePostgresServerHandlerBuilder::default() .query_handler(query_handler.clone()) + .copy_in_handler(copy_in_handler) .user_provider(user_provider.clone()) .force_tls(force_tls) .build() diff --git a/src/servers/src/postgres/types.rs b/src/servers/src/postgres/types.rs index 25cb94662c0..b0621d0a7b9 100644 --- a/src/servers/src/postgres/types.rs +++ b/src/servers/src/postgres/types.rs @@ -285,6 +285,11 @@ where schema: SchemaRef, record_batches: Pin>, encoder: S, + /// When set, replaces each batch's arrow schema for value encoding. + /// The COPY OUT path uses this to render timestamps with an explicit + /// session-timezone offset, so exported text re-imports identically + /// by the timezone-aware COPY IN decoder. + arrow_schema_override: Option>, } impl Stream for RecordBatchRowStream @@ -304,7 +309,10 @@ where return Poll::Ready(Some(Ok(vec![]))); } - let arrow_schema = record_batch.schema(); + let arrow_schema = self + .arrow_schema_override + .clone() + .unwrap_or_else(|| record_batch.schema()); let query_ctx = self.query_ctx.clone(); let pg_schema = self.pg_schema.clone(); let schema = self.schema.clone(); @@ -352,9 +360,19 @@ where schema, record_batches: Box::pin(record_batches), encoder, + arrow_schema_override: None, } } + /// See [`RecordBatchRowStream::arrow_schema_override`]. + pub(crate) fn with_arrow_schema_override( + mut self, + schema: Option>, + ) -> Self { + self.arrow_schema_override = schema; + self + } + fn encode_row( query_ctx: &QueryContextRef, pg_schema: &Arc>, @@ -2464,6 +2482,7 @@ mod test { PgSqlPlan { plan: SqlPlan::Empty, copy_to_stdout_format: None, + copy_from_stdin: None, }, client_param_types, )); diff --git a/src/servers/src/postgres/types/error.rs b/src/servers/src/postgres/types/error.rs index 00d1bbbb1c1..86317b91570 100644 --- a/src/servers/src/postgres/types/error.rs +++ b/src/servers/src/postgres/types/error.rs @@ -67,6 +67,9 @@ pub enum PgErrorCode { /// connection_exception #[snafu(display("connection_exception"))] Ec08000 = 400, + /// protocol_violation + #[snafu(display("protocol_violation"))] + Ec08P01 = 401, // === End of Class 08 — Connection Exception ===== // === Begin of Class 09 — Triggered Action Exception === @@ -130,6 +133,12 @@ pub enum PgErrorCode { /// invalid_parameter_value #[snafu(display("invalid_parameter_value"))] Ec22023 = 1401, + /// invalid_text_representation + #[snafu(display("invalid_text_representation"))] + Ec22P02 = 1402, + /// bad_copy_file_format + #[snafu(display("bad_copy_file_format"))] + Ec22P04 = 1403, // === End of Class 22 — Data Exception ===== // === Begin of Class 23 — Integrity Constraint Violation === @@ -295,6 +304,9 @@ pub enum PgErrorCode { /// operator_intervention #[snafu(display("operator_intervention"))] Ec57000 = 3600, + /// query_canceled + #[snafu(display("query_canceled"))] + Ec57014 = 3602, /// cannot_connect_now #[snafu(display("cannot_connect_now"))] diff --git a/src/servers/src/query_handler.rs b/src/servers/src/query_handler.rs index 12a56076d6f..57f55855ac1 100644 --- a/src/servers/src/query_handler.rs +++ b/src/servers/src/query_handler.rs @@ -43,6 +43,7 @@ use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetrics use pipeline::{GreptimePipelineParams, Pipeline, PipelineInfo, PipelineVersion, PipelineWay}; use serde_json::Value; use session::context::{QueryContext, QueryContextRef}; +use table::metadata::TableInfoRef; #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] pub struct DashboardDefinition { @@ -290,3 +291,33 @@ pub trait JaegerQueryHandler { query_params: QueryTraceParams, ) -> Result; } + +/// Handler for the PostgreSQL `COPY ... FROM STDIN` (copy-in) sub-protocol. +/// +/// The servers layer owns the wire protocol and data parsing; this trait is +/// the boundary to catalog lookup and ingestion, implemented by the +/// frontend instance. +#[async_trait] +pub trait CopyInHandler: Send + Sync { + /// Resolves the target table of a COPY and returns its metadata, or + /// `None` when the table does not exist. Implementations should also + /// enforce the write permission for the table here so unauthorized + /// copies fail before any data is transferred. + async fn copy_in_table( + &self, + catalog: &str, + schema: &str, + table: &str, + query_ctx: QueryContextRef, + ) -> Result>; + + /// Inserts a parsed batch produced by a COPY stream. Returns the + /// number of affected rows. + async fn copy_in_insert( + &self, + requests: RowInsertRequests, + query_ctx: QueryContextRef, + ) -> Result; +} + +pub type CopyInHandlerRef = Arc; diff --git a/src/servers/tests/mod.rs b/src/servers/tests/mod.rs index 699f3f76f1e..d78f57b7b83 100644 --- a/src/servers/tests/mod.rs +++ b/src/servers/tests/mod.rs @@ -14,13 +14,14 @@ use std::sync::Arc; +use api::v1::RowInsertRequests; use api::v1::greptime_request::Request; use api::v1::query_request::Query; use async_trait::async_trait; use catalog::memory::MemoryCatalogManager; use catalog::system_schema::SystemSchemaProvider; use catalog::system_schema::pg_catalog::PGCatalogProvider; -use catalog::{CatalogManager, RegisterTableRequest}; +use catalog::{CatalogManager, CatalogManagerRef, RegisterTableRequest}; use common_catalog::consts::{DEFAULT_CATALOG_NAME, PG_CATALOG_NAME}; use common_error::ext::BoxedError; use common_grpc::flight::do_put::DoPutResponse; @@ -32,6 +33,7 @@ use query::parser::{PromQuery, QueryStatement}; use query::query_engine::DescribeResult; use query::{QueryEngineFactory, QueryEngineRef}; use servers::error::{ExecuteQuerySnafu, NotSupportedSnafu, Result}; +use servers::query_handler::CopyInHandler; use servers::query_handler::grpc::GrpcQueryHandler; use servers::query_handler::sql::{ServerSqlQueryHandlerRef, SqlQueryHandler}; use session::context::QueryContextRef; @@ -39,6 +41,7 @@ use snafu::{ResultExt, ensure}; use sql::parser::{ParseOptions, ParserContext}; use sql::statements::statement::Statement; use table::TableRef; +use table::metadata::TableInfoRef; mod http; mod interceptor; @@ -54,6 +57,64 @@ impl DummyInstance { fn new(query_engine: QueryEngineRef) -> Self { Self { query_engine } } + + pub fn catalog_manager(&self) -> CatalogManagerRef { + self.query_engine.engine_state().catalog_manager().clone() + } +} + +/// A [`CopyInHandler`] test double: resolves tables through the catalog +/// manager and records the insert requests instead of writing them. +pub struct RecordingCopyInHandler { + catalog_manager: CatalogManagerRef, + pub requests: std::sync::Mutex>, +} + +impl RecordingCopyInHandler { + pub fn new(catalog_manager: CatalogManagerRef) -> Self { + Self { + catalog_manager, + requests: Default::default(), + } + } +} + +#[async_trait] +impl CopyInHandler for RecordingCopyInHandler { + async fn copy_in_table( + &self, + catalog: &str, + schema: &str, + table: &str, + _query_ctx: QueryContextRef, + ) -> Result> { + Ok(self + .catalog_manager + .table(catalog, schema, table, None) + .await + .unwrap() + .map(|table| table.table_info())) + } + + async fn copy_in_insert( + &self, + requests: RowInsertRequests, + _query_ctx: QueryContextRef, + ) -> Result { + let rows: usize = requests + .inserts + .iter() + .map(|insert| { + insert + .rows + .as_ref() + .map(|rows| rows.rows.len()) + .unwrap_or(0) + }) + .sum(); + self.requests.lock().unwrap().push(requests); + Ok(Output::new_with_affected_rows(rows)) + } } #[async_trait] diff --git a/src/servers/tests/postgres/mod.rs b/src/servers/tests/postgres/mod.rs index 23971145a86..1bc5e79c236 100644 --- a/src/servers/tests/postgres/mod.rs +++ b/src/servers/tests/postgres/mod.rs @@ -229,7 +229,18 @@ fn create_postgres_server_with_user_provider( tls: TlsOption, user_provider: Option, ) -> Result> { + create_postgres_server_inner(table, tls, user_provider).map(|(server, _)| server) +} + +fn create_postgres_server_inner( + table: TableRef, + tls: TlsOption, + user_provider: Option, +) -> Result<(Box, Arc)> { let instance = Arc::new(create_testing_instance(table)); + let copy_in_handler = Arc::new(crate::RecordingCopyInHandler::new( + instance.catalog_manager(), + )); let io_runtime = RuntimeBuilder::default() .worker_threads(4) .thread_name("postgres-io-handlers") @@ -241,15 +252,19 @@ fn create_postgres_server_with_user_provider( .expect("Failed to load certificates and keys"), ); - Ok(Box::new(PostgresServer::new( - instance, - tls.should_force_tls(), - tls_server_config, - 0, - io_runtime, - user_provider, - None, - ))) + Ok(( + Box::new(PostgresServer::new( + instance, + copy_in_handler.clone(), + tls.should_force_tls(), + tls_server_config, + 0, + io_runtime, + user_provider, + None, + )), + copy_in_handler, + )) } async fn start_test_server_with_user_provider( @@ -1149,3 +1164,237 @@ async fn do_simple_query_with_secure_server( do_simple_query(server_tls, client_tls).await } + +// --------------------------------------------------------------------------- +// COPY FROM STDIN tests +// --------------------------------------------------------------------------- + +use futures::SinkExt; + +fn copy_in_test_table() -> TableRef { + use common_recordbatch::RecordBatch; + use datatypes::prelude::ConcreteDataType; + use datatypes::schema::ColumnSchema; + use datatypes::types::TimestampType; + use datatypes::vectors::{Float64Vector, StringVector, TimestampMillisecondVector, VectorRef}; + + let column_schemas = vec![ + ColumnSchema::new( + "ts", + ConcreteDataType::Timestamp(TimestampType::Millisecond(Default::default())), + false, + ) + .with_time_index(true), + ColumnSchema::new("host", ConcreteDataType::string_datatype(), false), + ColumnSchema::new("val", ConcreteDataType::float64_datatype(), true), + ColumnSchema::new("note", ConcreteDataType::string_datatype(), true), + ]; + let schema = Arc::new(datatypes::schema::Schema::new(column_schemas)); + let columns: Vec = vec![ + Arc::new(TimestampMillisecondVector::from_slice([0i64])), + Arc::new(StringVector::from_slice(&["localhost"])), + Arc::new(Float64Vector::from_slice([1.0f64])), + Arc::new(StringVector::from_slice(&["seed"])), + ]; + let recordbatch = RecordBatch::new(schema, columns).unwrap(); + MemTable::table("metrics", recordbatch) +} + +async fn start_copy_in_test_server( + table: TableRef, +) -> Result<(Box, u16, Arc)> { + common_telemetry::init_default_ut_logging(); + let _ = install_default_crypto_provider(); + let (mut server, handler) = create_postgres_server_inner(table, Default::default(), None)?; + let listening = "127.0.0.1:0".parse::().unwrap(); + server.start(listening).await.unwrap(); + let port = server.bind_addr().unwrap().port(); + Ok((server, port, handler)) +} + +fn recorded_rows(handler: &crate::RecordingCopyInHandler) -> Vec { + handler + .requests + .lock() + .unwrap() + .iter() + .flat_map(|request| { + request.inserts.iter().flat_map(|insert| { + insert + .rows + .clone() + .map(|rows| rows.rows) + .unwrap_or_default() + }) + }) + .collect() +} + +async fn feed_copy_in( + client: &Client, + sql: &str, + chunks: Vec, +) -> std::result::Result { + let mut sink = Box::pin(client.copy_in(sql).await?); + for chunk in chunks { + sink.as_mut().send(chunk).await?; + } + sink.as_mut().finish().await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_copy_in_csv() -> Result<()> { + let (server, port, handler) = start_copy_in_test_server(copy_in_test_table()).await?; + let client = create_plain_connection(port, false).await.unwrap(); + + let rows = feed_copy_in( + &client, + "COPY metrics FROM STDIN WITH (FORMAT csv)", + vec![bytes::Bytes::from_static( + b"2023-11-14 22:13:20,host1,1.5,hello\n2023-11-14 22:13:21,host2,,\"quoted, comma\"\n", + )], + ) + .await + .unwrap(); + assert_eq!(rows, 2); + + let rows = recorded_rows(&handler); + assert_eq!(rows.len(), 2); + assert_eq!( + rows[0].values[1].value_data, + Some(api::v1::value::ValueData::StringValue("host1".to_string())) + ); + assert_eq!( + rows[0].values[2].value_data, + Some(api::v1::value::ValueData::F64Value(1.5)) + ); + // Unquoted empty field is NULL, quoted field keeps its content. + assert_eq!(rows[1].values[2].value_data, None); + assert_eq!( + rows[1].values[3].value_data, + Some(api::v1::value::ValueData::StringValue( + "quoted, comma".to_string() + )) + ); + + server.shutdown().await.unwrap(); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_copy_in_column_list_and_batches() -> Result<()> { + let (server, port, handler) = start_copy_in_test_server(copy_in_test_table()).await?; + let client = create_plain_connection(port, false).await.unwrap(); + + let mut chunks = Vec::new(); + let mut payload = String::new(); + for i in 0..70000u64 { + payload.push_str(&format!("{i},2023-11-14 22:13:20,host{i},0.5\n")); + if (i + 1) % 100 == 0 { + chunks.push(bytes::Bytes::from(std::mem::take(&mut payload))); + } + } + let rows = feed_copy_in( + &client, + "COPY metrics (note, ts, host, val) FROM STDIN (FORMAT csv)", + chunks, + ) + .await + .unwrap(); + assert_eq!(rows, 70000); + + let recorded = handler.requests.lock().unwrap().clone(); + // The 70000 rows must have been flushed in multiple batches of at + // most COPY_IN_FLUSH_ROWS. + assert!(recorded.len() > 1); + let total: usize = recorded + .iter() + .map(|request| { + request.inserts[0] + .rows + .as_ref() + .map(|rows| rows.rows.len()) + .unwrap_or(0) + }) + .sum(); + assert_eq!(total, 70000); + + let rows = recorded_rows(&handler); + // The column list reorders the fields: first field lands in `note`. + assert_eq!( + rows[0].values[0].value_data, + Some(api::v1::value::ValueData::StringValue("0".to_string())) + ); + + server.shutdown().await.unwrap(); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_copy_in_unknown_table() -> Result<()> { + let (server, port, _handler) = start_copy_in_test_server(copy_in_test_table()).await?; + let client = create_plain_connection(port, false).await.unwrap(); + + let err = match client + .copy_in::<_, bytes::Bytes>("COPY not_exist FROM STDIN WITH (FORMAT csv)") + .await + { + Ok(_) => panic!("expected COPY to unknown table to fail"), + Err(e) => e, + }; + let message = err.as_db_error().expect("expected db error").message(); + assert!(message.contains("does not exist"), "{message}"); + + server.shutdown().await.unwrap(); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_copy_in_bad_value() -> Result<()> { + let (server, port, handler) = start_copy_in_test_server(copy_in_test_table()).await?; + let client = create_plain_connection(port, false).await.unwrap(); + + let err = match feed_copy_in( + &client, + "COPY metrics FROM STDIN WITH (FORMAT csv)", + vec![bytes::Bytes::from_static(b"not_a_timestamp,host1,1.5,x\n")], + ) + .await + { + Ok(_) => panic!("expected COPY to fail"), + Err(e) => e, + }; + let message = err.as_db_error().expect("expected db error").message(); + assert!(message.contains("failed to parse"), "{message}"); + // The failed batch must not be written. + assert!(handler.requests.lock().unwrap().is_empty()); + + server.shutdown().await.unwrap(); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_copy_in_multi_statement_rejected() -> Result<()> { + let (server, port, _handler) = start_copy_in_test_server(copy_in_test_table()).await?; + let client = create_plain_connection(port, false).await.unwrap(); + + let err = match feed_copy_in( + &client, + "COPY metrics FROM STDIN WITH (FORMAT csv); SELECT 1", + vec![bytes::Bytes::from_static(b"1,2,3,4\n")], + ) + .await + { + Ok(_) => panic!("expected multi-statement COPY to fail"), + Err(e) => e, + }; + let message = err + .as_db_error() + .expect("expected db error") + .message() + .to_string(); + assert!(message.contains("alone"), "{message}"); + + server.shutdown().await.unwrap(); + Ok(()) +} diff --git a/tests-integration/Cargo.toml b/tests-integration/Cargo.toml index 4ddd9ee5a46..ebf009a5211 100644 --- a/tests-integration/Cargo.toml +++ b/tests-integration/Cargo.toml @@ -77,6 +77,7 @@ uuid.workspace = true [dev-dependencies] base64.workspace = true +bytes.workspace = true catalog = { workspace = true, features = ["testing"] } chrono.workspace = true clap.workspace = true diff --git a/tests-integration/src/test_util.rs b/tests-integration/src/test_util.rs index 0729e6ea96b..3cc820619c6 100644 --- a/tests-integration/src/test_util.rs +++ b/tests-integration/src/test_util.rs @@ -1205,7 +1205,8 @@ pub async fn setup_pg_server_with_slow_query_threshold( ); let mut pg_server = Box::new(PostgresServer::new( - fe_instance_ref, + fe_instance_ref.clone(), + fe_instance_ref.clone(), opts.tls.should_force_tls(), tls_server_config, 0, @@ -1248,7 +1249,8 @@ pub async fn setup_pg_server_with_user_provider( ); let mut pg_server = Box::new(PostgresServer::new( - fe_instance_ref, + fe_instance_ref.clone(), + fe_instance_ref.clone(), opts.tls.should_force_tls(), tls_server_config, 0, @@ -1303,6 +1305,7 @@ pub async fn setup_pg_server_with_prom_native_histogram( ); let mut pg_server = Box::new(PostgresServer::new( + instance.fe_instance().clone(), instance.fe_instance().clone(), opts.tls.should_force_tls(), tls_server_config, diff --git a/tests-integration/src/tests.rs b/tests-integration/src/tests.rs index 4a5b2f0c6c5..1984b849a7e 100644 --- a/tests-integration/src/tests.rs +++ b/tests-integration/src/tests.rs @@ -17,6 +17,7 @@ mod instance_kafka_wal_test; mod instance_noop_wal_test; mod instance_test; mod nested_merge_scan_capability_test; +mod postgres_copy_in_test; mod promql_test; mod reconcile_table; mod remote_dyn_filter_test; diff --git a/tests-integration/src/tests/postgres_copy_in_test.rs b/tests-integration/src/tests/postgres_copy_in_test.rs new file mode 100644 index 00000000000..82e076e7a65 --- /dev/null +++ b/tests-integration/src/tests/postgres_copy_in_test.rs @@ -0,0 +1,521 @@ +// 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. + +//! End-to-end tests for `COPY ... FROM STDIN` over the PostgreSQL protocol, +//! covering the full ingestion path (parser -> insert -> storage). + +use bytes::Bytes; +use common_telemetry::init_default_ut_logging; +use futures::SinkExt; +use tokio_postgres::NoTls; + +use crate::test_util::{StorageType, setup_pg_server}; + +async fn connect(port: u16) -> tokio_postgres::Client { + let url = format!("host=127.0.0.1 port={port} dbname=public connect_timeout=5"); + let (client, conn) = tokio_postgres::connect(&url, NoTls).await.unwrap(); + let _handle = tokio::spawn(conn); + client +} + +async fn feed_copy_in( + client: &tokio_postgres::Client, + sql: &str, + chunks: Vec, +) -> std::result::Result { + let mut sink = Box::pin(client.copy_in(sql).await?); + for chunk in chunks { + sink.as_mut().send(chunk).await?; + } + sink.as_mut().finish().await +} + +fn host_of(message: &tokio_postgres::SimpleQueryMessage, column: usize) -> String { + let tokio_postgres::SimpleQueryMessage::Row(row) = message else { + panic!("expected row, got {message:?}"); + }; + // simple_query renders NULL as `None`, not empty string. + row.get(column).map(str::to_string).unwrap_or_default() +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_postgres_copy_in_end_to_end() { + init_default_ut_logging(); + let (_guard, server) = setup_pg_server(StorageType::File, "pg_copy_in_e2e").await; + let port = server.bind_addr().unwrap().port(); + let client = connect(port).await; + + client + .simple_query( + "CREATE TABLE copy_demo (host STRING, val DOUBLE, note STRING, ts TIMESTAMP, TIME INDEX(ts))", + ) + .await + .unwrap(); + + // CSV format with quoted field and NULL (unquoted empty) value. + let rows = feed_copy_in( + &client, + "COPY copy_demo FROM STDIN WITH (FORMAT csv)", + vec![Bytes::from_static( + b"host1,1.5,\"hello, world\",2023-11-14 22:13:20\nhost2,,x,2023-11-14 22:13:21\n", + )], + ) + .await + .unwrap(); + assert_eq!(rows, 2); + + // A column list targeting a subset. + let rows = feed_copy_in( + &client, + "COPY copy_demo (note, ts) FROM STDIN WITH (FORMAT csv)", + vec![Bytes::from_static( + b"from_column_list,2023-11-14 22:13:24\n", + )], + ) + .await + .unwrap(); + assert_eq!(rows, 1); + + // Verify data written through the real insert path. + let messages = client + .simple_query("SELECT host, val, note FROM copy_demo ORDER BY ts") + .await + .unwrap(); + let rows: Vec<_> = messages + .iter() + .filter(|m| matches!(m, tokio_postgres::SimpleQueryMessage::Row(_))) + .collect(); + assert_eq!(rows.len(), 3); + assert_eq!(host_of(rows[0], 0), "host1"); + assert_eq!(host_of(rows[0], 2), "hello, world"); + assert_eq!(host_of(rows[1], 0), "host2"); + // NULL val and note round-trip as NULL. + assert_eq!(host_of(rows[2], 0), ""); + assert_eq!(host_of(rows[2], 2), "from_column_list"); + + server.shutdown().await.unwrap(); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_postgres_copy_in_error_keeps_table_clean() { + init_default_ut_logging(); + let (_guard, server) = setup_pg_server(StorageType::File, "pg_copy_in_error").await; + let port = server.bind_addr().unwrap().port(); + let client = connect(port).await; + + client + .simple_query( + "CREATE TABLE copy_err (host STRING, val DOUBLE, ts TIMESTAMP, TIME INDEX(ts))", + ) + .await + .unwrap(); + + // Bad double value aborts the copy. + let err = match feed_copy_in( + &client, + "COPY copy_err FROM STDIN WITH (FORMAT csv)", + vec![Bytes::from_static( + b"host1,not_a_double,2023-11-14 22:13:20\n", + )], + ) + .await + { + Ok(_) => panic!("expected COPY to fail"), + Err(e) => e, + }; + assert!( + err.as_db_error() + .expect("expected db error") + .message() + .contains("failed to parse"), + "{err}" + ); + + // A bare COPY (PostgreSQL's text default) fails with guidance. + let err = match client.copy_in::<_, Bytes>("COPY copy_err FROM STDIN").await { + Ok(_) => panic!("expected bare COPY to fail"), + Err(e) => e, + }; + let message = err + .as_db_error() + .expect("expected db error") + .message() + .to_string(); + assert!(message.contains("requires an explicit format"), "{message}"); + + // The connection stays usable and no rows were written. + let messages = client + .simple_query("SELECT host FROM copy_err") + .await + .unwrap(); + assert!( + messages + .iter() + .all(|m| !matches!(m, tokio_postgres::SimpleQueryMessage::Row(_))), + "expected no rows, got {messages:?}" + ); + + server.shutdown().await.unwrap(); +} + +/// Exercises the CSV format's options end-to-end through a Postgres +/// client driver: HEADER, custom delimiter, quoted fields with embedded +/// newlines and escaped quotes, CRLF record separators and NULL fields. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_postgres_copy_in_csv_options() { + init_default_ut_logging(); + let (_guard, server) = setup_pg_server(StorageType::File, "pg_copy_in_csv_options").await; + let port = server.bind_addr().unwrap().port(); + let client = connect(port).await; + + client + .simple_query( + "CREATE TABLE copy_csv_opts (host STRING, val DOUBLE, note STRING, ts TIMESTAMP, TIME INDEX(ts))", + ) + .await + .unwrap(); + + let rows = feed_copy_in( + &client, + // WITH omitted, matching PostgreSQL; HEADER skips the first record. + "COPY copy_csv_opts FROM STDIN (FORMAT csv, HEADER, DELIMITER ';')", + vec![Bytes::from_static( + b"host;val;note;ts\r\n\ + h1;1.5;\"line one\nline two\";2023-11-14 22:13:20\r\n\ + h2;;\"say \"\"hi\"\"\";2023-11-14 22:13:21\r\n", + )], + ) + .await + .unwrap(); + assert_eq!(rows, 2); + + // Verify the parsed values came through the real insert path. + let messages = client + .simple_query("SELECT host, val, note FROM copy_csv_opts ORDER BY ts") + .await + .unwrap(); + let rows: Vec<_> = messages + .iter() + .filter(|m| matches!(m, tokio_postgres::SimpleQueryMessage::Row(_))) + .collect(); + assert_eq!(rows.len(), 2); + assert_eq!(host_of(rows[0], 0), "h1"); + assert_eq!(host_of(rows[0], 1), "1.5"); + // Quoted field with an embedded newline survives the round trip. + assert_eq!(host_of(rows[0], 2), "line one\nline two"); + assert_eq!(host_of(rows[1], 0), "h2"); + // Unquoted empty field is NULL; escaped quotes decode to one quote. + assert_eq!(host_of(rows[1], 1), ""); + assert_eq!(host_of(rows[1], 2), "say \"hi\""); + + server.shutdown().await.unwrap(); +} + +/// COPY OUT → COPY IN round trip: PostgreSQL's COPY OUT emits `t`/`f` +/// for booleans, which COPY IN must accept back. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_postgres_copy_in_bool_round_trip() { + init_default_ut_logging(); + let (_guard, server) = setup_pg_server(StorageType::File, "pg_copy_in_bool_rt").await; + let port = server.bind_addr().unwrap().port(); + let client = connect(port).await; + + client + .simple_query( + "CREATE TABLE bool_src (flag BOOLEAN, ts TIMESTAMP, TIME INDEX(ts)); + CREATE TABLE bool_dst (flag BOOLEAN, ts TIMESTAMP, TIME INDEX(ts)); + INSERT INTO bool_src VALUES (true, '2024-01-01 00:00:00'), (false, '2024-01-01 00:00:01'), (true, '2024-01-01 00:00:02');", + ) + .await + .unwrap(); + + // COPY OUT: collect the exported bytes. + let mut exported = Vec::new(); + let mut out = Box::pin( + client + .copy_out( + "COPY (SELECT flag, ts FROM bool_src ORDER BY ts) TO STDOUT WITH (FORMAT csv)", + ) + .await + .unwrap(), + ); + use futures::StreamExt; + while let Some(chunk) = out.as_mut().next().await { + exported.extend_from_slice(&chunk.unwrap()); + } + // The wire encoder emits PostgreSQL boolean spellings. + let text = String::from_utf8(exported.clone()).unwrap(); + assert!(text.contains("t,"), "COPY OUT should emit t/f, got: {text}"); + assert!(text.contains("f,"), "COPY OUT should emit t/f, got: {text}"); + + // COPY IN: the exact bytes must import back. + let rows = feed_copy_in( + &client, + "COPY bool_dst FROM STDIN WITH (FORMAT csv)", + vec![Bytes::from(exported)], + ) + .await + .unwrap(); + assert_eq!(rows, 3); + + let messages = client + .simple_query("SELECT flag FROM bool_dst ORDER BY ts") + .await + .unwrap(); + let flags: Vec<&str> = messages + .iter() + .filter_map(|m| match m { + tokio_postgres::SimpleQueryMessage::Row(row) => row.get(0), + _ => None, + }) + .collect(); + // The simple-query path also renders booleans in the t/f spelling. + assert_eq!(flags, vec!["t", "f", "t"]); + + server.shutdown().await.unwrap(); +} + +/// A failed COPY must release its state so a later COPY on the same +/// connection works; this guards the error-path state drop. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_postgres_copy_in_failed_copy_then_new_copy() { + init_default_ut_logging(); + let (_guard, server) = setup_pg_server(StorageType::File, "pg_copy_in_fail_reuse").await; + let port = server.bind_addr().unwrap().port(); + let client = connect(port).await; + + client + .simple_query( + "CREATE TABLE copy_reuse (host STRING, val DOUBLE, ts TIMESTAMP, TIME INDEX(ts))", + ) + .await + .unwrap(); + + // A copy that fails at CopyDone (bad double parsed at flush time). + let err = match feed_copy_in( + &client, + "COPY copy_reuse FROM STDIN WITH (FORMAT csv)", + vec![Bytes::from_static( + b"host1,not_a_double,2024-01-01 00:00:00\n", + )], + ) + .await + { + Ok(_) => panic!("expected COPY to fail"), + Err(e) => e, + }; + assert!(err.as_db_error().is_some(), "{err}"); + + // The connection is usable and a NEW copy succeeds on it. + let rows = feed_copy_in( + &client, + "COPY copy_reuse FROM STDIN WITH (FORMAT csv)", + vec![Bytes::from_static(b"host1,1.5,2024-01-01 00:00:01\n")], + ) + .await + .unwrap(); + assert_eq!(rows, 1); + + let messages = client + .simple_query("SELECT host, val FROM copy_reuse") + .await + .unwrap(); + let found: Vec<_> = messages + .iter() + .filter(|m| matches!(m, tokio_postgres::SimpleQueryMessage::Row(_))) + .collect(); + assert_eq!(found.len(), 1); + + server.shutdown().await.unwrap(); +} + +/// COPY/INSERT parity in a non-UTC session: offset-free timestamp strings +/// must be interpreted in the session timezone by both paths. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_postgres_copy_in_session_timezone_parity() { + init_default_ut_logging(); + let (_guard, server) = setup_pg_server(StorageType::File, "pg_copy_in_tz").await; + let port = server.bind_addr().unwrap().port(); + let client = connect(port).await; + + client + .simple_query( + "CREATE TABLE copy_tz (case_id STRING PRIMARY KEY, ts TIMESTAMP, TIME INDEX(ts))", + ) + .await + .unwrap(); + + client + .simple_query("SET TIMEZONE TO 'Asia/Shanghai'") + .await + .unwrap(); + + client + .simple_query("INSERT INTO copy_tz VALUES ('insert', '2024-01-01 00:00:00')") + .await + .unwrap(); + let rows = feed_copy_in( + &client, + "COPY copy_tz (case_id, ts) FROM STDIN WITH (FORMAT csv)", + vec![Bytes::from_static(b"copy,2024-01-01 00:00:00\n")], + ) + .await + .unwrap(); + assert_eq!(rows, 1); + + let messages = client + .simple_query("SELECT case_id, CAST(ts AS BIGINT) FROM copy_tz ORDER BY case_id") + .await + .unwrap(); + let results: Vec<(String, String)> = messages + .iter() + .filter_map(|m| match m { + tokio_postgres::SimpleQueryMessage::Row(row) => { + Some((row.get(0)?.to_string(), row.get(1)?.to_string())) + } + _ => None, + }) + .collect(); + assert_eq!(results.len(), 2, "{results:?}"); + assert_eq!(results[0].0, "copy"); + assert_eq!(results[1].0, "insert"); + // The whole point: both paths store the same instant. + assert_eq!( + results[0].1, results[1].1, + "COPY and INSERT must agree on offset-free timestamps: {results:?}" + ); + // And it is the Asia/Shanghai interpretation, not UTC. + assert_eq!(results[0].1, "1704038400000"); + + server.shutdown().await.unwrap(); +} + +/// COPY OUT -> COPY IN must preserve timestamp instants in non-UTC +/// sessions: COPY OUT renders timestamps with an explicit session-tz +/// offset, which the timezone-aware COPY IN decoder honors. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_postgres_copy_out_in_timestamp_round_trip() { + init_default_ut_logging(); + let (_guard, server) = setup_pg_server(StorageType::File, "pg_copy_out_in_tz").await; + let port = server.bind_addr().unwrap().port(); + let client = connect(port).await; + + client + .simple_query( + "CREATE TABLE tz_src (host STRING, ts TIMESTAMP, TIME INDEX(ts)); + CREATE TABLE tz_dst (host STRING, ts TIMESTAMP, TIME INDEX(ts)); + INSERT INTO tz_src VALUES ('a', '2024-01-01 00:00:00'), ('b', '2024-07-01 12:30:00');", + ) + .await + .unwrap(); + + client + .simple_query("SET TIMEZONE TO 'Asia/Shanghai'") + .await + .unwrap(); + + let mut exported = Vec::new(); + let mut out = Box::pin( + client + .copy_out("COPY (SELECT host, ts FROM tz_src ORDER BY ts) TO STDOUT WITH (FORMAT csv)") + .await + .unwrap(), + ); + use futures::StreamExt; + while let Some(chunk) = out.as_mut().next().await { + exported.extend_from_slice(&chunk.unwrap()); + } + // The export now carries explicit offsets. + let text = String::from_utf8(exported.clone()).unwrap(); + assert!( + text.contains("+08"), + "COPY OUT should render timestamps with the session-tz offset, got: {text}" + ); + + let rows = feed_copy_in( + &client, + "COPY tz_dst (host, ts) FROM STDIN WITH (FORMAT csv)", + vec![Bytes::from(exported)], + ) + .await + .unwrap(); + assert_eq!(rows, 2); + + // Epochs must match the source exactly. + let compare = client + .simple_query( + "SELECT (SELECT CAST(ts AS BIGINT) FROM tz_src WHERE host='a') = (SELECT CAST(ts AS BIGINT) FROM tz_dst WHERE host='a') \ + AND (SELECT CAST(ts AS BIGINT) FROM tz_src WHERE host='b') = (SELECT CAST(ts AS BIGINT) FROM tz_dst WHERE host='b')", + ) + .await + .unwrap(); + let ok: Vec<_> = compare + .iter() + .filter_map(|m| match m { + tokio_postgres::SimpleQueryMessage::Row(row) => row.get(0).map(str::to_string), + _ => None, + }) + .collect(); + assert_eq!(ok, vec!["t"], "epochs must round-trip exactly"); + + server.shutdown().await.unwrap(); +} + +/// Empty fields — quoted or not — import as NULL, matching the +/// file-based COPY ... FROM 'file.csv' behavior exactly. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_postgres_copy_in_empty_field_is_null() { + init_default_ut_logging(); + let (_guard, server) = setup_pg_server(StorageType::File, "pg_copy_in_empty").await; + let port = server.bind_addr().unwrap().port(); + let client = connect(port).await; + + client + .simple_query("CREATE TABLE empty_str (note STRING, ts TIMESTAMP, TIME INDEX(ts))") + .await + .unwrap(); + + let rows = feed_copy_in( + &client, + "COPY empty_str (note, ts) FROM STDIN WITH (FORMAT csv)", + vec![Bytes::from_static( + b"plain,2024-01-01 00:00:00\n\"\",2024-01-01 00:00:01\n,2024-01-01 00:00:02\n", + )], + ) + .await + .unwrap(); + assert_eq!(rows, 3); + + let messages = client + .simple_query("SELECT note IS NULL AS is_null, note FROM empty_str ORDER BY ts") + .await + .unwrap(); + let rows: Vec<(bool, Option)> = messages + .iter() + .filter_map(|m| match m { + tokio_postgres::SimpleQueryMessage::Row(row) => { + Some((row.get(0) == Some("t"), row.get(1).map(str::to_string))) + } + _ => None, + }) + .collect(); + assert_eq!(rows.len(), 3); + assert_eq!(rows[0], (false, Some("plain".to_string()))); + // Quoted empty string: NULL, exactly like the file-based COPY route. + assert_eq!(rows[1], (true, None)); + // Unquoted empty field: NULL. + assert_eq!(rows[2], (true, None)); + + server.shutdown().await.unwrap(); +} diff --git a/tests-integration/tests/sql.rs b/tests-integration/tests/sql.rs index 4a01b0209f6..31a1c50e5d3 100644 --- a/tests-integration/tests/sql.rs +++ b/tests-integration/tests/sql.rs @@ -2908,3 +2908,112 @@ async fn test_sql_batcher_alignment() { instance.guard.remove_all().await; } } + +/// COPY FROM STDIN must route through the pending-rows batcher (with +/// synchronous acknowledgement) when PostgreSQL batching is enabled, and +/// keep direct writes when it is not. +#[tokio::test(flavor = "multi_thread")] +async fn test_postgres_copy_in_with_batching() { + use bytes::Bytes; + use common_telemetry::{dump_metrics, init_default_ut_logging}; + use frontend::server::Services; + use frontend::service_config::BatcherOptions; + use futures::SinkExt; + use servers::batcher::BatchingProtocol; + use tests_integration::standalone::GreptimeDbStandaloneBuilder; + + fn flushes() -> u64 { + dump_metrics() + .unwrap() + .lines() + .find_map(|line| { + line.strip_prefix("greptime_table_batcher_flush_total ") + .map(|value| value.parse().unwrap()) + }) + .unwrap_or(0) + } + + async fn feed_copy_in( + client: &Client, + sql: &str, + chunks: Vec, + ) -> std::result::Result { + let mut sink = Box::pin(client.copy_in(sql).await?); + for chunk in chunks { + sink.as_mut().send(chunk).await?; + } + sink.as_mut().finish().await + } + + init_default_ut_logging(); + for pg_enabled in [false, true] { + let protocols = if pg_enabled { + vec![BatchingProtocol::Postgres] + } else { + vec![] + }; + let mut instance = GreptimeDbStandaloneBuilder::new("pg_copy_in_batching") + .with_table_batcher(BatcherOptions { + protocols, + pending_rows_flush_interval: Duration::from_millis(10), + ..Default::default() + }) + .build() + .await; + let mut opts = instance.opts.clone(); + opts.http.addr = "127.0.0.1:0".into(); + opts.grpc.bind_addr = "127.0.0.1:0".into(); + opts.postgres.addr = "127.0.0.1:0".into(); + let mut servers = Services::new(opts, instance.fe_instance().clone(), Default::default()) + .build() + .unwrap(); + servers.start_all().await.unwrap(); + let (pg, connection) = tokio_postgres::connect( + &format!( + "postgres://{}/public", + servers.addr("POSTGRES_SERVER").unwrap() + ), + NoTls, + ) + .await + .unwrap(); + let pg_task = tokio::spawn(async move { connection.await.unwrap() }); + + pg.simple_query( + "CREATE TABLE batch_copy_in (ts TIMESTAMP TIME INDEX, host STRING, val INT)", + ) + .await + .unwrap(); + + let before = flushes(); + let rows = feed_copy_in( + &pg, + "COPY batch_copy_in FROM STDIN WITH (FORMAT csv)", + vec![Bytes::from_static( + b"2023-11-14 22:13:20,host1,1\n2023-11-14 22:13:21,host2,2\n", + )], + ) + .await + .unwrap(); + // Synchronous acknowledgement: the reported count is exact whether + // or not the batcher is active. + assert_eq!(rows, 2); + // The batch must have been flushed through the pending-rows batcher. + assert_eq!(flushes() - before, if pg_enabled { 1 } else { 0 }); + + // The flushed rows are visible immediately after the copy. + let rows: Vec = pg + .query("SELECT host FROM batch_copy_in ORDER BY host", &[]) + .await + .unwrap() + .iter() + .map(|row| row.get(0)) + .collect(); + assert_eq!(rows, vec!["host1", "host2"]); + + drop(pg); + pg_task.await.unwrap(); + servers.shutdown_all().await.unwrap(); + instance.guard.remove_all().await; + } +} From 2e37e19af495873254cbde375e425a8177b4ddd0 Mon Sep 17 00:00:00 2001 From: Dhruv Vaishnav Date: Mon, 26 Oct 2026 07:09:21 +0000 Subject: [PATCH 089/239] feat(datanode): expose ingest task panic and cancellation metrics (#9830) * feat(datanode): record ingest task join failures Signed-off-by: dhruvxvaishnav * test(datanode): remove extra ingest task metric tests Signed-off-by: dhruvxvaishnav --------- Signed-off-by: dhruvxvaishnav Co-authored-by: Weny Xu --- src/datanode/src/metrics.rs | 9 +++++++ src/datanode/src/region_server.rs | 40 ++++++++++++++++++++----------- 2 files changed, 35 insertions(+), 14 deletions(-) diff --git a/src/datanode/src/metrics.rs b/src/datanode/src/metrics.rs index e9f947ef9da..af9a82ed551 100644 --- a/src/datanode/src/metrics.rs +++ b/src/datanode/src/metrics.rs @@ -85,6 +85,15 @@ lazy_static! { ) .unwrap(); + /// Ingest-task join failures, labeled by request type and reason (`panic` or `cancelled`). + /// Counts tasks, not write entries: a failed task can leave entry outcomes unknown. + pub static ref INGEST_TASK_FAILURES_TOTAL: IntCounterVec = register_int_counter_vec!( + "greptime_datanode_ingest_task_failures_total", + "Ingest-task join failures, not failed write entries", + &[REGION_REQUEST_TYPE, "reason"] + ) + .unwrap(); + /// Failed region open attempts, including individual and batch requests. /// Counts every request in an engine group when no individual results are available. /// Failed retries count separately; explicitly ignored missing regions do not count. diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index 135d069fd77..3acf188ad19 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -49,7 +49,7 @@ use common_query::OutputData; use common_query::request::QueryRequest; use common_recordbatch::adapter::RecordBatchMetrics; use common_recordbatch::{OrderOption, RecordBatch, RecordBatchStream, SendableRecordBatchStream}; -use common_runtime::Runtime; +use common_runtime::{JoinHandle, Runtime}; use common_telemetry::tracing::{self, info_span}; use common_telemetry::tracing_context::{FutureExt, TracingContext}; use common_telemetry::{debug, error, info, warn}; @@ -250,11 +250,10 @@ impl RegionServer { if RegionServerInner::is_ingest_request(&request) { let inner = self.inner.clone(); let request_type = request.request_type(); - return common_runtime::spawn_ingest(async move { + let task = common_runtime::spawn_ingest(async move { inner.handle_request(region_id, request).await - }) - .await - .context(RuntimeJoinSnafu { request_type })?; + }); + return join_ingest_task(task, request_type).await?; } self.inner.handle_request(region_id, request).await @@ -679,14 +678,13 @@ impl RegionServer { "RegionServer::handle_metric_batch_puts", batch_size = batch_size )); - let results = common_runtime::spawn_ingest(async move { + let task = common_runtime::spawn_ingest(async move { metric_engine .put_regions_batch(put_requests) .trace(span) .await - }) - .await - .context(RuntimeJoinSnafu { request_type })?; + }); + let results = join_ingest_task(task, request_type).await?; let mut response = BatchRegionResponse { affected_rows: 0, extensions: HashMap::new(), @@ -746,7 +744,7 @@ impl RegionServer { .collect(); let started = Instant::now(); let tracing_context = TracingContext::from_current_span(); - common_runtime::spawn_ingest(async move { + let task = common_runtime::spawn_ingest(async move { let mut pending = engine .handle_batch_put_requests(puts) .into_iter() @@ -800,10 +798,8 @@ impl RegionServer { } response.failures.sort_unstable_by_key(|(index, _)| *index); response - }) - .await - .context(RuntimeJoinSnafu { request_type }) - .map(Either::Left) + }); + join_ingest_task(task, request_type).await.map(Either::Left) } /// Returns the mito engine when there are at least two requests, every one a put into a ready @@ -2292,6 +2288,22 @@ fn record_region_open_failures(engine: &str, status_code: StatusCode, count: usi .inc_by(count as u64); } +/// Awaits an ingest task without treating a join failure as an entry-level write result. +async fn join_ingest_task(task: JoinHandle, request_type: &'static str) -> Result { + task.await + .inspect_err(|error| { + let reason = if error.is_panic() { + "panic" + } else { + "cancelled" + }; + crate::metrics::INGEST_TASK_FAILURES_TOTAL + .with_label_values(&[request_type, reason]) + .inc(); + }) + .context(RuntimeJoinSnafu { request_type }) +} + /// Records the result of `count` ingest requests answered by one response: the rows they /// changed, or their failures. fn record_ingest_result(request_type: &str, count: usize, result: &Result) { From 88cb2a986ef0776a1eb84b0df9581b3e7ffb189d Mon Sep 17 00:00:00 2001 From: Dragon Roar <2181719471@qq.com> Date: Mon, 26 Oct 2026 07:13:41 +0000 Subject: [PATCH 090/239] fix: reclaim HDFS temporary files left behind by a datanode crash (#9616) * fix: reclaim HDFS temporary files left behind by a datanode crash Signed-off-by: lczllx <2181719471@qq.com> * fix(mito2): stop the temporary file sweep from starving GC jobs Signed-off-by: lczllx <2181719471@qq.com> * feat(mito2): add metrics for the temporary file sweep Signed-off-by: lczllx <2181719471@qq.com> * refactor(object-store): share one definition of the temporary file name Signed-off-by: lczllx <2181719471@qq.com> * fix(mito2): skip the temporary file sweep when no GC job is allowed Signed-off-by: lczllx <2181719471@qq.com> * refactor: gate the temporary file sweep behind the hdfs-object-store feature Signed-off-by: lczllx <2181719471@qq.com> * fix(mito2): sweep temporary files when a create opens an existing region Signed-off-by: lczllx <2181719471@qq.com> * docs(config): document the temporary file sweep in the GC settings Signed-off-by: lczllx <2181719471@qq.com> * refactor: align the temporary file sweep code with the surrounding style Signed-off-by: lczllx <2181719471@qq.com> --------- Signed-off-by: lczllx <2181719471@qq.com> --- config/config.md | 2 +- config/datanode.example.toml | 5 +- src/mito2/AGENTS.md | 1 + src/mito2/src/gc.rs | 15 +- src/mito2/src/lib.rs | 2 + src/mito2/src/metrics.rs | 21 ++ src/mito2/src/sst/location.rs | 7 +- src/mito2/src/temp_file_sweeper.rs | 152 ++++++++++ src/mito2/src/worker.rs | 45 +++ src/mito2/src/worker/handle_catchup.rs | 21 +- src/mito2/src/worker/handle_create.rs | 12 + src/object-store/src/factory.rs | 2 + src/object-store/src/layers/hdfs.rs | 48 +++- src/object-store/src/lib.rs | 2 + src/object-store/src/temp_file.rs | 373 +++++++++++++++++++++++++ 15 files changed, 686 insertions(+), 22 deletions(-) create mode 100644 src/mito2/src/temp_file_sweeper.rs create mode 100644 src/object-store/src/temp_file.rs diff --git a/config/config.md b/config/config.md index e370bb157a0..36c6f02baff 100644 --- a/config/config.md +++ b/config/config.md @@ -724,7 +724,7 @@ | `region_engine.mito.gc` | -- | -- | -- | | `region_engine.mito.gc.enable` | Bool | `false` | Whether GC is enabled. Need to be the same with metasrv's `gc.enable` or unexpected behavior will occur | | `region_engine.mito.gc.lingering_time` | String | `1h` | Lingering time before deleting files.
Should be long enough to allow long running queries to finish.
If set to None, then unused files will be deleted immediately. | -| `region_engine.mito.gc.unknown_file_lingering_time` | String | `1d` | Lingering time before deleting unknown files (files with undetermined expel time).
Only applies during full file listing GC.
This uses the object's last-modified timestamp as a heuristic (strict less-than comparison);
do not configure this value too small in production to avoid deleting pre-manifest files
from in-progress compaction or flush.
For active/open regions, an unknown file is deleted only if its object last-modified time exceeds this TTL.
If the object store does not provide a last-modified timestamp, the file is conservatively kept.
For dropped regions, unknown files are deleted immediately. | +| `region_engine.mito.gc.unknown_file_lingering_time` | String | `1d` | Lingering time before deleting unknown files (files with undetermined expel time).
It is also the age threshold of the sweep that reclaims temporary files left behind by
a crashed HDFS writer, and that sweep runs even if `enable` is false. HDFS advances the
modification time of a file under construction only as its blocks complete, so lowering
this value also shrinks the margin between the sweep and files that are still being written.
This uses the object's last-modified timestamp as a heuristic (strict less-than comparison);
do not configure this value too small in production to avoid deleting pre-manifest files
from in-progress compaction or flush.
For active/open regions, an unknown file is deleted only if its object last-modified time exceeds this TTL.
If the object store does not provide a last-modified timestamp, the file is conservatively kept.
For dropped regions, unknown files are deleted immediately. | | `region_engine.file` | -- | -- | Enable the file engine. | | `region_engine.metric` | -- | -- | Metric engine options. | | `logging` | -- | -- | The logging options. | diff --git a/config/datanode.example.toml b/config/datanode.example.toml index bfa3afc0bcf..f6e9a2571f8 100644 --- a/config/datanode.example.toml +++ b/config/datanode.example.toml @@ -831,7 +831,10 @@ enable = false lingering_time = "1h" ## Lingering time before deleting unknown files (files with undetermined expel time). -## Only applies during full file listing GC. +## It is also the age threshold of the sweep that reclaims temporary files left behind by +## a crashed HDFS writer, and that sweep runs even if `enable` is false. HDFS advances the +## modification time of a file under construction only as its blocks complete, so lowering +## this value also shrinks the margin between the sweep and files that are still being written. ## This uses the object's last-modified timestamp as a heuristic (strict less-than comparison); ## do not configure this value too small in production to avoid deleting pre-manifest files ## from in-progress compaction or flush. diff --git a/src/mito2/AGENTS.md b/src/mito2/AGENTS.md index 811632b1e26..a02fd4dfd43 100644 --- a/src/mito2/AGENTS.md +++ b/src/mito2/AGENTS.md @@ -34,6 +34,7 @@ snapshot isolation). It implements the `RegionEngine` trait from `store-api`. | `schedule` | `src/mito2/src/schedule/` | Local/remote background job scheduling | | `remap_manifest` | `src/mito2/src/remap_manifest.rs` | Manifest path remapping for region copy/migration | | `config` | `src/mito2/src/config.rs` | `MitoConfig` tuning knobs | +| `temp_file_sweeper` | `src/mito2/src/temp_file_sweeper.rs` | Reclaims temporary files a crashed writer left in a region, triggered when the region becomes a writable leader or a create opens an existing region (under the `hdfs-object-store` feature) | | `test_util` | `src/mito2/src/test_util.rs` | `TestEnv` and builders (under the `test` feature) | ## Write path diff --git a/src/mito2/src/gc.rs b/src/mito2/src/gc.rs index a5474099cbf..541bb4872b4 100644 --- a/src/mito2/src/gc.rs +++ b/src/mito2/src/gc.rs @@ -159,6 +159,12 @@ pub struct GcConfig { /// expel time is the time when the file is considered as removed, as in removed from the manifest. /// This should only occur rarely, as manifest keep tracks in `removed_files` field /// unless something goes wrong. + /// + /// It is also the age threshold of the sweep that reclaims temporary files left behind + /// by a crashed HDFS writer, and that sweep runs even if `enable` is false. HDFS advances + /// the modification time of a file under construction only as its blocks complete, so + /// lowering this value also shrinks the margin between the sweep and files that are + /// still being written. #[serde(with = "humantime_serde")] pub unknown_file_lingering_time: Duration, /// Maximum concurrent list operations per GC job. @@ -167,6 +173,8 @@ pub struct GcConfig { /// Maximum concurrent GC jobs. /// This is used to limit the number of concurrent GC jobs running on the datanode /// to prevent too many concurrent GC jobs from overwhelming the datanode. + /// Setting it to 0 disables GC jobs and the temporary file sweep, which is limited + /// to the same number of concurrent tasks. pub max_concurrent_gc_job: usize, } @@ -177,8 +185,9 @@ impl Default for GcConfig { // expect long running queries to be finished(or at least be able to notify it's using a deleted file) within a reasonable time lingering_time: Some(Duration::from_secs(60 * 60)), // 1 day, for unknown expel time, which is when this file get removed from manifest. - // Only applies to full-listing GC for active/open regions. A long default avoids - // accidentally deleting pre-manifest files (e.g. compaction/flush still in progress). + // A long default avoids accidentally deleting pre-manifest files (e.g. + // compaction/flush still in progress) or temporary files of a writer that is + // still alive; the temporary file sweep uses this value as well. unknown_file_lingering_time: Duration::from_secs(24 * 60 * 60), max_concurrent_lister_per_gc_job: 32, max_concurrent_gc_job: 4, @@ -806,7 +815,7 @@ impl LocalGcWorker { /// Only `.puffin` files (not subdirectories) are included. async fn list_region_index_files(&self, region_id: RegionId) -> Result> { let region_dir = self.access_layer.build_region_dir(region_id); - let index_dir = object_store::util::join_dir(®ion_dir, "index"); + let index_dir = location::index_dir(®ion_dir); let mut lister = match self .access_layer diff --git a/src/mito2/src/lib.rs b/src/mito2/src/lib.rs index dcbfe60dea5..a01bca47dd7 100644 --- a/src/mito2/src/lib.rs +++ b/src/mito2/src/lib.rs @@ -42,6 +42,8 @@ pub mod request; pub mod schedule; pub mod series_index; pub mod sst; +#[cfg(feature = "hdfs-object-store")] +mod temp_file_sweeper; mod time_provider; pub mod wal; pub mod worker; diff --git a/src/mito2/src/metrics.rs b/src/mito2/src/metrics.rs index cffa6de39a1..e8718f0e758 100644 --- a/src/mito2/src/metrics.rs +++ b/src/mito2/src/metrics.rs @@ -607,6 +607,27 @@ lazy_static! { register_int_counter!("greptime_mito_flush_file_total", "mito flushed file count").unwrap(); } +#[cfg(feature = "hdfs-object-store")] +lazy_static! { + /// Counter for the number of temporary files deleted by the sweeper. + pub static ref SWEEP_FILES_DELETED_TOTAL: IntCounter = register_int_counter!( + "greptime_mito_sweep_files_deleted_total", + "Total temporary files deleted by the sweeper", + ).unwrap(); + + /// Counter for the number of failures while sweeping temporary files. + pub static ref SWEEP_ERRORS_TOTAL: IntCounter = register_int_counter!( + "greptime_mito_sweep_errors_total", + "Total failures while sweeping temporary files", + ).unwrap(); + + /// Counter for the number of temporary file sweeps run. + pub static ref SWEEP_RUNS_TOTAL: IntCounter = register_int_counter!( + "greptime_mito_sweep_runs_total", + "Total temporary file sweeps run", + ).unwrap(); +} + /// Stager notifier to collect metrics. pub struct StagerMetrics { cache_hit: IntCounter, diff --git a/src/mito2/src/sst/location.rs b/src/mito2/src/sst/location.rs index 5bff4a6d4cf..379d3f3514e 100644 --- a/src/mito2/src/sst/location.rs +++ b/src/mito2/src/sst/location.rs @@ -39,6 +39,11 @@ pub fn region_dir_from_table_dir( } } +/// Generate index dir from region_dir +pub fn index_dir(region_dir: &str) -> String { + util::join_dir(region_dir, "index") +} + pub fn sst_file_path(table_dir: &str, region_file_id: RegionFileId, path_type: PathType) -> String { let region_dir = region_dir_from_table_dir(table_dir, region_file_id.region_id(), path_type); util::join_path( @@ -49,7 +54,7 @@ pub fn sst_file_path(table_dir: &str, region_file_id: RegionFileId, path_type: P pub fn index_file_path(table_dir: &str, index_id: RegionIndexId, path_type: PathType) -> String { let region_dir = region_dir_from_table_dir(table_dir, index_id.file_id.region_id(), path_type); - let index_dir = util::join_dir(®ion_dir, "index"); + let index_dir = index_dir(®ion_dir); let filename = if index_id.version == 0 { format!("{}.puffin", index_id.file_id.file_id()) diff --git a/src/mito2/src/temp_file_sweeper.rs b/src/mito2/src/temp_file_sweeper.rs new file mode 100644 index 00000000000..8cece62aef2 --- /dev/null +++ b/src/mito2/src/temp_file_sweeper.rs @@ -0,0 +1,152 @@ +// 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. + +//! Reclaims the temporary files that a crashed writer left in a region. + +use std::sync::Arc; +use std::time::Duration; + +use common_telemetry::info; +use object_store::temp_file::{SweepReport, sweep_temporary_files, writes_temporary_files}; +use store_api::storage::RegionId; +use tokio::sync::Semaphore; + +use crate::access_layer::{AccessLayer, AccessLayerRef}; +use crate::manifest::storage::manifest_dir; +use crate::manifest::storage::staging::{staging_blob_path, staging_manifest_path}; +use crate::metrics::{SWEEP_ERRORS_TOTAL, SWEEP_FILES_DELETED_TOTAL, SWEEP_RUNS_TOTAL}; +use crate::sst::location::index_dir; + +/// Returns the directories that `region_id` writes non-append files into. +/// +/// Every one of them can hold the temporary file of a crashed writer, and a +/// list is not recursive, so the sweep has to name each directory. +pub(crate) fn region_write_dirs(access_layer: &AccessLayer, region_id: RegionId) -> Vec { + let region_dir = access_layer.build_region_dir(region_id); + let manifest_path = manifest_dir(®ion_dir); + vec![ + region_dir.clone(), + index_dir(®ion_dir), + manifest_path.clone(), + staging_blob_path(&manifest_path), + staging_manifest_path(&manifest_path), + ] +} + +/// Sweeps the temporary files of `region_id` in the background. +/// +/// The caller is on the path that makes the region its writable leader, so +/// reclaiming leftovers must not slow it down or fail it: everything happens in +/// a spawned task that only logs. +/// +/// `older_than` is the margin the caller takes from GC's +/// `unknown_file_lingering_time`, which guards the same thing: a file this old +/// is not one a writer is still filling. +/// +/// A sweep waits for a permit of `sweep_limiter`, so a `None` means the +/// datanode allows no concurrent GC job and there is no sweep. +pub(crate) fn spawn_region_sweep( + access_layer: AccessLayerRef, + region_id: RegionId, + older_than: Duration, + sweep_limiter: Option>, +) { + // No write to this store leaves a temporary file behind, so don't spend a + // task on it: a sweep would still list every directory of the region to + // find that out. + if !writes_temporary_files(access_layer.object_store()) { + return; + } + + let Some(sweep_limiter) = sweep_limiter else { + return; + }; + + common_runtime::spawn_global(async move { + // A burst of regions becoming writable must not run more sweeps at once + // than the datanode allows GC jobs. Wait for a permit instead of giving + // up, the sweep has no deadline. + let Ok(_permit) = sweep_limiter.acquire_owned().await else { + return; + }; + + let dirs = region_write_dirs(&access_layer, region_id); + let report = sweep_temporary_files(access_layer.object_store(), &dirs, older_than).await; + record_sweep(&report); + if report.deleted > 0 { + info!( + "Reclaimed {} temporary file(s) of a crashed writer in region {region_id}", + report.deleted + ); + } + }); +} + +fn record_sweep(report: &SweepReport) { + SWEEP_RUNS_TOTAL.inc(); + SWEEP_FILES_DELETED_TOTAL.inc_by(report.deleted as u64); + SWEEP_ERRORS_TOTAL.inc_by(report.failed as u64); +} + +#[cfg(test)] +mod tests { + use object_store::ObjectStore; + use object_store::services::Memory; + use store_api::region_request::PathType; + + use super::*; + use crate::test_util::TestEnv; + + fn test_access_layer(env: &TestEnv) -> AccessLayer { + AccessLayer::new( + "test", + PathType::Bare, + ObjectStore::new(Memory::default()).unwrap(), + env.get_puffin_manager(), + env.get_intermediate_manager(), + ) + } + + #[tokio::test] + async fn test_region_write_dirs() { + let env = TestEnv::new().await; + let access_layer = test_access_layer(&env); + let region_id = RegionId::new(1, 2); + + let dirs = region_write_dirs(&access_layer, region_id); + + assert_eq!(5, dirs.len()); + assert!(dirs[0].ends_with("1_0000000002/")); + assert!(dirs[1].ends_with("1_0000000002/index/")); + assert!(dirs[2].ends_with("1_0000000002/manifest/")); + assert!(dirs[3].ends_with("1_0000000002/staging/blob/")); + assert!(dirs[4].ends_with("1_0000000002/staging/manifest/")); + } + + #[test] + fn test_record_sweep() { + let runs = SWEEP_RUNS_TOTAL.get(); + let deleted = SWEEP_FILES_DELETED_TOTAL.get(); + let errors = SWEEP_ERRORS_TOTAL.get(); + + record_sweep(&SweepReport { + deleted: 2, + failed: 1, + }); + + assert_eq!(runs + 1, SWEEP_RUNS_TOTAL.get()); + assert_eq!(deleted + 2, SWEEP_FILES_DELETED_TOTAL.get()); + assert_eq!(errors + 1, SWEEP_ERRORS_TOTAL.get()); + } +} diff --git a/src/mito2/src/worker.rs b/src/mito2/src/worker.rs index c92070ac4e7..789b0f539b1 100644 --- a/src/mito2/src/worker.rs +++ b/src/mito2/src/worker.rs @@ -87,6 +87,8 @@ use crate::sst::file_ref::FileReferenceManagerRef; use crate::sst::index::IndexBuildScheduler; use crate::sst::index::intermediate::IntermediateManager; use crate::sst::index::puffin_manager::PuffinManagerFactory; +#[cfg(feature = "hdfs-object-store")] +use crate::temp_file_sweeper; use crate::time_provider::{StdTimeProvider, TimeProviderRef}; use crate::wal::Wal; use crate::worker::handle_manifest::RegionEditQueues; @@ -241,6 +243,13 @@ impl WorkerGroup { let compaction_memory_manager = Arc::new(new_compaction_memory_manager(compaction_limit_bytes)); let gc_limiter = Arc::new(GcLimiter::new(config.gc.max_concurrent_gc_job)); + // A sweep waits for a permit while a GC job takes one without waiting, so + // sharing the GC limiter would fail GC jobs and report sweeps as GC jobs. + // A datanode that allows no concurrent GC job runs no sweep either: a + // limiter of zero would never hand out the permit the task waits for. + #[cfg(feature = "hdfs-object-store")] + let sweep_limiter = (config.gc.max_concurrent_gc_job > 0) + .then(|| Arc::new(Semaphore::new(config.gc.max_concurrent_gc_job))); let workers = (0..config.num_workers) .map(|id| { @@ -269,6 +278,8 @@ impl WorkerGroup { file_ref_manager: file_ref_manager.clone(), partition_expr_fetcher: partition_expr_fetcher.clone(), flush_semaphore: flush_semaphore.clone(), + #[cfg(feature = "hdfs-object-store")] + sweep_limiter: sweep_limiter.clone(), } .start() }) @@ -491,6 +502,13 @@ impl WorkerGroup { let compaction_memory_manager = Arc::new(new_compaction_memory_manager(compaction_limit_bytes)); let gc_limiter = Arc::new(GcLimiter::new(config.gc.max_concurrent_gc_job)); + // A sweep waits for a permit while a GC job takes one without waiting, so + // sharing the GC limiter would fail GC jobs and report sweeps as GC jobs. + // A datanode that allows no concurrent GC job runs no sweep either: a + // limiter of zero would never hand out the permit the task waits for. + #[cfg(feature = "hdfs-object-store")] + let sweep_limiter = (config.gc.max_concurrent_gc_job > 0) + .then(|| Arc::new(Semaphore::new(config.gc.max_concurrent_gc_job))); let workers = (0..config.num_workers) .map(|id| { WorkerStarter { @@ -518,6 +536,8 @@ impl WorkerGroup { file_ref_manager: file_ref_manager.clone(), partition_expr_fetcher: partition_expr_fetcher.clone(), flush_semaphore: flush_semaphore.clone(), + #[cfg(feature = "hdfs-object-store")] + sweep_limiter: sweep_limiter.clone(), } .start() }) @@ -628,6 +648,10 @@ struct WorkerStarter { file_ref_manager: FileReferenceManagerRef, partition_expr_fetcher: PartitionExprFetcherRef, flush_semaphore: Arc, + #[cfg(feature = "hdfs-object-store")] + /// Limiter for the temporary file sweeps, or `None` if the datanode allows + /// no concurrent GC job. + sweep_limiter: Option>, } impl WorkerStarter { @@ -728,6 +752,8 @@ impl WorkerStarter { partition_expr_fetcher: self.partition_expr_fetcher, flush_semaphore: self.flush_semaphore, plugins: self.plugins, + #[cfg(feature = "hdfs-object-store")] + sweep_limiter: self.sweep_limiter, }; let handle = common_runtime::spawn_global(async move { worker_thread.run().await; @@ -1055,6 +1081,10 @@ struct RegionWorkerLoop { flush_semaphore: Arc, /// Plugins for flush hooks. plugins: Plugins, + #[cfg(feature = "hdfs-object-store")] + /// Limiter for the temporary file sweeps, or `None` if the datanode allows + /// no concurrent GC job. + sweep_limiter: Option>, } impl RegionWorkerLoop { @@ -1475,6 +1505,10 @@ impl RegionWorkerLoop { // changes a writable or staging leader. region.set_role(RegionRole::DowngradingLeader); } + #[cfg(feature = "hdfs-object-store")] + let sweep_older_than = self.config.gc.unknown_file_lingering_time; + #[cfg(feature = "hdfs-object-store")] + let sweep_limiter = self.sweep_limiter.clone(); // We need to do this in background as we need the manifest lock. common_runtime::spawn_global(async move { match region.set_role_state_gracefully(region_role_state).await { @@ -1483,6 +1517,17 @@ impl RegionWorkerLoop { let _ = sender.send(SetRegionRoleStateResponse::success( SetRegionRoleStateSuccess::mito(last_entry_id), )); + // Only a writable leader reclaims what a writer that died + // before it renamed its temporary file left behind. + #[cfg(feature = "hdfs-object-store")] + if region_role_state == SettableRegionRoleState::Leader { + temp_file_sweeper::spawn_region_sweep( + region.access_layer(), + region_id, + sweep_older_than, + sweep_limiter, + ); + } } Err(e) => { error!(e; "Failed to set region {} role state to {:?}", region_id, region_role_state); diff --git a/src/mito2/src/worker/handle_catchup.rs b/src/mito2/src/worker/handle_catchup.rs index c526aab8c0c..f817c1309d8 100644 --- a/src/mito2/src/worker/handle_catchup.rs +++ b/src/mito2/src/worker/handle_catchup.rs @@ -28,6 +28,8 @@ use crate::region::MitoRegion; use crate::region::catchup::RegionCatchupTask; use crate::region::opener::RegionOpener; use crate::request::OptionOutputTx; +#[cfg(feature = "hdfs-object-store")] +use crate::temp_file_sweeper; use crate::wal::entry_distributor::WalEntryReceiver; use crate::worker::RegionWorkerLoop; @@ -69,6 +71,10 @@ impl RegionWorkerLoop { let catchup_regions = self.catchup_regions.clone(); let wal = self.wal.clone(); let allow_stale_entries = self.config.allow_stale_entries; + #[cfg(feature = "hdfs-object-store")] + let sweep_older_than = self.config.gc.unknown_file_lingering_time; + #[cfg(feature = "hdfs-object-store")] + let sweep_limiter = self.sweep_limiter.clone(); common_runtime::spawn_global(async move { let mut task = RegionCatchupTask::new(region.clone(), wal, allow_stale_entries) .with_entry_receiver(entry_receiver) @@ -80,11 +86,22 @@ impl RegionWorkerLoop { if request.set_writable { region.set_role(RegionRole::Leader); // Finalize leadership: persist backfilled metadata. - if let Err(err) = region + match region .set_role_state_gracefully(SettableRegionRoleState::Leader) .await { - error!(err; "Failed to set region {region_id} to leader"); + Ok(()) => { + #[cfg(feature = "hdfs-object-store")] + temp_file_sweeper::spawn_region_sweep( + region.access_layer(), + region_id, + sweep_older_than, + sweep_limiter, + ); + } + Err(err) => { + error!(err; "Failed to set region {region_id} to leader"); + } } } catchup_regions.remove_region(region_id); diff --git a/src/mito2/src/worker/handle_create.rs b/src/mito2/src/worker/handle_create.rs index cff77b50f5f..c31e9d4ac7c 100644 --- a/src/mito2/src/worker/handle_create.rs +++ b/src/mito2/src/worker/handle_create.rs @@ -22,6 +22,8 @@ use store_api::storage::RegionId; use crate::error::Result; use crate::region::opener::{RegionOpener, check_recovered_region}; +#[cfg(feature = "hdfs-object-store")] +use crate::temp_file_sweeper; use crate::worker::RegionWorkerLoop; impl RegionWorkerLoop { @@ -79,6 +81,16 @@ impl RegionWorkerLoop { let region = opener.create_or_open(&self.config, &self.wal).await?; + // A Create retried after a crash opens an existing region as a leader + // without a role-state change, so sweep here as well. + #[cfg(feature = "hdfs-object-store")] + temp_file_sweeper::spawn_region_sweep( + region.access_layer(), + region_id, + self.config.gc.unknown_file_lingering_time, + self.sweep_limiter.clone(), + ); + info!( "A new region created with requirement {:?}, worker: {}, region: {:?}", requirements, diff --git a/src/object-store/src/factory.rs b/src/object-store/src/factory.rs index 356d3aecb21..e9dc86aa96b 100644 --- a/src/object-store/src/factory.rs +++ b/src/object-store/src/factory.rs @@ -197,6 +197,8 @@ mod tests { let store = new_hdfs_object_store(&config).await.unwrap(); assert_eq!(HDFS_NATIVE_SCHEME, store.info().scheme()); assert_eq!("/greptimedb/", store.info().root()); + // The sweep recognizes this store as one whose writes leave temporary files. + assert!(crate::temp_file::writes_temporary_files(&store)); } #[tokio::test] diff --git a/src/object-store/src/layers/hdfs.rs b/src/object-store/src/layers/hdfs.rs index ed8b2019450..529c137cdff 100644 --- a/src/object-store/src/layers/hdfs.rs +++ b/src/object-store/src/layers/hdfs.rs @@ -23,7 +23,8 @@ use opendal::raw::{ OpWrite, RpCreateDir, RpPresign, RpRename, RpStat, Service, ServiceInfo, Servicer, oio, }; use opendal::{Buffer, Capability, ErrorKind, Metadata, OperationContext, Result}; -use uuid::Uuid; + +use crate::temp_file::temporary_path; /// Adds atomic writes and streaming copies to the native HDFS backend. #[derive(Debug, Clone)] @@ -219,7 +220,10 @@ impl oio::Write for HdfsWriter { Ok(metadata) => metadata, Err(error) => { drop(writer); - let _ = delete_path(inner, context, temporary_path).await; + if let Err(delete_error) = delete_path(inner, context, temporary_path).await + { + warn!(delete_error; "Failed to delete temporary file {temporary_path} after a failed HDFS close"); + } return Err(error); } }; @@ -229,7 +233,9 @@ impl oio::Write for HdfsWriter { .rename(inner, context, temporary_path, target_path) .await { - let _ = delete_path(inner, context, temporary_path).await; + if let Err(delete_error) = delete_path(inner, context, temporary_path).await { + warn!(delete_error; "Failed to delete temporary file {temporary_path} after a failed HDFS rename"); + } return Err(error); } @@ -424,7 +430,9 @@ async fn copy_via_read_write( Ok(metadata) => metadata, Err(error) => { drop(writer); - let _ = delete_path(&inner, context, &temporary_path).await; + if let Err(delete_error) = delete_path(&inner, context, &temporary_path).await { + warn!(delete_error; "Failed to delete temporary file {temporary_path} after a failed HDFS close"); + } return Err(error); } }; @@ -434,7 +442,9 @@ async fn copy_via_read_write( .rename(&inner, context, &temporary_path, target_path) .await { - let _ = delete_path(&inner, context, &temporary_path).await; + if let Err(delete_error) = delete_path(&inner, context, &temporary_path).await { + warn!(delete_error; "Failed to delete temporary file {temporary_path} after a failed HDFS rename"); + } return Err(error); } @@ -462,15 +472,6 @@ async fn delete_path(inner: &Servicer, context: &OperationContext, path: &str) - deleter.close().await } -// TODO(fengjiachun): Clean up temporary files left behind after a process crash. -fn temporary_path(path: &str) -> String { - let suffix = format!(".greptime-{}.tmp", Uuid::new_v4()); - match path.rsplit_once('/') { - Some((parent, name)) => format!("{parent}/.{name}{suffix}"), - None => format!(".{path}{suffix}"), - } -} - #[cfg(test)] mod tests { use opendal::services::Fs; @@ -565,6 +566,25 @@ mod tests { ); } + #[tokio::test] + async fn test_list_reports_temporary_file_of_crashed_writer() { + // Regression test for #9591: a temporary file left by a crashed writer + // must stay visible to a sweep over the directory of its target. + let (_directory, store) = test_store(); + store.create_dir("data/").await.unwrap(); + let temporary = temporary_path("data/manifest.json"); + assert!(temporary.starts_with("data/.manifest.json.greptime-")); + assert!(temporary.ends_with(".tmp")); + store.write(&temporary, "partial").await.unwrap(); + + let entries = store.list("data/").await.unwrap(); + let entry = entries + .iter() + .find(|entry| entry.path() == temporary) + .unwrap(); + assert!(entry.metadata().last_modified().is_some()); + } + #[tokio::test] async fn test_copy_fallback_streams_to_target() { let (_directory, store) = test_store(); diff --git a/src/object-store/src/lib.rs b/src/object-store/src/lib.rs index 6c8956445d8..e1c421797da 100644 --- a/src/object-store/src/lib.rs +++ b/src/object-store/src/lib.rs @@ -25,6 +25,8 @@ pub mod factory; pub mod layers; pub mod manager; pub mod secure_fs; +#[cfg(feature = "hdfs-object-store")] +pub mod temp_file; #[cfg(feature = "testing")] pub mod test_util; pub mod util; diff --git a/src/object-store/src/temp_file.rs b/src/object-store/src/temp_file.rs new file mode 100644 index 00000000000..ee169cbeefe --- /dev/null +++ b/src/object-store/src/temp_file.rs @@ -0,0 +1,373 @@ +// 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. + +//! Temporary files of non-append writes, and how to reclaim them. + +use std::time::Duration; + +use common_telemetry::warn; +use opendal::raw::Timestamp; +use uuid::Uuid; + +use crate::{ErrorKind, ObjectStore}; + +/// Prefix of a temporary file name. +const TEMPORARY_PREFIX: &str = "."; + +/// Marker that follows the name of the target in a temporary file name. +const TEMPORARY_MARKER: &str = ".greptime-"; + +/// Suffix of a temporary file name. +const TEMPORARY_SUFFIX: &str = ".tmp"; + +/// Scheme of the only service that publishes non-append writes with a rename. +const TEMPORARY_FILE_SCHEME: &str = crate::services::HDFS_NATIVE_SCHEME; + +/// Returns true if `store` publishes non-append writes with a rename, which is +/// what leaves temporary files behind when the writer dies. +/// +/// Only the HDFS compatibility layer does that; a write to any other service +/// appears under its final name. +pub fn writes_temporary_files(store: &ObjectStore) -> bool { + store.info().scheme() == TEMPORARY_FILE_SCHEME +} + +/// What a sweep did. +/// +/// A sweep goes through its directories and files one at a time, and a failure +/// on either doesn't stop it, so it reports what it reclaimed and what it +/// couldn't. +#[derive(Debug, Clone, PartialEq, Eq, Default)] +pub struct SweepReport { + /// Temporary files deleted. + pub deleted: usize, + /// List and delete operations that failed. + pub failed: usize, +} + +impl SweepReport { + fn merge(&mut self, other: SweepReport) { + self.deleted += other.deleted; + self.failed += other.failed; + } +} + +/// Deletes the temporary files in `dirs` that are older than `older_than`. +/// +/// A non-append write goes to a temporary file next to its target and is +/// renamed on close, so a writer that dies in between leaves the file behind +/// with nothing to reclaim it: GC can't parse the name, and the startup cleanup +/// of [`ATOMIC_WRITE_DIR`](crate::ATOMIC_WRITE_DIR) doesn't reach it. +/// +/// The caller runs this for a region that just became its writable leader, so a +/// file that is old enough belongs to a writer that is gone. Failures are +/// logged and don't stop the sweep. +pub async fn sweep_temporary_files( + store: &ObjectStore, + dirs: &[String], + older_than: Duration, +) -> SweepReport { + let mut report = SweepReport::default(); + for dir in dirs { + report.merge(sweep_dir(store, dir, older_than).await); + } + report +} + +async fn sweep_dir(store: &ObjectStore, dir: &str, older_than: Duration) -> SweepReport { + let entries = match store.list(dir).await { + Ok(entries) => entries, + // Not every region has every directory, for example one that was never + // repartitioned has no staging directory. + Err(error) if error.kind() == ErrorKind::NotFound => return SweepReport::default(), + Err(error) => { + warn!(error; "Failed to list {dir} for temporary files"); + return SweepReport { + failed: 1, + ..Default::default() + }; + } + }; + + let expired_before = Timestamp::now() + .into_inner() + .as_millisecond() + .saturating_sub(i64::try_from(older_than.as_millis()).unwrap_or(i64::MAX)); + + let mut report = SweepReport::default(); + for entry in entries { + if !is_temporary_file(entry.name()) { + continue; + } + + // HDFS advances the modification time of a file under construction only + // as its blocks complete, so an old file is not necessarily abandoned. + // That is what the threshold covers. + let is_expired = entry + .metadata() + .last_modified() + .map(|last_modified| last_modified.into_inner().as_millisecond() < expired_before) + // Without a modification time there is no way to tell whether the + // writer is gone, so keep the file. + .unwrap_or(false); + if !is_expired { + continue; + } + + // Delete one file at a time: the bulk delete stops at the first + // failure, and one file that can't be deleted must not keep the rest of + // the directory from being swept. + let path = entry.path(); + match store.delete(path).await { + Ok(()) => report.deleted += 1, + Err(error) => { + warn!(error; "Failed to delete temporary file {path}"); + report.failed += 1; + } + } + } + + report +} + +/// Returns the path a non-append write to `path` goes to before it is renamed. +/// +/// A writer that dies in between leaves the file behind, which is what +/// [`sweep_temporary_files`] reclaims. +pub(crate) fn temporary_path(path: &str) -> String { + let suffix = format!("{TEMPORARY_MARKER}{}{TEMPORARY_SUFFIX}", Uuid::new_v4()); + match path.rsplit_once('/') { + Some((parent, name)) => format!("{parent}/{TEMPORARY_PREFIX}{name}{suffix}"), + None => format!("{TEMPORARY_PREFIX}{path}{suffix}"), + } +} + +fn is_temporary_file(name: &str) -> bool { + // `name` is the last segment of a path, so a directory name ends with `/`. + name.starts_with(TEMPORARY_PREFIX) + && name.contains(TEMPORARY_MARKER) + && name.ends_with(TEMPORARY_SUFFIX) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use opendal::Operator; + use opendal::services::Fs; + use tempfile::TempDir; + + use super::*; + use crate::layers::HdfsCompatibilityLayer; + use crate::layers::mock::{MockLayer, MockLayerBuilder, OpDelete, oio}; + use crate::{Error, Result}; + + const EXPIRED: Duration = Duration::from_millis(1); + + fn test_store() -> (TempDir, ObjectStore) { + let directory = tempfile::tempdir().unwrap(); + let store = Operator::new(Fs::default().root(directory.path().to_str().unwrap())).unwrap(); + (directory, store) + } + + fn test_store_of(layer: MockLayer) -> (TempDir, ObjectStore) { + let directory = tempfile::tempdir().unwrap(); + let store = Operator::new(Fs::default().root(directory.path().to_str().unwrap())) + .unwrap() + .layer(layer); + (directory, store) + } + + fn test_store_of_failing_list(kind: ErrorKind) -> (TempDir, ObjectStore) { + struct FailingLister(ErrorKind); + + impl oio::List for FailingLister { + async fn next(&mut self) -> Result> { + Err(Error::new(self.0, "Failed to list")) + } + } + + let layer = MockLayerBuilder::default() + .lister_factory(Arc::new(move |_, _, _| Box::new(FailingLister(kind)))) + .build() + .unwrap(); + test_store_of(layer) + } + + // A store that fails to delete every path containing `marker`. + fn test_store_of_failing_delete(marker: &'static str) -> (TempDir, ObjectStore) { + struct FailingDeleter { + inner: oio::Deleter, + marker: &'static str, + } + + impl oio::Delete for FailingDeleter { + async fn delete(&mut self, path: &str, args: OpDelete) -> Result<()> { + if path.contains(self.marker) { + return Err(Error::new(ErrorKind::Unexpected, "Failed to delete")); + } + self.inner.delete(path, args).await + } + + async fn close(&mut self) -> Result<()> { + self.inner.close().await + } + } + + let layer = MockLayerBuilder::default() + .deleter_factory(Arc::new(move |inner| { + Box::new(FailingDeleter { inner, marker }) + })) + .build() + .unwrap(); + test_store_of(layer) + } + + async fn wait_until_expired() { + // A file written just before this call can carry the same modification + // time in milliseconds as the sweep that follows it, which would make the + // sweep keep it. + tokio::time::sleep(Duration::from_millis(10)).await; + } + + #[test] + fn test_is_temporary_file() { + assert!(is_temporary_file( + ".00000000000000000042.json.greptime-53f0f4a6-3f3e-4a2f-9a2f-6fe4b0d5b5b0.tmp" + )); + assert!(is_temporary_file(".a5f2.parquet.greptime-53f0f4a6.tmp")); + + assert!(!is_temporary_file("00000000000000000042.json")); + assert!(!is_temporary_file(".00000000000000000042.json")); + assert!(!is_temporary_file( + "00000000000000000042.json.greptime-53f0f4a6.tmp" + )); + assert!(!is_temporary_file( + ".00000000000000000042.json.greptime-53f0f4a6" + )); + assert!(!is_temporary_file("index/")); + } + + #[tokio::test] + async fn test_sweep_dir_deletes_expired_temporary_file() { + let (_directory, store) = test_store(); + let temporary = "data/.42.parquet.greptime-53f0f4a6.tmp"; + store.create_dir("data/").await.unwrap(); + store.write(temporary, "partial").await.unwrap(); + store.write("data/42.parquet", "data").await.unwrap(); + + wait_until_expired().await; + assert_eq!( + SweepReport { + deleted: 1, + failed: 0 + }, + sweep_dir(&store, "data/", EXPIRED).await + ); + + assert!(!store.exists(temporary).await.unwrap()); + assert!(store.exists("data/42.parquet").await.unwrap()); + } + + #[tokio::test] + async fn test_sweep_dir_keeps_young_temporary_file() { + let (_directory, store) = test_store(); + let temporary = "data/.42.parquet.greptime-53f0f4a6.tmp"; + store.create_dir("data/").await.unwrap(); + store.write(temporary, "partial").await.unwrap(); + + assert_eq!( + SweepReport::default(), + sweep_dir(&store, "data/", Duration::from_secs(24 * 60 * 60)).await + ); + + assert!(store.exists(temporary).await.unwrap()); + } + + #[tokio::test] + async fn test_sweep_dir_of_missing_directory() { + // Fs lists a directory it doesn't have as an empty one, but HDFS reports + // it as not found. + let (_directory, store) = test_store_of_failing_list(ErrorKind::NotFound); + + // A region that was never repartitioned has no staging directory, so a + // missing directory is an expected answer rather than a failure. + assert_eq!( + SweepReport::default(), + sweep_dir(&store, "data/staging/blob/", Duration::ZERO).await + ); + } + + #[tokio::test] + async fn test_sweep_dir_reports_failed_list() { + let (_directory, store) = test_store_of_failing_list(ErrorKind::Unexpected); + + assert_eq!( + SweepReport { + deleted: 0, + failed: 1 + }, + sweep_dir(&store, "data/", Duration::ZERO).await + ); + } + + #[tokio::test] + async fn test_sweep_dir_keeps_going_after_failed_delete() { + let (_directory, store) = test_store_of_failing_delete("fail"); + let failing = "data/.fail-42.parquet.greptime-53f0f4a6.tmp"; + let sweepable = "data/.ok-42.parquet.greptime-53f0f4a6.tmp"; + store.create_dir("data/").await.unwrap(); + store.write(failing, "partial").await.unwrap(); + store.write(sweepable, "partial").await.unwrap(); + + wait_until_expired().await; + assert_eq!( + SweepReport { + deleted: 1, + failed: 1 + }, + sweep_dir(&store, "data/", EXPIRED).await + ); + + assert!(store.exists(failing).await.unwrap()); + assert!(!store.exists(sweepable).await.unwrap()); + } + + #[tokio::test] + async fn test_sweep_temporary_files_on_compatibility_layer() { + let (_directory, store) = test_store(); + let temporary = "data/.42.parquet.greptime-53f0f4a6.tmp"; + store.create_dir("data/").await.unwrap(); + // What a writer that died before it renamed its file leaves behind. + store.write(temporary, "partial").await.unwrap(); + // The same layer HDFS writes through, on an Fs store whose scheme stays `fs`. + let store = store.layer(HdfsCompatibilityLayer::new_for_test()); + + wait_until_expired().await; + assert_eq!( + SweepReport { + deleted: 1, + failed: 0 + }, + sweep_temporary_files( + &store, + &["data/".to_string(), "data/staging/blob/".to_string()], + EXPIRED, + ) + .await + ); + + assert!(!store.exists(temporary).await.unwrap()); + } +} From 371c4e90044036d77a4e6f55e0287560d7613136 Mon Sep 17 00:00:00 2001 From: Yingwen Date: Mon, 26 Oct 2026 07:19:00 +0000 Subject: [PATCH 091/239] feat: support query concurrency limits in standalone (#9958) Signed-off-by: evenyag --- Cargo.lock | 1 + config/config.md | 2 + config/standalone.example.toml | 8 ++++ src/cmd/src/standalone.rs | 7 +++ src/standalone/Cargo.toml | 1 + src/standalone/src/options.rs | 81 +++++++++++++++++++++++++++++++++ tests-integration/tests/http.rs | 2 + 7 files changed, 102 insertions(+) diff --git a/Cargo.lock b/Cargo.lock index 16739107f16..639a771816a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -14706,6 +14706,7 @@ dependencies = [ "flow", "frontend", "hostname 0.4.1", + "humantime-serde", "log-store", "mito2", "pipeline", diff --git a/config/config.md b/config/config.md index 36c6f02baff..f04689d97f6 100644 --- a/config/config.md +++ b/config/config.md @@ -20,6 +20,8 @@ | `write_bytes_exhausted_policy` | String | `wait` | Policy when write bytes quota is exhausted.
Options: "wait" (default, 10s timeout), "wait()" (e.g., "wait(30s)"), "fail" | | `init_regions_in_background` | Bool | `false` | Initialize all regions in the background during the startup.
By default, it provides services after all regions have been initialized. | | `init_regions_parallelism` | Integer | `16` | Parallelism of initializing regions. | +| `max_concurrent_queries` | Integer | `0` | The maximum concurrent region-read streams. Zero means unlimited.
A single query can consume multiple permits. | +| `concurrent_query_limiter_timeout` | String | `100ms` | Timeout to acquire a permit when `max_concurrent_queries` is reached.
Only effective when the limiter is enabled. | | `enable_telemetry` | Bool | `true` | Enable telemetry to collect anonymous usage data. Enabled by default. | | `runtime` | -- | -- | The runtime options. | | `runtime.global_rt_size` | Integer | Auto | The number of threads to execute the runtime for global read operations.
Defaults to max(num_cpus, 2). | diff --git a/config/standalone.example.toml b/config/standalone.example.toml index f73a3f8a325..f569ff1b952 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -43,6 +43,14 @@ init_regions_in_background = false ## Parallelism of initializing regions. init_regions_parallelism = 16 +## The maximum concurrent region-read streams. Zero means unlimited. +## A single query can consume multiple permits. +max_concurrent_queries = 0 + +## Timeout to acquire a permit when `max_concurrent_queries` is reached. +## Only effective when the limiter is enabled. +concurrent_query_limiter_timeout = "100ms" + ## Enable telemetry to collect anonymous usage data. Enabled by default. #+ enable_telemetry = true diff --git a/src/cmd/src/standalone.rs b/src/cmd/src/standalone.rs index 3dca89f6c25..aff0b62181f 100644 --- a/src/cmd/src/standalone.rs +++ b/src/cmd/src/standalone.rs @@ -1324,6 +1324,8 @@ mod tests { let mut file = create_named_temp_file(); let toml_str = r#" enable_memory_catalog = true + max_concurrent_queries = 8 + concurrent_query_limiter_timeout = "250ms" [wal] provider = "raft_engine" @@ -1382,6 +1384,11 @@ mod tests { .component; let fe_opts = options.frontend_options(); let dn_opts = options.datanode_options(); + assert_eq!(dn_opts.max_concurrent_queries, 8); + assert_eq!( + dn_opts.concurrent_query_limiter_timeout, + Duration::from_millis(250) + ); let logging_opts = options.logging; assert_eq!("127.0.0.1:4000".to_string(), fe_opts.http.addr); assert_eq!(Duration::from_secs(33), fe_opts.http.timeout); diff --git a/src/standalone/Cargo.toml b/src/standalone/Cargo.toml index d89366d2e7e..60bd8f7991c 100644 --- a/src/standalone/Cargo.toml +++ b/src/standalone/Cargo.toml @@ -33,6 +33,7 @@ file-engine.workspace = true flow.workspace = true frontend.workspace = true hostname.workspace = true +humantime-serde.workspace = true log-store.workspace = true mito2.workspace = true pipeline.workspace = true diff --git a/src/standalone/src/options.rs b/src/standalone/src/options.rs index bbbad14b93d..211f21ca1e6 100644 --- a/src/standalone/src/options.rs +++ b/src/standalone/src/options.rs @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::time::Duration; + use common_base::readable_size::ReadableSize; use common_config::{Configurable, KvBackendConfig}; use common_event_recorder::EventRecorderOptions; @@ -73,6 +75,13 @@ pub struct StandaloneOptions { pub tracing: TracingOptions, pub init_regions_in_background: bool, pub init_regions_parallelism: usize, + /// Maximum concurrent region-read streams. Zero means unlimited. + /// A single query can consume multiple permits. + pub max_concurrent_queries: usize, + /// Timeout to acquire a permit when `max_concurrent_queries` is reached. + /// Only effective when the limiter is enabled. + #[serde(with = "humantime_serde")] + pub concurrent_query_limiter_timeout: Duration, pub slow_query: SlowQueryOptions, pub query: QueryOptions, pub memory: MemoryOptions, @@ -117,6 +126,8 @@ impl Default for StandaloneOptions { tracing: TracingOptions::default(), init_regions_in_background: false, init_regions_parallelism: 16, + max_concurrent_queries: 0, + concurrent_query_limiter_timeout: Duration::from_millis(100), slow_query: SlowQueryOptions::default(), query: QueryOptions::default(), memory: MemoryOptions::default(), @@ -199,6 +210,8 @@ impl StandaloneOptions { grpc: cloned_opts.grpc, init_regions_in_background: cloned_opts.init_regions_in_background, init_regions_parallelism: cloned_opts.init_regions_parallelism, + max_concurrent_queries: cloned_opts.max_concurrent_queries, + concurrent_query_limiter_timeout: cloned_opts.concurrent_query_limiter_timeout, query: cloned_opts.query, heartbeat_env_vars: cloned_opts.heartbeat_env_vars, ..Default::default() @@ -226,6 +239,74 @@ mod tests { use crate::options::*; + #[test] + fn test_concurrent_query_limiter_options() { + let defaults: StandaloneOptions = toml::from_str("").unwrap(); + let datanode_defaults = DatanodeOptions::default(); + assert_eq!(defaults.max_concurrent_queries, 0); + assert_eq!( + defaults.concurrent_query_limiter_timeout, + Duration::from_millis(100) + ); + assert_eq!( + defaults.max_concurrent_queries, + datanode_defaults.max_concurrent_queries + ); + assert_eq!( + defaults.concurrent_query_limiter_timeout, + datanode_defaults.concurrent_query_limiter_timeout + ); + + let options: StandaloneOptions = toml::from_str( + r#"max_concurrent_queries = 8 +concurrent_query_limiter_timeout = "250ms""#, + ) + .unwrap(); + let restored: StandaloneOptions = + toml::from_str(&toml::to_string(&options).unwrap()).unwrap(); + for options in [options, restored] { + assert_eq!(options.max_concurrent_queries, 8); + assert_eq!( + options.concurrent_query_limiter_timeout, + Duration::from_millis(250) + ); + let datanode = options.datanode_options(); + assert_eq!(datanode.max_concurrent_queries, 8); + assert_eq!( + datanode.concurrent_query_limiter_timeout, + Duration::from_millis(250) + ); + } + assert!( + toml::from_str::("concurrent_query_limiter_timeout = 'invalid'") + .is_err() + ); + } + + #[test] + fn test_concurrent_query_limiter_options_from_env() { + temp_env::with_vars( + [ + ("STANDALONE_LIMITER_TEST__MAX_CONCURRENT_QUERIES", Some("4")), + ( + "STANDALONE_LIMITER_TEST__CONCURRENT_QUERY_LIMITER_TIMEOUT", + Some("2s"), + ), + ], + || { + let options = + StandaloneOptions::load_layered_options(None, "STANDALONE_LIMITER_TEST") + .unwrap(); + let datanode = options.datanode_options(); + assert_eq!(datanode.max_concurrent_queries, 4); + assert_eq!( + datanode.concurrent_query_limiter_timeout, + Duration::from_secs(2) + ); + }, + ); + } + #[test] fn test_pipeline_host_access_config_forwarding() { let defaults: StandaloneOptions = toml::from_str("").unwrap(); diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 7ea4d3df2ea..1671f5a2921 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -3418,6 +3418,8 @@ max_in_flight_write_bytes = "0KiB" write_bytes_exhausted_policy = "wait" init_regions_in_background = false init_regions_parallelism = 16 +max_concurrent_queries = 0 +concurrent_query_limiter_timeout = "100ms" heartbeat_env_vars = [] [http] From e87573e2865607b277775c48bfe98581719cc39a Mon Sep 17 00:00:00 2001 From: "B. Naut" <2025llxe@gmail.com> Date: Mon, 26 Oct 2026 07:22:04 +0000 Subject: [PATCH 092/239] test(sqlness): cover query options in merge-scan regressions (#9956) * test(sqlness): cover query options on persistent SQL sessions Signed-off-by: Byte-Naut * test(query): cover merge-scan joins with query option controls Signed-off-by: Byte-Naut --------- Signed-off-by: Byte-Naut --- .../common/join/merge_scan_semi_join.result | 331 ++++++++++++++++++ .../common/join/merge_scan_semi_join.sql | 103 ++++++ .../standalone/common/query_options.result | 164 +++++++++ .../cases/standalone/common/query_options.sql | 70 ++++ 4 files changed, 668 insertions(+) create mode 100644 tests/cases/standalone/common/query_options.result create mode 100644 tests/cases/standalone/common/query_options.sql diff --git a/tests/cases/standalone/common/join/merge_scan_semi_join.result b/tests/cases/standalone/common/join/merge_scan_semi_join.result index 5eb70b1241b..598e2099662 100644 --- a/tests/cases/standalone/common/join/merge_scan_semi_join.result +++ b/tests/cases/standalone/common/join/merge_scan_semi_join.result @@ -98,6 +98,337 @@ GROUP BY key_id ORDER BY key_id; | 33 | 2026-10-09T00:00:00 | +--------+---------------------+ +-- Reproduce issue 9492 at parallelism 2 with remote dynamic filters enabled. +-- SQLNESS PROTOCOL MYSQL +SET query.parallelism = 2; + +affected_rows: 0 + +-- SQLNESS PROTOCOL MYSQL +SHOW VARIABLES query.parallelism; + ++-------------------+ +| query.parallelism | ++-------------------+ +| 2 | ++-------------------+ + +-- SQLNESS PROTOCOL MYSQL +SHOW VARIABLES query.enable_remote_dynamic_filter_pushdown; + ++---------------------------------------------+ +| query.enable_remote_dynamic_filter_pushdown | ++---------------------------------------------+ +| true | ++---------------------------------------------+ + +-- SQLNESS REPLACE (peers.*) peers=REDACTED +-- SQLNESS PROTOCOL MYSQL +EXPLAIN SELECT key_id, MIN(ts) AS first_ts +FROM merge_semi_data +WHERE key_id IN ( + SELECT key_id FROM merge_semi_data WHERE is_match = true +) +GROUP BY key_id ORDER BY key_id; + ++---------------+---------------------------------------------------------------------------------------------------------------------------------------+ +| plan_type | plan | ++---------------+---------------------------------------------------------------------------------------------------------------------------------------+ +| logical_plan | Sort: merge_semi_data.key_id ASC NULLS LAST | +| | Projection: merge_semi_data.key_id, min(merge_semi_data.ts) AS first_ts | +| | Aggregate: groupBy=[[merge_semi_data.key_id]], aggr=[[min(merge_semi_data.ts)]] | +| | LeftSemi Join: merge_semi_data.key_id = __correlated_sq_1.key_id | +| | Projection: merge_semi_data.ts, merge_semi_data.key_id | +| | MergeScan [is_placeholder=false, remote_input=[ | +| | Filter: merge_semi_data.key_id IS NOT NULL | +| | TableScan: merge_semi_data, partial_filters=[merge_semi_data.key_id IS NOT NULL] | +| | ]] | +| | MergeScan [is_placeholder=false, remote_input=[ | +| | SubqueryAlias: __correlated_sq_1 | +| | Projection: merge_semi_data.key_id | +| | Filter: merge_semi_data.is_match AND merge_semi_data.key_id IS NOT NULL | +| | TableScan: merge_semi_data, partial_filters=[merge_semi_data.is_match, merge_semi_data.key_id IS NOT NULL] | +| | ]] | +| physical_plan | SortPreservingMergeExec: [key_id@0 ASC NULLS LAST] | +| | ProjectionExec: expr=[key_id@0 as key_id, min(merge_semi_data.ts)@1 as first_ts] | +| | SortExec: expr=[key_id@0 ASC NULLS LAST], preserve_partitioning=[true] | +| | AggregateExec: mode=SinglePartitioned, gby=[key_id@1 as key_id], aggr=[min(merge_semi_data.ts)] | +| | HashJoinExec: mode=Partitioned, join_type=LeftSemi, on=[(key_id@1, key_id@0)] | +| | RepartitionExec: partitioning=Hash([key_id@1], 2), input_partitions=2 | +| | ProjectionExec: expr=[ts@0 as ts, key_id@1 as key_id] | +| | MergeScanExec: peers=REDACTED +| | RepartitionExec: partitioning=Hash([key_id@0], 2), input_partitions=2 | +| | MergeScanExec: peers=REDACTED +| | | ++---------------+---------------------------------------------------------------------------------------------------------------------------------------+ + +-- SQLNESS PROTOCOL MYSQL +SELECT key_id, MIN(ts) AS first_ts +FROM merge_semi_data +WHERE key_id IN ( + SELECT key_id FROM merge_semi_data WHERE is_match = true +) +GROUP BY key_id ORDER BY key_id; + ++--------+---------------------+ +| key_id | first_ts | ++--------+---------------------+ +| 1 | 2026-10-09 00:00:00 | +| 2 | 2026-10-09 00:00:00 | +| 3 | 2026-10-09 00:00:00 | +| 11 | 2026-10-09 00:00:00 | +| 12 | 2026-10-09 00:00:00 | +| 13 | 2026-10-09 00:00:00 | +| 21 | 2026-10-09 00:00:00 | +| 22 | 2026-10-09 00:00:00 | +| 23 | 2026-10-09 00:00:00 | +| 31 | 2026-10-09 00:00:00 | +| 32 | 2026-10-09 00:00:00 | +| 33 | 2026-10-09 00:00:00 | ++--------+---------------------+ + +-- Use one partition as the result control. +-- SQLNESS PROTOCOL MYSQL +SET query.parallelism = 1; + +affected_rows: 0 + +-- SQLNESS PROTOCOL MYSQL +SHOW VARIABLES query.parallelism; + ++-------------------+ +| query.parallelism | ++-------------------+ +| 1 | ++-------------------+ + +-- SQLNESS PROTOCOL MYSQL +SHOW VARIABLES query.enable_remote_dynamic_filter_pushdown; + ++---------------------------------------------+ +| query.enable_remote_dynamic_filter_pushdown | ++---------------------------------------------+ +| true | ++---------------------------------------------+ + +-- SQLNESS REPLACE (peers.*) peers=REDACTED +-- SQLNESS PROTOCOL MYSQL +EXPLAIN SELECT key_id, MIN(ts) AS first_ts +FROM merge_semi_data +WHERE key_id IN ( + SELECT key_id FROM merge_semi_data WHERE is_match = true +) +GROUP BY key_id ORDER BY key_id; + ++---------------+-------------------------------------------------------------------------------------------------------------------------------------+ +| plan_type | plan | ++---------------+-------------------------------------------------------------------------------------------------------------------------------------+ +| logical_plan | Sort: merge_semi_data.key_id ASC NULLS LAST | +| | Projection: merge_semi_data.key_id, min(merge_semi_data.ts) AS first_ts | +| | Aggregate: groupBy=[[merge_semi_data.key_id]], aggr=[[min(merge_semi_data.ts)]] | +| | LeftSemi Join: merge_semi_data.key_id = __correlated_sq_1.key_id | +| | Projection: merge_semi_data.ts, merge_semi_data.key_id | +| | MergeScan [is_placeholder=false, remote_input=[ | +| | Filter: merge_semi_data.key_id IS NOT NULL | +| | TableScan: merge_semi_data, partial_filters=[merge_semi_data.key_id IS NOT NULL] | +| | ]] | +| | MergeScan [is_placeholder=false, remote_input=[ | +| | SubqueryAlias: __correlated_sq_1 | +| | Projection: merge_semi_data.key_id | +| | Filter: merge_semi_data.is_match AND merge_semi_data.key_id IS NOT NULL | +| | TableScan: merge_semi_data, partial_filters=[merge_semi_data.is_match, merge_semi_data.key_id IS NOT NULL] | +| | ]] | +| physical_plan | ProjectionExec: expr=[key_id@0 as key_id, min(merge_semi_data.ts)@1 as first_ts] | +| | SortExec: expr=[key_id@0 ASC NULLS LAST], preserve_partitioning=[false] | +| | AggregateExec: mode=Single, gby=[key_id@1 as key_id], aggr=[min(merge_semi_data.ts)] | +| | HashJoinExec: mode=CollectLeft, join_type=LeftSemi, on=[(key_id@1, key_id@0)] | +| | ProjectionExec: expr=[ts@0 as ts, key_id@1 as key_id] | +| | CooperativeExec | +| | MergeScanExec: peers=REDACTED +| | CooperativeExec | +| | MergeScanExec: peers=REDACTED +| | | ++---------------+-------------------------------------------------------------------------------------------------------------------------------------+ + +-- SQLNESS PROTOCOL MYSQL +SELECT key_id, MIN(ts) AS first_ts +FROM merge_semi_data +WHERE key_id IN ( + SELECT key_id FROM merge_semi_data WHERE is_match = true +) +GROUP BY key_id ORDER BY key_id; + ++--------+---------------------+ +| key_id | first_ts | ++--------+---------------------+ +| 1 | 2026-10-09 00:00:00 | +| 2 | 2026-10-09 00:00:00 | +| 3 | 2026-10-09 00:00:00 | +| 11 | 2026-10-09 00:00:00 | +| 12 | 2026-10-09 00:00:00 | +| 13 | 2026-10-09 00:00:00 | +| 21 | 2026-10-09 00:00:00 | +| 22 | 2026-10-09 00:00:00 | +| 23 | 2026-10-09 00:00:00 | +| 31 | 2026-10-09 00:00:00 | +| 32 | 2026-10-09 00:00:00 | +| 33 | 2026-10-09 00:00:00 | ++--------+---------------------+ + +-- Keep two partitions and disable remote dynamic filters as the second control. +-- SQLNESS PROTOCOL MYSQL +SET query.parallelism = 2; + +affected_rows: 0 + +-- SQLNESS PROTOCOL MYSQL +SET query.enable_remote_dynamic_filter_pushdown = false; + +affected_rows: 0 + +-- SQLNESS PROTOCOL MYSQL +SHOW VARIABLES query.parallelism; + ++-------------------+ +| query.parallelism | ++-------------------+ +| 2 | ++-------------------+ + +-- SQLNESS PROTOCOL MYSQL +SHOW VARIABLES query.enable_remote_dynamic_filter_pushdown; + ++---------------------------------------------+ +| query.enable_remote_dynamic_filter_pushdown | ++---------------------------------------------+ +| false | ++---------------------------------------------+ + +-- SQLNESS REPLACE (peers.*) peers=REDACTED +-- SQLNESS PROTOCOL MYSQL +EXPLAIN SELECT key_id, MIN(ts) AS first_ts +FROM merge_semi_data +WHERE key_id IN ( + SELECT key_id FROM merge_semi_data WHERE is_match = true +) +GROUP BY key_id ORDER BY key_id; + ++---------------+---------------------------------------------------------------------------------------------------------------------------------------+ +| plan_type | plan | ++---------------+---------------------------------------------------------------------------------------------------------------------------------------+ +| logical_plan | Sort: merge_semi_data.key_id ASC NULLS LAST | +| | Projection: merge_semi_data.key_id, min(merge_semi_data.ts) AS first_ts | +| | Aggregate: groupBy=[[merge_semi_data.key_id]], aggr=[[min(merge_semi_data.ts)]] | +| | LeftSemi Join: merge_semi_data.key_id = __correlated_sq_1.key_id | +| | Projection: merge_semi_data.ts, merge_semi_data.key_id | +| | MergeScan [is_placeholder=false, remote_input=[ | +| | Filter: merge_semi_data.key_id IS NOT NULL | +| | TableScan: merge_semi_data, partial_filters=[merge_semi_data.key_id IS NOT NULL] | +| | ]] | +| | MergeScan [is_placeholder=false, remote_input=[ | +| | SubqueryAlias: __correlated_sq_1 | +| | Projection: merge_semi_data.key_id | +| | Filter: merge_semi_data.is_match AND merge_semi_data.key_id IS NOT NULL | +| | TableScan: merge_semi_data, partial_filters=[merge_semi_data.is_match, merge_semi_data.key_id IS NOT NULL] | +| | ]] | +| physical_plan | SortPreservingMergeExec: [key_id@0 ASC NULLS LAST] | +| | ProjectionExec: expr=[key_id@0 as key_id, min(merge_semi_data.ts)@1 as first_ts] | +| | SortExec: expr=[key_id@0 ASC NULLS LAST], preserve_partitioning=[true] | +| | AggregateExec: mode=SinglePartitioned, gby=[key_id@1 as key_id], aggr=[min(merge_semi_data.ts)] | +| | HashJoinExec: mode=Partitioned, join_type=LeftSemi, on=[(key_id@1, key_id@0)] | +| | RepartitionExec: partitioning=Hash([key_id@1], 2), input_partitions=2 | +| | ProjectionExec: expr=[ts@0 as ts, key_id@1 as key_id] | +| | MergeScanExec: peers=REDACTED +| | RepartitionExec: partitioning=Hash([key_id@0], 2), input_partitions=2 | +| | MergeScanExec: peers=REDACTED +| | | ++---------------+---------------------------------------------------------------------------------------------------------------------------------------+ + +-- SQLNESS PROTOCOL MYSQL +SELECT key_id, MIN(ts) AS first_ts +FROM merge_semi_data +WHERE key_id IN ( + SELECT key_id FROM merge_semi_data WHERE is_match = true +) +GROUP BY key_id ORDER BY key_id; + ++--------+---------------------+ +| key_id | first_ts | ++--------+---------------------+ +| 1 | 2026-10-09 00:00:00 | +| 2 | 2026-10-09 00:00:00 | +| 3 | 2026-10-09 00:00:00 | +| 11 | 2026-10-09 00:00:00 | +| 12 | 2026-10-09 00:00:00 | +| 13 | 2026-10-09 00:00:00 | +| 21 | 2026-10-09 00:00:00 | +| 22 | 2026-10-09 00:00:00 | +| 23 | 2026-10-09 00:00:00 | +| 31 | 2026-10-09 00:00:00 | +| 32 | 2026-10-09 00:00:00 | +| 33 | 2026-10-09 00:00:00 | ++--------+---------------------+ + +-- Restore SQLness defaults on the shared session and check its next plan. +-- SQLNESS PROTOCOL MYSQL +SET query.parallelism = 4; + +affected_rows: 0 + +-- SQLNESS PROTOCOL MYSQL +SET query.enable_remote_dynamic_filter_pushdown = true; + +affected_rows: 0 + +-- SQLNESS PROTOCOL MYSQL +SHOW VARIABLES query.enable_remote_dynamic_filter_pushdown; + ++---------------------------------------------+ +| query.enable_remote_dynamic_filter_pushdown | ++---------------------------------------------+ +| true | ++---------------------------------------------+ + +-- SQLNESS REPLACE (peers.*) peers=REDACTED +-- SQLNESS PROTOCOL MYSQL +EXPLAIN SELECT key_id, MIN(ts) AS first_ts +FROM merge_semi_data +WHERE key_id IN ( + SELECT key_id FROM merge_semi_data WHERE is_match = true +) +GROUP BY key_id ORDER BY key_id; + ++---------------+---------------------------------------------------------------------------------------------------------------------------------------+ +| plan_type | plan | ++---------------+---------------------------------------------------------------------------------------------------------------------------------------+ +| logical_plan | Sort: merge_semi_data.key_id ASC NULLS LAST | +| | Projection: merge_semi_data.key_id, min(merge_semi_data.ts) AS first_ts | +| | Aggregate: groupBy=[[merge_semi_data.key_id]], aggr=[[min(merge_semi_data.ts)]] | +| | LeftSemi Join: merge_semi_data.key_id = __correlated_sq_1.key_id | +| | Projection: merge_semi_data.ts, merge_semi_data.key_id | +| | MergeScan [is_placeholder=false, remote_input=[ | +| | Filter: merge_semi_data.key_id IS NOT NULL | +| | TableScan: merge_semi_data, partial_filters=[merge_semi_data.key_id IS NOT NULL] | +| | ]] | +| | MergeScan [is_placeholder=false, remote_input=[ | +| | SubqueryAlias: __correlated_sq_1 | +| | Projection: merge_semi_data.key_id | +| | Filter: merge_semi_data.is_match AND merge_semi_data.key_id IS NOT NULL | +| | TableScan: merge_semi_data, partial_filters=[merge_semi_data.is_match, merge_semi_data.key_id IS NOT NULL] | +| | ]] | +| physical_plan | SortPreservingMergeExec: [key_id@0 ASC NULLS LAST] | +| | ProjectionExec: expr=[key_id@0 as key_id, min(merge_semi_data.ts)@1 as first_ts] | +| | SortExec: expr=[key_id@0 ASC NULLS LAST], preserve_partitioning=[true] | +| | AggregateExec: mode=SinglePartitioned, gby=[key_id@1 as key_id], aggr=[min(merge_semi_data.ts)] | +| | HashJoinExec: mode=Partitioned, join_type=LeftSemi, on=[(key_id@1, key_id@0)] | +| | RepartitionExec: partitioning=Hash([key_id@1], 4), input_partitions=4 | +| | ProjectionExec: expr=[ts@0 as ts, key_id@1 as key_id] | +| | MergeScanExec: peers=REDACTED +| | RepartitionExec: partitioning=Hash([key_id@0], 4), input_partitions=4 | +| | MergeScanExec: peers=REDACTED +| | | ++---------------+---------------------------------------------------------------------------------------------------------------------------------------+ + DROP TABLE merge_semi_data; Affected Rows: 0 diff --git a/tests/cases/standalone/common/join/merge_scan_semi_join.sql b/tests/cases/standalone/common/join/merge_scan_semi_join.sql index 34e7b226592..6c3ad12baa1 100644 --- a/tests/cases/standalone/common/join/merge_scan_semi_join.sql +++ b/tests/cases/standalone/common/join/merge_scan_semi_join.sql @@ -46,4 +46,107 @@ WHERE key_id IN ( ) GROUP BY key_id ORDER BY key_id; +-- Reproduce issue 9492 at parallelism 2 with remote dynamic filters enabled. +-- SQLNESS PROTOCOL MYSQL +SET query.parallelism = 2; + +-- SQLNESS PROTOCOL MYSQL +SHOW VARIABLES query.parallelism; + +-- SQLNESS PROTOCOL MYSQL +SHOW VARIABLES query.enable_remote_dynamic_filter_pushdown; + +-- SQLNESS REPLACE (peers.*) peers=REDACTED +-- SQLNESS PROTOCOL MYSQL +EXPLAIN SELECT key_id, MIN(ts) AS first_ts +FROM merge_semi_data +WHERE key_id IN ( + SELECT key_id FROM merge_semi_data WHERE is_match = true +) +GROUP BY key_id ORDER BY key_id; + +-- SQLNESS PROTOCOL MYSQL +SELECT key_id, MIN(ts) AS first_ts +FROM merge_semi_data +WHERE key_id IN ( + SELECT key_id FROM merge_semi_data WHERE is_match = true +) +GROUP BY key_id ORDER BY key_id; + +-- Use one partition as the result control. +-- SQLNESS PROTOCOL MYSQL +SET query.parallelism = 1; + +-- SQLNESS PROTOCOL MYSQL +SHOW VARIABLES query.parallelism; + +-- SQLNESS PROTOCOL MYSQL +SHOW VARIABLES query.enable_remote_dynamic_filter_pushdown; + +-- SQLNESS REPLACE (peers.*) peers=REDACTED +-- SQLNESS PROTOCOL MYSQL +EXPLAIN SELECT key_id, MIN(ts) AS first_ts +FROM merge_semi_data +WHERE key_id IN ( + SELECT key_id FROM merge_semi_data WHERE is_match = true +) +GROUP BY key_id ORDER BY key_id; + +-- SQLNESS PROTOCOL MYSQL +SELECT key_id, MIN(ts) AS first_ts +FROM merge_semi_data +WHERE key_id IN ( + SELECT key_id FROM merge_semi_data WHERE is_match = true +) +GROUP BY key_id ORDER BY key_id; + +-- Keep two partitions and disable remote dynamic filters as the second control. +-- SQLNESS PROTOCOL MYSQL +SET query.parallelism = 2; + +-- SQLNESS PROTOCOL MYSQL +SET query.enable_remote_dynamic_filter_pushdown = false; + +-- SQLNESS PROTOCOL MYSQL +SHOW VARIABLES query.parallelism; + +-- SQLNESS PROTOCOL MYSQL +SHOW VARIABLES query.enable_remote_dynamic_filter_pushdown; + +-- SQLNESS REPLACE (peers.*) peers=REDACTED +-- SQLNESS PROTOCOL MYSQL +EXPLAIN SELECT key_id, MIN(ts) AS first_ts +FROM merge_semi_data +WHERE key_id IN ( + SELECT key_id FROM merge_semi_data WHERE is_match = true +) +GROUP BY key_id ORDER BY key_id; + +-- SQLNESS PROTOCOL MYSQL +SELECT key_id, MIN(ts) AS first_ts +FROM merge_semi_data +WHERE key_id IN ( + SELECT key_id FROM merge_semi_data WHERE is_match = true +) +GROUP BY key_id ORDER BY key_id; + +-- Restore SQLness defaults on the shared session and check its next plan. +-- SQLNESS PROTOCOL MYSQL +SET query.parallelism = 4; + +-- SQLNESS PROTOCOL MYSQL +SET query.enable_remote_dynamic_filter_pushdown = true; + +-- SQLNESS PROTOCOL MYSQL +SHOW VARIABLES query.enable_remote_dynamic_filter_pushdown; + +-- SQLNESS REPLACE (peers.*) peers=REDACTED +-- SQLNESS PROTOCOL MYSQL +EXPLAIN SELECT key_id, MIN(ts) AS first_ts +FROM merge_semi_data +WHERE key_id IN ( + SELECT key_id FROM merge_semi_data WHERE is_match = true +) +GROUP BY key_id ORDER BY key_id; + DROP TABLE merge_semi_data; diff --git a/tests/cases/standalone/common/query_options.result b/tests/cases/standalone/common/query_options.result new file mode 100644 index 00000000000..7fe377944a1 --- /dev/null +++ b/tests/cases/standalone/common/query_options.result @@ -0,0 +1,164 @@ +-- Query options persist on each protocol's connection. SQLness pins N=4. +-- SQLNESS PROTOCOL MYSQL +SHOW VARIABLES query.parallelism; + ++-------------------+ +| query.parallelism | ++-------------------+ +| 4 | ++-------------------+ + +-- SQLNESS PROTOCOL MYSQL +SET query_parallelism = 2; + +affected_rows: 0 + +-- SQLNESS PROTOCOL MYSQL +SHOW VARIABLES query.parallelism; + ++-------------------+ +| query.parallelism | ++-------------------+ +| 2 | ++-------------------+ + +-- Invalid SET preserves the preceding value. +-- SQLNESS PROTOCOL MYSQL +SET query.parallelism = 0; + +Failed to execute query, err: MySqlError { ERROR 1210 (HY000): (InvalidArguments): Invalid value for parameter "query.parallelism": 0 +Hint: query.parallelism must be between 1 and 1024 } + +-- SQLNESS PROTOCOL MYSQL +SHOW VARIABLES query.parallelism; + ++-------------------+ +| query.parallelism | ++-------------------+ +| 2 | ++-------------------+ + +-- SQLNESS PROTOCOL MYSQL +SET query.enable_remote_dynamic_filter_pushdown = false; + +affected_rows: 0 + +-- SQLNESS PROTOCOL MYSQL +SHOW VARIABLES query.enable_remote_dynamic_filter_pushdown; + ++---------------------------------------------+ +| query.enable_remote_dynamic_filter_pushdown | ++---------------------------------------------+ +| false | ++---------------------------------------------+ + +-- Restore the shared MySQL connection before another case uses it. +-- SQLNESS PROTOCOL MYSQL +SET query.parallelism = 4; + +affected_rows: 0 + +-- SQLNESS PROTOCOL MYSQL +SET query.enable_remote_dynamic_filter_pushdown = true; + +affected_rows: 0 + +-- SQLNESS PROTOCOL MYSQL +SHOW VARIABLES query.parallelism; + ++-------------------+ +| query.parallelism | ++-------------------+ +| 4 | ++-------------------+ + +-- SQLNESS PROTOCOL MYSQL +SHOW VARIABLES query.enable_remote_dynamic_filter_pushdown; + ++---------------------------------------------+ +| query.enable_remote_dynamic_filter_pushdown | ++---------------------------------------------+ +| true | ++---------------------------------------------+ + +-- PostgreSQL owns a separate session and accepts the same query options. +-- SQLNESS PROTOCOL POSTGRES +SHOW VARIABLES query.parallelism; + ++-------------------+ +| query.parallelism | ++-------------------+ +| 4 | ++-------------------+ + +-- SQLNESS PROTOCOL POSTGRES +SET query.parallelism TO 2; + +Affected Rows: 0 + +-- SQLNESS PROTOCOL POSTGRES +SHOW VARIABLES query.parallelism; + ++-------------------+ +| query.parallelism | ++-------------------+ +| 2 | ++-------------------+ + +-- SQLNESS PROTOCOL POSTGRES +SET query.parallelism = 0; + +Failed to execute query, encountered: Error { kind: Db, cause: Some(DbError { severity: "ERROR", parsed_severity: None, code: SqlState(E22023), message: "Invalid value for parameter \"query.parallelism\": 0\nHint: query.parallelism must be between 1 and 1024", detail: None, hint: None, position: None, where_: None, schema: None, table: None, column: None, datatype: None, constraint: None, file: None, line: None, routine: None }) } + +-- SQLNESS PROTOCOL POSTGRES +SHOW VARIABLES query.parallelism; + ++-------------------+ +| query.parallelism | ++-------------------+ +| 2 | ++-------------------+ + +-- SQLNESS PROTOCOL POSTGRES +SET query.enable_remote_dynamic_filter_pushdown = false; + +Affected Rows: 0 + +-- SQLNESS PROTOCOL POSTGRES +SHOW VARIABLES query.enable_remote_dynamic_filter_pushdown; + ++---------------------------------------------+ +| query.enable_remote_dynamic_filter_pushdown | ++---------------------------------------------+ +| false | ++---------------------------------------------+ + +-- Restore the shared PostgreSQL connection as well. +-- SQLNESS PROTOCOL POSTGRES +SET query.parallelism = 4; + +Affected Rows: 0 + +-- SQLNESS PROTOCOL POSTGRES +SET query.enable_remote_dynamic_filter_pushdown = true; + +Affected Rows: 0 + +-- SQLNESS PROTOCOL POSTGRES +SHOW VARIABLES query.parallelism; + ++-------------------+ +| query.parallelism | ++-------------------+ +| 4 | ++-------------------+ + +-- SQLNESS PROTOCOL POSTGRES +SHOW VARIABLES query.enable_remote_dynamic_filter_pushdown; + ++---------------------------------------------+ +| query.enable_remote_dynamic_filter_pushdown | ++---------------------------------------------+ +| true | ++---------------------------------------------+ + diff --git a/tests/cases/standalone/common/query_options.sql b/tests/cases/standalone/common/query_options.sql new file mode 100644 index 00000000000..b2578072632 --- /dev/null +++ b/tests/cases/standalone/common/query_options.sql @@ -0,0 +1,70 @@ +-- Query options persist on each protocol's connection. SQLness pins N=4. +-- SQLNESS PROTOCOL MYSQL +SHOW VARIABLES query.parallelism; + +-- SQLNESS PROTOCOL MYSQL +SET query_parallelism = 2; + +-- SQLNESS PROTOCOL MYSQL +SHOW VARIABLES query.parallelism; + +-- Invalid SET preserves the preceding value. +-- SQLNESS PROTOCOL MYSQL +SET query.parallelism = 0; + +-- SQLNESS PROTOCOL MYSQL +SHOW VARIABLES query.parallelism; + +-- SQLNESS PROTOCOL MYSQL +SET query.enable_remote_dynamic_filter_pushdown = false; + +-- SQLNESS PROTOCOL MYSQL +SHOW VARIABLES query.enable_remote_dynamic_filter_pushdown; + +-- Restore the shared MySQL connection before another case uses it. +-- SQLNESS PROTOCOL MYSQL +SET query.parallelism = 4; + +-- SQLNESS PROTOCOL MYSQL +SET query.enable_remote_dynamic_filter_pushdown = true; + +-- SQLNESS PROTOCOL MYSQL +SHOW VARIABLES query.parallelism; + +-- SQLNESS PROTOCOL MYSQL +SHOW VARIABLES query.enable_remote_dynamic_filter_pushdown; + +-- PostgreSQL owns a separate session and accepts the same query options. +-- SQLNESS PROTOCOL POSTGRES +SHOW VARIABLES query.parallelism; + +-- SQLNESS PROTOCOL POSTGRES +SET query.parallelism TO 2; + +-- SQLNESS PROTOCOL POSTGRES +SHOW VARIABLES query.parallelism; + +-- SQLNESS PROTOCOL POSTGRES +SET query.parallelism = 0; + +-- SQLNESS PROTOCOL POSTGRES +SHOW VARIABLES query.parallelism; + +-- SQLNESS PROTOCOL POSTGRES +SET query.enable_remote_dynamic_filter_pushdown = false; + +-- SQLNESS PROTOCOL POSTGRES +SHOW VARIABLES query.enable_remote_dynamic_filter_pushdown; + +-- Restore the shared PostgreSQL connection as well. +-- SQLNESS PROTOCOL POSTGRES +SET query.parallelism = 4; + +-- SQLNESS PROTOCOL POSTGRES +SET query.enable_remote_dynamic_filter_pushdown = true; + +-- SQLNESS PROTOCOL POSTGRES +SHOW VARIABLES query.parallelism; + +-- SQLNESS PROTOCOL POSTGRES +SHOW VARIABLES query.enable_remote_dynamic_filter_pushdown; From 513b9cee58590e26ea4a23747df22140eeceb11d Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Mon, 26 Oct 2026 07:23:30 +0000 Subject: [PATCH 093/239] feat(catalog): add public-global information schema scope (#9957) Signed-off-by: shuiyisong --- .../src/system_schema/information_schema.rs | 106 +++++++++++------- src/catalog/src/system_schema/memory_table.rs | 14 ++- src/table/src/requests/semantic.rs | 11 +- .../information_schema/table_semantics.result | 24 ++-- .../information_schema/table_semantics.sql | 4 +- .../common/show/show_databases_tables.result | 10 +- .../common/system/information_schema.result | 10 +- .../standalone/common/view/create.result | 10 +- 8 files changed, 117 insertions(+), 72 deletions(-) diff --git a/src/catalog/src/system_schema/information_schema.rs b/src/catalog/src/system_schema/information_schema.rs index b2145df9cd3..1f65652289d 100644 --- a/src/catalog/src/system_schema/information_schema.rs +++ b/src/catalog/src/system_schema/information_schema.rs @@ -139,7 +139,7 @@ lazy_static! { } macro_rules! setup_memory_table { - ($name: expr) => { + ($name: expr, $scope: ident) => { paste! { { let (schema, columns) = get_schema_columns($name); @@ -147,7 +147,8 @@ macro_rules! setup_memory_table { consts::[], $name, schema, - columns + columns, + InformationSchemaTableScope::$scope )) as _) } } @@ -165,7 +166,9 @@ pub struct MakeInformationTableRequest { pub enum InformationSchemaTableScope { /// Every query is filtered to the catalog supplied when constructing the provider. Catalog, - /// The table contains catalog-independent or shared data, including compatibility stubs. + /// Catalog-independent data that is safe for authenticated non-admin users to read. + PublicGlobal, + /// Shared data without a non-admin visibility guarantee, including compatibility stubs. Global, } @@ -173,6 +176,7 @@ impl InformationSchemaTableScope { pub const fn as_str(self) -> &'static str { match self { Self::Catalog => "catalog", + Self::PublicGlobal => "public_global", Self::Global => "global", } } @@ -228,30 +232,30 @@ impl InformationSchemaProvider { self.catalog_name.clone(), self.catalog_manager.clone(), )) as _), - ENGINES => setup_memory_table!(ENGINES), - COLUMN_PRIVILEGES => setup_memory_table!(COLUMN_PRIVILEGES), - COLUMN_STATISTICS => setup_memory_table!(COLUMN_STATISTICS), - BUILD_INFO => setup_memory_table!(BUILD_INFO), - CHARACTER_SETS => setup_memory_table!(CHARACTER_SETS), - COLLATIONS => setup_memory_table!(COLLATIONS), + ENGINES => setup_memory_table!(ENGINES, PublicGlobal), + COLUMN_PRIVILEGES => setup_memory_table!(COLUMN_PRIVILEGES, Global), + COLUMN_STATISTICS => setup_memory_table!(COLUMN_STATISTICS, Global), + BUILD_INFO => setup_memory_table!(BUILD_INFO, PublicGlobal), + CHARACTER_SETS => setup_memory_table!(CHARACTER_SETS, PublicGlobal), + COLLATIONS => setup_memory_table!(COLLATIONS, PublicGlobal), COLLATION_CHARACTER_SET_APPLICABILITY => { - setup_memory_table!(COLLATION_CHARACTER_SET_APPLICABILITY) + setup_memory_table!(COLLATION_CHARACTER_SET_APPLICABILITY, PublicGlobal) } - CHECK_CONSTRAINTS => setup_memory_table!(CHECK_CONSTRAINTS), - EVENTS => setup_memory_table!(EVENTS), - FILES => setup_memory_table!(FILES), - OPTIMIZER_TRACE => setup_memory_table!(OPTIMIZER_TRACE), - PARAMETERS => setup_memory_table!(PARAMETERS), - PROFILING => setup_memory_table!(PROFILING), - REFERENTIAL_CONSTRAINTS => setup_memory_table!(REFERENTIAL_CONSTRAINTS), - ROUTINES => setup_memory_table!(ROUTINES), - SCHEMA_PRIVILEGES => setup_memory_table!(SCHEMA_PRIVILEGES), - TABLE_PRIVILEGES => setup_memory_table!(TABLE_PRIVILEGES), - GLOBAL_STATUS => setup_memory_table!(GLOBAL_STATUS), - SESSION_STATUS => setup_memory_table!(SESSION_STATUS), - PLUGINS => setup_memory_table!(PLUGINS), - USER_PRIVILEGES => setup_memory_table!(USER_PRIVILEGES), - PROCESSLIST => setup_memory_table!(PROCESSLIST), + CHECK_CONSTRAINTS => setup_memory_table!(CHECK_CONSTRAINTS, Global), + EVENTS => setup_memory_table!(EVENTS, Global), + FILES => setup_memory_table!(FILES, Global), + OPTIMIZER_TRACE => setup_memory_table!(OPTIMIZER_TRACE, Global), + PARAMETERS => setup_memory_table!(PARAMETERS, Global), + PROFILING => setup_memory_table!(PROFILING, Global), + REFERENTIAL_CONSTRAINTS => setup_memory_table!(REFERENTIAL_CONSTRAINTS, Global), + ROUTINES => setup_memory_table!(ROUTINES, Global), + SCHEMA_PRIVILEGES => setup_memory_table!(SCHEMA_PRIVILEGES, Global), + TABLE_PRIVILEGES => setup_memory_table!(TABLE_PRIVILEGES, Global), + GLOBAL_STATUS => setup_memory_table!(GLOBAL_STATUS, Global), + SESSION_STATUS => setup_memory_table!(SESSION_STATUS, Global), + PLUGINS => setup_memory_table!(PLUGINS, Global), + USER_PRIVILEGES => setup_memory_table!(USER_PRIVILEGES, Global), + PROCESSLIST => setup_memory_table!(PROCESSLIST, Global), KEY_COLUMN_USAGE => Some(Arc::new(InformationSchemaKeyColumnUsage::new( self.catalog_name.clone(), self.catalog_manager.clone(), @@ -374,7 +378,6 @@ impl InformationSchemaProvider { // Only put these tables under `greptime` catalog to prevent info leak. if self.catalog_name == DEFAULT_CATALOG_NAME { table_names.extend([ - BUILD_INFO, REGION_PEERS, CLUSTER_INFO, PROCEDURE_INFO, @@ -387,6 +390,7 @@ impl InformationSchemaProvider { } table_names.extend([ + BUILD_INFO, TABLES, VIEWS, SCHEMATA, @@ -605,14 +609,15 @@ mod tests { #[test] fn every_registered_information_schema_table_has_scope() { - use InformationSchemaTableScope::{Catalog, Global}; + use InformationSchemaTableScope::{Catalog, Global, PublicGlobal}; let manager: Arc = MemoryCatalogManager::new(); let backend = Arc::new(MemoryKvBackend::default()); let factories: HashMap<_, InformationSchemaTableFactoryRef> = [ ("future_catalog_table", Catalog), ("future_global_table", Global), - (ENGINES, Catalog), + ("future_public_table", PublicGlobal), + (BUILD_INFO, Global), ] .into_iter() .map(|(name, scope)| { @@ -638,7 +643,13 @@ mod tests { #[cfg(feature = "enterprise")] RECYCLE_BIN, "future_catalog_table", + ]; + let public_global_tables = [ ENGINES, + CHARACTER_SETS, + COLLATIONS, + COLLATION_CHARACTER_SET_APPLICABILITY, + "future_public_table", ]; for (catalog, with_process_manager) in [ @@ -654,31 +665,42 @@ mod tests { with_process_manager.then(|| Arc::new(ProcessManager::new(String::new(), None))), backend.clone(), ); - assert_eq!( - provider - .table(ENGINES) - .unwrap() - .table_info() - .meta - .options - .extra_options - .get(SEMANTIC_TABLE_SCOPE) - .map(String::as_str), - Some(Global.as_str()), - ); + for name in [ + BUILD_INFO, + ENGINES, + CHARACTER_SETS, + COLLATIONS, + COLLATION_CHARACTER_SET_APPLICABILITY, + ] { + assert_eq!( + provider + .table(name) + .unwrap() + .table_info() + .meta + .options + .extra_options + .get(SEMANTIC_TABLE_SCOPE) + .map(String::as_str), + Some(PublicGlobal.as_str()), + "{catalog}.{name}" + ); + } // An extension replacing a built-in supplies the replacement table's scope. let provider = provider.with_extra_table_factories(factories.clone()); let expected_count = - 43 + usize::from(with_process_manager) + usize::from(cfg!(feature = "enterprise")) + 44 + usize::from(with_process_manager) + usize::from(cfg!(feature = "enterprise")) - if catalog == DEFAULT_CATALOG_NAME { 0 } else { - 9 + 8 }; assert_eq!(provider.tables().len(), expected_count); for (name, table) in provider.tables() { let scope = if catalog_tables.contains(&name.as_str()) { Catalog + } else if public_global_tables.contains(&name.as_str()) { + PublicGlobal } else { Global }; diff --git a/src/catalog/src/system_schema/memory_table.rs b/src/catalog/src/system_schema/memory_table.rs index c7246939c23..e01487a2045 100644 --- a/src/catalog/src/system_schema/memory_table.rs +++ b/src/catalog/src/system_schema/memory_table.rs @@ -40,6 +40,7 @@ pub(crate) struct MemoryTable { pub(crate) table_name: &'static str, pub(crate) schema: SchemaRef, pub(crate) columns: Vec, + scope: InformationSchemaTableScope, } impl MemoryTable { @@ -49,12 +50,14 @@ impl MemoryTable { table_name: &'static str, schema: SchemaRef, columns: Vec, + scope: InformationSchemaTableScope, ) -> Self { Self { table_id, table_name, schema, columns, + scope, } } @@ -107,7 +110,7 @@ impl DfPartitionStream for MemoryTable { impl InformationTable for MemoryTable { fn scope(&self) -> InformationSchemaTableScope { - InformationSchemaTableScope::Global + self.scope } } @@ -172,6 +175,7 @@ mod tests { Arc::new(StringVector::from(vec!["a1", "a2"])), Arc::new(StringVector::from(vec!["b1", "b2"])), ], + InformationSchemaTableScope::Global, ); assert_eq!(42, table.table_id()); @@ -201,7 +205,13 @@ mod tests { ColumnSchema::new("b", ConcreteDataType::string_datatype(), false), ])); - let table = MemoryTable::new(42, "test", schema.clone(), vec![]); + let table = MemoryTable::new( + 42, + "test", + schema.clone(), + vec![], + InformationSchemaTableScope::Global, + ); assert_eq!(42, table.table_id()); assert_eq!("test", table.table_name()); diff --git a/src/table/src/requests/semantic.rs b/src/table/src/requests/semantic.rs index 05b1edb56f5..bf15e4cb2ac 100644 --- a/src/table/src/requests/semantic.rs +++ b/src/table/src/requests/semantic.rs @@ -50,8 +50,9 @@ pub const SEMANTIC_PER_TABLE_INDEX_KEY: &str = "greptime.internal.semantic.per_t // ---- Common keys (all signals) ---- /// Information-schema visibility: `catalog` for providers that filter rows to -/// their catalog, or `global` for catalog-independent data. Only registered -/// system-table metadata is trusted for authorization. +/// their catalog, `public_global` for shared data safe for non-admin reads, or +/// `global` for shared data without that guarantee. Only registered system-table +/// metadata is trusted for authorization. pub const SEMANTIC_TABLE_SCOPE: &str = "greptime.semantic.table_scope"; /// Signal kind: one of [`SIGNAL_TYPE_TRACE`] / [`SIGNAL_TYPE_LOG`] / @@ -288,7 +289,7 @@ pub fn validate_semantic_option(key: &str, value: &str) -> bool { return !value.is_empty() && value.split(',').all(|column| !column.trim().is_empty()); } match key { - SEMANTIC_TABLE_SCOPE => matches!(value, "catalog" | "global"), + SEMANTIC_TABLE_SCOPE => matches!(value, "catalog" | "public_global" | "global"), SEMANTIC_PIPELINE | SEMANTIC_SOURCE_VERSION | SEMANTIC_METRIC_UNIT @@ -412,6 +413,10 @@ mod tests { fn test_validate_semantic_option() { // Enum keys reject out-of-domain values. assert!(validate_semantic_option(SEMANTIC_TABLE_SCOPE, "catalog")); + assert!(validate_semantic_option( + SEMANTIC_TABLE_SCOPE, + "public_global" + )); assert!(validate_semantic_option(SEMANTIC_TABLE_SCOPE, "global")); assert!(!validate_semantic_option(SEMANTIC_TABLE_SCOPE, "cluster")); assert!(!validate_semantic_option(SEMANTIC_TABLE_SCOPE, "unknown")); diff --git a/tests/cases/standalone/common/information_schema/table_semantics.result b/tests/cases/standalone/common/information_schema/table_semantics.result index 6aac406f7a0..e1edfc94db2 100644 --- a/tests/cases/standalone/common/information_schema/table_semantics.result +++ b/tests/cases/standalone/common/information_schema/table_semantics.result @@ -68,17 +68,23 @@ ORDER BY table_name; SELECT table_name, semantic_options FROM information_schema.table_semantics WHERE table_schema = 'information_schema' - AND table_name IN ('build_info', 'partitions', 'process_list', 'table_semantics') + AND table_name IN ('build_info', 'engines', 'character_sets', 'collations', + 'collation_character_set_applicability', 'partitions', + 'process_list', 'table_semantics') ORDER BY table_name; -+-----------------+---------------------------+ -| table_name | semantic_options | -+-----------------+---------------------------+ -| build_info | {"table_scope":"global"} | -| partitions | {"table_scope":"catalog"} | -| process_list | {"table_scope":"global"} | -| table_semantics | {"table_scope":"catalog"} | -+-----------------+---------------------------+ ++---------------------------------------+---------------------------------+ +| table_name | semantic_options | ++---------------------------------------+---------------------------------+ +| build_info | {"table_scope":"public_global"} | +| character_sets | {"table_scope":"public_global"} | +| collation_character_set_applicability | {"table_scope":"public_global"} | +| collations | {"table_scope":"public_global"} | +| engines | {"table_scope":"public_global"} | +| partitions | {"table_scope":"catalog"} | +| process_list | {"table_scope":"global"} | +| table_semantics | {"table_scope":"catalog"} | ++---------------------------------------+---------------------------------+ -- Predicate pushdown on a promoted column. SELECT table_name, signal_type diff --git a/tests/cases/standalone/common/information_schema/table_semantics.sql b/tests/cases/standalone/common/information_schema/table_semantics.sql index d20d47e64e0..9899bbe7a5b 100644 --- a/tests/cases/standalone/common/information_schema/table_semantics.sql +++ b/tests/cases/standalone/common/information_schema/table_semantics.sql @@ -39,7 +39,9 @@ ORDER BY table_name; SELECT table_name, semantic_options FROM information_schema.table_semantics WHERE table_schema = 'information_schema' - AND table_name IN ('build_info', 'partitions', 'process_list', 'table_semantics') + AND table_name IN ('build_info', 'engines', 'character_sets', 'collations', + 'collation_character_set_applicability', 'partitions', + 'process_list', 'table_semantics') ORDER BY table_name; -- Predicate pushdown on a promoted column. diff --git a/tests/cases/standalone/common/show/show_databases_tables.result b/tests/cases/standalone/common/show/show_databases_tables.result index 29b11ff23f7..d9eaaaa3cfe 100644 --- a/tests/cases/standalone/common/show/show_databases_tables.result +++ b/tests/cases/standalone/common/show/show_databases_tables.result @@ -135,16 +135,16 @@ SHOW TABLE STATUS; +++++++++++++++++++ |Name|Engine|Version|Row_format|Rows|Avg_row_length|Data_length|Max_data_length|Index_length|Data_free|Auto_increment|Create_time|Update_time|Check_time|Collation|Checksum|Create_options|Comment| +++++++++++++++++++ -|build_info||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| -|character_sets||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| +|build_info||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=public_global|| +|character_sets||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=public_global|| |check_constraints||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| |cluster_info||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| -|collation_character_set_applicability||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| -|collations||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| +|collation_character_set_applicability||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=public_global|| +|collations||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=public_global|| |column_privileges||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| |column_statistics||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| |columns||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=catalog|| -|engines||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| +|engines||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=public_global|| |events||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| |files||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global|| |flow_statistics||11|Fixed|0|0|0|0|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=catalog|| diff --git a/tests/cases/standalone/common/system/information_schema.result b/tests/cases/standalone/common/system/information_schema.result index a83d66f4109..8fb02828fea 100644 --- a/tests/cases/standalone/common/system/information_schema.result +++ b/tests/cases/standalone/common/system/information_schema.result @@ -15,16 +15,16 @@ order by table_schema, table_name; +++++++++++++++++++++++++ |table_catalog|table_schema|table_name|table_type|table_id|data_length|max_data_length|index_length|max_index_length|avg_row_length|engine|version|row_format|table_rows|data_free|auto_increment|create_time|update_time|check_time|table_collation|checksum|create_options|table_comment|temporary| +++++++++++++++++++++++++ -|greptime|information_schema|build_info|LOCALTEMPORARY|8|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| -|greptime|information_schema|character_sets|LOCALTEMPORARY|9|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|build_info|LOCALTEMPORARY|8|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=public_global||Y| +|greptime|information_schema|character_sets|LOCALTEMPORARY|9|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=public_global||Y| |greptime|information_schema|check_constraints|LOCALTEMPORARY|12|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| |greptime|information_schema|cluster_info|LOCALTEMPORARY|31|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| -|greptime|information_schema|collation_character_set_applicability|LOCALTEMPORARY|11|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| -|greptime|information_schema|collations|LOCALTEMPORARY|10|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|collation_character_set_applicability|LOCALTEMPORARY|11|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=public_global||Y| +|greptime|information_schema|collations|LOCALTEMPORARY|10|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=public_global||Y| |greptime|information_schema|column_privileges|LOCALTEMPORARY|6|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| |greptime|information_schema|column_statistics|LOCALTEMPORARY|7|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| |greptime|information_schema|columns|LOCALTEMPORARY|4|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=catalog||Y| -|greptime|information_schema|engines|LOCALTEMPORARY|5|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|engines|LOCALTEMPORARY|5|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=public_global||Y| |greptime|information_schema|events|LOCALTEMPORARY|13|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| |greptime|information_schema|files|LOCALTEMPORARY|14|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=global||Y| |greptime|information_schema|flow_statistics|LOCALTEMPORARY|45|0|0|0|0|0||11|Fixed|0|0|0|DATETIME|DATETIME||utf8_bin|0|greptime.semantic.table_scope=catalog||Y| diff --git a/tests/cases/standalone/common/view/create.result b/tests/cases/standalone/common/view/create.result index 82d6765ded8..e8c53070189 100644 --- a/tests/cases/standalone/common/view/create.result +++ b/tests/cases/standalone/common/view/create.result @@ -96,16 +96,16 @@ ORDER BY TABLE_NAME, TABLE_TYPE; +++++++++++++++++++++++++ |table_catalog|table_schema|table_name|table_type|table_id|data_length|max_data_length|index_length|max_index_length|avg_row_length|engine|version|row_format|table_rows|data_free|auto_increment|create_time|update_time|check_time|table_collation|checksum|create_options|table_comment|temporary| +++++++++++++++++++++++++ -|greptime|information_schema|build_info|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| -|greptime|information_schema|character_sets|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|build_info|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=public_global||Y| +|greptime|information_schema|character_sets|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=public_global||Y| |greptime|information_schema|check_constraints|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| |greptime|information_schema|cluster_info|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| -|greptime|information_schema|collation_character_set_applicability|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| -|greptime|information_schema|collations|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|collation_character_set_applicability|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=public_global||Y| +|greptime|information_schema|collations|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=public_global||Y| |greptime|information_schema|column_privileges|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| |greptime|information_schema|column_statistics|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| |greptime|information_schema|columns|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=catalog||Y| -|greptime|information_schema|engines|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| +|greptime|information_schema|engines|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=public_global||Y| |greptime|information_schema|events|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| |greptime|information_schema|files|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=global||Y| |greptime|information_schema|flow_statistics|LOCALTEMPORARY|ID|ID|ID|ID|ID|ID||ID|Fixed|ID|ID|ID|DATETIME|DATETIME||utf8_bin|ID|greptime.semantic.table_scope=catalog||Y| From c96cdd0026bfda6bb1c1dc0d71d2ada192275339 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Mon, 26 Oct 2026 08:54:02 +0000 Subject: [PATCH 094/239] perf: stream and prune SST manifest scans (#9931) * refactor: add reusable metadata pruning and projected SST batches Signed-off-by: WenyXu * perf(mito2): stream and prune manifest SST metadata Signed-off-by: WenyXu * perf(datanode): stream inspection queries through reusable providers Signed-off-by: WenyXu * perf(catalog): push manifest filters through inspection queries Signed-off-by: WenyXu * refactor: align manifest column constants with module conventions Signed-off-by: WenyXu * perf: reduce manifest inspection allocations Signed-off-by: WenyXu * refactor: narrow manifest inspection API visibility Signed-off-by: WenyXu * refactor: retain residual filters for manifest inspection Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- src/catalog/src/information_extension.rs | 119 ++++++ .../system_schema/information_schema/ssts.rs | 7 + src/common/recordbatch/src/filter.rs | 38 +- src/common/recordbatch/src/filter/tests.rs | 88 ++++ src/datanode/src/region_server/catalog.rs | 18 +- .../src/region_server/catalog/inspect.rs | 311 ++++++++++++++ .../region_server/catalog/tests/manifest.rs | 257 ++++++++++++ src/metric-engine/src/engine/bulk_insert.rs | 11 +- src/mito2/src/engine.rs | 116 +++++- src/mito2/src/engine/basic_test.rs | 244 +++++++++++ src/mito2/src/engine/staging_test.rs | 98 +++++ src/mito2/src/region.rs | 355 ++++++++++++++-- src/store-api/src/sst_entry.rs | 381 ++++++++++++------ .../information_schema/ssts_limit.result | 53 +++ .../information_schema/ssts_limit.sql | 22 + .../common/information_schema/ssts.result | 78 ++++ .../common/information_schema/ssts.sql | 35 ++ 17 files changed, 2054 insertions(+), 177 deletions(-) create mode 100644 src/datanode/src/region_server/catalog/inspect.rs create mode 100644 src/datanode/src/region_server/catalog/tests/manifest.rs diff --git a/src/catalog/src/information_extension.rs b/src/catalog/src/information_extension.rs index 9d55f9f7a1e..f22c169df13 100644 --- a/src/catalog/src/information_extension.rs +++ b/src/catalog/src/information_extension.rs @@ -310,3 +310,122 @@ impl InformationExtension for DistributedInformationExtension { )))) } } + +#[cfg(test)] +mod tests { + use common_meta::node_manager::{DatanodeManager, DatanodeRef}; + use common_meta::peer::Peer; + use datafusion::common::tree_node::TreeNode; + use datafusion::logical_expr::LogicalPlan; + use datafusion::physical_plan::sorts::sort::SortExec; + use meta_client::client::MetaClientBuilder; + use store_api::sst_entry::ManifestSstEntry; + use store_api::storage::ScanRequest; + + use super::*; + use crate::information_schema::DatanodeInspectKind; + + struct UnusedDatanodeManager; + + #[async_trait::async_trait] + impl DatanodeManager for UnusedDatanodeManager { + async fn datanode(&self, _node: &Peer) -> DatanodeRef { + panic!("plan tests must not contact datanodes") + } + } + + #[derive(Debug)] + struct ManifestScanProvider; + + #[async_trait::async_trait] + impl datafusion::catalog::TableProvider for ManifestScanProvider { + fn schema(&self) -> ArrowSchemaRef { + ManifestSstEntry::schema().arrow_schema().clone() + } + + fn table_type(&self) -> datafusion::logical_expr::TableType { + datafusion::logical_expr::TableType::View + } + + fn supports_filters_pushdown( + &self, + filters: &[&datafusion::logical_expr::Expr], + ) -> Result> { + Ok(vec![datafusion::logical_expr::TableProviderFilterPushDown::Inexact; filters.len()]) + } + + async fn scan( + &self, + _state: &dyn datafusion::catalog::Session, + projection: Option<&Vec>, + filters: &[datafusion::logical_expr::Expr], + limit: Option, + ) -> Result> { + let schema = ManifestSstEntry::schema(); + let schema = match projection { + Some(projection) => Arc::new(schema.try_project(projection).unwrap()), + None => schema, + }; + Ok(Arc::new(DistributedInspectExec::new( + DatanodeInspectRequest { + kind: DatanodeInspectKind::SstManifest, + scan: ScanRequest { + projection: projection.cloned(), + filters: filters.to_vec(), + limit, + ..Default::default() + }, + }, + schema, + Arc::new(MetaClientBuilder::new(0, api::v1::meta::Role::Frontend).build()), + Arc::new(UnusedDatanodeManager), + ))) + } + } + + #[tokio::test] + async fn test_sql_manifest_sort_stays_on_frontend() { + let ctx = datafusion::prelude::SessionContext::new(); + ctx.register_table("ssts_manifest", Arc::new(ManifestScanProvider)) + .unwrap(); + for sql in [ + "SELECT region_id, file_id, min_ts, max_ts, file_size FROM ssts_manifest WHERE region_id = 42 ORDER BY min_ts LIMIT 100", + "SELECT file_id FROM ssts_manifest WHERE region_id = 42 ORDER BY max_ts DESC LIMIT 10 OFFSET 3", + "SELECT min_ts AS ts, file_id FROM ssts_manifest WHERE region_id = 42 ORDER BY ts DESC NULLS FIRST LIMIT 10", + "SELECT file_id FROM ssts_manifest WHERE region_id = 42 ORDER BY min_ts, max_ts DESC LIMIT 10", + ] { + let plan = ctx + .sql(sql) + .await + .unwrap() + .create_physical_plan() + .await + .unwrap(); + let mut found_scan = false; + let mut found_sort = false; + plan.apply(|plan| { + found_sort |= plan.is::(); + if let Some(scan) = plan.downcast_ref::() { + found_scan = true; + // All matching rows must reach the frontend global sort. + assert_eq!(scan.request.scan.limit, None); + let remote = scan.request.clone().build_plan().unwrap(); + remote.apply(|node| { + assert!(!matches!( + node, + LogicalPlan::Sort(_) | LogicalPlan::Limit(_) + )); + Ok(TreeNodeRecursion::Continue) + })?; + } + Ok(TreeNodeRecursion::Continue) + }) + .unwrap(); + assert!( + found_scan && found_sort, + "{}", + datafusion::physical_plan::displayable(plan.as_ref()).indent(true) + ); + } + } +} diff --git a/src/catalog/src/system_schema/information_schema/ssts.rs b/src/catalog/src/system_schema/information_schema/ssts.rs index 9006ab6fed9..ace6a87bce5 100644 --- a/src/catalog/src/system_schema/information_schema/ssts.rs +++ b/src/catalog/src/system_schema/information_schema/ssts.rs @@ -349,6 +349,13 @@ mod tests { .unwrap() .unwrap(); + let filter = + datafusion::logical_expr::col("region_id").eq(datafusion::logical_expr::lit(42_u64)); + assert_eq!( + table.supports_filters_pushdown(&[&filter]).unwrap(), + vec![table::metadata::FilterPushDownType::Inexact], + ); + let request = ScanRequest { projection: Some(vec![0]), ..Default::default() diff --git a/src/common/recordbatch/src/filter.rs b/src/common/recordbatch/src/filter.rs index 60ff399601f..0893d2aa3ad 100644 --- a/src/common/recordbatch/src/filter.rs +++ b/src/common/recordbatch/src/filter.rs @@ -19,10 +19,11 @@ mod predicate; mod tests; mod timestamp; +use std::borrow::Borrow; use std::sync::Arc; use datafusion::error::Result as DfResult; -use datafusion::logical_expr::{Expr, Literal, Operator}; +use datafusion::logical_expr::{Expr, Literal, Operator, utils}; use datafusion::physical_plan::PhysicalExpr; use datafusion_common::arrow::array::{ArrayRef, Datum}; use datafusion_common::arrow::buffer::BooleanBuffer; @@ -41,6 +42,41 @@ pub use self::timestamp::timestamp_scalar_value; use crate::error::{Result, ToArrowScalarSnafu}; use crate::filter::predicate::SimplePredicate; +/// Prunes candidates using known metadata columns. Missing columns and unsupported +/// predicates cannot exclude a candidate. Complete SQL predicates must still run +/// after pruning. +#[derive(Debug, Clone, Default)] +pub struct MetadataFilter { + filters: Arc<[SimpleFilterEvaluator]>, +} + +impl MetadataFilter { + /// Compiles supported conjuncts once for reuse across scan stages. OR is + /// never split: an unknown alternative might match the candidate. + pub fn new(expressions: &[Expr]) -> Self { + Self { + filters: expressions + .iter() + .flat_map(utils::split_conjunction) + .filter_map(SimpleFilterEvaluator::try_new) + .collect(), + } + } + + /// Returns false only when known metadata proves a conjunct cannot match. + /// `None` means unavailable metadata, distinct from a known SQL NULL scalar. + /// Values may be borrowed so repeated predicates can reuse allocated strings. + pub fn matches>( + &self, + mut value: impl FnMut(&str) -> Option, + ) -> bool { + self.filters.iter().all(|filter| { + value(filter.column_name()) + .is_none_or(|value| filter.evaluate_scalar(value.borrow()).unwrap_or(true)) + }) + } +} + /// Evaluates supported single-column predicates with precompiled regular expressions. /// Compound predicates preserve SQL null semantics until conversion to a scan mask. #[derive(Debug, Clone)] diff --git a/src/common/recordbatch/src/filter/tests.rs b/src/common/recordbatch/src/filter/tests.rs index 77abf59b404..f48b4a3c126 100644 --- a/src/common/recordbatch/src/filter/tests.rs +++ b/src/common/recordbatch/src/filter/tests.rs @@ -70,3 +70,91 @@ fn test_nullable_boolean_predicate_becomes_scan_mask() { boolean_array_to_scan_mask(&predicate) ); } + +#[test] +fn test_partial_metadata() { + use datafusion_common::ScalarValue; + + use crate::filter::MetadataFilter; + + let node = col("node_id").eq(lit(7u64)); + let path = col("file_path").eq(lit("table/region/file.parquet")); + let filters = MetadataFilter::new(&[node.clone().and(path.clone())]); + assert!(!filters.matches(|name| (name == "node_id").then_some(ScalarValue::UInt64(Some(8))))); + assert!(filters.matches(|name| (name == "node_id").then_some(ScalarValue::UInt64(Some(7))))); + assert!(!filters.matches(|name| { + (name == "file_path").then(|| ScalarValue::Utf8(Some("other/file.parquet".into()))) + })); + assert!(filters.matches(|name| { + (name == "file_path").then(|| ScalarValue::Utf8(Some("table/region/file.parquet".into()))) + })); + // Unknown alternatives cannot be discarded from OR, nor extracted from NOT. + let disjunction = MetadataFilter::new(&[node.clone().or(path.clone())]); + assert!( + disjunction.matches(|name| (name == "node_id").then_some(ScalarValue::UInt64(Some(8)))) + ); + let negation = MetadataFilter::new(&[datafusion::logical_expr::Expr::Not(Box::new( + node.and(path), + ))]); + assert!(negation.matches(|_| None::)); + // Unsupported LIKE remains a residual even when its column is available. + let residual = MetadataFilter::new(&[col("file_path").like(lit("table/%"))]); + assert!(residual.matches(|_| Some(ScalarValue::Utf8(Some("other/file".into()))))); + // Missing metadata differs from known NULL; type errors conservatively pass. + let null_filter = MetadataFilter::new(&[col("node_id").is_not_null()]); + assert!(null_filter.matches(|_| None::)); + assert!(!null_filter.matches(|_| Some(ScalarValue::UInt64(None)))); + let typed = MetadataFilter::new(&[col("node_id").eq(lit(7u64))]); + assert!(typed.matches(|_| Some(ScalarValue::Utf8(Some("invalid".into()))))); +} + +#[test] +fn test_borrowed_and_owned_values() { + use datafusion_common::ScalarValue; + + use crate::filter::MetadataFilter; + + let pruner = MetadataFilter::new(&[col("path").gt_eq(lit("a")).and(col("path").lt(lit("z")))]); + for (value, expected) in [ + (ScalarValue::Utf8(Some("a".into())), true), + (ScalarValue::Utf8(Some("z".into())), false), + (ScalarValue::Utf8(None), false), + (ScalarValue::UInt64(Some(1)), true), + ] { + assert_eq!(pruner.matches(|_| Some(&value)), expected); + assert_eq!(pruner.matches(|_| Some(value.clone())), expected); + } +} + +#[test] +fn test_borrowed_lazy_metadata() { + use std::cell::OnceCell; + + use datafusion_common::ScalarValue; + + use crate::filter::MetadataFilter; + + let pruner = MetadataFilter::new(&[col("path").gt(lit("a")).and(col("path").lt(lit("z")))]); + let path = OnceCell::new(); + let mut builds = 0; + assert!(pruner.matches(|name| { + (name == "path").then(|| { + path.get_or_init(|| { + builds += 1; + ScalarValue::Utf8(Some("region/file.parquet".into())) + }) + }) + })); + assert_eq!(builds, 1); + assert_eq!( + path.into_inner(), + Some(ScalarValue::Utf8(Some("region/file.parquet".into()))) + ); + + let path = OnceCell::new(); + assert!( + MetadataFilter::default() + .matches(|_| { Some(path.get_or_init(|| ScalarValue::Utf8(Some("unused".into())))) }) + ); + assert!(path.get().is_none()); +} diff --git a/src/datanode/src/region_server/catalog.rs b/src/datanode/src/region_server/catalog.rs index 264408342e9..789aa514785 100644 --- a/src/datanode/src/region_server/catalog.rs +++ b/src/datanode/src/region_server/catalog.rs @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +mod inspect; + use std::collections::HashMap; use std::sync::Arc; @@ -33,6 +35,7 @@ use store_api::storage::RegionId; use crate::error::{DataFusionSnafu, ListStorageSstsSnafu, Result, UnexpectedSnafu}; use crate::region_server::RegionServer; +use crate::region_server::catalog::inspect::{InspectTable, ManifestSstSource}; /// Reserved internal table kinds used. /// These are recognized by reserved table names and mapped to providers. @@ -75,7 +78,7 @@ impl InternalTableKind { } impl RegionServer { - /// Expose SSTs listed in Manifest as an in-memory table for inspection. + /// Exposes manifest SSTs through a lazy table with metadata predicate pruning. pub async fn inspect_sst_manifest_provider(&self) -> Result> { let mito = { let guard = self.inner.mito_engine.read().unwrap(); @@ -84,14 +87,9 @@ impl RegionServer { })? }; - let entries = mito.all_ssts_from_manifest().await; - let schema = ManifestSstEntry::schema().arrow_schema().clone(); - let batch = ManifestSstEntry::to_record_batch(&entries) - .map_err(DataFusionError::from) - .context(DataFusionSnafu)?; - - let table = MemTable::try_new(schema, vec![vec![batch]]).context(DataFusionSnafu)?; - Ok(Arc::new(table)) + Ok(Arc::new(InspectTable::new(Arc::new( + ManifestSstSource::new(mito), + )))) } /// Expose SSTs found in storage as an in-memory table for inspection. @@ -355,6 +353,8 @@ impl TreeNodeRewriter for NameAwareDataSourceInjector { #[cfg(test)] mod tests { + mod manifest; + use std::sync::Arc; use datafusion::catalog::MemTable as DfMemTable; diff --git a/src/datanode/src/region_server/catalog/inspect.rs b/src/datanode/src/region_server/catalog/inspect.rs new file mode 100644 index 00000000000..307c8b42883 --- /dev/null +++ b/src/datanode/src/region_server/catalog/inspect.rs @@ -0,0 +1,311 @@ +// 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 providers for internal inspection tables. + +use std::sync::Arc; + +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::catalog::{Session, TableProvider}; +use datafusion::error::Result; +use datafusion::execution::TaskContext; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::streaming::{PartitionStream, StreamingTableExec}; +use datafusion::physical_plan::{ExecutionPlan, SendableRecordBatchStream}; +use datafusion_common::DataFusionError; +use datafusion_expr::{Expr, TableProviderFilterPushDown, TableType}; +use futures::StreamExt; +use mito2::engine::MitoEngine; +use store_api::sst_entry::{ManifestSstBatchBuilder, ManifestSstEntry}; + +/// Supplies batches for an internal inspection table. Implementations may prune +/// conservatively; DataFusion retains complete residual predicates. +pub(crate) trait InspectSource: std::fmt::Debug + Send + Sync { + fn schema(&self) -> SchemaRef; + + /// Creates a fresh scan at execution time, including on repeated execution. + /// Projection includes columns required by residual predicates. + fn scan(&self, request: &InspectScan, ctx: Arc) -> SendableRecordBatchStream; +} + +/// Parameters shared by inspection sources. Limits are enforced by the physical +/// plan, after residual filters, rather than by individual sources. +#[derive(Debug, Clone)] +pub(crate) struct InspectScan { + schema: SchemaRef, + projection: Option>, + filters: Vec, + batch_size: usize, +} + +/// Adapts a lazy inspection source to DataFusion's streaming execution plan. +#[derive(Debug)] +pub(crate) struct InspectTable { + source: Arc, +} + +impl InspectTable { + /// Wraps a source without opening a scan. + pub(crate) fn new(source: Arc) -> Self { + Self { source } + } +} + +#[async_trait::async_trait] +impl TableProvider for InspectTable { + fn schema(&self) -> SchemaRef { + self.source.schema() + } + + fn table_type(&self) -> TableType { + TableType::View + } + + fn supports_filters_pushdown( + &self, + filters: &[&Expr], + ) -> Result> { + Ok(vec![TableProviderFilterPushDown::Inexact; filters.len()]) + } + + async fn scan( + &self, + state: &dyn Session, + projection: Option<&Vec>, + filters: &[Expr], + limit: Option, + ) -> Result> { + // An inexact predicate must be checked before selecting a limited prefix. + let limit = limit.filter(|_| filters.is_empty()); + let configured_batch_size: usize = state.config_options().execution.batch_size.into(); + let schema = match projection { + Some(projection) => Arc::new(self.schema().project(projection)?), + None => self.schema(), + }; + let partition = InspectPartition { + source: self.source.clone(), + request: InspectScan { + schema: schema.clone(), + projection: projection.cloned(), + filters: filters.to_vec(), + batch_size: limit + .unwrap_or(configured_batch_size) + .min(configured_batch_size) + .max(1), + }, + }; + Ok(Arc::new(StreamingTableExec::try_new( + schema, + vec![Arc::new(partition)], + None, + None, + false, + limit, + )?)) + } +} + +/// Reads manifest snapshots without materializing the full SST inventory. +pub(crate) struct ManifestSstSource { + engine: MitoEngine, +} + +impl ManifestSstSource { + /// Creates a source backed by the engine's open-region manifests. + pub(crate) fn new(engine: MitoEngine) -> Self { + Self { engine } + } +} + +impl std::fmt::Debug for ManifestSstSource { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ManifestSstSource").finish_non_exhaustive() + } +} + +impl InspectSource for ManifestSstSource { + fn schema(&self) -> SchemaRef { + ManifestSstEntry::schema().arrow_schema().clone() + } + + fn scan(&self, request: &InspectScan, _ctx: Arc) -> SendableRecordBatchStream { + let builder = ManifestSstBatchBuilder::new(request.schema.clone()); + let stream = self + .engine + .scan_ssts_from_manifest_projected( + &request.filters, + request.batch_size, + request.projection.as_deref(), + ) + .map(move |entries| builder.build(&entries).map_err(DataFusionError::from)); + Box::pin(RecordBatchStreamAdapter::new( + request.schema.clone(), + stream, + )) + } +} + +#[derive(Debug)] +struct InspectPartition { + source: Arc, + request: InspectScan, +} + +impl PartitionStream for InspectPartition { + fn schema(&self) -> &SchemaRef { + &self.request.schema + } + + fn execute(&self, ctx: Arc) -> SendableRecordBatchStream { + self.source.scan(&self.request, ctx) + } +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use datafusion::arrow::array::{ArrayRef, Int32Array}; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::arrow::record_batch::RecordBatch; + use datafusion::prelude::SessionContext; + + use super::*; + + #[derive(Debug, Default)] + struct CountingSource { + scans: Arc, + rows: Arc, + drops: Arc, + } + + struct ScanGuard(Arc); + + impl Drop for ScanGuard { + fn drop(&mut self) { + self.0.fetch_add(1, Ordering::SeqCst); + } + } + + impl InspectSource for CountingSource { + fn schema(&self) -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("value", DataType::Int32, false), + ])) + } + + fn scan(&self, request: &InspectScan, _ctx: Arc) -> SendableRecordBatchStream { + self.scans.fetch_add(1, Ordering::SeqCst); + let guard = ScanGuard(self.drops.clone()); + let rows = self.rows.clone(); + let schema = request.schema.clone(); + let projection = request.projection.clone().unwrap_or_else(|| vec![0, 1]); + // Deliberately ignore source predicates: the provider must preserve + // residual filtering before limit, including hidden predicate columns. + let stream = futures::stream::iter(0..100).map(move |id| { + let _guard = &guard; + rows.fetch_add(1, Ordering::SeqCst); + let columns: Vec = projection + .iter() + .map(|column| { + Arc::new(Int32Array::from(vec![if *column == 0 { + id + } else { + id * 10 + }])) as ArrayRef + }) + .collect(); + RecordBatch::try_new(schema.clone(), columns).map_err(DataFusionError::from) + }); + Box::pin(RecordBatchStreamAdapter::new( + request.schema.clone(), + stream, + )) + } + } + + #[tokio::test] + async fn test_lazy_repeatable_scan_and_limit_cancellation() { + let source = Arc::new(CountingSource::default()); + let ctx = SessionContext::new(); + ctx.register_table("inspection", Arc::new(InspectTable::new(source.clone()))) + .unwrap(); + let frame = ctx + .sql("SELECT value FROM inspection LIMIT 2") + .await + .unwrap(); + let plan = frame.create_physical_plan().await.unwrap(); + assert_eq!(source.scans.load(Ordering::SeqCst), 0); + assert_eq!(source.rows.load(Ordering::SeqCst), 0); + for run in 1..=2 { + let batches = datafusion::physical_plan::collect(plan.clone(), ctx.task_ctx()) + .await + .unwrap(); + let values = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect::>(); + assert_eq!(values, vec![0, 10]); + assert_eq!(source.scans.load(Ordering::SeqCst), run); + assert_eq!(source.rows.load(Ordering::SeqCst), run * 2); + assert_eq!(source.drops.load(Ordering::SeqCst), run); + } + } + + #[tokio::test] + async fn test_hidden_residual_filter_before_limit() { + let source = Arc::new(CountingSource::default()); + let ctx = SessionContext::new(); + ctx.register_table("inspection", Arc::new(InspectTable::new(source.clone()))) + .unwrap(); + let batches = ctx + .sql("SELECT value FROM inspection WHERE id >= 8 LIMIT 2") + .await + .unwrap() + .collect() + .await + .unwrap(); + let values = batches + .iter() + .flat_map(|batch| { + assert_eq!(batch.num_columns(), 1); + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect::>(); + // No ORDER BY: repartitioning may choose any two qualifying rows. + assert_eq!(values.len(), 2); + assert!( + values + .iter() + .all(|value| (80..1000).contains(value) && value % 10 == 0) + ); + assert_eq!(source.scans.load(Ordering::SeqCst), 1); + assert_eq!(source.drops.load(Ordering::SeqCst), 1); + } +} diff --git a/src/datanode/src/region_server/catalog/tests/manifest.rs b/src/datanode/src/region_server/catalog/tests/manifest.rs new file mode 100644 index 00000000000..91b9f757a3d --- /dev/null +++ b/src/datanode/src/region_server/catalog/tests/manifest.rs @@ -0,0 +1,257 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use common_time::Timestamp; +use datafusion::arrow::compute::concat_batches; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::catalog::{MemTable, TableProvider}; +use datafusion::prelude::{SessionConfig, SessionContext}; +use mito2::config::MitoConfig; +use mito2::engine::MitoEngine; +use mito2::manifest::action::RegionEdit; +use mito2::sst::file::FileMeta; +use mito2::test_util::{CreateRequestBuilder, TestEnv}; +use store_api::region_engine::RegionEngine; +use store_api::region_request::RegionRequest; +use store_api::sst_entry::ManifestSstEntry; +use store_api::storage::{FileId, RegionId}; + +use crate::region_server::catalog::inspect::{InspectTable, ManifestSstSource}; + +/// Installs deterministic metadata only. No SST object is read by inspection. +async fn manifest_engine(regions: u32, files_per_region: usize) -> (TestEnv, MitoEngine) { + let mut env = TestEnv::new().await; + let engine = env + .create_engine(MitoConfig { + schedule_compaction_after_edit: false, + ..Default::default() + }) + .await; + for number in 0..regions { + let region_id = RegionId::new(11, number); + engine + .handle_request( + region_id, + RegionRequest::Create(CreateRequestBuilder::new().build()), + ) + .await + .unwrap(); + engine + .edit_region( + region_id, + RegionEdit { + files_to_add: (0..files_per_region) + .map(|index| FileMeta { + region_id, + file_id: FileId::from_bytes( + (((number as u128) << 64) | index as u128).to_be_bytes(), + ), + time_range: ( + Timestamp::new_millisecond(index as i64), + Timestamp::new_millisecond(index as i64 + 1), + ), + file_size: 1024 + index as u64, + sequence: None, + ..Default::default() + }) + .collect(), + files_to_remove: vec![], + timestamp_ms: None, + compaction_time_window: None, + flushed_entry_id: None, + flushed_sequence: None, + committed_sequence: None, + }, + ) + .await + .unwrap(); + } + (env, engine) +} + +/// Baseline preserves eager full-row materialization, then applies SQL. Both +/// variants use the current engine enumerator, so this does not measure removal +/// of the old manifest-map cloning independently. +async fn inspect_query(engine: &MitoEngine, streaming: bool, sql: &str) -> RecordBatch { + let schema = ManifestSstEntry::schema().arrow_schema().clone(); + let table: Arc = if streaming { + Arc::new(InspectTable::new(Arc::new(ManifestSstSource::new( + engine.clone(), + )))) + } else { + let entries = engine.all_ssts_from_manifest().await; + let batch = ManifestSstEntry::to_record_batch(&entries).unwrap(); + Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()) + }; + let context = SessionContext::new_with_config( + SessionConfig::new() + .with_batch_size(17) + .with_target_partitions(1), + ); + context.register_table("manifest", table).unwrap(); + let frame = context.sql(sql).await.unwrap(); + let schema = frame.schema().as_arrow().clone(); + let batches = frame.collect().await.unwrap(); + concat_batches(&Arc::new(schema), &batches).unwrap() +} + +#[tokio::test] +async fn test_manifest_provider_matches_eager_scan() { + let (_env, engine) = manifest_engine(3, 64).await; + let region = RegionId::new(11, 1).as_u64(); + for sql in [ + "SELECT file_id, min_ts FROM manifest WHERE table_id = 11 ORDER BY min_ts, file_id LIMIT 7" + .to_string(), + "SELECT file_id FROM manifest WHERE table_id = 12 LIMIT 7".to_string(), + format!( + "SELECT file_id FROM manifest WHERE table_id = 11 AND region_id = {region} ORDER BY min_ts DESC LIMIT 7" + ), + format!("SELECT * FROM manifest WHERE region_id = {region} ORDER BY file_id"), + format!( + "SELECT min_ts, max_ts FROM manifest WHERE region_id = {region} ORDER BY min_ts ASC NULLS LAST, file_id LIMIT 7" + ), + format!( + "SELECT file_size FROM manifest WHERE region_id = {region} ORDER BY max_ts DESC NULLS FIRST, file_id LIMIT 7 OFFSET 3" + ), + // Arithmetic is deliberately left as a residual predicate. Limiting + // metadata enumeration before it would return too few rows. + format!( + "SELECT file_id FROM manifest WHERE region_id = {region} AND file_size % 3 = 0 ORDER BY min_ts, file_id LIMIT 7" + ), + format!( + "SELECT file_id FROM manifest WHERE region_id = {region} AND min_ts >= TIMESTAMP '1970-01-01 00:00:00.020' ORDER BY min_ts, file_id LIMIT 7" + ), + // Variable-sized metadata needed only by a residual filter must still + // be materialized even when it is absent from the SELECT list. + format!( + "SELECT file_id FROM manifest WHERE region_id = {region} AND length(file_path) > 0 ORDER BY min_ts, file_id LIMIT 7" + ), + // Repeated path predicates reuse metadata while both path columns stay + // hidden from the output. The range selects only region 1's 64 files. + "SELECT file_id FROM manifest WHERE table_dir >= 'test' AND table_dir < 'tesu' AND file_path >= 'test/11_0000000001/' AND file_path < 'test/11_0000000002/' ORDER BY file_id LIMIT 7".to_string(), + // Visibility snapshots are needed for residual predicates even when + // the SELECT list does not include visible. + format!( + "SELECT file_id FROM manifest WHERE region_id = {region} AND visible ORDER BY file_id LIMIT 7" + ), + format!( + "SELECT file_id FROM manifest WHERE region_id = {region} AND NOT visible ORDER BY file_id LIMIT 7" + ), + "SELECT file_id FROM manifest WHERE region_id = 0 ORDER BY min_ts LIMIT 1".to_string(), + "SELECT min_ts FROM manifest ORDER BY min_ts LIMIT 0".to_string(), + "SELECT count(*) FROM manifest".to_string(), + // Unsupported compound predicates must not prune matching regions. + format!( + "SELECT file_id FROM manifest WHERE region_id = {region} OR file_size = 1024 ORDER BY file_id" + ), + ] { + let expected = inspect_query(&engine, false, &sql).await; + let actual = inspect_query(&engine, true, &sql).await; + assert_eq!(actual, expected, "{sql}"); + } +} + +#[tokio::test] +async fn test_manifest_referring_region_and_origin_path() { + use datafusion::arrow::array::{StringArray, UInt64Array}; + + let (_env, engine) = manifest_engine(2, 2).await; + let referring = RegionId::new(11, 1); + // The creating region need not be open on this datanode. + let origin = RegionId::new(11, 99); + let file_id = FileId::from_bytes(999u128.to_be_bytes()); + engine + .edit_region( + referring, + RegionEdit { + files_to_add: vec![FileMeta { + region_id: origin, + file_id, + file_size: 4096, + ..Default::default() + }], + files_to_remove: vec![], + timestamp_ms: None, + compaction_time_window: None, + flushed_entry_id: None, + flushed_sequence: None, + committed_sequence: None, + }, + ) + .await + .unwrap(); + let sql = format!( + "SELECT region_id, origin_region_id, file_path FROM manifest WHERE table_id = 11 AND region_id = {} AND origin_region_id = {} AND file_path LIKE '%/11_0000000099/%' LIMIT 1", + referring.as_u64(), + origin.as_u64(), + ); + // Check source pruning without SQL residuals: the origin is not a loaded + // region, but its file must be found in the referring region's manifest. + use datafusion_expr::{col, lit}; + use futures::StreamExt; + for (predicate, expected) in [ + (col("origin_region_id").eq(lit(origin.as_u64())), 1), + ( + col("file_path").eq(lit(format!("test/11_0000000099/{file_id}.parquet"))), + 1, + ), + (col("file_path").eq(lit("missing/file.parquet")), 0), + (col("table_dir").eq(lit("missing/")), 0), + (col("node_id").is_not_null(), 0), + ] { + let batches = engine + .scan_ssts_from_manifest_projected(std::slice::from_ref(&predicate), 2, None) + .collect::>() + .await; + assert_eq!( + batches.iter().map(Vec::len).sum::(), + expected, + "{predicate}" + ); + } + let actual = inspect_query(&engine, true, &sql).await; + assert_eq!(actual, inspect_query(&engine, false, &sql).await); + assert_eq!(actual.num_rows(), 1); + assert_eq!( + actual + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + referring.as_u64() + ); + assert_eq!( + actual + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + origin.as_u64() + ); + assert!( + actual + .column(2) + .as_any() + .downcast_ref::() + .unwrap() + .value(0) + .ends_with(&format!("11_0000000099/{file_id}.parquet")) + ); + let missing = sql.replace("%/11_0000000099/%", "%/11_0000000001/%"); + assert_eq!(inspect_query(&engine, true, &missing).await.num_rows(), 0); +} diff --git a/src/metric-engine/src/engine/bulk_insert.rs b/src/metric-engine/src/engine/bulk_insert.rs index a9080745e2e..6467c710501 100644 --- a/src/metric-engine/src/engine/bulk_insert.rs +++ b/src/metric-engine/src/engine/bulk_insert.rs @@ -768,11 +768,16 @@ mod tests { .unwrap(); let region = env.mito().find_region(data_region_id).unwrap(); - let entry = region - .manifest_sst_entries() + let entry = env + .mito() + .all_ssts_from_manifest() .await .into_iter() - .find(|entry| entry.visible && entry.file_path.ends_with(".parquet")) + .find(|entry| { + entry.region_id == data_region_id + && entry.visible + && entry.file_path.ends_with(".parquet") + }) .unwrap(); let mut cache_metrics = MetadataCacheMetrics::default(); let footer = MetadataLoader::new( diff --git a/src/mito2/src/engine.rs b/src/mito2/src/engine.rs index 715357a39b3..92d848f5677 100644 --- a/src/mito2/src/engine.rs +++ b/src/mito2/src/engine.rs @@ -104,10 +104,13 @@ use common_base::Plugins; use common_error::ext::BoxedError; use common_meta::error::UnexpectedSnafu; use common_meta::key::SchemaMetadataManagerRef; +use common_recordbatch::filter::MetadataFilter; use common_recordbatch::{QueryMemoryTracker, SendableRecordBatchStream}; use common_stat::get_total_memory_bytes; use common_telemetry::{debug, info, tracing, warn}; use common_wal::options::WalOptions; +use datafusion_common::ScalarValue; +use datafusion_expr::Expr; use futures::FutureExt; use futures::future::{self, join_all, try_join_all}; use futures::stream::{self, Stream, StreamExt}; @@ -132,7 +135,7 @@ use store_api::region_info::RegionInfoEntry; use store_api::region_request::{ AffectedRows, RegionCatchupRequest, RegionOpenRequest, RegionPutRequest, RegionRequest, }; -use store_api::sst_entry::{ManifestSstEntry, PuffinIndexMetaEntry, StorageSstEntry}; +use store_api::sst_entry::{self, ManifestSstEntry, PuffinIndexMetaEntry, StorageSstEntry}; use store_api::storage::{FileId, FileRefsManifest, RegionId, ScanRequest, SequenceNumber}; use tokio::sync::{Semaphore, oneshot}; @@ -578,19 +581,65 @@ impl MitoEngine { /// Lists all SSTs from the manifest of all regions in the engine. pub async fn all_ssts_from_manifest(&self) -> Vec { - let node_id = self.inner.workers.file_ref_manager().node_id(); - let regions = self.inner.workers.all_regions(); + self.scan_ssts_from_manifest(&[], 1024) + .flat_map(stream::iter) + .collect() + .await + } - let mut results = Vec::new(); - for region in regions { - let mut entries = region.manifest_sst_entries().await; - for e in &mut entries { - e.node_id = node_id; + /// Lazily emits bounded batches of manifest entries, pruning regions before + /// opening their manifest snapshots and files before constructing entries. + /// Filters are conservative; callers must retain their residual predicates. + pub(crate) fn scan_ssts_from_manifest( + &self, + filters: &[Expr], + batch_size: usize, + ) -> futures::stream::BoxStream<'static, Vec> { + self.scan_ssts_from_manifest_projected(filters, batch_size, None) + } + + /// Like [`Self::scan_ssts_from_manifest`], but avoids allocating unrequested + /// variable-sized fields. The projection must include residual filter columns. + pub fn scan_ssts_from_manifest_projected( + &self, + filters: &[Expr], + batch_size: usize, + projection: Option<&[usize]>, + ) -> futures::stream::BoxStream<'static, Vec> { + let engine = self.clone(); + let projection = crate::region::ManifestSstProjection::new(projection); + let filters = MetadataFilter::new(filters); + Box::pin(async_stream::stream! { + let node_id = engine.inner.workers.file_ref_manager().node_id(); + if !filters.matches(|column| match column { + sst_entry::NODE_ID => Some(ScalarValue::UInt64(node_id)), + _ => None, + }) { + return; } - results.extend(entries); - } - - results + let mut batch = Vec::with_capacity(batch_size.max(1)); + for region in engine.inner.workers.all_regions() { + if !manifest_region_matches(region.region_id, node_id, &filters) + || !filters.matches(|column| match column { + sst_entry::TABLE_DIR => Some(ScalarValue::Utf8(Some(region.table_dir().to_string()))), + _ => None, + }) + { + continue; + } + let mut entries = region.manifest_sst_stream(filters.clone(), projection); + while let Some(mut entry) = entries.next().await { + entry.node_id = node_id; + batch.push(entry); + if batch.len() == batch_size.max(1) { + yield std::mem::replace(&mut batch, Vec::with_capacity(batch_size.max(1))); + } + } + } + if !batch.is_empty() { + yield batch; + } + }) } /// Lists metadata about all puffin index targets stored in the engine. @@ -1700,6 +1749,25 @@ impl MitoEngine { } } +/// Evaluates region-constant predicates before visiting any of its files. +fn manifest_region_matches( + region_id: RegionId, + node_id: Option, + filters: &MetadataFilter, +) -> bool { + filters.matches(|column| { + Some(match column { + sst_entry::REGION_ID => ScalarValue::UInt64(Some(region_id.as_u64())), + sst_entry::TABLE_ID => ScalarValue::UInt32(Some(region_id.table_id())), + sst_entry::REGION_NUMBER => ScalarValue::UInt32(Some(region_id.region_number())), + sst_entry::REGION_GROUP => ScalarValue::UInt8(Some(region_id.region_group())), + sst_entry::REGION_SEQUENCE => ScalarValue::UInt32(Some(region_id.region_sequence())), + sst_entry::NODE_ID => ScalarValue::UInt64(node_id), + _ => return None, + }) + }) +} + #[cfg(test)] mod tests { use std::time::Duration; @@ -1707,6 +1775,30 @@ mod tests { use super::*; use crate::sst::file::FileMeta; + #[test] + fn test_manifest_region_predicates() { + use super::manifest_region_matches; + use common_recordbatch::filter::MetadataFilter; + use datafusion_expr::{col, lit}; + let region = store_api::storage::RegionId::new(42, 3); + for (expr, expected) in [ + (col("region_id").eq(lit(region.as_u64())), true), + (col("region_id").eq(lit(region.as_u64() + 1)), false), + (col("table_id").eq(lit(42u32)), true), + (col("node_id").is_null(), true), + (col("file_size").gt(lit(100u64)), true), + // A type mismatch is left to the residual filter, never pruned. + (col("region_id").eq(lit("invalid")), true), + ] { + let filter = MetadataFilter::new(std::slice::from_ref(&expr)); + assert_eq!( + manifest_region_matches(region, None, &filter), + expected, + "{expr}" + ); + } + } + #[test] fn test_is_valid_region_edit() { // Valid: has only "files_to_add" diff --git a/src/mito2/src/engine/basic_test.rs b/src/mito2/src/engine/basic_test.rs index 2576158bab6..1b3d27120c8 100644 --- a/src/mito2/src/engine/basic_test.rs +++ b/src/mito2/src/engine/basic_test.rs @@ -2047,3 +2047,247 @@ async fn test_bulk_insert_encode_failure_keeps_worker_batch() { assert_eq!(3, scan_num_rows(&engine, region_id).await); assert_eq!(2, scan_num_rows(&engine, other_region_id).await); } + +#[tokio::test] +async fn test_manifest_scan_region_filter() { + use datafusion_expr::{col, lit}; + use futures::StreamExt; + + let regions = 3; + let files_per_region = 64; + let mut env = TestEnv::new().await; + let engine = env.create_engine(MitoConfig::default()).await; + for number in 0..regions { + add_manifest_scan_region(&engine, RegionId::new(11, number), files_per_region).await; + } + + let target = RegionId::new(11, 1); + let batch_size = 17; + let mut reference = None; + for pushed in [false, true] { + let filters = if pushed { + vec![col("region_id").eq(lit(target.as_u64()))] + } else { + vec![] + }; + let mut stream = engine.scan_ssts_from_manifest(&filters, batch_size); + let mut examined_entries = 0; + let mut selected = Vec::new(); + while let Some(batch) = stream.next().await { + assert!(batch.len() <= batch_size); + examined_entries += batch.len(); + selected.extend(batch.into_iter().filter(|entry| entry.region_id == target)); + } + assert_eq!(selected.len(), files_per_region); + assert_eq!( + examined_entries, + if pushed { + files_per_region + } else { + regions as usize * files_per_region + } + ); + selected.sort_unstable_by(|a, b| a.file_id.cmp(&b.file_id)); + if let Some(reference) = &reference { + assert_eq!(&selected, reference); + } else { + reference = Some(selected); + } + } + // Region pruning must happen before acquiring a manifest snapshot. Holding + // another region's write lock would deadlock an unpruned enumeration. + let unrelated = engine.get_region(RegionId::new(11, 0)).unwrap(); + let guard = unrelated.manifest_ctx.manifest_manager.write().await; + // Node and directory predicates also reject candidates before opening any + // manifest; otherwise the locked region would block these empty scans. + for predicate in [ + col("node_id").is_not_null(), + col("table_dir").eq(lit("missing/")), + ] { + let batches = tokio::time::timeout( + Duration::from_secs(5), + engine + .scan_ssts_from_manifest(&[predicate], batch_size) + .collect::>(), + ) + .await + .expect("pruned scan must not acquire a manifest lock"); + assert!(batches.is_empty()); + } + let filtered = + engine.scan_ssts_from_manifest(&[col("region_id").eq(lit(target.as_u64()))], batch_size); + let batches = tokio::time::timeout(Duration::from_secs(5), filtered.collect::>()) + .await + .expect("filtered scan must not acquire an unrelated manifest lock"); + assert_eq!( + batches.iter().map(Vec::len).sum::(), + files_per_region + ); + drop(guard); + + let schema = store_api::sst_entry::ManifestSstEntry::schema(); + let projection = ["file_id", "min_ts"] + .map(|name| { + schema + .column_schemas() + .iter() + .position(|c| c.name == name) + .unwrap() + }) + .to_vec(); + let mut projected = engine.scan_ssts_from_manifest_projected( + &[col("region_id").eq(lit(target.as_u64()))], + batch_size, + Some(&projection), + ); + let mut projected_entries = Vec::new(); + while let Some(batch) = projected.next().await { + assert!( + batch + .iter() + .all(|entry| entry.table_dir.is_empty() && entry.file_path.is_empty()) + ); + projected_entries.extend(batch); + } + projected_entries.sort_unstable_by(|a, b| a.file_id.cmp(&b.file_id)); + let builder = store_api::sst_entry::ManifestSstBatchBuilder::new(Arc::new( + store_api::sst_entry::ManifestSstEntry::schema() + .arrow_schema() + .project(&projection) + .unwrap(), + )); + let expected = builder.build(&reference.unwrap()).unwrap(); + let actual = builder.build(&projected_entries).unwrap(); + assert_eq!(actual, expected); + assert!( + engine + .scan_ssts_from_manifest(&[col("region_id").eq(lit(0u64))], batch_size) + .next() + .await + .is_none() + ); +} + +/// Installs synthetic immutable file metadata without creating physical SSTs. +async fn add_manifest_scan_region( + engine: &MitoEngine, + region_id: RegionId, + files_per_region: usize, +) { + use crate::manifest::action::{RegionEdit, RegionMetaAction, RegionMetaActionList}; + use crate::sst::file::FileMeta; + use store_api::storage::FileId; + engine + .handle_request( + region_id, + RegionRequest::Create(CreateRequestBuilder::new().build()), + ) + .await + .unwrap(); + let region = engine.get_region(region_id).unwrap(); + let edit = RegionEdit { + files_to_add: (0..files_per_region) + .map(|index| FileMeta { + region_id, + file_id: FileId::from_bytes( + (((region_id.region_number() as u128) << 64) | index as u128).to_be_bytes(), + ), + time_range: ( + Timestamp::new_millisecond(index as i64), + Timestamp::new_millisecond(index as i64 + 1), + ), + file_size: 1024 + index as u64, + ..Default::default() + }) + .collect(), + files_to_remove: vec![], + timestamp_ms: None, + compaction_time_window: None, + flushed_entry_id: None, + flushed_sequence: None, + committed_sequence: None, + }; + region + .manifest_ctx + .manifest_manager + .write() + .await + .update( + RegionMetaActionList::new(vec![RegionMetaAction::Edit(edit)]), + false, + ) + .await + .unwrap(); +} + +#[tokio::test] +async fn test_manifest_scan_table_filter() { + use datafusion_expr::{col, lit}; + use futures::StreamExt; + use store_api::sst_entry::ManifestSstEntry; + + let mut env = TestEnv::new().await; + let engine = env.create_engine(MitoConfig::default()).await; + let target_regions = [RegionId::new(11, 0), RegionId::new(11, 1)]; + let unrelated_id = RegionId::new(12, 0); + for region_id in target_regions.into_iter().chain([unrelated_id]) { + add_manifest_scan_region(&engine, region_id, 3).await; + } + let mut expected = engine.all_ssts_from_manifest().await; + expected.retain(|entry| entry.table_id == 11); + expected.sort_unstable_by_key(|entry| (entry.region_id, entry.file_id.clone())); + assert_eq!(expected.len(), 6); + + // A table predicate must exclude other tables before locking their manifests, + // while keeping every region of the selected table (including the same number). + let unrelated = engine.get_region(unrelated_id).unwrap(); + let guard = unrelated.manifest_ctx.manifest_manager.write().await; + let batches = tokio::time::timeout( + Duration::from_secs(5), + engine + .scan_ssts_from_manifest(&[col("table_id").eq(lit(11u32))], 2) + .collect::>(), + ) + .await + .expect("table filtering must not acquire another table's manifest lock"); + assert!(batches.iter().all(|batch| batch.len() <= 2)); + let mut actual = batches.into_iter().flatten().collect::>(); + actual.sort_unstable_by_key(|entry| (entry.region_id, entry.file_id.clone())); + assert_eq!(actual, expected); + + let schema = ManifestSstEntry::schema(); + let projection = ["region_id", "file_id"] + .map(|name| { + schema + .column_schemas() + .iter() + .position(|column| column.name == name) + .unwrap() + }) + .to_vec(); + let batches = tokio::time::timeout( + Duration::from_secs(5), + engine + .scan_ssts_from_manifest_projected( + &[col("table_id").eq(lit(11u32))], + 2, + Some(&projection), + ) + .collect::>(), + ) + .await + .expect("hidden table predicate must prune unrelated manifests"); + let mut actual = batches.into_iter().flatten().collect::>(); + actual.sort_unstable_by_key(|entry| (entry.region_id, entry.file_id.clone())); + let builder = store_api::sst_entry::ManifestSstBatchBuilder::new(Arc::new( + ManifestSstEntry::schema() + .arrow_schema() + .project(&projection) + .unwrap(), + )); + assert_eq!( + builder.build(&actual).unwrap(), + builder.build(&expected).unwrap() + ); + drop(guard); +} diff --git a/src/mito2/src/engine/staging_test.rs b/src/mito2/src/engine/staging_test.rs index 1b2a7ffee13..98973de4f30 100644 --- a/src/mito2/src/engine/staging_test.rs +++ b/src/mito2/src/engine/staging_test.rs @@ -843,6 +843,104 @@ async fn test_staging_exit_success_with_manifests_with_format(flat_format: bool) assert_eq!(sst_entries.iter().filter(|e| e.visible).count(), 1); assert_eq!(sst_entries.iter().filter(|e| !e.visible).count(), 2); + // Streaming inspection must retain staging entries and their visibility. + { + use futures::StreamExt; + + let mut batches = engine.scan_ssts_from_manifest( + &[datafusion_expr::col("region_id").eq(datafusion_expr::lit(region_id.as_u64()))], + 1, + ); + let mut actual = Vec::new(); + while let Some(batch) = batches.next().await { + assert_eq!(batch.len(), 1); + actual.extend(batch); + } + actual.sort_unstable_by(|a, b| a.file_id.cmp(&b.file_id)); + let mut expected = sst_entries; + expected.sort_unstable_by(|a, b| a.file_id.cmp(&b.file_id)); + assert_eq!(actual, expected); + } + + // Staging metadata overrides the same file in the normal manifest. Filter + // evaluation must use that winning metadata, even if it is not projected. + { + use futures::StreamExt; + + let region = engine.get_region(region_id).unwrap(); + let original = region + .manifest_ctx + .manifest_manager + .read() + .await + .manifest() + .files + .values() + .next() + .unwrap() + .clone(); + let mut overridden = original.clone(); + overridden.file_size = u64::MAX - 1; + for metadata in [overridden, original.clone()] { + region + .manifest_ctx + .manifest_manager + .write() + .await + .update( + RegionMetaActionList::new(vec![RegionMetaAction::Edit(RegionEdit { + files_to_add: vec![metadata.clone()], + files_to_remove: vec![], + timestamp_ms: None, + compaction_time_window: None, + flushed_entry_id: None, + flushed_sequence: None, + committed_sequence: None, + })]), + true, + ) + .await + .unwrap(); + let entries = engine.all_ssts_from_manifest().await; + assert_eq!(entries.len(), 3); + let matches = entries + .iter() + .filter(|entry| entry.file_id == original.file_id.to_string()) + .collect::>(); + assert_eq!(matches.len(), 1); + assert_eq!(matches[0].file_size, metadata.file_size); + assert!(matches[0].visible); + + let schema = store_api::sst_entry::ManifestSstEntry::schema(); + let projection = ["file_id", "visible"] + .map(|name| { + schema + .column_schemas() + .iter() + .position(|c| c.name == name) + .unwrap() + }) + .to_vec(); + let mut stream = engine.scan_ssts_from_manifest_projected( + &[datafusion_expr::col("file_size").eq(datafusion_expr::lit(metadata.file_size))], + 1, + Some(&projection), + ); + let mut filtered = Vec::new(); + while let Some(batch) = stream.next().await { + filtered.extend(batch); + } + assert!( + filtered + .iter() + .any(|entry| entry.file_id == original.file_id.to_string()) + ); + if metadata.file_size == u64::MAX - 1 { + assert_eq!(filtered.len(), 1); + } + } + } + // Exit staging mode successfully engine .set_region_role_state_gracefully(region_id, SettableRegionRoleState::Leader) diff --git a/src/mito2/src/region.rs b/src/mito2/src/region.rs index ba32f7aff56..613b1ce428e 100644 --- a/src/mito2/src/region.rs +++ b/src/mito2/src/region.rs @@ -20,6 +20,7 @@ pub mod options; pub mod utils; pub(crate) mod version; +use std::cell::OnceCell; use std::collections::hash_map::Entry; use std::collections::{HashMap, HashSet}; use std::sync::atomic::{AtomicI64, AtomicU64, Ordering}; @@ -27,8 +28,12 @@ use std::sync::{Arc, Mutex, RwLock}; use common_base::hash::partition_expr_version; use common_recordbatch::adapter::RegionQueryStatCounters; +use common_recordbatch::filter::MetadataFilter; use common_telemetry::{error, info, warn}; +use common_time::timestamp::TimeUnit; use crossbeam_utils::atomic::AtomicCell; +use datafusion_common::ScalarValue; +use futures::stream::{BoxStream, StreamExt}; use object_store::ObjectStore; use partition::expr::PartitionExpr; use snafu::{OptionExt, ResultExt, ensure}; @@ -42,7 +47,7 @@ use store_api::region_engine::{ }; use store_api::region_info::RegionInfoEntry; use store_api::region_request::{PathType, StagingPartitionDirective}; -use store_api::sst_entry::ManifestSstEntry; +use store_api::sst_entry::{self, ManifestSstEntry}; use store_api::storage::{FileId, RegionId, SequenceNumber}; use tokio::sync::RwLockWriteGuard; pub use utils::*; @@ -823,58 +828,85 @@ impl MitoRegion { } /// Returns the SST entries of the region. - pub async fn manifest_sst_entries(&self) -> Vec { - let table_dir = self.table_dir(); - let path_type = self.access_layer.path_type(); - - let visible_ssts = self - .version() - .ssts - .levels() - .iter() - .flat_map(|level| level.files().map(|file| file.file_id().file_id())) - .collect::>(); - - let manifest_files = self.manifest_ctx.manifest().await.files.clone(); - let staging_files = self - .manifest_ctx - .staging_manifest() + pub(crate) async fn manifest_sst_entries(&self) -> Vec { + self.manifest_sst_stream(MetadataFilter::default(), ManifestSstProjection::new(None)) + .collect() .await - .map(|m| m.files.clone()) - .unwrap_or_default(); - let files = manifest_files - .into_iter() - .chain(staging_files) - .collect::>(); + } - files - .values() - .map(|meta| { + /// Traverses manifest snapshots without copying the file maps. Staging metadata + /// overrides the normal manifest for the same file. + pub(crate) fn manifest_sst_stream( + &self, + filters: MetadataFilter, + projection: ManifestSstProjection, + ) -> BoxStream<'_, ManifestSstEntry> { + Box::pin(async_stream::stream! { + let table_dir = self.table_dir(); + let path_type = self.access_layer.path_type(); + let visible_files = self.manifest_visible_file_ids(projection); + let (manifest, staging) = { + let manager = self.manifest_ctx.manifest_manager.read().await; + (manager.manifest(), manager.staging_manifest()) + }; + let normal_files = manifest.files.values().map(|meta| (false, meta)); + let staging_files = staging.iter().flat_map(|s| s.files.values()).map(|meta| (true, meta)); + for (index, (is_staging, meta)) in normal_files.chain(staging_files).enumerate() { + // A selective filter may otherwise traverse millions of files in + // a single poll without yielding any rows. + if index > 0 && index % 1024 == 0 { + tokio::task::yield_now().await; + } + if !is_staging && staging.as_ref().is_some_and(|s| s.files.contains_key(&meta.file_id)) { + continue; + } + if !manifest_file_matches(meta, &filters) { + continue; + } + // Paths describe the creating region, not the referring region. + // Construct them only when a path predicate or projection needs them. + let file_path = OnceCell::new(); + if !filters.matches(|column| match column { + sst_entry::FILE_PATH => Some(file_path.get_or_init(|| { + ScalarValue::Utf8(Some(sst_file_path(table_dir, meta.file_id(), path_type))) + })), + _ => None, + }) { + continue; + } let region_id = self.region_id; let origin_region_id = meta.region_id; let (index_version, index_file_path, index_file_size) = if meta.index_file_size > 0 { - let index_file_path = index_file_path(table_dir, meta.index_id(), path_type); + let index_file_path = projection.index_file_path.then(|| index_file_path(table_dir, meta.index_id(), path_type)); ( meta.index_version, - Some(index_file_path), + index_file_path, Some(meta.index_file_size), ) } else { (0, None, None) }; - let visible = visible_ssts.contains(&meta.file_id); - ManifestSstEntry { - table_dir: table_dir.to_string(), + let visible = visible_files.as_ref().is_some_and(|files| files.contains(&meta.file_id)); + let file_path = if projection.file_path { + match file_path.into_inner() { + Some(ScalarValue::Utf8(Some(path))) => path, + _ => sst_file_path(table_dir, meta.file_id(), path_type), + } + } else { + String::new() + }; + yield ManifestSstEntry { + table_dir: if projection.table_dir { table_dir.to_string() } else { String::new() }, region_id, table_id: region_id.table_id(), region_number: region_id.region_number(), region_group: region_id.region_group(), region_sequence: region_id.region_sequence(), - file_id: meta.file_id.to_string(), + file_id: if projection.file_id { meta.file_id.to_string() } else { String::new() }, index_version, level: meta.level, - file_path: sst_file_path(table_dir, meta.file_id(), path_type), + file_path, file_size: meta.file_size, max_row_group_uncompressed_size: meta.max_row_group_uncompressed_size, index_file_path, @@ -885,15 +917,31 @@ impl MitoRegion { min_ts: meta.time_range.0, max_ts: meta.time_range.1, sequence: meta.sequence.map(|s| s.get()), - partition_expr: meta.partition_expr.as_ref().map(ToString::to_string), + partition_expr: if projection.partition_expr { meta.partition_expr.as_ref().map(ToString::to_string) } else { None }, origin_region_id, node_id: None, visible, - primary_key_min: meta.primary_key_min.clone(), - primary_key_max: meta.primary_key_max.clone(), - } - }) - .collect() + primary_key_min: if projection.primary_key_min { meta.primary_key_min.clone() } else { None }, + primary_key_max: if projection.primary_key_max { meta.primary_key_max.clone() } else { None }, + }; + } + }) + } + + /// Copies only file IDs so a paused inspection cannot delay SST purging. + /// Hidden residual columns are included in the query's projection. + fn manifest_visible_file_ids( + &self, + projection: ManifestSstProjection, + ) -> Option> { + projection.visible.then(|| { + self.version() + .ssts + .levels() + .iter() + .flat_map(|level| level.files().map(|file| file.file_id().file_id())) + .collect() + }) } /// Returns the file metas of the region by file ids. @@ -1993,6 +2041,75 @@ pub fn parse_partition_expr(partition_expr_str: Option<&str>) -> Result) -> Self { + let schema = ManifestSstEntry::schema(); + let included = |name: &str| { + projection.is_none_or(|indices| { + indices.iter().any(|index| { + schema + .column_schemas() + .get(*index) + .is_some_and(|column| column.name == name) + }) + }) + }; + Self { + visible: included(sst_entry::VISIBLE), + table_dir: included(sst_entry::TABLE_DIR), + file_id: included(sst_entry::FILE_ID), + file_path: included(sst_entry::FILE_PATH), + index_file_path: included(sst_entry::INDEX_FILE_PATH), + partition_expr: included(sst_entry::PARTITION_EXPR), + primary_key_min: included(sst_entry::PRIMARY_KEY_MIN), + primary_key_max: included(sst_entry::PRIMARY_KEY_MAX), + } + } +} + +/// Applies supported file predicates before allocating an inspection row. Unknown +/// predicates and evaluation failures are left to the query's residual filter. +fn manifest_file_matches(meta: &FileMeta, filters: &MetadataFilter) -> bool { + filters.matches(|column| { + Some(match column { + sst_entry::FILE_SIZE => ScalarValue::UInt64(Some(meta.file_size)), + sst_entry::LEVEL => ScalarValue::UInt8(Some(meta.level)), + sst_entry::NUM_ROWS => ScalarValue::UInt64(Some(meta.num_rows)), + sst_entry::NUM_ROW_GROUPS => ScalarValue::UInt64(Some(meta.num_row_groups)), + sst_entry::ORIGIN_REGION_ID => ScalarValue::UInt64(Some(meta.region_id.as_u64())), + sst_entry::MIN_TS => ScalarValue::TimestampNanosecond( + meta.time_range + .0 + .convert_to(TimeUnit::Nanosecond) + .map(|ts| ts.value()), + None, + ), + sst_entry::MAX_TS => ScalarValue::TimestampNanosecond( + meta.time_range + .1 + .convert_to(TimeUnit::Nanosecond) + .map(|ts| ts.value()), + None, + ), + _ => return None, + }) + }) +} + #[cfg(test)] mod tests { use std::sync::Arc; @@ -2024,6 +2141,66 @@ mod tests { use crate::test_util::version_util::VersionControlBuilder; use crate::time_provider::StdTimeProvider; + #[test] + fn test_manifest_file_predicates() { + use common_recordbatch::filter::MetadataFilter; + use common_time::Timestamp; + use datafusion_common::ScalarValue; + use datafusion_expr::{col, lit}; + let meta = super::FileMeta { + file_size: 100, + time_range: ( + Timestamp::new_millisecond(10), + Timestamp::new_millisecond(20), + ), + ..Default::default() + }; + for (expr, expected) in [ + (col("file_size").lt(lit(101u64)), true), + (col("file_size").lt(lit(100u64)), false), + ( + col("min_ts").eq(lit(ScalarValue::TimestampNanosecond( + Some(10_000_000), + None, + ))), + true, + ), + ( + col("max_ts").lt(lit(ScalarValue::TimestampNanosecond( + Some(20_000_000), + None, + ))), + false, + ), + (col("file_id").eq(lit("unknown")), true), + ] { + let filter = MetadataFilter::new(std::slice::from_ref(&expr)); + assert_eq!( + super::manifest_file_matches(&meta, &filter), + expected, + "{expr}" + ); + } + } + + #[test] + fn test_manifest_timestamp_overflow_is_null() { + use common_recordbatch::filter::MetadataFilter; + use common_time::Timestamp; + use datafusion_expr::col; + let meta = super::FileMeta { + time_range: ( + Timestamp::new_second(i64::MAX), + Timestamp::new_second(i64::MAX), + ), + ..Default::default() + }; + let filter = MetadataFilter::new(&[col("min_ts").is_null()]); + assert!(super::manifest_file_matches(&meta, &filter)); + let filter = MetadataFilter::new(&[col("max_ts").is_not_null()]); + assert!(!super::manifest_file_matches(&meta, &filter)); + } + #[test] fn test_region_state_lock_free() { assert!(AtomicCell::::is_lock_free()); @@ -2092,6 +2269,106 @@ mod tests { } } + #[tokio::test] + async fn test_manifest_stream_does_not_pin_sst_handles() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use common_recordbatch::filter::MetadataFilter; + use datafusion_expr::{col, lit}; + use futures::StreamExt; + + use super::ManifestSstProjection; + use crate::sst::file::FileMeta; + use crate::sst::file_purger::FilePurger; + + #[derive(Debug, Default)] + struct CountingPurger(AtomicUsize); + + impl FilePurger for CountingPurger { + fn remove_file(&self, _: FileMeta, is_delete: bool, _: bool) { + if is_delete { + self.0.fetch_add(1, Ordering::SeqCst); + } + } + } + + for include_visible in [false, true] { + let env = SchedulerEnv::new().await; + let region = build_test_region(&env).await; + let purger = Arc::new(CountingPurger::default()); + let files = (0..2) + .map(|_| FileMeta { + region_id: region.region_id, + file_id: FileId::random(), + ..Default::default() + }) + .collect::>(); + let mut edit = empty_edit(); + edit.files_to_add = files.clone(); + region + .manifest_ctx + .manifest_manager + .write() + .await + .update( + RegionMetaActionList::with_action(RegionMetaAction::Edit(edit.clone())), + false, + ) + .await + .unwrap(); + region + .version_control + .apply_edit(Some(edit), &[], purger.clone()); + + let schema = store_api::sst_entry::ManifestSstEntry::schema(); + let columns = if include_visible { + vec!["file_id", "visible"] + } else { + vec!["file_id"] + }; + let projection = columns + .iter() + .map(|name| { + schema + .column_schemas() + .iter() + .position(|column| column.name == *name) + .unwrap() + }) + .collect::>(); + let projection = ManifestSstProjection::new(Some(&projection)); + let snapshot = region.manifest_visible_file_ids(projection); + assert_eq!(snapshot.is_some(), include_visible); + if let Some(snapshot) = snapshot { + assert_eq!(snapshot, files.iter().map(|file| file.file_id).collect()); + } + // A direct caller may omit a residual column. The source must not + // interpret an unmaterialized visible value as grounds for pruning. + let filters = MetadataFilter::new(&[col("visible").eq(lit(true))]); + let mut stream = region.manifest_sst_stream(filters, projection); + let first = stream.next().await.unwrap(); + if include_visible { + assert!(first.visible); + } + let mut removal = empty_edit(); + removal.files_to_remove = files; + region + .version_control + .apply_edit(Some(removal), &[], purger.clone()); + // Both final FileHandles must reach the purger while the stream is + // suspended at its first yield, rather than waiting for cancellation. + assert_eq!(purger.0.load(Ordering::SeqCst), 2); + let second = stream.next().await.unwrap(); + if include_visible { + assert!( + second.visible, + "visibility must retain the original snapshot" + ); + } + assert!(stream.next().await.is_none()); + } + } + fn empty_edit() -> RegionEdit { RegionEdit { files_to_add: Vec::new(), diff --git a/src/store-api/src/sst_entry.rs b/src/store-api/src/sst_entry.rs index c70e909ebd6..6017f26bed3 100644 --- a/src/store-api/src/sst_entry.rs +++ b/src/store-api/src/sst_entry.rs @@ -24,6 +24,7 @@ use datatypes::arrow::array::{ ArrayRef, BinaryArray, BooleanArray, TimestampMillisecondArray, TimestampNanosecondArray, UInt8Array, UInt32Array, UInt64Array, }; +use datatypes::arrow::datatypes::SchemaRef as ArrowSchemaRef; use datatypes::arrow::error::ArrowError; use datatypes::arrow_array::StringArray; use datatypes::schema::{ColumnSchema, Schema, SchemaRef}; @@ -31,6 +32,34 @@ use serde::{Deserialize, Serialize}; use crate::storage::{RegionGroup, RegionId, RegionNumber, RegionSeq, ScanRequest, TableId}; +// Manifest SST column names. +pub const TABLE_DIR: &str = "table_dir"; +pub const REGION_ID: &str = "region_id"; +pub const TABLE_ID: &str = "table_id"; +pub const REGION_NUMBER: &str = "region_number"; +pub const REGION_GROUP: &str = "region_group"; +pub const REGION_SEQUENCE: &str = "region_sequence"; +pub const FILE_ID: &str = "file_id"; +pub const INDEX_VERSION: &str = "index_version"; +pub const LEVEL: &str = "level"; +pub const FILE_PATH: &str = "file_path"; +pub const FILE_SIZE: &str = "file_size"; +pub const INDEX_FILE_PATH: &str = "index_file_path"; +pub const INDEX_FILE_SIZE: &str = "index_file_size"; +pub const NUM_ROWS: &str = "num_rows"; +pub const NUM_ROW_GROUPS: &str = "num_row_groups"; +pub const NUM_SERIES: &str = "num_series"; +pub const MIN_TS: &str = "min_ts"; +pub const MAX_TS: &str = "max_ts"; +pub const SEQUENCE: &str = "sequence"; +pub const ORIGIN_REGION_ID: &str = "origin_region_id"; +pub const NODE_ID: &str = "node_id"; +pub const VISIBLE: &str = "visible"; +pub const PRIMARY_KEY_MIN: &str = "primary_key_min"; +pub const PRIMARY_KEY_MAX: &str = "primary_key_max"; +pub const MAX_ROW_GROUP_UNCOMPRESSED_SIZE: &str = "max_row_group_uncompressed_size"; +pub const PARTITION_EXPR: &str = "partition_expr"; + /// Puffin index type for bloom filter indexes. pub const PUFFIN_INDEX_TYPE_BLOOM_FILTER: &str = "bloom_filter"; /// Puffin index type for bloom-backed full-text indexes. @@ -102,110 +131,55 @@ impl ManifestSstEntry { pub fn schema() -> SchemaRef { use datatypes::prelude::ConcreteDataType as Ty; Arc::new(Schema::new(vec![ - ColumnSchema::new("table_dir", Ty::string_datatype(), false), - ColumnSchema::new("region_id", Ty::uint64_datatype(), false), - ColumnSchema::new("table_id", Ty::uint32_datatype(), false), - ColumnSchema::new("region_number", Ty::uint32_datatype(), false), - ColumnSchema::new("region_group", Ty::uint8_datatype(), false), - ColumnSchema::new("region_sequence", Ty::uint32_datatype(), false), - ColumnSchema::new("file_id", Ty::string_datatype(), false), - ColumnSchema::new("index_version", Ty::uint64_datatype(), false), - ColumnSchema::new("level", Ty::uint8_datatype(), false), - ColumnSchema::new("file_path", Ty::string_datatype(), false), - ColumnSchema::new("file_size", Ty::uint64_datatype(), false), - ColumnSchema::new("index_file_path", Ty::string_datatype(), true), - ColumnSchema::new("index_file_size", Ty::uint64_datatype(), true), - ColumnSchema::new("num_rows", Ty::uint64_datatype(), false), - ColumnSchema::new("num_row_groups", Ty::uint64_datatype(), false), - ColumnSchema::new("num_series", Ty::uint64_datatype(), true), - ColumnSchema::new("min_ts", Ty::timestamp_nanosecond_datatype(), true), - ColumnSchema::new("max_ts", Ty::timestamp_nanosecond_datatype(), true), - ColumnSchema::new("sequence", Ty::uint64_datatype(), true), - ColumnSchema::new("origin_region_id", Ty::uint64_datatype(), false), - ColumnSchema::new("node_id", Ty::uint64_datatype(), true), - ColumnSchema::new("visible", Ty::boolean_datatype(), false), - ColumnSchema::new("primary_key_min", Ty::binary_datatype(), true), - ColumnSchema::new("primary_key_max", Ty::binary_datatype(), true), + ColumnSchema::new(TABLE_DIR, Ty::string_datatype(), false), + ColumnSchema::new(REGION_ID, Ty::uint64_datatype(), false), + ColumnSchema::new(TABLE_ID, Ty::uint32_datatype(), false), + ColumnSchema::new(REGION_NUMBER, Ty::uint32_datatype(), false), + ColumnSchema::new(REGION_GROUP, Ty::uint8_datatype(), false), + ColumnSchema::new(REGION_SEQUENCE, Ty::uint32_datatype(), false), + ColumnSchema::new(FILE_ID, Ty::string_datatype(), false), + ColumnSchema::new(INDEX_VERSION, Ty::uint64_datatype(), false), + ColumnSchema::new(LEVEL, Ty::uint8_datatype(), false), + ColumnSchema::new(FILE_PATH, Ty::string_datatype(), false), + ColumnSchema::new(FILE_SIZE, Ty::uint64_datatype(), false), + ColumnSchema::new(INDEX_FILE_PATH, Ty::string_datatype(), true), + ColumnSchema::new(INDEX_FILE_SIZE, Ty::uint64_datatype(), true), + ColumnSchema::new(NUM_ROWS, Ty::uint64_datatype(), false), + ColumnSchema::new(NUM_ROW_GROUPS, Ty::uint64_datatype(), false), + ColumnSchema::new(NUM_SERIES, Ty::uint64_datatype(), true), + ColumnSchema::new(MIN_TS, Ty::timestamp_nanosecond_datatype(), true), + ColumnSchema::new(MAX_TS, Ty::timestamp_nanosecond_datatype(), true), + ColumnSchema::new(SEQUENCE, Ty::uint64_datatype(), true), + ColumnSchema::new(ORIGIN_REGION_ID, Ty::uint64_datatype(), false), + ColumnSchema::new(NODE_ID, Ty::uint64_datatype(), true), + ColumnSchema::new(VISIBLE, Ty::boolean_datatype(), false), + ColumnSchema::new(PRIMARY_KEY_MIN, Ty::binary_datatype(), true), + ColumnSchema::new(PRIMARY_KEY_MAX, Ty::binary_datatype(), true), ColumnSchema::new( - "max_row_group_uncompressed_size", + MAX_ROW_GROUP_UNCOMPRESSED_SIZE, Ty::uint64_datatype(), false, ), - ColumnSchema::new("partition_expr", Ty::string_datatype(), true), + ColumnSchema::new(PARTITION_EXPR, Ty::string_datatype(), true), ])) } /// Converts a list of manifest sst entries to a record batch. pub fn to_record_batch(entries: &[Self]) -> std::result::Result { - let schema = Self::schema(); - let table_dirs = entries.iter().map(|e| e.table_dir.as_str()); - let region_ids = entries.iter().map(|e| e.region_id.as_u64()); - let table_ids = entries.iter().map(|e| e.table_id); - let region_numbers = entries.iter().map(|e| e.region_number); - let region_groups = entries.iter().map(|e| e.region_group); - let region_sequences = entries.iter().map(|e| e.region_sequence); - let file_ids = entries.iter().map(|e| e.file_id.as_str()); - let index_versions = entries.iter().map(|e| e.index_version); - let levels = entries.iter().map(|e| e.level); - let file_paths = entries.iter().map(|e| e.file_path.as_str()); - let file_sizes = entries.iter().map(|e| e.file_size); - let max_row_group_uncompressed_sizes = - entries.iter().map(|e| e.max_row_group_uncompressed_size); - let index_file_paths = entries.iter().map(|e| e.index_file_path.as_ref()); - let index_file_sizes = entries.iter().map(|e| e.index_file_size); - let num_rows = entries.iter().map(|e| e.num_rows); - let num_row_groups = entries.iter().map(|e| e.num_row_groups); - let num_series = entries.iter().map(|e| e.num_series); - let min_ts = entries.iter().map(|e| { - e.min_ts - .convert_to(TimeUnit::Nanosecond) - .map(|ts| ts.value()) - }); - let max_ts = entries.iter().map(|e| { - e.max_ts - .convert_to(TimeUnit::Nanosecond) - .map(|ts| ts.value()) - }); - let sequences = entries.iter().map(|e| e.sequence); - let partition_exprs = entries.iter().map(|e| e.partition_expr.as_ref()); - let origin_region_ids = entries.iter().map(|e| e.origin_region_id.as_u64()); - let node_ids = entries.iter().map(|e| e.node_id); - let visible_flags = entries.iter().map(|e| Some(e.visible)); - let primary_key_min = entries.iter().map(|e| e.primary_key_min.as_deref()); - let primary_key_max = entries.iter().map(|e| e.primary_key_max.as_deref()); + Self::to_projected_record_batch(entries, None) + } - let columns: Vec = vec![ - Arc::new(StringArray::from_iter_values(table_dirs)), - Arc::new(UInt64Array::from_iter_values(region_ids)), - Arc::new(UInt32Array::from_iter_values(table_ids)), - Arc::new(UInt32Array::from_iter_values(region_numbers)), - Arc::new(UInt8Array::from_iter_values(region_groups)), - Arc::new(UInt32Array::from_iter_values(region_sequences)), - Arc::new(StringArray::from_iter_values(file_ids)), - Arc::new(UInt64Array::from_iter(index_versions)), - Arc::new(UInt8Array::from_iter_values(levels)), - Arc::new(StringArray::from_iter_values(file_paths)), - Arc::new(UInt64Array::from_iter_values(file_sizes)), - Arc::new(StringArray::from_iter(index_file_paths)), - Arc::new(UInt64Array::from_iter(index_file_sizes)), - Arc::new(UInt64Array::from_iter_values(num_rows)), - Arc::new(UInt64Array::from_iter_values(num_row_groups)), - Arc::new(UInt64Array::from_iter(num_series)), - Arc::new(TimestampNanosecondArray::from_iter(min_ts)), - Arc::new(TimestampNanosecondArray::from_iter(max_ts)), - Arc::new(UInt64Array::from_iter(sequences)), - Arc::new(UInt64Array::from_iter_values(origin_region_ids)), - Arc::new(UInt64Array::from_iter(node_ids)), - Arc::new(BooleanArray::from_iter(visible_flags)), - Arc::new(BinaryArray::from_iter(primary_key_min)), - Arc::new(BinaryArray::from_iter(primary_key_max)), - Arc::new(UInt64Array::from_iter_values( - max_row_group_uncompressed_sizes, - )), - Arc::new(StringArray::from_iter(partition_exprs)), - ]; - - DfRecordBatch::try_new(schema.arrow_schema().clone(), columns) + /// Builds only the requested columns, including empty and repeated projections. + fn to_projected_record_batch( + entries: &[Self], + projection: Option<&[usize]>, + ) -> std::result::Result { + let schema = Self::schema().arrow_schema().clone(); + let schema = match projection { + Some(projection) => Arc::new(schema.project(projection)?), + None => schema, + }; + ManifestSstBatchBuilder::new(schema).build(entries) } /// Reserved internal inspect table name. @@ -227,6 +201,140 @@ impl ManifestSstEntry { } } +/// Converts manifest SST batches using a schema prepared once per scan. +/// Repeated projected columns share the same Arrow array. +pub struct ManifestSstBatchBuilder { + schema: ArrowSchemaRef, + duplicates: Vec>, +} + +impl ManifestSstBatchBuilder { + /// Creates a builder for a projected manifest SST schema. + pub fn new(schema: ArrowSchemaRef) -> Self { + let duplicates = schema + .fields() + .iter() + .enumerate() + .map(|(index, field)| { + schema.fields()[..index] + .iter() + .position(|previous| previous == field) + }) + .collect(); + Self { schema, duplicates } + } + + /// Builds the projected columns, copying string contents only into Arrow buffers. + pub fn build( + &self, + entries: &[ManifestSstEntry], + ) -> std::result::Result { + let mut columns: Vec = Vec::with_capacity(self.schema.fields().len()); + for (index, field) in self.schema.fields().iter().enumerate() { + if let Some(previous) = self.duplicates[index] { + columns.push(columns[previous].clone()); + continue; + } + let array: ArrayRef = match field.name().as_str() { + TABLE_DIR => Arc::new(StringArray::from_iter_values( + entries.iter().map(|e| e.table_dir.as_str()), + )), + REGION_ID => Arc::new(UInt64Array::from_iter_values( + entries.iter().map(|e| e.region_id.as_u64()), + )), + TABLE_ID => Arc::new(UInt32Array::from_iter_values( + entries.iter().map(|e| e.table_id), + )), + REGION_NUMBER => Arc::new(UInt32Array::from_iter_values( + entries.iter().map(|e| e.region_number), + )), + REGION_GROUP => Arc::new(UInt8Array::from_iter_values( + entries.iter().map(|e| e.region_group), + )), + REGION_SEQUENCE => Arc::new(UInt32Array::from_iter_values( + entries.iter().map(|e| e.region_sequence), + )), + FILE_ID => Arc::new(StringArray::from_iter_values( + entries.iter().map(|e| e.file_id.as_str()), + )), + INDEX_VERSION => Arc::new(UInt64Array::from_iter( + entries.iter().map(|e| e.index_version), + )), + LEVEL => Arc::new(UInt8Array::from_iter_values( + entries.iter().map(|e| e.level), + )), + FILE_PATH => Arc::new(StringArray::from_iter_values( + entries.iter().map(|e| e.file_path.as_str()), + )), + FILE_SIZE => Arc::new(UInt64Array::from_iter_values( + entries.iter().map(|e| e.file_size), + )), + INDEX_FILE_PATH => Arc::new(StringArray::from_iter( + entries.iter().map(|e| e.index_file_path.as_ref()), + )), + INDEX_FILE_SIZE => Arc::new(UInt64Array::from_iter( + entries.iter().map(|e| e.index_file_size), + )), + NUM_ROWS => Arc::new(UInt64Array::from_iter_values( + entries.iter().map(|e| e.num_rows), + )), + NUM_ROW_GROUPS => Arc::new(UInt64Array::from_iter_values( + entries.iter().map(|e| e.num_row_groups), + )), + NUM_SERIES => { + Arc::new(UInt64Array::from_iter(entries.iter().map(|e| e.num_series))) + } + MIN_TS => Arc::new(TimestampNanosecondArray::from_iter(entries.iter().map( + |e| { + e.min_ts + .convert_to(TimeUnit::Nanosecond) + .map(|ts| ts.value()) + }, + ))), + MAX_TS => Arc::new(TimestampNanosecondArray::from_iter(entries.iter().map( + |e| { + e.max_ts + .convert_to(TimeUnit::Nanosecond) + .map(|ts| ts.value()) + }, + ))), + SEQUENCE => Arc::new(UInt64Array::from_iter(entries.iter().map(|e| e.sequence))), + ORIGIN_REGION_ID => Arc::new(UInt64Array::from_iter_values( + entries.iter().map(|e| e.origin_region_id.as_u64()), + )), + NODE_ID => Arc::new(UInt64Array::from_iter(entries.iter().map(|e| e.node_id))), + VISIBLE => Arc::new(BooleanArray::from_iter( + entries.iter().map(|e| Some(e.visible)), + )), + PRIMARY_KEY_MIN => Arc::new(BinaryArray::from_iter( + entries.iter().map(|e| e.primary_key_min.as_deref()), + )), + PRIMARY_KEY_MAX => Arc::new(BinaryArray::from_iter( + entries.iter().map(|e| e.primary_key_max.as_deref()), + )), + MAX_ROW_GROUP_UNCOMPRESSED_SIZE => Arc::new(UInt64Array::from_iter_values( + entries.iter().map(|e| e.max_row_group_uncompressed_size), + )), + PARTITION_EXPR => Arc::new(StringArray::from_iter( + entries.iter().map(|e| e.partition_expr.as_ref()), + )), + name => { + return Err(ArrowError::SchemaError(format!( + "Unknown manifest SST column: {name}" + ))); + } + }; + columns.push(array); + } + DfRecordBatch::try_new_with_options( + self.schema.clone(), + columns, + &datatypes::arrow::record_batch::RecordBatchOptions::new() + .with_row_count(Some(entries.len())), + ) + } +} + /// An entry describing a SST file listed from storage layer directly. #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] pub struct StorageSstEntry { @@ -421,13 +529,25 @@ fn build_plan_helper( ) -> Result { let table_source = LogicalTableSource::new(schema.arrow_schema().clone()); + // Filters may reference columns omitted from the requested output. Apply + // them to the full schema and let the optimizer push the minimal projection. let projection = scan_request.projection; - let mut builder = LogicalPlanBuilder::scan(table_name, Arc::new(table_source), projection)?; + let mut builder = LogicalPlanBuilder::scan(table_name, Arc::new(table_source), None)?; for filter in scan_request.filters { builder = builder.filter(filter)?; } + if let Some(projection) = projection { + let projected = schema.arrow_schema().project(&projection)?; + builder = builder.project( + projected + .fields() + .iter() + .map(|field| datafusion_expr::col(field.name())), + )?; + } + if let Some(limit) = scan_request.limit { builder = builder.limit(0, Some(limit))?; } @@ -438,7 +558,7 @@ fn build_plan_helper( #[cfg(test)] mod tests { use datafusion_common::TableReference; - use datafusion_expr::{LogicalPlan, Operator, binary_expr, col, lit}; + use datafusion_expr::{Expr, LogicalPlan, Operator, binary_expr, col, lit}; use datatypes::arrow::array::{ Array, BinaryArray, TimestampMillisecondArray, TimestampNanosecondArray, UInt8Array, UInt32Array, UInt64Array, @@ -523,6 +643,29 @@ mod tests { let schema = ManifestSstEntry::schema(); let batch = ManifestSstEntry::to_record_batch(&entries).unwrap(); + for projection in [vec![], vec![16, 6, 10], vec![1, 1], (0..26).rev().collect()] { + let projected = + ManifestSstEntry::to_projected_record_batch(&entries, Some(&projection)).unwrap(); + assert_eq!(projected, batch.project(&projection).unwrap()); + assert_eq!(projected.num_rows(), entries.len()); + } + assert!(ManifestSstEntry::to_projected_record_batch(&entries, Some(&[26])).is_err()); + + let projected_schema = Arc::new(batch.schema().project(&[6, 9, 6, 11, 25, 9]).unwrap()); + let builder = ManifestSstBatchBuilder::new(projected_schema.clone()); + for rows in [entries.as_slice(), &entries[..1], &entries[..0]] { + let projected = builder.build(rows).unwrap(); + assert_eq!( + projected, + ManifestSstEntry::to_record_batch(rows) + .unwrap() + .project(&[6, 9, 6, 11, 25, 9]) + .unwrap() + ); + assert!(Arc::ptr_eq(&projected.schema(), &projected_schema)); + assert!(Arc::ptr_eq(projected.column(0), projected.column(2))); + assert!(Arc::ptr_eq(projected.column(1), projected.column(5))); + } // Schema checks assert_eq!(schema.arrow_schema().fields().len(), batch.num_columns()); @@ -1021,35 +1164,46 @@ mod tests { #[test] fn test_manifest_build_plan() { - // Note: filter must reference a column in the projected schema - let projection = Some(vec![0, 1, 2]); let request = ScanRequest { - projection, - filters: vec![binary_expr(col("table_id"), Operator::Gt, lit(0))], + projection: Some(vec![0, 1]), + filters: vec![col("table_id").gt(lit(0u32))], limit: Some(5), ..Default::default() }; let plan = ManifestSstEntry::build_plan(request).unwrap(); - // Expect plan to be Filter -> Limit -> TableScan or Filter+Limit wrapped. - // We'll pattern match to reach TableScan and verify key fields. - let (scan, has_filter, has_limit) = extract_scan(&plan); - - assert!(has_filter); - assert!(has_limit); + // The hidden table_id predicate must run before projection and limit. + let LogicalPlan::Limit(limit) = &plan else { + panic!("expected limit") + }; + let LogicalPlan::Projection(projection) = limit.input.as_ref() else { + panic!("expected projection") + }; + let LogicalPlan::Filter(filter) = projection.input.as_ref() else { + panic!("expected filter") + }; + let LogicalPlan::TableScan(scan) = filter.input.as_ref() else { + panic!("expected table scan") + }; assert_eq!( scan.table_name, TableReference::bare(ManifestSstEntry::reserved_table_name_for_inspection()) ); - assert_eq!(scan.projection, Some(vec![0, 1, 2])); + assert_eq!(scan.projection, None); + assert_eq!( + filter.predicate, + Expr::Column(datafusion_common::Column::new( + Some(scan.table_name.clone()), + "table_id", + )) + .gt(lit(0u32)) + ); - // projected schema should match projection - let fields = scan.projected_schema.fields(); - assert_eq!(fields.len(), 3); + let fields = plan.schema().fields(); + assert_eq!(fields.len(), 2); assert_eq!(fields[0].name(), "table_dir"); assert_eq!(fields[1].name(), "region_id"); - assert_eq!(fields[2].name(), "table_id"); } #[test] @@ -1070,9 +1224,9 @@ mod tests { scan.table_name, TableReference::bare(StorageSstEntry::reserved_table_name_for_inspection()) ); - assert_eq!(scan.projection, Some(vec![0, 2])); + assert_eq!(scan.projection, None); - let fields = scan.projected_schema.fields(); + let fields = plan.schema().fields(); assert_eq!(fields.len(), 2); assert_eq!(fields[0].name(), "file_path"); assert_eq!(fields[1].name(), "last_modified_ms"); @@ -1091,6 +1245,7 @@ mod tests { let (scan, has_filter, _) = extract_scan(input); (scan, has_filter, true) } + LogicalPlan::Projection(projection) => extract_scan(&projection.input), LogicalPlan::TableScan(scan) => (scan, false, false), other => panic!("unexpected plan: {other:?}"), } diff --git a/tests/cases/distributed/information_schema/ssts_limit.result b/tests/cases/distributed/information_schema/ssts_limit.result index 1b21a84c7af..67c6abf3cb1 100644 --- a/tests/cases/distributed/information_schema/ssts_limit.result +++ b/tests/cases/distributed/information_schema/ssts_limit.result @@ -95,6 +95,59 @@ FROM ( | | | +---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +-- The frontend orders results from all datanodes before applying the global offset. +SELECT min_ts, max_ts +FROM information_schema.ssts_manifest +WHERE region_id > 0 +ORDER BY min_ts DESC +LIMIT 2 OFFSET 1; + ++-------------------------+-------------------------+ +| min_ts | max_ts | ++-------------------------+-------------------------+ +| 1970-01-01T00:00:00.002 | 1970-01-01T00:00:00.002 | +| 1970-01-01T00:00:00.001 | 1970-01-01T00:00:00.001 | ++-------------------------+-------------------------+ + +-- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED +EXPLAIN SELECT min_ts +FROM information_schema.ssts_manifest +WHERE region_id > 0 +ORDER BY min_ts DESC +LIMIT 1; + ++---------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| plan_type | plan | ++---------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| logical_plan | MergeScan [is_placeholder=false, remote_input=[ | +| | Limit: skip=0, fetch=1 | +| | Sort: information_schema.ssts_manifest.min_ts DESC NULLS FIRST | +| | Projection: information_schema.ssts_manifest.min_ts | +| | Filter: information_schema.ssts_manifest.region_id > UInt64(0) | +| | TableScan: information_schema.ssts_manifest, partial_filters=[information_schema.ssts_manifest.region_id > UInt64(0)] | +| | ]] | +| physical_plan | SortPreservingMergeExec: [min_ts@0 DESC], fetch=1 | +| | SortExec: TopK(fetch=1), expr=[min_ts@0 DESC], preserve_partitioning=[true] | +| | FilterExec: region_id@0 > 0, projection=[min_ts@1] | +| | RepartitionExec: REDACTED +| | DistributedInspectExec: kind=SstManifest, scan=ScanRequest { projection: [1, 16], filters: [region_id > UInt64(0)] }, schema=Schema { fields: [Field { name: "region_id", data_type: UInt64 }, Field { name: "min_ts", data_type: Timestamp(Nanosecond, None), nullable: true }], metadata: {"greptime:version": "0"} } | +| | | ++---------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ + +-- Timestamp filtering happens before global ordering and offset. The sort key +-- must remain available even when it is absent from the output projection. +SELECT max_ts +FROM information_schema.ssts_manifest +WHERE region_id > 0 AND min_ts >= '1970-01-01T00:00:00.002Z'::TIMESTAMP_NS +ORDER BY min_ts DESC +LIMIT 1 OFFSET 1; + ++-------------------------+ +| max_ts | ++-------------------------+ +| 1970-01-01T00:00:00.002 | ++-------------------------+ + DROP TABLE ssts_limit_case; Affected Rows: 0 diff --git a/tests/cases/distributed/information_schema/ssts_limit.sql b/tests/cases/distributed/information_schema/ssts_limit.sql index d548bd283a5..e634f7a2378 100644 --- a/tests/cases/distributed/information_schema/ssts_limit.sql +++ b/tests/cases/distributed/information_schema/ssts_limit.sql @@ -44,4 +44,26 @@ FROM ( LIMIT 1 ); +-- The frontend orders results from all datanodes before applying the global offset. +SELECT min_ts, max_ts +FROM information_schema.ssts_manifest +WHERE region_id > 0 +ORDER BY min_ts DESC +LIMIT 2 OFFSET 1; + +-- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED +EXPLAIN SELECT min_ts +FROM information_schema.ssts_manifest +WHERE region_id > 0 +ORDER BY min_ts DESC +LIMIT 1; + +-- Timestamp filtering happens before global ordering and offset. The sort key +-- must remain available even when it is absent from the output projection. +SELECT max_ts +FROM information_schema.ssts_manifest +WHERE region_id > 0 AND min_ts >= '1970-01-01T00:00:00.002Z'::TIMESTAMP_NS +ORDER BY min_ts DESC +LIMIT 1 OFFSET 1; + DROP TABLE ssts_limit_case; diff --git a/tests/cases/standalone/common/information_schema/ssts.result b/tests/cases/standalone/common/information_schema/ssts.result index d7c7814a995..d8149d27769 100644 --- a/tests/cases/standalone/common/information_schema/ssts.result +++ b/tests/cases/standalone/common/information_schema/ssts.result @@ -243,6 +243,84 @@ FROM ( | 1 | +-----------------------+ +-- Timestamp filters must run before ordering and limiting. Sort columns must +-- remain available even when they are absent from the output projection. +SELECT min_ts, max_ts +FROM information_schema.ssts_manifest +WHERE region_id > 0 AND min_ts >= '1970-01-01T00:00:00.002Z'::TIMESTAMP_NS +ORDER BY min_ts ASC NULLS LAST, max_ts ASC +LIMIT 2; + ++-------------------------+-------------------------+ +| min_ts | max_ts | ++-------------------------+-------------------------+ +| 1970-01-01T00:00:00.002 | 1970-01-01T00:00:00.002 | +| 1970-01-01T00:00:00.003 | 1970-01-01T00:00:00.003 | ++-------------------------+-------------------------+ + +SELECT min_ts, max_ts +FROM information_schema.ssts_manifest +WHERE region_id > 0 +ORDER BY max_ts DESC NULLS FIRST, min_ts DESC +LIMIT 2 OFFSET 1; + ++-------------------------+-------------------------+ +| min_ts | max_ts | ++-------------------------+-------------------------+ +| 1970-01-01T00:00:00.100 | 1970-01-01T00:00:00.200 | +| 1970-01-01T00:00:00.003 | 1970-01-01T00:00:00.003 | ++-------------------------+-------------------------+ + +SELECT visible +FROM information_schema.ssts_manifest +WHERE region_id > 0 AND min_ts >= '1970-01-01T00:00:00.002Z'::TIMESTAMP_NS +ORDER BY min_ts DESC +LIMIT 2; + ++---------+ +| visible | ++---------+ +| true | +| true | ++---------+ + +-- Table predicates share region metadata pruning, including when table_id +-- is not part of the output projection. +SELECT min_ts, max_ts FROM information_schema.ssts_manifest +WHERE table_id > 0 ORDER BY min_ts, max_ts LIMIT 2; + ++-------------------------+-------------------------+ +| min_ts | max_ts | ++-------------------------+-------------------------+ +| 1970-01-01T00:00:00.001 | 1970-01-01T00:00:00.001 | +| 1970-01-01T00:00:00.002 | 1970-01-01T00:00:00.002 | ++-------------------------+-------------------------+ + +SELECT min_ts FROM information_schema.ssts_manifest +WHERE table_id = 0 ORDER BY min_ts LIMIT 1; + ++--------+ +| min_ts | ++--------+ ++--------+ + +-- A false region predicate and a zero limit must produce empty results. +SELECT min_ts FROM information_schema.ssts_manifest +WHERE region_id = 0 ORDER BY min_ts LIMIT 1; + ++--------+ +| min_ts | ++--------+ ++--------+ + +SELECT min_ts FROM information_schema.ssts_manifest +ORDER BY min_ts LIMIT 0; + ++--------+ +| min_ts | ++--------+ ++--------+ + DROP TABLE sst_case; Affected Rows: 0 diff --git a/tests/cases/standalone/common/information_schema/ssts.sql b/tests/cases/standalone/common/information_schema/ssts.sql index 7b40b0056bb..e630ab9a0ba 100644 --- a/tests/cases/standalone/common/information_schema/ssts.sql +++ b/tests/cases/standalone/common/information_schema/ssts.sql @@ -80,4 +80,39 @@ FROM ( LIMIT 1 ); +-- Timestamp filters must run before ordering and limiting. Sort columns must +-- remain available even when they are absent from the output projection. +SELECT min_ts, max_ts +FROM information_schema.ssts_manifest +WHERE region_id > 0 AND min_ts >= '1970-01-01T00:00:00.002Z'::TIMESTAMP_NS +ORDER BY min_ts ASC NULLS LAST, max_ts ASC +LIMIT 2; + +SELECT min_ts, max_ts +FROM information_schema.ssts_manifest +WHERE region_id > 0 +ORDER BY max_ts DESC NULLS FIRST, min_ts DESC +LIMIT 2 OFFSET 1; + +SELECT visible +FROM information_schema.ssts_manifest +WHERE region_id > 0 AND min_ts >= '1970-01-01T00:00:00.002Z'::TIMESTAMP_NS +ORDER BY min_ts DESC +LIMIT 2; + +-- Table predicates share region metadata pruning, including when table_id +-- is not part of the output projection. +SELECT min_ts, max_ts FROM information_schema.ssts_manifest +WHERE table_id > 0 ORDER BY min_ts, max_ts LIMIT 2; + +SELECT min_ts FROM information_schema.ssts_manifest +WHERE table_id = 0 ORDER BY min_ts LIMIT 1; + +-- A false region predicate and a zero limit must produce empty results. +SELECT min_ts FROM information_schema.ssts_manifest +WHERE region_id = 0 ORDER BY min_ts LIMIT 1; + +SELECT min_ts FROM information_schema.ssts_manifest +ORDER BY min_ts LIMIT 0; + DROP TABLE sst_case; From 6f5c064f543eedfa65d312bebe64f33870c3ca60 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Mon, 26 Oct 2026 09:02:40 +0000 Subject: [PATCH 095/239] fix(meta): fence datanode processes replaced by a newer process (#9878) * fix(meta): fence datanode processes replaced by a newer process Two datanode processes can run with the same node id, and metasrv grants leases and delivers instructions by node id only. Every datanode heartbeat already carries node_epoch, the time the process started. The metasrv leader records the highest accepted epoch per node id under __datanode_epoch/{node_id}, a prefix of the leader cache, and raises it with compare-and-put. The heartbeat session checks each datanode request with a non-zero epoch before the handler chain: a lower epoch gets a reply without running the chain, so no lease, keep-alive, stats or mailbox reply. An epoch more than 5 minutes ahead of the leader's clock is not recorded and is fenced too. Datanode pushers carry the epoch of their process: a lower epoch cannot replace a higher one, and a stream removes only its own pusher. Signed-off-by: jeremyhi * fix(meta): keep the leader cache load and pusher updates coherent A write made while the leader cache is disabled, as during the load at leader start, went to the store without the cache noticing, so the load could install a value read before the write. Such writes now create new cache versions around the store call, and the load drops what it loaded for a prefix if a write happened meanwhile. This also covers node addresses, which heartbeats write during the load. A stream's pusher that takes its first epoch is updated in place instead of being replaced, so messages waiting on it are not cancelled, and the connection gauge counts only a pusher that fills an empty slot. The epoch check moves into the heartbeat service. Signed-off-by: jeremyhi * fix(meta): keep the leader cache coherent with writes while it is disabled A write that started while the cache was disabled could commit after the load had read the old value and the cache had been enabled, leaving the old value cached. Writes now maintain the cache whether or not it is enabled, as enabled writes already did: new versions around the store call, then the key is updated or invalidated. Signed-off-by: jeremyhi * test(meta): pin the epoch 0 exemption and pusher order, and update docs The datanode test without an epoch now runs after another process has recorded one, and the pusher test checks that a stream's first accepted request receives what its handlers broadcast. The leader cache load no longer claims that writes must not run while it loads, and the meta-srv guide states that epoch 0 is not checked. Signed-off-by: jeremyhi * fix(meta): register a datanode pusher only after its first request passes A datanode stream registered an epoch-0 pusher on its handshake, so an older process's reconnecting stream could take an empty slot, or a newer process's pusher before its first heartbeat, for about a round trip. A datanode stream now registers its pusher on its first request that passes the epoch check (epoch-0 requests after the handshake included), still before the handlers run. Pusher registration is logged only when the insert happens. Tests cover a fenced stream removing its own pusher, a handshake not taking an empty slot, and a KV error failing the request; two steps that could not fail are removed. Signed-off-by: jeremyhi * fix(meta): check a datanode epoch again before replying A request could pass the epoch check, wait in the handler chain while a newer process of the node was recorded, and still return a region lease. After the handlers run, the session checks the epoch again against the cached record and, if a newer process has been recorded, replies as to a fenced request and removes the stream's own pusher. This costs one more read of the cached record per checked datanode request. Signed-off-by: jeremyhi * fix(meta): handle a datanode's requests one at a time A request of an older process could pass the epoch check, wait in the handlers while a newer process was admitted, and then keep the node alive with its own address. Checking the epoch again after the handlers only dropped the reply. The heartbeat session now takes an in-memory lock for the node id before checking a request that carries an epoch, and releases it after the handlers, before the reply is sent. A newer process is therefore admitted only after the handlers of an older request finish. The check after the handlers is removed. Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- src/common/meta/src/key.rs | 12 +- src/common/meta/src/key/datanode_epoch.rs | 66 ++ src/meta-srv/AGENTS.md | 8 + src/meta-srv/src/handler.rs | 95 ++- src/meta-srv/src/service/heartbeat.rs | 633 +++++++++++++++++- src/meta-srv/src/service/store/cached_kv.rs | 97 ++- .../meta/datanode_epoch/value.json | 1 + 7 files changed, 841 insertions(+), 71 deletions(-) create mode 100644 src/common/meta/src/key/datanode_epoch.rs create mode 100644 tests/compatibility/meta/datanode_epoch/value.json diff --git a/src/common/meta/src/key.rs b/src/common/meta/src/key.rs index 1ed8f8df6f0..189c427e9ec 100644 --- a/src/common/meta/src/key.rs +++ b/src/common/meta/src/key.rs @@ -65,6 +65,9 @@ //! 13. Topic name to region map key `__topic_region/{topic_name}/{region_id}` //! - Mapping {topic_name} to {region_id} //! +//! 14. Datanode epoch key: `__datanode_epoch/{datanode_id}` +//! - The highest `node_epoch` the metasrv leader has accepted from the datanode +//! //! All keys have related managers. The managers take care of the serialization and deserialization //! of keys and values, and the interaction with the underlying KV store backend. //! @@ -98,6 +101,7 @@ //! {partition_id} pub mod catalog_name; +pub mod datanode_epoch; pub mod datanode_table; pub mod flow; pub mod node_address; @@ -149,6 +153,7 @@ use topic_region::{TopicRegionKey, TopicRegionManager}; use view_info::{ViewInfoKey, ViewInfoManager, ViewInfoValue}; use self::catalog_name::{CatalogManager, CatalogNameKey, CatalogNameValue}; +use self::datanode_epoch::DatanodeEpochValue; use self::datanode_table::RegionInfo; use self::flow::flow_info::FlowInfoValue; use self::flow::flow_name::FlowNameValue; @@ -179,6 +184,7 @@ pub const PAUSE_PROCEDURE_KEY: &str = "__switches/pause_procedure"; pub const RECOVERY_MODE_KEY: &str = "__switches/recovery"; pub const DATANODE_TABLE_KEY_PREFIX: &str = "__dn_table"; +pub const DATANODE_EPOCH_KEY_PREFIX: &str = "__datanode_epoch"; pub const TABLE_INFO_KEY_PREFIX: &str = "__table_info"; pub const VIEW_INFO_KEY_PREFIX: &str = "__view_info"; pub const TABLE_NAME_KEY_PREFIX: &str = "__table_name"; @@ -198,12 +204,13 @@ pub const ELECTION_KEY: &str = "__metasrv_election"; pub const CANDIDATES_ROOT: &str = "__metasrv_election_candidates/"; /// The keys with these prefixes will be loaded into the cache when the leader starts. -pub const CACHE_KEY_PREFIXES: [&str; 5] = [ +pub const CACHE_KEY_PREFIXES: [&str; 6] = [ TABLE_NAME_KEY_PREFIX, CATALOG_NAME_KEY_PREFIX, SCHEMA_NAME_KEY_PREFIX, TABLE_ROUTE_PREFIX, NODE_ADDRESS_PREFIX, + DATANODE_EPOCH_KEY_PREFIX, ]; /// A set of regions with the same role. @@ -1886,7 +1893,8 @@ impl_metadata_value! { SchemaNameValue, FlowStateValue, PoisonValue, - TopicRegionValue + TopicRegionValue, + DatanodeEpochValue } impl_optional_metadata_value! { diff --git a/src/common/meta/src/key/datanode_epoch.rs b/src/common/meta/src/key/datanode_epoch.rs new file mode 100644 index 00000000000..9003ff3dde4 --- /dev/null +++ b/src/common/meta/src/key/datanode_epoch.rs @@ -0,0 +1,66 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::fmt::Display; + +use serde::{Deserialize, Serialize}; + +use crate::DatanodeId; +use crate::key::DATANODE_EPOCH_KEY_PREFIX; + +/// The key of the highest `node_epoch` the metasrv leader has accepted from a datanode. +/// +/// The layout: `__datanode_epoch/{node_id}` +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct DatanodeEpochKey { + pub node_id: DatanodeId, +} + +impl DatanodeEpochKey { + pub fn new(node_id: DatanodeId) -> Self { + Self { node_id } + } + + pub fn to_bytes(&self) -> Vec { + self.to_string().into_bytes() + } +} + +impl Display for DatanodeEpochKey { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}/{}", DATANODE_EPOCH_KEY_PREFIX, self.node_id) + } +} + +/// The value of [DatanodeEpochKey]. It only grows. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub struct DatanodeEpochValue { + pub node_epoch: u64, +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::key::MetadataValue; + + #[test] + fn test_frozen_datanode_epoch_record() { + let value = + include_bytes!("../../../../../tests/compatibility/meta/datanode_epoch/value.json"); + assert_eq!(DatanodeEpochKey::new(7).to_bytes(), b"__datanode_epoch/7"); + let decoded = DatanodeEpochValue::try_from_raw_value(value).unwrap(); + assert_eq!(decoded.node_epoch, 1761091200000); + assert_eq!(decoded.try_as_raw_value().unwrap(), value); + } +} diff --git a/src/meta-srv/AGENTS.md b/src/meta-srv/AGENTS.md index c1aff378496..2f758780bf1 100644 --- a/src/meta-srv/AGENTS.md +++ b/src/meta-srv/AGENTS.md @@ -43,6 +43,14 @@ abstraction, election, key encoding, and the DDL manager live in `common-meta`; `handler/utils.rs` for per-node address persistence. Datanode address updates retain the statistics handler's per-node ordering. Address epochs are confirmed only after a successful save; cache invalidation remains best effort. +- **Datanode process fence**: before the handler chain, the leader's heartbeat + session checks a datanode request's non-zero `node_epoch` against the highest + one recorded under `__datanode_epoch`, read and raised through the leader + cache. A lower epoch gets a reply without running the chain; requests that + carry an epoch take a per-node lock through the check and the chain; epoch 0 + (handshakes, older datanodes) is not checked. A datanode stream registers + its pusher on its first request that passes, not on the handshake; pushers + carry the epoch, and a lower one cannot replace a higher one. - **DDL**: `service/procedure.rs` (`ddl`) is leader-only and hands off to `common-meta`'s `DdlManager`, executed via `common-procedure`. - **Region migration**: `procedure/region_migration/manager.rs` plus per-step diff --git a/src/meta-srv/src/handler.rs b/src/meta-srv/src/handler.rs index 24c1764db2a..a452e3a79b5 100644 --- a/src/meta-srv/src/handler.rs +++ b/src/meta-srv/src/handler.rs @@ -51,7 +51,7 @@ use response_header_handler::ResponseHeaderHandler; use snafu::{OptionExt, ResultExt}; use store_api::storage::RegionId; use tokio::sync::mpsc::Sender; -use tokio::sync::{Notify, RwLock, oneshot, watch}; +use tokio::sync::{Notify, OwnedMutexGuard, RwLock, oneshot, watch}; use crate::error::{self, DeserializeFromJsonSnafu, Result, UnexpectedInstructionReplySnafu}; use crate::handler::collect_topic_stats_handler::CollectTopicStatsHandler; @@ -198,6 +198,8 @@ pub struct Pusher { deregister_signal_receiver: DeregisterSignalReceiver, res_header: ResponseHeader, + /// The `node_epoch` of the datanode process that registered the pusher, or 0. + node_epoch: u64, } impl Drop for Pusher { @@ -220,9 +222,15 @@ impl Pusher { deregister_signal_sender, deregister_signal_receiver, res_header, + node_epoch: 0, } } + pub(crate) fn with_node_epoch(mut self, node_epoch: u64) -> Self { + self.node_epoch = node_epoch; + self + } + #[inline] pub fn header(&self) -> ResponseHeader { self.res_header.clone() @@ -327,12 +335,38 @@ impl Pushers { Ok(()) } - pub(crate) async fn insert(&self, pusher_id: PusherId, pusher: Pusher) -> Option { - self.0.write().await.insert(pusher_id, pusher) + /// Inserts the pusher unless the slot holds one with a higher node epoch, so an older + /// datanode process cannot take over a newer one's pusher. A pusher of the same stream + /// only takes the new epoch, keeping the messages waiting on it. Returns `None` if refused, + /// otherwise whether the slot was empty. + pub(crate) async fn insert(&self, pusher_id: PusherId, pusher: Pusher) -> Option { + let mut pushers = self.0.write().await; + match pushers.get_mut(&pusher_id) { + None => { + pushers.insert(pusher_id, pusher); + return Some(true); + } + Some(current) if current.node_epoch > pusher.node_epoch => return None, + Some(current) if current.sender.same_channel(&pusher.sender) => { + current.node_epoch = pusher.node_epoch; + } + Some(current) => *current = pusher, + } + Some(false) } - async fn remove(&self, pusher_id: PusherId) -> Option { - self.0.write().await.remove(&pusher_id) + /// Removes the pusher if it pushes to `sender`'s channel, keeping one that another stream + /// has registered since. + async fn remove( + &self, + pusher_id: PusherId, + sender: &Sender>, + ) -> bool { + let mut pushers = self.0.write().await; + let own = pushers + .get(&pusher_id) + .is_some_and(|pusher| pusher.sender.same_channel(sender)); + own && pushers.remove(&pusher_id).is_some() } } @@ -357,24 +391,43 @@ pub type HeartbeatHandlerGroupRef = Arc; pub struct HeartbeatHandlerGroup { handlers: Vec, pushers: Pushers, + /// One lock per datanode id seen, kept for the life of the process. + datanode_locks: DashMap>>, } impl HeartbeatHandlerGroup { - /// Registers the heartbeat response [`Pusher`] with the given key to the group. + /// Registers the heartbeat response [`Pusher`] with the given key to the group, unless the + /// group holds one with a higher node epoch. pub async fn register_pusher(&self, pusher_id: PusherId, pusher: Pusher) { - METRIC_META_HEARTBEAT_CONNECTION_NUM.inc(); + let Some(vacant) = self.pushers.insert(pusher_id, pusher).await else { + return; + }; info!("Pusher register: {}", pusher_id); - let _ = self.pushers.insert(pusher_id, pusher).await; + if vacant { + METRIC_META_HEARTBEAT_CONNECTION_NUM.inc(); + } } - /// Deregisters the heartbeat response [`Pusher`] with the given key from the group. - pub async fn deregister_push(&self, pusher_id: PusherId) { - if self.pushers.remove(pusher_id).await.is_some() { + /// Deregisters the heartbeat response [`Pusher`] with the given key from the group if + /// `sender` registered it. + pub async fn deregister_push( + &self, + pusher_id: PusherId, + sender: &Sender>, + ) { + if self.pushers.remove(pusher_id, sender).await { info!("Pusher unregister: {}", pusher_id); METRIC_META_HEARTBEAT_CONNECTION_NUM.dec(); } } + /// Locks the datanode with the given id until the guard is dropped, so its requests are + /// checked and handled one at a time. + pub(crate) async fn lock_datanode(&self, node_id: u64) -> OwnedMutexGuard<()> { + let lock = self.datanode_locks.entry(node_id).or_default().clone(); + lock.lock_owned().await + } + #[cfg(test)] /// Returns whether the group contains the heartbeat response [`Pusher`] with the given key. pub async fn contains_pusher(&self, pusher_id: &PusherId) -> bool { @@ -783,6 +836,7 @@ impl HeartbeatHandlerGroupBuilder { Ok(HeartbeatHandlerGroup { handlers: self.handlers, pushers: self.pushers, + datanode_locks: DashMap::new(), }) } @@ -1477,6 +1531,25 @@ mod tests { assert!(healthy_rx.try_recv().is_ok()); } + #[tokio::test] + async fn test_pusher_insert_reports_an_empty_slot() { + let pushers = Pushers::default(); + let id = PusherId::new(Role::Datanode, 1); + let (tx1, _rx1) = mpsc::channel(1); + let (tx2, _rx2) = mpsc::channel(1); + assert_eq!( + pushers.insert(id, Pusher::new(tx1.clone())).await, + Some(true) + ); + // Taking an epoch and replacing another stream's pusher fill no slot. + let in_place = Pusher::new(tx1).with_node_epoch(2); + assert_eq!(pushers.insert(id, in_place).await, Some(false)); + let lower = Pusher::new(tx2.clone()).with_node_epoch(1); + assert_eq!(pushers.insert(id, lower).await, None); + let other = Pusher::new(tx2).with_node_epoch(2); + assert_eq!(pushers.insert(id, other).await, Some(false)); + } + #[tokio::test] async fn test_pushers_broadcast_does_not_fence_the_registry() { let pushers = Pushers::default(); diff --git a/src/meta-srv/src/service/heartbeat.rs b/src/meta-srv/src/service/heartbeat.rs index 782daa289de..c4f4978e746 100644 --- a/src/meta-srv/src/service/heartbeat.rs +++ b/src/meta-srv/src/service/heartbeat.rs @@ -12,16 +12,23 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::cmp::Ordering; use std::io::ErrorKind; use std::sync::Arc; use std::sync::atomic::AtomicU64; +use std::time::Duration; use api::v1::meta::{ - AskLeaderRequest, AskLeaderResponse, HeartbeatRequest, HeartbeatResponse, Peer, RequestHeader, - ResponseHeader, Role, heartbeat_server, + AskLeaderRequest, AskLeaderResponse, Error, HeartbeatRequest, HeartbeatResponse, + PROTOCOL_VERSION, Peer, RequestHeader, ResponseHeader, Role, heartbeat_server, }; use common_meta::election::LeaderChangeMessage; +use common_meta::key::MetadataValue; +use common_meta::key::datanode_epoch::{DatanodeEpochKey, DatanodeEpochValue}; +use common_meta::kv_backend::KvBackend; +use common_meta::rpc::store::CompareAndPutRequest; use common_telemetry::{debug, error, info, warn}; +use common_time::util::current_time_millis; use futures::StreamExt; use once_cell::sync::OnceCell; use snafu::{OptionExt, ResultExt}; @@ -39,6 +46,9 @@ use crate::service::{GrpcResult, GrpcStream}; type HeartbeatResponseResult = std::result::Result; +/// How far ahead of the leader's clock a datanode epoch can be recorded. +const MAX_NODE_EPOCH_AHEAD: Duration = Duration::from_secs(5 * 60); + #[async_trait::async_trait] trait HeartbeatRequestStream { async fn next(&mut self) -> Option>; @@ -109,6 +119,8 @@ struct HeartbeatSession { handler_group: Arc, ctx: Context, sender_id: PusherId, + /// The node epoch the stream's pusher was registered with, if any. + pusher_epoch: Option, } impl HeartbeatSession @@ -144,7 +156,12 @@ where return None; }; - let sender_id = register_pusher(&handler_group, header, tx.clone()).await; + let sender_id = PusherId::new(header.role(), get_node_id(header)); + // A datanode stream registers its pusher on its first request that passes the check. + if sender_id.role != Role::Datanode { + let pusher = Pusher::new(tx.clone()); + handler_group.register_pusher(sender_id, pusher).await; + } let mut session = Self { requests, tx, @@ -152,6 +169,7 @@ where handler_group, ctx, sender_id, + pusher_epoch: None, }; if session.handle_request(req, true).await { @@ -219,10 +237,23 @@ where .with_label_values(&[sender_id.as_str()]) .inc(); - let res = self - .handler_group - .handle(req, self.ctx.clone().with_handshake(is_handshake)) - .await + // A datanode's requests that carry an epoch are checked and handled one at a time, so + // an older process's request cannot run its handlers after a newer process is admitted. + let datanode_lock = if self.sender_id.role == Role::Datanode && req.node_epoch != 0 { + Some(self.handler_group.lock_datanode(self.sender_id.id).await) + } else { + None + }; + let res = match self.check_node_epoch(&req, is_handshake).await { + Some(res) => res, + None => { + self.handler_group + .handle(req, self.ctx.clone().with_handshake(is_handshake)) + .await + } + }; + drop(datanode_lock); + let res = res .inspect_err( |e| warn!(e; "Failed to handle heartbeat request, sender: {}", self.sender_id), ) @@ -257,6 +288,48 @@ where true } + /// Checks the node epoch of a datanode request, and registers the stream's pusher once + /// the request passes. Returns the response to send instead of running the handlers, if any. + async fn check_node_epoch( + &mut self, + req: &HeartbeatRequest, + is_handshake: bool, + ) -> Option> { + let node_epoch = req.node_epoch; + if self.sender_id.role != Role::Datanode || (node_epoch == 0 && is_handshake) { + return None; + } + // Datanodes that do not report their epoch are not checked. + if node_epoch != 0 { + if let Some(election) = &self.ctx.election + && !election.is_leader() + { + return Some(Ok(header_only(Some(Error::is_not_leader())))); + } + let kv_backend = &*self.ctx.leader_cached_kv_backend; + match accept_node_epoch(kv_backend, self.sender_id.id, node_epoch).await { + Ok(true) => {} + // No error, so the datanode does not reconnect. + Ok(false) => { + self.handler_group + .deregister_push(self.sender_id, &self.tx) + .await; + return Some(Ok(header_only(None))); + } + Err(e) => return Some(Err(e)), + } + } + // Before the handlers run, so the stream receives what they broadcast. + if self.pusher_epoch.is_none_or(|epoch| node_epoch > epoch) { + self.pusher_epoch = Some(node_epoch); + let pusher = Pusher::new(self.tx.clone()).with_node_epoch(node_epoch); + self.handler_group + .register_pusher(self.sender_id, pusher) + .await; + } + None + } + async fn send_not_leader_error(&mut self) { let _ = self .tx @@ -279,7 +352,9 @@ where async fn cleanup(&self) { info!("Heartbeat stream closed, sender: {}", self.sender_id); - let _ = self.handler_group.deregister_push(self.sender_id).await; + self.handler_group + .deregister_push(self.sender_id, &self.tx) + .await; } } @@ -400,17 +475,62 @@ fn get_node_id(header: &RequestHeader) -> u64 { } } -async fn register_pusher( - handler_group: &HeartbeatHandlerGroup, - header: &RequestHeader, - sender: Sender>, -) -> PusherId { - let role = header.role(); - let id = get_node_id(header); - let pusher_id = PusherId::new(role, id); - let pusher = Pusher::new(sender); - handler_group.register_pusher(pusher_id, pusher).await; - pusher_id +/// Returns whether `node_epoch` is not lower than the highest epoch recorded for the datanode, +/// raising the record when it is higher. An epoch too far ahead of this metasrv's clock is +/// refused, since recording it would fence every later process. +async fn accept_node_epoch( + kv_backend: &dyn KvBackend, + node_id: u64, + node_epoch: u64, +) -> Result { + let key = DatanodeEpochKey::new(node_id).to_bytes(); + let mut current = kv_backend.get(&key).await.context(error::KvBackendSnafu)?; + let limit = current_time_millis().max(0) as u64 + MAX_NODE_EPOCH_AHEAD.as_millis() as u64; + loop { + let (expect, highest) = match current { + Some(kv) => { + let value = DatanodeEpochValue::try_from_raw_value(&kv.value) + .context(error::KvBackendSnafu)?; + (kv.value, value.node_epoch) + } + None => (vec![], 0), + }; + match node_epoch.cmp(&highest) { + Ordering::Equal => return Ok(true), + Ordering::Greater if node_epoch <= limit => {} + _ => { + warn!("Fenced datanode {node_id} with epoch {node_epoch}, recorded {highest}"); + return Ok(false); + } + } + let value = DatanodeEpochValue { node_epoch } + .try_as_raw_value() + .context(error::KvBackendSnafu)?; + let req = CompareAndPutRequest { + key: key.clone(), + expect, + value, + }; + let resp = kv_backend + .compare_and_put(req) + .await + .context(error::KvBackendSnafu)?; + if resp.success { + info!("Recorded epoch {node_epoch} of datanode {node_id}"); + return Ok(true); + } + current = resp.prev_kv; + } +} + +fn header_only(error: Option) -> HeartbeatResponse { + HeartbeatResponse { + header: Some(ResponseHeader { + protocol_version: PROTOCOL_VERSION, + error, + }), + ..Default::default() + } } #[cfg(test)] @@ -596,9 +716,9 @@ mod tests { async fn test_heartbeat_session_init_registers_sender() { let (tx, mut rx) = mpsc::channel(8); let handler_group = test_handler_group(); - let sender_id = sender_id(Role::Datanode, 42); + let sender_id = sender_id(Role::Flownode, 42); let requests = - MockHeartbeatRequestStream::new(vec![Ok(heartbeat_request(Role::Datanode, 42))]); + MockHeartbeatRequestStream::new(vec![Ok(heartbeat_request(Role::Flownode, 42))]); let session = init_session( requests, @@ -619,9 +739,9 @@ mod tests { async fn test_heartbeat_session_run_deregisters_sender_on_stream_close() { let (tx, mut rx) = mpsc::channel(8); let handler_group = test_handler_group(); - let sender_id = sender_id(Role::Datanode, 42); + let sender_id = sender_id(Role::Flownode, 42); let requests = - MockHeartbeatRequestStream::new(vec![Ok(heartbeat_request(Role::Datanode, 42))]); + MockHeartbeatRequestStream::new(vec![Ok(heartbeat_request(Role::Flownode, 42))]); let session = init_session( requests, tx, @@ -641,9 +761,9 @@ mod tests { async fn test_heartbeat_session_run_forwards_stream_error_after_init() { let (tx, mut rx) = mpsc::channel(8); let handler_group = test_handler_group(); - let sender_id = sender_id(Role::Datanode, 42); + let sender_id = sender_id(Role::Flownode, 42); let requests = MockHeartbeatRequestStream::new(vec![ - Ok(heartbeat_request(Role::Datanode, 42)), + Ok(heartbeat_request(Role::Flownode, 42)), Err(Status::unavailable("temporary")), ]); let session = init_session( @@ -668,9 +788,9 @@ mod tests { async fn test_heartbeat_session_leader_step_down_sends_aborted_and_deregisters() { let (tx, mut rx) = mpsc::channel(8); let handler_group = test_handler_group(); - let sender_id = sender_id(Role::Datanode, 42); + let sender_id = sender_id(Role::Flownode, 42); let requests = MockHeartbeatRequestStream::pending_after(vec![Ok(heartbeat_request( - Role::Datanode, + Role::Flownode, 42, ))]); let session = init_session( @@ -694,9 +814,9 @@ mod tests { async fn test_heartbeat_session_leader_watcher_closed_sends_unavailable_and_deregisters() { let (tx, mut rx) = mpsc::channel(8); let handler_group = test_handler_group(); - let sender_id = sender_id(Role::Datanode, 42); + let sender_id = sender_id(Role::Flownode, 42); let requests = MockHeartbeatRequestStream::pending_after(vec![Ok(heartbeat_request( - Role::Datanode, + Role::Flownode, 42, ))]); let session = init_session( @@ -768,4 +888,459 @@ mod tests { assert_eq!(i, get_node_id(&header)); } } + + mod node_epoch { + use std::sync::Arc; + use std::time::Duration; + + use api::v1::meta::{HeartbeatRequest, HeartbeatResponse, MailboxMessage, Peer, Role}; + use common_meta::error::UnexpectedSnafu; + use common_meta::key::datanode_epoch::{DatanodeEpochKey, DatanodeEpochValue}; + use common_meta::key::{MetadataValue, TableMetadataManager}; + use common_meta::kv_backend::memory::MemoryKvBackend; + use common_meta::kv_backend::test_util::MockKvBackendBuilder; + use common_meta::kv_backend::{KvBackend, KvBackendRef, ResettableKvBackend}; + use common_meta::region_keeper::MemoryRegionKeeper; + use common_meta::rpc::store::{PutRequest, RangeResponse}; + use common_meta::sequence::SequenceBuilder; + use common_time::util::current_time_millis; + use tokio::sync::{Notify, mpsc}; + + use super::{MockHeartbeatRequestStream, MockLeaderStepDown, heartbeat_request}; + use crate::discovery::utils::find_datanode_lease_value; + use crate::handler::region_lease_handler::RegionLeaseHandler; + use crate::handler::test_utils::TestEnv; + use crate::handler::{ + HandleControl, HeartbeatAccumulator, HeartbeatHandler, HeartbeatHandlerGroup, + HeartbeatHandlerGroupBuilder, HeartbeatMailbox, Pusher, PusherId, Pushers, + }; + use crate::metasrv::Context; + use crate::service::heartbeat::{ + HeartbeatResponseResult, HeartbeatSession, accept_node_epoch, + }; + use crate::service::mailbox::{BroadcastChannel, Channel, MailboxRef}; + use crate::service::store::cached_kv::LeaderCachedKvBackend; + + const NODE_ID: u64 = 1; + + /// A datanode stream that has sent its handshake, which carries no epoch. + struct Stream { + session: HeartbeatSession, + rx: mpsc::Receiver, + } + + impl Stream { + async fn open(group: &Arc, ctx: &Context) -> Self { + let (tx, mut rx) = mpsc::channel(8); + let handshake = heartbeat_request(Role::Datanode, NODE_ID); + let session = HeartbeatSession::init( + MockHeartbeatRequestStream::new(vec![Ok(handshake)]), + tx, + None, + group.clone(), + ctx.clone(), + ) + .await + .unwrap(); + rx.recv().await.unwrap().unwrap(); + Self { session, rx } + } + + async fn send(&mut self, req: HeartbeatRequest) -> HeartbeatResponse { + assert!(self.session.handle_request(req, false).await); + self.rx.recv().await.unwrap().unwrap() + } + + /// Returns whether the stream has received a pushed message. + fn received_push(&mut self) -> bool { + self.rx.try_recv().is_ok() + } + } + + fn group_builder() -> HeartbeatHandlerGroupBuilder { + let lease_handler = RegionLeaseHandler::new( + 10, + Arc::new(TableMetadataManager::new(Arc::new(MemoryKvBackend::new()))), + Arc::new(MemoryRegionKeeper::default()), + None, + ); + let mut builder = HeartbeatHandlerGroupBuilder::new(Pushers::default()) + .with_region_lease_handler(Some(lease_handler)) + .add_default_handlers(); + builder + .add_handler_after("ResponseHeaderHandler", BroadcastHandler) + .unwrap(); + builder + } + + fn new_group() -> Arc { + Arc::new(group_builder().build().unwrap()) + } + + /// Holds a request from the address "pause": notifies `entered`, then waits for `release`. + struct PauseHandler { + entered: Arc, + release: Arc, + } + + #[async_trait::async_trait] + impl HeartbeatHandler for PauseHandler { + fn is_acceptable(&self, role: Role) -> bool { + role == Role::Datanode + } + + async fn handle( + &self, + req: &HeartbeatRequest, + _: &mut Context, + _: &mut HeartbeatAccumulator, + ) -> crate::error::Result { + if req.peer.as_ref().is_some_and(|peer| peer.addr == "pause") { + self.entered.notify_one(); + self.release.notified().await; + } + Ok(HandleControl::Continue) + } + } + + /// Broadcasts to datanodes while handling a request from the address "broadcast". + struct BroadcastHandler; + + #[async_trait::async_trait] + impl HeartbeatHandler for BroadcastHandler { + fn is_acceptable(&self, role: Role) -> bool { + role == Role::Datanode + } + + async fn handle( + &self, + req: &HeartbeatRequest, + ctx: &mut Context, + _: &mut HeartbeatAccumulator, + ) -> crate::error::Result { + if req + .peer + .as_ref() + .is_some_and(|peer| peer.addr == "broadcast") + { + let msg = MailboxMessage::default(); + ctx.mailbox + .broadcast(&BroadcastChannel::Datanode, &msg) + .await?; + } + Ok(HandleControl::Continue) + } + } + + fn heartbeat(node_epoch: u64, addr: &str) -> HeartbeatRequest { + HeartbeatRequest { + peer: Some(Peer::new(NODE_ID, addr)), + node_epoch, + ..heartbeat_request(Role::Datanode, NODE_ID) + } + } + + fn mailbox(pushers: Pushers) -> MailboxRef { + let seq = SequenceBuilder::new("test_seq", Arc::new(MemoryKvBackend::new())).build(); + HeartbeatMailbox::create(pushers, seq) + } + + async fn push(group: &HeartbeatHandlerGroup) { + mailbox(group.pushers()) + .send_oneway(&Channel::Datanode(NODE_ID), MailboxMessage::default()) + .await + .unwrap(); + } + + async fn lease_addr(ctx: &Context) -> Option { + let in_memory = ctx.in_memory.clone() as KvBackendRef; + find_datanode_lease_value(&in_memory, NODE_ID) + .await + .unwrap() + .map(|lease| lease.node_addr) + } + + /// Records `node_epoch` as the highest epoch of the node, through the leader cache. + async fn record(ctx: &Context, node_epoch: u64) { + let key = DatanodeEpochKey::new(NODE_ID).to_bytes(); + let value = DatanodeEpochValue { node_epoch } + .try_as_raw_value() + .unwrap(); + let req = PutRequest::new().with_key(key).with_value(value); + ctx.leader_cached_kv_backend.put(req).await.unwrap(); + } + + fn epochs() -> (u64, u64) { + let now = current_time_millis() as u64; + (now - 1000, now) + } + + #[tokio::test] + async fn test_newer_process_fences_the_older_one() { + let group = new_group(); + let mut ctx = TestEnv::new().ctx(); + // A mailbox waiting for replies from the datanode. + let pushers = Pushers::default(); + let (tx, _rx) = mpsc::channel(8); + pushers + .insert(PusherId::new(Role::Datanode, NODE_ID), Pusher::new(tx)) + .await; + ctx.mailbox = mailbox(pushers); + let wait_for_reply = || { + ctx.mailbox.send( + &Channel::Datanode(NODE_ID), + MailboxMessage::default(), + Duration::from_millis(500), + ) + }; + let (e1, e2) = epochs(); + let mut a1 = Stream::open(&group, &ctx).await; + let mut a2 = Stream::open(&group, &ctx).await; + + let res = a1.send(heartbeat(e1, "a1")).await; + assert!(res.region_lease.is_some()); + // A2 replaces A1 and records its epoch. + let res = a2.send(heartbeat(e2, "a2")).await; + assert!(res.region_lease.is_some()); + + // A1 is fenced: no lease, no keep-alive, no error that would make it reconnect. + let res = a1.send(heartbeat(e1, "a1")).await; + assert!(res.region_lease.is_none()); + assert!(res.header.unwrap().error.is_none()); + assert_eq!(lease_addr(&ctx).await.as_deref(), Some("a2")); + + // A1's mailbox reply is dropped, A2's is delivered. + let a1_reply = wait_for_reply().await.unwrap(); + let a2_reply = wait_for_reply().await.unwrap(); + for (stream, reply, epoch) in [(&mut a1, &a1_reply, e1), (&mut a2, &a2_reply, e2)] { + let mut req = heartbeat(epoch, "a"); + req.mailbox_message = Some(MailboxMessage { + id: reply.message_id(), + ..Default::default() + }); + stream.send(req).await; + } + assert!(a1_reply.await.is_err()); + assert!(a2_reply.await.is_ok()); + } + + #[tokio::test] + async fn test_pusher_follows_the_newer_process() { + let group = new_group(); + let mut ctx = TestEnv::new().ctx(); + ctx.mailbox = mailbox(group.pushers()); + let (e1, e2) = epochs(); + let mut a1 = Stream::open(&group, &ctx).await; + a1.send(heartbeat(e1, "a1")).await; + // A2 takes over before the handlers of its first heartbeat run, so it receives + // what they broadcast ahead of the response. + let mut a2 = Stream::open(&group, &ctx).await; + assert!( + a2.session + .handle_request(heartbeat(e2, "broadcast"), false) + .await + ); + let broadcast = a2.rx.recv().await.unwrap().unwrap(); + assert!(broadcast.mailbox_message.is_some()); + let res = a2.rx.recv().await.unwrap().unwrap(); + assert!(res.region_lease.is_some()); + assert!(!a1.received_push()); + push(&group).await; + assert!(a2.received_push()); + + // A1's new stream does not take the pusher back. + let mut a1_again = Stream::open(&group, &ctx).await; + a1_again.send(heartbeat(e1, "a1")).await; + push(&group).await; + assert!(a2.received_push()); + + // A stream ending removes only its own pusher. + a1.session.run().await; + a1_again.session.run().await; + push(&group).await; + assert!(a2.received_push()); + // A2's new stream takes over with the same epoch, and A2's old stream ending + // keeps it. + let mut a2_again = Stream::open(&group, &ctx).await; + a2_again.send(heartbeat(e2, "a2")).await; + a2.session.run().await; + push(&group).await; + assert!(a2_again.received_push()); + + // A fenced stream removes its own pusher, and a handshake does not take the slot. + record(&ctx, e2 + 1).await; + a2_again.send(heartbeat(e2, "a2")).await; + let pusher_id = PusherId::new(Role::Datanode, NODE_ID); + assert!(!group.contains_pusher(&pusher_id).await); + let _a1_reconnected = Stream::open(&group, &ctx).await; + assert!(!group.contains_pusher(&pusher_id).await); + } + + #[tokio::test] + async fn test_newer_process_waits_for_an_older_request() { + let (entered, release) = (Arc::new(Notify::new()), Arc::new(Notify::new())); + let mut builder = group_builder(); + let pause = PauseHandler { + entered: entered.clone(), + release: release.clone(), + }; + builder + .add_handler_before("DatanodeKeepLeaseHandler", pause) + .unwrap(); + let group = Arc::new(builder.build().unwrap()); + let ctx = TestEnv::new().ctx(); + let (e1, e2) = epochs(); + let mut a1 = Stream::open(&group, &ctx).await; + let mut a2 = Stream::open(&group, &ctx).await; + + // A1's request passes the check and is held before the keep-alive handler, and + // A2's request waits for it. + let a1_req = a1.session.handle_request(heartbeat(e1, "pause"), false); + let a2_req = async { + entered.notified().await; + let mut a2_req = std::pin::pin!(a2.send(heartbeat(e2, "a2"))); + let waiting = tokio::time::timeout(Duration::from_millis(100), &mut a2_req).await; + assert!(waiting.is_err()); + release.notify_one(); + a2_req.await + }; + let (open, res) = tokio::join!(a1_req, a2_req); + assert!(open); + assert!(res.region_lease.is_some()); + assert_eq!(lease_addr(&ctx).await.as_deref(), Some("a2")); + + // A1 can no longer overwrite the address or get a lease. + a1.rx.recv().await.unwrap().unwrap(); + let res = a1.send(heartbeat(e1, "a1")).await; + assert!(res.region_lease.is_none()); + assert_eq!(lease_addr(&ctx).await.as_deref(), Some("a2")); + } + + #[tokio::test] + async fn test_kv_error_fails_the_request() { + let group = new_group(); + let mut ctx = TestEnv::new().ctx(); + let kv_backend = MockKvBackendBuilder::default() + .range_fn(Arc::new(|_| { + Ok(RangeResponse { + kvs: vec![], + more: false, + }) + })) + .txn(Arc::new(|_| { + UnexpectedSnafu { + err_msg: "injected", + } + .fail() + })) + .build() + .unwrap(); + let cached = LeaderCachedKvBackend::with_always_leader(Arc::new(kv_backend)); + ctx.leader_cached_kv_backend = Arc::new(cached); + let (e1, _) = epochs(); + let mut a1 = Stream::open(&group, &ctx).await; + + let req = HeartbeatRequest { + node_epoch: e1, + ..heartbeat_request(Role::Datanode, NODE_ID) + }; + assert!(a1.session.handle_request(req, false).await); + assert!(a1.rx.recv().await.unwrap().is_err()); + } + + #[tokio::test] + async fn test_first_epoch_keeps_waiting_messages() { + let group = new_group(); + let mut ctx = TestEnv::new().ctx(); + ctx.mailbox = mailbox(group.pushers()); + let (e1, _) = epochs(); + let mut a1 = Stream::open(&group, &ctx).await; + // A request without an epoch registers the pusher first. + a1.send(heartbeat(0, "a1")).await; + let reply = ctx + .mailbox + .send( + &Channel::Datanode(NODE_ID), + MailboxMessage::default(), + Duration::from_secs(10), + ) + .await + .unwrap(); + let id = a1 + .rx + .recv() + .await + .unwrap() + .unwrap() + .mailbox_message + .unwrap() + .id; + tokio::pin!(reply); + + // The stream's pusher takes the epoch of its first heartbeat without cancelling + // the message, which a later heartbeat answers. + a1.send(heartbeat(e1, "a1")).await; + assert!(futures::poll!(reply.as_mut()).is_pending()); + let mut req = heartbeat(e1, "a1"); + req.mailbox_message = Some(MailboxMessage { + id, + ..Default::default() + }); + a1.send(req).await; + assert!(reply.await.is_ok()); + } + + #[tokio::test] + async fn test_record_is_reloaded_on_leader_change() { + let store: KvBackendRef = Arc::new(MemoryKvBackend::new()); + let cached = LeaderCachedKvBackend::with_always_leader(store.clone()); + let now = current_time_millis() as u64; + let (e1, e2) = (now - 1000, now); + assert!(accept_node_epoch(&cached, 1, e1).await.unwrap()); + assert!(accept_node_epoch(&cached, 1, e2).await.unwrap()); + + // The cache is reset and loaded again when this metasrv becomes the leader. The record + // is then served from the loaded cache, even without the store. + cached.reset(); + cached.load().await.unwrap(); + store + .delete(&DatanodeEpochKey::new(1).to_bytes(), false) + .await + .unwrap(); + assert!(!accept_node_epoch(&cached, 1, e1).await.unwrap()); + assert!(accept_node_epoch(&cached, 1, e2).await.unwrap()); + } + + #[tokio::test] + async fn test_epoch_far_ahead_is_not_recorded() { + let kv_backend = MemoryKvBackend::new(); + let now = current_time_millis() as u64; + let ahead = now + Duration::from_secs(10 * 60).as_millis() as u64; + assert!(!accept_node_epoch(&kv_backend, 1, ahead).await.unwrap()); + let key = DatanodeEpochKey::new(1).to_bytes(); + assert!(kv_backend.get(&key).await.unwrap().is_none()); + + let slightly_ahead = now + Duration::from_secs(60).as_millis() as u64; + assert!( + accept_node_epoch(&kv_backend, 1, slightly_ahead) + .await + .unwrap() + ); + } + + #[tokio::test] + async fn test_datanode_without_epoch_is_not_checked() { + let group = new_group(); + let ctx = TestEnv::new().ctx(); + // Even after another process of the node has recorded its epoch. + let (_, e2) = epochs(); + record(&ctx, e2).await; + let mut a1 = Stream::open(&group, &ctx).await; + + let res = a1.send(heartbeat(0, "a1")).await; + assert!(res.region_lease.is_some()); + push(&group).await; + assert!(a1.received_push()); + } + } } diff --git a/src/meta-srv/src/service/store/cached_kv.rs b/src/meta-srv/src/service/store/cached_kv.rs index 1c2990f5446..deccf67e021 100644 --- a/src/meta-srv/src/service/store/cached_kv.rs +++ b/src/meta-srv/src/service/store/cached_kv.rs @@ -69,9 +69,10 @@ pub struct LeaderCachedKvBackend { check_leader: CheckLeaderRef, store: KvBackendRef, cache: ResettableKvBackendRef, - /// Incremented before and after each store mutation. The first increment - /// detects overlapping writers; the second invalidates cache fills that - /// read the old value while the store mutation was in flight. + /// Incremented before and after each store mutation, also while the cache is + /// disabled. The first increment detects overlapping writers; the second + /// invalidates cache fills and loads that read the old value while the + /// store mutation was in flight. version: AtomicUsize, name: String, } @@ -94,11 +95,13 @@ impl LeaderCachedKvBackend { Self::new(Arc::new(AlwaysLeader), store) } - /// The caller MUST ensure during the loading, there are no mutation requests reaching the `LeaderCachedKvStore`. + /// Loads the cached prefixes. Writes may run meanwhile: a prefix whose load overlapped a + /// write is dropped from the cache again, and read through later. pub async fn load(&self) -> Result<()> { for prefix in &CACHE_KEY_PREFIXES[..] { let _timer = metrics::METRIC_META_LEADER_CACHED_KV_LOAD_ELAPSED.with_label_values(&[prefix]); + let ver = self.get_version(); // TODO(weny): Refactors PaginationStream's output to unary output. let stream = PaginationStream::new( @@ -110,6 +113,7 @@ impl LeaderCachedKvBackend { .into_stream(); let kvs = stream.try_collect::>().await?; + let keys = kvs.iter().map(|kv| kv.key.clone()).collect::>(); self.cache .batch_put(BatchPutRequest { @@ -117,6 +121,10 @@ impl LeaderCachedKvBackend { prev_kv: false, }) .await?; + // A write while loading may have changed what was read. + if !self.validate_version(ver) { + self.invalid_keys(keys).await?; + } } Ok(()) @@ -210,10 +218,6 @@ impl KvBackend for LeaderCachedKvBackend { } async fn put(&self, req: PutRequest) -> Result { - if !self.is_leader() { - return self.store.put(req).await; - } - let ver = self.create_new_version(); let res = self.store.put(req.clone()).await?; @@ -229,10 +233,6 @@ impl KvBackend for LeaderCachedKvBackend { } async fn batch_put(&self, req: BatchPutRequest) -> Result { - if !self.is_leader() { - return self.store.batch_put(req).await; - } - let ver = self.create_new_version(); let res = self.store.batch_put(req.clone()).await?; @@ -305,10 +305,6 @@ impl KvBackend for LeaderCachedKvBackend { } async fn compare_and_put(&self, req: CompareAndPutRequest) -> Result { - if !self.is_leader() { - return self.store.compare_and_put(req).await; - } - let _ = self.create_new_version(); let key = req.key.clone(); @@ -323,10 +319,6 @@ impl KvBackend for LeaderCachedKvBackend { } async fn delete_range(&self, req: DeleteRangeRequest) -> Result { - if !self.is_leader() { - return self.store.delete_range(req).await; - } - let _ = self.create_new_version(); let res = self.store.delete_range(req.clone()).await?; @@ -336,10 +328,6 @@ impl KvBackend for LeaderCachedKvBackend { } async fn batch_delete(&self, req: BatchDeleteRequest) -> Result { - if !self.is_leader() { - return self.store.batch_delete(req).await; - } - let _ = self.create_new_version(); let res = self.store.batch_delete(req.clone()).await?; @@ -354,10 +342,6 @@ impl TxnService for LeaderCachedKvBackend { type Error = Error; async fn txn(&self, txn: Txn) -> Result { - if !self.is_leader() { - return self.store.txn(txn).await; - } - let _ = self.create_new_version(); let res = self.store.txn(txn.clone()).await?; @@ -531,7 +515,8 @@ mod tests { Write::DeleteRange, Write::BatchDelete, ]; - const KEY: &[u8] = b"route"; + /// Under the first prefix the leader cache loads. + const KEY: &[u8] = b"__table_name/1"; impl Write { async fn execute(self, store: &LeaderCachedKvBackend, value: &[u8]) { @@ -733,6 +718,60 @@ mod tests { } } + struct EnableCache(std::sync::atomic::AtomicBool); + + impl CheckLeader for EnableCache { + fn check(&self) -> bool { + self.0.load(Ordering::Relaxed) + } + } + + /// A write made while the cache is disabled overlaps the load: the load reads the old value + /// either before the write commits and installs it afterwards, or while the write is in + /// flight and completes, with the cache enabled, before the write commits. + #[tokio::test] + async fn test_load_overlaps_write() { + for write in WRITES { + for write_first in [false, true] { + let backing = Arc::new(ControlledBackend::default()); + backing + .inner + .put(PutRequest::new().with_key(KEY).with_value(b"old")) + .await + .unwrap(); + let enabled = Arc::new(EnableCache(Default::default())); + let store = LeaderCachedKvBackend::new(enabled.clone(), backing.clone()); + + if write_first { + let commit = backing.before_write.arm(); + let writer = write.execute(&store, b"new"); + tokio::pin!(writer); + assert!(poll!(writer.as_mut()).is_pending()); + store.load().await.unwrap(); + enabled.0.store(true, Ordering::Relaxed); + commit.send(()).unwrap(); + writer.await; + } else { + let return_read = backing.after_read.arm(); + let loader = store.load(); + tokio::pin!(loader); + assert!(poll!(loader.as_mut()).is_pending()); + write.execute(&store, b"new").await; + return_read.send(()).unwrap(); + loader.await.unwrap(); + enabled.0.store(true, Ordering::Relaxed); + } + + let durable = backing.inner.get(KEY).await.unwrap(); + assert_eq!( + read(&store, false).await, + durable, + "{write:?}, write_first={write_first}" + ); + } + } + } + #[tokio::test] async fn test_cache_miss_overlaps_reset() { for batch in [true, false] { diff --git a/tests/compatibility/meta/datanode_epoch/value.json b/tests/compatibility/meta/datanode_epoch/value.json new file mode 100644 index 00000000000..e6bed6e9b87 --- /dev/null +++ b/tests/compatibility/meta/datanode_epoch/value.json @@ -0,0 +1 @@ +{"node_epoch":1761091200000} \ No newline at end of file From 1e97a0e7c3e9c2e6023c12bbfe4b04e283f0aaa5 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Mon, 26 Oct 2026 09:08:13 +0000 Subject: [PATCH 096/239] fix(log-store): refuse an object store that ignores conditional create (#9909) * fix(log-store): refuse an object store that ignores conditional create The object store WAL checks only the write_with_if_not_exists capability before it opens, and OpenDAL declares that capability per service. A store behind the S3 service can accept If-None-Match and overwrite the object anyway, so the WAL could overwrite its own objects and lose acknowledged writes without an error. Before recovery, the store now creates a probe key, unique to the open, twice under /probes/, outside the objects it lists, and deletes it. The second create must be refused with a condition failure, otherwise the open fails with an error naming the object store. The probe costs two writes and one delete per open. Signed-off-by: jeremyhi * test(log-store): drop a comment that restates the keys helper Signed-off-by: jeremyhi * fix(log-store): accept AlreadyExists from the second probe create A backend may report an existing conditional-create target as AlreadyExists. The first create stored the probe key, so that error can only refer to it. Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- src/log-store/Cargo.toml | 2 +- src/log-store/src/object_store_wal/io.rs | 59 +++++++++- src/log-store/src/object_store_wal/store.rs | 116 +++++++++++++++++++- 3 files changed, 171 insertions(+), 6 deletions(-) diff --git a/src/log-store/Cargo.toml b/src/log-store/Cargo.toml index 54e0b48c41d..ee147fbcd4b 100644 --- a/src/log-store/Cargo.toml +++ b/src/log-store/Cargo.toml @@ -46,6 +46,7 @@ snafu.workspace = true store-api.workspace = true tokio.workspace = true tokio-util = { workspace = true, features = ["rt"] } +uuid.workspace = true [dev-dependencies] common-meta = { workspace = true, features = ["testing"] } @@ -54,7 +55,6 @@ common-wal = { workspace = true, features = ["testing"] } object-store = { workspace = true, features = ["testing"] } rand.workspace = true rskafka = { workspace = true, features = ["unstable-fuzzing"] } -uuid.workspace = true [package.metadata.cargo-shear] # `protobuf` is required by the code generated by `build.rs`: the generated diff --git a/src/log-store/src/object_store_wal/io.rs b/src/log-store/src/object_store_wal/io.rs index 43f018a7ecb..0c087141c52 100644 --- a/src/log-store/src/object_store_wal/io.rs +++ b/src/log-store/src/object_store_wal/io.rs @@ -17,12 +17,14 @@ use async_stream::try_stream; use bytes::Bytes; +use common_telemetry::warn; use futures::TryStreamExt; use futures::stream::BoxStream; use object_store::{Entry, ErrorKind, ObjectStore}; use snafu::{OptionExt, ResultExt, ensure}; use tokio::sync::mpsc; use tokio_util::task::AbortOnDropHandle; +use uuid::Uuid; use crate::error::{ CorruptedWalObjectSnafu, InvalidWalObjectStoreSnafu, Result, WalObjectConflictSnafu, @@ -33,6 +35,7 @@ use crate::error::{ /// [`u64::MAX`] so that lexicographic and numeric order agree. const OBJECT_SEQ_WIDTH: usize = 20; const OBJECT_SUFFIX: &str = ".wal"; +const PROBE_CONTENT: &[u8] = b"conditional create probe"; /// Objects a listing may hold ahead of the reads that take them, four LIST /// pages of S3, so the listing stays pages ahead of the reads while what it /// holds is bounded. @@ -61,6 +64,9 @@ pub(crate) struct ListedObject { pub(crate) struct ObjectStoreIo { store: ObjectStore, object_prefix: String, + /// Holds the keys of [`Self::probe_conditional_create`], outside + /// `object_prefix` so that no listing of the objects returns one. + probe_prefix: String, } impl ObjectStoreIo { @@ -68,7 +74,7 @@ impl ObjectStoreIo { /// conditional creates, which is what keeps a retried write from /// overwriting a durable object. pub(crate) fn new(store: ObjectStore, prefix: impl AsRef) -> Result { - let object_prefix = normalize_prefix(prefix.as_ref())?; + let prefix = normalize_prefix(prefix.as_ref())?; ensure!( store.info().capability().write_with_if_not_exists, InvalidWalObjectStoreSnafu { @@ -77,10 +83,57 @@ impl ObjectStoreIo { ); Ok(Self { store, - object_prefix, + object_prefix: format!("{prefix}/objects/"), + probe_prefix: format!("{prefix}/probes/"), }) } + /// Checks that the store refuses a conditional create of an object that + /// exists. The capability checked by [`Self::new`] is declared per + /// service, and a store behind the S3 service may accept the condition and + /// overwrite the object anyway, which would let a retried write replace a + /// durable object. The probe creates a key of its own, unique to this + /// call, twice, and then deletes it. + pub(crate) async fn probe_conditional_create(&self) -> Result<()> { + let path = format!("{}{}", self.probe_prefix, Uuid::new_v4()); + let create = || { + self.store + .write_with(&path, PROBE_CONTENT) + .if_not_exists(true) + }; + create().await.context(WalObjectStoreSnafu { + operation: "write", + path: path.clone(), + })?; + let recreated = create().await; + if let Err(error) = self.store.delete(&path).await { + warn!(error; "Failed to delete the conditional create probe {path} of the object store WAL"); + } + match recreated { + // The first create stored this key, so `AlreadyExists` can only + // refer to it. + Err(error) + if matches!( + error.kind(), + ErrorKind::ConditionNotMatch | ErrorKind::AlreadyExists + ) => + { + Ok(()) + } + Err(error) => Err(error).context(WalObjectStoreSnafu { + operation: "write", + path, + }), + Ok(_) => InvalidWalObjectStoreSnafu { + reason: format!( + "the object store ({}) overwrote the existing object {path} on a conditional create, but the object store WAL requires conditional create support; a versioned bucket or the S3-compatible endpoint of GCS can cause this", + self.storage() + ), + } + .fail(), + } + } + /// Creates the object `object_seq` unless it already exists. A retry with /// the same content is accepted, different content is rejected. /// @@ -309,7 +362,7 @@ fn normalize_prefix(prefix: &str) -> Result { reason: format!("object prefix {prefix:?} has an empty or relative component"), } ); - Ok(format!("{prefix}/objects/")) + Ok(prefix.to_string()) } #[cfg(test)] diff --git a/src/log-store/src/object_store_wal/store.rs b/src/log-store/src/object_store_wal/store.rs index 40d6e75326c..a741c337d88 100644 --- a/src/log-store/src/object_store_wal/store.rs +++ b/src/log-store/src/object_store_wal/store.rs @@ -295,8 +295,10 @@ impl fmt::Debug for ObjectStoreLogStore { impl ObjectStoreLogStore { /// Builds the store under the node and generation prefix derived from `config`, /// recovering the catalog from the objects that already exist and writing - /// the object that starts the epoch of this instance. Recovery fails on the - /// first corrupted or conflicting object. + /// the object that starts the epoch of this instance. Before it recovers, + /// it fails unless the object store refuses to overwrite an object on a + /// conditional create. Recovery fails on the first corrupted or + /// conflicting object. pub async fn try_new( object_store: ObjectStore, config: &ObjectStoreWalConfig, @@ -305,6 +307,7 @@ impl ObjectStoreLogStore { ) -> Result> { let prefix = config.node_prefix(node_id, generation); let io = ObjectStoreIo::new(object_store, &prefix)?; + io.probe_conditional_create().await?; Self::open(Arc::new(io), config, prefix).await } @@ -4011,6 +4014,115 @@ mod tests { } } + /// A writer that stores the object whether it exists or not, as an object + /// store that ignores the condition of a create does. + struct OverwritingWriter { + store: ObjectStore, + path: String, + content: Vec, + } + + impl oio::Write for OverwritingWriter { + async fn write(&mut self, buffer: mock::Buffer) -> mock::Result<()> { + self.content.extend_from_slice(&buffer.to_bytes()); + Ok(()) + } + + async fn close(&mut self) -> mock::Result { + let content = std::mem::take(&mut self.content); + self.store.write(&self.path, content).await + } + + async fn abort(&mut self) -> mock::Result<()> { + Ok(()) + } + } + + async fn keys(object_store: &ObjectStore) -> Vec { + let mut keys = object_store + .list_with(&format!("{PREFIX}/")) + .recursive(true) + .await + .unwrap() + .into_iter() + .filter(|entry| entry.metadata().is_file()) + .map(|entry| entry.path().to_string()) + .collect::>(); + keys.sort_unstable(); + keys + } + + #[tokio::test] + async fn test_store_open_deletes_its_conditional_create_probe() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + let store = open(object_store.clone(), &eager()).await; + store.stop().await.unwrap(); + + assert_eq!( + vec![object_path(&object_store, 0)], + keys(&object_store).await + ); + } + + #[tokio::test] + async fn test_store_open_refuses_an_object_store_that_ignores_conditional_create() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + let target = object_store.clone(); + let ignoring = object_store.clone().layer( + MockLayerBuilder::default() + .writer_factory(Arc::new(move |path, _, _| { + Box::new(OverwritingWriter { + store: target.clone(), + path: path.to_string(), + content: Vec::new(), + }) + })) + .build() + .unwrap(), + ); + + let error = ObjectStoreLogStore::try_new(ignoring, &eager(), 1, 2) + .await + .unwrap_err(); + assert!( + matches!(&error, Error::InvalidWalObjectStore { reason, .. } + if reason.contains("overwrote the existing object")), + "unexpected error: {error:?}" + ); + // The open stops before it writes a start object, and deletes the + // probe. + assert!(keys(&object_store).await.is_empty()); + } + + #[tokio::test] + async fn test_store_recovery_ignores_a_probe_left_behind() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + // OpenDAL's unit deleter refuses every delete, so the probe stays. + let undeletable = object_store.clone().layer( + MockLayerBuilder::default() + .deleter_factory(Arc::new(|_| Box::new(()))) + .build() + .unwrap(), + ); + open(undeletable, &eager()).await.stop().await.unwrap(); + let probe_prefix = format!("{PREFIX}/probes/"); + let probes = keys(&object_store) + .await + .into_iter() + .filter(|key| key.starts_with(&probe_prefix)) + .collect::>(); + assert_eq!(1, probes.len(), "{probes:?}"); + + // The next open recovers the start object of the first one only. + let store = open(object_store.clone(), &eager()).await; + let io = ObjectStoreIo::new(object_store, PREFIX).unwrap(); + assert_eq!(vec![0, 1], object_seqs(&io).await); + store.stop().await.unwrap(); + } + /// Rebuilds the catalog by decoding whole objects as a recovery oracle. async fn recover_by_decoding(io: &dyn WalObjectIo) -> Result { let mut objects = Vec::new(); From 8170d0458a17fc4df553b99d3b027f5f7bf0b609 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Mon, 26 Oct 2026 09:36:49 +0000 Subject: [PATCH 097/239] fix(ci): scope Aliyun janitor sweeps to the owning repository (#9966) Signed-off-by: WenyXu --- .../query-regression/README.md | 9 +++- .../scripts/aliyun-ecs-runner-provision.py | 2 + .github/scripts/aliyun-ecs-runner-teardown.py | 13 +++-- tests/perf/test_aliyun_ecs_runner_scripts.py | 49 +++++++++++++++++-- 4 files changed, 61 insertions(+), 12 deletions(-) diff --git a/.github/runner-scale-sets/query-regression/README.md b/.github/runner-scale-sets/query-regression/README.md index 44500a8e50f..748821966d6 100644 --- a/.github/runner-scale-sets/query-regression/README.md +++ b/.github/runner-scale-sets/query-regression/README.md @@ -80,8 +80,13 @@ after its TTL, so finish the inspection within that window. Trust model on the ECS path: the instance receives only the one-hour runner registration token via user data and holds no cloud credentials; the Aliyun AK/SK exist only in the control-plane jobs. Instance tags -(`managed-by=query-regression-ci`, `query-regression-run-id`) feed the janitor -sweep. +(`managed-by=query-regression-ci`, `query-regression-run-id`, +`github-repository`) feed the janitor sweep. The sweep only selects instances +whose `github-repository` matches the requested repository. Older instances +without that tag are excluded: verify their ownership before adding the tag, +or clean them up by explicit instance ID. Per-run teardown remains unchanged. +Update every repository sharing the account and region to the scoped helper; +an older janitor can still sweep instances belonging to other repositories. ### Building and updating the ECS image diff --git a/.github/scripts/aliyun-ecs-runner-provision.py b/.github/scripts/aliyun-ecs-runner-provision.py index 46efbffd3ec..2d026c00bc0 100644 --- a/.github/scripts/aliyun-ecs-runner-provision.py +++ b/.github/scripts/aliyun-ecs-runner-provision.py @@ -59,6 +59,7 @@ MANAGED_BY_TAG_KEY = "managed-by" MANAGED_BY_TAG_VALUE = "query-regression-ci" RUN_TAG_KEY = "query-regression-run-id" TTL_TAG_KEY = "runner-ttl-hours" +REPO_TAG_KEY = "github-repository" # Runner cache paths on the instance system disk. They are created empty # every provision and deleted with the VM. @@ -359,6 +360,7 @@ def run_instance(client, config: ProvisionConfig, user_data: str) -> str: key=MANAGED_BY_TAG_KEY, value=MANAGED_BY_TAG_VALUE ), ecs_models.RunInstancesRequestTag(key=RUN_TAG_KEY, value=config.run_id), + ecs_models.RunInstancesRequestTag(key=REPO_TAG_KEY, value=config.repo), ] + ( [ diff --git a/.github/scripts/aliyun-ecs-runner-teardown.py b/.github/scripts/aliyun-ecs-runner-teardown.py index d5b30b2922d..85742c1f694 100644 --- a/.github/scripts/aliyun-ecs-runner-teardown.py +++ b/.github/scripts/aliyun-ecs-runner-teardown.py @@ -32,8 +32,8 @@ Two modes: runner by name. Both steps are idempotent and best-effort; a missing instance or runner is not an error. - Sweep (scheduled janitor): delete every instance tagged as managed by - query-regression CI whose creation time is older than its tagged TTL (or - the given fallback TTL for untagged instances), and + query-regression CI in the requested repository whose creation time is older + than its tagged TTL (or the fallback TTL for instances without a TTL tag), and deregister the matching runners. This is the safety net for runs whose teardown job never executed. """ @@ -162,7 +162,7 @@ def delete_instance(client, instance_id: str, region_id: str | None = None) -> b def list_managed_instances( - client, region_id: str + client, region_id: str, repo: str ) -> list[tuple[str, str, str, str | None]]: from alibabacloud_ecs20140526 import models as ecs_models @@ -175,6 +175,9 @@ def list_managed_instances( ecs_models.DescribeInstancesRequestTag( key=provision.MANAGED_BY_TAG_KEY, value=provision.MANAGED_BY_TAG_VALUE, + ), + ecs_models.DescribeInstancesRequestTag( + key=provision.REPO_TAG_KEY, value=repo ) ], max_results=100, @@ -201,7 +204,7 @@ def list_managed_instances( def sweep(client, region_id: str, repo: str, github_token: str, ttl: timedelta) -> int: - instances = list_managed_instances(client, region_id) + instances = list_managed_instances(client, region_id, repo) print(f"Found {len(instances)} managed instance(s) in {region_id}", flush=True) expired = expired_instance_names(instances, datetime.now(timezone.utc), ttl) ok = True @@ -227,7 +230,7 @@ def main() -> int: parser.add_argument( "--sweep", action="store_true", - help="Janitor mode: delete all managed instances older than --sweep-ttl-hours.", + help="Janitor mode: delete expired managed instances tagged for --repo.", ) parser.add_argument("--sweep-ttl-hours", type=float, default=SWEEP_TTL.total_seconds() / 3600) args = parser.parse_args() diff --git a/tests/perf/test_aliyun_ecs_runner_scripts.py b/tests/perf/test_aliyun_ecs_runner_scripts.py index 652eae77cfe..5ba73857add 100644 --- a/tests/perf/test_aliyun_ecs_runner_scripts.py +++ b/tests/perf/test_aliyun_ecs_runner_scripts.py @@ -143,6 +143,7 @@ class ProvisionResourcesTest(unittest.TestCase): expected = { provision.MANAGED_BY_TAG_KEY: provision.MANAGED_BY_TAG_VALUE, provision.RUN_TAG_KEY: "12345", + "github-repository": "owner/repo", } if ttl is not None: expected[provision.TTL_TAG_KEY] = ttl @@ -423,7 +424,7 @@ class TeardownExpiryTest(unittest.TestCase): sys.modules, {"alibabacloud_ecs20140526": SimpleNamespace(models=models)} ): self.assertEqual( - teardown.list_managed_instances(client, "region"), + teardown.list_managed_instances(client, "region", "owner/repo"), [ ("i-1", "one", "2026-08-17T01:00Z", "8"), ("i-2", "two", "2026-08-17T02:00Z", None), @@ -432,11 +433,49 @@ class TeardownExpiryTest(unittest.TestCase): self.assertEqual( client.describe_instances.call_args_list[1].args[0].next_token, "next" ) - tag = client.describe_instances.call_args_list[0].args[0].tag[0] - self.assertEqual( - (tag.key, tag.value), - (provision.MANAGED_BY_TAG_KEY, provision.MANAGED_BY_TAG_VALUE), + for call in client.describe_instances.call_args_list: + self.assertEqual( + {tag.key: tag.value for tag in call.args[0].tag}, + {"managed-by": "query-regression-ci", "github-repository": "owner/repo"}, + ) + + def test_sweep_isolates_repositories_and_untagged_instances(self): + models = SimpleNamespace( + DescribeInstancesRequest=SimpleNamespace, + DescribeInstancesRequestTag=SimpleNamespace, ) + instances = [ + ("own", {"managed-by": "query-regression-ci", "github-repository": "owner/repo"}), + ("other", {"managed-by": "query-regression-ci", "github-repository": "owner/other"}), + ("legacy", {"managed-by": "query-regression-ci"}), + ("unmanaged", {"github-repository": "owner/repo"}), + ] + + def describe(request): + # Model ECS tag intersection using the actual request filters. + selected = [ + SimpleNamespace(instance_id=identity, instance_name=identity, + creation_time="2020-01-01T00:00Z", tags=None) + for identity, tags in instances + if all(tags.get(tag.key) == tag.value for tag in request.tag) + ] + return SimpleNamespace(body=SimpleNamespace( + instances=SimpleNamespace(instance=selected), next_token=None)) + + with ( + patch.dict(sys.modules, {"alibabacloud_ecs20140526": SimpleNamespace(models=models)}), + patch.object(teardown, "delete_instance", return_value=True) as delete, + patch.object(teardown, "deregister_runner", return_value=True) as deregister, + ): + client = Mock() + client.describe_instances.side_effect = describe + for repo, expected in (("owner/repo", "own"), ("owner/other", "other")): + with self.subTest(repo=repo): + delete.reset_mock() + deregister.reset_mock() + self.assertEqual(teardown.sweep(client, "region", repo, "token", self.TTL), 0) + delete.assert_called_once_with(client, expected, "region") + deregister.assert_called_once_with("token", repo, expected) if __name__ == "__main__": From 91a59b82fbff9d1443968024dc5df1c96e661d4e Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Mon, 26 Oct 2026 09:49:13 +0000 Subject: [PATCH 098/239] fix: acknowledge writes only within the region's leader lease (#9904) * fix: acknowledge writes only within the region's leader lease Signed-off-by: jeremyhi * fix: drop cached logical columns after a failed logical drop Signed-off-by: jeremyhi * fix: keep the prepared region for the lease check after the WAL append Signed-off-by: jeremyhi * fix: serialize a logical drop with fills of the logical column cache Signed-off-by: jeremyhi * test: check the drop lock and grant long leases deterministically Signed-off-by: jeremyhi * test: cover a bulk insert that loses its lease during the WAL append Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- src/datanode/src/alive_keeper.rs | 87 +++++- src/datanode/src/datanode.rs | 5 + src/datanode/src/region_server.rs | 17 ++ src/datanode/src/region_server/flow_write.rs | 9 + src/datanode/src/tests.rs | 9 + src/file-engine/src/engine.rs | 11 + src/metric-engine/src/engine.rs | 24 ++ src/metric-engine/src/engine/alter.rs | 22 +- src/metric-engine/src/engine/create.rs | 14 +- src/metric-engine/src/engine/drop.rs | 20 +- src/metric-engine/src/engine/put.rs | 136 ++++++++- src/metric-engine/src/metadata_region.rs | 8 +- src/mito2/src/engine.rs | 33 +++ src/mito2/src/engine/lease_test.rs | 284 +++++++++++++++++++ src/mito2/src/engine/listener.rs | 40 ++- src/mito2/src/error.rs | 20 ++ src/mito2/src/region.rs | 83 +++++- src/mito2/src/region/opener.rs | 2 + src/mito2/src/region_write_ctx.rs | 16 +- src/mito2/src/worker.rs | 26 +- src/mito2/src/worker/handle_catchup.rs | 4 + src/mito2/src/worker/handle_create.rs | 2 + src/mito2/src/worker/handle_write.rs | 59 +++- src/query/src/optimizer/test_util.rs | 9 + src/store-api/src/region_engine.rs | 13 + 25 files changed, 917 insertions(+), 36 deletions(-) create mode 100644 src/mito2/src/engine/lease_test.rs diff --git a/src/datanode/src/alive_keeper.rs b/src/datanode/src/alive_keeper.rs index dbfb699036c..bef59d715c5 100644 --- a/src/datanode/src/alive_keeper.rs +++ b/src/datanode/src/alive_keeper.rs @@ -278,6 +278,9 @@ impl HeartbeatResponseHandler for RegionAliveKeeper { let start_instant = self.epoch + Duration::from_millis(region_lease.duration_since_epoch); let deadline = start_instant + Duration::from_secs(region_lease.lease_seconds); + if let Some(mito) = self.region_server.mito_engine() { + mito.renew_node_lease(deadline.into_std()); + } self.renew_region_leases(®ion_lease.regions, deadline) .await; self.close_staled_regions(®ion_lease.closeable_region_ids) @@ -435,7 +438,7 @@ impl CountdownTask { }, Some(CountdownCommand::Reset((role, deadline, extension_info))) => { let prev_role_leader = self.region_server.is_region_leader(self.region_id).unwrap_or(false); - if let Err(err) = self.region_server.set_region_role(self.region_id, role) { + if let Err(err) = self.region_server.set_region_role_with_lease(self.region_id, role, deadline.into_std()) { if err.status_code() == StatusCode::RegionNotFound { // Table metadata in metasrv is deleted after its regions are dropped. // The datanode may still receive lease renewal responses that depend on the metadata @@ -500,14 +503,92 @@ impl CountdownTask { #[cfg(test)] mod test { + use api::v1::Rows; + use api::v1::meta::{HeartbeatResponse, RegionLease}; + use common_meta::heartbeat::mailbox::HeartbeatMailbox; use mito2::config::MitoConfig; - use mito2::test_util::{CreateRequestBuilder, TestEnv}; + use mito2::test_util::{CreateRequestBuilder, TestEnv, build_rows, rows_schema}; use store_api::region_engine::RegionEngine; - use store_api::region_request::{EnterStagingRequest, StagingPartitionDirective}; + use store_api::region_request::{ + EnterStagingRequest, RegionPutRequest, StagingPartitionDirective, + }; use super::*; use crate::tests::mock_region_server; + /// Handles a heartbeat response that grants `regions` for a heartbeat sent now. + async fn handle_lease_response(alive_keeper: &RegionAliveKeeper, regions: Vec) { + let (tx, _rx) = mpsc::channel(1); + let response = HeartbeatResponse { + region_lease: Some(RegionLease { + regions, + duration_since_epoch: alive_keeper.epoch().elapsed().as_millis() as u64, + // The countdown of a granted region does not expire during the test. + lease_seconds: 86400, + closeable_region_ids: vec![], + }), + ..Default::default() + }; + let mut ctx = + HeartbeatResponseHandlerContext::new(Arc::new(HeartbeatMailbox::new(tx)), response); + alive_keeper.handle(&mut ctx).await.unwrap(); + } + + #[tokio::test(flavor = "multi_thread")] + async fn lease_responses_let_regions_write() { + common_telemetry::init_default_ut_logging(); + let mut region_server = mock_region_server(); + let mut engine_env = TestEnv::with_prefix("region-alive-keeper-write-lease").await; + let engine = engine_env.create_engine(MitoConfig::default()).await; + engine.require_write_lease(); + // The clock of the lease checks stands still, so only the lease propagation matters. + engine.freeze_lease_clock(std::time::Instant::now()); + region_server.register_engine(Arc::new(engine)); + let alive_keeper = + RegionAliveKeeper::new(region_server.clone(), None, Duration::from_millis(100)); + let create = |region_id| { + let request = RegionRequest::Create(CreateRequestBuilder::new().build()); + region_server.handle_request(region_id, request) + }; + let schema = rows_schema(&CreateRequestBuilder::new().build()); + let put = |region_id| { + let rows = Rows { + schema: schema.clone(), + rows: build_rows(0, 1), + }; + let request = RegionRequest::Put(RegionPutRequest { + rows, + hint: None, + partition_expr_version: None, + skip_wal: false, + }); + region_server.handle_request(region_id, request) + }; + + // Created before any heartbeat response, the region has no lease. + let region_id = RegionId::new(1024, 1); + create(region_id).await.unwrap(); + assert!(put(region_id).await.unwrap_err().is_retryable()); + + // A response renews the lease of the node, under which regions created later write. + handle_lease_response(&alive_keeper, vec![]).await; + let created_region_id = RegionId::new(1024, 2); + create(created_region_id).await.unwrap(); + put(created_region_id).await.unwrap(); + + // A Leader grant renews the lease of the region. + alive_keeper.register_region(region_id).await; + let grant = GrantedRegion { + region_id: region_id.as_u64(), + role: api::v1::meta::RegionRole::Leader.into(), + extensions: HashMap::new(), + }; + handle_lease_response(&alive_keeper, vec![grant]).await; + // The countdown task has applied the grant once it answers. + alive_keeper.deadline(region_id).await.unwrap(); + put(region_id).await.unwrap(); + } + #[tokio::test(flavor = "multi_thread")] async fn region_alive_keeper() { common_telemetry::init_default_ut_logging(); diff --git a/src/datanode/src/datanode.rs b/src/datanode/src/datanode.rs index df12cf4f2a8..c304aaa6e00 100644 --- a/src/datanode/src/datanode.rs +++ b/src/datanode/src/datanode.rs @@ -370,6 +370,11 @@ impl DatanodeBuilder { cache_registry.clone(), ) .await?; + // A region acknowledges only writes whose WAL append completed within its lease from + // the metasrv. + if controlled_by_metasrv && let Some(mito) = region_server.mito_engine() { + mito.require_write_lease(); + } // TODO(weny): Considering introducing a readonly kv_backend trait. let runtime_switch_manager = RuntimeSwitchManager::new(self.kv_backend.clone()); diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index 3acf188ad19..06b9d42fd00 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -471,6 +471,23 @@ impl RegionServer { .with_context(|_| HandleRegionRequestSnafu { region_id }) } + /// Sets the role of a region from a lease the metasrv grants until `deadline`. + pub fn set_region_role_with_lease( + &self, + region_id: RegionId, + role: RegionRole, + deadline: std::time::Instant, + ) -> Result<()> { + let engine = self + .inner + .region_map + .get(®ion_id) + .with_context(|| RegionNotFoundSnafu { region_id })?; + engine + .set_region_role_with_lease(region_id, role, deadline) + .with_context(|_| HandleRegionRequestSnafu { region_id }) + } + /// Set region role state gracefully. /// /// For [SettableRegionRoleState::Follower]: diff --git a/src/datanode/src/region_server/flow_write.rs b/src/datanode/src/region_server/flow_write.rs index 6bcb112bc78..bc9bf65d69d 100644 --- a/src/datanode/src/region_server/flow_write.rs +++ b/src/datanode/src/region_server/flow_write.rs @@ -430,6 +430,15 @@ mod tests { ) -> Result<(), common_error::ext::BoxedError> { self.inner.set_region_role(id, role) } + + fn set_region_role_with_lease( + &self, + id: RegionId, + role: store_api::region_engine::RegionRole, + deadline: std::time::Instant, + ) -> Result<(), common_error::ext::BoxedError> { + self.inner.set_region_role_with_lease(id, role, deadline) + } async fn sync_region( &self, id: RegionId, diff --git a/src/datanode/src/tests.rs b/src/datanode/src/tests.rs index 793ae057076..4fa247d07ca 100644 --- a/src/datanode/src/tests.rs +++ b/src/datanode/src/tests.rs @@ -289,6 +289,15 @@ impl RegionEngine for MockRegionEngine { Ok(()) } + fn set_region_role_with_lease( + &self, + _region_id: RegionId, + _role: RegionRole, + _deadline: std::time::Instant, + ) -> Result<(), BoxedError> { + Ok(()) + } + async fn set_region_role_state_gracefully( &self, region_id: RegionId, diff --git a/src/file-engine/src/engine.rs b/src/file-engine/src/engine.rs index 05be2937703..be1d141bcbd 100644 --- a/src/file-engine/src/engine.rs +++ b/src/file-engine/src/engine.rs @@ -15,6 +15,7 @@ use std::any::Any; use std::collections::HashMap; use std::sync::{Arc, RwLock}; +use std::time::Instant; use api::region::RegionResponse; use async_trait::async_trait; @@ -131,6 +132,16 @@ impl RegionEngine for FileRegionEngine { .map_err(BoxedError::new) } + /// Regions of this engine take no writes, so the lease is not needed. + fn set_region_role_with_lease( + &self, + region_id: RegionId, + role: RegionRole, + _deadline: Instant, + ) -> Result<(), BoxedError> { + self.set_region_role(region_id, role) + } + async fn set_region_role_state_gracefully( &self, region_id: RegionId, diff --git a/src/metric-engine/src/engine.rs b/src/metric-engine/src/engine.rs index af15ea087c9..0c1b1114448 100644 --- a/src/metric-engine/src/engine.rs +++ b/src/metric-engine/src/engine.rs @@ -31,6 +31,7 @@ mod sync; use std::any::Any; use std::collections::HashMap; use std::sync::{Arc, RwLock}; +use std::time::Instant; use api::region::RegionResponse; use async_trait::async_trait; @@ -383,6 +384,29 @@ impl RegionEngine for MetricEngine { Ok(()) } + fn set_region_role_with_lease( + &self, + region_id: RegionId, + role: RegionRole, + deadline: Instant, + ) -> Result<(), BoxedError> { + // ignore the region not found error + for x in [ + utils::to_metadata_region_id(region_id), + utils::to_data_region_id(region_id), + ] { + if let Err(e) = self + .inner + .mito + .set_region_role_with_lease(x, role, deadline) + && e.status_code() != StatusCode::RegionNotFound + { + return Err(e); + } + } + Ok(()) + } + async fn sync_region( &self, region_id: RegionId, diff --git a/src/metric-engine/src/engine/alter.rs b/src/metric-engine/src/engine/alter.rs index 2bf59edda0f..e940fa4cbbb 100644 --- a/src/metric-engine/src/engine/alter.rs +++ b/src/metric-engine/src/engine/alter.rs @@ -196,20 +196,30 @@ impl MetricEngineInner { (name.to_string(), column_metadata.clone()) }); + // The data region has the new columns, whatever happens to the metadata write. + self.state + .write() + .unwrap() + .add_physical_columns(data_region_id, new_add_columns); + // Writes logical regions metadata to metadata region - self.metadata_region + let result = self + .metadata_region .add_logical_regions(physical_region_id, false, logical_region_columns) - .await?; + .await; + // A failed write may still be applied, e.g. when the region lease expires after the WAL + // append, so the cached logical columns are dropped either way. + self.state + .write() + .unwrap() + .invalid_logical_regions_cache(requests.iter().map(|(region_id, _)| *region_id)); + result?; extension_return_value.insert( ALTER_PHYSICAL_EXTENSION_KEY.to_string(), ColumnMetadata::encode_list(&physical_columns).context(SerializeColumnMetadataSnafu)?, ); - let mut state = self.state.write().unwrap(); - state.add_physical_columns(data_region_id, new_add_columns); - state.invalid_logical_regions_cache(requests.iter().map(|(region_id, _)| *region_id)); - Ok(0) } diff --git a/src/metric-engine/src/engine/create.rs b/src/metric-engine/src/engine/create.rs index 2f27a405b36..1061565a80e 100644 --- a/src/metric-engine/src/engine/create.rs +++ b/src/metric-engine/src/engine/create.rs @@ -324,6 +324,11 @@ impl MetricEngineInner { let column_metadata = *physical_schema_map.get(name).unwrap(); (name.to_string(), column_metadata.clone()) }); + // The data region has the new columns, whatever happens to the metadata write. + self.state + .write() + .unwrap() + .add_physical_columns(data_region_id, new_add_columns); extension_return_value.insert( ALTER_PHYSICAL_EXTENSION_KEY.to_string(), @@ -335,11 +340,10 @@ impl MetricEngineInner { .add_logical_regions(physical_region_id, true, logical_region_columns) .await?; - { - let mut state = self.state.write().unwrap(); - state.add_physical_columns(data_region_id, new_add_columns); - state.add_logical_regions(physical_region_id, logical_regions.clone()); - } + self.state + .write() + .unwrap() + .add_logical_regions(physical_region_id, logical_regions.clone()); for logical_region_id in logical_regions { self.metadata_region .open_logical_region(logical_region_id) diff --git a/src/metric-engine/src/engine/drop.rs b/src/metric-engine/src/engine/drop.rs index 7e98effdeee..f7db294478c 100644 --- a/src/metric-engine/src/engine/drop.rs +++ b/src/metric-engine/src/engine/drop.rs @@ -157,10 +157,24 @@ impl MetricEngineInner { logical_region_id: RegionId, physical_region_id: RegionId, ) -> Result { - // Update metadata - self.metadata_region - .remove_logical_region(physical_region_id, logical_region_id) + // A fill of the cached logical columns holds the read lock until it publishes them, so + // none publishes columns read before the delete after the invalidation below. + let _write_guard = self + .metadata_region + .write_lock_logical_region(logical_region_id) .await?; + // Update metadata + let result = self + .metadata_region + .remove_logical_region(physical_region_id, logical_region_id) + .await; + // A failed delete may still be applied, e.g. when the region lease expires after the WAL + // append, so the cached logical columns are dropped either way. + self.state + .write() + .unwrap() + .invalid_logical_regions_cache([logical_region_id]); + result?; // Update engine state self.state diff --git a/src/metric-engine/src/engine/put.rs b/src/metric-engine/src/engine/put.rs index 27ee3fba56c..e8098131612 100644 --- a/src/metric-engine/src/engine/put.rs +++ b/src/metric-engine/src/engine/put.rs @@ -785,6 +785,8 @@ impl MetricEngineInner { #[cfg(test)] mod tests { use std::collections::HashSet; + use std::sync::Arc; + use std::time::{Duration, Instant}; use api::v1::value::ValueData; use api::v1::{ColumnDataType, ColumnSchema as PbColumnSchema}; @@ -799,6 +801,10 @@ mod tests { use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnDefaultConstraint, ColumnSchema}; use datatypes::value::Value as PartitionValue; + use futures_util::FutureExt; + use mito2::config::MitoConfig; + use mito2::engine::listener::WriteWalListener; + use mito2::test_util::TestEnv as MitoTestEnv; use partition::expr::col; use store_api::metadata::ColumnMetadata; use store_api::metric_engine_consts::{ @@ -806,10 +812,10 @@ mod tests { PHYSICAL_TABLE_METADATA_KEY, PRIMARY_KEY_ENCODING, }; use store_api::path_utils::table_dir; - use store_api::region_engine::RegionEngine; + use store_api::region_engine::{RegionEngine, RegionRole}; use store_api::region_request::{ - EnterStagingRequest, PathType, RegionCloseRequest, RegionOpenRequest, RegionRequest, - StagingPartitionDirective, + EnterStagingRequest, PathType, RegionCloseRequest, RegionDropRequest, RegionOpenRequest, + RegionRequest, StagingPartitionDirective, }; use store_api::storage::ScanRequest; use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME; @@ -2350,4 +2356,128 @@ mod tests { "expected impure-default rejection, got: {err}" ); } + + #[tokio::test] + async fn test_logical_region_ddl_needs_leader_lease() { + common_telemetry::init_default_ut_logging(); + let mut mito_env = MitoTestEnv::with_prefix("logical-region-lease").await; + let listener = Arc::new(WriteWalListener::default()); + let mito = mito_env + .create_engine_with(MitoConfig::default(), None, Some(listener.clone()), None) + .await; + mito.require_write_lease(); + let start = Instant::now(); + mito.freeze_lease_clock(start); + let env = TestEnv::with_mito_engine(mito_env, mito.clone()); + let engine = env.metric(); + // Created before the node has a lease, the physical region has none. + let physical_region_id = env.default_physical_region_id(); + env.create_physical_region(physical_region_id, &TestEnv::default_table_dir(), vec![]) + .await; + // The clock of the lease checks stands at `start` until `expire(n)` moves it to `at(n)`, + // where a lease from `grant(n)` expires. + let at = |n| start + n * Duration::from_secs(10); + let grant = |n| { + engine + .set_region_role_with_lease(physical_region_id, RegionRole::Leader, at(n)) + .unwrap() + }; + let expire = |n| mito.freeze_lease_clock(at(n)); + let logical_region_id = env.default_logical_region_id(); + let alter = |tag| RegionRequest::Alter(test_util::alter_logical_region_request(&[tag])); + let drop_request = || { + RegionRequest::Drop(RegionDropRequest { + fast_path: false, + force: false, + partial_drop: false, + }) + }; + let put = |tags: &[&str]| { + let request = RegionRequest::Put(RegionPutRequest { + rows: Rows { + schema: test_util::row_schema_with_tags(tags), + rows: test_util::build_rows(tags.len(), 5), + }, + hint: None, + partition_expr_version: None, + skip_wal: false, + }); + engine.handle_request(logical_region_id, request) + }; + // Holds `request` after its metadata region append while the lease from `grant(n)` + // expires. With `blocks_fills`, a fill of the cached logical columns, which takes the read + // lock of the logical region, cannot start until the request finishes. + let expire_after_append = async |request, n, blocks_fills: bool| { + listener.pause_next(); + let mut task = { + let engine = engine.clone(); + tokio::spawn(async move { engine.handle_request(logical_region_id, request).await }) + }; + tokio::select! { + () = listener.wait_appended() => {} + result = &mut task => panic!("the request finished before its append: {result:?}"), + } + if blocks_fills { + let metadata_region = &engine.inner.metadata_region; + let fill = metadata_region.read_lock_logical_region(logical_region_id); + assert!(fill.now_or_never().is_none()); + } + expire(n); + listener.resume(); + task.await.unwrap() + }; + + // Each DDL below gives the data region a new column before the metadata region rejects + // its write; the retry must not add the column again. + let create = || { + RegionRequest::Create(test_util::create_logical_region_request( + &["job"], + physical_region_id, + &table_dir("test", logical_region_id.table_id()), + )) + }; + let err = engine + .handle_request(logical_region_id, create()) + .await + .unwrap_err(); + assert!(err.is_retryable(), "{err:?}"); + grant(1); + engine + .handle_request(logical_region_id, create()) + .await + .unwrap(); + put(&["job"]).await.unwrap(); + + // The metadata write of the ALTER is applied, so the cached metadata must not hide it. + let err = expire_after_append(alter("tag1"), 1, false) + .await + .unwrap_err(); + assert!(err.is_retryable(), "{err:?}"); + let metadata = engine.get_metadata(logical_region_id).await.unwrap(); + assert!(metadata.column_by_name("tag1").is_some()); + + let err = engine + .handle_request(logical_region_id, alter("tag2")) + .await + .unwrap_err(); + assert!(err.is_retryable(), "{err:?}"); + grant(2); + engine + .handle_request(logical_region_id, alter("tag2")) + .await + .unwrap(); + put(&["job", "tag1", "tag2"]).await.unwrap(); + + // The metadata delete of the DROP is applied, so writes must not use the cached columns. + let err = expire_after_append(drop_request(), 2, true) + .await + .unwrap_err(); + assert!(err.is_retryable(), "{err:?}"); + grant(3); + put(&["job"]).await.unwrap_err(); + engine + .handle_request(logical_region_id, drop_request()) + .await + .unwrap(); + } } diff --git a/src/metric-engine/src/metadata_region.rs b/src/metric-engine/src/metadata_region.rs index 18fa51f3597..7f797ff9727 100644 --- a/src/metric-engine/src/metadata_region.rs +++ b/src/metric-engine/src/metadata_region.rs @@ -538,11 +538,15 @@ impl MetadataRegion { ) -> Result<()> { let cache_access_lock = self.cache_access_lock(metadata_region_id); let _cache_guard = cache_access_lock.write().await; - self.mito + let result = self + .mito .handle_request(metadata_region_id, request) .await - .context(MitoWriteOperationSnafu)?; + .context(MitoWriteOperationSnafu); + // A failed write may still be applied, e.g. when the region lease expires after the + // WAL append. self.cache.invalidate(&metadata_region_id).await; + result?; Ok(()) } diff --git a/src/mito2/src/engine.rs b/src/mito2/src/engine.rs index 92d848f5677..0147fb71e2a 100644 --- a/src/mito2/src/engine.rs +++ b/src/mito2/src/engine.rs @@ -48,6 +48,8 @@ mod filter_deleted_test; mod flush_test; #[cfg(test)] mod index_build_test; +#[cfg(test)] +mod lease_test; #[cfg(any(test, feature = "test"))] pub mod listener; #[cfg(test)] @@ -327,6 +329,22 @@ impl MitoEngine { self.inner.workers.gc_limiter() } + /// Makes regions acknowledge only writes whose WAL append completes within their lease. + pub fn require_write_lease(&self) { + self.inner.workers.node_lease().enable(); + } + + /// Renews the lease of this node from a heartbeat response that carries region leases. + pub fn renew_node_lease(&self, deadline: Instant) { + self.inner.workers.node_lease().renew(deadline); + } + + /// Freezes the clock of the lease checks at `at`. + #[cfg(any(test, feature = "test"))] + pub fn freeze_lease_clock(&self, at: Instant) { + self.inner.workers.node_lease().freeze_clock(at); + } + pub fn object_store_manager(&self) -> &ObjectStoreManagerRef { self.inner.workers.object_store_manager() } @@ -1581,6 +1599,21 @@ impl RegionEngine for MitoEngine { .map_err(BoxedError::new) } + fn set_region_role_with_lease( + &self, + region_id: RegionId, + role: RegionRole, + deadline: Instant, + ) -> Result<(), BoxedError> { + let region = self.inner.find_region(region_id).map_err(BoxedError::new)?; + // Before the role, and even if the role leaves the region's state unchanged. + if role.writable() { + region.leader_deadline.extend(deadline); + } + region.set_role(role); + Ok(()) + } + async fn set_region_role_state_gracefully( &self, region_id: RegionId, diff --git a/src/mito2/src/engine/lease_test.rs b/src/mito2/src/engine/lease_test.rs new file mode 100644 index 00000000000..a0e68f7e94b --- /dev/null +++ b/src/mito2/src/engine/lease_test.rs @@ -0,0 +1,284 @@ +// 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. + +//! Tests that a region acknowledges only writes whose WAL append completed within its leader +//! lease. + +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use api::region::RegionResponse; +use api::v1::{ColumnSchema, Rows}; +use common_error::ext::{BoxedError, ErrorExt}; +use common_recordbatch::RecordBatches; +use rstest::rstest; +use rstest_reuse::apply; +use store_api::region_engine::{RegionEngine, RegionRole}; +use store_api::region_request::{RegionCatchupRequest, RegionPutRequest, RegionRequest}; +use store_api::storage::{RegionId, ScanRequest}; +use tokio::task::JoinHandle; + +use crate::config::MitoConfig; +use crate::engine::MitoEngine; +use crate::engine::basic_test::build_bulk_insert_request; +use crate::engine::listener::WriteWalListener; +use crate::error::Error; +use crate::test_util::{ + CreateRequestBuilder, LogStoreFactory, TestEnv, build_rows, kafka_log_store_factory, + multiple_log_store_factories, prepare_test_for_kafka_log_store, raft_engine_log_store_factory, + rows_schema, +}; + +const LEASE: Duration = Duration::from_secs(10); +const REGION_ID: RegionId = RegionId::new(1, 1); + +/// Creates [REGION_ID] on an engine whose writes need a lease. The node has no lease yet, and +/// the clock of the lease checks stands still until a test moves it. +async fn create_region( + env: &mut TestEnv, + topic: Option, +) -> (Arc, MitoEngine, Vec) { + let listener = Arc::new(WriteWalListener::default()); + let engine = env + .create_engine_with(MitoConfig::default(), None, Some(listener.clone()), None) + .await; + engine.require_write_lease(); + engine.freeze_lease_clock(Instant::now()); + let request = CreateRequestBuilder::new().kafka_topic(topic).build(); + let schema = rows_schema(&request); + engine + .handle_request(REGION_ID, RegionRequest::Create(request)) + .await + .unwrap(); + (listener, engine, schema) +} + +fn grant(engine: &MitoEngine, role: RegionRole, deadline: Instant) { + engine + .set_region_role_with_lease(REGION_ID, role, deadline) + .unwrap(); +} + +/// Returns a request that writes rows `start..start + 3`. +fn put_request(schema: &[ColumnSchema], start: usize) -> RegionRequest { + let rows = Rows { + schema: schema.to_vec(), + rows: build_rows(start, start + 3), + }; + RegionRequest::Put(RegionPutRequest { + rows, + hint: None, + partition_expr_version: None, + skip_wal: false, + }) +} + +fn put( + engine: &MitoEngine, + region_id: RegionId, + schema: &[ColumnSchema], + start: usize, +) -> JoinHandle> { + let engine = engine.clone(); + let request = put_request(schema, start); + tokio::spawn(async move { engine.handle_request(region_id, request).await }) +} + +/// Sends `request` to [REGION_ID] and waits until it is held after its WAL append. +async fn hold( + listener: &WriteWalListener, + engine: &MitoEngine, + request: RegionRequest, +) -> JoinHandle> { + listener.pause_next(); + let engine = engine.clone(); + let mut write = tokio::spawn(async move { engine.handle_request(REGION_ID, request).await }); + tokio::select! { + () = listener.wait_appended() => {} + result = &mut write => panic!("the write finished before its WAL append: {result:?}"), + } + write +} + +async fn count_rows(engine: &MitoEngine) -> usize { + let stream = engine + .scan_to_stream(REGION_ID, ScanRequest::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + batches.iter().map(|batch| batch.num_rows()).sum() +} + +fn assert_retryable(result: Result, expected: fn(&Error) -> bool) { + let err = result.unwrap_err(); + assert!(err.retry_hint().is_retryable(), "{err:?}"); + let err = err.into_inner(); + let err = match err.as_any().downcast_ref::().unwrap() { + Error::WriteGroup { source } => source.as_ref(), + err => err, + }; + assert!(expected(err), "{err:?}"); +} + +fn lease_expired(err: &Error) -> bool { + matches!(err, Error::RegionLeaseExpired { .. }) +} + +fn unacknowledged(err: &Error) -> bool { + matches!(err, Error::UnacknowledgedWrite { .. }) +} + +#[apply(multiple_log_store_factories)] +async fn test_lease_expires_while_writing_wal(factory: Option) { + common_telemetry::init_default_ut_logging(); + let Some(factory) = factory else { + return; + }; + let mut env = TestEnv::with_prefix("lease-expires-while-writing-wal") + .await + .with_log_store_factory(factory.clone()); + let topic = prepare_test_for_kafka_log_store(&factory).await; + let (listener, engine, schema) = create_region(&mut env, topic).await; + let deadline = Instant::now() + LEASE; + grant(&engine, RegionRole::Leader, deadline); + + let write = hold(&listener, &engine, put_request(&schema, 0)).await; + engine.freeze_lease_clock(deadline); + listener.resume(); + assert_retryable(write.await.unwrap(), unacknowledged); + // The write is in the WAL, so it is applied rather than undone. + assert_eq!(3, count_rows(&engine).await); + + // A new write fails before it reaches the WAL. + assert_retryable( + put(&engine, REGION_ID, &schema, 3).await.unwrap(), + lease_expired, + ); + assert_eq!(3, count_rows(&engine).await); + + // The same holds for a bulk insert. + let deadline = deadline + LEASE; + grant(&engine, RegionRole::Leader, deadline); + let bulk_insert = RegionRequest::BulkInserts(build_bulk_insert_request(REGION_ID, 3)); + let write = hold(&listener, &engine, bulk_insert).await; + let committed_sequence = || { + let region = engine.get_region(REGION_ID).unwrap(); + region.version_control.current().committed_sequence + }; + let sequence = committed_sequence(); + engine.freeze_lease_clock(deadline); + listener.resume(); + assert_retryable(write.await.unwrap(), unacknowledged); + assert_eq!(6, count_rows(&engine).await); + assert!(committed_sequence() > sequence); +} + +#[tokio::test] +async fn test_lease_renewed_while_writing_wal() { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::with_prefix("lease-renewed-while-writing-wal").await; + let (listener, engine, schema) = create_region(&mut env, None).await; + let deadline = Instant::now() + LEASE; + grant(&engine, RegionRole::Leader, deadline); + + let write = hold(&listener, &engine, put_request(&schema, 0)).await; + // The grant keeps the region writable rather than staging, and still extends the deadline. + grant(&engine, RegionRole::StagingLeader, deadline + LEASE); + engine.freeze_lease_clock(deadline); + listener.resume(); + assert_eq!(3, write.await.unwrap().unwrap().affected_rows); +} + +#[rstest] +#[case::within_leader_lease(false)] +#[case::after_leader_lease(true)] +#[tokio::test] +async fn test_downgrade_while_writing_wal(#[case] lease_expires: bool) { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::with_prefix("downgrade-while-writing-wal").await; + let (listener, engine, schema) = create_region(&mut env, None).await; + let deadline = Instant::now() + LEASE; + grant(&engine, RegionRole::Leader, deadline); + + let write = hold(&listener, &engine, put_request(&schema, 0)).await; + // The metasrv renews the lease as DowngradingLeader once the route marks the region as + // downgrading. A write already in the WAL is acknowledged, but only within the Leader + // deadline, which such grants do not extend. + grant(&engine, RegionRole::DowngradingLeader, deadline + LEASE); + if lease_expires { + engine.freeze_lease_clock(deadline); + } + listener.resume(); + let result = write.await.unwrap(); + if lease_expires { + assert_retryable(result, unacknowledged); + } else { + assert_eq!(3, result.unwrap().affected_rows); + } +} + +#[tokio::test] +async fn test_created_region_writes_under_node_lease() { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::with_prefix("created-region-node-lease").await; + let (_listener, engine, schema) = create_region(&mut env, None).await; + let bulk_insert = |region_id| { + engine.handle_request( + region_id, + RegionRequest::BulkInserts(build_bulk_insert_request(region_id, 3)), + ) + }; + // Created before the node has a lease, the region has no deadline. + assert_retryable( + put(&engine, REGION_ID, &schema, 0).await.unwrap(), + lease_expired, + ); + assert_retryable(bulk_insert(REGION_ID).await, lease_expired); + + // Renewing the node lease does not extend the deadline of a region created before. + engine.renew_node_lease(Instant::now() + LEASE); + assert!(put(&engine, REGION_ID, &schema, 0).await.unwrap().is_err()); + + // A region created under the node lease takes writes before its first Leader grant. + let region_id = RegionId::new(1, 2); + engine + .handle_request( + region_id, + RegionRequest::Create(CreateRequestBuilder::new().build()), + ) + .await + .unwrap(); + put(&engine, region_id, &schema, 0).await.unwrap().unwrap(); + bulk_insert(region_id).await.unwrap(); +} + +#[tokio::test] +async fn test_upgraded_region_writes_before_leader_grant() { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::with_prefix("upgraded-region-node-lease").await; + let (_listener, engine, schema) = create_region(&mut env, None).await; + // A migration candidate is a follower until the upgrade. + grant(&engine, RegionRole::Follower, Instant::now() + LEASE); + engine.renew_node_lease(Instant::now() + LEASE); + + let request = RegionCatchupRequest { + set_writable: true, + ..Default::default() + }; + engine + .handle_request(REGION_ID, RegionRequest::Catchup(request)) + .await + .unwrap(); + put(&engine, REGION_ID, &schema, 0).await.unwrap().unwrap(); +} diff --git a/src/mito2/src/engine/listener.rs b/src/mito2/src/engine/listener.rs index d5e4b8e8854..f754d4faf00 100644 --- a/src/mito2/src/engine/listener.rs +++ b/src/mito2/src/engine/listener.rs @@ -15,7 +15,7 @@ //! Engine event listener for tests. use std::sync::Arc; -use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::time::Duration; use async_trait::async_trait; @@ -39,6 +39,9 @@ pub trait EventListener: Send + Sync { /// Notifies the listener that the engine is stalled. fn on_write_stall(&self) {} + /// Notifies the listener after a batch of writes is appended to the WAL. + async fn on_write_wal_end(&self) {} + /// Notifies the listener that the region starts to do flush. async fn on_flush_begin(&self, region_id: RegionId) { let _ = region_id; @@ -617,3 +620,38 @@ impl EventListener for GateIndexBuildListener { self.stop_notify.notify_one(); } } + +/// Listener to pause a batch of writes after it is appended to the WAL. +#[derive(Default)] +pub struct WriteWalListener { + pause: AtomicBool, + appended_notify: Notify, + resume_notify: Notify, +} + +impl WriteWalListener { + /// Pauses the next batch of writes after it is appended to the WAL. + pub fn pause_next(&self) { + self.pause.store(true, Ordering::Relaxed); + } + + /// Waits until the paused batch is appended to the WAL. + pub async fn wait_appended(&self) { + self.appended_notify.notified().await; + } + + /// Resumes the paused batch. + pub fn resume(&self) { + self.resume_notify.notify_one(); + } +} + +#[async_trait] +impl EventListener for WriteWalListener { + async fn on_write_wal_end(&self) { + if self.pause.swap(false, Ordering::Relaxed) { + self.appended_notify.notify_one(); + self.resume_notify.notified().await; + } + } +} diff --git a/src/mito2/src/error.rs b/src/mito2/src/error.rs index 9bb908720fd..a336fb89673 100644 --- a/src/mito2/src/error.rs +++ b/src/mito2/src/error.rs @@ -656,6 +656,23 @@ pub enum Error { location: Location, }, + #[snafu(display("Region {} does not hold a valid leader lease", region_id))] + RegionLeaseExpired { + region_id: RegionId, + #[snafu(implicit)] + location: Location, + }, + + #[snafu(display( + "Region {} lost its leader lease while writing to the WAL, the write may have been applied", + region_id + ))] + UnacknowledgedWrite { + region_id: RegionId, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Failed to compact region {}", region_id))] CompactRegion { region_id: RegionId, @@ -1584,6 +1601,7 @@ impl ErrorExt for Error { CompatReader { .. } => StatusCode::Unexpected, InvalidRegionRequest { source, .. } => source.status_code(), RegionState { .. } | UpdateManifest { .. } => StatusCode::RegionNotReady, + RegionLeaseExpired { .. } | UnacknowledgedWrite { .. } => StatusCode::RegionNotReady, JsonOptions { .. } => StatusCode::InvalidArguments, EmptyRegionDir { .. } | EmptyManifestDir { .. } => StatusCode::RegionNotFound, ConvertValue { source, .. } => source.status_code(), @@ -1681,6 +1699,8 @@ impl ErrorExt for Error { | RegionNotFound { .. } | RegionClosed { .. } | RejectWriteInStaging { .. } + | RegionLeaseExpired { .. } + | UnacknowledgedWrite { .. } | RegionState { state: RegionRoleState::Follower, .. diff --git a/src/mito2/src/region.rs b/src/mito2/src/region.rs index 613b1ce428e..e6e3de2df1e 100644 --- a/src/mito2/src/region.rs +++ b/src/mito2/src/region.rs @@ -23,8 +23,9 @@ pub(crate) mod version; use std::cell::OnceCell; use std::collections::hash_map::Entry; use std::collections::{HashMap, HashSet}; -use std::sync::atomic::{AtomicI64, AtomicU64, Ordering}; -use std::sync::{Arc, Mutex, RwLock}; +use std::sync::atomic::{AtomicBool, AtomicI64, AtomicU64, Ordering}; +use std::sync::{Arc, LazyLock, Mutex, RwLock}; +use std::time::Instant; use common_base::hash::partition_expr_version; use common_recordbatch::adapter::RegionQueryStatCounters; @@ -185,6 +186,8 @@ pub struct MitoRegion { /// There are no WAL entries in range [flushed_entry_id, topic_latest_entry_id] for current region, /// which means these WAL entries maybe able to be pruned up to `topic_latest_entry_id`. pub(crate) topic_latest_entry_id: AtomicU64, + /// Deadline of the region's latest Leader or StagingLeader lease from the metasrv. + pub(crate) leader_deadline: Deadline, /// Region stats. pub(crate) region_stats: RegionStats, /// manifest stats @@ -221,6 +224,80 @@ impl RegionStats { pub type MitoRegionRef = Arc; +/// The instant lease deadlines count from, so that a deadline fits in an [AtomicU64]. +static LEASE_EPOCH: LazyLock = LazyLock::new(Instant::now); + +fn lease_nanos(instant: Instant) -> u64 { + instant.saturating_duration_since(*LEASE_EPOCH).as_nanos() as u64 +} + +/// A lease deadline in nanoseconds since [LEASE_EPOCH]. It only moves forward; 0 means none. +#[derive(Debug, Default)] +pub(crate) struct Deadline(AtomicU64); + +impl Deadline { + pub(crate) fn extend(&self, deadline: Instant) { + self.0.fetch_max(lease_nanos(deadline), Ordering::Relaxed); + } + + pub(crate) fn get(&self) -> u64 { + self.0.load(Ordering::Relaxed) + } +} + +/// Lease of this node: the latest deadline of the heartbeat responses that carry region leases. +#[derive(Debug, Default)] +pub(crate) struct NodeLease { + /// Whether writes need a lease. Only a datanode the metasrv controls sets it. + enabled: AtomicBool, + deadline: Deadline, + /// The time of the lease checks while a test freezes it, in the unit of [Deadline]; 0 when + /// it is not frozen. + #[cfg(any(test, feature = "test"))] + frozen_now: AtomicU64, +} + +impl NodeLease { + pub(crate) fn enable(&self) { + self.enabled.store(true, Ordering::Relaxed); + } + + pub(crate) fn is_enabled(&self) -> bool { + self.enabled.load(Ordering::Relaxed) + } + + pub(crate) fn renew(&self, deadline: Instant) { + self.deadline.extend(deadline); + } + + /// Returns the current time in the unit of [Deadline]. + pub(crate) fn now(&self) -> u64 { + #[cfg(any(test, feature = "test"))] + { + let frozen_now = self.frozen_now.load(Ordering::Relaxed); + if frozen_now != 0 { + return frozen_now; + } + } + lease_nanos(Instant::now()) + } + + #[cfg(any(test, feature = "test"))] + pub(crate) fn freeze_clock(&self, at: Instant) { + self.frozen_now + .store(lease_nanos(at).max(1), Ordering::Relaxed); + } + + /// Lets a region created or upgraded on this node write until the node's deadline. The + /// metasrv granted that deadline before the creation or upgrade, so a later migration of the + /// region waits past it. Later responses do not extend it; the region's Leader grants do. + pub(crate) fn seed(&self, leader_deadline: &Deadline) { + leader_deadline + .0 + .fetch_max(self.deadline.get(), Ordering::Relaxed); + } +} + #[derive(Debug, Clone)] pub(crate) struct StagingPartitionInfo { pub(crate) partition_directive: StagingPartitionDirective, @@ -2264,6 +2341,7 @@ mod tests { last_schedule_compaction_millis: Default::default(), time_provider: Arc::new(StdTimeProvider), topic_latest_entry_id: Default::default(), + leader_deadline: Default::default(), region_stats: RegionStats::new(), stats: ManifestStats::default(), } @@ -2869,6 +2947,7 @@ mod tests { last_schedule_compaction_millis: Default::default(), time_provider: Arc::new(StdTimeProvider), topic_latest_entry_id: Default::default(), + leader_deadline: Default::default(), region_stats: RegionStats::new(), stats: ManifestStats::default(), }; diff --git a/src/mito2/src/region/opener.rs b/src/mito2/src/region/opener.rs index d10545198f7..947844c8f20 100644 --- a/src/mito2/src/region/opener.rs +++ b/src/mito2/src/region/opener.rs @@ -459,6 +459,7 @@ impl RegionOpener { last_schedule_compaction_millis: AtomicI64::new(now), time_provider: self.time_provider.clone(), topic_latest_entry_id: AtomicU64::new(flushed_entry_id), + leader_deadline: Default::default(), region_stats: RegionStats::new(), stats: self.stats, })) @@ -697,6 +698,7 @@ impl RegionOpener { last_schedule_compaction_millis: AtomicI64::new(now), time_provider: self.time_provider.clone(), topic_latest_entry_id: AtomicU64::new(topic_latest_entry_id), + leader_deadline: Default::default(), region_stats: RegionStats::new(), stats: self.stats.clone(), }; diff --git a/src/mito2/src/region_write_ctx.rs b/src/mito2/src/region_write_ctx.rs index b0b781eb43c..e4e1e505326 100644 --- a/src/mito2/src/region_write_ctx.rs +++ b/src/mito2/src/region_write_ctx.rs @@ -27,6 +27,7 @@ use crate::error::{Error, Result, UnexpectedSnafu, WriteGroupSnafu}; use crate::memtable::KeyValues; use crate::memtable::bulk::part::BulkPart; use crate::metrics; +use crate::region::MitoRegionRef; use crate::region::version::{VersionControlData, VersionControlRef, VersionRef}; use crate::request::OptionOutputTx; use crate::wal::{EntryId, WalWriter}; @@ -131,6 +132,9 @@ pub(crate) struct RegionWriteCtx { pub(crate) delete_num: usize, /// The total bytes written to the region. pub(crate) written_bytes: Option>, + /// The region and its Leader deadline when the write was prepared, if the write needs a + /// lease. + pub(crate) lease: Option<(MitoRegionRef, u64)>, } impl RegionWriteCtx { @@ -167,6 +171,7 @@ impl RegionWriteCtx { delete_num: 0, bulk_parts: vec![], written_bytes, + lease: None, } } @@ -247,6 +252,14 @@ impl RegionWriteCtx { /// Sets error and marks all write operations are failed. pub(crate) fn set_error(&mut self, err: Arc) { + self.fail_waiters(err); + + // Fail the whole write operation. + self.failed = true; + } + + /// Sets error for all notifiers. The writes still go to the memtables. + pub(crate) fn fail_waiters(&mut self, err: Arc) { // Set error for all notifiers. for notify in self .wal_notifiers @@ -258,9 +271,6 @@ impl RegionWriteCtx { for notify in &mut self.bulk_notifiers { notify.err = Some(err.clone()); } - - // Fail the whole write operation. - self.failed = true; } /// Returns whether the write operation is already marked as failed. diff --git a/src/mito2/src/worker.rs b/src/mito2/src/worker.rs index 789b0f539b1..50fb5bc0eb1 100644 --- a/src/mito2/src/worker.rs +++ b/src/mito2/src/worker.rs @@ -71,8 +71,8 @@ use crate::memtable::MemtableBuilderProvider; use crate::metrics::{REGION_COUNT, REQUEST_WAIT_TIME, WRITE_STALLING}; use crate::region::opener::PartitionExprFetcherRef; use crate::region::{ - CatchupRegions, CatchupRegionsRef, MitoRegionRef, OpeningRegions, OpeningRegionsRef, RegionMap, - RegionMapRef, + CatchupRegions, CatchupRegionsRef, MitoRegionRef, NodeLease, OpeningRegions, OpeningRegionsRef, + RegionMap, RegionMapRef, }; use crate::request::{ BackgroundNotify, BulkInsertRequest, DdlRequest, OptionOutputTx, SenderBulkRequest, @@ -157,6 +157,7 @@ pub(crate) struct WorkerGroup { file_ref_manager: FileReferenceManagerRef, /// Gc limiter to limit concurrent gc jobs. gc_limiter: GcLimiterRef, + node_lease: Arc, /// Object store manager. object_store_manager: ObjectStoreManagerRef, /// Puffin manager factory. @@ -250,6 +251,7 @@ impl WorkerGroup { #[cfg(feature = "hdfs-object-store")] let sweep_limiter = (config.gc.max_concurrent_gc_job > 0) .then(|| Arc::new(Semaphore::new(config.gc.max_concurrent_gc_job))); + let node_lease = Arc::new(NodeLease::default()); let workers = (0..config.num_workers) .map(|id| { @@ -280,6 +282,7 @@ impl WorkerGroup { flush_semaphore: flush_semaphore.clone(), #[cfg(feature = "hdfs-object-store")] sweep_limiter: sweep_limiter.clone(), + node_lease: node_lease.clone(), } .start() }) @@ -294,6 +297,7 @@ impl WorkerGroup { cache_manager, file_ref_manager, gc_limiter, + node_lease, object_store_manager, puffin_manager_factory, intermediate_manager, @@ -398,6 +402,10 @@ impl WorkerGroup { self.gc_limiter.clone() } + pub(crate) fn node_lease(&self) -> &NodeLease { + &self.node_lease + } + /// Get worker for specific `region_id`. pub(crate) fn worker(&self, region_id: RegionId) -> &RegionWorker { let index = region_id_to_index(region_id, self.workers.len()); @@ -509,6 +517,7 @@ impl WorkerGroup { #[cfg(feature = "hdfs-object-store")] let sweep_limiter = (config.gc.max_concurrent_gc_job > 0) .then(|| Arc::new(Semaphore::new(config.gc.max_concurrent_gc_job))); + let node_lease = Arc::new(NodeLease::default()); let workers = (0..config.num_workers) .map(|id| { WorkerStarter { @@ -538,6 +547,7 @@ impl WorkerGroup { flush_semaphore: flush_semaphore.clone(), #[cfg(feature = "hdfs-object-store")] sweep_limiter: sweep_limiter.clone(), + node_lease: node_lease.clone(), } .start() }) @@ -552,6 +562,7 @@ impl WorkerGroup { cache_manager, file_ref_manager, gc_limiter, + node_lease, object_store_manager, puffin_manager_factory, intermediate_manager, @@ -652,6 +663,7 @@ struct WorkerStarter { /// Limiter for the temporary file sweeps, or `None` if the datanode allows /// no concurrent GC job. sweep_limiter: Option>, + node_lease: Arc, } impl WorkerStarter { @@ -754,6 +766,7 @@ impl WorkerStarter { plugins: self.plugins, #[cfg(feature = "hdfs-object-store")] sweep_limiter: self.sweep_limiter, + node_lease: self.node_lease, }; let handle = common_runtime::spawn_global(async move { worker_thread.run().await; @@ -1085,6 +1098,7 @@ struct RegionWorkerLoop { /// Limiter for the temporary file sweeps, or `None` if the datanode allows /// no concurrent GC job. sweep_limiter: Option>, + node_lease: Arc, } impl RegionWorkerLoop { @@ -1606,6 +1620,14 @@ impl WorkerListener { } } + /// A batch of writes is appended to the WAL. + pub(crate) async fn on_write_wal_end(&self) { + #[cfg(any(test, feature = "test"))] + if let Some(listener) = &self.listener { + listener.on_write_wal_end().await; + } + } + pub(crate) async fn on_flush_begin(&self, region_id: RegionId) { #[cfg(any(test, feature = "test"))] if let Some(listener) = &self.listener { diff --git a/src/mito2/src/worker/handle_catchup.rs b/src/mito2/src/worker/handle_catchup.rs index f817c1309d8..92874672789 100644 --- a/src/mito2/src/worker/handle_catchup.rs +++ b/src/mito2/src/worker/handle_catchup.rs @@ -75,6 +75,7 @@ impl RegionWorkerLoop { let sweep_older_than = self.config.gc.unknown_file_lingering_time; #[cfg(feature = "hdfs-object-store")] let sweep_limiter = self.sweep_limiter.clone(); + let node_lease = self.node_lease.clone(); common_runtime::spawn_global(async move { let mut task = RegionCatchupTask::new(region.clone(), wal, allow_stale_entries) .with_entry_receiver(entry_receiver) @@ -84,6 +85,9 @@ impl RegionWorkerLoop { match task.run().await { Ok(_) => { if request.set_writable { + // The region takes writes under the node lease until its next Leader + // grant. + node_lease.seed(®ion.leader_deadline); region.set_role(RegionRole::Leader); // Finalize leadership: persist backfilled metadata. match region diff --git a/src/mito2/src/worker/handle_create.rs b/src/mito2/src/worker/handle_create.rs index c31e9d4ac7c..15641385e59 100644 --- a/src/mito2/src/worker/handle_create.rs +++ b/src/mito2/src/worker/handle_create.rs @@ -106,6 +106,8 @@ impl RegionWorkerLoop { hook.on_region_opened(region_id, ®ion.metadata()).await; } + // The region takes writes under the node lease until its first Leader grant. + self.node_lease.seed(®ion.leader_deadline); // Insert the MitoRegion into the RegionMap. self.register_region(region); if let Some(state) = &self.series_index_task_state { diff --git a/src/mito2/src/worker/handle_write.rs b/src/mito2/src/worker/handle_write.rs index 5de0439e26e..7e290a768fb 100644 --- a/src/mito2/src/worker/handle_write.rs +++ b/src/mito2/src/worker/handle_write.rs @@ -25,8 +25,9 @@ use store_api::logstore::LogStore; use store_api::storage::RegionId; use crate::error::{ - InvalidRequestSnafu, PartitionExprVersionMismatchSnafu, RegionNotFoundSnafu, RegionStateSnafu, - RejectWriteInStagingSnafu, RejectWriteSnafu, Result, + InvalidRequestSnafu, PartitionExprVersionMismatchSnafu, RegionLeaseExpiredSnafu, + RegionNotFoundSnafu, RegionStateSnafu, RejectWriteInStagingSnafu, RejectWriteSnafu, Result, + UnacknowledgedWriteSnafu, }; use crate::metrics; use crate::metrics::{ @@ -118,6 +119,8 @@ impl RegionWorkerLoop { return; } } + self.listener.on_write_wal_end().await; + self.fail_writes_after_lease_expiry(&mut region_ctxs); let (mut put_rows, mut delete_rows) = (0, 0); // Write to memtables. @@ -294,8 +297,10 @@ impl RegionWorkerLoop { ) -> HashMap { // Initialize region write context map. let mut region_ctxs = HashMap::new(); - self.process_write_requests(&mut region_ctxs, write_requests); - self.process_bulk_requests(&mut region_ctxs, bulk_requests); + // One reading serves the batch, as preparing it does not wait. + let lease_now = self.node_lease.is_enabled().then(|| self.node_lease.now()); + self.process_write_requests(&mut region_ctxs, write_requests, lease_now); + self.process_bulk_requests(&mut region_ctxs, bulk_requests, lease_now); region_ctxs } @@ -303,6 +308,7 @@ impl RegionWorkerLoop { &mut self, region_ctxs: &mut HashMap, write_requests: &mut Vec, + lease_now: Option, ) { for mut sender_req in write_requests.drain(..) { let region_id = sender_req.request.region_id; @@ -334,6 +340,13 @@ impl RegionWorkerLoop { match region.state() { RegionRoleState::Leader(RegionLeaderState::Writable) | RegionRoleState::Leader(RegionLeaderState::Staging) => { + let leader_deadline = region.leader_deadline.get(); + if lease_now.is_some_and(|now| now >= leader_deadline) { + sender_req + .sender + .send(RegionLeaseExpiredSnafu { region_id }.fail()); + continue; + } if region.reject_all_writes_in_staging() { sender_req .sender @@ -341,13 +354,14 @@ impl RegionWorkerLoop { continue; } - let region_ctx = RegionWriteCtx::new( + let mut region_ctx = RegionWriteCtx::new( region.region_id, ®ion.version_control, region.provider.clone(), Some(region.region_stats.written_bytes.clone()), WriteSource::Request, ); + region_ctx.lease = lease_now.map(|_| (region, leader_deadline)); e.insert(region_ctx); } @@ -460,6 +474,7 @@ impl RegionWorkerLoop { &mut self, region_ctxs: &mut HashMap, requests: &mut Vec, + lease_now: Option, ) { let _timer = metrics::REGION_WORKER_HANDLE_WRITE_ELAPSED .with_label_values(&["prepare_bulk_request"]) @@ -482,19 +497,27 @@ impl RegionWorkerLoop { match region.state() { RegionRoleState::Leader(RegionLeaderState::Writable) | RegionRoleState::Leader(RegionLeaderState::Staging) => { + let leader_deadline = region.leader_deadline.get(); + if lease_now.is_some_and(|now| now >= leader_deadline) { + bulk_req + .sender + .send(RegionLeaseExpiredSnafu { region_id }.fail()); + continue; + } if region.reject_all_writes_in_staging() { bulk_req .sender .send(RejectWriteInStagingSnafu { region_id }.fail()); continue; } - let region_ctx = RegionWriteCtx::new( + let mut region_ctx = RegionWriteCtx::new( region.region_id, ®ion.version_control, region.provider.clone(), Some(region.region_stats.written_bytes.clone()), WriteSource::Request, ); + region_ctx.lease = lease_now.map(|_| (region, leader_deadline)); e.insert(region_ctx); } @@ -573,6 +596,30 @@ impl RegionWorkerLoop { } } + /// Fails the waiters of writes whose WAL append completed at or after their region's Leader + /// deadline. The writes still go to the memtables, as their entries are in the WAL. + fn fail_writes_after_lease_expiry(&self, region_ctxs: &mut HashMap) { + if !self.node_lease.is_enabled() { + return; + } + let now = self.node_lease.now(); + for (region_id, region_ctx) in region_ctxs.iter_mut() { + // A Leader grant may have extended the deadline during the append. + let expired = region_ctx.lease.as_ref().is_some_and(|(region, deadline)| { + now >= *deadline && now >= region.leader_deadline.get() + }); + if !expired || region_ctx.is_failed() { + continue; + } + region_ctx.fail_waiters(Arc::new( + UnacknowledgedWriteSnafu { + region_id: *region_id, + } + .build(), + )); + } + } + /// Returns true if the engine needs to reject some write requests. pub(crate) fn should_reject_write(&self) -> bool { // If memory usage reaches high threshold (we should also consider stalled requests) returns true. diff --git a/src/query/src/optimizer/test_util.rs b/src/query/src/optimizer/test_util.rs index 6542539570c..a6b95ffdeba 100644 --- a/src/query/src/optimizer/test_util.rs +++ b/src/query/src/optimizer/test_util.rs @@ -103,6 +103,15 @@ impl RegionEngine for MetaRegionEngine { unimplemented!() } + fn set_region_role_with_lease( + &self, + _region_id: RegionId, + _role: RegionRole, + _deadline: std::time::Instant, + ) -> Result<(), BoxedError> { + unimplemented!() + } + async fn set_region_role_state_gracefully( &self, _region_id: RegionId, diff --git a/src/store-api/src/region_engine.rs b/src/store-api/src/region_engine.rs index c95bf6d3b21..6dffdd03c3b 100644 --- a/src/store-api/src/region_engine.rs +++ b/src/store-api/src/region_engine.rs @@ -18,6 +18,7 @@ use std::any::Any; use std::collections::HashMap; use std::fmt::{Debug, Display}; use std::sync::{Arc, Mutex}; +use std::time::Instant; use api::greptime_proto::v1::meta::{GrantedRegion as PbGrantedRegion, RegionRole as PbRegionRole}; use api::region::RegionResponse; @@ -1081,6 +1082,18 @@ pub trait RegionEngine: Send + Sync { /// take effect. fn set_region_role(&self, region_id: RegionId, role: RegionRole) -> Result<(), BoxedError>; + /// Sets [RegionRole] for a region from a lease the metasrv grants until `deadline`. + /// + /// The region must not acknowledge a write whose WAL append completes at or after the + /// deadline of its latest Leader or StagingLeader lease. Leases of other roles do not extend + /// that deadline. + fn set_region_role_with_lease( + &self, + region_id: RegionId, + role: RegionRole, + deadline: Instant, + ) -> Result<(), BoxedError>; + /// Syncs the region manifest to the given manifest version. async fn sync_region( &self, From df8f6848cc36fdffd1839c41e52841b9275d77f4 Mon Sep 17 00:00:00 2001 From: discord9 Date: Mon, 26 Oct 2026 09:53:32 +0000 Subject: [PATCH 099/239] fix(query): avoid lossy casts in partition pruning (#9600) * fix(query): keep lossy casts out of partition pruning Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(query): align cast pruning expectations with conservative fallback Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(query): record lossy cast partition pruning regression Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(query): retain safe pruning conjuncts beside unsupported casts Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(query): refresh cast pruning plans after rebase Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- .../src/dist_plan/predicate_extractor.rs | 393 ++++++++++++------ .../optimizer/partition_pruning_cast.result | 175 ++++++++ .../optimizer/partition_pruning_cast.sql | 75 ++++ 3 files changed, 526 insertions(+), 117 deletions(-) create mode 100644 tests/cases/distributed/optimizer/partition_pruning_cast.result create mode 100644 tests/cases/distributed/optimizer/partition_pruning_cast.sql diff --git a/src/query/src/dist_plan/predicate_extractor.rs b/src/query/src/dist_plan/predicate_extractor.rs index 778f5b9fc94..020e6d169ef 100644 --- a/src/query/src/dist_plan/predicate_extractor.rs +++ b/src/query/src/dist_plan/predicate_extractor.rs @@ -18,10 +18,9 @@ use std::collections::HashSet; -use arrow::datatypes::DataType; use common_telemetry::debug; use datafusion_common::Result as DfResult; -use datafusion_expr::{Expr, LogicalPlan, Operator}; +use datafusion_expr::{Expr, LogicalPlan, Operator, utils::split_conjunction}; use datatypes::value::Value; use partition::expr::{Operand, PartitionExpr, RestrictedOp}; @@ -48,33 +47,46 @@ impl PredicateExtractor { let partition_set: HashSet = partition_columns.iter().cloned().collect(); for filter_expr in filter_exprs { - match DataFusionExprConverter::convert(&filter_expr) { - Ok(partition_expr) => { - // Check expression for safe partition pruning - match ExpressionChecker::check_expression_for_pruning( - &partition_expr, - &partition_set, - ) { - ExpressionCheckResult::UseAsIs(expr) => { - partition_exprs.push(expr); - } - ExpressionCheckResult::UsePartial(exprs) => { - partition_exprs.extend(exprs); - } - ExpressionCheckResult::Drop => { - debug!( - "Dropping mixed expression for correctness: {}", - partition_expr - ); - } - } - } + let partition_expr = match DataFusionExprConverter::convert(&filter_expr) { + Ok(expr) => Some(expr), Err(err) => { debug!( - "Failed to convert filter expression to PartitionExpr: {}, skipping", + "Failed to convert filter expression to PartitionExpr: {}, trying conjuncts", err ); - continue; + split_conjunction(&filter_expr) + .into_iter() + .filter_map(|conjunct| match DataFusionExprConverter::convert(conjunct) { + Ok(expr) => Some(expr), + Err(err) => { + debug!( + "Failed to convert filter conjunct to PartitionExpr: {}, skipping", + err + ); + None + } + }) + // Keep salvaged conjuncts in one AND expression, not as alternatives. + .reduce(PartitionExpr::and) + } + }; + let Some(partition_expr) = partition_expr else { + continue; + }; + + // Check expression for safe partition pruning + match ExpressionChecker::check_expression_for_pruning(&partition_expr, &partition_set) { + ExpressionCheckResult::UseAsIs(expr) => { + partition_exprs.push(expr); + } + ExpressionCheckResult::UsePartial(exprs) => { + partition_exprs.extend(exprs); + } + ExpressionCheckResult::Drop => { + debug!( + "Dropping mixed expression for correctness: {}", + partition_expr + ); } } } @@ -423,7 +435,11 @@ impl DataFusionExprConverter { Ok(Operand::Column(column_name)) } Expr::Literal(scalar_value, _) => { - let value = Value::try_from(scalar_value.clone()).unwrap(); + let value = Value::try_from(scalar_value.clone()).map_err(|err| { + datafusion_common::DataFusionError::Plan(format!( + "Failed to convert literal for partition pruning: {err}" + )) + })?; Ok(Operand::Value(value)) } Expr::Alias(alias_expr) => { @@ -431,19 +447,20 @@ impl DataFusionExprConverter { Self::convert_to_operand(&alias_expr.expr) } Expr::Cast(cast_expr) => { - // For safe casts, unwrap to the inner expression - // For unsafe casts, skip with debug logging - if Self::is_safe_cast_for_partition_pruning(cast_expr.field.data_type()) { - Self::convert_to_operand(&cast_expr.expr) + // Column casts need input type/value equivalence proof; literals must use the actual cast value. + if let Expr::Literal(value, _) = cast_expr.expr.as_ref() { + let value = value.cast_to(cast_expr.field.data_type())?; + let value = Value::try_from(value).map_err(|err| { + datafusion_common::DataFusionError::Plan(format!( + "Failed to convert cast literal for partition pruning: {err}" + )) + })?; + Ok(Operand::Value(value)) } else { - debug!( - "Skipping unsafe cast for partition pruning: {:?}", - cast_expr.field.data_type() - ); - Err(datafusion_common::DataFusionError::Plan(format!( - "Cast to {:?} not supported for partition pruning", - cast_expr.field.data_type() - ))) + debug!("Skipping cast expression for partition pruning: {:?}", expr); + Err(datafusion_common::DataFusionError::Plan( + "Cast expression not supported for partition pruning".to_string(), + )) } } other => { @@ -481,46 +498,6 @@ impl DataFusionExprConverter { }; Self::convert_operator(&negated) } - - /// Determine if a cast is safe for partition pruning - /// Safe casts don't change the logical meaning of constraints - fn is_safe_cast_for_partition_pruning(data_type: &DataType) -> bool { - match data_type { - // Integer widening casts are generally safe - DataType::Int8 => true, - DataType::Int16 => true, - DataType::Int32 => true, - DataType::Int64 => true, - DataType::UInt8 => true, - DataType::UInt16 => true, - DataType::UInt32 => true, - DataType::UInt64 => true, - - // Float casts are generally safe for equality/inequality comparisons - DataType::Float32 => true, - DataType::Float64 => true, - - // String casts might be safe in some cases - DataType::Utf8 => true, - DataType::LargeUtf8 => true, - // Dictionary encoding does not change the logical value. DataFusion inserts this cast - // when a predicate literal is compared with a dictionary-encoded scan column. - DataType::Dictionary(_, value_type) => { - Self::is_safe_cast_for_partition_pruning(value_type) - } - - // Date/time casts might be safe if they don't change precision significantly - DataType::Date32 => true, - DataType::Date64 => true, - DataType::Timestamp(_, _) => true, - - // Boolean casts are straightforward - DataType::Boolean => true, - - // For other types, be conservative and skip - _ => false, - } - } } #[cfg(test)] @@ -654,6 +631,190 @@ mod tests { ); } + #[test] + fn test_cast_literals_are_evaluated_and_column_casts_are_skipped() { + let casts = [ + (DataType::Int64, lit(1i64)), + (DataType::Float32, lit(1.0f32)), + (DataType::Utf8, lit("1")), + (DataType::Date32, lit(1i32)), + ( + DataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Nanosecond, None), + lit(1i64), + ), + ]; + for (data_type, value) in casts { + let cast = Expr::Cast(datafusion_expr::Cast::new( + Box::new(col("user_id")), + data_type, + )); + assert!(DataFusionExprConverter::convert(&cast.eq(value)).is_err()); + } + + let chain = Expr::Cast(datafusion_expr::Cast::new( + Box::new(Expr::Cast(datafusion_expr::Cast::new( + Box::new(col("user_id")), + DataType::Float64, + ))), + DataType::Int64, + )); + assert!(DataFusionExprConverter::convert(&chain.eq(lit(1i64))).is_err()); + + let cast_int = Expr::Cast(datafusion_expr::Cast::new( + Box::new(lit(1.5f64)), + DataType::Int64, + )); + assert_eq!( + DataFusionExprConverter::convert(&col("user_id").eq(cast_int)) + .unwrap() + .rhs(), + &Operand::Value(Value::Int64(1)) + ); + + let cast_string = Expr::Cast(datafusion_expr::Cast::new( + Box::new(lit("12")), + DataType::Int64, + )); + assert_eq!( + DataFusionExprConverter::convert(&col("user_id").eq(cast_string)) + .unwrap() + .rhs(), + &Operand::Value(Value::Int64(12)) + ); + + let unsupported = Expr::Cast(datafusion_expr::Cast::new( + Box::new(lit(1i64)), + DataType::Decimal256(10, 2), + )); + assert!(DataFusionExprConverter::convert_to_operand(&unsupported).is_err()); + } + + #[test] + fn test_cast_predicates_are_skipped_in_and_and_or() { + let cast = Expr::Cast(datafusion_expr::Cast::new( + Box::new(col("user_id")), + DataType::Float32, + )); + let lower = PartitionExpr::new( + Operand::Column("user_id".to_string()), + RestrictedOp::GtEq, + Operand::Value(Value::Int64(1)), + ); + let upper = PartitionExpr::new( + Operand::Column("user_id".to_string()), + RestrictedOp::Lt, + Operand::Value(Value::Int64(10)), + ); + let bounds = PartitionExpr::new( + Operand::Expr(lower.clone()), + RestrictedOp::And, + Operand::Expr(upper.clone()), + ); + check_partition_expressions(vec![ + FilterTestCase::new( + "cast_in_and_preserves_partition_bound", + col("user_id") + .eq(lit(1i64)) + .and(cast.clone().eq(lit(1.0f32))), + vec![PartitionExpr::new( + Operand::Column("user_id".to_string()), + RestrictedOp::Eq, + Operand::Value(Value::Int64(1)), + )], + vec!["user_id"], + ), + FilterTestCase::new( + "cast_between_two_bounds_preserves_both", + col("user_id") + .gt_eq(lit(1i64)) + .and(col("user_id").lt(lit(10i64))) + .and(cast.clone().eq(lit(1.0f32))), + vec![bounds.clone()], + vec!["user_id"], + ), + FilterTestCase::new( + "cast_left_of_bounds_preserves_both", + cast.clone() + .eq(lit(1.0f32)) + .and(col("user_id").gt_eq(lit(1i64))) + .and(col("user_id").lt(lit(10i64))), + vec![bounds.clone()], + vec!["user_id"], + ), + FilterTestCase::new( + "cast_between_bounds_preserves_both", + col("user_id") + .gt_eq(lit(1i64)) + .and(cast.clone().eq(lit(1.0f32))) + .and(col("user_id").lt(lit(10i64))), + vec![bounds], + vec!["user_id"], + ), + FilterTestCase::new( + "is_not_null_cast_preserves_partition_bound", + col("user_id").eq(lit(1i64)).and( + Expr::Cast(datafusion_expr::Cast::new( + Box::new(col("value")), + DataType::Int32, + )) + .is_not_null(), + ), + vec![PartitionExpr::new( + Operand::Column("user_id".to_string()), + RestrictedOp::Eq, + Operand::Value(Value::Int64(1)), + )], + vec!["user_id"], + ), + FilterTestCase::new( + "all_unsupported_conjunction_is_dropped", + Expr::Cast(datafusion_expr::Cast::new( + Box::new(col("value")), + DataType::Int32, + )) + .eq(lit(1i32)) + .and( + Expr::Cast(datafusion_expr::Cast::new( + Box::new(col("timestamp")), + DataType::Float32, + )) + .eq(lit(1.0f32)), + ), + vec![], + vec!["user_id"], + ), + FilterTestCase::new( + "nested_and_with_or_side_preserves_bound", + col("user_id").eq(lit(1i64)).and( + col("user_id") + .eq(lit(2i64)) + .or(cast.clone().eq(lit(1.0f32))), + ), + vec![PartitionExpr::new( + Operand::Column("user_id".to_string()), + RestrictedOp::Eq, + Operand::Value(Value::Int64(1)), + )], + vec!["user_id"], + ), + FilterTestCase::new( + "or_with_nested_and_cast_drops_whole_expression", + col("user_id") + .eq(lit(1i64)) + .and(cast.clone().eq(lit(1.0f32))) + .or(col("user_id").eq(lit(2i64))), + vec![], + vec!["user_id"], + ), + FilterTestCase::new( + "cast_in_or_drops_entire_predicate", + col("user_id").eq(lit(1i64)).or(cast.eq(lit(1.0f32))), + vec![], + vec!["user_id"], + ), + ]); + } + #[test] fn test_combines_table_scan_filters_as_conjunction() { let table_scan = create_test_table_scan(); @@ -663,32 +824,42 @@ mod tests { panic!("expected test table scan"); }; let plan = LogicalPlan::TableScan(datafusion_expr::logical_plan::TableScan { - filters: vec![filter_a, filter_b], + filters: vec![filter_a.clone(), filter_b.clone()], + ..scan.clone() + }); + let partition_columns = ["user_id".to_string(), "value".to_string()]; + let expected = vec![PartitionExpr::new( + Operand::Expr(PartitionExpr::new( + Operand::Column("user_id".to_string()), + RestrictedOp::Eq, + Operand::Value(Value::Int64(10)), + )), + RestrictedOp::And, + Operand::Expr(PartitionExpr::new( + Operand::Column("value".to_string()), + RestrictedOp::Eq, + Operand::Value(Value::Int64(20)), + )), + )]; + + let partition_exprs = + PredicateExtractor::extract_partition_expressions(&plan, &partition_columns).unwrap(); + assert_eq!(partition_exprs, expected); + + let unsupported = Expr::Cast(datafusion_expr::Cast::new( + Box::new(col("timestamp")), + DataType::Float32, + )) + .eq(lit(1.0f32)); + let plan = LogicalPlan::TableScan(datafusion_expr::logical_plan::TableScan { + filters: vec![filter_a, unsupported, filter_b], ..scan }); - let partition_exprs = PredicateExtractor::extract_partition_expressions( - &plan, - &["user_id".to_string(), "value".to_string()], - ) - .unwrap(); + let partition_exprs = + PredicateExtractor::extract_partition_expressions(&plan, &partition_columns).unwrap(); - assert_eq!( - partition_exprs, - vec![PartitionExpr::new( - Operand::Expr(PartitionExpr::new( - Operand::Column("user_id".to_string()), - RestrictedOp::Eq, - Operand::Value(Value::Int64(10)), - )), - RestrictedOp::And, - Operand::Expr(PartitionExpr::new( - Operand::Column("value".to_string()), - RestrictedOp::Eq, - Operand::Value(Value::Int64(20)), - )), - )] - ); + assert_eq!(partition_exprs, expected); } #[test] @@ -1010,31 +1181,23 @@ mod tests { fn test_cast_expressions() { let cases = vec![ FilterTestCase::new( - "safe_cast", + "column_cast_requires_type_proof", Expr::Cast(datafusion_expr::Cast::new( Box::new(col("user_id")), DataType::Int64, )) .eq(lit(100i64)), - vec![PartitionExpr::new( - Operand::Column("user_id".to_string()), - RestrictedOp::Eq, - Operand::Value(Value::Int64(100)), - )], + vec![], vec!["user_id"], ), FilterTestCase::new( - "cast_with_alias", + "cast_with_alias_requires_type_proof", Expr::Cast(datafusion_expr::Cast::new( Box::new(col("user_id").alias("uid")), DataType::Int64, )) .eq(lit(100i64)), - vec![PartitionExpr::new( - Operand::Column("user_id".to_string()), - RestrictedOp::Eq, - Operand::Value(Value::Int64(100)), - )], + vec![], vec!["user_id"], ), FilterTestCase::new( @@ -1122,11 +1285,7 @@ mod tests { let in_expr = col("user_id") .alias("uid") .in_list(vec![lit(100i64), lit(200i64)], false); - let cast_expr = Expr::Cast(datafusion_expr::Cast::new( - Box::new(col("user_id")), - DataType::Int64, - )); - let between_expr = cast_expr.between(lit(300i64), lit(400i64)); + let between_expr = col("user_id").between(lit(300i64), lit(400i64)); in_expr.or(between_expr) }, vec![PartitionExpr::new( diff --git a/tests/cases/distributed/optimizer/partition_pruning_cast.result b/tests/cases/distributed/optimizer/partition_pruning_cast.result new file mode 100644 index 00000000000..2b18475fda8 --- /dev/null +++ b/tests/cases/distributed/optimizer/partition_pruning_cast.result @@ -0,0 +1,175 @@ +CREATE TABLE partition_pruning_cast_partitioned ( + ts TIMESTAMP TIME INDEX, + k FLOAT PRIMARY KEY, + v BIGINT, +) +PARTITION ON COLUMNS (k) ( + k < 1.25, + k >= 1.25 +); + +Affected Rows: 0 + +CREATE TABLE partition_pruning_cast_control ( + ts TIMESTAMP TIME INDEX, + k FLOAT PRIMARY KEY, + v BIGINT, +); + +Affected Rows: 0 + +INSERT INTO partition_pruning_cast_partitioned VALUES + (1000, 1.0, 10), + (1000, 1.5, 15), + (1000, 2.5, 25); + +Affected Rows: 3 + +INSERT INTO partition_pruning_cast_control VALUES + (1000, 1.0, 10), + (1000, 1.5, 15), + (1000, 2.5, 25); + +Affected Rows: 3 + +SELECT k, v, CAST(k AS BIGINT) AS cast_k +FROM partition_pruning_cast_partitioned ORDER BY k; + ++-----+----+--------+ +| k | v | cast_k | ++-----+----+--------+ +| 1.0 | 10 | 1 | +| 1.5 | 15 | 1 | +| 2.5 | 25 | 2 | ++-----+----+--------+ + +SELECT k, v, CAST(k AS BIGINT) AS cast_k +FROM partition_pruning_cast_control ORDER BY k; + ++-----+----+--------+ +| k | v | cast_k | ++-----+----+--------+ +| 1.0 | 10 | 1 | +| 1.5 | 15 | 1 | +| 2.5 | 25 | 2 | ++-----+----+--------+ + +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (partitioning.*) REDACTED +EXPLAIN SELECT k, v, CAST(k AS BIGINT) AS cast_k +FROM partition_pruning_cast_partitioned +WHERE CAST(k AS BIGINT) = 1; + ++---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------+ +| plan_type | plan | ++---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------+ +| logical_plan | MergeScan [is_placeholder=false, remote_input=[ | +| | Projection: partition_pruning_cast_partitioned.k, partition_pruning_cast_partitioned.v, CAST(partition_pruning_cast_partitioned.k AS Int64) AS cast_k | +| | Filter: CAST(partition_pruning_cast_partitioned.k AS Int64) = Int64(1) | +| | TableScan: partition_pruning_cast_partitioned, partial_filters=[CAST(partition_pruning_cast_partitioned.k AS Int64) = Int64(1)] | +| | ]] | +| physical_plan | CooperativeExec | +| | MergeScanExec: REDACTED +| | | ++---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------+ + +SELECT ts, k, v, CAST(k AS BIGINT) AS cast_k +FROM partition_pruning_cast_partitioned +WHERE CAST(k AS BIGINT) = 1 +ORDER BY k; + ++---------------------+-----+----+--------+ +| ts | k | v | cast_k | ++---------------------+-----+----+--------+ +| 1970-01-01T00:00:01 | 1.0 | 10 | 1 | +| 1970-01-01T00:00:01 | 1.5 | 15 | 1 | ++---------------------+-----+----+--------+ + +SELECT ts, k, v, CAST(k AS BIGINT) AS cast_k +FROM partition_pruning_cast_control +WHERE CAST(k AS BIGINT) = 1 +ORDER BY k; + ++---------------------+-----+----+--------+ +| ts | k | v | cast_k | ++---------------------+-----+----+--------+ +| 1970-01-01T00:00:01 | 1.0 | 10 | 1 | +| 1970-01-01T00:00:01 | 1.5 | 15 | 1 | ++---------------------+-----+----+--------+ + +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (elapsed_compute.*) REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (metrics.*) REDACTED +-- SQLNESS REPLACE (partitioning.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED +EXPLAIN ANALYZE SELECT k, v +FROM partition_pruning_cast_partitioned +WHERE k >= 1.25::FLOAT AND k < 2.5::FLOAT; + ++-+-+-+ +| stage | node | plan_| ++-+-+-+ +| 0_| 0_|_CooperativeExec REDACTED +|_|_|_MergeScanExec: REDACTED +|_|_|_| +| 1_| 0_|_CooperativeExec REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["k", "v"], "filters": ["k >= Float32(1.25)", "k < Float32(2.5)", "k >= Float32(1.25)"], "time_range":REDACTED +|_|_|_| +|_|_| Total rows: 1_| ++-+-+-+ + +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (elapsed_compute.*) REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (metrics.*) REDACTED +-- SQLNESS REPLACE (partitioning.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED +EXPLAIN ANALYZE SELECT k, v +FROM partition_pruning_cast_partitioned +WHERE k >= 1.25::FLOAT AND k < 2.5::FLOAT AND CAST(v AS INT) IS NOT NULL; + ++-+-+-+ +| stage | node | plan_| ++-+-+-+ +| 0_| 0_|_CooperativeExec REDACTED +|_|_|_MergeScanExec: REDACTED +|_|_|_| +| 1_| 0_|_FilterExec: CAST(v@1 AS Int32) IS NOT NULL REDACTED +|_|_|_CooperativeExec REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["k", "v"], "filters": ["k >= Float32(1.25)", "k < Float32(2.5)", "CAST(v AS Int32) IS NOT NULL", "k >= Float32(1.25)"], "time_range":REDACTED +|_|_|_| +|_|_| Total rows: 1_| ++-+-+-+ + +SELECT k, v FROM partition_pruning_cast_partitioned +WHERE k >= 1.25::FLOAT AND k < 2.5::FLOAT AND CAST(v AS INT) IS NOT NULL; + ++-----+----+ +| k | v | ++-----+----+ +| 1.5 | 15 | ++-----+----+ + +SELECT k FROM partition_pruning_cast_partitioned +WHERE k = 1.0 OR CAST(k AS BIGINT) = 1 ORDER BY k; + ++-----+ +| k | ++-----+ +| 1.0 | +| 1.5 | ++-----+ + +DROP TABLE partition_pruning_cast_partitioned; + +Affected Rows: 0 + +DROP TABLE partition_pruning_cast_control; + +Affected Rows: 0 + diff --git a/tests/cases/distributed/optimizer/partition_pruning_cast.sql b/tests/cases/distributed/optimizer/partition_pruning_cast.sql new file mode 100644 index 00000000000..058045f733a --- /dev/null +++ b/tests/cases/distributed/optimizer/partition_pruning_cast.sql @@ -0,0 +1,75 @@ +CREATE TABLE partition_pruning_cast_partitioned ( + ts TIMESTAMP TIME INDEX, + k FLOAT PRIMARY KEY, + v BIGINT, +) +PARTITION ON COLUMNS (k) ( + k < 1.25, + k >= 1.25 +); + +CREATE TABLE partition_pruning_cast_control ( + ts TIMESTAMP TIME INDEX, + k FLOAT PRIMARY KEY, + v BIGINT, +); + +INSERT INTO partition_pruning_cast_partitioned VALUES + (1000, 1.0, 10), + (1000, 1.5, 15), + (1000, 2.5, 25); +INSERT INTO partition_pruning_cast_control VALUES + (1000, 1.0, 10), + (1000, 1.5, 15), + (1000, 2.5, 25); + +SELECT k, v, CAST(k AS BIGINT) AS cast_k +FROM partition_pruning_cast_partitioned ORDER BY k; +SELECT k, v, CAST(k AS BIGINT) AS cast_k +FROM partition_pruning_cast_control ORDER BY k; + +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (partitioning.*) REDACTED +EXPLAIN SELECT k, v, CAST(k AS BIGINT) AS cast_k +FROM partition_pruning_cast_partitioned +WHERE CAST(k AS BIGINT) = 1; + +SELECT ts, k, v, CAST(k AS BIGINT) AS cast_k +FROM partition_pruning_cast_partitioned +WHERE CAST(k AS BIGINT) = 1 +ORDER BY k; +SELECT ts, k, v, CAST(k AS BIGINT) AS cast_k +FROM partition_pruning_cast_control +WHERE CAST(k AS BIGINT) = 1 +ORDER BY k; + +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (elapsed_compute.*) REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (metrics.*) REDACTED +-- SQLNESS REPLACE (partitioning.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED +EXPLAIN ANALYZE SELECT k, v +FROM partition_pruning_cast_partitioned +WHERE k >= 1.25::FLOAT AND k < 2.5::FLOAT; +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (elapsed_compute.*) REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (metrics.*) REDACTED +-- SQLNESS REPLACE (partitioning.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED +EXPLAIN ANALYZE SELECT k, v +FROM partition_pruning_cast_partitioned +WHERE k >= 1.25::FLOAT AND k < 2.5::FLOAT AND CAST(v AS INT) IS NOT NULL; +SELECT k, v FROM partition_pruning_cast_partitioned +WHERE k >= 1.25::FLOAT AND k < 2.5::FLOAT AND CAST(v AS INT) IS NOT NULL; + +SELECT k FROM partition_pruning_cast_partitioned +WHERE k = 1.0 OR CAST(k AS BIGINT) = 1 ORDER BY k; + +DROP TABLE partition_pruning_cast_partitioned; +DROP TABLE partition_pruning_cast_control; From 0dab3aef21cf2bb8846b278c25bab17ed9da1c17 Mon Sep 17 00:00:00 2001 From: discord9 Date: Mon, 26 Oct 2026 10:25:25 +0000 Subject: [PATCH 100/239] fix(query): preserve IGNORE NULLS semantics in distributed queries (#9601) * fix(query): skip LastRow pushdown for IGNORE NULLS Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(query): keep IGNORE NULLS aggregates above remote scans Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(query): record IGNORE NULLS results before and after flush Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(query): share IGNORE NULLS coverage across deployments Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(query): refresh last-value plans after rebase Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/query/src/dist_plan/analyzer/test.rs | 71 ++++++++- src/query/src/dist_plan/commutativity.rs | 13 +- src/query/src/optimizer/scan_hint.rs | 54 ++++++- .../aggregate/last_value_ignore_nulls.result | 142 ++++++++++++++++++ .../aggregate/last_value_ignore_nulls.sql | 76 ++++++++++ .../standalone/optimizer/last_value.result | 30 ++++ .../cases/standalone/optimizer/last_value.sql | 15 ++ 7 files changed, 395 insertions(+), 6 deletions(-) create mode 100644 tests/cases/standalone/common/aggregate/last_value_ignore_nulls.result create mode 100644 tests/cases/standalone/common/aggregate/last_value_ignore_nulls.sql diff --git a/src/query/src/dist_plan/analyzer/test.rs b/src/query/src/dist_plan/analyzer/test.rs index 207de7d43fd..f644a9d93fb 100644 --- a/src/query/src/dist_plan/analyzer/test.rs +++ b/src/query/src/dist_plan/analyzer/test.rs @@ -32,7 +32,7 @@ use datafusion::functions_nested::expr_fn::make_array; use datafusion::prelude::SessionContext; use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{ExprSchema, JoinType, ScalarValue, TableReference}; -use datafusion_expr::expr::{Exists, ScalarFunction}; +use datafusion_expr::expr::{Exists, NullTreatment, ScalarFunction}; use datafusion_expr::utils::split_conjunction; use datafusion_expr::{ AggregateUDF, Expr, ExprSchemable as _, Extension, LogicalPlanBuilder, Operator, Subquery, @@ -2484,6 +2484,75 @@ fn date_bin_ts_group_by() { assert_eq!(expected, result.to_string()); } +#[test] +fn ignore_nulls_aggregate_stays_on_frontend() { + use datafusion::functions_aggregate::first_last::last_value; + + let table = TestTable::table_with_name(0, "t".to_string()); + let source = Arc::new(DefaultTableSource::new(Arc::new( + DfTableProviderAdapter::new(table), + ))); + let input = LogicalPlanBuilder::scan_with_filters("t", source, None, vec![]) + .unwrap() + .build() + .unwrap(); + for (group_by, null_treatment) in [ + ( + vec![col("pk1"), col("pk2")], + Some(NullTreatment::IgnoreNulls), + ), + (vec![col("number")], Some(NullTreatment::IgnoreNulls)), + (vec![col("pk1"), col("pk2")], None), + ( + vec![col("pk1"), col("pk2")], + Some(NullTreatment::RespectNulls), + ), + ] { + let mut ignored = last_value(col("number"), vec![]); + let Expr::AggregateFunction(func) = &mut ignored else { + panic!("expected aggregate expression"); + }; + func.params.null_treatment = null_treatment; + let aggregate_exprs = if null_treatment == Some(NullTreatment::IgnoreNulls) { + let mut ordinary = last_value(col("number"), vec![]); + let Expr::AggregateFunction(func) = &mut ordinary else { + panic!("expected aggregate expression"); + }; + func.params.null_treatment = None; + vec![ordinary.alias("ordinary"), ignored.alias("ignored")] + } else { + vec![ignored] + }; + let plan = LogicalPlanBuilder::from(input.clone()) + .aggregate(group_by, aggregate_exprs) + .unwrap() + .build() + .unwrap(); + let analyzed = DistPlannerAnalyzer {} + .analyze(plan.clone(), &ConfigOptions::default()) + .unwrap(); + if null_treatment == Some(NullTreatment::IgnoreNulls) { + let LogicalPlan::Aggregate(aggregate) = &analyzed else { + panic!("IGNORE NULLS aggregate must remain on the frontend: {analyzed}"); + }; + let remote_input = find_merge_scan(&analyzed).unwrap().input(); + assert!( + !matches!(remote_input, LogicalPlan::Aggregate(_)), + "remote input contains aggregate: {remote_input}" + ); + let Some(func) = + common_function::aggrs::aggr_wrapper::get_aggr_func(&aggregate.aggr_expr[1]) + else { + unreachable!(); + }; + assert_eq!(Some(NullTreatment::IgnoreNulls), func.params.null_treatment); + } else { + let remote_input = find_merge_scan(&analyzed).unwrap().input(); + assert!(matches!(remote_input, LogicalPlan::Aggregate(_))); + } + } +} + #[test] fn test_last_value_order_by() { init_default_ut_logging(); diff --git a/src/query/src/dist_plan/commutativity.rs b/src/query/src/dist_plan/commutativity.rs index b69589a7cfe..ae7bec52297 100644 --- a/src/query/src/dist_plan/commutativity.rs +++ b/src/query/src/dist_plan/commutativity.rs @@ -15,10 +15,13 @@ use std::collections::HashSet; use std::sync::Arc; -use common_function::aggrs::aggr_wrapper::{StateMergeHelper, is_all_aggr_exprs_steppable}; +use common_function::aggrs::aggr_wrapper::{ + StateMergeHelper, get_aggr_func, is_all_aggr_exprs_steppable, +}; use common_telemetry::debug; use datafusion::error::Result as DfResult; use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; +use datafusion_expr::expr::NullTreatment; use datafusion_expr::{Expr, LogicalPlan, UserDefinedLogicalNode}; use promql::extension_plan::{ EmptyMetric, InstantManipulate, RangeManipulate, SeriesDivide, SeriesNormalize, @@ -116,6 +119,14 @@ impl Categorizer { LogicalPlan::Filter(filter) => Self::check_expr(&filter.predicate), LogicalPlan::Window(_) => Commutativity::Unimplemented, LogicalPlan::Aggregate(aggr) => { + // The pinned Substrait codec drops IGNORE NULLS, so keep these aggregates on the frontend. + if aggr.aggr_expr.iter().any(|expr| { + get_aggr_func(expr).is_some_and(|func| { + func.params.null_treatment == Some(NullTreatment::IgnoreNulls) + }) + }) { + return Ok(Commutativity::NonCommutative); + } // The state/merge split maps each group expression to one output column, // which doesn't hold for grouping sets. let has_grouping_set = aggr diff --git a/src/query/src/optimizer/scan_hint.rs b/src/query/src/optimizer/scan_hint.rs index 4685889c998..bc3c04ec1bf 100644 --- a/src/query/src/optimizer/scan_hint.rs +++ b/src/query/src/optimizer/scan_hint.rs @@ -22,7 +22,7 @@ use common_recordbatch::filter::SimpleFilterEvaluator; use datafusion::datasource::DefaultTableSource; use datafusion_common::tree_node::{Transformed, TreeNodeRewriter}; use datafusion_common::{Column, Result}; -use datafusion_expr::expr::Sort; +use datafusion_expr::expr::{NullTreatment, Sort}; use datafusion_expr::{Expr, LogicalPlan, utils}; use datafusion_optimizer::{OptimizerConfig, OptimizerRule}; use datatypes::arrow::datatypes::{DataType, TimeUnit as ArrowTimeUnit}; @@ -413,6 +413,8 @@ impl ScanHintRewriter { && func.func.name() != aggr_state_func_name("last_value")) || func.params.filter.is_some() || func.params.distinct + // A shared final-row selector would discard earlier nonnull values. + || func.params.null_treatment == Some(NullTreatment::IgnoreNulls) { return None; } @@ -455,7 +457,7 @@ mod test { use datafusion::prelude::JoinType; use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_expr::expr::{ - AggregateFunction, AggregateFunctionParams, Cast, WindowFunctionParams, + AggregateFunction, AggregateFunctionParams, Cast, NullTreatment, WindowFunctionParams, }; use datafusion_expr::expr_fn::scalar_subquery; use datafusion_expr::{Extension, LogicalPlan, LogicalPlanBuilder, col, lit}; @@ -561,6 +563,13 @@ mod test { } fn last_value_aggregate(input: LogicalPlan) -> LogicalPlan { + last_value_aggregate_with_null_treatments(input, [None, None]) + } + + fn last_value_aggregate_with_null_treatments( + input: LogicalPlan, + null_treatments: [Option; 2], + ) -> LogicalPlan { let aggregate = LogicalPlanBuilder::from(input) .aggregate( vec![col("k0")], @@ -576,7 +585,7 @@ mod test { asc: true, nulls_first: true, }], - null_treatment: None, + null_treatment: null_treatments[0], }, }), Expr::AggregateFunction(AggregateFunction { @@ -590,7 +599,7 @@ mod test { asc: true, nulls_first: true, }], - null_treatment: None, + null_treatment: null_treatments[1], }, }), ], @@ -901,6 +910,43 @@ mod test { } } + #[test] + fn last_value_null_treatment_controls_selector() { + let provider = Arc::new(mock_table_provider(RegionId::new(1, 1))); + let context = OptimizerContext::default(); + + for treatment in [None, Some(NullTreatment::RespectNulls)] { + let rewritten = ScanHintRule + .rewrite( + last_value_aggregate_with_null_treatments( + scan_plan(provider.clone(), "t"), + [treatment, treatment], + ), + &context, + ) + .unwrap() + .data; + assert_eq!( + scan_requests(&rewritten)[0].series_row_selector, + Some(TimeSeriesRowSelector::LastRow { after_merge: false }) + ); + } + + // One IGNORE NULLS member prevents the shared scan selector for every + // aggregate in the group, including the member that respects nulls. + let rewritten = ScanHintRule + .rewrite( + last_value_aggregate_with_null_treatments( + scan_plan(provider, "t"), + [None, Some(NullTreatment::IgnoreNulls)], + ), + &context, + ) + .unwrap() + .data; + assert_eq!(scan_requests(&rewritten)[0].series_row_selector, None); + } + #[test] fn single_evaluation_last_value_aggregate_keeps_legacy_selector() { let provider = Arc::new(mock_table_provider(RegionId::new(1, 1))); diff --git a/tests/cases/standalone/common/aggregate/last_value_ignore_nulls.result b/tests/cases/standalone/common/aggregate/last_value_ignore_nulls.result new file mode 100644 index 00000000000..75f70c6a58f --- /dev/null +++ b/tests/cases/standalone/common/aggregate/last_value_ignore_nulls.result @@ -0,0 +1,142 @@ +create table last_value_ignore_nulls ( + ts timestamp time index, + host string primary key, + v1 double, + v2 double, +); + +Affected Rows: 0 + +insert into last_value_ignore_nulls values + (1000, 'a', 42.0, null), + (2000, 'a', null, 20.0), + (1000, 'b', null, null), + (2000, 'b', null, null); + +Affected Rows: 4 + +select + host, + last_value(v1 order by ts) ignore nulls as v1_ignore, + last_value(v2 order by ts) ignore nulls as v2_ignore, + last_value(v1 order by ts) as v1_default, + last_value(v1 order by ts) respect nulls as v1_respect +from last_value_ignore_nulls +group by host +order by host; + ++------+-----------+-----------+------------+------------+ +| host | v1_ignore | v2_ignore | v1_default | v1_respect | ++------+-----------+-----------+------------+------------+ +| a | 42.0 | 20.0 | | | +| b | | | | | ++------+-----------+-----------+------------+------------+ + +-- Default and RESPECT NULLS return the same results in both modes. +select + host, + last_value(v1 order by ts) as v1_default, + last_value(v1 order by ts) respect nulls as v1_respect +from last_value_ignore_nulls +group by host +order by host; + ++------+------------+------------+ +| host | v1_default | v1_respect | ++------+------------+------------+ +| a | | | +| b | | | ++------+------------+------------+ + +-- The IGNORE NULLS member rejects the shared selector, including when it is second. +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (elapsed_compute.*) REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (metrics.*) REDACTED +-- SQLNESS REPLACE (partitioning.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED +explain analyze select + last_value(v1 order by ts), + last_value(v2 order by ts) ignore nulls +from last_value_ignore_nulls +group by host; + ++-+-+-+ +| stage | node | plan_| ++-+-+-+ +| 0_| 0_|_ProjectionExec: expr=[last_value(last_value_ignore_nulls.v1) ORDER BY [last_value_ignore_nulls.ts ASC NULLS LAST]@1 as last_value(last_value_ignore_nulls.v1) ORDER BY [last_value_ignore_nulls.ts ASC NULLS LAST], last_value(last_value_ignore_nulls.v2) IGNORE NULLS ORDER BY [last_value_ignore_nulls.ts ASC NULLS LAST]@2 as last_value(last_value_ignore_nulls.v2) IGNORE NULLS ORDER BY [last_value_ignore_nulls.ts ASC NULLS LAST]] REDACTED +|_|_|_AggregateExec: mode=FinalPartitioned, gby=[host@0 as host], aggr=[last_value(last_value_ignore_nulls.v1) ORDER BY [last_value_ignore_nulls.ts ASC NULLS LAST], last_value(last_value_ignore_nulls.v2) IGNORE NULLS ORDER BY [last_value_ignore_nulls.ts ASC NULLS LAST]] REDACTED +|_|_|_RepartitionExec: REDACTED +|_|_|_AggregateExec: mode=Partial, gby=[host@1 as host], aggr=[last_value(last_value_ignore_nulls.v1) ORDER BY [last_value_ignore_nulls.ts ASC NULLS LAST], last_value(last_value_ignore_nulls.v2) IGNORE NULLS ORDER BY [last_value_ignore_nulls.ts ASC NULLS LAST]] REDACTED +|_|_|_RepartitionExec: REDACTED +|_|_|_MergeScanExec: REDACTED +|_|_|_| +| 1_| 0_|_CooperativeExec REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["ts", "host", "v1", "v2"], "time_range":REDACTED +|_|_|_| +|_|_| Total rows: 2_| ++-+-+-+ + +admin flush_table('last_value_ignore_nulls'); + ++----------------------------------------------+ +| ADMIN flush_table('last_value_ignore_nulls') | ++----------------------------------------------+ +| 0 | ++----------------------------------------------+ + +select + host, + last_value(v1 order by ts) ignore nulls as v1_ignore, + last_value(v2 order by ts) ignore nulls as v2_ignore, + last_value(v1 order by ts) as v1_default, + last_value(v1 order by ts) respect nulls as v1_respect +from last_value_ignore_nulls +group by host +order by host; + ++------+-----------+-----------+------------+------------+ +| host | v1_ignore | v2_ignore | v1_default | v1_respect | ++------+-----------+-----------+------------+------------+ +| a | 42.0 | 20.0 | | | +| b | | | | | ++------+-----------+-----------+------------+------------+ + +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (elapsed_compute.*) REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (metrics.*) REDACTED +-- SQLNESS REPLACE (partitioning.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED +explain analyze select + last_value(v1 order by ts), + last_value(v2 order by ts) ignore nulls +from last_value_ignore_nulls +group by host; + ++-+-+-+ +| stage | node | plan_| ++-+-+-+ +| 0_| 0_|_ProjectionExec: expr=[last_value(last_value_ignore_nulls.v1) ORDER BY [last_value_ignore_nulls.ts ASC NULLS LAST]@1 as last_value(last_value_ignore_nulls.v1) ORDER BY [last_value_ignore_nulls.ts ASC NULLS LAST], last_value(last_value_ignore_nulls.v2) IGNORE NULLS ORDER BY [last_value_ignore_nulls.ts ASC NULLS LAST]@2 as last_value(last_value_ignore_nulls.v2) IGNORE NULLS ORDER BY [last_value_ignore_nulls.ts ASC NULLS LAST]] REDACTED +|_|_|_AggregateExec: mode=FinalPartitioned, gby=[host@0 as host], aggr=[last_value(last_value_ignore_nulls.v1) ORDER BY [last_value_ignore_nulls.ts ASC NULLS LAST], last_value(last_value_ignore_nulls.v2) IGNORE NULLS ORDER BY [last_value_ignore_nulls.ts ASC NULLS LAST]] REDACTED +|_|_|_RepartitionExec: REDACTED +|_|_|_AggregateExec: mode=Partial, gby=[host@1 as host], aggr=[last_value(last_value_ignore_nulls.v1) ORDER BY [last_value_ignore_nulls.ts ASC NULLS LAST], last_value(last_value_ignore_nulls.v2) IGNORE NULLS ORDER BY [last_value_ignore_nulls.ts ASC NULLS LAST]] REDACTED +|_|_|_RepartitionExec: REDACTED +|_|_|_MergeScanExec: REDACTED +|_|_|_| +| 1_| 0_|_CooperativeExec REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["ts", "host", "v1", "v2"], "time_range":REDACTED +|_|_|_| +|_|_| Total rows: 2_| ++-+-+-+ + +drop table last_value_ignore_nulls; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/aggregate/last_value_ignore_nulls.sql b/tests/cases/standalone/common/aggregate/last_value_ignore_nulls.sql new file mode 100644 index 00000000000..659dd4fec39 --- /dev/null +++ b/tests/cases/standalone/common/aggregate/last_value_ignore_nulls.sql @@ -0,0 +1,76 @@ +create table last_value_ignore_nulls ( + ts timestamp time index, + host string primary key, + v1 double, + v2 double, +); + +insert into last_value_ignore_nulls values + (1000, 'a', 42.0, null), + (2000, 'a', null, 20.0), + (1000, 'b', null, null), + (2000, 'b', null, null); + +select + host, + last_value(v1 order by ts) ignore nulls as v1_ignore, + last_value(v2 order by ts) ignore nulls as v2_ignore, + last_value(v1 order by ts) as v1_default, + last_value(v1 order by ts) respect nulls as v1_respect +from last_value_ignore_nulls +group by host +order by host; + +-- Default and RESPECT NULLS return the same results in both modes. +select + host, + last_value(v1 order by ts) as v1_default, + last_value(v1 order by ts) respect nulls as v1_respect +from last_value_ignore_nulls +group by host +order by host; + +-- The IGNORE NULLS member rejects the shared selector, including when it is second. +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (elapsed_compute.*) REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (metrics.*) REDACTED +-- SQLNESS REPLACE (partitioning.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED +explain analyze select + last_value(v1 order by ts), + last_value(v2 order by ts) ignore nulls +from last_value_ignore_nulls +group by host; + +admin flush_table('last_value_ignore_nulls'); + +select + host, + last_value(v1 order by ts) ignore nulls as v1_ignore, + last_value(v2 order by ts) ignore nulls as v2_ignore, + last_value(v1 order by ts) as v1_default, + last_value(v1 order by ts) respect nulls as v1_respect +from last_value_ignore_nulls +group by host +order by host; + +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (elapsed_compute.*) REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (metrics.*) REDACTED +-- SQLNESS REPLACE (partitioning.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED +explain analyze select + last_value(v1 order by ts), + last_value(v2 order by ts) ignore nulls +from last_value_ignore_nulls +group by host; + +drop table last_value_ignore_nulls; diff --git a/tests/cases/standalone/optimizer/last_value.result b/tests/cases/standalone/optimizer/last_value.result index d2765209006..4131a2240d2 100644 --- a/tests/cases/standalone/optimizer/last_value.result +++ b/tests/cases/standalone/optimizer/last_value.result @@ -146,6 +146,36 @@ explain analyze |_|_| Total rows: 1_| +-+-+-+ +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (elapsed_compute.*) REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (metrics.*) REDACTED +-- SQLNESS REPLACE (partitioning.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED +explain analyze select + last_value(val order by ts), + last_value(val order by ts) respect nulls +from t +group by host; + ++-+-+-+ +| stage | node | plan_| ++-+-+-+ +| 0_| 0_|_CooperativeExec REDACTED +|_|_|_MergeScanExec: REDACTED +|_|_|_| +| 1_| 0_|_ProjectionExec: expr=[last_value(t.val) ORDER BY [t.ts ASC NULLS LAST]@1 as last_value(t.val) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) RESPECT NULLS ORDER BY [t.ts ASC NULLS LAST]@2 as last_value(t.val) RESPECT NULLS ORDER BY [t.ts ASC NULLS LAST]] REDACTED +|_|_|_AggregateExec: mode=FinalPartitioned, gby=[host@0 as host], aggr=[last_value(t.val) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) RESPECT NULLS ORDER BY [t.ts ASC NULLS LAST]] REDACTED +|_|_|_RepartitionExec: REDACTED +|_|_|_AggregateExec: mode=Partial, gby=[host@1 as host], aggr=[last_value(t.val) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) RESPECT NULLS ORDER BY [t.ts ASC NULLS LAST]] REDACTED +|_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "selector":"LastRow { after_merge: false }", "projection": ["ts", "host", "val"], "time_range":REDACTED +|_|_|_| +|_|_| Total rows: 4_| ++-+-+-+ + drop table t; Affected Rows: 0 diff --git a/tests/cases/standalone/optimizer/last_value.sql b/tests/cases/standalone/optimizer/last_value.sql index 4b6ecd09ee2..7618276217b 100644 --- a/tests/cases/standalone/optimizer/last_value.sql +++ b/tests/cases/standalone/optimizer/last_value.sql @@ -72,4 +72,19 @@ explain select last_value(ts order by ts) from t; explain analyze select last_value(ts order by ts) from t; +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (elapsed_compute.*) REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (metrics.*) REDACTED +-- SQLNESS REPLACE (partitioning.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \\?"(dyn_filters|time_range|file_summary|max_part_range)\\?":.* "$1":REDACTED +explain analyze select + last_value(val order by ts), + last_value(val order by ts) respect nulls +from t +group by host; + drop table t; \ No newline at end of file From 9f0e58e2af02befdf0c7804c2d9003ebc2b19213 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Mon, 26 Oct 2026 11:34:45 +0000 Subject: [PATCH 101/239] fix(meta): share peer discovery with plugin hooks (#9977) Signed-off-by: WenyXu --- src/meta-srv/src/metasrv.rs | 7 +++- src/meta-srv/src/metasrv/builder.rs | 52 ++++++++++++++++++++++++++--- 2 files changed, 54 insertions(+), 5 deletions(-) diff --git a/src/meta-srv/src/metasrv.rs b/src/meta-srv/src/metasrv.rs index fd63deaa75a..f59e15cd5cf 100644 --- a/src/meta-srv/src/metasrv.rs +++ b/src/meta-srv/src/metasrv.rs @@ -67,6 +67,7 @@ use tokio::sync::broadcast::error::RecvError; use crate::cluster::MetaPeerClientRef; use crate::discovery; +use crate::discovery::node_info::NodeInfoAccessor; use crate::error::{ self, InitMetadataSnafu, KvBackendSnafu, Result, StartProcedureManagerSnafu, StartTelemetryTaskSnafu, StopProcedureManagerSnafu, @@ -515,9 +516,13 @@ pub trait SelectorFactory: Send + Sync { fn build(&self, ctx: SelectorFactoryContext) -> SelectorRef; /// Builds the resolver shared by allocation and failover callers. + /// + /// The accessor includes last-known records of inactive source nodes; discovery + /// is the shared scheduling-aware view of active nodes. fn build_affinity_resolver( &self, - _meta_peer_client: MetaPeerClientRef, + _node_info_accessor: Arc, + _peer_discovery: PeerDiscoveryRef, ) -> Option { None } diff --git a/src/meta-srv/src/metasrv/builder.rs b/src/meta-srv/src/metasrv/builder.rs index 55fdf457044..0a0e1d339e8 100644 --- a/src/meta-srv/src/metasrv/builder.rs +++ b/src/meta-srv/src/metasrv/builder.rs @@ -39,6 +39,7 @@ use common_meta::key::runtime_switch::{RuntimeSwitchManager, RuntimeSwitchManage use common_meta::kv_backend::memory::MemoryKvBackend; use common_meta::kv_backend::{KvBackendRef, ResettableKvBackendRef}; use common_meta::node_manager::NodeManagerRef; +use common_meta::peer::PeerDiscoveryRef; use common_meta::reconciliation::manager::ReconciliationManager; use common_meta::region_keeper::MemoryRegionKeeper; use common_meta::region_registry::LeaderRegionRegistry; @@ -412,7 +413,12 @@ impl MetasrvBuilder { let affinity_resolver = plugins .as_ref() .and_then(|plugins| plugins.get::()) - .and_then(|factory| factory.build_affinity_resolver(meta_peer_client.clone())); + .and_then(|factory| { + factory.build_affinity_resolver( + meta_peer_client.clone(), + selector_ctx.peer_discovery.clone(), + ) + }); let region_supervisor_selector = plugins .as_ref() .and_then(|plugins| plugins.get::()); @@ -497,7 +503,7 @@ impl MetasrvBuilder { { let ctx = DdlManagerConfigureContext { kv_backend: kv_backend.clone(), - meta_peer_client: meta_peer_client.clone(), + peer_discovery: selector_ctx.peer_discovery.clone(), }; configurator .configure(ddl_manager, ctx) @@ -759,7 +765,7 @@ impl Default for MetasrvBuilder { /// The context for [`DdlManagerConfiguratorRef`]. pub struct DdlManagerConfigureContext { pub kv_backend: KvBackendRef, - pub meta_peer_client: MetaPeerClientRef, + pub peer_discovery: PeerDiscoveryRef, } #[cfg(test)] @@ -854,6 +860,9 @@ mod tests { use common_meta::cluster::NodeInfo; use common_meta::error::Result as MetaResult; + use common_meta::ddl_manager::DdlManagerConfigurator; + + use crate::discovery::node_info::NodeInfoAccessor; use crate::metasrv::{SelectorFactory, SelectorFactoryContext}; struct TestAffinity; @@ -879,10 +888,32 @@ mod tests { struct RecordingFactory { calls: Arc, resolver: AffinityResolverRef, + discovery: Arc>>, + } + + struct RecordingConfigurator { + discovery: Arc>>, + } + + #[async_trait::async_trait] + impl DdlManagerConfigurator for RecordingConfigurator { + async fn configure( + &self, + ddl_manager: DdlManager, + ctx: DdlManagerConfigureContext, + ) -> std::result::Result { + *self.discovery.lock().unwrap() = Some(ctx.peer_discovery); + Ok(ddl_manager) + } } impl SelectorFactory for RecordingFactory { - fn build_affinity_resolver(&self, _: MetaPeerClientRef) -> Option { + fn build_affinity_resolver( + &self, + _: Arc, + peer_discovery: PeerDiscoveryRef, + ) -> Option { + *self.discovery.lock().unwrap() = Some(peer_discovery); self.calls.fetch_add(1, Ordering::Relaxed); Some(self.resolver.clone()) } @@ -892,6 +923,8 @@ mod tests { } } + let factory_discovery = Arc::new(std::sync::Mutex::new(None)); + let configurator_discovery = Arc::new(std::sync::Mutex::new(None)); let calls = Arc::new(AtomicUsize::new(0)); let resolver: AffinityResolverRef = Arc::new(TestResolver); let affinity: AffinityRef = Arc::new(TestAffinity); @@ -901,7 +934,12 @@ mod tests { plugins.insert(Arc::new(RecordingFactory { calls: calls.clone(), resolver: resolver.clone(), + discovery: factory_discovery.clone(), }) as SelectorFactoryRef); + plugins.insert(Arc::new(RecordingConfigurator { + discovery: configurator_discovery.clone(), + }) + as DdlManagerConfiguratorRef); let metasrv = MetasrvBuilder::new() .options(MetasrvOptions { enable_telemetry: false, @@ -921,6 +959,12 @@ mod tests { AffinityFallback::Unconstrained ); assert_eq!(calls.load(Ordering::Relaxed), 1); + for discovery in [factory_discovery, configurator_discovery] { + assert!(Arc::ptr_eq( + discovery.lock().unwrap().as_ref().unwrap(), + &metasrv.selector_ctx().peer_discovery, + )); + } assert!(Arc::ptr_eq(&metasrv.selector_ctx().affinity, &affinity)); assert!(Arc::ptr_eq( metasrv From 2d00ae25c82e363eaba62f63ec1a21e85a348337 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Mon, 26 Oct 2026 11:52:41 +0000 Subject: [PATCH 102/239] feat(meta): fence a draining datanode and plan the drain (#9960) Introduce the metasrv side procedure that fences a datanode before it is drained, and collect the regions placed on it: - mark the datanode unschedulable through the node scheduling manager - build the drain plan from the datanode table index and the authoritative table routes, splitting leader regions from follower regions - keep the plan in the volatile context: it is derived from the table routes, so a restarted procedure rebuilds it instead of reading a persisted copy The leader regions are not migrated yet: the migration state machine lands separately. The procedure has no caller yet, so it stays inert. Signed-off-by: WenyXu --- src/meta-srv/src/error.rs | 8 + src/meta-srv/src/metasrv.rs | 7 + src/meta-srv/src/metasrv/builder.rs | 9 + src/meta-srv/src/procedure.rs | 1 + src/meta-srv/src/procedure/drain_datanode.rs | 652 ++++++++++++++++++ .../drain_datanode/drain_datanode_end.rs | 41 ++ .../drain_datanode/drain_datanode_start.rs | 61 ++ .../src/procedure/drain_datanode/manager.rs | 187 +++++ .../src/procedure/drain_datanode/utils.rs | 83 +++ 9 files changed, 1049 insertions(+) create mode 100644 src/meta-srv/src/procedure/drain_datanode.rs create mode 100644 src/meta-srv/src/procedure/drain_datanode/drain_datanode_end.rs create mode 100644 src/meta-srv/src/procedure/drain_datanode/drain_datanode_start.rs create mode 100644 src/meta-srv/src/procedure/drain_datanode/manager.rs create mode 100644 src/meta-srv/src/procedure/drain_datanode/utils.rs diff --git a/src/meta-srv/src/error.rs b/src/meta-srv/src/error.rs index 011576f9b27..8ceb39ed56e 100644 --- a/src/meta-srv/src/error.rs +++ b/src/meta-srv/src/error.rs @@ -91,6 +91,13 @@ pub enum Error { region_id: RegionId, }, + #[snafu(display("Another drain procedure is running for datanode: {}", node_id))] + DrainDatanodeRunning { + #[snafu(implicit)] + location: Location, + node_id: DatanodeId, + }, + #[snafu(display( "The region migration procedure is completed for region: {}, target_peer: {}", region_id, @@ -1254,6 +1261,7 @@ impl ErrorExt for Error { | Error::RegionRouteNotFound { .. } | Error::MigrationAbort { .. } | Error::MigrationRunning { .. } + | Error::DrainDatanodeRunning { .. } | Error::RegionMigrated { .. } => StatusCode::Unexpected, Error::TableNotFound { .. } => StatusCode::TableNotFound, Error::SaveClusterInfo { source, .. } diff --git a/src/meta-srv/src/metasrv.rs b/src/meta-srv/src/metasrv.rs index f59e15cd5cf..37d8d3ebb2c 100644 --- a/src/meta-srv/src/metasrv.rs +++ b/src/meta-srv/src/metasrv.rs @@ -77,6 +77,7 @@ use crate::gc::{GcSchedulerOptions, GcTickerRef}; use crate::handler::{HeartbeatHandlerGroupBuilder, HeartbeatHandlerGroupRef}; use crate::node_scheduling::NodeSchedulingManagerRef; use crate::procedure::ProcedureManagerListenerAdapter; +use crate::procedure::drain_datanode::manager::DrainDatanodeManagerRef; use crate::procedure::region_migration::manager::RegionMigrationManagerRef; use crate::procedure::repartition::gc_requirement::RepartitionGcRequirementManagerRef; use crate::procedure::wal_prune::manager::WalPruneTickerRef; @@ -606,6 +607,7 @@ pub struct Metasrv { memory_region_keeper: MemoryRegionKeeperRef, greptimedb_telemetry_task: Arc, region_migration_manager: RegionMigrationManagerRef, + drain_datanode_manager: DrainDatanodeManagerRef, region_supervisor_ticker: Option, cache_invalidator: CacheInvalidatorRef, leader_region_registry: LeaderRegionRegistryRef, @@ -916,6 +918,11 @@ impl Metasrv { &self.region_migration_manager } + /// Returns the [`DrainDatanodeManagerRef`]. + pub fn drain_datanode_manager(&self) -> &DrainDatanodeManagerRef { + &self.drain_datanode_manager + } + pub fn publish(&self) -> Option { self.plugins.get::() } diff --git a/src/meta-srv/src/metasrv/builder.rs b/src/meta-srv/src/metasrv/builder.rs index 0a0e1d339e8..bc10a27af85 100644 --- a/src/meta-srv/src/metasrv/builder.rs +++ b/src/meta-srv/src/metasrv/builder.rs @@ -74,6 +74,7 @@ use crate::metasrv::{ }; use crate::node_scheduling::NodeSchedulingManager; use crate::peer::MetasrvPeerAllocator; +use crate::procedure::drain_datanode::manager::DrainDatanodeManager; use crate::procedure::region_migration::DefaultContextFactory; use crate::procedure::region_migration::manager::RegionMigrationManager; use crate::procedure::repartition::gc_requirement::RepartitionGcRequirementManager; @@ -410,6 +411,13 @@ impl MetasrvBuilder { ), )); region_migration_manager.try_start()?; + + let drain_datanode_manager = Arc::new(DrainDatanodeManager::new( + procedure_manager.clone(), + table_metadata_manager.clone(), + node_scheduling_manager.clone(), + )); + drain_datanode_manager.try_start()?; let affinity_resolver = plugins .as_ref() .and_then(|plugins| plugins.get::()) @@ -688,6 +696,7 @@ impl MetasrvBuilder { plugins: plugins.unwrap_or_else(Plugins::default), memory_region_keeper, region_migration_manager, + drain_datanode_manager, region_supervisor_ticker, cache_invalidator, leader_region_registry, diff --git a/src/meta-srv/src/procedure.rs b/src/meta-srv/src/procedure.rs index da1a1b00e7a..231671eb683 100644 --- a/src/meta-srv/src/procedure.rs +++ b/src/meta-srv/src/procedure.rs @@ -18,6 +18,7 @@ use common_meta::leadership_notifier::LeadershipChangeListener; use common_procedure::ProcedureManagerRef; use snafu::ResultExt; +pub mod drain_datanode; pub mod region_migration; pub mod repartition; #[cfg(any(test, feature = "testing"))] diff --git a/src/meta-srv/src/procedure/drain_datanode.rs b/src/meta-srv/src/procedure/drain_datanode.rs new file mode 100644 index 00000000000..41bdb85c217 --- /dev/null +++ b/src/meta-srv/src/procedure/drain_datanode.rs @@ -0,0 +1,652 @@ +// 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. + +pub(crate) mod drain_datanode_end; +pub(crate) mod drain_datanode_start; +pub(crate) mod manager; +pub(crate) mod utils; + +use std::any::Any; +use std::fmt::Debug; + +use common_meta::DatanodeId; +use common_meta::key::TableMetadataManagerRef; +use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; +use common_procedure::{ + Context as ProcedureContext, Error as ProcedureError, LockKey, Procedure, + Result as ProcedureResult, Status, +}; +use common_telemetry::{error, info}; +use serde::{Deserialize, Serialize}; +use snafu::ResultExt; +use store_api::storage::RegionId; + +use crate::error::Result; +use crate::node_scheduling::NodeSchedulingManagerRef; +use crate::procedure::drain_datanode::drain_datanode_start::DrainDatanodeStart; +use crate::procedure::drain_datanode::manager::{ + DrainDatanodeProcedureGuard, DrainDatanodeProcedureTracker, +}; + +/// The regions placed on a datanode, split by role. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct DrainDatanodePlan { + /// The leader regions still to migrate, sorted and deduplicated. + pub leader_regions: Vec, + /// The follower regions of the datanode, sorted and deduplicated. + pub follower_regions: Vec, +} + +/// The persistent data of [`DrainDatanodeProcedure`]. +/// +/// Only what cannot be derived from the cluster metadata is persisted: the plan is rebuilt +/// from the authoritative table routes after a restart. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct PersistentContext { + /// The datanode to drain. + pub node_id: DatanodeId, + /// The pinned destination of leader regions. + pub target_node_id: Option, +} + +/// The in-memory data of [`DrainDatanodeProcedure`]. +/// +/// It is not persisted: a restarted procedure rebuilds the plan by rescanning the datanode. +#[derive(Debug, Clone, Default)] +pub struct VolatileContext { + /// The drain plan, `None` until the first state builds it. + pub plan: Option, +} + +/// The task of [`DrainDatanodeProcedure`]. +#[derive(Debug, Clone)] +pub struct DrainDatanodeProcedureTask { + /// The datanode to drain. + pub node_id: DatanodeId, + /// The pinned destination of leader regions. + pub target_node_id: Option, +} + +/// The context of [`DrainDatanodeProcedure`]. +#[derive(Clone)] +pub struct Context { + /// The persistent data of the procedure. + pub persistent_ctx: PersistentContext, + /// The in-memory data of the procedure. + pub volatile_ctx: VolatileContext, + /// The table metadata manager. + pub table_metadata_manager: TableMetadataManagerRef, + /// The node scheduling manager, used to fence the drained datanode. + pub node_scheduling_manager: NodeSchedulingManagerRef, +} + +impl Context { + /// Builds the context of a run. + /// + /// The per-run state starts empty: the plan is rebuilt from the table routes when the + /// procedure needs it. + pub(crate) fn new( + table_metadata_manager: TableMetadataManagerRef, + node_scheduling_manager: NodeSchedulingManagerRef, + persistent_ctx: PersistentContext, + ) -> Self { + Self { + persistent_ctx, + volatile_ctx: VolatileContext::default(), + table_metadata_manager, + node_scheduling_manager, + } + } +} + +impl Context { + /// Builds the drain plan when it is not available yet. + /// + /// The plan is derived from the authoritative table routes, so a restarted procedure + /// rebuilds it instead of reading a persisted copy. + pub(crate) async fn ensure_plan(&mut self) -> Result<()> { + if self.volatile_ctx.plan.is_some() { + return Ok(()); + } + + let node_id = self.persistent_ctx.node_id; + let plan = utils::build_datanode_plan(&self.table_metadata_manager, node_id).await?; + self.volatile_ctx.plan = Some(plan); + + Ok(()) + } +} + +impl VolatileContext { + /// Returns the leader regions still to migrate. + pub(crate) fn leader_regions(&self) -> &[RegionId] { + self.plan + .as_ref() + .map_or(&[], |plan| plan.leader_regions.as_slice()) + } + + /// Returns the follower regions collected by the plan. + pub(crate) fn follower_regions(&self) -> &[RegionId] { + self.plan + .as_ref() + .map_or(&[], |plan| plan.follower_regions.as_slice()) + } +} + +/// The persistent data of [`DrainDatanodeProcedure`] as read by `from_json`. +#[derive(Debug, Deserialize)] +pub struct DrainDatanodeDataOwned { + persistent_ctx: PersistentContext, + state: Box, +} + +/// The persistent data of [`DrainDatanodeProcedure`] as written by `dump`. +#[derive(Debug, Serialize)] +pub struct DrainDatanodeData<'a> { + persistent_ctx: &'a PersistentContext, + state: &'a dyn State, +} + +/// A state of [`DrainDatanodeProcedure`]. +#[async_trait::async_trait] +#[typetag::serde(tag = "drain_datanode_state")] +pub(crate) trait State: Sync + Send + Debug { + /// Returns the short name of this state. + fn name(&self) -> &'static str { + let type_name = std::any::type_name::(); + type_name.split("::").last().unwrap_or(type_name) + } + + /// Yields the next [`State`] and [`Status`]. + async fn next( + &mut self, + ctx: &mut Context, + procedure_ctx: &ProcedureContext, + ) -> Result<(Box, Status)>; + + /// Returns as [`Any`]. + fn as_any(&self) -> &dyn Any; +} + +/// The procedure draining a datanode. +/// +/// It fences the datanode and collects the leader and follower regions placed on it. +/// Migrating the leader regions off the datanode lands in a follow-up. +pub(crate) struct DrainDatanodeProcedure { + state: Box, + context: Context, + /// Keeps the drained datanode in the running set of the manager until the procedure ends, + /// so a second drain of the datanode is rejected at submission instead of being queued. + _guards: Vec, +} + +impl DrainDatanodeProcedure { + const TYPE_NAME: &'static str = "metasrv-procedure::DrainDatanode"; + + /// Returns a new [`DrainDatanodeProcedure`] in its initial state. + pub fn new(context: Context, guards: Vec) -> Self { + Self { + state: Box::new(DrainDatanodeStart), + context, + _guards: guards, + } + } + + /// Returns a new [`DrainDatanodeProcedure`] from the given serialized data. + /// + /// The procedure takes its place back in `tracker`, so a reloaded drain still rejects a + /// second drain of the same datanode. + pub fn from_json( + json: &str, + tracker: &DrainDatanodeProcedureTracker, + ctx_factory: F, + ) -> ProcedureResult + where + F: FnOnce(PersistentContext) -> Context, + { + let DrainDatanodeDataOwned { + persistent_ctx, + state, + } = serde_json::from_str(json).context(FromJsonSnafu)?; + let guards = tracker + .insert_running_procedure(&DrainDatanodeProcedureTask { + node_id: persistent_ctx.node_id, + target_node_id: persistent_ctx.target_node_id, + }) + .into_iter() + .collect(); + + Ok(Self { + state, + context: ctx_factory(persistent_ctx), + _guards: guards, + }) + } +} + +#[async_trait::async_trait] +impl Procedure for DrainDatanodeProcedure { + fn type_name(&self) -> &str { + Self::TYPE_NAME + } + + #[tracing::instrument(skip_all, fields( + state = %self.state.name(), + node_id = %self.context.persistent_ctx.node_id + ))] + async fn execute(&mut self, ctx: &ProcedureContext) -> ProcedureResult { + let state_name = self.state.name(); + info!( + "Drain datanode procedure executing state: {}, node_id: {}", + state_name, self.context.persistent_ctx.node_id + ); + match self.state.next(&mut self.context, ctx).await { + Ok((next, status)) => { + self.state = next; + Ok(status) + } + Err(e) => { + if e.is_retryable() { + Err(ProcedureError::retry_later(e)) + } else { + error!( + e; + "Drain datanode procedure failed, node_id: {}", + self.context.persistent_ctx.node_id + ); + Err(ProcedureError::external(e)) + } + } + } + } + + fn dump(&self) -> ProcedureResult { + let data = DrainDatanodeData { + persistent_ctx: &self.context.persistent_ctx, + state: self.state.as_ref(), + }; + serde_json::to_string(&data).context(ToJsonSnafu) + } + + fn lock_key(&self) -> LockKey { + LockKey::single_exclusive(format!( + "__drain_datanode/{}", + self.context.persistent_ctx.node_id + )) + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use common_meta::key::TableMetadataManagerRef; + use common_meta::key::table_route::TableRouteValue; + use common_meta::key::test_utils::new_test_table_info; + use common_meta::peer::Peer; + use common_meta::rpc::router::{Region, RegionRoute}; + + use super::*; + use crate::error; + use crate::key::NodeSchedulingRole; + use crate::procedure::drain_datanode::drain_datanode_end::DrainDatanodeEnd; + use crate::procedure::drain_datanode::manager::{ + DrainDatanodeManager, DrainDatanodeProcedureTracker, + }; + use crate::procedure::region_migration::test_util::TestingEnv; + use crate::test_util::{create_meta_peer_client, new_node_scheduling_discovery, put_datanodes}; + + fn new_persistent_context(target_node_id: Option) -> PersistentContext { + PersistentContext { + node_id: 1, + target_node_id, + } + } + + fn leader_region(region_id: RegionId, leader_id: DatanodeId) -> RegionRoute { + RegionRoute { + region: Region::new_test(region_id), + leader_peer: Some(Peer::empty(leader_id)), + follower_peers: vec![], + ..Default::default() + } + } + + fn follower_region( + region_id: RegionId, + leader_id: DatanodeId, + follower_id: DatanodeId, + ) -> RegionRoute { + RegionRoute { + region: Region::new_test(region_id), + leader_peer: Some(Peer::empty(leader_id)), + follower_peers: vec![Peer::empty(follower_id)], + ..Default::default() + } + } + + struct TestEnv { + env: TestingEnv, + table_metadata_manager: TableMetadataManagerRef, + node_scheduling_manager: NodeSchedulingManagerRef, + } + + impl TestEnv { + async fn new() -> Self { + let env = TestingEnv::new(); + // Shares one backend between the scheduling manager and the peer discovery, so + // the selector observes the scheduling marks written by the procedure. + let meta_peer_client = create_meta_peer_client(); + put_datanodes( + &meta_peer_client, + vec![ + Peer { + id: 1, + addr: "dn1:3001".to_string(), + }, + Peer { + id: 2, + addr: "dn2:3001".to_string(), + }, + ], + ) + .await; + let (node_scheduling_manager, _discovery) = + new_node_scheduling_discovery(&meta_peer_client); + let table_metadata_manager = env.table_metadata_manager().clone(); + + Self { + env, + table_metadata_manager, + node_scheduling_manager, + } + } + + fn context(&self, persistent_ctx: PersistentContext) -> Context { + Context::new( + self.table_metadata_manager.clone(), + self.node_scheduling_manager.clone(), + persistent_ctx, + ) + } + + fn table_metadata_manager(&self) -> &TableMetadataManagerRef { + &self.table_metadata_manager + } + + async fn create_table(&self, region_routes: Vec) { + create_table(&self.table_metadata_manager, region_routes).await; + } + } + + async fn create_table( + table_metadata_manager: &TableMetadataManagerRef, + region_routes: Vec, + ) { + table_metadata_manager + .create_table_metadata( + new_test_table_info(1024), + TableRouteValue::physical(region_routes), + HashMap::default(), + ) + .await + .unwrap(); + } + + async fn run_until_done( + state: &mut Box, + ctx: &mut Context, + procedure_ctx: &ProcedureContext, + ) -> Result { + for _ in 0..10 { + let (next, status) = state.next(ctx, procedure_ctx).await?; + *state = next; + if !matches!(status, Status::Executing { .. }) { + return Ok(status); + } + } + + panic!("the drain procedure did not finish in 10 steps"); + } + + #[tokio::test] + async fn test_build_datanode_plan() { + let env = TestEnv::new().await; + let first = RegionId::new(1024, 1); + let second = RegionId::new(1024, 2); + env.create_table(vec![leader_region(first, 1), follower_region(second, 2, 1)]) + .await; + + let plan = utils::build_datanode_plan(env.table_metadata_manager(), 1) + .await + .unwrap(); + + assert_eq!(vec![first], plan.leader_regions); + assert_eq!(vec![second], plan.follower_regions); + } + + #[tokio::test] + async fn test_build_datanode_plan_of_empty_datanode() { + let env = TestEnv::new().await; + + let plan = utils::build_datanode_plan(env.table_metadata_manager(), 1) + .await + .unwrap(); + + assert!(plan.leader_regions.is_empty()); + assert!(plan.follower_regions.is_empty()); + } + + #[tokio::test] + async fn test_build_datanode_plan_reads_the_table_route() { + let env = TestEnv::new().await; + let region_id = RegionId::new(1024, 1); + // The datanode is a follower of the region, so it is not a leader even though the + // datanode table index of the datanode contains the table. + env.create_table(vec![follower_region(region_id, 2, 1)]) + .await; + + let plan = utils::build_datanode_plan(env.table_metadata_manager(), 1) + .await + .unwrap(); + + assert!(plan.leader_regions.is_empty()); + assert_eq!(vec![region_id], plan.follower_regions); + } + + #[tokio::test] + async fn test_drain_datanode_start_transitions_to_end() { + let env = TestEnv::new().await; + env.create_table(vec![leader_region(RegionId::new(1024, 1), 1)]) + .await; + + let mut ctx = env.context(new_persistent_context(Some(2))); + let procedure_ctx = TestingEnv::procedure_context(); + + let (next, status) = DrainDatanodeStart + .next(&mut ctx, &procedure_ctx) + .await + .unwrap(); + + assert!(next.as_any().downcast_ref::().is_some()); + // Registering the procedure is the only persisted step of the run. + assert!(matches!(status, Status::Executing { persist: true, .. })); + assert_eq!( + vec![RegionId::new(1024, 1)], + ctx.volatile_ctx.leader_regions() + ); + assert!( + !ctx.node_scheduling_manager + .is_schedulable(NodeSchedulingRole::Datanode, 1) + .await + .unwrap() + ); + } + + #[tokio::test] + async fn test_drain_datanode_end_transitions_to_done() { + let env = TestEnv::new().await; + let mut ctx = env.context(new_persistent_context(Some(2))); + let procedure_ctx = TestingEnv::procedure_context(); + + let (next, status) = DrainDatanodeEnd + .next(&mut ctx, &procedure_ctx) + .await + .unwrap(); + + assert!(next.as_any().downcast_ref::().is_some()); + assert!(matches!(status, Status::Done { output: None })); + } + + #[tokio::test] + async fn test_drain_datanode_fences_and_collects_the_plan() { + let env = TestEnv::new().await; + let region_id = RegionId::new(1024, 1); + let follower_region_id = RegionId::new(1024, 2); + env.create_table(vec![ + follower_region(follower_region_id, 2, 1), + leader_region(region_id, 1), + ]) + .await; + + let mut ctx = env.context(new_persistent_context(Some(2))); + let procedure_ctx = TestingEnv::procedure_context(); + let mut state: Box = Box::new(DrainDatanodeStart); + + let status = run_until_done(&mut state, &mut ctx, &procedure_ctx) + .await + .unwrap(); + + assert!(matches!(status, Status::Done { .. })); + assert_eq!(vec![region_id], ctx.volatile_ctx.leader_regions()); + assert_eq!( + vec![follower_region_id], + ctx.volatile_ctx.follower_regions() + ); + assert!( + !ctx.node_scheduling_manager + .is_schedulable(NodeSchedulingRole::Datanode, 1) + .await + .unwrap() + ); + } + + #[tokio::test] + async fn test_drain_datanode_dump_roundtrip() { + let env = TestEnv::new().await; + let persistent_ctx = new_persistent_context(Some(2)); + + let procedure = DrainDatanodeProcedure::new(env.context(persistent_ctx), vec![]); + let json = procedure.dump().unwrap(); + let tracker = DrainDatanodeProcedureTracker::default(); + let reloaded = DrainDatanodeProcedure::from_json(&json, &tracker, |persistent_ctx| { + env.context(persistent_ctx) + }) + .unwrap(); + + assert_eq!("DrainDatanodeStart", reloaded.state.name()); + let ctx = &reloaded.context.persistent_ctx; + assert_eq!(1, ctx.node_id); + assert_eq!(Some(2), ctx.target_node_id); + // The plan is volatile, so it is not dumped. + assert!(!json.contains("leader_regions")); + } + + #[tokio::test] + async fn test_drain_datanode_end_finishes() { + let env = TestEnv::new().await; + let mut ctx = env.context(new_persistent_context(Some(2))); + let procedure_ctx = TestingEnv::procedure_context(); + + let (next, status) = DrainDatanodeEnd + .next(&mut ctx, &procedure_ctx) + .await + .unwrap(); + + assert!(next.as_any().downcast_ref::().is_some()); + assert!(matches!(status, Status::Done { .. })); + } + + #[tokio::test] + async fn test_submit_procedure_rejects_same_target() { + let env = TestEnv::new().await; + let manager = DrainDatanodeManager::new( + env.env.procedure_manager().clone(), + env.table_metadata_manager.clone(), + env.node_scheduling_manager.clone(), + ); + + let error = manager + .submit_procedure( + common_procedure::ProcedureContext::default(), + DrainDatanodeProcedureTask { + node_id: 1, + target_node_id: Some(1), + }, + ) + .await + .unwrap_err(); + + assert!(matches!(error, error::Error::InvalidArguments { .. })); + } + + #[tokio::test] + async fn test_submit_procedure_rejects_a_running_drain() { + let env = TestEnv::new().await; + let manager = DrainDatanodeManager::new( + env.env.procedure_manager().clone(), + env.table_metadata_manager.clone(), + env.node_scheduling_manager.clone(), + ); + let task = DrainDatanodeProcedureTask { + node_id: 1, + target_node_id: Some(2), + }; + + // A drain of datanode 1 is running, for example one recovered after a restart. + let _guard = manager.tracker().insert_running_procedure(&task).unwrap(); + + let error = manager + .submit_procedure(common_procedure::ProcedureContext::default(), task) + .await + .unwrap_err(); + + assert!(matches!(error, error::Error::DrainDatanodeRunning { .. })); + assert!(!error.is_retryable()); + } + + #[tokio::test] + async fn test_drain_datanode_tracker_takes_the_place_back_on_reload() { + let env = TestEnv::new().await; + let procedure = + DrainDatanodeProcedure::new(env.context(new_persistent_context(Some(2))), vec![]); + let json = procedure.dump().unwrap(); + let tracker = DrainDatanodeProcedureTracker::default(); + + // Reloading takes the place back, so a second drain of the datanode is rejected. + let reloaded = DrainDatanodeProcedure::from_json(&json, &tracker, |persistent_ctx| { + env.context(persistent_ctx) + }) + .unwrap(); + let task = DrainDatanodeProcedureTask { + node_id: 1, + target_node_id: Some(2), + }; + assert!(tracker.insert_running_procedure(&task).is_none()); + + // Dropping the reloaded procedure releases the place. + drop(reloaded); + assert!(tracker.insert_running_procedure(&task).is_some()); + } +} diff --git a/src/meta-srv/src/procedure/drain_datanode/drain_datanode_end.rs b/src/meta-srv/src/procedure/drain_datanode/drain_datanode_end.rs new file mode 100644 index 00000000000..1ff5e3041e1 --- /dev/null +++ b/src/meta-srv/src/procedure/drain_datanode/drain_datanode_end.rs @@ -0,0 +1,41 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::any::Any; + +use common_procedure::{Context as ProcedureContext, Status}; +use serde::{Deserialize, Serialize}; + +use crate::error::Result; +use crate::procedure::drain_datanode::{Context, State}; + +/// Finishes the procedure. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub(crate) struct DrainDatanodeEnd; + +#[async_trait::async_trait] +#[typetag::serde] +impl State for DrainDatanodeEnd { + async fn next( + &mut self, + _ctx: &mut Context, + _procedure_ctx: &ProcedureContext, + ) -> Result<(Box, Status)> { + Ok((Box::new(DrainDatanodeEnd), Status::done())) + } + + fn as_any(&self) -> &dyn Any { + self + } +} diff --git a/src/meta-srv/src/procedure/drain_datanode/drain_datanode_start.rs b/src/meta-srv/src/procedure/drain_datanode/drain_datanode_start.rs new file mode 100644 index 00000000000..0e657ebbd00 --- /dev/null +++ b/src/meta-srv/src/procedure/drain_datanode/drain_datanode_start.rs @@ -0,0 +1,61 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::any::Any; + +use common_procedure::{Context as ProcedureContext, Status}; +use common_telemetry::info; +use serde::{Deserialize, Serialize}; + +use crate::error::Result; +use crate::key::NodeSchedulingRole; +use crate::procedure::drain_datanode::drain_datanode_end::DrainDatanodeEnd; +use crate::procedure::drain_datanode::{Context, State}; + +/// Fences the datanode and builds the drain plan. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub(crate) struct DrainDatanodeStart; + +#[async_trait::async_trait] +#[typetag::serde] +impl State for DrainDatanodeStart { + async fn next( + &mut self, + ctx: &mut Context, + _procedure_ctx: &ProcedureContext, + ) -> Result<(Box, Status)> { + let node_id = ctx.persistent_ctx.node_id; + let marked = ctx + .node_scheduling_manager + .mark_unschedulable(NodeSchedulingRole::Datanode, node_id) + .await?; + if marked { + info!("Marked datanode {} as unschedulable", node_id); + } + + ctx.ensure_plan().await?; + info!( + "Planned drain of datanode {}, leader regions: {}, follower regions: {}", + node_id, + ctx.volatile_ctx.leader_regions().len(), + ctx.volatile_ctx.follower_regions().len() + ); + + Ok((Box::new(DrainDatanodeEnd), Status::executing(true))) + } + + fn as_any(&self) -> &dyn Any { + self + } +} diff --git a/src/meta-srv/src/procedure/drain_datanode/manager.rs b/src/meta-srv/src/procedure/drain_datanode/manager.rs new file mode 100644 index 00000000000..1211356bea0 --- /dev/null +++ b/src/meta-srv/src/procedure/drain_datanode/manager.rs @@ -0,0 +1,187 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::HashMap; +use std::collections::hash_map::Entry; +use std::sync::{Arc, RwLock}; + +use common_meta::DatanodeId; +use common_meta::key::TableMetadataManagerRef; +use common_procedure::{ProcedureContext, ProcedureId, ProcedureManagerRef, ProcedureWithId}; +use common_telemetry::info; +use snafu::ResultExt; + +use crate::error::{self, Result}; +use crate::node_scheduling::NodeSchedulingManagerRef; +use crate::procedure::drain_datanode::{ + Context as DrainDatanodeContext, DrainDatanodeProcedure, DrainDatanodeProcedureTask, + PersistentContext, +}; + +pub type DrainDatanodeManagerRef = Arc; + +/// Tracks the running drains, indexed by the drained datanode. +/// +/// Mirrors `RegionMigrationProcedureTracker`: a drain of a datanode that is already being +/// drained is rejected at submission instead of waiting for the procedure lock. +#[derive(Default, Clone)] +pub struct DrainDatanodeProcedureTracker { + running_procedures: Arc>>, +} + +impl DrainDatanodeProcedureTracker { + /// Returns the guard of a drain of `task.node_id`, or `None` if that datanode is already + /// being drained. + pub(crate) fn insert_running_procedure( + &self, + task: &DrainDatanodeProcedureTask, + ) -> Option { + let mut procedures = self.running_procedures.write().unwrap(); + match procedures.entry(task.node_id) { + Entry::Occupied(_) => None, + Entry::Vacant(v) => { + v.insert(task.clone()); + Some(DrainDatanodeProcedureGuard { + node_id: task.node_id, + running_procedures: self.running_procedures.clone(), + }) + } + } + } +} + +/// Removes the drained datanode from [`DrainDatanodeProcedureTracker`] when dropped. +pub(crate) struct DrainDatanodeProcedureGuard { + node_id: DatanodeId, + running_procedures: Arc>>, +} + +impl Drop for DrainDatanodeProcedureGuard { + fn drop(&mut self) { + self.running_procedures + .write() + .unwrap() + .remove(&self.node_id); + } +} + +/// Manages the [`DrainDatanodeProcedure`]s. +/// +/// It registers the procedure loader and submits drain tasks. A datanode is drained by at +/// most one procedure at a time: [`DrainDatanodeManager::submit_procedure`] rejects a second +/// drain through the [`DrainDatanodeProcedureTracker`]. +pub struct DrainDatanodeManager { + procedure_manager: ProcedureManagerRef, + table_metadata_manager: TableMetadataManagerRef, + node_scheduling_manager: NodeSchedulingManagerRef, + tracker: DrainDatanodeProcedureTracker, +} + +impl DrainDatanodeManager { + /// Returns a new [`DrainDatanodeManager`]. + pub fn new( + procedure_manager: ProcedureManagerRef, + table_metadata_manager: TableMetadataManagerRef, + node_scheduling_manager: NodeSchedulingManagerRef, + ) -> Self { + Self { + procedure_manager, + table_metadata_manager, + node_scheduling_manager, + tracker: DrainDatanodeProcedureTracker::default(), + } + } + + /// Returns the tracker of the running drains. + #[cfg(test)] + pub(crate) fn tracker(&self) -> DrainDatanodeProcedureTracker { + self.tracker.clone() + } + + /// Returns the context of a drain run. + fn context(&self, persistent_ctx: PersistentContext) -> DrainDatanodeContext { + DrainDatanodeContext::new( + self.table_metadata_manager.clone(), + self.node_scheduling_manager.clone(), + persistent_ctx, + ) + } + + /// Submits a [`DrainDatanodeProcedure`] and returns its id. + pub async fn submit_procedure( + &self, + procedure_context: ProcedureContext, + task: DrainDatanodeProcedureTask, + ) -> Result { + if task.target_node_id == Some(task.node_id) { + return error::InvalidArgumentsSnafu { + err_msg: "the target datanode must differ from the drained datanode", + } + .fail(); + } + + let Some(guard) = self.tracker.insert_running_procedure(&task) else { + return error::DrainDatanodeRunningSnafu { + node_id: task.node_id, + } + .fail(); + }; + + let persistent_ctx = PersistentContext { + node_id: task.node_id, + target_node_id: task.target_node_id, + }; + let procedure = DrainDatanodeProcedure::new(self.context(persistent_ctx), vec![guard]); + let procedure_with_id = + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); + let procedure_id = procedure_with_id.id; + info!( + "Starting drain datanode procedure {procedure_id} for datanode {}", + task.node_id + ); + self.procedure_manager + .submit(procedure_with_id) + .await + .context(error::SubmitProcedureSnafu)?; + + Ok(procedure_id) + } + + /// Registers the loader of [`DrainDatanodeProcedure`]. + pub(crate) fn try_start(&self) -> Result<()> { + let table_metadata_manager = self.table_metadata_manager.clone(); + let node_scheduling_manager = self.node_scheduling_manager.clone(); + let tracker = self.tracker.clone(); + self.procedure_manager + .register_loader( + DrainDatanodeProcedure::TYPE_NAME, + Box::new(move |json| { + let table_metadata_manager = table_metadata_manager.clone(); + let node_scheduling_manager = node_scheduling_manager.clone(); + let factory = move |persistent_ctx| { + DrainDatanodeContext::new( + table_metadata_manager, + node_scheduling_manager, + persistent_ctx, + ) + }; + DrainDatanodeProcedure::from_json(json, &tracker, factory) + .map(|procedure| Box::new(procedure) as _) + }), + ) + .context(error::RegisterProcedureLoaderSnafu { + type_name: DrainDatanodeProcedure::TYPE_NAME, + }) + } +} diff --git a/src/meta-srv/src/procedure/drain_datanode/utils.rs b/src/meta-srv/src/procedure/drain_datanode/utils.rs new file mode 100644 index 00000000000..a43ba815e18 --- /dev/null +++ b/src/meta-srv/src/procedure/drain_datanode/utils.rs @@ -0,0 +1,83 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::BTreeSet; + +use common_meta::DatanodeId; +use common_meta::key::TableMetadataManagerRef; +use common_meta::key::table_route::TableRouteValue; +use common_meta::rpc::router::region_distribution; +use futures::TryStreamExt; +use snafu::ResultExt; +use store_api::storage::RegionId; + +use crate::error::{self, Result}; +use crate::procedure::drain_datanode::DrainDatanodePlan; + +/// Builds the drain plan of `node_id`. +/// +/// Walks the datanode table index and reads the authoritative table routes, so the plan +/// reflects the routes rather than the index. Each table route is split by role with +/// [`region_distribution`], which keys the regions by datanode id. +pub(crate) async fn build_datanode_plan( + table_metadata_manager: &TableMetadataManagerRef, + node_id: DatanodeId, +) -> Result { + let mut leader_regions = BTreeSet::new(); + let mut follower_regions = BTreeSet::new(); + + let mut tables = table_metadata_manager + .datanode_table_manager() + .tables(node_id); + while let Some(value) = tables + .try_next() + .await + .context(error::TableMetadataManagerSnafu)? + { + let Some(table_route) = table_metadata_manager + .table_route_manager() + .table_route_storage() + .get(value.table_id) + .await + .context(error::TableMetadataManagerSnafu)? + else { + continue; + }; + let TableRouteValue::Physical(table_route) = table_route else { + continue; + }; + + let mut distribution = region_distribution(&table_route.region_routes); + let Some(roles) = distribution.remove(&node_id) else { + continue; + }; + leader_regions.extend( + roles + .leader_regions + .into_iter() + .map(|number| RegionId::new(value.table_id, number)), + ); + follower_regions.extend( + roles + .follower_regions + .into_iter() + .map(|number| RegionId::new(value.table_id, number)), + ); + } + + Ok(DrainDatanodePlan { + leader_regions: leader_regions.into_iter().collect(), + follower_regions: follower_regions.into_iter().collect(), + }) +} From 5e9bb3c689234b3c0a88fd3ece43987548345df8 Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Tue, 27 Oct 2026 03:09:13 +0000 Subject: [PATCH 103/239] feat: support native histograms in sampled remote read (#9971) * feat: support native histograms in sampled remote read Signed-off-by: shuiyisong * fix: align native histogram ingestion with remote read validation Signed-off-by: shuiyisong --------- Signed-off-by: shuiyisong --- docs/rfcs/2026-08-04-native-histograms.md | 76 +- src/common/query/src/native_histogram.rs | 36 +- .../query/src/native_histogram/encoding.rs | 1101 +++++++++++++++-- src/frontend/src/instance/prom_store.rs | 22 + src/servers/Cargo.toml | 1 + src/servers/src/prom_store.rs | 342 ++++- src/servers/tests/http/prom_store_test.rs | 58 +- tests-integration/src/otlp.rs | 80 ++ tests-integration/src/prom_store.rs | 323 +++++ tests-integration/tests/http.rs | 145 ++- tests/compatibility/README.md | 8 + .../native_histogram_remote_read/README.md | 184 +++ .../native_histogram_remote_read/response.pb | Bin 0 -> 420 bytes 13 files changed, 2270 insertions(+), 106 deletions(-) create mode 100644 tests/compatibility/wire/native_histogram_remote_read/README.md create mode 100644 tests/compatibility/wire/native_histogram_remote_read/response.pb diff --git a/docs/rfcs/2026-08-04-native-histograms.md b/docs/rfcs/2026-08-04-native-histograms.md index 6b0a9ebb757..cc31a8df5e4 100644 --- a/docs/rfcs/2026-08-04-native-histograms.md +++ b/docs/rfcs/2026-08-04-native-histograms.md @@ -16,8 +16,8 @@ Struct before persistence and are queried only as native histograms. Prometheus Remote Write 2.0 native histograms and cumulative OTLP/HTTP exponential histograms are enabled by default. Remote Write 1.0 histogram payloads are rejected instead of being acknowledged and dropped. -Native-histogram Remote Read is deferred; the existing Remote Read path -continues to return scalar samples only. +Remote Read returns scalar samples and native histograms through `SAMPLES`, +including exact stored integer counts. Streamed responses remain deferred. # Goals @@ -25,13 +25,15 @@ continues to return scalar samples only. 2. Keep one stable Prometheus-compatible persisted representation across Prometheus and OTLP ingestion. 3. Match Prometheus query behavior where it is observable and practical. -4. State intentional limitations explicitly so incomplete behavior is not +4. Return stored native histograms through sampled Remote Read without rounding + integer counts or losing supported histogram fields. +5. State intentional limitations explicitly so incomplete behavior is not mistaken for support. # Non-Goals - Supporting native histograms in Remote Write 1.0. -- Returning native histograms through Prometheus Remote Read. +- Streaming Remote Read responses (`STREAMED_XOR_CHUNKS`). - Persisting Remote Write metric metadata. - Persisting exemplars. - Converting OTLP delta values to cumulative values during ingestion. Supported @@ -107,19 +109,57 @@ beyond it. The overflow index is `(1024 << schema) + 1` for non-negative schemas and `(1024 >> -schema) + 1` for negative schemas. Inputs beyond that index are rejected rather than clamped or merged into infinity. +Ingestion and sampled Remote Read share histogram layout validation. An initial +exponential span offset of `i32::MIN` is rejected on either side, including for +empty spans; other negative initial offsets remain supported. For custom +histograms with `N` bounds, bucket index `N` is the overflow bucket, and empty +spans may end at `N + 1`. Span extents beyond that boundary are rejected. + +This tightens ingestion acceptance: these layouts were previously accepted on +write but rejected by Remote Read. A validation failure rejects the entire +Remote Write 2.0 request with HTTP 400 and all written-count headers set to zero, +before any samples or histograms are persisted. Healthy samples in the same +request are also rejected; senders that do not retry HTTP 400 can drop them. + Exemplars and declared metric metadata are accepted on the wire but are not persisted by this feature. ## Remote Read -Native-histogram Remote Read is deferred. GreptimeDB's existing Remote Read path -continues to negotiate `SAMPLES` responses and serialize scalar samples only. +`/v1/prometheus/read` returns scalar samples and integer or float native +histograms through the existing Snappy-compressed `SAMPLES` protobuf response. +Omitting accepted types, requesting `SAMPLES`, or requesting +`[STREAMED_XOR_CHUNKS, SAMPLES]` selects this format. Streaming-only requests +remain unsupported; the response is still buffered as a whole. -A follow-up can map the stored Struct into `TimeSeries.histograms`. It must -preserve integer counts without passing through `f64`, retain the stored -histogram fields, and decide sampled versus streamed response coverage. The -sampled protobuf has no start-timestamp field, so that omission must remain -explicit if sampled responses are implemented. +One request can query both scalar and histogram metrics. A histogram series +can contain integer and float histogram samples, but the one-sample-kind-per- +metric-table rule remains unchanged. The selected field must be `Float64` or +the canonical native-histogram Struct, regardless of its stored column name. + +Integer counts remain exact within the existing nonnegative `i64` storage +range, including above `2^53`. Stored absolute integer buckets become signed +wire deltas, retaining their sequence across span boundaries. Float bucket +counts remain absolute. Schema, spans, custom bounds, zero threshold, sum +(including stale-NaN bits), reset hint, and millisecond timestamps are retained. +The sampled protobuf has no `start_timestamp` field, so that value remains in +storage but is omitted from the response. Raw-delta OTLP histograms retain +their temporality labels and gauge hints; Remote Read does not accumulate them +or reconstruct the original OTLP point. + +Malformed or wire-incompatible non-null stored histograms fail the entire +`ReadRequest` with HTTP 500, including healthy queries in that request. Read +validation remains necessary for older rows and typed inserts that bypass the +histogram ingestion encoder. Incompatible stored rows are not repaired by the +stricter ingestion checks. Null histograms remain absent samples. + +Custom-bucket interoperability is exercised against Prometheus revision +`e471c7708c6b0b13e2050222bbd5a907b36f4113`, including `AtHistogram` and +`AtFloatHistogram` on integer custom buckets. This is a tested revision, not +a minimum-release claim; older readers can lose custom bounds during integer- +to-float conversion. Integer payloads are not downgraded for those readers. +The [wire fixture and iterator check](../../tests/compatibility/wire/native_histogram_remote_read/README.md) +record the independent expectations and reproduction commands. ## OTLP @@ -340,16 +380,20 @@ partial-success behavior, kind exclusivity, stale-marker selector semantics, synthetic-zero rates, incompatible empty layouts, overflow-safe averages, layout-sensitive equality, infinite custom midpoints, and exponential overflow indices for schemas `-4`, `0`, and `8`. -Behavioral coverage also verifies frontend PromQL warning and info responses. +Behavioral coverage also verifies frontend PromQL warning and info responses, +sampled Remote Read before and after flush in standalone and distributed +setups, exact large integers, stale-marker bits, timestamp units, and OTLP +normalization. Frozen upstream protobuf fixtures guard the wire encoder. This work does not change GreptimeDB's persisted Struct or protobuf -dependencies. It adds the disabled-by-default public OTLP gate; existing -native-histogram data remains readable. +dependencies. Remote Read uses the existing endpoint and negotiation without +a new configuration option; existing native-histogram data remains readable. # Future Work -- Native-histogram Remote Read, including exact integer round-trips and streamed - chunks with start timestamps. +- TODO: Add native histogram and float histogram chunks to + `STREAMED_XOR_CHUNKS`, with checksummed framing and bounded buffering. Decide + separately whether to include start-timestamp-bearing chunk encodings. - Raw-delta OTLP exponential-histogram storage and temporality-aware native histogram queries, as described above; no delta-to-cumulative conversion. - Persistent Remote Write metadata and accurate help/unit updates. diff --git a/src/common/query/src/native_histogram.rs b/src/common/query/src/native_histogram.rs index d8a559a2dc5..38a7895c27a 100644 --- a/src/common/query/src/native_histogram.rs +++ b/src/common/query/src/native_histogram.rs @@ -36,7 +36,10 @@ use datafusion::arrow::datatypes::{ use datafusion_common::{DataFusionError, Result as DfResult}; use datatypes::data_type::{ConcreteDataType, DataType}; use datatypes::types::{StructField, StructType}; -pub use encoding::{NativeHistogramError, encode_native_histogram, native_histogram_column_schema}; +pub use encoding::{ + NativeHistogramError, decode_native_histogram_for_remote_read, encode_native_histogram, + native_histogram_column_schema, +}; use once_cell::sync::Lazy; use crate::prelude::greptime_native_histogram; @@ -136,6 +139,37 @@ pub fn is_native_histogram_value_schema(name: &str, data_type: &ConcreteDataType name == greptime_native_histogram() && is_native_histogram_value_type(data_type) } +#[cfg(any(test, feature = "testing"))] +pub mod testing { + use api::helper::pb_value_to_value_ref; + use api::v1::{Value, value::ValueData}; + use datafusion::arrow::array::{AsArray, StructArray}; + use datatypes::vectors::Helper; + + use crate::native_histogram::{native_histogram_column_schema, native_histogram_value_type}; + + /// Builds stored arrays through the same protobuf-to-vector path as ingestion. + pub fn histogram_array(values: &[Option]) -> StructArray { + let schema = native_histogram_column_schema().unwrap(); + let rows = values + .iter() + .map(|value_data| { + datatypes::value::Value::from(pb_value_to_value_ref( + &Value { + value_data: value_data.clone(), + }, + schema.datatype_extension.as_ref(), + )) + }) + .collect::>(); + Helper::try_from_row_into_vector(&rows, native_histogram_value_type()) + .unwrap() + .to_arrow_array() + .as_struct() + .clone() + } +} + /// Schema identifier for native histograms with explicit custom bucket bounds. pub const CUSTOM_BUCKETS_SCHEMA: i32 = -53; const MIN_EXPONENTIAL_SCHEMA: i32 = -4; diff --git a/src/common/query/src/native_histogram/encoding.rs b/src/common/query/src/native_histogram/encoding.rs index be7b35ec454..d5e1e545f84 100644 --- a/src/common/query/src/native_histogram/encoding.rs +++ b/src/common/query/src/native_histogram/encoding.rs @@ -15,12 +15,15 @@ use api::greptime_proto::io::prometheus::write::v2::histogram::{Count, ZeroCount}; use api::greptime_proto::io::prometheus::write::v2::{BucketSpan, Histogram}; use api::helper::ColumnDataTypeWrapper; +use api::prom_store::remote; use api::v1::value::ValueData; use api::v1::{ColumnSchema, ListValue, SemanticType, Value}; +use datafusion::arrow::array::{Array, StructArray}; +use datafusion::arrow::datatypes::{Float64Type, Int32Type, Int64Type, TimestampMillisecondType}; use snafu::{Snafu, ensure}; use crate::native_histogram::{ - CUSTOM_BUCKETS_SCHEMA, MAX_EXPONENTIAL_SCHEMA, NATIVE_HISTOGRAM_FIELD_NAMES, + self, CUSTOM_BUCKETS_SCHEMA, MAX_EXPONENTIAL_SCHEMA, NATIVE_HISTOGRAM_FIELD_NAMES, exponential_overflow_bucket_index, native_histogram_value_type, }; use crate::prelude::greptime_native_histogram; @@ -78,7 +81,6 @@ pub fn encode_native_histogram(histogram: &Histogram) -> Result { ]); if uses_float_counts { - validate_float_native_histogram_counts(histogram)?; let count = match histogram.count.as_ref() { Some(Count::CountFloat(count)) => *count, _ => 0.0, @@ -87,6 +89,12 @@ pub fn encode_native_histogram(histogram: &Histogram) -> Result { Some(ZeroCount::ZeroCountFloat(zero_count)) => *zero_count, _ => 0.0, }; + validate_float_native_histogram_counts( + count, + zero_count, + &histogram.positive_counts, + &histogram.negative_counts, + )?; items.extend([ null_pb_value(), null_pb_value(), @@ -108,7 +116,13 @@ pub fn encode_native_histogram(histogram: &Histogram) -> Result { }; let positive_buckets = bucket_counts_from_deltas(&histogram.positive_deltas)?; let negative_buckets = bucket_counts_from_deltas(&histogram.negative_deltas)?; - validate_integer_native_histogram_counts(histogram, &positive_buckets, &negative_buckets)?; + validate_integer_native_histogram_counts( + count, + zero_count, + histogram.sum, + &positive_buckets, + &negative_buckets, + )?; let count = i64::try_from(count).map_err(|_| NativeHistogramError { message: format!("native histogram integer count {count} overflows i64"), })?; @@ -130,27 +144,215 @@ pub fn encode_native_histogram(histogram: &Histogram) -> Result { Ok(ValueData::StructValue(api::v1::StructValue { items })) } -fn validate_native_histogram(histogram: &Histogram, uses_float_counts: bool) -> Result<()> { - let exponential_overflow_index = validate_native_histogram_schema(histogram.schema)?; - validate_native_histogram_custom_values(histogram)?; +/// Decodes a stored histogram without rounding integer counts through `f64`. +/// +/// Null structs are absent samples. Malformed or sampled-wire-incompatible +/// values are errors. The sampled format has no field for `start_timestamp`. +pub fn decode_native_histogram_for_remote_read( + array: &StructArray, + row: usize, + timestamp: i64, +) -> Result> { + use crate::native_histogram::{ + COUNT_F64_FIELD, COUNT_I64_FIELD, CUSTOM_VALUES_FIELD, NEGATIVE_BUCKETS_F64_FIELD, + NEGATIVE_BUCKETS_I64_FIELD, NEGATIVE_SPAN_LENGTHS_FIELD, NEGATIVE_SPAN_OFFSETS_FIELD, + POSITIVE_BUCKETS_F64_FIELD, POSITIVE_BUCKETS_I64_FIELD, POSITIVE_SPAN_LENGTHS_FIELD, + POSITIVE_SPAN_OFFSETS_FIELD, RESET_HINT_FIELD, SCHEMA_FIELD, START_TIMESTAMP_FIELD, + SUM_FIELD, ZERO_COUNT_F64_FIELD, ZERO_COUNT_I64_FIELD, ZERO_THRESHOLD_FIELD, + }; - if histogram.schema == CUSTOM_BUCKETS_SCHEMA { - ensure!( - histogram.zero_threshold == 0.0 && native_histogram_zero_count_is_zero(histogram), - NativeHistogramSnafu { - message: "custom native histogram must not use a zero bucket" - } - ); - ensure!( - histogram.negative_spans.is_empty() - && histogram.negative_deltas.is_empty() - && histogram.negative_counts.is_empty(), - NativeHistogramSnafu { - message: "custom native histogram must not use negative buckets" - } - ); + ensure!( + row < array.len(), + NativeHistogramSnafu { + message: "native histogram row is out of bounds" + } + ); + if array.is_null(row) { + return Ok(None); } + let invalid_field = |error| NativeHistogramError { + message: format!("invalid stored native histogram: {error}"), + }; + let read_spans = |name, offsets, lengths| -> Result> { + let offsets = native_histogram::list_values::(array, offsets, row) + .map_err(invalid_field)?; + let lengths = native_histogram::list_values::(array, lengths, row) + .map_err(invalid_field)?; + native_histogram::read_spans(offsets, lengths, name) + .map_err(invalid_field)? + .into_iter() + .map(|span| { + Ok(remote::BucketSpan { + offset: span.offset, + length: u32::try_from(span.length).map_err(|_| NativeHistogramError { + message: format!("native histogram {name} span length is negative"), + })?, + }) + }) + .collect() + }; + let mut histogram = remote::Histogram { + schema: native_histogram::required_primitive::(array, SCHEMA_FIELD, row) + .map_err(invalid_field)?, + zero_threshold: native_histogram::required_primitive::( + array, + ZERO_THRESHOLD_FIELD, + row, + ) + .map_err(invalid_field)?, + sum: native_histogram::required_primitive::(array, SUM_FIELD, row) + .map_err(invalid_field)?, + reset_hint: native_histogram::required_primitive::(array, RESET_HINT_FIELD, row) + .map_err(invalid_field)?, + custom_values: native_histogram::list_values::( + array, + CUSTOM_VALUES_FIELD, + row, + ) + .map_err(invalid_field)?, + positive_spans: read_spans( + "positive", + POSITIVE_SPAN_OFFSETS_FIELD, + POSITIVE_SPAN_LENGTHS_FIELD, + )?, + negative_spans: read_spans( + "negative", + NEGATIVE_SPAN_OFFSETS_FIELD, + NEGATIVE_SPAN_LENGTHS_FIELD, + )?, + timestamp, + ..Default::default() + }; + // Validate the stored child even though the sampled wire format cannot carry it. + native_histogram::optional_primitive::( + array, + START_TIMESTAMP_FIELD, + row, + ) + .map_err(invalid_field)?; + let count_i64 = native_histogram::optional_primitive::(array, COUNT_I64_FIELD, row) + .map_err(invalid_field)?; + let count_f64 = + native_histogram::optional_primitive::(array, COUNT_F64_FIELD, row) + .map_err(invalid_field)?; + let zero_i64 = + native_histogram::optional_primitive::(array, ZERO_COUNT_I64_FIELD, row) + .map_err(invalid_field)?; + let zero_f64 = + native_histogram::optional_primitive::(array, ZERO_COUNT_F64_FIELD, row) + .map_err(invalid_field)?; + let positive_i64 = + native_histogram::list_values::(array, POSITIVE_BUCKETS_I64_FIELD, row) + .map_err(invalid_field)?; + let negative_i64 = + native_histogram::list_values::(array, NEGATIVE_BUCKETS_I64_FIELD, row) + .map_err(invalid_field)?; + let positive_f64 = + native_histogram::list_values::(array, POSITIVE_BUCKETS_F64_FIELD, row) + .map_err(invalid_field)?; + let negative_f64 = + native_histogram::list_values::(array, NEGATIVE_BUCKETS_F64_FIELD, row) + .map_err(invalid_field)?; + let (zero_count_is_zero, positive_len, negative_len) = match (count_i64, count_f64) { + (Some(count), None) => { + ensure!( + zero_f64.is_none() && positive_f64.is_empty() && negative_f64.is_empty(), + NativeHistogramSnafu { + message: "integer native histogram has populated float fields" + } + ); + let count = u64::try_from(count).map_err(|_| NativeHistogramError { + message: "native histogram integer count is negative".to_string(), + })?; + let zero_count = + u64::try_from(zero_i64.unwrap_or_default()).map_err(|_| NativeHistogramError { + message: "native histogram integer zero_count is negative".to_string(), + })?; + validate_integer_native_histogram_counts( + count, + zero_count, + histogram.sum, + &positive_i64, + &negative_i64, + )?; + histogram.count = Some(remote::histogram::Count::CountInt(count)); + histogram.zero_count = Some(remote::histogram::ZeroCount::ZeroCountInt(zero_count)); + histogram.positive_deltas = bucket_deltas_from_counts(&positive_i64)?; + histogram.negative_deltas = bucket_deltas_from_counts(&negative_i64)?; + (zero_count == 0, positive_i64.len(), negative_i64.len()) + } + (None, Some(count)) => { + ensure!( + zero_i64.is_none() && positive_i64.is_empty() && negative_i64.is_empty(), + NativeHistogramSnafu { + message: "float native histogram has populated integer fields" + } + ); + let zero_count = zero_f64.unwrap_or_default(); + validate_float_native_histogram_counts( + count, + zero_count, + &positive_f64, + &negative_f64, + )?; + histogram.count = Some(remote::histogram::Count::CountFloat(count)); + histogram.zero_count = Some(remote::histogram::ZeroCount::ZeroCountFloat(zero_count)); + histogram.positive_counts = positive_f64; + histogram.negative_counts = negative_f64; + ( + zero_count == 0.0, + histogram.positive_counts.len(), + histogram.negative_counts.len(), + ) + } + _ => { + return NativeHistogramSnafu { + message: "stored native histogram must have exactly one populated count family", + } + .fail(); + } + }; + validate_native_histogram_layout( + histogram.schema, + histogram.zero_threshold, + zero_count_is_zero, + &histogram.custom_values, + ( + histogram + .positive_spans + .iter() + .map(|s| (s.offset, s.length)), + positive_len, + ), + ( + histogram + .negative_spans + .iter() + .map(|s| (s.offset, s.length)), + negative_len, + ), + )?; + Ok(Some(histogram)) +} + +fn bucket_deltas_from_counts(counts: &[i64]) -> Result> { + let mut previous = 0i64; + counts + .iter() + .map(|&count| { + let delta = count + .checked_sub(previous) + .ok_or_else(|| NativeHistogramError { + message: "native histogram bucket delta overflows i64".to_string(), + })?; + previous = count; + Ok(delta) + }) + .collect() +} + +fn validate_native_histogram(histogram: &Histogram, uses_float_counts: bool) -> Result<()> { let (positive_buckets, negative_buckets) = if uses_float_counts { ( histogram.positive_counts.len(), @@ -162,28 +364,82 @@ fn validate_native_histogram(histogram: &Histogram, uses_float_counts: bool) -> histogram.negative_deltas.len(), ) }; + validate_native_histogram_layout( + histogram.schema, + histogram.zero_threshold, + native_histogram_zero_count_is_zero(histogram), + &histogram.custom_values, + ( + histogram + .positive_spans + .iter() + .map(|s| (s.offset, s.length)), + positive_buckets, + ), + ( + histogram + .negative_spans + .iter() + .map(|s| (s.offset, s.length)), + negative_buckets, + ), + ) +} + +fn validate_native_histogram_layout( + schema: i32, + zero_threshold: f64, + zero_count_is_zero: bool, + custom_values: &[f64], + positive: (impl ExactSizeIterator + Clone, usize), + negative: (impl ExactSizeIterator + Clone, usize), +) -> Result<()> { + let exponential_overflow_index = validate_native_histogram_schema(schema)?; + validate_native_histogram_custom_values(schema, custom_values)?; + if schema == CUSTOM_BUCKETS_SCHEMA { + ensure!( + zero_threshold == 0.0 && zero_count_is_zero, + NativeHistogramSnafu { + message: "custom native histogram must not use a zero bucket" + } + ); + ensure!( + negative.0.len() == 0 && negative.1 == 0, + NativeHistogramSnafu { + message: "custom native histogram must not use negative buckets" + } + ); + } + let bucket_index_range = if let Some(overflow_index) = exponential_overflow_index { (i32::MIN, overflow_index) } else { ( 0, - i32::try_from(histogram.custom_values.len()).map_err(|_| NativeHistogramError { + i32::try_from(custom_values.len()).map_err(|_| NativeHistogramError { message: "custom native histogram has too many custom_values".to_string(), })?, ) }; - validate_native_histogram_spans( - "positive", - &histogram.positive_spans, - positive_buckets, - bucket_index_range, - )?; - validate_native_histogram_spans( - "negative", - &histogram.negative_spans, - negative_buckets, - bucket_index_range, - )?; + let positive_extent = + validate_native_histogram_spans("positive", positive.0, positive.1, bucket_index_range)?; + validate_native_histogram_spans("negative", negative.0, negative.1, bucket_index_range)?; + + if schema == CUSTOM_BUCKETS_SCHEMA { + let max_extent = + bucket_index_range + .1 + .checked_add(1) + .ok_or_else(|| NativeHistogramError { + message: "custom native histogram extent overflows i32".to_string(), + })?; + ensure!( + positive_extent <= max_extent, + NativeHistogramSnafu { + message: "custom native histogram span extent exceeds custom_values plus the overflow bucket" + } + ); + } Ok(()) } @@ -208,10 +464,10 @@ fn validate_native_histogram_schema(schema: i32) -> Result> { } } -fn validate_native_histogram_custom_values(histogram: &Histogram) -> Result<()> { - if histogram.schema != CUSTOM_BUCKETS_SCHEMA { +fn validate_native_histogram_custom_values(schema: i32, custom_values: &[f64]) -> Result<()> { + if schema != CUSTOM_BUCKETS_SCHEMA { ensure!( - histogram.custom_values.is_empty(), + custom_values.is_empty(), NativeHistogramSnafu { message: "standard native histogram must not use custom_values" } @@ -219,7 +475,7 @@ fn validate_native_histogram_custom_values(histogram: &Histogram) -> Result<()> return Ok(()); } - for value in &histogram.custom_values { + for value in custom_values { ensure!( !value.is_nan() && *value != f64::INFINITY, NativeHistogramSnafu { @@ -227,7 +483,7 @@ fn validate_native_histogram_custom_values(histogram: &Histogram) -> Result<()> } ); } - for values in histogram.custom_values.windows(2) { + for values in custom_values.windows(2) { ensure!( values[0] < values[1], NativeHistogramSnafu { @@ -241,12 +497,12 @@ fn validate_native_histogram_custom_values(histogram: &Histogram) -> Result<()> fn validate_native_histogram_spans( name: &str, - spans: &[BucketSpan], + spans: impl Iterator + Clone, bucket_count: usize, bucket_index_range: (i32, i32), -) -> Result<()> { - let span_len = spans.iter().try_fold(0usize, |sum, span| { - let length = usize::try_from(span.length).map_err(|_| NativeHistogramError { +) -> Result { + let span_len = spans.clone().try_fold(0usize, |sum, (_, length)| { + let length = usize::try_from(length).map_err(|_| NativeHistogramError { message: format!("native histogram {name} span length exceeds usize"), })?; sum.checked_add(length).ok_or_else(|| NativeHistogramError { @@ -263,28 +519,36 @@ fn validate_native_histogram_spans( ); let mut current_index = 0i32; - for (span_index, span) in spans.iter().enumerate() { + for (span_index, (offset, length)) in spans.enumerate() { ensure!( - span.offset >= 0 || (span_index == 0 && bucket_index_range.0 == i32::MIN), + offset >= 0 || (span_index == 0 && bucket_index_range.0 == i32::MIN), NativeHistogramSnafu { message: format!( "native histogram {name} span {} has negative offset {}", span_index + 1, - span.offset + offset + ) + } + ); + ensure!( + span_index != 0 || offset != i32::MIN, + NativeHistogramSnafu { + message: format!( + "native histogram {name} initial span offset i32::MIN is incompatible with Remote Read" ) } ); current_index = if span_index == 0 { - span.offset + offset } else { current_index - .checked_add(span.offset) + .checked_add(offset) .ok_or_else(|| NativeHistogramError { message: format!("native histogram {name} span index overflows i32"), })? }; - for _ in 0..span.length { + for _ in 0..length { ensure!( (bucket_index_range.0..=bucket_index_range.1).contains(¤t_index), NativeHistogramSnafu { @@ -301,14 +565,15 @@ fn validate_native_histogram_spans( } } - Ok(()) + Ok(current_index) } -fn validate_float_native_histogram_counts(histogram: &Histogram) -> Result<()> { - let count = match histogram.count.as_ref() { - Some(Count::CountFloat(count)) => *count, - _ => 0.0, - }; +fn validate_float_native_histogram_counts( + count: f64, + zero_count: f64, + positive_buckets: &[f64], + negative_buckets: &[f64], +) -> Result<()> { ensure!( count >= 0.0 || count.is_nan(), NativeHistogramSnafu { @@ -316,10 +581,6 @@ fn validate_float_native_histogram_counts(histogram: &Histogram) -> Result<()> { } ); - let zero_count = match histogram.zero_count.as_ref() { - Some(ZeroCount::ZeroCountFloat(zero_count)) => *zero_count, - _ => 0.0, - }; ensure!( zero_count >= 0.0 || zero_count.is_nan(), NativeHistogramSnafu { @@ -328,8 +589,8 @@ fn validate_float_native_histogram_counts(histogram: &Histogram) -> Result<()> { ); for (name, counts) in [ - ("positive", &histogram.positive_counts), - ("negative", &histogram.negative_counts), + ("positive", positive_buckets), + ("negative", negative_buckets), ] { for (index, count) in counts.iter().enumerate() { ensure!( @@ -348,18 +609,12 @@ fn validate_float_native_histogram_counts(histogram: &Histogram) -> Result<()> { } fn validate_integer_native_histogram_counts( - histogram: &Histogram, + count: u64, + zero_count: u64, + sum: f64, positive_buckets: &[i64], negative_buckets: &[i64], ) -> Result<()> { - let count = match histogram.count.as_ref() { - Some(Count::CountInt(count)) => *count, - _ => 0, - }; - let zero_count = match histogram.zero_count.as_ref() { - Some(ZeroCount::ZeroCountInt(zero_count)) => *zero_count, - _ => 0, - }; let bucket_count = positive_buckets .iter() @@ -375,7 +630,7 @@ fn validate_integer_native_histogram_counts( }) })?; ensure!( - if histogram.sum.is_nan() { + if sum.is_nan() { bucket_count <= count } else { bucket_count == count @@ -510,6 +765,710 @@ fn bucket_counts_from_deltas(deltas: &[i64]) -> Result> { mod tests { use super::*; + use std::sync::Arc; + + use datafusion::arrow::array::{ArrayRef, Float64Array, Int32Array}; + use datafusion::arrow::datatypes::Field; + + use crate::native_histogram::testing::histogram_array; + use crate::native_histogram::*; + + fn empty_histogram(float: bool) -> Histogram { + Histogram { + count: Some(if float { + Count::CountFloat(0.0) + } else { + Count::CountInt(0) + }), + zero_count: Some(if float { + ZeroCount::ZeroCountFloat(0.0) + } else { + ZeroCount::ZeroCountInt(0) + }), + ..Default::default() + } + } + + fn stored(histogram: &Histogram) -> StructArray { + histogram_array(&[Some(encode_native_histogram(histogram).unwrap())]) + } + + fn modified(histogram: &Histogram, changes: &[(&str, Value)]) -> StructArray { + let ValueData::StructValue(mut value) = encode_native_histogram(histogram).unwrap() else { + unreachable!() + }; + for (name, replacement) in changes { + let index = NATIVE_HISTOGRAM_FIELD_NAMES + .iter() + .position(|field| field == name) + .unwrap(); + value.items[index] = replacement.clone(); + } + histogram_array(&[Some(ValueData::StructValue(value))]) + } + + fn assert_invalid(array: &StructArray, message: &str) { + let error = decode_native_histogram_for_remote_read(array, 0, 123).unwrap_err(); + assert!( + error.to_string().contains(message), + "{error}: expected {message}" + ); + } + + #[test] + fn native_histogram_remote_read_preserves_integer_counts_and_deltas() { + for count in [(1u64 << 53) + 1, i64::MAX as u64] { + let histogram = Histogram { + count: Some(Count::CountInt(count)), + zero_count: Some(ZeroCount::ZeroCountInt(count - 12)), + positive_spans: vec![ + BucketSpan { + offset: -2, + length: 2, + }, + BucketSpan { + offset: 3, + length: 1, + }, + ], + negative_spans: vec![ + BucketSpan { + offset: -3, + length: 1, + }, + BucketSpan { + offset: 2, + length: 2, + }, + ], + positive_deltas: vec![1, 2, -1], + negative_deltas: vec![1, 2, -1], + sum: 10.0, + reset_hint: 42, + start_timestamp: 123, + ..Default::default() + }; + let array = stored(&histogram); + assert_eq!( + native_histogram::list_values::(&array, POSITIVE_BUCKETS_I64_FIELD, 0) + .unwrap(), + vec![1, 3, 2] + ); + assert_eq!( + native_histogram::list_values::(&array, NEGATIVE_BUCKETS_I64_FIELD, 0) + .unwrap(), + vec![1, 3, 2] + ); + let decoded = decode_native_histogram_for_remote_read(&array, 0, -2) + .unwrap() + .unwrap(); + assert_eq!( + decoded.count, + Some(remote::histogram::Count::CountInt(count)) + ); + assert_eq!( + decoded.zero_count, + Some(remote::histogram::ZeroCount::ZeroCountInt(count - 12)) + ); + assert_eq!(decoded.positive_deltas, vec![1, 2, -1]); + assert_eq!(decoded.negative_deltas, vec![1, 2, -1]); + assert_eq!( + decoded + .positive_spans + .iter() + .map(|s| (s.offset, s.length)) + .collect::>(), + vec![(-2, 2), (3, 1)] + ); + assert_eq!(decoded.reset_hint, 42); + assert_eq!(decoded.timestamp, -2); + } + // A NaN sum permits unobserved NaNs; an incorrect [3, 2] delta vector + // would also pass a total check, so assert the deltas independently. + let histogram = Histogram { + count: Some(Count::CountInt(8)), + sum: f64::NAN, + positive_spans: vec![BucketSpan { + offset: 0, + length: 2, + }], + positive_deltas: vec![3, -1], + ..Default::default() + }; + let decoded = decode_native_histogram_for_remote_read(&stored(&histogram), 0, 0) + .unwrap() + .unwrap(); + assert_eq!(decoded.positive_deltas, vec![3, -1]); + } + + #[test] + fn native_histogram_remote_read_zero_families_and_float_bits() { + for float in [false, true] { + let histogram = empty_histogram(float); + let array = modified( + &histogram, + &[ + ( + if float { + ZERO_COUNT_F64_FIELD + } else { + ZERO_COUNT_I64_FIELD + }, + null_pb_value(), + ), + (POSITIVE_BUCKETS_I64_FIELD, null_pb_value()), + (POSITIVE_BUCKETS_F64_FIELD, null_pb_value()), + ], + ); + let decoded = decode_native_histogram_for_remote_read(&array, 0, 1) + .unwrap() + .unwrap(); + assert_eq!( + matches!( + decoded.count, + Some(remote::histogram::Count::CountFloat(0.0)) + ), + float + ); + assert_eq!( + matches!(decoded.count, Some(remote::histogram::Count::CountInt(0))), + !float + ); + assert_eq!( + matches!( + decoded.zero_count, + Some(remote::histogram::ZeroCount::ZeroCountFloat(0.0)) + ), + float + ); + assert_eq!( + matches!( + decoded.zero_count, + Some(remote::histogram::ZeroCount::ZeroCountInt(0)) + ), + !float + ); + for bits in [0x7ff8_0000_0000_0123, 0x7ff0_0000_0000_0002] { + let histogram = Histogram { + sum: f64::from_bits(bits), + ..histogram.clone() + }; + let decoded = decode_native_histogram_for_remote_read(&stored(&histogram), 0, 1) + .unwrap() + .unwrap(); + assert_eq!(decoded.sum.to_bits(), bits); + } + } + let histogram = Histogram { + count: Some(Count::CountFloat(0.3)), + positive_spans: vec![BucketSpan { + offset: 0, + length: 2, + }], + positive_counts: vec![0.1, 0.2], + ..empty_histogram(true) + }; + let decoded = decode_native_histogram_for_remote_read(&stored(&histogram), 0, 1) + .unwrap() + .unwrap(); + assert_eq!(decoded.positive_counts, vec![0.1, 0.2]); + for value in [f64::INFINITY, f64::from_bits(0x7ff8_0000_0000_0123)] { + let histogram = Histogram { + count: Some(Count::CountFloat(value)), + zero_count: Some(ZeroCount::ZeroCountFloat(value)), + positive_counts: vec![value, value], + ..histogram.clone() + }; + let decoded = decode_native_histogram_for_remote_read(&stored(&histogram), 0, 1) + .unwrap() + .unwrap(); + let Some(remote::histogram::Count::CountFloat(count)) = decoded.count else { + panic!() + }; + assert_eq!(count.to_bits(), value.to_bits()); + assert_eq!(decoded.positive_counts[0].to_bits(), value.to_bits()); + } + assert!( + decode_native_histogram_for_remote_read(&histogram_array(&[None]), 0, 0) + .unwrap() + .is_none() + ); + } + + #[test] + fn native_histogram_remote_read_rejects_invalid_stored_children() { + let integer = empty_histogram(false); + let float = empty_histogram(true); + for field in [ + SCHEMA_FIELD, + ZERO_THRESHOLD_FIELD, + SUM_FIELD, + RESET_HINT_FIELD, + ] { + assert_invalid(&modified(&integer, &[(field, null_pb_value())]), "is null"); + } + for (histogram, changes, message) in [ + ( + &integer, + vec![(COUNT_I64_FIELD, null_pb_value())], + "exactly one", + ), + ( + &integer, + vec![(COUNT_F64_FIELD, pb_value(ValueData::F64Value(0.0)))], + "exactly one", + ), + ( + &integer, + vec![(ZERO_COUNT_F64_FIELD, pb_value(ValueData::F64Value(0.0)))], + "populated float", + ), + ( + &integer, + vec![(POSITIVE_BUCKETS_F64_FIELD, f64_list_value([0.0]))], + "populated float", + ), + ( + &integer, + vec![(NEGATIVE_BUCKETS_F64_FIELD, f64_list_value([0.0]))], + "populated float", + ), + ( + &float, + vec![(ZERO_COUNT_I64_FIELD, pb_value(ValueData::I64Value(0)))], + "populated integer", + ), + ( + &float, + vec![(POSITIVE_BUCKETS_I64_FIELD, i64_list_value([0]))], + "populated integer", + ), + ( + &float, + vec![(NEGATIVE_BUCKETS_I64_FIELD, i64_list_value([0]))], + "populated integer", + ), + ( + &integer, + vec![(COUNT_I64_FIELD, pb_value(ValueData::I64Value(-1)))], + "count is negative", + ), + ( + &integer, + vec![(ZERO_COUNT_I64_FIELD, pb_value(ValueData::I64Value(-1)))], + "zero_count is negative", + ), + ( + &integer, + vec![(POSITIVE_BUCKETS_I64_FIELD, i64_list_value([-1]))], + "count is negative", + ), + ( + &integer, + vec![(NEGATIVE_BUCKETS_I64_FIELD, i64_list_value([-1]))], + "count is negative", + ), + ( + &float, + vec![(COUNT_F64_FIELD, pb_value(ValueData::F64Value(-1.0)))], + "must not be negative", + ), + ( + &float, + vec![(ZERO_COUNT_F64_FIELD, pb_value(ValueData::F64Value(-1.0)))], + "must not be negative", + ), + ( + &float, + vec![(POSITIVE_BUCKETS_F64_FIELD, f64_list_value([-1.0]))], + "must not be negative", + ), + ( + &float, + vec![(NEGATIVE_BUCKETS_F64_FIELD, f64_list_value([-1.0]))], + "must not be negative", + ), + ( + &integer, + vec![ + (COUNT_I64_FIELD, pb_value(ValueData::I64Value(1))), + ( + POSITIVE_BUCKETS_I64_FIELD, + i64_list_value([i64::MAX, i64::MAX, 3]), + ), + ], + "overflows u64", + ), + ( + &integer, + vec![(COUNT_I64_FIELD, pb_value(ValueData::I64Value(1)))], + "observations", + ), + ( + &integer, + vec![(SCHEMA_FIELD, pb_value(ValueData::I32Value(53)))], + "unsupported", + ), + ( + &integer, + vec![(POSITIVE_SPAN_OFFSETS_FIELD, i32_list_value([0]))], + "offsets and lengths mismatch", + ), + ( + &integer, + vec![ + (POSITIVE_SPAN_OFFSETS_FIELD, i32_list_value([0])), + (POSITIVE_SPAN_LENGTHS_FIELD, i32_list_value([-1])), + ], + "negative length", + ), + ( + &integer, + vec![ + (POSITIVE_SPAN_OFFSETS_FIELD, i32_list_value([0])), + (POSITIVE_SPAN_LENGTHS_FIELD, i32_list_value([1])), + ], + "describe 1 buckets", + ), + ( + &integer, + vec![ + (POSITIVE_SPAN_OFFSETS_FIELD, i32_list_value([i32::MAX, 1])), + (POSITIVE_SPAN_LENGTHS_FIELD, i32_list_value([0, 0])), + ], + "overflows i32", + ), + ( + &integer, + vec![( + CUSTOM_VALUES_FIELD, + pb_value(ValueData::ListValue(ListValue { + items: vec![null_pb_value()], + })), + )], + "contains null", + ), + ] { + assert_invalid(&modified(histogram, &changes), message); + } + let array = stored(&integer); + let mut fields = array.fields().to_vec(); + let mut children = array.columns().to_vec(); + fields.remove(0); + children.remove(0); + assert_invalid( + &StructArray::new(fields.into(), children, None), + "missing field schema", + ); + let mut fields = array.fields().to_vec(); + let mut children = array.columns().to_vec(); + children[0] = Arc::new(Float64Array::from(vec![0.0])) as ArrayRef; + fields[0] = Arc::new(Field::new( + SCHEMA_FIELD, + children[0].data_type().clone(), + true, + )); + assert_invalid( + &StructArray::new(fields.into(), children, None), + "invalid type", + ); + // Even inactive children have to retain their declared types. + let mut fields = array.fields().to_vec(); + let mut children = array.columns().to_vec(); + children[14] = Arc::new(Int32Array::from(vec![None])) as ArrayRef; + fields[14] = Arc::new(Field::new( + COUNT_F64_FIELD, + children[14].data_type().clone(), + true, + )); + assert_invalid( + &StructArray::new(fields.into(), children, None), + "invalid type", + ); + } + + #[test] + fn native_histogram_remote_read_schema_rules_match_ingestion() { + for float in [false, true] { + let base = empty_histogram(float); + for bounds in [ + vec![f64::NAN], + vec![f64::INFINITY], + vec![1.0, 1.0], + vec![2.0, 1.0], + ] { + let invalid = Histogram { + schema: CUSTOM_BUCKETS_SCHEMA, + custom_values: bounds.clone(), + ..base.clone() + }; + assert!(encode_native_histogram(&invalid).is_err()); + assert_invalid( + &modified( + &base, + &[ + ( + SCHEMA_FIELD, + pb_value(ValueData::I32Value(CUSTOM_BUCKETS_SCHEMA)), + ), + (CUSTOM_VALUES_FIELD, f64_list_value(bounds)), + ], + ), + "custom_values", + ); + } + let invalid = Histogram { + custom_values: vec![1.0], + ..base.clone() + }; + assert!(encode_native_histogram(&invalid).is_err()); + assert_invalid( + &modified(&base, &[(CUSTOM_VALUES_FIELD, f64_list_value([1.0]))]), + "must not use custom_values", + ); + let custom = Histogram { + schema: CUSTOM_BUCKETS_SCHEMA, + custom_values: vec![-f64::INFINITY, 1.0], + ..base + }; + let decoded = decode_native_histogram_for_remote_read(&stored(&custom), 0, 0) + .unwrap() + .unwrap(); + assert_eq!(decoded.custom_values, vec![-f64::INFINITY, 1.0]); + let invalid = Histogram { + zero_threshold: 1.0, + ..custom.clone() + }; + assert!(encode_native_histogram(&invalid).is_err()); + assert_invalid( + &modified( + &custom, + &[(ZERO_THRESHOLD_FIELD, pb_value(ValueData::F64Value(1.0)))], + ), + "must not use a zero bucket", + ); + let (count_field, zero_field, bucket_field, one, buckets) = if float { + ( + COUNT_F64_FIELD, + ZERO_COUNT_F64_FIELD, + NEGATIVE_BUCKETS_F64_FIELD, + pb_value(ValueData::F64Value(1.0)), + f64_list_value([1.0]), + ) + } else { + ( + COUNT_I64_FIELD, + ZERO_COUNT_I64_FIELD, + NEGATIVE_BUCKETS_I64_FIELD, + pb_value(ValueData::I64Value(1)), + i64_list_value([1]), + ) + }; + let invalid = Histogram { + count: Some(if float { + Count::CountFloat(1.0) + } else { + Count::CountInt(1) + }), + zero_count: Some(if float { + ZeroCount::ZeroCountFloat(1.0) + } else { + ZeroCount::ZeroCountInt(1) + }), + ..custom.clone() + }; + assert!(encode_native_histogram(&invalid).is_err()); + assert_invalid( + &modified( + &custom, + &[(count_field, one.clone()), (zero_field, one.clone())], + ), + "must not use a zero bucket", + ); + for length in [0, 1] { + let mut invalid = custom.clone(); + invalid.negative_spans = vec![BucketSpan { offset: 1, length }]; + let mut changes = vec![ + (NEGATIVE_SPAN_OFFSETS_FIELD, i32_list_value([1])), + (NEGATIVE_SPAN_LENGTHS_FIELD, i32_list_value([length as i32])), + ]; + if length == 1 { + invalid.count = Some(if float { + Count::CountFloat(1.0) + } else { + Count::CountInt(1) + }); + if float { + invalid.negative_counts = vec![1.0]; + } else { + invalid.negative_deltas = vec![1]; + } + changes.extend([(count_field, one.clone()), (bucket_field, buckets.clone())]); + } + assert!(encode_native_histogram(&invalid).is_err()); + assert_invalid( + &modified(&custom, &changes), + "must not use negative buckets", + ); + } + } + } + + #[test] + fn native_histogram_remote_read_span_rules_match_ingestion() { + for float in [false, true] { + let custom = Histogram { + schema: CUSTOM_BUCKETS_SCHEMA, + custom_values: vec![1.0, 2.0], + ..empty_histogram(float) + }; + for custom_values in [vec![], custom.custom_values.clone()] { + let overflow_index = custom_values.len() as i32; + for (spans, accepted) in [ + (vec![(-1, 0)], false), + (vec![(overflow_index + 2, 0)], false), + (vec![(overflow_index, 0), (2, 0)], false), + (vec![(overflow_index + 1, 0)], true), + (vec![(overflow_index, 0), (1, 0)], true), + (vec![(overflow_index, 1)], true), + (vec![(overflow_index + 1, 1)], false), + ] { + let bucket_count = spans.iter().map(|(_, length)| *length).sum::(); + let base = Histogram { + custom_values: custom_values.clone(), + positive_spans: vec![BucketSpan { + offset: 0, + length: bucket_count, + }], + positive_counts: if float { + vec![0.0; bucket_count as usize] + } else { + vec![] + }, + positive_deltas: if float { + vec![] + } else { + vec![0; bucket_count as usize] + }, + ..custom.clone() + }; + let write = Histogram { + positive_spans: spans + .iter() + .map(|&(offset, length)| BucketSpan { offset, length }) + .collect(), + ..base.clone() + }; + assert_eq!(encode_native_histogram(&write).is_ok(), accepted); + let array = modified( + &base, + &[ + ( + POSITIVE_SPAN_OFFSETS_FIELD, + i32_list_value(spans.iter().map(|s| s.0)), + ), + ( + POSITIVE_SPAN_LENGTHS_FIELD, + i32_list_value(spans.iter().map(|s| s.1 as i32)), + ), + ], + ); + assert_eq!( + decode_native_histogram_for_remote_read(&array, 0, 0).is_ok(), + accepted + ); + } + } + for negative in [false, true] { + for length in [0, 1] { + let mut base = empty_histogram(float); + let spans = vec![BucketSpan { offset: 0, length }]; + if negative { + base.negative_spans = spans; + } else { + base.positive_spans = spans; + } + if length == 1 { + if float { + if negative { + base.negative_counts = vec![0.0]; + } else { + base.positive_counts = vec![0.0]; + } + } else if negative { + base.negative_deltas = vec![0]; + } else { + base.positive_deltas = vec![0]; + } + } + for offset in [i32::MIN, i32::MIN + 1, -2, 0] { + let mut write = base.clone(); + let offset_field = if negative { + write.negative_spans[0].offset = offset; + NEGATIVE_SPAN_OFFSETS_FIELD + } else { + write.positive_spans[0].offset = offset; + POSITIVE_SPAN_OFFSETS_FIELD + }; + if offset == i32::MIN { + assert!(encode_native_histogram(&write).is_err()); + assert_invalid( + &modified(&base, &[(offset_field, i32_list_value([offset]))]), + "initial span offset i32::MIN", + ); + } else { + assert!( + decode_native_histogram_for_remote_read(&stored(&write), 0, 0) + .unwrap() + .is_some() + ); + } + } + } + } + let base = empty_histogram(float); + // Out-of-range buckets are different from valid empty end cursors. + for offset in [-1, 3] { + let custom = Histogram { + positive_spans: vec![BucketSpan { + offset: 0, + length: 1, + }], + positive_counts: if float { vec![0.0] } else { vec![] }, + positive_deltas: if float { vec![] } else { vec![0] }, + ..custom.clone() + }; + assert_invalid( + &modified( + &custom, + &[(POSITIVE_SPAN_OFFSETS_FIELD, i32_list_value([offset]))], + ), + if offset < 0 { + "negative offset" + } else { + "out of range" + }, + ); + } + let exponential = Histogram { + positive_spans: vec![BucketSpan { + offset: 0, + length: 1, + }], + positive_counts: if float { vec![0.0] } else { vec![] }, + positive_deltas: if float { vec![] } else { vec![0] }, + ..base + }; + assert_invalid( + &modified( + &exponential, + &[(POSITIVE_SPAN_OFFSETS_FIELD, i32_list_value([i32::MAX]))], + ), + "out of range", + ); + } + } + #[test] fn shared_errors_are_protocol_neutral() { let error = encode_native_histogram(&Histogram { diff --git a/src/frontend/src/instance/prom_store.rs b/src/frontend/src/instance/prom_store.rs index 7238c38f83d..94687238ed0 100644 --- a/src/frontend/src/instance/prom_store.rs +++ b/src/frontend/src/instance/prom_store.rs @@ -850,6 +850,24 @@ mod tests { use super::*; + #[test] + fn test_native_histogram_remote_read_response_negotiation() { + for accepted in [ + vec![], + vec![ResponseType::Samples as i32], + vec![ + ResponseType::StreamedXorChunks as i32, + ResponseType::Samples as i32, + ], + ] { + assert_eq!( + negotiate_response_type(&accepted).unwrap(), + ResponseType::Samples + ); + } + assert!(negotiate_response_type(&[ResponseType::StreamedXorChunks as i32]).is_err()); + } + #[test] fn test_remote_read_label_types() { use datatypes::schema::ColumnSchema; @@ -864,6 +882,10 @@ mod tests { (ConcreteDataType::vector_datatype(3), false), (ConcreteDataType::int64_datatype(), false), (ConcreteDataType::boolean_datatype(), false), + ( + common_query::native_histogram::native_histogram_value_type().clone(), + false, + ), ] { let schema = Schema::new(vec![ ColumnSchema::new( diff --git a/src/servers/Cargo.toml b/src/servers/Cargo.toml index 1f433318938..4383896e1f8 100644 --- a/src/servers/Cargo.toml +++ b/src/servers/Cargo.toml @@ -155,6 +155,7 @@ catalog = { workspace = true, features = ["testing"] } client = { workspace = true, features = ["testing"] } common-base.workspace = true common-meta = { workspace = true, features = ["testing"] } +common-query = { workspace = true, features = ["testing"] } common-test-util.workspace = true criterion = "0.5" json5 = "0.4" diff --git a/src/servers/src/prom_store.rs b/src/servers/src/prom_store.rs index b991cc9c524..d45cabceaf7 100644 --- a/src/servers/src/prom_store.rs +++ b/src/servers/src/prom_store.rs @@ -33,6 +33,7 @@ use arrow::datatypes::{ }; use common_grpc::precision::Precision; use common_memory_manager::MemoryGuard; +use common_query::native_histogram; use common_query::prelude::{greptime_timestamp, greptime_value}; use common_recordbatch::{RecordBatch, RecordBatches}; use common_telemetry::{tracing, warn}; @@ -458,6 +459,7 @@ pub fn recordbatches_to_timeseries( for ts in &mut timeseries { ts.samples.sort_unstable_by_key(|s| s.timestamp); + ts.histograms.sort_unstable_by_key(|h| h.timestamp); } timeseries @@ -531,21 +533,26 @@ fn append_recordbatch_to_timeseries( ), })?; - // TODO: Add native-histogram encoding when Prometheus Remote Read support is prioritized. - // The current path intentionally returns scalar samples only. let field_column = recordbatch .column_by_name(value_column_name) .with_context(|| error::InvalidPromRemoteReadQueryResultSnafu { msg: format!("missing value column '{value_column_name}' in query result"), })?; - let field_column = field_column - .as_primitive_opt::() - .with_context(|| error::InvalidPromRemoteReadQueryResultSnafu { + let scalar_column = field_column.as_primitive_opt::(); + let histogram_column = recordbatch + .schema + .column_schema_by_name(value_column_name) + .filter(|column| native_histogram::is_native_histogram_value_type(&column.data_type)) + .and_then(|_| field_column.as_struct_opt()); + ensure!( + scalar_column.is_some() || histogram_column.is_some(), + error::InvalidPromRemoteReadQueryResultSnafu { msg: format!( - "Expect value column of datatype Float64, actual {:?}", + "Expect value column of datatype Float64 or canonical native histogram Struct, actual {:?}", field_column.data_type() ), - })?; + } + ); let columns = label_columns(&recordbatch, timestamp_column_name, value_column_name)?; let mut previous_timeseries: Option = None; @@ -577,11 +584,25 @@ fn append_recordbatch_to_timeseries( continue; } - let value = field_column.value(row); let timestamp = ts_column.value(row); - let sample = Sample { value, timestamp }; - - timeseries[timeseries_index].samples.push(sample); + if let Some(column) = scalar_column { + timeseries[timeseries_index].samples.push(Sample { + value: column.value(row), + timestamp, + }); + } else if let Some(column) = histogram_column { + let histogram = + native_histogram::decode_native_histogram_for_remote_read(column, row, timestamp) + .map_err(|error| { + error::InvalidPromRemoteReadQueryResultSnafu { + msg: error.to_string(), + } + .build() + })?; + if let Some(histogram) = histogram { + timeseries[timeseries_index].histograms.push(histogram); + } + } } Ok(()) @@ -1540,6 +1561,305 @@ mod tests { ); } + #[test] + fn test_native_histogram_remote_read_wire_fixture() { + use api::greptime_proto::io::prometheus::write::v2::{ + BucketSpan, Histogram, + histogram::{Count, ZeroCount}, + }; + use api::prom_store::remote::{QueryResult, ReadResponse}; + use common_query::native_histogram::testing::histogram_array; + use prost::Message; + + let histograms = vec![ + Histogram { + count: Some(Count::CountInt((1 << 53) + 13)), + zero_count: Some(ZeroCount::ZeroCountInt((1 << 53) + 1)), + sum: 10.0, + schema: 1, + zero_threshold: 0.001, + reset_hint: 2, + start_timestamp: 123, + positive_spans: vec![ + BucketSpan { + offset: -2, + length: 2, + }, + BucketSpan { + offset: 3, + length: 1, + }, + ], + positive_deltas: vec![1, 2, -1], + negative_spans: vec![ + BucketSpan { + offset: -3, + length: 1, + }, + BucketSpan { + offset: 2, + length: 2, + }, + ], + negative_deltas: vec![1, 2, -1], + ..Default::default() + }, + Histogram { + count: Some(Count::CountFloat(6.0)), + zero_count: Some(ZeroCount::ZeroCountFloat(0.5)), + sum: 20.0, + schema: 2, + zero_threshold: 0.002, + reset_hint: 3, + positive_spans: vec![BucketSpan { + offset: 3, + length: 2, + }], + positive_counts: vec![2.0, 3.5], + ..Default::default() + }, + Histogram { + count: Some(Count::CountInt(1)), + sum: 1.0, + schema: -53, + custom_values: vec![1.0, 2.0], + positive_spans: vec![BucketSpan { + offset: 0, + length: 1, + }], + positive_deltas: vec![1], + ..Default::default() + }, + Histogram { + count: Some(Count::CountFloat(3.5)), + sum: 4.0, + schema: -53, + custom_values: vec![1.0, 2.0], + reset_hint: 3, + positive_spans: vec![BucketSpan { + offset: 0, + length: 2, + }], + positive_counts: vec![1.0, 2.5], + ..Default::default() + }, + Histogram { + count: Some(Count::CountInt(0)), + ..Default::default() + }, + Histogram { + count: Some(Count::CountFloat(0.0)), + ..Default::default() + }, + Histogram { + count: Some(Count::CountInt(8)), + sum: f64::from_bits(0x7ff8_0000_0000_0123), + positive_spans: vec![BucketSpan { + offset: 0, + length: 2, + }], + positive_deltas: vec![3, -1], + ..Default::default() + }, + Histogram { + count: Some(Count::CountInt(0)), + sum: f64::from_bits(0x7ff0_0000_0000_0002), + ..Default::default() + }, + Histogram { + count: Some(Count::CountInt(i64::MAX as u64)), + zero_count: Some(ZeroCount::ZeroCountInt(i64::MAX as u64)), + ..Default::default() + }, + ]; + let values = histograms + .iter() + .map(|h| Some(native_histogram::encode_native_histogram(h).unwrap())) + .collect::>(); + let array = histogram_array(&values); + let buckets = array + .column_by_name("positive_buckets_i64") + .unwrap() + .as_list::() + .value(0); + assert_eq!( + buckets + .as_primitive::() + .values() + .as_ref(), + &[1, 3, 2] + ); + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + Field::new( + "ts", + ArrowDataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new("histogram", array.data_type().clone(), true), + ])); + let schema = Arc::new(Schema::try_from(arrow_schema.clone()).unwrap()); + let batch = DfRecordBatch::try_new( + arrow_schema, + vec![ + Arc::new(TimestampMillisecondArray::from( + (1..=9).map(|i| i * 1000).collect::>(), + )), + Arc::new(array), + ], + ) + .unwrap(); + let batches = RecordBatches::try_new( + schema.clone(), + vec![RecordBatch::from_df_record_batch(schema, batch)], + ) + .unwrap(); + let response = ReadResponse { + results: vec![QueryResult { + timeseries: recordbatches_to_timeseries( + "native_histogram_remote_read", + "ts", + "histogram", + batches, + ) + .unwrap(), + }], + }; + let body = snappy_compress(&response.encode_to_vec()).unwrap(); + let mut decoded = + ReadResponse::decode(snappy_decompress(&body).unwrap().as_slice()).unwrap(); + let mut expected = ReadResponse::decode( + include_bytes!( + "../../../tests/compatibility/wire/native_histogram_remote_read/response.pb" + ) + .as_slice(), + ) + .unwrap(); + let actual_histograms = &mut decoded.results[0].timeseries[0].histograms; + assert_eq!(actual_histograms[0].positive_deltas, vec![1, 2, -1]); + assert_eq!(actual_histograms[0].negative_deltas, vec![1, 2, -1]); + assert_eq!(actual_histograms[6].positive_deltas, vec![3, -1]); + for (actual, expected) in actual_histograms + .iter_mut() + .zip(&mut expected.results[0].timeseries[0].histograms) + { + assert_eq!(actual.sum.to_bits(), expected.sum.to_bits()); + actual.sum = 0.0; + expected.sum = 0.0; + } + assert_eq!(decoded, expected); + // The optional output is consumed by the pinned Prometheus iterator check. + if let Some(path) = std::env::var_os("GREPTIME_NATIVE_HISTOGRAM_REMOTE_READ_OUTPUT") { + std::fs::write(path, body).unwrap(); + } + } + + #[test] + fn test_native_histogram_remote_read_groups_and_sorts_batches() { + use api::greptime_proto::io::prometheus::write::v2::{Histogram, histogram::Count}; + use common_query::native_histogram::testing::histogram_array; + + let integer = native_histogram::encode_native_histogram(&Histogram { + count: Some(Count::CountInt(0)), + sum: f64::from_bits(0x7ff0_0000_0000_0002), + ..Default::default() + }) + .unwrap(); + let float = native_histogram::encode_native_histogram(&Histogram { + count: Some(Count::CountFloat(0.0)), + ..Default::default() + }) + .unwrap(); + let histograms = histogram_array(&[Some(integer), Some(float.clone()), None, Some(float)]); + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + Field::new( + "event_time", + ArrowDataType::Timestamp(TimeUnit::Microsecond, None), + true, + ), + Field::new("legacy_histogram", histograms.data_type().clone(), true), + Field::new_dictionary("instance", ArrowDataType::UInt32, ArrowDataType::Utf8, true), + ])); + let schema = Arc::new(Schema::try_from(arrow_schema.clone()).unwrap()); + let batch = |timestamps: Vec>, labels: Vec>| { + RecordBatch::from_df_record_batch( + schema.clone(), + DfRecordBatch::try_new( + arrow_schema.clone(), + vec![ + Arc::new(arrow::array::TimestampMicrosecondArray::from(timestamps)), + Arc::new(histograms.clone()), + Arc::new(DictionaryArray::::new( + UInt32Array::from(labels), + Arc::new(StringArray::from(vec![Some("b"), Some("a"), None])), + )), + ], + ) + .unwrap(), + ) + }; + let batches = vec![ + batch( + vec![Some(4000), Some(3000), Some(5000), Some(-1001)], + vec![Some(0), None, Some(0), Some(0)], + ), + batch( + vec![Some(2000), Some(1000), Some(6000), None], + vec![Some(0), Some(2), Some(1), Some(0)], + ), + ]; + let series = recordbatches_to_timeseries( + "latency", + "event_time", + "legacy_histogram", + RecordBatches::try_new(schema.clone(), batches).unwrap(), + ) + .unwrap(); + assert_eq!(series.len(), 3); + assert_eq!( + series[0].labels, + vec![new_label( + METRIC_NAME_LABEL.to_string(), + "latency".to_string() + )] + ); + assert_eq!( + series[0] + .histograms + .iter() + .map(|h| h.timestamp) + .collect::>(), + vec![1, 3] + ); + assert_eq!(series[1].labels[1].value, "a"); + assert!(series[1].histograms.is_empty()); + assert_eq!(series[2].labels[1].value, "b"); + assert_eq!( + series[2] + .histograms + .iter() + .map(|h| h.timestamp) + .collect::>(), + vec![-2, 2, 4] + ); + assert!(matches!( + series[2].histograms[0].count, + Some(api::prom_store::remote::histogram::Count::CountFloat(0.0)) + )); + for histogram in &series[2].histograms[1..] { + assert!(matches!( + histogram.count, + Some(api::prom_store::remote::histogram::Count::CountInt(0)) + )); + assert_eq!(histogram.sum.to_bits(), 0x7ff0_0000_0000_0002); + } + assert!( + series + .iter() + .all(|s| s.samples.is_empty() + && s.labels.iter().all(|l| l.name != "legacy_histogram")) + ); + } + #[test] fn test_recordbatches_to_timeseries() { let schema = Arc::new(Schema::new(vec![ diff --git a/src/servers/tests/http/prom_store_test.rs b/src/servers/tests/http/prom_store_test.rs index c7a288ccee0..e4ef6219b30 100644 --- a/src/servers/tests/http/prom_store_test.rs +++ b/src/servers/tests/http/prom_store_test.rs @@ -16,8 +16,8 @@ use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; use api::greptime_proto::io::prometheus::write::v2::{ - Request as RemoteWriteV2Request, Sample as RemoteWriteV2Sample, - TimeSeries as RemoteWriteV2TimeSeries, + BucketSpan as RemoteWriteV2BucketSpan, Request as RemoteWriteV2Request, + Sample as RemoteWriteV2Sample, TimeSeries as RemoteWriteV2TimeSeries, }; use api::prom_store::remote::{ LabelMatcher, Query, QueryResult, ReadRequest, ReadResponse, WriteRequest, @@ -28,6 +28,7 @@ use async_trait::async_trait; use axum::Router; use axum::http::HeaderMap; use common_query::Output; +use common_query::native_histogram::CUSTOM_BUCKETS_SCHEMA; use common_query::prelude::{ GREPTIME_PHYSICAL_TABLE, greptime_native_histogram, greptime_timestamp, greptime_value, }; @@ -317,6 +318,59 @@ async fn test_prometheus_remote_write_v2_convert_error_has_written_headers() { assert!(write_rx.try_recv().is_err()); } +#[tokio::test] +async fn test_prometheus_remote_write_v2_rejects_incompatible_histogram_before_writes() { + common_telemetry::init_default_ut_logging(); + let (read_tx, _read_rx) = mpsc::channel(100); + let (write_tx, mut write_rx) = mpsc::channel(100); + let app = make_test_app_with_write_capture(read_tx, write_tx); + let client = TestClient::new(app).await; + + for (schema, offset, message) in [ + (0, i32::MIN, "initial span offset i32::MIN"), + ( + CUSTOM_BUCKETS_SCHEMA, + 2, + "span extent exceeds custom_values", + ), + ] { + let mut histogram = remote_write_v2::histogram(2000); + histogram.schema = schema; + histogram.positive_spans = vec![RemoteWriteV2BucketSpan { offset, length: 0 }]; + let write_request = RemoteWriteV2Request { + symbols: vec![ + String::new(), + prom_store::METRIC_NAME_LABEL.to_string(), + "http_requests_total".to_string(), + "http_request_duration_seconds".to_string(), + ], + timeseries: vec![ + RemoteWriteV2TimeSeries { + labels_refs: vec![1, 2], + samples: vec![RemoteWriteV2Sample { + value: 42.0, + timestamp: 1000, + start_timestamp: 0, + }], + ..Default::default() + }, + RemoteWriteV2TimeSeries { + labels_refs: vec![1, 3], + histograms: vec![remote_write_v2::histogram(1000), histogram], + ..Default::default() + }, + ], + }; + + let result = post_remote_write_v2(&client, &write_request).await; + + assert_eq!(result.status(), 400); + assert_remote_write_v2_written_headers(&result.headers(), "0"); + assert!(result.text().await.contains(message)); + assert!(write_rx.try_recv().is_err()); + } +} + #[tokio::test] async fn test_prometheus_remote_write_v2_write_error_has_partial_written_headers() { common_telemetry::init_default_ut_logging(); diff --git a/tests-integration/src/otlp.rs b/tests-integration/src/otlp.rs index 4998e29e0c8..5ab264d23d0 100644 --- a/tests-integration/src/otlp.rs +++ b/tests-integration/src/otlp.rs @@ -2838,6 +2838,32 @@ WITH( .remove(0) .unwrap(); } + let series = + native_histogram_remote_read(&instance, "delta_distribution", ctx.clone()).await; + assert_eq!(series.len(), 1); + assert!( + series[0] + .labels + .iter() + .any(|l| l.name == "otlp_aggregation_temporality" && l.value == "delta") + ); + assert_eq!(series[0].histograms.len(), 3); + for (histogram, (timestamp, count)) in + series[0] + .histograms + .iter() + .zip([(60_000, 10), (120_000, 20), (180_000, 15)]) + { + assert_eq!(histogram.timestamp, timestamp); + assert_eq!( + histogram.count, + Some(api::prom_store::remote::histogram::Count::CountInt(count)) + ); + assert_eq!(histogram.positive_deltas, vec![count as i64]); + assert_eq!(histogram.positive_spans[0].offset, 1); + assert_eq!(histogram.sum, count as f64 * 1.5); + assert_eq!(histogram.reset_hint, 3); + } for (expression, expected) in [ ( "histogram_count(sum_over_time(delta_distribution[3m]))", @@ -3139,6 +3165,42 @@ WITH( } } + async fn native_histogram_remote_read( + instance: &Arc, + table: &str, + ctx: Arc, + ) -> Vec { + use api::prom_store::remote::{LabelMatcher, Query, ReadRequest, ReadResponse}; + use prost::Message; + use servers::query_handler::PromStoreProtocolHandler; + + let response = instance + .read( + ReadRequest { + queries: vec![Query { + start_timestamp_ms: 0, + end_timestamp_ms: 200_000, + matchers: vec![LabelMatcher { + name: "__name__".to_string(), + value: table.to_string(), + r#type: 0, + }], + ..Default::default() + }], + ..Default::default() + }, + ctx, + ) + .await + .unwrap(); + let body = servers::prom_store::snappy_decompress(&response.body).unwrap(); + ReadResponse::decode(body.as_slice()) + .unwrap() + .results + .remove(0) + .timeseries + } + #[tokio::test(flavor = "multi_thread")] async fn test_otlp_delta_exponential_storage_precision_and_replacement() { use datatypes::arrow::datatypes::Int64Type; @@ -3194,6 +3256,24 @@ WITH( .remove(0) .unwrap(); } + let series = + native_histogram_remote_read(instance, "interval", ctx.clone()).await; + assert_eq!(series.len(), 1); + assert!(series[0].labels.iter().any(|l| l.name == "otlp_aggregation_temporality" && l.value == "delta")); + let histogram = &series[0].histograms[0]; + assert_eq!(histogram.timestamp, 60_000); + assert_eq!( + histogram.count, + Some(api::prom_store::remote::histogram::Count::CountInt(count)) + ); + assert_eq!( + histogram.zero_count, + Some(api::prom_store::remote::histogram::ZeroCount::ZeroCountInt( + count + )) + ); + assert_eq!(histogram.reset_hint, 3); + assert_eq!(histogram.sum, 0.0); let output = instance .do_query( "SELECT greptime_timestamp,greptime_native_histogram FROM interval", diff --git a/tests-integration/src/prom_store.rs b/tests-integration/src/prom_store.rs index 87798c282c1..0608ff88f9c 100644 --- a/tests-integration/src/prom_store.rs +++ b/tests-integration/src/prom_store.rs @@ -778,6 +778,329 @@ mod tests { } } + #[tokio::test(flavor = "multi_thread")] + async fn test_standalone_native_histogram_remote_read() { + let standalone = + GreptimeDbStandaloneBuilder::new("standalone_native_histogram_remote_read") + .build() + .await; + check_native_histogram_remote_read(standalone.fe_instance()).await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_distributed_native_histogram_remote_read() { + let cluster = + tests::create_distributed_instance("distributed_native_histogram_remote_read").await; + check_native_histogram_remote_read(&cluster.frontend()).await; + } + + async fn check_native_histogram_remote_read(instance: &Arc) { + use api::greptime_proto::io::prometheus::write::v2::{ + Histogram, + histogram::{Count, ZeroCount}, + }; + use common_query::native_histogram::{ + encode_native_histogram, native_histogram_column_schema, + }; + use servers::prom_remote_write::v2::test_util as remote_write_v2; + + let ctx = Arc::new(QueryContext::with(DEFAULT_CATALOG_NAME, "prometheus")); + instance + .do_query("CREATE DATABASE prometheus", ctx.clone()) + .await + .remove(0) + .unwrap(); + let server = HttpServerBuilder::new(HttpOptions::default()) + .with_prom_handler(instance.clone(), None, true, Default::default(), None) + .build(); + let client = TestClient::new(server.build(server.make_app()).unwrap()).await; + let large = (1u64 << 53) + 1; + let histograms = vec![ + Histogram { + count: Some(Count::CountInt(6)), + timestamp: 1000, + positive_spans: vec![api::greptime_proto::io::prometheus::write::v2::BucketSpan { + offset: 0, + length: 3, + }], + positive_deltas: vec![1, 2, -1], + sum: 10.0, + ..Default::default() + }, + Histogram { + count: Some(Count::CountInt(large)), + zero_count: Some(ZeroCount::ZeroCountInt(large)), + timestamp: 2000, + ..Default::default() + }, + Histogram { + count: Some(Count::CountFloat(6.0)), + zero_count: Some(ZeroCount::ZeroCountFloat(0.5)), + positive_spans: vec![api::greptime_proto::io::prometheus::write::v2::BucketSpan { + offset: 3, + length: 2, + }], + positive_counts: vec![2.0, 3.5], + timestamp: 3000, + ..Default::default() + }, + Histogram { + count: Some(Count::CountInt(0)), + sum: f64::from_bits(0x7ff0_0000_0000_0002), + timestamp: 4000, + ..Default::default() + }, + ]; + let request = remote_write_v2::request_with_labels_and_histograms( + vec![(prom_store::METRIC_NAME_LABEL, "latency"), ("job", "api")], + histograms, + ); + let response = client + .post("/v1/prometheus/write?db=prometheus") + .header( + "content-type", + "application/x-protobuf;proto=io.prometheus.write.v2.Request", + ) + .header("content-encoding", "snappy") + .body(prom_store::snappy_compress(&request.encode_to_vec()).unwrap()) + .send() + .await; + assert_eq!(response.status(), 204, "{}", response.text().await); + let (scalar, _) = to_grpc_row_insert_requests(&WriteRequest { + timeseries: vec![prom_store::mock_timeseries()[0].clone()], + ..Default::default() + }) + .unwrap(); + instance.write(scalar, ctx.clone(), true).await.unwrap(); + + for flushed in [false, true] { + if flushed { + instance + .do_query("ADMIN FLUSH_TABLE('latency')", ctx.clone()) + .await + .remove(0) + .unwrap(); + } + for accepted_response_types in [vec![], vec![0], vec![1, 0]] { + let mut query = remote_query("latency"); + query.end_timestamp_ms = 4000; + query.matchers.push(LabelMatcher { + name: "job".to_string(), + value: "api".to_string(), + r#type: MatcherType::Eq as i32, + }); + let mut empty = query.clone(); + empty.start_timestamp_ms = 5000; + empty.end_timestamp_ms = 6000; + let request = ReadRequest { + queries: vec![ + remote_query("metric1"), + query, + remote_query("missing_metric"), + empty, + ], + accepted_response_types, + }; + let response = client + .post("/v1/prometheus/read?db=prometheus") + .body(prom_store::snappy_compress(&request.encode_to_vec()).unwrap()) + .send() + .await; + assert_eq!(response.status(), 200, "{}", response.text().await); + assert_eq!(response.headers()["content-type"], "application/x-protobuf"); + assert_eq!(response.headers()["content-encoding"], "snappy"); + let body = prom_store::snappy_decompress(&response.bytes().await).unwrap(); + let response = ReadResponse::decode(body.as_slice()).unwrap(); + assert_eq!(response.results.len(), 4); + assert!(!response.results[0].timeseries[0].samples.is_empty()); + assert!(response.results[2].timeseries.is_empty()); + assert!(response.results[3].timeseries.is_empty()); + let series = &response.results[1].timeseries[0]; + assert_eq!( + series.labels, + vec![ + Label { + name: "__name__".to_string(), + value: "latency".to_string() + }, + Label { + name: "job".to_string(), + value: "api".to_string() + } + ] + ); + assert_eq!(series.histograms.len(), 4); + assert_eq!(series.histograms[0].positive_deltas, vec![1, 2, -1]); + assert_eq!(series.histograms[2].positive_counts, vec![2.0, 3.5]); + assert_eq!( + series.histograms[1].count, + Some(api::prom_store::remote::histogram::Count::CountInt(large)) + ); + assert_eq!( + series.histograms[1].zero_count, + Some(api::prom_store::remote::histogram::ZeroCount::ZeroCountInt( + large + )) + ); + assert_eq!( + series.histograms[2].count, + Some(api::prom_store::remote::histogram::Count::CountFloat(6.0)) + ); + assert_eq!(series.histograms[3].sum.to_bits(), 0x7ff0_0000_0000_0002); + } + } + let response = client + .post("/v1/prometheus/read?db=prometheus") + .body( + prom_store::snappy_compress( + &ReadRequest { + queries: vec![remote_query("latency")], + accepted_response_types: vec![1], + } + .encode_to_vec(), + ) + .unwrap(), + ) + .send() + .await; + assert!(!response.status().is_success()); + assert!( + response + .text() + .await + .contains("does not support any of the requested response types") + ); + + // Typed inserts supply sub-millisecond pre-epoch times that neither OTLP + // nor millisecond Remote Write can express, and use a legacy field name. + for (table, datatype, timestamp) in [ + ( + "legacy_us", + ColumnDataType::TimestampMicrosecond, + ValueData::TimestampMicrosecondValue(-1001), + ), + ( + "legacy_ns", + ColumnDataType::TimestampNanosecond, + ValueData::TimestampNanosecondValue(-1_001_000), + ), + ] { + let mut histogram_schema = native_histogram_column_schema().unwrap(); + histogram_schema.column_name = "payload".to_string(); + let histogram = encode_native_histogram(&Histogram { + count: Some(Count::CountInt(large)), + zero_count: Some(ZeroCount::ZeroCountInt(large)), + ..Default::default() + }) + .unwrap(); + instance + .write( + RowInsertRequests { + inserts: vec![RowInsertRequest { + table_name: table.to_string(), + rows: Some(Rows { + schema: vec![ + api::v1::ColumnSchema { + column_name: "event_time".to_string(), + datatype: datatype as i32, + semantic_type: SemanticType::Timestamp as i32, + ..Default::default() + }, + histogram_schema, + ], + rows: vec![Row { + values: vec![ + Value { + value_data: Some(timestamp), + }, + Value { + value_data: Some(histogram), + }, + ], + }], + }), + }], + }, + ctx.clone(), + false, + ) + .await + .unwrap(); + for flushed in [false, true] { + if flushed { + instance + .do_query(&format!("ADMIN FLUSH_TABLE('{table}')"), ctx.clone()) + .await + .remove(0) + .unwrap(); + } + let mut query = remote_query(table); + query.start_timestamp_ms = -2; + query.end_timestamp_ms = 0; + let response = post_remote_read(&client, vec![query]).await; + assert_eq!(response.status(), 200, "{}", response.text().await); + let body = prom_store::snappy_decompress(&response.bytes().await).unwrap(); + let response = ReadResponse::decode(body.as_slice()).unwrap(); + let series = &response.results[0].timeseries[0]; + assert_eq!(series.labels.len(), 1); + assert_eq!(series.histograms[0].timestamp, -2); + assert_eq!( + series.histograms[0].count, + Some(api::prom_store::remote::histogram::Count::CountInt(large)) + ); + } + } + // Bypass ingestion's histogram codec to exercise whole-request failure. + let ValueData::StructValue(mut malformed) = + encode_native_histogram(&Histogram::default()).unwrap() + else { + unreachable!() + }; + malformed.items[14].value_data = Some(ValueData::F64Value(0.0)); + instance + .write( + RowInsertRequests { + inserts: vec![RowInsertRequest { + table_name: "malformed_histogram".to_string(), + rows: Some(Rows { + schema: vec![ + api::v1::ColumnSchema { + column_name: "ts".to_string(), + datatype: ColumnDataType::TimestampMillisecond as i32, + semantic_type: SemanticType::Timestamp as i32, + ..Default::default() + }, + native_histogram_column_schema().unwrap(), + ], + rows: vec![Row { + values: vec![ + Value { + value_data: Some(ValueData::TimestampMillisecondValue( + 1000, + )), + }, + Value { + value_data: Some(ValueData::StructValue(malformed)), + }, + ], + }], + }), + }], + }, + ctx.clone(), + false, + ) + .await + .unwrap(); + for queries in [ + vec![remote_query("metric1"), remote_query("malformed_histogram")], + vec![remote_query("malformed_histogram"), remote_query("metric1")], + ] { + let response = post_remote_read(&client, queries).await; + assert_eq!(response.status(), 500); + } + } + fn remote_query(metric: &str) -> Query { Query { start_timestamp_ms: 1000, diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 1671f5a2921..11c41431fbc 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -4262,6 +4262,27 @@ pub async fn test_prometheus_remote_write_v2_native_histogram(store_type: Storag ], ..Default::default() }); + let large = (1u64 << 53) + 1; + write_request.timeseries[1].histograms.extend([ + Histogram { + count: Some(Count::CountInt(large)), + zero_count: Some(ZeroCount::ZeroCountInt(large)), + timestamp: 5000, + ..Default::default() + }, + Histogram { + count: Some(Count::CountInt(0)), + sum: f64::from_bits(0x7ff0_0000_0000_0002), + timestamp: 6000, + ..Default::default() + }, + Histogram { + count: Some(Count::CountFloat(0.0)), + sum: f64::from_bits(0x7ff8_0000_0000_0123), + timestamp: 7000, + ..Default::default() + }, + ]); let serialized_request = write_request.encode_to_vec(); let compressed_request = prom_store::snappy_compress(&serialized_request).expect("failed to encode snappy"); @@ -4278,7 +4299,7 @@ pub async fn test_prometheus_remote_write_v2_native_histogram(store_type: Storag .await; assert_eq!(res.status(), StatusCode::NO_CONTENT); let headers = res.headers(); - assert_remote_write_v2_written_headers_with_histograms(&headers, "1", "2"); + assert_remote_write_v2_written_headers_with_histograms(&headers, "1", "5"); assert!(res.text().await.is_empty()); validate_data( @@ -4289,7 +4310,7 @@ pub async fn test_prometheus_remote_write_v2_native_histogram(store_type: Storag ) .await; - let histogram_sql = "select greptime_timestamp, greptime_native_histogram, job, instance from remote_write_v2_latency_seconds order by greptime_timestamp;"; + let histogram_sql = "select greptime_timestamp, greptime_native_histogram, job, instance from remote_write_v2_latency_seconds where greptime_timestamp < 5000::TIMESTAMP order by greptime_timestamp;"; let histogram_rows = "[[3000,{\"count_f64\":null,\"count_i64\":8,\"custom_values\":[],\"negative_buckets_f64\":[],\"negative_buckets_i64\":[1],\"negative_span_lengths\":[1],\"negative_span_offsets\":[-2],\"positive_buckets_f64\":[],\"positive_buckets_i64\":[1,3,2],\"positive_span_lengths\":[3],\"positive_span_offsets\":[0],\"reset_hint\":2,\"schema\":1,\"start_timestamp\":1500,\"sum\":10.0,\"zero_count_f64\":null,\"zero_count_i64\":1,\"zero_threshold\":0.001},\"api\",\"localhost:9090\"],[4000,{\"count_f64\":6.0,\"count_i64\":null,\"custom_values\":[],\"negative_buckets_f64\":[],\"negative_buckets_i64\":[],\"negative_span_lengths\":[],\"negative_span_offsets\":[],\"positive_buckets_f64\":[2.0,3.5],\"positive_buckets_i64\":[],\"positive_span_lengths\":[2],\"positive_span_offsets\":[3],\"reset_hint\":3,\"schema\":2,\"start_timestamp\":2500,\"sum\":20.0,\"zero_count_f64\":0.5,\"zero_count_i64\":null,\"zero_threshold\":0.002},\"api\",\"localhost:9090\"]]"; validate_data( "prometheus_remote_write_v2_native_histogram_rows", @@ -4299,6 +4320,76 @@ pub async fn test_prometheus_remote_write_v2_native_histogram(store_type: Storag ) .await; + for flushed in [false, true] { + if flushed { + flush_table(&client, "remote_write_v2_latency_seconds").await; + } + for accepted_response_types in [vec![], vec![0], vec![1, 0]] { + let query = |name: &str| Query { + start_timestamp_ms: 0, + end_timestamp_ms: 10_000, + matchers: vec![LabelMatcher { + name: prom_store::METRIC_NAME_LABEL.to_string(), + value: name.to_string(), + r#type: 0, + }], + ..Default::default() + }; + let request = ReadRequest { + queries: vec![ + query("remote_write_v2_sample_seconds"), + query("missing_native_histogram"), + query("remote_write_v2_latency_seconds"), + ], + accepted_response_types, + }; + let response = client + .post("/v1/prometheus/read") + .body(prom_store::snappy_compress(&request.encode_to_vec()).unwrap()) + .send() + .await; + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(response.headers()["content-type"], "application/x-protobuf"); + assert_eq!(response.headers()["content-encoding"], "snappy"); + let body = prom_store::snappy_decompress(&response.bytes().await).unwrap(); + let response = ReadResponse::decode(body.as_slice()).unwrap(); + assert_eq!(response.results.len(), 3); + assert_eq!( + response.results[0].timeseries[0].samples, + vec![Sample { + value: 5.0, + timestamp: 2500 + }] + ); + assert!(response.results[1].timeseries.is_empty()); + let series = &response.results[2].timeseries[0]; + assert!(series.samples.is_empty()); + assert_eq!( + series + .histograms + .iter() + .map(|h| h.timestamp) + .collect::>(), + vec![3000, 4000, 5000, 6000, 7000] + ); + assert_eq!(series.histograms[0].positive_deltas, vec![1, 2, -1]); + assert_eq!(series.histograms[0].negative_deltas, vec![1]); + assert_eq!(series.histograms[1].positive_counts, vec![2.0, 3.5]); + assert_eq!( + series.histograms[2].count, + Some(api::prom_store::remote::histogram::Count::CountInt(large)) + ); + assert_eq!( + series.histograms[2].zero_count, + Some(api::prom_store::remote::histogram::ZeroCount::ZeroCountInt( + large + )) + ); + assert_eq!(series.histograms[3].sum.to_bits(), 0x7ff0_0000_0000_0002); + assert_eq!(series.histograms[4].sum.to_bits(), 0x7ff8_0000_0000_0123); + } + } + // The native histogram column is a struct of nested lists, so its SST // encoding differs from its memtable representation. flush_table(&client, "remote_write_v2_latency_seconds").await; @@ -9133,9 +9224,17 @@ pub async fn test_otlp_exponential_histogram(store_type: StorageType) { )] }; - let (app, mut guard) = - setup_test_http_app_with_frontend(store_type, "test_otlp_exponential_histogram").await; - let client = TestClient::new(app).await; + let standalone = tests_integration::standalone::GreptimeDbStandaloneBuilder::new( + "test_otlp_exponential_histogram", + ) + .with_default_store_type(store_type) + .build() + .await; + let server = build_test_prom_server(standalone.fe_instance().clone(), false, None) + .with_otlp_handler(standalone.fe_instance().clone(), true) + .build(); + let client = TestClient::new(server.build(server.make_app()).unwrap()).await; + let mut guard = standalone.guard; let res = send_req(&client, headers(), "/v1/otlp/v1/metrics", body, false).await; assert_eq!(StatusCode::OK, res.status()); let response = ExportMetricsServiceResponse::decode(res.bytes().await).unwrap(); @@ -9158,6 +9257,42 @@ pub async fn test_otlp_exponential_histogram(store_type: StorageType) { ) .await; + for flushed in [false, true] { + if flushed { + flush_table(&client, "otlp_exponential_latency").await; + } + let response = post_remote_read( + &client, + vec![Query { + start_timestamp_ms: 0, + end_timestamp_ms: 4000, + matchers: vec![LabelMatcher { + name: "__name__".to_string(), + value: "otlp_exponential_latency".to_string(), + r#type: 0, + }], + ..Default::default() + }], + ) + .await; + let histogram = &response.results[0].timeseries[0].histograms[0]; + assert_eq!(histogram.timestamp, 3000); + assert_eq!( + histogram.count, + Some(api::prom_store::remote::histogram::Count::CountInt(4)) + ); + assert_eq!( + histogram.zero_count, + Some(api::prom_store::remote::histogram::ZeroCount::ZeroCountInt( + 1 + )) + ); + assert_eq!(histogram.positive_deltas, vec![1, 1]); + assert_eq!(histogram.positive_spans[0].offset, 0); + assert_eq!(histogram.reset_hint, 0); + assert_eq!(histogram.sum, 8.0); + } + for query in [ "histogram_count(otlp_exponential_latency)", "histogram_sum(otlp_exponential_latency)", diff --git a/tests/compatibility/README.md b/tests/compatibility/README.md index 1daf42cc736..53ecda81048 100644 --- a/tests/compatibility/README.md +++ b/tests/compatibility/README.md @@ -15,6 +15,14 @@ read-preference defaulting and unrelated metadata preservation: cargo nextest run -p session -E 'test(test_legacy_query_context_wire_defaults_to_leader)' ``` +The [native histogram Remote Read fixture](wire/native_histogram_remote_read/README.md) +checks the sampled encoder against independent Prometheus protobuf values, +including exact integer counts, custom bounds, and NaN payloads: + +```shell +cargo nextest run -p servers -E 'test(native_histogram_remote_read_wire_fixture)' +``` + ## Object store WAL fixtures Frozen objects of the object store WAL are checked byte for byte by the format tests of diff --git a/tests/compatibility/wire/native_histogram_remote_read/README.md b/tests/compatibility/wire/native_histogram_remote_read/README.md new file mode 100644 index 00000000000..9db185efcd4 --- /dev/null +++ b/tests/compatibility/wire/native_histogram_remote_read/README.md @@ -0,0 +1,184 @@ +# Native histogram sampled Remote Read wire compatibility + +`response.pb` is a frozen `prompb.ReadResponse` generated from Prometheus +revision `e471c7708c6b0b13e2050222bbd5a907b36f4113`, independently of +`greptime-proto`. It contains one query result, one series named +`native_histogram_remote_read`, and nine samples at 1000ms intervals. + +| Timestamp (ms) | Independent expectation | +| --- | --- | +| 1000 | Integer count `2^53 + 13`, zero count `2^53 + 1`, both bucket sides `[1,3,2]`, wire deltas `[1,2,-1]` across span boundaries, schema 1, reset hint 2 | +| 2000 | Float count 6, zero count 0.5, absolute buckets `[2,3.5]`, schema 2, gauge hint | +| 3000 | Integer custom buckets `[1,2]`, one observation in `(-Inf,1]` | +| 4000 | Float custom buckets `[1,2]`, absolute counts `[1,2.5]`, gauge hint | +| 5000 | Empty integer histogram with explicit count and zero-count oneofs | +| 6000 | Empty float histogram with explicit count and zero-count oneofs | +| 7000 | Integer count 8, buckets `[3,2]`, wire deltas `[3,-1]`, ordinary NaN sum bits `0x7ff8000000000123` | +| 8000 | Empty integer histogram, stale NaN sum bits `0x7ff0000000000002` | +| 9000 | Integer count and zero count `i64::MAX` | + +The 420-byte fixture has SHA-256 `f1414bfa6f0821e920a88e325964abfad0269b05b7c0643a84ca6e908729d1eb`. +Generation used Go `go1.27.3 darwin/arm64`, the pinned Prometheus `go.mod` and +`go.sum` (declared Go version `1.26.10`), and its generated `prompb` types. +Direct wire dependencies were `github.com/gogo/protobuf v1.3.2` and +`github.com/golang/snappy v1.0.0`. No local `protoc` or Greptime protobuf +generator is involved. + +## Automated regression + +From the GreptimeDB checkout: + +```shell +cargo nextest run -p servers -E 'test(native_histogram_remote_read_wire_fixture)' +``` + +This constructs stored arrays from independent Rust inputs, checks fixed +absolute bucket counts, runs the actual record-batch encoder, Snappy-compresses +and protobuf-decodes its output, and compares all fields with the upstream +fixture. NaN payloads are compared by bits. It also checks fixed wire deltas; +converting those deltas back with an upstream helper is not the expected-value +generator. Tests never regenerate the frozen fixture. + +## Reproduce and check with Prometheus + +Use a scratch module to preserve the Prometheus checkout. Set the two checkout +paths, create a scratch directory, and save the Go program below as `main.go` +in that directory: + +```shell +PROMETHEUS_CHECKOUT=/path/to/prometheus +GREPTIMEDB_CHECKOUT=/path/to/greptimedb +NH_SCRATCH=$(mktemp -d) +git -C "$PROMETHEUS_CHECKOUT" rev-parse HEAD +cp "$PROMETHEUS_CHECKOUT/go.mod" "$PROMETHEUS_CHECKOUT/go.sum" "$NH_SCRATCH/" +cd "$NH_SCRATCH" +go mod edit -module=native-histogram-interop +go mod edit -require=github.com/prometheus/prometheus@v0.0.0 +go mod edit -replace="github.com/prometheus/prometheus=$PROMETHEUS_CHECKOUT" +# Save the Go program below as main.go here before running these commands. +go run -mod=mod . generate "$NH_SCRATCH/response.pb" +shasum -a 256 "$NH_SCRATCH/response.pb" +cmp "$NH_SCRATCH/response.pb" "$GREPTIMEDB_CHECKOUT/tests/compatibility/wire/native_histogram_remote_read/response.pb" +cd "$GREPTIMEDB_CHECKOUT" +GREPTIME_NATIVE_HISTOGRAM_REMOTE_READ_OUTPUT="$NH_SCRATCH/greptime-response.snappy" \ + cargo nextest run -p servers -E 'test(native_histogram_remote_read_wire_fixture)' +cd "$NH_SCRATCH" +go run -mod=mod . check "$NH_SCRATCH/greptime-response.snappy" +``` + +The check consumes `storage/remote.FromQueryResult` iterators to exhaustion, +checks both iterator and series-set errors, integer/float types, exact integer +counts, timestamps, spans, bucket placement, custom bounds, reset hints, and +NaN bits. It calls both `AtHistogram` and `AtFloatHistogram` on integer custom +samples, including a single finite-bound bucket. Intentional integer-to-float +conversion is not expected to retain large-integer precision. + +Verified on 2026-10-26 against the response emitted by the Rust test: +`PASS: 9 histogram samples, both iterator accessors, no series or iterator errors`. + +The tested baseline is the exact revision above, not a claimed minimum +Prometheus release. This Go check is a manual implementation check; it is not +part of the existing Go CI. The Rust regression is the automated guard. + +```go +package main + +import ( + "fmt" + "math" + "os" + "reflect" + + "github.com/golang/snappy" + "github.com/prometheus/prometheus/prompb" + "github.com/prometheus/prometheus/storage/remote" + "github.com/prometheus/prometheus/tsdb/chunkenc" +) + +func expected() *prompb.ReadResponse { + hs := []prompb.Histogram{ + {Count: &prompb.Histogram_CountInt{CountInt: (1<<53)+13}, ZeroCount: &prompb.Histogram_ZeroCountInt{ZeroCountInt: (1<<53)+1}, Sum: 10, Schema: 1, ZeroThreshold: 0.001, ResetHint: prompb.Histogram_NO, + PositiveSpans: []prompb.BucketSpan{{Offset: -2, Length: 2}, {Offset: 3, Length: 1}}, PositiveDeltas: []int64{1,2,-1}, NegativeSpans: []prompb.BucketSpan{{Offset: -3, Length: 1}, {Offset: 2, Length: 2}}, NegativeDeltas: []int64{1,2,-1}}, + {Count: &prompb.Histogram_CountFloat{CountFloat: 6}, ZeroCount: &prompb.Histogram_ZeroCountFloat{ZeroCountFloat: 0.5}, Sum: 20, Schema: 2, ZeroThreshold: 0.002, ResetHint: prompb.Histogram_GAUGE, + PositiveSpans: []prompb.BucketSpan{{Offset: 3, Length: 2}}, PositiveCounts: []float64{2,3.5}}, + {Count: &prompb.Histogram_CountInt{CountInt: 1}, ZeroCount: &prompb.Histogram_ZeroCountInt{}, Sum: 1, Schema: -53, CustomValues: []float64{1,2}, PositiveSpans: []prompb.BucketSpan{{Length: 1}}, PositiveDeltas: []int64{1}}, + {Count: &prompb.Histogram_CountFloat{CountFloat: 3.5}, ZeroCount: &prompb.Histogram_ZeroCountFloat{}, Sum: 4, Schema: -53, CustomValues: []float64{1,2}, PositiveSpans: []prompb.BucketSpan{{Length: 2}}, PositiveCounts: []float64{1,2.5}, ResetHint: prompb.Histogram_GAUGE}, + {Count: &prompb.Histogram_CountInt{}, ZeroCount: &prompb.Histogram_ZeroCountInt{}}, + {Count: &prompb.Histogram_CountFloat{}, ZeroCount: &prompb.Histogram_ZeroCountFloat{}}, + {Count: &prompb.Histogram_CountInt{CountInt: 8}, ZeroCount: &prompb.Histogram_ZeroCountInt{}, Sum: math.Float64frombits(0x7ff8000000000123), PositiveSpans: []prompb.BucketSpan{{Length: 2}}, PositiveDeltas: []int64{3,-1}}, + {Count: &prompb.Histogram_CountInt{}, ZeroCount: &prompb.Histogram_ZeroCountInt{}, Sum: math.Float64frombits(0x7ff0000000000002)}, + {Count: &prompb.Histogram_CountInt{CountInt: math.MaxInt64}, ZeroCount: &prompb.Histogram_ZeroCountInt{ZeroCountInt: math.MaxInt64}}, + } + for i := range hs { hs[i].Timestamp = int64(i+1)*1000 } + return &prompb.ReadResponse{Results: []*prompb.QueryResult{{Timeseries: []*prompb.TimeSeries{{Labels: []prompb.Label{{Name:"__name__", Value:"native_histogram_remote_read"}}, Histograms:hs}}}}} +} + +func check(ok bool, args ...any) { if !ok { panic(fmt.Sprint(args...)) } } +func main() { + want := expected() + if os.Args[1] == "generate" { + data, err := want.Marshal(); if err != nil { panic(err) } + if err := os.WriteFile(os.Args[2], data, 0644); err != nil { panic(err) } + fmt.Printf("generated %d bytes\n", len(data)); return + } + data, err := os.ReadFile(os.Args[2]); if err != nil { panic(err) } + data, err = snappy.Decode(nil, data); if err != nil { panic(err) } + var got prompb.ReadResponse + if err = got.Unmarshal(data); err != nil { panic(err) } + check(len(got.Results)==1 && len(got.Results[0].Timeseries)==1, "result/series count") + hs := got.Results[0].Timeseries[0].Histograms + expectedHs := want.Results[0].Timeseries[0].Histograms + check(len(hs)==len(expectedHs), "histogram count") + for i := range hs { + g,w := hs[i], expectedHs[i] + check(math.Float64bits(g.Sum)==math.Float64bits(w.Sum), "sum bits ", i) + g.Sum,w.Sum=0,0 + check(reflect.DeepEqual(g,w), "protobuf fields ", i, ": ", g, " expected ", w) + } + ss := remote.FromQueryResult(true, got.Results[0]) + samples,series := 0,0 + for ss.Next() { + series++ + it := ss.At().Iterator(nil) + for typ := it.Next(); typ != chunkenc.ValNone; typ=it.Next() { + w := expectedHs[samples] + if w.IsFloatHistogram() { + check(typ==chunkenc.ValFloatHistogram, "float type ", samples) + t,h:=it.AtFloatHistogram(nil) + check(t==w.Timestamp && h.Count==w.GetCountFloat() && h.ZeroCount==w.GetZeroCountFloat(), "float count/time ", samples) + check(reflect.DeepEqual(h.PositiveBuckets,w.PositiveCounts), "float buckets ", samples) + check(reflect.DeepEqual(h.CustomValues,w.CustomValues), "float bounds ", samples) + check(math.Float64bits(h.Sum)==math.Float64bits(w.Sum), "float sum ", samples) + check(int32(h.CounterResetHint)==int32(w.ResetHint), "float reset ", samples) + } else { + check(typ==chunkenc.ValHistogram, "integer type ", samples) + t,h:=it.AtHistogram(nil) + check(t==w.Timestamp && h.Count==w.GetCountInt() && h.ZeroCount==w.GetZeroCountInt(), "integer count/time ", samples) + check(reflect.DeepEqual(h.PositiveBuckets,w.PositiveDeltas) && reflect.DeepEqual(h.NegativeBuckets,w.NegativeDeltas), "integer deltas ", samples) + check(reflect.DeepEqual(h.CustomValues,w.CustomValues), "integer bounds ", samples) + check(math.Float64bits(h.Sum)==math.Float64bits(w.Sum), "integer sum ", samples) + check(int32(h.CounterResetHint)==int32(w.ResetHint), "integer reset ", samples) + _,f:=it.AtFloatHistogram(nil) + check(reflect.DeepEqual(f.CustomValues,w.CustomValues), "converted bounds ", samples) + if w.Schema == -53 { + buckets:=f.PositiveBucketIterator() + check(buckets.Next(), "custom bucket missing") + b:=buckets.At() + check(b.Lower==math.Inf(-1) && b.Upper==1 && b.Count==1, "custom bucket placement: ",b) + check(!buckets.Next(), "extra custom bucket") + } + } + _,f:=it.AtFloatHistogram(nil) + check(f.Schema==w.Schema && f.ZeroThreshold==w.ZeroThreshold, "schema/threshold ",samples) + check(len(f.PositiveSpans)==len(w.PositiveSpans) && len(f.NegativeSpans)==len(w.NegativeSpans), "span counts ",samples) + for i,s:=range f.PositiveSpans { check(s.Offset==w.PositiveSpans[i].Offset && s.Length==w.PositiveSpans[i].Length,"positive spans") } + for i,s:=range f.NegativeSpans { check(s.Offset==w.NegativeSpans[i].Offset && s.Length==w.NegativeSpans[i].Length,"negative spans") } + samples++ + } + check(it.Err()==nil,"iterator error: ",it.Err()) + } + check(ss.Err()==nil,"series error: ",ss.Err()) + check(series==1 && samples==9,"series/sample count ",series,"/",samples) + fmt.Println("PASS: 9 histogram samples, both iterator accessors, no series or iterator errors") +} +``` diff --git a/tests/compatibility/wire/native_histogram_remote_read/response.pb b/tests/compatibility/wire/native_histogram_remote_read/response.pb new file mode 100644 index 0000000000000000000000000000000000000000..42d761854204c5403b596993cf7476255e7d13d5 GIT binary patch literal 420 zcmd;D$jmj5nM;FQ^gB*C1*j1DH0A)<4hbBSfIuw*jqv*AJ~KSF=Yur_zVsZn#ZAlx#9x9k`4zW zRQCsag-ip6C?=pCj7$|f#G4qQT0Yprv?_T+t>l4ODd3=x3AKp%=;o{|Iyhomwi$nGEYV23d!F|jdL+%QuT;9!ts0z2vh WP^RL5osu}me<-LofU_!AxB&p8<5}$h literal 0 HcmV?d00001 From e6b60644b3ce10b96d8be3d8eb11c43c8d18d0ed Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Tue, 27 Oct 2026 03:42:27 +0000 Subject: [PATCH 104/239] fix(mito2): stop the series distributor from stalling on unread partitions (#9991) When every partition channel is full, the distributor walked the senders round-robin with a 100us send timeout each. Tokio timers fire on a 1ms tick, so a consumer that leaves some partitions unpolled costs about one millisecond per unpolled partition for every series batch. SeriesScan::build_stream and scan_all_partitions drain partitions one after another, and SortPreservingMergeExec above a multi-partition SeriesScan only pulls the partition holding the current minimum key. With 16 partitions and one batch per series after a merge, a 40-series scan read partition by partition took ~480ms (measured on the native PromQL two-input route). Wait on all senders at once and send to the first one that frees a slot. The send-timeout metric goes away with the timeout. Signed-off-by: Dennis Zhuang --- src/mito2/src/read/scan_util.rs | 10 ---- src/mito2/src/read/series_scan.rs | 99 +++++++++++++++++-------------- 2 files changed, 56 insertions(+), 53 deletions(-) diff --git a/src/mito2/src/read/scan_util.rs b/src/mito2/src/read/scan_util.rs index b122704a71e..a67e2449baa 100644 --- a/src/mito2/src/read/scan_util.rs +++ b/src/mito2/src/read/scan_util.rs @@ -202,8 +202,6 @@ pub(crate) struct ScanMetricsSet { /// Elapsed time before the first poll operation. first_poll: Duration, - /// Number of send timeout in SeriesScan. - num_series_send_timeout: usize, /// Number of send full in SeriesScan. num_series_send_full: usize, /// Number of rows the series distributor scanned. @@ -325,7 +323,6 @@ impl fmt::Debug for ScanMetricsSet { num_sst_batches, num_sst_rows, first_poll, - num_series_send_timeout, num_series_send_full, num_distributor_rows, num_distributor_batches, @@ -459,9 +456,6 @@ impl fmt::Debug for ScanMetricsSet { } // Write non-zero distributor metrics - if *num_series_send_timeout > 0 { - write!(f, ", \"num_series_send_timeout\":{num_series_send_timeout}")?; - } if *num_series_send_full > 0 { write!(f, ", \"num_series_send_full\":{num_series_send_full}")?; } @@ -775,7 +769,6 @@ impl ScanMetricsSet { /// Sets distributor metrics. fn set_distributor_metrics(&mut self, distributor_metrics: &SeriesDistributorMetrics) { let SeriesDistributorMetrics { - num_series_send_timeout, num_series_send_full, num_rows, num_batches, @@ -784,7 +777,6 @@ impl ScanMetricsSet { divider_cost, } = distributor_metrics; - self.num_series_send_timeout += *num_series_send_timeout; self.num_series_send_full += *num_series_send_full; self.num_distributor_rows += *num_rows; self.num_distributor_batches += *num_batches; @@ -1144,8 +1136,6 @@ impl fmt::Debug for PartitionMetrics { /// Metrics for the series distributor. #[derive(Default)] pub(crate) struct SeriesDistributorMetrics { - /// Number of send timeout in SeriesScan. - pub(crate) num_series_send_timeout: usize, /// Number of send full in SeriesScan. pub(crate) num_series_send_full: usize, /// Number of rows the series distributor scanned. diff --git a/src/mito2/src/read/series_scan.rs b/src/mito2/src/read/series_scan.rs index f8af9e51b91..23418686cb9 100644 --- a/src/mito2/src/read/series_scan.rs +++ b/src/mito2/src/read/series_scan.rs @@ -17,7 +17,7 @@ use std::fmt; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; -use std::time::{Duration, Instant}; +use std::time::Instant; use async_stream::try_stream; use common_error::ext::BoxedError; @@ -39,7 +39,7 @@ use store_api::region_engine::{ PartitionRange, PrepareRequest, QueryScanContext, RegionScanner, ScannerProperties, }; use tokio::sync::Semaphore; -use tokio::sync::mpsc::error::{SendTimeoutError, TrySendError}; +use tokio::sync::mpsc::error::TrySendError; use tokio::sync::mpsc::{self, Receiver, Sender}; use crate::error::{ @@ -61,9 +61,6 @@ use crate::read::stream::{ConvertBatchStream, ScanBatch, ScanBatchStream}; use crate::sst::parquet::flat_format::primary_key_column_index; use crate::sst::parquet::format::PrimaryKeyArray; -/// Timeout to send a batch to a sender. -const SEND_TIMEOUT: Duration = Duration::from_micros(100); - /// Maximum number of candidate series retained before distributing assignments. const CANDIDATE_SERIES_ASSIGNMENT_THRESHOLD: usize = 1_000_000; @@ -843,7 +840,6 @@ impl SeriesDistributor { } metrics.scan_cost += fetch_start.elapsed(); - metrics.num_series_send_timeout = self.senders.num_timeout; metrics.num_series_send_full = self.senders.num_full; part_metrics.set_distributor_metrics(&metrics); @@ -996,8 +992,6 @@ struct SenderList { num_nones: usize, /// Index of the current partition to send. sender_idx: usize, - /// Number of timeout. - num_timeout: usize, /// Number of full senders. num_full: usize, } @@ -1009,7 +1003,6 @@ impl SenderList { senders, num_nones, sender_idx: 0, - num_timeout: 0, num_full: 0, } } @@ -1045,46 +1038,35 @@ impl SenderList { } /// Finds a partition and sends the batch to the partition. - async fn send_batch(&mut self, mut batch: SeriesBatch) -> Result<()> { - // Sends the batch without blocking first. - match self.try_send_batch(batch)? { - Some(b) => { - // Unable to send batch to partition. - batch = b; - } - None => { - return Ok(()); - } - } + async fn send_batch(&mut self, batch: SeriesBatch) -> Result<()> { + let Some(batch) = self.try_send_batch(batch)? else { + return Ok(()); + }; loop { ensure!(self.num_nones < self.senders.len(), InvalidSenderSnafu); - let sender_idx = self.fetch_add_sender_idx(); - let Some(sender) = &self.senders[sender_idx] else { - continue; + // Waits on every partition at once and sends to whichever frees a slot first. A + // consumer may leave partitions unpolled for a long time, e.g. one that drains them one + // after another, and waiting on such a partition would stall the scan. + let closed = { + let reservations = self.senders.iter().enumerate().filter_map(|(idx, sender)| { + sender + .as_ref() + .map(|sender| Box::pin(async move { (idx, sender.reserve().await) })) + }); + let ((idx, reserved), _, _) = futures::future::select_all(reservations).await; + match reserved { + Ok(permit) => { + permit.send(Ok(batch)); + return Ok(()); + } + Err(_) => idx, + } }; - // Adds a timeout to avoid blocking indefinitely and sending - // the batch in a round-robin fashion when some partitions - // don't poll their inputs. This may happen if we have a - // node like sort merging. But it is rare when we are using SeriesScan. - match sender.send_timeout(Ok(batch), SEND_TIMEOUT).await { - Ok(()) => break, - Err(SendTimeoutError::Timeout(res)) => { - self.num_timeout += 1; - // Safety: we send Ok. - batch = res.unwrap(); - } - Err(SendTimeoutError::Closed(res)) => { - self.senders[sender_idx] = None; - self.num_nones += 1; - // Safety: we send Ok. - batch = res.unwrap(); - } - } + self.senders[closed] = None; + self.num_nones += 1; } - - Ok(()) } async fn send_error(&self, error: Error) { @@ -1294,6 +1276,37 @@ mod tests { assert_eq!(0, active_receivers.load(Ordering::Relaxed)); } + #[tokio::test] + async fn distributor_does_not_stall_on_partitions_read_one_after_another() { + const NUM_BATCHES: usize = 2000; + let (mut senders, receivers) = new_legacy_channel_list(16); + let producer = tokio::spawn(async move { + for _ in 0..NUM_BATCHES { + senders + .send_batch(SeriesBatch::Flat(FlatSeriesBatch::default())) + .await + .unwrap(); + } + }); + // Reads partition 0 to its end before touching the next one, so all other partitions stay + // full while the producer still has batches to send. + let consume = async { + let mut received = 0; + for mut receiver in receivers.into_iter().flatten() { + while let Some(batch) = receiver.recv().await { + batch.unwrap(); + received += 1; + } + } + received + }; + let received = tokio::time::timeout(std::time::Duration::from_secs(1), consume) + .await + .expect("the distributor stalled on partitions the consumer was not reading"); + producer.await.unwrap(); + assert_eq!(NUM_BATCHES, received); + } + fn new_test_record_batch( primary_keys: &[&[u8]], timestamps: &[i64], From 2117f8b059e6c6c5a518dc4d07e71c42c6854cc3 Mon Sep 17 00:00:00 2001 From: Chenglong Zheng Date: Tue, 27 Oct 2026 07:27:55 +0000 Subject: [PATCH 105/239] feat: Support procedure_state() using the last procedure ID (#9876) * feat: Support procedure_state() using the last procedure ID Signed-off-by: hahaha * feat: Fix and add tests according to review comments Signed-off-by: hahaha * fix: the zero-argument signature in ADMIN swallows explicit arguments Signed-off-by: hahaha * fix: fix the CI failures caused by SQL version issues Signed-off-by: hahaha * fix: apply the changes based on the review Signed-off-by: hahaha --------- Signed-off-by: hahaha --- .../function/src/admin/migrate_region.rs | 11 +- .../function/src/admin/reconcile_catalog.rs | 11 +- .../function/src/admin/reconcile_database.rs | 11 +- .../function/src/admin/reconcile_table.rs | 11 +- src/common/function/src/state.rs | 16 +- .../function/src/system/procedure_state.rs | 122 +++++-- src/common/macro/src/admin_fn.rs | 27 +- src/operator/src/statement/admin.rs | 72 +++- src/operator/src/statement/ddl.rs | 15 +- src/session/src/context.rs | 18 + src/session/src/lib.rs | 4 + tests-integration/tests/repartition.rs | 58 +++- tests-integration/tests/truncate.rs | 308 ++++++++++++++++-- .../function/procedure_state_session.result | 124 +++++++ .../function/procedure_state_session.sql | 61 ++++ 15 files changed, 788 insertions(+), 81 deletions(-) create mode 100644 tests/cases/standalone/common/function/procedure_state_session.result create mode 100644 tests/cases/standalone/common/function/procedure_state_session.sql diff --git a/src/common/function/src/admin/migrate_region.rs b/src/common/function/src/admin/migrate_region.rs index 5d258e14946..a359e0ed4ef 100644 --- a/src/common/function/src/admin/migrate_region.rs +++ b/src/common/function/src/admin/migrate_region.rs @@ -94,7 +94,10 @@ pub(crate) async fn migrate_region( .await?; match pid { - Some(pid) => Ok(Value::from(pid)), + Some(pid) => { + query_ctx.set_last_procedure_id(pid.clone()); + Ok(Value::from(pid)) + } None => Ok(Value::Null), } } @@ -184,7 +187,10 @@ mod tests { #[tokio::test] async fn test_migrate_region() { let factory: ScalarFunctionFactory = MigrateRegionFunction::factory().into(); - let provider = factory.provide(FunctionContext::mock()); + let context = FunctionContext::mock(); + let query_ctx = context.query_ctx.clone(); + query_ctx.set_last_procedure_id("previous-id".to_string()); + let provider = factory.provide(context); let f = provider.as_async().unwrap(); let func_args = datafusion::logical_expr::ScalarFunctionArgs { @@ -203,6 +209,7 @@ mod tests { config_options: Arc::new(datafusion_common::config::ConfigOptions::default()), }; let result = f.invoke_async_with_args(func_args).await.unwrap(); + assert_eq!(query_ctx.last_procedure_id().as_deref(), Some("test_pid")); match result { ColumnarValue::Array(array) => { diff --git a/src/common/function/src/admin/reconcile_catalog.rs b/src/common/function/src/admin/reconcile_catalog.rs index d860970d610..6a709c759b7 100644 --- a/src/common/function/src/admin/reconcile_catalog.rs +++ b/src/common/function/src/admin/reconcile_catalog.rs @@ -96,7 +96,10 @@ pub(crate) async fn reconcile_catalog( }) .await?; match pid { - Some(pid) => Ok(Value::from(pid)), + Some(pid) => { + query_ctx.set_last_procedure_id(pid.clone()); + Ok(Value::from(pid)) + } None => Ok(Value::Null), } } @@ -138,7 +141,10 @@ mod tests { // reconcile_catalog() let factory: ScalarFunctionFactory = ReconcileCatalogFunction::factory().into(); - let provider = factory.provide(FunctionContext::mock()); + let context = FunctionContext::mock(); + let query_ctx = context.query_ctx.clone(); + query_ctx.set_last_procedure_id("previous-id".to_string()); + let provider = factory.provide(context); let f = provider.as_async().unwrap(); let func_args = datafusion::logical_expr::ScalarFunctionArgs { @@ -150,6 +156,7 @@ mod tests { }; let result = f.invoke_async_with_args(func_args).await.unwrap(); + assert_eq!(query_ctx.last_procedure_id().as_deref(), Some("test_pid")); match result { ColumnarValue::Array(array) => { let result_array = array.as_any().downcast_ref::().unwrap(); diff --git a/src/common/function/src/admin/reconcile_database.rs b/src/common/function/src/admin/reconcile_database.rs index 020a7f4cf04..5d07d1d87c7 100644 --- a/src/common/function/src/admin/reconcile_database.rs +++ b/src/common/function/src/admin/reconcile_database.rs @@ -105,7 +105,10 @@ pub(crate) async fn reconcile_database( }) .await?; match pid { - Some(pid) => Ok(Value::from(pid)), + Some(pid) => { + query_ctx.set_last_procedure_id(pid.clone()); + Ok(Value::from(pid)) + } None => Ok(Value::Null), } } @@ -148,7 +151,10 @@ mod tests { // reconcile_database(database_name) let factory: ScalarFunctionFactory = ReconcileDatabaseFunction::factory().into(); - let provider = factory.provide(FunctionContext::mock()); + let context = FunctionContext::mock(); + let query_ctx = context.query_ctx.clone(); + query_ctx.set_last_procedure_id("previous-id".to_string()); + let provider = factory.provide(context); let f = provider.as_async().unwrap(); let func_args = datafusion::logical_expr::ScalarFunctionArgs { @@ -161,6 +167,7 @@ mod tests { config_options: Arc::new(datafusion_common::config::ConfigOptions::default()), }; let result = f.invoke_async_with_args(func_args).await.unwrap(); + assert_eq!(query_ctx.last_procedure_id().as_deref(), Some("test_pid")); match result { ColumnarValue::Array(array) => { let result_array = array.as_any().downcast_ref::().unwrap(); diff --git a/src/common/function/src/admin/reconcile_table.rs b/src/common/function/src/admin/reconcile_table.rs index 0c82c2ddc9c..37176541370 100644 --- a/src/common/function/src/admin/reconcile_table.rs +++ b/src/common/function/src/admin/reconcile_table.rs @@ -86,7 +86,10 @@ pub(crate) async fn reconcile_table( }) .await?; match pid { - Some(pid) => Ok(Value::from(pid)), + Some(pid) => { + query_ctx.set_last_procedure_id(pid.clone()); + Ok(Value::from(pid)) + } None => Ok(Value::Null), } } @@ -121,7 +124,10 @@ mod tests { // reconcile_table(table_name) let factory: ScalarFunctionFactory = ReconcileTableFunction::factory().into(); - let provider = factory.provide(FunctionContext::mock()); + let context = FunctionContext::mock(); + let query_ctx = context.query_ctx.clone(); + query_ctx.set_last_procedure_id("previous-id".to_string()); + let provider = factory.provide(context); let f = provider.as_async().unwrap(); let func_args = datafusion::logical_expr::ScalarFunctionArgs { @@ -134,6 +140,7 @@ mod tests { config_options: Arc::new(datafusion_common::config::ConfigOptions::default()), }; let result = f.invoke_async_with_args(func_args).await.unwrap(); + assert_eq!(query_ctx.last_procedure_id().as_deref(), Some("test_pid")); match result { ColumnarValue::Array(array) => { let result_array = array.as_any().downcast_ref::().unwrap(); diff --git a/src/common/function/src/state.rs b/src/common/function/src/state.rs index bf2405930c3..d7e4f5e85b5 100644 --- a/src/common/function/src/state.rs +++ b/src/common/function/src/state.rs @@ -34,7 +34,16 @@ impl FunctionState { /// Create a mock [`FunctionState`] for test. #[cfg(any(test, feature = "testing"))] pub fn mock() -> Self { + Self::mock_with_procedure_state_counter(Default::default()) + } + + /// Creates a mock state that counts procedure state queries. + #[cfg(any(test, feature = "testing"))] + pub fn mock_with_procedure_state_counter( + calls: std::sync::Arc, + ) -> Self { use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; use api::v1::meta::{ProcedureStatus, ReconcileRequest}; use async_trait::async_trait; @@ -55,7 +64,9 @@ impl FunctionState { use table::table_name::TableName; use crate::handlers::{FlowServiceHandler, ProcedureServiceHandler, TableMutationHandler}; - struct MockProcedureServiceHandler; + struct MockProcedureServiceHandler { + calls: Arc, + } struct MockTableMutationHandler; struct MockFlowServiceHandler; const ROWS: usize = 42; @@ -83,6 +94,7 @@ impl FunctionState { } async fn query_procedure_state(&self, _pid: &str) -> Result { + self.calls.fetch_add(1, Ordering::Relaxed); Ok(ProcedureStateResponse { status: ProcedureStatus::Done.into(), error: "OK".to_string(), @@ -219,7 +231,7 @@ impl FunctionState { Self { plugins: Plugins::default(), table_mutation_handler: Some(Arc::new(MockTableMutationHandler)), - procedure_service_handler: Some(Arc::new(MockProcedureServiceHandler)), + procedure_service_handler: Some(Arc::new(MockProcedureServiceHandler { calls })), flow_service_handler: Some(Arc::new(MockFlowServiceHandler)), } } diff --git a/src/common/function/src/system/procedure_state.rs b/src/common/function/src/system/procedure_state.rs index 7f7388edae9..b9b2f453fdb 100644 --- a/src/common/function/src/system/procedure_state.rs +++ b/src/common/function/src/system/procedure_state.rs @@ -20,11 +20,11 @@ use common_query::error::{ InvalidFuncArgsSnafu, MissingProcedureServiceHandlerSnafu, Result, UnsupportedInputDataTypeSnafu, }; -use datafusion_expr::{Signature, Volatility}; +use datafusion_expr::{Signature, TypeSignature, Volatility}; use datatypes::prelude::*; use serde::Serialize; use session::context::QueryContextRef; -use snafu::ensure; +use snafu::{OptionExt, ensure}; use crate::handlers::ProcedureServiceHandlerRef; @@ -35,8 +35,10 @@ struct ProcedureStateJson { error: Option, } -/// A function to query procedure state by its id. -/// Such as `procedure_state(pid)`. +/// Queries procedure state using an explicit procedure ID, or the latest +/// procedure ID recorded in the Session. +/// The zero-argument form reads the procedure ID from the Session at execution +/// time. State lookups do not update the recorded procedure ID. #[admin_fn( name = ProcedureStateFunction, display_name = procedure_state, @@ -45,25 +47,35 @@ struct ProcedureStateJson { )] pub(crate) async fn procedure_state( procedure_service_handler: &ProcedureServiceHandlerRef, - _ctx: &QueryContextRef, + ctx: &QueryContextRef, params: &[ValueRef<'_>], ) -> Result { ensure!( - params.len() == 1, + params.len() <= 1, InvalidFuncArgsSnafu { err_msg: format!( - "The length of the args is not correct, expect 1, have: {}", + "The length of the args is not correct, expect 0 or 1, have: {}", params.len() ), } ); - let ValueRef::String(pid) = params[0] else { - return UnsupportedInputDataTypeSnafu { - function: "procedure_state", - datatypes: params.iter().map(|v| v.data_type()).collect::>(), + let last_procedure_id; + let pid = match params { + [] => { + last_procedure_id = ctx.last_procedure_id().context(InvalidFuncArgsSnafu { + err_msg: "No procedure ID has been returned in this session; pass an explicit procedure ID", + })?; + last_procedure_id.as_str() + } + [ValueRef::String(pid)] => *pid, + _ => { + return UnsupportedInputDataTypeSnafu { + function: "procedure_state", + datatypes: params.iter().map(|v| v.data_type()).collect::>(), + } + .fail(); } - .fail(); }; let ProcedureStateResponse { status, error, .. } = @@ -82,20 +94,29 @@ pub(crate) async fn procedure_state( } fn signature() -> Signature { - Signature::uniform(1, vec![ArrowDataType::Utf8], Volatility::Immutable) + Signature::one_of( + vec![ + TypeSignature::Exact(vec![]), + TypeSignature::Exact(vec![ArrowDataType::Utf8]), + ], + Volatility::Volatile, + ) } #[cfg(test)] mod tests { use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; use arrow::array::StringArray; use arrow::datatypes::{DataType, Field}; use datafusion_expr::ColumnarValue; + use datafusion_expr::type_coercion::functions::fields_with_udf; use super::*; use crate::function::FunctionContext; use crate::function_factory::ScalarFunctionFactory; + use crate::state::FunctionState; #[test] fn test_procedure_state_misc() { @@ -103,12 +124,75 @@ mod tests { let f = factory.provide(FunctionContext::mock()); assert_eq!("procedure_state", f.name()); assert_eq!(DataType::Utf8, f.return_type(&[]).unwrap()); - assert!(matches!(f.signature(), - datafusion_expr::Signature { - type_signature: datafusion_expr::TypeSignature::Uniform(1, valid_types), - volatility: datafusion_expr::Volatility::Immutable, - .. - } if valid_types == &vec![ArrowDataType::Utf8])); + assert_eq!( + f.signature(), + &Signature::one_of( + vec![ + TypeSignature::Exact(vec![]), + TypeSignature::Exact(vec![ArrowDataType::Utf8]), + ], + Volatility::Volatile, + ) + ); + let field = Arc::new(Field::new("pid", DataType::Utf8, false)); + assert!(fields_with_udf(&[field.clone(), field], &f).is_err()); + } + + #[tokio::test] + async fn test_procedure_state_args() { + let context = FunctionContext::mock(); + let handler = context.state.procedure_service_handler.as_ref().unwrap(); + let ctx = &context.query_ctx; + let err = procedure_state(handler, ctx, &[]).await.unwrap_err(); + assert!( + err.to_string() + .contains("No procedure ID has been returned") + ); + + ctx.set_last_procedure_id("session-id".to_string()); + for (params, succeeds) in [ + (vec![], true), + (vec![ValueRef::String("explicit-id")], true), + (vec![ValueRef::Null], false), + (vec![ValueRef::String("a"), ValueRef::String("b")], false), + ] { + let result = procedure_state(handler, ctx, ¶ms).await; + assert_eq!(result.is_ok(), succeeds, "{params:?}: {result:?}"); + assert_eq!(ctx.last_procedure_id().as_deref(), Some("session-id")); + } + } + + #[tokio::test] + async fn test_procedure_state_zero_arg_rows() { + let calls = Arc::new(AtomicUsize::new(0)); + let context = FunctionContext { + state: Arc::new(FunctionState::mock_with_procedure_state_counter( + calls.clone(), + )), + ..Default::default() + }; + context + .query_ctx + .set_last_procedure_id("session-id".to_string()); + let factory: ScalarFunctionFactory = ProcedureStateFunction::factory().into(); + let function = factory.provide(context); + let f = function.as_async().unwrap(); + + for number_rows in [0, 1, 2] { + calls.store(0, Ordering::Relaxed); + let args = datafusion_expr::ScalarFunctionArgs { + args: vec![], + arg_fields: vec![], + return_field: Arc::new(Field::new("result", DataType::Utf8, true)), + number_rows, + config_options: Arc::new(datafusion_common::config::ConfigOptions::default()), + }; + let ColumnarValue::Array(result) = f.invoke_async_with_args(args).await.unwrap() else { + panic!("expected array output"); + }; + assert_eq!(result.len(), number_rows); + assert_eq!(calls.load(Ordering::Relaxed), number_rows); + } } #[tokio::test] diff --git a/src/common/macro/src/admin_fn.rs b/src/common/macro/src/admin_fn.rs index 3575ebfe6c2..43b4671fb77 100644 --- a/src/common/macro/src/admin_fn.rs +++ b/src/common/macro/src/admin_fn.rs @@ -311,12 +311,8 @@ fn build_struct( // Safety check: Ensure under the `greptime` catalog for security #user_path::ensure_greptime!(self.func_ctx); - let columns_num = columns.len(); - let rows_num = if columns.is_empty() { - 1 - } else { - columns[0].len() - }; + // Zero-argument functions also execute once per input row. + let rows_num = columns.first().map_or(args.number_rows, |column| column.len()); #validate_rows use snafu::{OptionExt, ResultExt}; @@ -333,22 +329,15 @@ fn build_struct( let mut builder = store_api::storage::ConcreteDataType::#ret() .create_mutable_vector(rows_num); - if columns_num == 0 { - let result = #fn_name(handler, query_ctx, &[]).await + for i in 0..rows_num { + let args: Vec<_> = columns.iter() + .map(|vector| vector.get_ref(i)) + .collect(); + + let result = #fn_name(handler, query_ctx, &args).await .map_err(|e| datafusion_common::DataFusionError::Execution(e.output_msg()))?; builder.push_value_ref(&result.as_value_ref()); - } else { - for i in 0..rows_num { - let args: Vec<_> = columns.iter() - .map(|vector| vector.get_ref(i)) - .collect(); - - let result = #fn_name(handler, query_ctx, &args).await - .map_err(|e| datafusion_common::DataFusionError::Execution(e.output_msg()))?; - - builder.push_value_ref(&result.as_value_ref()); - } } let result_vector = builder.to_vector(); diff --git a/src/operator/src/statement/admin.rs b/src/operator/src/statement/admin.rs index ff8b26915cc..0e488d0caaa 100644 --- a/src/operator/src/statement/admin.rs +++ b/src/operator/src/statement/admin.rs @@ -306,6 +306,17 @@ fn args_to_vector( let tz = query_ctx.timezone(); match type_signature { + TypeSignature::Nullary => { + ensure!( + args.is_empty(), + error::FunctionArityMismatchSnafu { + actual: args.len(), + expected: 0_usize, + } + ); + Ok(vec![]) + } + TypeSignature::Variadic(valid_types) => { values_to_vectors_by_valid_types(valid_types, args, Some(&tz)) } @@ -323,6 +334,13 @@ fn args_to_vector( } TypeSignature::Exact(data_types) => { + ensure!( + args.len() == data_types.len(), + error::FunctionArityMismatchSnafu { + actual: args.len(), + expected: data_types.len(), + } + ); values_to_vectors_by_exact_types(data_types, args, Some(&tz)) } @@ -450,9 +468,61 @@ fn try_get_data_type_for_sql_value(value: &SqlValue) -> Result { mod tests { use std::sync::Arc; + use datafusion_expr::TypeSignature; + use datatypes::arrow::datatypes::DataType as ArrowDataType; + use datatypes::value::Value; use datatypes::vectors::{Int32Vector, VectorRef}; + use session::context::QueryContext; + use sql::ast::Value as SqlValue; + use sql::dialect::GreptimeDbDialect; + use sql::parser::{ParseOptions, ParserContext}; + use sql::statements::statement::Statement; - use crate::statement::admin::immediate_result; + use crate::error::Error; + use crate::statement::admin::{admin_output_schema, args_to_vector, immediate_result}; + + #[test] + fn test_args_to_vector_arity() { + let query_ctx = QueryContext::arc(); + let value = SqlValue::SingleQuotedString("explicit-id".to_string()); + for (signature, expected) in [ + (TypeSignature::Nullary, 0), + (TypeSignature::Exact(vec![]), 0), + (TypeSignature::Exact(vec![ArrowDataType::Utf8]), 1), + (TypeSignature::Exact(vec![ArrowDataType::Utf8; 2]), 2), + ] { + for actual in 0..=3 { + let result = args_to_vector(&signature, &vec![&value; actual], &query_ctx); + if actual == expected { + let vectors = result.unwrap(); + assert_eq!(vectors.len(), expected); + for vector in vectors { + assert_eq!(vector.get(0), Value::from("explicit-id")); + } + } else { + assert!(matches!( + result, + Err(Error::FunctionArityMismatch { actual: a, expected: e }) + if a == actual && e == expected + )); + } + } + } + } + + #[test] + fn test_admin_reconcile_catalog_nullary() { + let Statement::Admin(statement) = ParserContext::create_with_dialect( + "ADMIN reconcile_catalog()", + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .unwrap() + .remove(0) else { + panic!("expected ADMIN statement"); + }; + assert!(admin_output_schema(&statement, &QueryContext::arc()).is_some()); + } #[test] fn empty_admin_function_result_has_no_immediate_value() { diff --git a/src/operator/src/statement/ddl.rs b/src/operator/src/statement/ddl.rs index 19cc294d1bf..e14256aaffe 100644 --- a/src/operator/src/statement/ddl.rs +++ b/src/operator/src/statement/ddl.rs @@ -146,8 +146,7 @@ const ALLOWED_FLOW_OPTIONS: [&str; 2] = [ FLOW_EXPERIMENTAL_ENABLE_INCREMENTAL_READ_KEY, ]; -fn build_procedure_id_output(procedure_id: Vec) -> Result { - let procedure_id = String::from_utf8_lossy(&procedure_id).to_string(); +fn build_procedure_id_output(procedure_id: &str) -> Result { let vector: VectorRef = Arc::new(StringVector::from(vec![procedure_id])); let schema = Arc::new(Schema::new(vec![ColumnSchema::new( "Procedure ID", @@ -1775,12 +1774,15 @@ impl StatementExecutor { table_id, time_ranges, ddl_options, - query_context, + query_context.clone(), ) .await?; if !ddl_options.wait { - return build_procedure_id_output(response.key); + let procedure_id = String::from_utf8_lossy(&response.key).into_owned(); + let output = build_procedure_id_output(&procedure_id)?; + query_context.set_last_procedure_id(procedure_id); + return Ok(output); } Ok(Output::new_with_affected_rows(0)) @@ -2062,7 +2064,10 @@ impl StatementExecutor { .context(error::ExecuteDdlSnafu)?; if !ddl_options.wait { - return build_procedure_id_output(response.key); + let procedure_id = String::from_utf8_lossy(&response.key).into_owned(); + let output = build_procedure_id_output(&procedure_id)?; + query_context.set_last_procedure_id(procedure_id); + return Ok(output); } // Only invalidate cache if wait is true. diff --git a/src/session/src/context.rs b/src/session/src/context.rs index 43cd5434ca3..f46a2a590b2 100644 --- a/src/session/src/context.rs +++ b/src/session/src/context.rs @@ -437,6 +437,24 @@ impl QueryContext { self.mutable_session_data.write().unwrap().read_preference = read_preference; } + /// Returns the latest procedure ID recorded in this Session. + pub fn last_procedure_id(&self) -> Option { + self.mutable_session_data + .read() + .unwrap() + .last_procedure_id + .clone() + } + + /// Records a procedure ID after successful statement execution and + /// construction of its ID result. + /// + /// Recording precedes `post_execute` hooks and protocol transmission, which + /// may still fail. It does not confirm that the client received the ID. + pub fn set_last_procedure_id(&self, procedure_id: String) { + self.mutable_session_data.write().unwrap().last_procedure_id = Some(procedure_id); + } + pub fn current_user(&self) -> UserInfoRef { self.mutable_session_data.read().unwrap().user_info.clone() } diff --git a/src/session/src/lib.rs b/src/session/src/lib.rs index 478458180ef..af0a4918e4b 100644 --- a/src/session/src/lib.rs +++ b/src/session/src/lib.rs @@ -61,6 +61,9 @@ pub(crate) struct MutableInner { timezone: Timezone, query_timeout: Option, read_preference: ReadPreference, + /// Latest procedure ID recorded during statement execution, before + /// `post_execute` hooks and protocol transmission. + last_procedure_id: Option, /// Request-level WAL policy for ordinary inserts. skip_wal: bool, #[debug(skip)] @@ -77,6 +80,7 @@ impl Default for MutableInner { timezone: get_timezone(None).clone(), query_timeout: None, read_preference: ReadPreference::Leader, + last_procedure_id: None, skip_wal: false, cursors: HashMap::with_capacity(0), warnings: VecDeque::new(), diff --git a/tests-integration/tests/repartition.rs b/tests-integration/tests/repartition.rs index 1198e6d23e5..880bd715c57 100644 --- a/tests-integration/tests/repartition.rs +++ b/tests-integration/tests/repartition.rs @@ -16,23 +16,25 @@ use std::collections::BTreeSet; use std::sync::Arc; use std::time::Duration; -use client::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; +use client::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, OutputData}; use common_error::root_source; use common_event_recorder::{PersistentEventContext, TriggerReason}; use common_meta::key::table_name::TableNameKey; -use common_procedure::{ProcedureContext, ProcedureWithId, watcher}; +use common_procedure::{ProcedureContext, ProcedureId, ProcedureWithId, watcher}; use common_query::Output; use common_telemetry::info; use common_test_util::recordbatch::check_output_stream; use common_test_util::temp_dir::{TempDir, create_temp_dir}; use common_wal::config::DatanodeWalConfig; +use datatypes::arrow::array::AsArray; use frontend::instance::Instance; use meta_srv::gc::{self, BatchGcProcedure, GcSchedulerOptions, GcTickerRef}; use meta_srv::metasrv::Metasrv; use mito2::gc::GcConfig; use servers::error::Result as ServerResult; use servers::query_handler::sql::SqlQueryHandler; -use session::context::{QueryContext, QueryContextRef}; +use session::Session; +use session::context::{Channel, QueryContext, QueryContextRef}; use store_api::codec::PrimaryKeyEncoding; use store_api::storage::{RegionId, TableId}; use tests_integration::cluster::{GreptimeDbCluster, GreptimeDbClusterBuilder}; @@ -149,6 +151,48 @@ macro_rules! repartition_tests { }; } +#[tokio::test(flavor = "multi_thread")] +async fn test_repartition_procedure_id_session() { + let (cluster, _home_guard) = repartition_cluster("repartition_procedure_id_session").await; + let instance = cluster.fe_instance(); + let session = Session::new(None, Channel::Postgres, Default::default(), 0); + let query_ctx = session.new_query_context(); + run_sql( + instance, + r#"CREATE TABLE repartition_session_test ("id" INT PRIMARY KEY, ts TIMESTAMP TIME INDEX) + PARTITION ON COLUMNS ("id") ("id" < 10, "id" >= 10)"#, + session.new_query_context(), + ) + .await + .unwrap(); + query_ctx.set_last_procedure_id("previous-id".to_string()); + let sql = r#"ALTER TABLE repartition_session_test REPARTITION ("id" < 10) + INTO ("id" < 5, "id" >= 5 AND "id" < 10) WITH (wait=false)"#; + let OutputData::RecordBatches(batches) = run_sql(instance, sql, session.new_query_context()) + .await + .unwrap() + .data + else { + panic!("expected Procedure ID result"); + }; + let batch = batches.iter().next().unwrap(); + let id = batch + .column_by_name("Procedure ID") + .unwrap() + .as_string::() + .value(0); + assert_eq!(query_ctx.last_procedure_id().as_deref(), Some(id)); + let mut watcher = cluster + .metasrv + .procedure_manager() + .procedure_watcher(ProcedureId::parse_str(id).unwrap()) + .unwrap(); + tokio::time::timeout(Duration::from_secs(60), watcher::wait(&mut watcher)) + .await + .unwrap() + .unwrap(); +} + #[tokio::test(flavor = "multi_thread")] async fn test_repartition_physical_metric_without_logical_table() { common_telemetry::init_default_ut_logging(); @@ -247,7 +291,7 @@ async fn test_repartition_physical_metric_without_logical_table() { /// COUNT must use visible rows rather than the full row counts of shared SSTs. #[tokio::test(flavor = "multi_thread")] async fn test_repartition_append_count_file() { - let (cluster, _home_guard) = append_count_cluster("repartition_append_count").await; + let (cluster, _home_guard) = repartition_cluster("repartition_append_count").await; let instance = cluster.fe_instance(); let table = "count_repartition"; prepare_append_count_table(instance, table, "").await; @@ -265,7 +309,7 @@ async fn test_repartition_append_count_file() { #[tokio::test(flavor = "multi_thread")] async fn test_split_append_count_file() { - let (cluster, _home_guard) = append_count_cluster("split_append_count").await; + let (cluster, _home_guard) = repartition_cluster("split_append_count").await; let instance = cluster.fe_instance(); let table = "count_split"; prepare_append_count_table( @@ -288,7 +332,7 @@ async fn test_split_append_count_file() { #[tokio::test(flavor = "multi_thread")] async fn test_repartition_append_count_memtable_file() { - let (cluster, home_guard) = append_count_cluster("repartition_append_count_memtable").await; + let (cluster, home_guard) = repartition_cluster("repartition_append_count_memtable").await; assert!(home_guard.path().is_dir()); let instance = cluster.fe_instance(); let table = "count_repartition_memtable"; @@ -305,7 +349,7 @@ async fn test_repartition_append_count_memtable_file() { assert!(home_guard.path().is_dir()); } -async fn append_count_cluster(name: &str) -> (GreptimeDbCluster, Arc) { +async fn repartition_cluster(name: &str) -> (GreptimeDbCluster, Arc) { common_telemetry::init_default_ut_logging(); let (store_config, _guard) = get_test_store_config(&StorageType::File); let home_dir = Arc::new(create_temp_dir(name)); diff --git a/tests-integration/tests/truncate.rs b/tests-integration/tests/truncate.rs index c1e7735b633..e75dde80fd1 100644 --- a/tests-integration/tests/truncate.rs +++ b/tests-integration/tests/truncate.rs @@ -17,15 +17,23 @@ use std::time::Duration; use client::OutputData; use common_error::ext::ErrorExt; +use common_error::snafu::ResultExt; use common_error::status_code::StatusCode; use common_meta::lock_key::TableLock; use common_procedure::{ LockKey, Procedure, ProcedureId, ProcedureState, ProcedureWithId, Status, watcher, }; +use common_recordbatch::RecordBatches; use datatypes::arrow::array::AsArray; +use frontend::instance::Instance; +use servers::error::{CollectRecordbatchSnafu, Result as ServerResult}; +use servers::query_handler::sql::SqlQueryHandler; +use session::Session; +use session::context::Channel; +use sqlx::{Connection, Executor, Row, Statement}; use tests_integration::cluster::GreptimeDbClusterBuilder; use tests_integration::standalone::GreptimeDbStandaloneBuilder; -use tests_integration::test_util::{execute_sql, execute_sql_and_expect, try_execute_sql}; +use tests_integration::test_util::{execute_sql, execute_sql_and_expect}; use tokio::sync::Notify; /// Holds the table lock until explicitly released, keeping truncate pending. @@ -126,37 +134,41 @@ async fn test_truncate_wait_and_timeout(distributed: bool) { .await .unwrap(); - let output = execute_sql( + let session = Session::new(None, Channel::Mysql, Default::default(), 0); + let query_ctx = session.new_query_context(); + let output = run_sql( &frontend, "TRUNCATE TABLE truncate_wait_test WITH (wait=false, timeout='10s')", + &session, ) - .await; - let OutputData::RecordBatches(batches) = output.data else { - panic!("expected Procedure ID result"); - }; - let batch = batches.iter().next().unwrap(); - let id = batch - .column_by_name("Procedure ID") - .unwrap() - .as_string::() - .value(0); - let id = ProcedureId::parse_str(id).unwrap(); - let state = execute_sql(&frontend, &format!("SELECT procedure_state('{id}')")) - .await - .data - .pretty_print() - .await; - assert!(state.contains(r#""status":"Running""#), "{state}"); + .await + .unwrap(); + let id = procedure_id_from_output(output).to_string(); + assert_eq!(query_ctx.last_procedure_id(), Some(id.clone())); + for sql in [ + format!("SELECT procedure_state('{id}')"), + "SELECT procedure_state()".to_string(), + ] { + let state = run_sql(&frontend, &sql, &session) + .await + .unwrap() + .pretty_print() + .await; + assert!(state.contains(r#""status":"Running""#), "{state}"); + assert_eq!(query_ctx.last_procedure_id(), Some(id.clone())); + } // The same table lock holds a synchronous truncate past its wait deadline. // Give metadata reads and submission a generous budget. Using the same budget // also verifies that the async procedure keeps running after it expires. - let err = try_execute_sql( + let err = run_sql( &frontend, "TRUNCATE TABLE truncate_wait_test WITH (wait=true, timeout='10s')", + &session, ) .await .unwrap_err(); + assert_eq!(query_ctx.last_procedure_id(), Some(id.clone())); if distributed { assert!(matches!( err.status_code(), @@ -222,4 +234,260 @@ async fn test_truncate_wait_and_timeout(distributed: bool) { .pretty_print() .await ); + + for (sql, succeeds) in [ + ("SELECT 1", true), + ("TRUNCATE TABLE truncate_wait_test WITH (wait=true)", true), + ( + "TRUNCATE TABLE nonexistent_session_table WITH (wait=false)", + false, + ), + ("SELECT procedure_state(NULL)", false), + ("SELECT procedure_state('')", false), + ("SELECT procedure_state('invalid-id')", false), + ] { + let result = run_sql(&frontend, sql, &session).await; + assert_eq!(result.is_ok(), succeeds, "{sql}: {result:?}"); + assert_eq!(query_ctx.last_procedure_id(), Some(id.clone())); + } + + let other_session = Session::new(None, Channel::Mysql, Default::default(), 1); + let state = run_sql( + &frontend, + &format!("SELECT procedure_state('{id}')"), + &other_session, + ) + .await + .unwrap() + .pretty_print() + .await; + assert!(state.contains(r#""status":"Done""#), "{state}"); + let err = run_sql(&frontend, "SELECT procedure_state()", &other_session) + .await + .unwrap_err(); + assert!( + err.output_msg() + .contains("No procedure ID has been returned") + ); + + let output = run_sql( + &frontend, + "TRUNCATE TABLE truncate_wait_test WITH (wait=false)", + &session, + ) + .await + .unwrap(); + let new_id = procedure_id_from_output(output); + assert_ne!(new_id.to_string(), id); + assert_eq!(query_ctx.last_procedure_id(), Some(new_id.to_string())); + let mut watcher = manager.procedure_watcher(new_id).unwrap(); + tokio::time::timeout(Duration::from_secs(10), watcher::wait(&mut watcher)) + .await + .unwrap() + .unwrap(); +} + +#[rstest::rstest] +#[case::mysql("MYSQL_SERVER", "mysql")] +#[case::postgres("POSTGRES_SERVER", "postgres")] +#[tokio::test(flavor = "multi_thread")] +async fn test_procedure_state_prepared_and_admin( + #[case] server_name: &str, + #[case] protocol: &str, +) { + let instance = GreptimeDbStandaloneBuilder::new("procedure_state_prepared") + .build() + .await; + let mut opts = instance.opts.clone(); + opts.http.addr = "127.0.0.1:0".into(); + opts.grpc.bind_addr = "127.0.0.1:0".into(); + opts.mysql.addr = "127.0.0.1:0".into(); + opts.postgres.addr = "127.0.0.1:0".into(); + let mut servers = + frontend::server::Services::new(opts, instance.fe_instance().clone(), Default::default()) + .build() + .unwrap(); + servers.start_all().await.unwrap(); + sqlx::any::install_default_drivers(); + let url = format!("{protocol}://{}/public", servers.addr(server_name).unwrap()); + let mut connection = sqlx::AnyConnection::connect(&url).await.unwrap(); + for table in ["prepared_async", "prepared_sync"] { + // Safety: the table names are fixed test-local identifiers. + connection + .execute(sqlx::AssertSqlSafe(format!( + "CREATE TABLE {table} (ts TIMESTAMP TIME INDEX)" + ))) + .await + .unwrap(); + } + + // Prepare once, before any ID exists, and reuse the same statement below. + let statement = connection + .prepare(sqlx::SqlStr::from_static( + "SELECT procedure_state() AS state", + )) + .await + .unwrap(); + let manager = &instance.procedure_manager; + let id: String = connection + .fetch_one("TRUNCATE TABLE prepared_async WITH (wait=false)") + .await + .unwrap() + .get(0); + let mut completed = manager + .procedure_watcher(ProcedureId::parse_str(&id).unwrap()) + .unwrap(); + tokio::time::timeout(Duration::from_secs(10), watcher::wait(&mut completed)) + .await + .unwrap() + .unwrap(); + assert_eq!(prepared_state(&mut connection, &statement).await, "Done"); + + // An explicit ADMIN ID works without history and does not record that ID. + let mut other_connection = sqlx::AnyConnection::connect(&url).await.unwrap(); + for explicit in [false, true, false] { + let sql = if explicit { + format!("ADMIN procedure_state('{id}')") + } else { + "ADMIN procedure_state()".to_string() + }; + // Safety: the SQL contains only a procedure ID returned by this test. + let result = sqlx::query(sqlx::AssertSqlSafe(sql.as_str())) + .fetch_one(&mut other_connection) + .await; + if explicit { + assert_eq!(result.unwrap().get::(0), r#"{"status":"Done"}"#); + } else { + let err = result.err().expect("expected missing procedure ID error"); + assert!( + err.to_string() + .contains("No procedure ID has been returned"), + "{err}" + ); + } + } + other_connection.close().await.unwrap(); + + let table_id = instance + .fe_instance() + .catalog_manager() + .table("greptime", "public", "prepared_async", None) + .await + .unwrap() + .unwrap() + .table_info() + .table_id(); + let started = Arc::new(Notify::new()); + let release = Arc::new(Notify::new()); + let mut blocker = manager + .submit(ProcedureWithId::with_random_id(Box::new( + HoldTruncateTableLock { + table_id, + started: started.clone(), + release: release.clone(), + }, + ))) + .await + .unwrap(); + tokio::time::timeout(Duration::from_secs(10), started.notified()) + .await + .unwrap(); + let new_id: String = connection + .fetch_one("TRUNCATE TABLE prepared_async WITH (wait=false)") + .await + .unwrap() + .get(0); + assert_ne!(id, new_id); + let mut pending = manager + .procedure_watcher(ProcedureId::parse_str(&new_id).unwrap()) + .unwrap(); + // A is Done and B is Running, so retaining A cannot pass this assertion. + assert_eq!(prepared_state(&mut connection, &statement).await, "Running"); + for (sql, expected) in [ + ("ADMIN procedure_state()".to_string(), "Running"), + (format!("ADMIN procedure_state('{id}')"), "Done"), + ("ADMIN procedure_state()".to_string(), "Running"), + ] { + // Safety: the SQL contains only a procedure ID returned by this test. + let state: String = connection + .fetch_one(sqlx::AssertSqlSafe(sql.as_str())) + .await + .unwrap() + .get(0); + assert_eq!(state, format!(r#"{{"status":"{expected}"}}"#), "{sql}"); + } + for sql in [ + "ADMIN procedure_state('a', 'b')", + "ADMIN procedure_state(NULL)", + "ADMIN procedure_state('')", + "ADMIN procedure_state('invalid-id')", + ] { + assert!(connection.fetch_one(sql).await.is_err(), "{sql}"); + } + connection + .execute("TRUNCATE TABLE prepared_sync WITH (wait=true)") + .await + .unwrap(); + assert_eq!(prepared_state(&mut connection, &statement).await, "Running"); + + release.notify_one(); + tokio::time::timeout(Duration::from_secs(10), async { + watcher::wait(&mut blocker).await.unwrap(); + watcher::wait(&mut pending).await.unwrap(); + }) + .await + .unwrap(); + // The ID stays B; only its state changes between executions. + assert_eq!(prepared_state(&mut connection, &statement).await, "Done"); + connection.close().await.unwrap(); + servers.shutdown_all().await.unwrap(); +} + +async fn prepared_state( + connection: &mut sqlx::AnyConnection, + statement: &sqlx::any::AnyStatement, +) -> String { + let state: String = statement + .query() + .fetch_one(connection) + .await + .unwrap() + .get(0); + serde_json::from_str::(&state).unwrap()["status"] + .as_str() + .unwrap() + .to_string() +} + +async fn run_sql( + frontend: &Arc, + sql: &str, + session: &Session, +) -> ServerResult { + let data = frontend + .do_query(sql, session.new_query_context()) + .await + .remove(0)? + .data; + match data { + OutputData::Stream(stream) => Ok(OutputData::RecordBatches( + RecordBatches::try_collect(stream) + .await + .context(CollectRecordbatchSnafu)?, + )), + data => Ok(data), + } +} + +fn procedure_id_from_output(output: OutputData) -> ProcedureId { + let OutputData::RecordBatches(batches) = output else { + panic!("expected Procedure ID result"); + }; + let batch = batches.iter().next().unwrap(); + let id = batch + .column_by_name("Procedure ID") + .unwrap() + .as_string::() + .value(0); + ProcedureId::parse_str(id).unwrap() } diff --git a/tests/cases/standalone/common/function/procedure_state_session.result b/tests/cases/standalone/common/function/procedure_state_session.result new file mode 100644 index 00000000000..2cc17302d2e --- /dev/null +++ b/tests/cases/standalone/common/function/procedure_state_session.result @@ -0,0 +1,124 @@ +-- Each protocol reuses its own connection across requests. +-- SQLNESS PROTOCOL MYSQL +SELECT procedure_state(); + +Failed to parse query result, err: MySqlError { ERROR 1815 (HY000): (EngineExecuteQuery): Execution error: Invalid function args: No procedure ID has been returned in this session; pass an explicit procedure ID } + +-- SQLNESS PROTOCOL POSTGRES +SELECT procedure_state(); + +Failed to execute query, encountered: Error { kind: Db, cause: Some(DbError { severity: "ERROR", parsed_severity: None, code: SqlState(EP0000), message: "Execution error: Invalid function args: No procedure ID has been returned in this session; pass an explicit procedure ID", detail: None, hint: None, position: None, where_: None, schema: None, table: None, column: None, datatype: None, constraint: None, file: None, line: None, routine: None }) } + +CREATE TABLE procedure_state_session (ts TIMESTAMP TIME INDEX); + +Affected Rows: 0 + +-- Synchronous DDL must not record a Procedure ID in the Session. +-- SQLNESS PROTOCOL MYSQL +TRUNCATE TABLE procedure_state_session WITH (wait=true); + +affected_rows: 0 + +-- SQLNESS PROTOCOL MYSQL +SELECT procedure_state(); + +Failed to parse query result, err: MySqlError { ERROR 1815 (HY000): (EngineExecuteQuery): Execution error: Invalid function args: No procedure ID has been returned in this session; pass an explicit procedure ID } + +-- SQLNESS PROTOCOL POSTGRES +TRUNCATE TABLE procedure_state_session WITH (wait=true); + +Affected Rows: 0 + +-- SQLNESS PROTOCOL POSTGRES +SELECT procedure_state(); + +Failed to execute query, encountered: Error { kind: Db, cause: Some(DbError { severity: "ERROR", parsed_severity: None, code: SqlState(EP0000), message: "Execution error: Invalid function args: No procedure ID has been returned in this session; pass an explicit procedure ID", detail: None, hint: None, position: None, where_: None, schema: None, table: None, column: None, datatype: None, constraint: None, file: None, line: None, routine: None }) } + +-- Synchronous DDL preserves the ID returned by the previous request. +-- SQLNESS PROTOCOL MYSQL +-- SQLNESS REPLACE [0-9a-f]{8}-[0-9a-f-]{27} PROC_ID +TRUNCATE TABLE procedure_state_session WITH (wait=false); + ++--------------------------------------+ +| Procedure ID | ++--------------------------------------+ +| PROC_ID | ++--------------------------------------+ + +-- SQLNESS PROTOCOL MYSQL +TRUNCATE TABLE procedure_state_session WITH (wait=true); + +affected_rows: 0 + +-- SQLNESS PROTOCOL MYSQL +SELECT procedure_state() IN ('{"status":"Running"}', '{"status":"Done"}') AS has_state; + ++-----------+ +| has_state | ++-----------+ +| 1 | ++-----------+ + +-- A different physical connection still has no history. +-- SQLNESS PROTOCOL POSTGRES +SELECT procedure_state(); + +Failed to execute query, encountered: Error { kind: Db, cause: Some(DbError { severity: "ERROR", parsed_severity: None, code: SqlState(EP0000), message: "Execution error: Invalid function args: No procedure ID has been returned in this session; pass an explicit procedure ID", detail: None, hint: None, position: None, where_: None, schema: None, table: None, column: None, datatype: None, constraint: None, file: None, line: None, routine: None }) } + +-- SQLNESS PROTOCOL POSTGRES +-- SQLNESS REPLACE [0-9a-f]{8}-[0-9a-f-]{27} PROC_ID +TRUNCATE TABLE procedure_state_session WITH (wait=false); + ++--------------------------------------+ +| Procedure ID | ++--------------------------------------+ +| PROC_ID | ++--------------------------------------+ + +-- SQLNESS PROTOCOL POSTGRES +TRUNCATE TABLE procedure_state_session WITH (wait=true); + +Affected Rows: 0 + +-- SQLNESS PROTOCOL POSTGRES +SELECT procedure_state() IN ('{"status":"Running"}', '{"status":"Done"}') AS has_state; + ++-----------+ +| has_state | ++-----------+ +| t | ++-----------+ + +-- Failed submissions must preserve history. +-- SQLNESS PROTOCOL MYSQL +TRUNCATE TABLE nonexistent_session_table WITH (wait=false); + +Failed to execute query, err: MySqlError { ERROR 1146 (42S02): (TableNotFound): Table not found: greptime.public.nonexistent_session_table } + +-- SQLNESS PROTOCOL MYSQL +SELECT procedure_state() IN ('{"status":"Running"}', '{"status":"Done"}') AS has_state; + ++-----------+ +| has_state | ++-----------+ +| 1 | ++-----------+ + +-- SQLNESS PROTOCOL POSTGRES +TRUNCATE TABLE nonexistent_session_table WITH (wait=false); + +Failed to execute query, encountered: Error { kind: Db, cause: Some(DbError { severity: "ERROR", parsed_severity: None, code: SqlState(E42P01), message: "Table not found: greptime.public.nonexistent_session_table", detail: None, hint: None, position: None, where_: None, schema: None, table: None, column: None, datatype: None, constraint: None, file: None, line: None, routine: None }) } + +-- SQLNESS PROTOCOL POSTGRES +SELECT procedure_state() IN ('{"status":"Running"}', '{"status":"Done"}') AS has_state; + ++-----------+ +| has_state | ++-----------+ +| t | ++-----------+ + +DROP TABLE procedure_state_session; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/function/procedure_state_session.sql b/tests/cases/standalone/common/function/procedure_state_session.sql new file mode 100644 index 00000000000..a225d3a2b03 --- /dev/null +++ b/tests/cases/standalone/common/function/procedure_state_session.sql @@ -0,0 +1,61 @@ +-- Each protocol reuses its own connection across requests. +-- SQLNESS PROTOCOL MYSQL +SELECT procedure_state(); + +-- SQLNESS PROTOCOL POSTGRES +SELECT procedure_state(); + +CREATE TABLE procedure_state_session (ts TIMESTAMP TIME INDEX); + +-- Synchronous DDL must not record a Procedure ID in the Session. +-- SQLNESS PROTOCOL MYSQL +TRUNCATE TABLE procedure_state_session WITH (wait=true); + +-- SQLNESS PROTOCOL MYSQL +SELECT procedure_state(); + +-- SQLNESS PROTOCOL POSTGRES +TRUNCATE TABLE procedure_state_session WITH (wait=true); + +-- SQLNESS PROTOCOL POSTGRES +SELECT procedure_state(); + +-- Synchronous DDL preserves the ID returned by the previous request. +-- SQLNESS PROTOCOL MYSQL +-- SQLNESS REPLACE [0-9a-f]{8}-[0-9a-f-]{27} PROC_ID +TRUNCATE TABLE procedure_state_session WITH (wait=false); + +-- SQLNESS PROTOCOL MYSQL +TRUNCATE TABLE procedure_state_session WITH (wait=true); + +-- SQLNESS PROTOCOL MYSQL +SELECT procedure_state() IN ('{"status":"Running"}', '{"status":"Done"}') AS has_state; + +-- A different physical connection still has no history. +-- SQLNESS PROTOCOL POSTGRES +SELECT procedure_state(); + +-- SQLNESS PROTOCOL POSTGRES +-- SQLNESS REPLACE [0-9a-f]{8}-[0-9a-f-]{27} PROC_ID +TRUNCATE TABLE procedure_state_session WITH (wait=false); + +-- SQLNESS PROTOCOL POSTGRES +TRUNCATE TABLE procedure_state_session WITH (wait=true); + +-- SQLNESS PROTOCOL POSTGRES +SELECT procedure_state() IN ('{"status":"Running"}', '{"status":"Done"}') AS has_state; + +-- Failed submissions must preserve history. +-- SQLNESS PROTOCOL MYSQL +TRUNCATE TABLE nonexistent_session_table WITH (wait=false); + +-- SQLNESS PROTOCOL MYSQL +SELECT procedure_state() IN ('{"status":"Running"}', '{"status":"Done"}') AS has_state; + +-- SQLNESS PROTOCOL POSTGRES +TRUNCATE TABLE nonexistent_session_table WITH (wait=false); + +-- SQLNESS PROTOCOL POSTGRES +SELECT procedure_state() IN ('{"status":"Running"}', '{"status":"Done"}') AS has_state; + +DROP TABLE procedure_state_session; From 2e84833ac8ee199477c57a9fa300a1be6d2d7066 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Tue, 27 Oct 2026 09:23:00 +0000 Subject: [PATCH 106/239] feat(partition): support iceberg_hash in partition rules (#9693) * feat(partition): support iceberg_hash in partition rules Signed-off-by: Dennis Zhuang * fix(partition): reject partition column type changes and quote rule identifiers Signed-off-by: Dennis Zhuang * fix(table): keep time index widening for partition columns Signed-off-by: Dennis Zhuang * refactor(partition): replace PartitionFunction with IcebergHash Signed-off-by: Dennis Zhuang * docs(partition): describe iceberg_hash in the function module Signed-off-by: Dennis Zhuang * fix(meta): validate iceberg_hash partition bounds on ALTER and trim comments Signed-off-by: Dennis Zhuang * test: quote partition rule identifiers in SHOW CREATE TABLE expectations Signed-off-by: Dennis Zhuang * test(partition): check hash rules in information_schema and repartition target columns Signed-off-by: Dennis Zhuang * fix(partition): coerce hash partition rules against the evaluation schema Signed-off-by: Dennis Zhuang * fix(partition): keep quotes in partition column names when building filters Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/common/meta/src/ddl/alter_table.rs | 10 +- .../src/ddl/alter_table/partition_compat.rs | 173 +++------ src/mito2/src/engine/partition_filter_test.rs | 116 +++++- src/operator/src/statement/ddl.rs | 190 +++++++++- .../benches/bench_split_record_batch.rs | 62 +++- src/partition/src/checker.rs | 71 +++- src/partition/src/collider.rs | 133 +++++-- src/partition/src/expr.rs | 205 +++++++---- src/partition/src/function.rs | 345 ++++++++++++++++++ src/partition/src/lib.rs | 1 + src/partition/src/multi_dim.rs | 309 +++++++++++++++- src/partition/src/simplify.rs | 46 ++- .../src/dist_plan/predicate_extractor.rs | 3 +- src/query/src/dist_plan/region_pruner.rs | 55 +++ src/sql/src/parsers/create_parser.rs | 15 +- src/sql/src/statements/create.rs | 11 +- src/table/src/metadata.rs | 54 +++ tests-integration/src/otlp.rs | 32 +- tests-integration/src/tests/instance_test.rs | 8 +- tests-integration/tests/http.rs | 8 +- .../repartition/repartition.result | 24 +- .../create/metric_engine_partition.result | 18 +- .../common/partition_iceberg_hash.result | 184 ++++++++++ .../common/partition_iceberg_hash.sql | 78 ++++ .../standalone/common/show/show_create.result | 18 +- .../cases/partition_iceberg_hash/case.toml | 9 + .../cases/partition_iceberg_hash/setup.sql | 22 ++ .../partition_iceberg_hash/verify.result | 79 ++++ .../cases/partition_iceberg_hash/verify.sql | 19 + 29 files changed, 1947 insertions(+), 351 deletions(-) create mode 100644 src/partition/src/function.rs create mode 100644 tests/cases/standalone/common/partition_iceberg_hash.result create mode 100644 tests/cases/standalone/common/partition_iceberg_hash.sql create mode 100644 tests/compatibility/cases/partition_iceberg_hash/case.toml create mode 100644 tests/compatibility/cases/partition_iceberg_hash/setup.sql create mode 100644 tests/compatibility/cases/partition_iceberg_hash/verify.result create mode 100644 tests/compatibility/cases/partition_iceberg_hash/verify.sql diff --git a/src/common/meta/src/ddl/alter_table.rs b/src/common/meta/src/ddl/alter_table.rs index 232191def88..648f99c624a 100644 --- a/src/common/meta/src/ddl/alter_table.rs +++ b/src/common/meta/src/ddl/alter_table.rs @@ -150,8 +150,6 @@ impl AlterTableProcedure { )?; self.new_table_info = Some(new_table_info); - // Reject schema changes that the persisted partition expressions of some - // region cannot survive, before any region is altered. self.validate_region_partition_exprs().await?; // Safety: Checked in `AlterTableProcedure::new`. @@ -248,9 +246,8 @@ impl AlterTableProcedure { Ok(Status::executing(true)) } - /// Validates the persisted partition expressions of every region against the - /// schema proposed by this ALTER TABLE, so an ALTER that would break a region's - /// partition predicate is rejected before any region is modified. + /// Rejects type changes the persisted partition expressions can't take, before + /// any region is altered. async fn validate_region_partition_exprs(&self) -> Result<()> { // Safety: filled in `fill_table_info`. let old_schema = &self @@ -275,8 +272,7 @@ impl AlterTableProcedure { .get_physical_table_route(self.data.table_id()) .await?; if physical_table_id != self.data.table_id() { - // Partition expressions are persisted on the regions of the physical - // table; the altered schema here belongs to a logical table. + // Partition expressions live on the physical table's regions. return Ok(()); } diff --git a/src/common/meta/src/ddl/alter_table/partition_compat.rs b/src/common/meta/src/ddl/alter_table/partition_compat.rs index 9fc94004d5c..43d30092673 100644 --- a/src/common/meta/src/ddl/alter_table/partition_compat.rs +++ b/src/common/meta/src/ddl/alter_table/partition_compat.rs @@ -12,39 +12,17 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! Preflight validation of persisted partition expressions against an ALTER TABLE schema change. +//! Rejects an ALTER TABLE whose column type changes would break the persisted +//! partition expressions of a region, before any region is altered. //! -//! Persisted partition bounds keep the types they were written with. A column whose -//! type changes while a partition expression references it breaks the region on both -//! paths that consume the bounds: +//! Bounds keep the type they were written with. Row routing compares them with +//! row values as `Value`s, which only stays valid across a timestamp unit +//! widening; scans coerce the bound to the common comparison type, which can +//! overflow (a millisecond bound in year 3000 cast to nanoseconds). //! -//! * On scan, the region coerces the bound literal and the new column type to a -//! common comparison type while building the partition predicate. If no common type -//! exists, or the bound cannot be cast to it (a classic example is a millisecond -//! timestamp bound that overflows when the time index column is widened to -//! nanoseconds), every scan of that region fails at that point. -//! * On write, the region routes rows through `MultiDimPartitionRule::find_region`, -//! which compares the persisted bound against row values directly (`Value::cmp`): a -//! bound whose variant no longer matches the column's panics on ordered comparisons -//! and misroutes silently on equality. Only a timestamp unit widening keeps the two -//! comparable — `Timestamp` comparisons convert units before comparing — so every -//! other type change of a referenced column is rejected outright, and the widening -//! keeps the scan check below. -//! -//! Such an ALTER must be rejected before it is submitted to any region, otherwise -//! already-altered regions keep failing while the rest of the table is untouched. -//! -//! The `partition` crate depends on `common-meta`, so this crate cannot depend on -//! `partition` back and cannot reuse [`partition::expr::PartitionExpr`]. The `Persisted*` -//! types below mirror the JSON representation of `partition::partition::PartitionBound` and -//! `partition::expr::PartitionExpr` and must be kept in sync with `partition/src/partition.rs` -//! and `partition/src/expr.rs`; round-trip tests in `partition` run real serialized -//! expressions through this validator so drift fails in CI. The compatibility check itself -//! deliberately repeats the same conversion (value -> [`ScalarValue`]) and the same -//! comparison coercion the region performs when it builds the partition predicate -//! (`ExprSimplifier::coerce` / [`comparison_coercion`]): both sides are cast to the common -//! comparison type, and the persisted bound must be castable to that type with -//! [`ScalarValue::cast_to`]. A bound that would fail on scan is rejected here as well. +//! `partition` depends on this crate, so the `Persisted*` types mirror the JSON of +//! `partition::partition::PartitionBound` and `partition::expr::PartitionExpr`. +//! A test in `partition` runs real serialized expressions through this validator. use std::collections::HashMap; @@ -62,10 +40,9 @@ use store_api::storage::RegionId; use crate::error::{PartitionExprIncompatibleSnafu, Result}; use crate::rpc::router::RegionRoute; -/// Mirrors `partition::partition::PartitionBound`. #[derive(Debug, Deserialize)] enum PersistedPartitionBound { - /// Deprecated since 0.9.0. Only deserialized to keep unknown old formats parseable. + /// Deprecated since 0.9.0. #[allow(dead_code)] Value(Value), /// Deprecated since 0.15.0. @@ -73,7 +50,6 @@ enum PersistedPartitionBound { Expr(PersistedPartitionExpr), } -/// Mirrors `partition::expr::PartitionExpr`. #[derive(Debug, Deserialize)] struct PersistedPartitionExpr { lhs: Box, @@ -81,15 +57,15 @@ struct PersistedPartitionExpr { rhs: Box, } -/// Mirrors `partition::expr::Operand`. #[derive(Debug, Deserialize)] enum PersistedOperand { Column(String), Value(Value), Expr(PersistedPartitionExpr), + /// `iceberg_hash(column)`. + IcebergHash(String), } -/// Mirrors `partition::expr::RestrictedOp`. #[derive(Debug, Deserialize)] enum PersistedOp { Eq, @@ -102,10 +78,7 @@ enum PersistedOp { Or, } -/// Returns the columns whose data type differs between `old_schema` and `new_schema`, -/// mapped to (old type, new type). Columns only present in one of the schemas (added or -/// dropped) are not included; dropping a column referenced by a partition expression is -/// already rejected by the table metadata validation. +/// Returns the columns present in both schemas whose type differs, as (old, new). pub(crate) fn changed_column_types( old_schema: &Schema, new_schema: &Schema, @@ -126,12 +99,10 @@ pub(crate) fn changed_column_types( changed } -/// Validates that the persisted partition bounds of every region are still representable -/// with the types of `changed_columns`. Returns an error naming the first offending -/// region without modifying anything. +/// Returns an error naming the first region whose partition expression can't take +/// the type changes in `changed_columns`. /// -/// Public so that round-trip tests in the `partition` crate can run real serialized -/// partition expressions through the same mirror used here. +/// Public so that tests in `partition` can run real serialized expressions through it. pub fn validate_region_partition_bounds( changed_columns: &HashMap, region_routes: &[RegionRoute], @@ -149,8 +120,7 @@ pub fn validate_region_partition_bounds( let bound: PersistedPartitionBound = match serde_json::from_str(&expr_json) { Ok(bound) => bound, Err(err) => { - // A malformed persisted expression already fails every scan of this - // region. It is unrelated to the schema change, so don't block the ALTER. + // Such a region already fails every scan, regardless of this ALTER. warn!( "Skipping partition expression validation for region {}, \ the persisted expression is not parseable: {}, expr: {}", @@ -159,9 +129,8 @@ pub fn validate_region_partition_bounds( continue; } }; + // Deprecated bound formats are not used to build partition predicates. let PersistedPartitionBound::Expr(expr) = bound else { - // Deprecated bound formats don't carry expressions and are not used - // when the region builds its partition predicate. continue; }; @@ -187,7 +156,6 @@ fn validate_expr( | PersistedOp::LtEq | PersistedOp::Gt | PersistedOp::GtEq => { - // Comparison between a column and a bound literal in either order. validate_comparison(&expr.lhs, &expr.rhs, changed_columns, region_id)?; validate_comparison(&expr.rhs, &expr.lhs, changed_columns, region_id)?; } @@ -203,7 +171,9 @@ fn validate_operand( ) -> Result<()> { match operand { PersistedOperand::Expr(expr) => validate_expr(expr, changed_columns, region_id), - PersistedOperand::Column(_) | PersistedOperand::Value(_) => Ok(()), + PersistedOperand::Column(_) + | PersistedOperand::Value(_) + | PersistedOperand::IcebergHash(_) => Ok(()), } } @@ -213,83 +183,56 @@ fn validate_comparison( changed_columns: &HashMap, region_id: RegionId, ) -> Result<()> { - let (PersistedOperand::Column(column), PersistedOperand::Value(value)) = (operand, other) - else { - return Ok(()); + let (column, value) = match (operand, other) { + (PersistedOperand::Column(column), PersistedOperand::Value(value)) => (column, value), + // The hash depends on the column type, so any change reroutes rows. Table + // metadata rejects such changes first; this covers stale partition columns. + (PersistedOperand::IcebergHash(column), PersistedOperand::Value(_)) => { + return match changed_columns.get(column) { + Some((_, target_type)) => incompatible(region_id, column, target_type), + None => Ok(()), + }; + } + _ => return Ok(()), }; let Some((old_type, target_type)) = changed_columns.get(column) else { return Ok(()); }; - // Values not representable as a partition bound (lists, structs, json) are - // rejected by the partition parser and cannot appear in a persisted expression. let Some(scalar) = value_to_scalar_value(value) else { return Ok(()); }; - // The write path (`MultiDimPartitionRule::find_region`) compares the persisted - // bound against row values directly: `Value::cmp` panics on mismatched variants - // and equality misroutes silently. Only a timestamp unit widening keeps the - // bound and the rows comparable (`Timestamp` comparisons convert units first), - // so every other type change of a referenced column is rejected here. if !old_type.is_timestamp_unit_widening_to(target_type) { - warn!( - "Persisted partition bound would be compared against values of a different \ - type after the ALTER, region_id: {}, column: '{}', target_type: {}", - region_id, column, target_type - ); - return PartitionExprIncompatibleSnafu { - region_id, - column: column.clone(), - target_type: target_type.to_string(), - } - .fail(); + return incompatible(region_id, column, target_type); } - // For the widening, mirror `ExprSimplifier::coerce` in `PredicateGroup::new`: - // the region casts both the column and the persisted bound to the common - // comparison type. The bound is not cast to the column type itself, so the - // widening is accepted as long as the bound survives the common type; a bound - // that cannot be cast to it — or a pair of types without a common comparison - // type — fails every scan. - let target_arrow_type = target_type.as_arrow_type(); - let Some(common_type) = comparison_coercion(&target_arrow_type, &scalar.data_type()) else { - warn!( - "Persisted partition bound has no common comparison type with the altered \ - column type, region_id: {}, column: '{}', target_type: {}", - region_id, column, target_type - ); - return PartitionExprIncompatibleSnafu { - region_id, - column: column.clone(), - target_type: target_type.to_string(), - } - .fail(); + // Mirrors the coercion `PredicateGroup::new` applies on scan: the bound must + // cast to the common comparison type. + let Some(common_type) = comparison_coercion(&target_type.as_arrow_type(), &scalar.data_type()) + else { + return incompatible(region_id, column, target_type); }; if scalar.data_type() != common_type && let Err(error) = scalar.cast_to(&common_type) { - warn!( - error; - "Persisted partition bound is not representable with the altered column type, \ - region_id: {}, column: '{}', target_type: {}", - region_id, - column, - target_type - ); - return PartitionExprIncompatibleSnafu { - region_id, - column: column.clone(), - target_type: target_type.to_string(), - } - .fail(); + warn!(error; "Partition bound of region {} can't cast to {}", region_id, common_type); + return incompatible(region_id, column, target_type); } Ok(()) } -/// Converts a persisted bound value to a [`ScalarValue`], mirroring -/// `partition::expr::Operand::try_as_logical_expr`. Returns `None` for values that -/// cannot appear in a persisted partition expression. +fn incompatible(region_id: RegionId, column: &str, target_type: &ConcreteDataType) -> Result<()> { + PartitionExprIncompatibleSnafu { + region_id, + column, + target_type: target_type.to_string(), + } + .fail() +} + +/// Converts a bound the same way as `partition::expr::Operand::try_as_logical_expr`. +/// Returns `None` for values that can't be partition bounds. fn value_to_scalar_value(value: &Value) -> Option { Some(match value { Value::Boolean(v) => ScalarValue::Boolean(Some(*v)), @@ -508,10 +451,8 @@ mod tests { #[test] fn test_field_narrowing_is_rejected_for_persisted_bounds() { - // BIGINT bound outside the INT range with the column narrowed to INT: the - // scan would coerce both sides to Int64, but the region's write path - // compares the persisted Int64 bound against Int32 rows directly and - // panics, so the ALTER must be rejected. + // Scans would coerce both sides, but row routing compares Int64 bounds with + // Int32 rows. let bound = r#"{"Expr":{"lhs":{"Column":"v"},"op":"Lt","rhs":{"Value":{"Int64":3000000000}}}}"#; let err = validate_region_partition_bounds( @@ -528,9 +469,7 @@ mod tests { #[test] fn test_non_timestamp_widening_is_rejected_for_persisted_bounds() { - // Widening INT to BIGINT still changes the variant rows carry: the - // region's write path compares the persisted Int32 bound against Int64 - // rows directly and panics, so the ALTER must be rejected as well. + // Row routing would compare Int32 bounds with Int64 rows. let bound = r#"{"Expr":{"lhs":{"Column":"v"},"op":"Lt","rhs":{"Value":{"Int32":10}}}}"#; let err = validate_region_partition_bounds( &changed( @@ -546,9 +485,7 @@ mod tests { #[test] fn test_bound_without_common_comparison_type_is_rejected() { - // A boolean bound has no comparison coercion with a timestamp column, so - // even a widening whose bounds are normally representable fails every - // predicate build; the ALTER must be rejected. + // No common comparison type between a boolean and a timestamp. let bound = r#"{"Expr":{"lhs":{"Column":"ts"},"op":"Lt","rhs":{"Value":{"Boolean":true}}}}"#; let err = validate_region_partition_bounds( diff --git a/src/mito2/src/engine/partition_filter_test.rs b/src/mito2/src/engine/partition_filter_test.rs index b71559c15e9..70a4b7c487f 100644 --- a/src/mito2/src/engine/partition_filter_test.rs +++ b/src/mito2/src/engine/partition_filter_test.rs @@ -16,8 +16,10 @@ use api::v1::Rows; use common_recordbatch::RecordBatches; +use datatypes::arrow::array::AsArray; use datatypes::value::Value; -use partition::expr::col; +use partition::expr::{Operand, col}; +use partition::function::IcebergHash; use store_api::region_engine::RegionEngine; use store_api::region_request::{ EnterStagingRequest, RegionFlushRequest, RegionRequest, StagingPartitionDirective, @@ -168,3 +170,115 @@ async fn test_partition_filter_basic_with_format(flat_format: bool) { flat_format ); } + +#[tokio::test] +async fn test_partition_filter_iceberg_hash() { + test_partition_filter_iceberg_hash_with_format(false).await; + test_partition_filter_iceberg_hash_with_format(true).await; +} + +/// Narrows a region to `iceberg_hash(tag_0) < 0` and checks rows flushed under +/// the old partition expr are filtered by the hash. +async fn test_partition_filter_iceberg_hash_with_format(flat_format: bool) { + let mut env = TestEnv::new().await; + let engine = env + .create_engine(MitoConfig { + default_flat_format: flat_format, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1024, 0); + let request = CreateRequestBuilder::new() + .partition_expr_json(Some(range_expr_string("field_0", 0., 99.))) + .build(); + let column_schemas = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + let flush = || async { + engine + .handle_request( + region_id, + RegionRequest::Flush(RegionFlushRequest::default()), + ) + .await + .unwrap(); + }; + put_rows( + &engine, + region_id, + Rows { + schema: column_schemas.clone(), + rows: build_rows(0, 20), + }, + ) + .await; + flush().await; + + let in_partition = |i: &usize| { + IcebergHash::evaluate_value(&Value::String(i.to_string().as_str().into())).unwrap() + < Value::Int32(0) + }; + let hash_expr = Operand::IcebergHash(IcebergHash::new("tag_0")).lt(Value::Int32(0)); + engine + .handle_request( + region_id, + RegionRequest::EnterStaging(EnterStagingRequest { + partition_directive: StagingPartitionDirective::UpdatePartitionExpr( + hash_expr.as_json_str().unwrap(), + ), + }), + ) + .await + .unwrap(); + put_rows( + &engine, + region_id, + Rows { + schema: column_schemas, + rows: (20..40) + .filter(in_partition) + .flat_map(|i| build_rows(i, i + 1)) + .collect(), + }, + ) + .await; + flush().await; + use store_api::region_engine::SettableRegionRoleState; + engine + .set_region_role_state_gracefully(region_id, SettableRegionRoleState::Leader) + .await + .unwrap(); + + let request = ScanRequest { + projection: Some(vec![0]), + ..Default::default() + }; + let scanner = engine.scanner(region_id, request).await.unwrap(); + let batches = RecordBatches::try_collect(scanner.scan().await.unwrap()) + .await + .unwrap(); + let mut actual = vec![]; + for batch in batches.iter() { + let tags = datatypes::arrow::compute::cast( + batch.column(0), + &datatypes::arrow::datatypes::DataType::Utf8, + ) + .unwrap(); + actual.extend( + tags.as_string::() + .iter() + .map(|v| v.unwrap().to_string()), + ); + } + actual.sort(); + let mut expected = (0..40) + .filter(in_partition) + .map(|i| i.to_string()) + .collect::>(); + expected.sort(); + // Rows flushed before staging fall on both sides of the hash bound. + assert!((0..20).any(|i| in_partition(&i)) && !(0..20).all(|i| in_partition(&i))); + assert_eq!(actual, expected, "flat_format: {flat_format}"); +} diff --git a/src/operator/src/statement/ddl.rs b/src/operator/src/statement/ddl.rs index e14256aaffe..c19be914a73 100644 --- a/src/operator/src/statement/ddl.rs +++ b/src/operator/src/statement/ddl.rs @@ -72,6 +72,7 @@ use datatypes::value::Value; use datatypes::vectors::{StringVector, VectorRef}; use humantime::parse_duration; use partition::expr::{Operand, PartitionExpr, RestrictedOp}; +use partition::function::IcebergHash; use partition::multi_dim::MultiDimPartitionRule; use query::parser::QueryStatement; use query::plan::extract_and_rewrite_full_table_names; @@ -92,7 +93,10 @@ use sql::statements::create::{ CreateExternalTable, CreateFlow, CreateTable, CreateTableLike, CreateView, Partitions, }; use sql::statements::statement::Statement; -use sqlparser::ast::{Expr, Ident, UnaryOperator, Value as ParserValue}; +use sqlparser::ast::{ + Expr, Function, FunctionArg, FunctionArgExpr, FunctionArguments, Ident, UnaryOperator, + Value as ParserValue, +}; use store_api::metric_engine_consts::{ LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME, PHYSICAL_TABLE_METADATA_KEY, }; @@ -2919,6 +2923,15 @@ fn ensure_partition_operand_columns_in_target( ensure_partition_operand_columns_in_target(&expr.lhs, target_partition_columns)?; ensure_partition_operand_columns_in_target(&expr.rhs, target_partition_columns)?; } + Operand::IcebergHash(hash) => ensure!( + target_partition_columns.contains(&hash.column().to_string()), + InvalidPartitionRuleSnafu { + reason: format!( + "partition expression references column '{}' that is not in target partition columns", + hash.column() + ) + } + ), Operand::Value(_) => {} } @@ -2970,6 +2983,22 @@ fn convert_one_expr( let value = convert_value(&v.value, data_type, timezone, Some(*unary_op))?; (Operand::Value(value), op, Operand::Column(column_name)) } + // func, val + (Expr::Function(function), literal) + if !matches!(op, RestrictedOp::And | RestrictedOp::Or) => + { + let hash = convert_iceberg_hash(function, column_name_and_type)?; + let value = convert_literal(literal, IcebergHash::output_type(), timezone)?; + (Operand::IcebergHash(hash), op, Operand::Value(value)) + } + // val, func + (literal, Expr::Function(function)) + if !matches!(op, RestrictedOp::And | RestrictedOp::Or) => + { + let hash = convert_iceberg_hash(function, column_name_and_type)?; + let value = convert_literal(literal, IcebergHash::output_type(), timezone)?; + (Operand::Value(value), op, Operand::IcebergHash(hash)) + } (Expr::BinaryOp { .. }, Expr::BinaryOp { .. }) => { // sub-expr must against another sub-expr let lhs = convert_one_expr(left, column_name_and_type, timezone)?; @@ -2987,6 +3016,70 @@ fn convert_one_expr( Ok(PartitionExpr::new(lhs, op, rhs)) } +/// Converts an `iceberg_hash` call, which takes exactly one partition column. +fn convert_iceberg_hash( + function: &Function, + column_name_and_type: &HashMap<&String, ConcreteDataType>, +) -> Result { + let invalid = || { + InvalidPartitionRuleSnafu { + reason: format!("unsupported partition function call {function}"), + } + .build() + }; + let FunctionArguments::List(args) = &function.args else { + return Err(invalid()); + }; + let [FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::Identifier(ident)))] = + args.args.as_slice() + else { + return Err(invalid()); + }; + if function.filter.is_some() + || function.over.is_some() + || function.null_treatment.is_some() + || !function.within_group.is_empty() + || !matches!(function.parameters, FunctionArguments::None) + || args.duplicate_treatment.is_some() + || !args.clauses.is_empty() + { + return Err(invalid()); + } + ensure!( + function + .name + .to_string() + .eq_ignore_ascii_case(IcebergHash::NAME), + InvalidPartitionRuleSnafu { + reason: format!("unsupported partition function {}", function.name), + } + ); + let (column_name, data_type) = convert_identifier(ident, column_name_and_type)?; + ensure!( + IcebergHash::accepts(&data_type), + InvalidPartitionRuleSnafu { + reason: format!( + "partition function {} does not support column type {data_type}", + IcebergHash::NAME + ), + } + ); + Ok(IcebergHash::new(column_name)) +} + +fn convert_literal(expr: &Expr, data_type: ConcreteDataType, timezone: &Timezone) -> Result { + match expr { + Expr::Value(value) => convert_value(&value.value, data_type, timezone, None), + Expr::UnaryOp { op, expr } if let Expr::Value(value) = &**expr => { + convert_value(&value.value, data_type, timezone, Some(*op)) + } + _ => InvalidPartitionRuleSnafu { + reason: format!("expect a literal in partition expr, got {expr}"), + } + .fail(), + } +} + fn convert_identifier( ident: &Ident, column_name_and_type: &HashMap<&String, ConcreteDataType>, @@ -3748,6 +3841,66 @@ SELECT max(c1), min(c2) FROM schema_2.table_2;"; assert!(!NAME_PATTERN_REG.is_match("#")); } + #[test] + fn test_convert_iceberg_hash() { + let names = ["host", "count", "device", "measure", "host-name"].map(String::from); + let column_name_and_type = HashMap::from([ + (&names[0], ConcreteDataType::string_datatype()), + (&names[1], ConcreteDataType::int16_datatype()), + (&names[2], ConcreteDataType::uint64_datatype()), + (&names[3], ConcreteDataType::float64_datatype()), + (&names[4], ConcreteDataType::string_datatype()), + ]); + let timezone = Timezone::from_tz_string("UTC").unwrap(); + let dialect = GreptimeDbDialect {}; + let parse = |sql: &str| { + Parser::new(&dialect) + .try_with_sql(sql) + .unwrap() + .parse_expr() + .unwrap() + }; + + for sql in [ + "iceberg_hash(host) < 0", + "ICEBERG_HASH(device) >= -1073741824", + "0 > iceberg_hash(count)", + "iceberg_hash(host) < 100 AND host < 'm'", + r#"iceberg_hash("host-name") < 0 AND "host-name" < 'm'"#, + ] { + let expr = convert_one_expr(&parse(sql), &column_name_and_type, &timezone).unwrap(); + // SHOW CREATE TABLE prints the rule back as SQL, quoting identifiers. + let mut partitions = Partitions { + column_list: vec![], + exprs: vec![expr.to_parser_expr()], + }; + partitions.set_quote('"'); + let restored = parse(&partitions.exprs[0].to_string()); + assert_eq!( + convert_one_expr(&restored, &column_name_and_type, &timezone).unwrap(), + expr, + "{sql}" + ); + } + + for sql in [ + "iceberg_hash(measure) < 0", + "iceberg_hash(host, count) < 0", + "iceberg_hash('a') < 0", + "iceberg_hash(missing) < 0", + "iceberg_hash(DISTINCT host) < 0", + "md5(host) < 0", + "iceberg_hash(host) < host", + "iceberg_hash(host) < 'abc'", + "iceberg_hash(host) < 2147483648", + ] { + assert!( + convert_one_expr(&parse(sql), &column_name_and_type, &timezone).is_err(), + "{sql}" + ); + } + } + #[test] fn test_partition_expr_equivalence_with_swapped_operands() { let column_name = "device_id".to_string(); @@ -3818,25 +3971,28 @@ SELECT max(c1), min(c2) FROM schema_2.table_2;"; (&area, ConcreteDataType::string_datatype()), ]); let dialect = GreptimeDbDialect {}; - let mut parser = Parser::new(&dialect) - .try_with_sql("device_id >= 100") - .unwrap(); - let remaining_old_expr = convert_one_expr( - &parser.parse_expr().unwrap(), - &column_name_and_type, - &timezone, - ) - .unwrap(); let target_partition_columns = HashSet::from([&area]); + for sql in ["device_id >= 100", "iceberg_hash(device_id) >= 0"] { + let mut parser = Parser::new(&dialect).try_with_sql(sql).unwrap(); + let remaining_old_expr = convert_one_expr( + &parser.parse_expr().unwrap(), + &column_name_and_type, + &timezone, + ) + .unwrap(); - let err = ensure_partition_expr_columns_in_target( - &[remaining_old_expr], - &target_partition_columns, - ) - .unwrap_err(); + let err = ensure_partition_expr_columns_in_target( + &[remaining_old_expr], + &target_partition_columns, + ) + .unwrap_err(); - assert!(err.to_string().contains("device_id")); - assert!(err.to_string().contains("target partition columns")); + assert!(err.to_string().contains("device_id"), "{sql}"); + assert!( + err.to_string().contains("target partition columns"), + "{sql}" + ); + } } #[test] diff --git a/src/partition/benches/bench_split_record_batch.rs b/src/partition/benches/bench_split_record_batch.rs index 86eed0c981e..caee564c7cb 100644 --- a/src/partition/benches/bench_split_record_batch.rs +++ b/src/partition/benches/bench_split_record_batch.rs @@ -22,7 +22,8 @@ use datatypes::arrow::datatypes::{DataType, Field, Schema, TimeUnit}; use datatypes::arrow::record_batch::RecordBatch; use datatypes::value::Value; use partition::PartitionRule; -use partition::expr::{Operand, col}; +use partition::expr::{Operand, PartitionExpr, col}; +use partition::function::IcebergHash; use partition::multi_dim::MultiDimPartitionRule; use rand::Rng; use store_api::storage::RegionNumber; @@ -219,9 +220,66 @@ fn bench_split_record_batch_vs_row(c: &mut Criterion) { group.finish(); } +/// Splits `dim` into `num_partitions` ranges, `bounds` being the inner boundaries. +fn range_rule( + column: &str, + dim: Operand, + bounds: impl IntoIterator, +) -> MultiDimPartitionRule { + let bounds = bounds.into_iter().map(Value::Int32).collect::>(); + let mut exprs: Vec = vec![dim.clone().lt(bounds[0].clone())]; + for pair in bounds.windows(2) { + exprs.push( + dim.clone() + .gt_eq(pair[0].clone()) + .and(dim.clone().lt(pair[1].clone())), + ); + } + exprs.push(dim.gt_eq(bounds[bounds.len() - 1].clone())); + let regions = (0..exprs.len()).map(|v| v as u32).collect(); + MultiDimPartitionRule::try_new(vec![column.to_string()], regions, exprs, true).unwrap() +} + +fn hash_rule(column: &str, num_partitions: i64) -> MultiDimPartitionRule { + let width = (1i64 << 32) / num_partitions; + range_rule( + column, + Operand::IcebergHash(IcebergHash::new(column)), + (1..num_partitions).map(|k| (i32::MIN as i64 + k * width) as i32), + ) +} + +fn bench_split_record_batch_hash(c: &mut Criterion) { + let mut group = c.benchmark_group("split_record_batch_hash"); + + for num_partitions in [4, 16, 64] { + for num_rows in [1000, 100000] { + let batch = create_test_batch(num_rows); + // Baseline: the same number of plain column ranges over a0 in [0, 100). + let baseline = range_rule( + "a0", + col("a0"), + (1..num_partitions).map(|k| (k * 100 / num_partitions) as i32), + ); + group.bench_function(format!("column_range_{num_partitions}_{num_rows}"), |b| { + b.iter(|| black_box(baseline.split_record_batch(black_box(&batch))).unwrap()) + }); + for (key, column) in [("int", "a0"), ("string", "a1")] { + let rule = hash_rule(column, num_partitions); + group.bench_function(format!("hash_{key}_{num_partitions}_{num_rows}"), |b| { + b.iter(|| black_box(rule.split_record_batch(black_box(&batch))).unwrap()) + }); + } + } + } + + group.finish(); +} + criterion_group!( benches, bench_split_record_batch_naive_vs_optimized, - bench_split_record_batch_vs_row + bench_split_record_batch_vs_row, + bench_split_record_batch_hash ); criterion_main!(benches); diff --git a/src/partition/src/checker.rs b/src/partition/src/checker.rs index c88941e4df7..7f4d55036d9 100644 --- a/src/partition/src/checker.rs +++ b/src/partition/src/checker.rs @@ -19,7 +19,7 @@ use datatypes::arrow::array::{BooleanArray, Float64Array, Float64Builder, Record use datatypes::arrow::datatypes::{DataType, Field, Schema}; use datatypes::value::OrderedF64; -use crate::collider::{CHECK_STEP, Collider, NORMALIZE_STEP}; +use crate::collider::{CHECK_STEP, Collider, NORMALIZE_STEP, PartitionDim}; use crate::error::{ CheckpointNotCoveredSnafu, CheckpointOverlappedSnafu, DuplicateExprSnafu, Result, }; @@ -29,12 +29,26 @@ use crate::multi_dim::MultiDimPartitionRule; pub struct PartitionChecker<'a> { rule: &'a MultiDimPartitionRule, collider: Collider<'a>, + /// Dimensions in the matrix. A dimension's field name in the matrix is its + /// index here, since a column may be named like a function call. + dims: Vec, } impl<'a> PartitionChecker<'a> { pub fn try_new(rule: &'a MultiDimPartitionRule) -> Result { let collider = Collider::new(rule.exprs())?; - Ok(Self { rule, collider }) + let mut dims = collider + .normalized_values + .iter() + .filter(|(_, values)| !values.is_empty()) + .map(|(dim, _)| dim.clone()) + .collect::>(); + dims.sort_unstable(); + Ok(Self { + rule, + collider, + dims, + }) } pub fn check(&self) -> Result<()> { @@ -76,19 +90,19 @@ impl<'a> PartitionChecker<'a> { // TODO(ruihang): merge atomic exprs to improve checker's performance // matrix test + let field_names = (0..self.dims.len()) + .map(|index| index.to_string()) + .collect::>(); let mut matrix_foundation = HashMap::new(); - for (col, values) in self.collider.normalized_values.iter() { - if values.is_empty() { - continue; - } - + for (dim, field_name) in self.dims.iter().zip(&field_names) { + let values = &self.collider.normalized_values[dim]; let mut cornerstones = Vec::with_capacity(values.len() * 2 + 1); cornerstones.push(values[0].1 - CHECK_STEP); for value in values { cornerstones.push(value.1); cornerstones.push(value.1 + CHECK_STEP); } - matrix_foundation.insert(col.as_str(), cornerstones); + matrix_foundation.insert(field_name.as_str(), cornerstones); } // If there are no values, the rule is empty and valid. @@ -104,7 +118,12 @@ impl<'a> PartitionChecker<'a> { .collider .atomic_exprs .iter() - .map(|expr| expr.to_physical_expr(matrix_generator.schema())) + .map(|expr| { + expr.to_physical_expr(matrix_generator.schema(), |dim| { + // Every nucleon's dimension has a value, so it is in `dims`. + self.dims.binary_search(dim).unwrap().to_string() + }) + }) .collect::>(); for batch in matrix_generator { results.clear(); @@ -158,14 +177,20 @@ impl<'a> PartitionChecker<'a> { let mut check_point = String::new(); let schema = batch.schema(); for (col_index, normalized_value) in normalized_row.iter().enumerate() { - let col_name = schema.field(col_index).name(); + let dim = schema + .field(col_index) + .name() + .parse::() + .ok() + .and_then(|index| self.dims.get(index)); + let col_name = dim.map(ToString::to_string).unwrap_or_default(); if col_index > 0 { check_point.push_str(", "); } // Check if point is on NORMALIZE_STEP or between steps - if let Some(values) = self.collider.normalized_values.get(col_name) { + if let Some(values) = dim.and_then(|dim| self.collider.normalized_values.get(dim)) { let normalize_step = NORMALIZE_STEP.0; // Check if the normalized value is on a NORMALIZE_STEP boundary @@ -350,6 +375,26 @@ mod tests { use crate::expr::col; use crate::multi_dim::MultiDimPartitionRule; + #[test] + fn test_hash_and_same_named_column_are_distinct() { + use crate::expr::Operand; + use crate::function::IcebergHash; + + let hash = IcebergHash::new("host"); + // A column can be named like the hash call, but it is another dimension, so these + // two rules don't cover each other's complement. + let result = MultiDimPartitionRule::try_new( + vec!["host".to_string(), hash.to_string()], + vec![1, 2], + vec![ + Operand::IcebergHash(hash.clone()).lt(Value::Int32(0)), + col(hash.to_string()).gt_eq(Value::Int32(0)), + ], + true, + ); + assert!(result.is_err()); + } + #[test] fn test_matrix_generator_single_column() { let mut matrix_foundation = HashMap::new(); @@ -582,8 +627,8 @@ mod tests { let checker = PartitionChecker::try_new(&rule).unwrap(); let schema = Arc::new(Schema::new(vec![ - Field::new("host", DataType::Float64, false), - Field::new("value", DataType::Float64, false), + Field::new("0", DataType::Float64, false), + Field::new("1", DataType::Float64, false), ])); let host_array = Float64Array::from(vec![-0.5, 0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5]); let value_array = Float64Array::from(vec![-0.5, 0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5]); diff --git a/src/partition/src/collider.rs b/src/partition/src/collider.rs index c426e84575f..fa7a02ddf09 100644 --- a/src/partition/src/collider.rs +++ b/src/partition/src/collider.rs @@ -26,7 +26,7 @@ //! force inside nucleons. use std::collections::HashMap; -use std::fmt::Debug; +use std::fmt::{Debug, Display, Formatter}; use std::sync::Arc; use datafusion_expr::Operator; @@ -38,11 +38,45 @@ use datatypes::value::{OrderedF64, OrderedFloat, Value}; use crate::error; use crate::error::Result; use crate::expr::{Operand, PartitionExpr, RestrictedOp}; +use crate::function::IcebergHash; const ZERO: OrderedF64 = OrderedFloat(0.0f64); pub(crate) const NORMALIZE_STEP: OrderedF64 = OrderedFloat(1.0f64); pub(crate) const CHECK_STEP: OrderedF64 = OrderedFloat(0.5f64); +/// A dimension of the partition space: a column, or the hash of a column. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum PartitionDim { + Column(String), + IcebergHash(IcebergHash), +} + +impl PartitionDim { + fn from_operand(operand: &Operand) -> Option { + match operand { + Operand::Column(column) => Some(Self::Column(column.clone())), + Operand::IcebergHash(hash) => Some(Self::IcebergHash(hash.clone())), + Operand::Value(_) | Operand::Expr(_) => None, + } + } + + pub fn to_operand(&self) -> Operand { + match self { + Self::Column(column) => Operand::Column(column.clone()), + Self::IcebergHash(hash) => Operand::IcebergHash(hash.clone()), + } + } +} + +impl Display for PartitionDim { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + match self { + Self::Column(column) => write!(f, "{column}"), + Self::IcebergHash(hash) => write!(f, "{hash}"), + } + } +} + /// Represents an "atomic" Expression, which isn't composed (OR-ed) of other expressions. #[allow(unused)] #[derive(Debug, Clone, PartialEq, Eq)] @@ -55,10 +89,15 @@ pub struct AtomicExpr { } impl AtomicExpr { - pub fn to_physical_expr(&self, schema: &Schema) -> Arc { + /// Builds the conjunction over `schema`, where `field_name` maps each dimension to its field. + pub fn to_physical_expr( + &self, + schema: &Schema, + field_name: impl Fn(&PartitionDim) -> String, + ) -> Arc { let mut exprs = Vec::with_capacity(self.nucleons.len()); for nucleon in &self.nucleons { - exprs.push(nucleon.to_physical_expr(schema)); + exprs.push(nucleon.to_physical_expr(schema, &field_name(&nucleon.column))); } let result: Arc = exprs .into_iter() @@ -79,23 +118,23 @@ impl PartialOrd for AtomicExpr { /// This struct is used to compose [`AtomicExpr`], hence "nucleon". #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] pub struct NucleonExpr { - column: String, + column: PartitionDim, op: GluonOp, /// Normalized [`Value`]. value: OrderedF64, } impl NucleonExpr { - pub fn to_physical_expr(&self, schema: &Schema) -> Arc { + pub fn to_physical_expr(&self, schema: &Schema, field_name: &str) -> Arc { Arc::new(BinaryExpr::new( - col(&self.column, schema).unwrap(), + col(field_name, schema).unwrap(), self.op.to_operator(), lit(*self.value.as_ref()), )) } - /// Get the column name - pub fn column(&self) -> &str { + /// Get the dimension + pub fn column(&self) -> &PartitionDim { &self.column } @@ -109,12 +148,8 @@ impl NucleonExpr { &self.op } - pub fn new(column: impl Into, op: GluonOp, value: OrderedF64) -> Self { - Self { - column: column.into(), - op, - value, - } + pub fn new(column: PartitionDim, op: GluonOp, value: OrderedF64) -> Self { + Self { column, op, value } } } @@ -153,22 +188,22 @@ pub struct Collider<'a> { source_exprs: &'a [PartitionExpr], pub atomic_exprs: Vec, - /// A map of column name to a list of `(value, normalized value)` pairs. + /// A map of dimension to a list of `(value, normalized value)` pairs. /// /// The normalized value is used for comparison. The normalization process keeps the order of the values. - pub normalized_values: HashMap>, + pub normalized_values: HashMap>, } impl<'a> Collider<'a> { pub fn new(source_exprs: &'a [PartitionExpr]) -> Result { // first walk to collect all values - let mut values: HashMap> = HashMap::new(); + let mut values: HashMap> = HashMap::new(); for expr in source_exprs { Self::collect_column_values_from_expr(expr, &mut values)?; } // normalize values, assumes all values on a column are the same type - let mut normalized_values: HashMap> = + let mut normalized_values: HashMap> = HashMap::with_capacity(values.len()); for (column, mut column_values) in values { column_values.sort_unstable(); @@ -215,13 +250,14 @@ impl<'a> Collider<'a> { /// Helper to collect values with their associated columns from an expression fn collect_column_values_from_expr( expr: &PartitionExpr, - values: &mut HashMap>, + values: &mut HashMap>, ) -> Result<()> { // Handle binary operations between column and value match (&*expr.lhs, &*expr.rhs) { - (Operand::Column(col), Operand::Value(val)) - | (Operand::Value(val), Operand::Column(col)) => { - values.entry(col.clone()).or_default().push(val.clone()); + (dim, Operand::Value(val)) | (Operand::Value(val), dim) + if let Some(dim) = PartitionDim::from_operand(dim) => + { + values.entry(dim).or_default().push(val.clone()); Ok(()) } (Operand::Expr(left_expr), Operand::Expr(right_expr)) => { @@ -241,7 +277,7 @@ impl<'a> Collider<'a> { fn collide_expr( expr: &PartitionExpr, index: usize, - normalized_values: &HashMap>, + normalized_values: &HashMap>, result: &mut Vec, ) -> Result<()> { match expr.op { @@ -300,7 +336,7 @@ impl<'a> Collider<'a> { fn collect_nucleons_from_expr( expr: &PartitionExpr, nucleons: &mut Vec, - normalized_values: &HashMap>, + normalized_values: &HashMap>, ) -> Result<()> { match expr.op { RestrictedOp::And => { @@ -325,7 +361,7 @@ impl<'a> Collider<'a> { fn collect_nucleons_from_operand( operand: &Operand, nucleons: &mut Vec, - normalized_values: &HashMap>, + normalized_values: &HashMap>, ) -> Result<()> { match operand { Operand::Expr(expr) => { @@ -346,7 +382,7 @@ impl<'a> Collider<'a> { lhs: &Operand, op: &RestrictedOp, rhs: &Operand, - normalized_values: &HashMap>, + normalized_values: &HashMap>, ) -> Result { let gluon_op = match op { RestrictedOp::Eq => GluonOp::Eq, @@ -365,19 +401,19 @@ impl<'a> Collider<'a> { }; match (lhs, rhs) { - (Operand::Column(col), Operand::Value(val)) => { - if let Some(column_values) = normalized_values.get(col) + (dim, Operand::Value(val)) if let Some(dim) = PartitionDim::from_operand(dim) => { + if let Some(column_values) = normalized_values.get(&dim) && let Some(&normalized_val) = column_values.get(val) { return Ok(NucleonExpr { - column: col.clone(), + column: dim, op: gluon_op, value: normalized_val, }); } } - (Operand::Value(val), Operand::Column(col)) => { - if let Some(column_values) = normalized_values.get(col) + (Operand::Value(val), dim) if let Some(dim) = PartitionDim::from_operand(dim) => { + if let Some(column_values) = normalized_values.get(&dim) && let Some(&normalized_val) = column_values.get(val) { // Flip the operation for value op column @@ -389,7 +425,7 @@ impl<'a> Collider<'a> { op => op, // Eq and NotEq remain the same }; return Ok(NucleonExpr { - column: col.clone(), + column: dim, op: flipped_op, value: normalized_val, }); @@ -436,7 +472,7 @@ mod test { assert_eq!(collider.normalized_values.len(), 4); // Check age column - should have 2 unique values (25, 30) - let age_values = &collider.normalized_values["age"]; + let age_values = &collider.normalized_values[&PartitionDim::Column("age".to_string())]; assert_eq!(age_values.len(), 2); assert_eq!( age_values, @@ -447,7 +483,7 @@ mod test { ); // Check name column - should have 2 values - let name_values = &collider.normalized_values["name"]; + let name_values = &collider.normalized_values[&PartitionDim::Column("name".to_string())]; assert_eq!(name_values.len(), 2); assert_eq!( name_values, @@ -458,7 +494,8 @@ mod test { ); // Check active column - should have 2 values - let active_values = &collider.normalized_values["active"]; + let active_values = + &collider.normalized_values[&PartitionDim::Column("active".to_string())]; assert_eq!(active_values.len(), 2); assert_eq!( active_values, @@ -469,7 +506,7 @@ mod test { ); // Check score column - should have 2 values - let score_values = &collider.normalized_values["score"]; + let score_values = &collider.normalized_values[&PartitionDim::Column("score".to_string())]; assert_eq!(score_values.len(), 2); assert_eq!( score_values, @@ -665,10 +702,26 @@ mod test { // Should have normalized values for all 4 columns assert_eq!(collider.normalized_values.len(), 4); - assert!(collider.normalized_values.contains_key("name")); - assert!(collider.normalized_values.contains_key("age")); - assert!(collider.normalized_values.contains_key("status")); - assert!(collider.normalized_values.contains_key("score")); + assert!( + collider + .normalized_values + .contains_key(&PartitionDim::Column("name".to_string())) + ); + assert!( + collider + .normalized_values + .contains_key(&PartitionDim::Column("age".to_string())) + ); + assert!( + collider + .normalized_values + .contains_key(&PartitionDim::Column("status".to_string())) + ); + assert!( + collider + .normalized_values + .contains_key(&PartitionDim::Column("score".to_string())) + ); } #[test] diff --git a/src/partition/src/expr.rs b/src/partition/src/expr.rs index 03e7249797b..55139cd9dca 100644 --- a/src/partition/src/expr.rs +++ b/src/partition/src/expr.rs @@ -31,9 +31,13 @@ use datatypes::value::{ use serde::{Deserialize, Serialize}; use snafu::ResultExt; use sql::statements::value_to_sql_value; -use sqlparser::ast::{BinaryOperator as ParserBinaryOperator, Expr as ParserExpr, Ident}; +use sqlparser::ast::{ + BinaryOperator as ParserBinaryOperator, Expr as ParserExpr, Function, FunctionArg, + FunctionArgExpr, FunctionArgumentList, FunctionArguments, Ident, ObjectName, +}; use crate::error; +use crate::function::IcebergHash; use crate::partition::PartitionBound; /// Struct for partition expression. This can be converted back to sqlparser's [Expr]. @@ -52,6 +56,7 @@ pub enum Operand { Column(String), Value(Value), Expr(PartitionExpr), + IcebergHash(IcebergHash), } pub fn col(column_name: impl Into) -> Operand { @@ -67,7 +72,7 @@ impl From for Operand { impl Operand { pub fn try_as_logical_expr(&self) -> error::Result { match self { - Self::Column(c) => Ok(datafusion_expr::col(format!(r#""{}""#, c))), + Self::Column(c) => Ok(Expr::Column(datafusion_common::Column::from_name(c))), Self::Value(v) => { let scalar_value = match v { Value::Boolean(v) => ScalarValue::Boolean(Some(*v)), @@ -108,6 +113,38 @@ impl Operand { Ok(datafusion_expr::lit(scalar_value)) } Self::Expr(e) => e.try_as_logical_expr(), + Self::IcebergHash(f) => Ok(f.to_logical_expr()), + } + } + + /// Whether this operand is a partition dimension that values are compared against. + fn is_dimension(&self) -> bool { + matches!(self, Self::Column(_) | Self::IcebergHash(_)) + } + + fn to_parser_expr(&self) -> ParserExpr { + // Safety: Partition rule won't contains unsupported value type. + // Otherwise it will be rejected by the parser. + match self { + Self::Column(c) => ParserExpr::Identifier(Ident::new(c.clone())), + Self::Value(v) => ParserExpr::Value(value_to_sql_value(v).unwrap().into()), + Self::Expr(e) => e.to_parser_expr(), + Self::IcebergHash(f) => ParserExpr::Function(Function { + name: ObjectName::from(vec![Ident::new(IcebergHash::NAME)]), + uses_odbc_syntax: false, + parameters: FunctionArguments::None, + args: FunctionArguments::List(FunctionArgumentList { + duplicate_treatment: None, + args: vec![FunctionArg::Unnamed(FunctionArgExpr::Expr( + ParserExpr::Identifier(Ident::new(f.column())), + ))], + clauses: vec![], + }), + filter: None, + null_treatment: None, + over: None, + within_group: vec![], + }), } } @@ -142,6 +179,7 @@ impl Display for Operand { Self::Column(v) => write!(f, "{v}"), Self::Value(v) => write!(f, "{v}"), Self::Expr(v) => write!(f, "{v}"), + Self::IcebergHash(v) => write!(f, "{v}"), } } } @@ -239,7 +277,7 @@ impl PartitionExpr { rhs: Box::new(rhs), }; - if matches!(&*expr.lhs, Operand::Value(_)) && matches!(&*expr.rhs, Operand::Column(_)) { + if matches!(&*expr.lhs, Operand::Value(_)) && expr.rhs.is_dimension() { std::mem::swap(&mut expr.lhs, &mut expr.rhs); expr.op = expr.op.invert_for_swap(); } @@ -258,24 +296,10 @@ impl PartitionExpr { /// /// [Expr]: ParserExpr pub fn to_parser_expr(&self) -> ParserExpr { - // Safety: Partition rule won't contains unsupported value type. - // Otherwise it will be rejected by the parser. - let lhs = match &*self.lhs { - Operand::Column(c) => ParserExpr::Identifier(Ident::new(c.clone())), - Operand::Value(v) => ParserExpr::Value(value_to_sql_value(v).unwrap().into()), - Operand::Expr(e) => e.to_parser_expr(), - }; - - let rhs = match &*self.rhs { - Operand::Column(c) => ParserExpr::Identifier(Ident::new(c.clone())), - Operand::Value(v) => ParserExpr::Value(value_to_sql_value(v).unwrap().into()), - Operand::Expr(e) => e.to_parser_expr(), - }; - ParserExpr::BinaryOp { - left: Box::new(lhs), + left: Box::new(self.lhs.to_parser_expr()), op: self.op.to_parser_op(), - right: Box::new(rhs), + right: Box::new(self.rhs.to_parser_expr()), } } @@ -284,10 +308,10 @@ impl PartitionExpr { // Batch routing and region filters treat a NULL result as false, so only the // comparisons that hold for NULL need an explicit `IS NULL`. let comparison = match (self.lhs.as_ref(), self.rhs.as_ref()) { - (column @ Operand::Column(_), Operand::Value(value)) => { + (column, Operand::Value(value)) if column.is_dimension() => { Some((column, value, self.op.clone())) } - (Operand::Value(value), column @ Operand::Column(_)) => { + (Operand::Value(value), column) if column.is_dimension() => { Some((column, value, self.op.invert_for_swap())) } _ => None, @@ -355,32 +379,7 @@ impl PartitionExpr { &self, schema: &arrow::datatypes::SchemaRef, ) -> error::Result> { - let df_schema = schema - .clone() - .to_dfschema_ref() - .context(error::ToDFSchemaSnafu)?; - let expr = self.try_as_logical_expr()?; - // Persisted partition bounds keep the unit (and type) they were created - // with, so after a schema change like widening the time index unit the - // expression may mix types with the columns it is evaluated against. - // Coerce it against the schema so comparisons share a common type, the - // same way `PredicateGroup` does for region partition filters. - let simplifier = ExprSimplifier::new( - SimplifyContext::builder() - .with_schema(df_schema.clone()) - .build(), - ); - let expr = simplifier - .coerce(expr, &df_schema) - .context(error::CoercePartitionExprSnafu)?; - let execution_props = &ExecutionProps::default(); - create_physical_expr( - &expr, - &df_schema, - execution_props, - &PhysicalPlanningContext::default(), - ) - .context(error::CreatePhysicalExprSnafu) + plan_physical_expr(self.try_as_logical_expr()?, schema) } pub fn and(self, rhs: PartitionExpr) -> PartitionExpr { @@ -432,11 +431,45 @@ impl PartitionExpr { Operand::Expr(e) => { e.collect_column_names(columns); } + Operand::IcebergHash(f) => { + columns.insert(f.column().to_string()); + } Operand::Value(_) => {} } } } +/// Plans a logical partition expression for evaluation against `schema`. +pub(crate) fn plan_physical_expr( + expr: Expr, + schema: &arrow::datatypes::SchemaRef, +) -> error::Result> { + let df_schema = schema + .clone() + .to_dfschema_ref() + .context(error::ToDFSchemaSnafu)?; + // Persisted partition bounds keep the unit (and type) they were created + // with, so after a schema change like widening the time index unit the + // expression may mix types with the columns it is evaluated against. + // Coerce it against the schema so comparisons share a common type, the + // same way `PredicateGroup` does for region partition filters. + let simplifier = ExprSimplifier::new( + SimplifyContext::builder() + .with_schema(df_schema.clone()) + .build(), + ); + let expr = simplifier + .coerce(expr, &df_schema) + .context(error::CoercePartitionExprSnafu)?; + create_physical_expr( + &expr, + &df_schema, + &ExecutionProps::default(), + &PhysicalPlanningContext::default(), + ) + .context(error::CreatePhysicalExprSnafu) +} + impl Display for PartitionExpr { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { write!(f, "{} {} {}", self.lhs, self.op, self.rhs) @@ -813,6 +846,42 @@ mod tests { assert!(expr5.is_none()); } + #[test] + fn test_serde_iceberg_hash_partition_expr() { + let hash = Operand::IcebergHash(IcebergHash::new("host")); + let expr = hash.clone().lt(Value::Int32(0)); + let json = expr.as_json_str().unwrap(); + assert_eq!( + json, + r#"{"Expr":{"lhs":{"IcebergHash":"host"},"op":"Lt","rhs":{"Value":{"Int32":0}}}}"# + ); + assert_eq!(PartitionExpr::from_json_str(&json).unwrap().unwrap(), expr); + + let swapped = + r#"{"Expr":{"lhs":{"Value":{"Int32":0}},"op":"Gt","rhs":{"IcebergHash":"host"}}}"#; + assert_eq!( + PartitionExpr::from_json_str(swapped).unwrap().unwrap(), + expr + ); + assert_eq!(expr.to_string(), "iceberg_hash(host) < 0"); + } + + #[test] + fn test_logical_expr_keeps_quotes_in_column_names() { + use datatypes::arrow::datatypes::{DataType, Field, Schema}; + + // Declared as `"a""b"`; the stored name contains the quote. + let name = r#"a"b"#; + let schema = Arc::new(Schema::new(vec![Field::new(name, DataType::Utf8, true)])); + for expr in [ + col(name).lt(Value::from("m")), + Operand::IcebergHash(IcebergHash::new(name)).lt(Value::Int32(0)), + ] { + expr.try_as_physical_expr(&schema) + .unwrap_or_else(|e| panic!("{expr}: {e}")); + } + } + #[test] fn test_collect_column_names() { // Simple expression: col_a = 1 should give {col_a} @@ -854,10 +923,8 @@ mod tests { assert!(columns.contains("c")); } - /// Real serialized expressions must survive the round trip through - /// `common-meta`'s persisted-expression mirror: whenever that mirror drifts - /// from these types, deserialization fails, validation is skipped and the - /// overflowing cases below stop producing `PartitionExprIncompatible`. + /// `common-meta` mirrors these types to validate ALTERs. If the mirror drifts, + /// it can't parse the expressions and skips validation, which fails these cases. #[test] fn test_serialized_expressions_validate_against_type_changes() { let overflow_ms = serde_json::from_str::( @@ -884,10 +951,7 @@ mod tests { HashMap::from([(column.to_string(), (old_type, new_type))]) }; - // Compile-time tripwire: exhaustively matching the operand and operator - // variants makes a new variant break this test at compile time, so the - // `partition_compat::Persisted*` mirror and these cases must be extended - // together instead of silently skipping validation. + // Fails to compile on a new variant, which the mirror must learn too. fn walk(expr: &PartitionExpr) { walk_operand(&expr.lhs); walk_operand(&expr.rhs); @@ -904,13 +968,11 @@ mod tests { } fn walk_operand(operand: &Operand) { match operand { - Operand::Column(_) | Operand::Value(_) => {} + Operand::Column(_) | Operand::Value(_) | Operand::IcebergHash(_) => {} Operand::Expr(expr) => walk(expr), } } - // Every expression shape must be rejected while the bound overflows - // nanoseconds, with the error naming the offending column. let simple = col("ts").lt(overflow_ms.clone()); let nested_and = col("ts") .gt_eq(start_ms.clone()) @@ -920,7 +982,10 @@ mod tests { RestrictedOp::Or, Operand::Expr(col("ts").lt(overflow_ms.clone())), ); - for expr in [&simple, &nested_and, &nested_or] { + let with_hash = col("ts") + .lt(overflow_ms.clone()) + .and(Operand::IcebergHash(IcebergHash::new("host")).lt(Value::Int32(0))); + for expr in [&simple, &nested_and, &nested_or, &with_hash] { walk(expr); let err = validate_region_partition_bounds( &changed("ts", ts_ms.clone(), ts_ns.clone()), @@ -941,9 +1006,6 @@ mod tests { validate_region_partition_bounds(&changed("ts", ts_ms, ts_ns), &routes(&compatible)) .unwrap(); - // Narrowing a field column must be rejected as well: the region's write - // path compares the persisted Int64 bound against Int32 rows directly - // and panics, even though the scan would coerce both sides. let narrowing = col("v").lt(Value::Int64(3_000_000_000)); let err = validate_region_partition_bounds( &changed( @@ -961,5 +1023,24 @@ mod tests { ), "unexpected error: {err}" ); + + // The hash of a column depends on its type. + let hashed = Operand::IcebergHash(IcebergHash::new("v")).lt(Value::Int32(0)); + let err = validate_region_partition_bounds( + &changed( + "v", + ConcreteDataType::int32_datatype(), + ConcreteDataType::int64_datatype(), + ), + &routes(&hashed), + ) + .unwrap_err(); + assert!( + matches!( + err, + Error::PartitionExprIncompatible { ref column, .. } if column == "v" + ), + "unexpected error: {err}" + ); } } diff --git a/src/partition/src/function.rs b/src/partition/src/function.rs new file mode 100644 index 00000000000..cd121deb7f4 --- /dev/null +++ b/src/partition/src/function.rs @@ -0,0 +1,345 @@ +// 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. + +//! `iceberg_hash`, the hash allowed on the column side of a partition rule. +//! +//! The hash decides which region a row is written to and is persisted in +//! partition metadata, so its result for a given input must never change across +//! versions. A different algorithm needs a new type. +//! +//! Rule simplification and coverage checks treat the hash as an independent +//! dimension, which holds because it is defined for every value of the accepted +//! column types. + +use std::fmt::{Display, Formatter}; +use std::sync::Arc; + +use datafusion_common::arrow::array::{Array, ArrayRef, AsArray, Int32Array}; +use datafusion_common::arrow::compute::take; +use datafusion_common::arrow::datatypes::{ + DataType, Int8Type, Int16Type, Int32Type, Int64Type, UInt8Type, UInt16Type, UInt32Type, + UInt64Type, +}; +use datafusion_common::{DataFusionError, ScalarValue}; +use datafusion_expr::{ + ColumnarValue, Expr, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility, +}; +use datatypes::prelude::ConcreteDataType; +use datatypes::value::Value; +use serde::{Deserialize, Serialize}; + +use crate::error::{self, Result}; + +/// `iceberg_hash(column)`: the 32-bit Murmur3 hash defined by the Iceberg table +/// spec (appendix B) over a string or integer column, as an `INT`. +/// +/// Integers of every width hash as Iceberg longs. `UINT64` values hash their bits +/// as a long, so values up to `i64::MAX` match Iceberg. `NULL` hashes to `NULL`. +/// The field is the hashed column. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] +pub struct IcebergHash(String); + +impl IcebergHash { + pub const NAME: &str = "iceberg_hash"; + + pub fn new(column: impl Into) -> Self { + Self(column.into()) + } + + pub fn column(&self) -> &str { + &self.0 + } + + pub fn accepts(data_type: &ConcreteDataType) -> bool { + data_type.is_string() || data_type.is_unsigned() || is_signed_integer(data_type) + } + + /// The type of the result, which is also the type of the literals it is compared with. + pub fn output_type() -> ConcreteDataType { + ConcreteDataType::int32_datatype() + } + + pub fn evaluate_value(value: &Value) -> Result { + let hash = match value { + Value::Null => return Ok(Value::Null), + Value::String(v) => murmur3_32(v.as_utf8().as_bytes()), + Value::Int8(v) => hash_long(*v as i64), + Value::Int16(v) => hash_long(*v as i64), + Value::Int32(v) => hash_long(*v as i64), + Value::Int64(v) => hash_long(*v), + Value::UInt8(v) => hash_long(*v as i64), + Value::UInt16(v) => hash_long(*v as i64), + Value::UInt32(v) => hash_long(*v as i64), + Value::UInt64(v) => hash_long(*v as i64), + other => { + return error::UnexpectedSnafu { + err_msg: format!("Unsupported value for {}: {other:?}", Self::NAME), + } + .fail(); + } + }; + Ok(Value::Int32(hash)) + } + + pub fn evaluate_array(array: &dyn Array) -> datafusion_common::Result { + iceberg_hash_array(array) + } + + pub fn to_logical_expr(&self) -> Expr { + let column = Expr::Column(datafusion_common::Column::from_name(self.column())); + ScalarUDF::new_from_impl(IcebergHashUdf::default()).call(vec![column]) + } +} + +impl Display for IcebergHash { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "{}({})", Self::NAME, self.column()) + } +} + +fn is_signed_integer(data_type: &ConcreteDataType) -> bool { + matches!( + data_type, + ConcreteDataType::Int8(_) + | ConcreteDataType::Int16(_) + | ConcreteDataType::Int32(_) + | ConcreteDataType::Int64(_) + ) +} + +fn iceberg_hash_array(array: &dyn Array) -> datafusion_common::Result { + let hashes: Int32Array = match array.data_type() { + DataType::Utf8 => array + .as_string::() + .iter() + .map(|v| v.map(|v| murmur3_32(v.as_bytes()))) + .collect(), + DataType::LargeUtf8 => array + .as_string::() + .iter() + .map(|v| v.map(|v| murmur3_32(v.as_bytes()))) + .collect(), + DataType::Utf8View => array + .as_string_view() + .iter() + .map(|v| v.map(|v| murmur3_32(v.as_bytes()))) + .collect(), + DataType::Int8 => array + .as_primitive::() + .unary(|v| hash_long(v as i64)), + DataType::Int16 => array + .as_primitive::() + .unary(|v| hash_long(v as i64)), + DataType::Int32 => array + .as_primitive::() + .unary(|v| hash_long(v as i64)), + DataType::Int64 => array.as_primitive::().unary(hash_long), + DataType::UInt8 => array + .as_primitive::() + .unary(|v| hash_long(v as i64)), + DataType::UInt16 => array + .as_primitive::() + .unary(|v| hash_long(v as i64)), + DataType::UInt32 => array + .as_primitive::() + .unary(|v| hash_long(v as i64)), + DataType::UInt64 => array + .as_primitive::() + .unary(|v| hash_long(v as i64)), + DataType::Dictionary(_, _) => { + let dictionary = array.as_any_dictionary(); + let values = iceberg_hash_array(dictionary.values().as_ref())?; + return Ok(take(values.as_ref(), dictionary.keys(), None)?); + } + other => { + return Err(DataFusionError::Execution(format!( + "{} does not support type {other}", + IcebergHash::NAME + ))); + } + }; + Ok(Arc::new(hashes)) +} + +fn hash_long(value: i64) -> i32 { + murmur3_32(&value.to_le_bytes()) +} + +/// MurmurHash3 x86 32-bit with seed 0, returned as a signed integer like Iceberg. +fn murmur3_32(data: &[u8]) -> i32 { + const C1: u32 = 0xcc9e_2d51; + const C2: u32 = 0x1b87_3593; + + fn mix_k(k: u32) -> u32 { + k.wrapping_mul(C1).rotate_left(15).wrapping_mul(C2) + } + + let mut h = 0u32; + let (blocks, tail) = data.as_chunks::<4>(); + for block in blocks { + h = (h ^ mix_k(u32::from_le_bytes(*block))) + .rotate_left(13) + .wrapping_mul(5) + .wrapping_add(0xe654_6b64); + } + if !tail.is_empty() { + let k = tail + .iter() + .rev() + .fold(0u32, |k, byte| (k << 8) | *byte as u32); + h ^= mix_k(k); + } + + // Murmur3 mixes in the length modulo 2^32. + h ^= data.len() as u32; + h ^= h >> 16; + h = h.wrapping_mul(0x85eb_ca6b); + h ^= h >> 13; + h = h.wrapping_mul(0xc2b2_ae35); + h ^= h >> 16; + h as i32 +} + +/// DataFusion form of [`IcebergHash`], so storage engines +/// can evaluate a region's partition expression as a filter. +#[derive(Debug, PartialEq, Eq, Hash)] +struct IcebergHashUdf { + signature: Signature, +} + +impl Default for IcebergHashUdf { + fn default() -> Self { + Self { + signature: Signature::any(1, Volatility::Immutable), + } + } +} + +impl ScalarUDFImpl for IcebergHashUdf { + fn name(&self) -> &str { + IcebergHash::NAME + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> datafusion_common::Result { + Ok(DataType::Int32) + } + + fn invoke_with_args( + &self, + args: ScalarFunctionArgs, + ) -> datafusion_common::Result { + let [arg] = args.args.as_slice() else { + return Err(DataFusionError::Internal(format!( + "{} expects one argument, got {}", + IcebergHash::NAME, + args.args.len() + ))); + }; + match arg { + ColumnarValue::Array(array) => { + Ok(ColumnarValue::Array(iceberg_hash_array(array.as_ref())?)) + } + ColumnarValue::Scalar(scalar) => { + let hashes = iceberg_hash_array(scalar.to_array()?.as_ref())?; + Ok(ColumnarValue::Scalar(ScalarValue::try_from_array( + &hashes, 0, + )?)) + } + } + } +} + +#[cfg(test)] +mod tests { + use datafusion_common::arrow::array::{ + Int64Array, StringArray, StringDictionaryBuilder, StringViewArray, UInt64Array, + }; + use datafusion_common::arrow::datatypes::UInt16Type; + + use super::*; + + #[test] + fn test_iceberg_spec_vectors() { + // Iceberg table spec, appendix B. + assert_eq!(hash_long(34), 2017239379); + // date 2017-11-16 hashes as int 17486. + assert_eq!(hash_long(17486), -653330422); + // time 22:31:08 hashes as long 81068000000 micros. + assert_eq!(hash_long(81068000000), -662762989); + assert_eq!(murmur3_32("iceberg".as_bytes()), 1210000089); + assert_eq!(murmur3_32(&[0, 1, 2, 3]), -188683207); + } + + #[test] + fn test_unsigned_policy() { + assert_eq!( + IcebergHash::evaluate_value(&Value::UInt64(34)).unwrap(), + Value::Int32(2017239379) + ); + assert_eq!( + IcebergHash::evaluate_value(&Value::UInt64(u64::MAX)).unwrap(), + Value::Int32(hash_long(-1)) + ); + } + + #[test] + fn test_array_matches_value() { + let strings = [Some("iceberg"), None, Some(""), Some("中🙂")]; + let mut dictionary = StringDictionaryBuilder::::new(); + for v in strings { + dictionary.append_option(v); + } + let string_arrays: Vec = vec![ + Arc::new(StringArray::from(strings.to_vec())), + Arc::new(StringViewArray::from(strings.to_vec())), + Arc::new(dictionary.finish()), + ]; + let string_values = strings.map(|v| v.map_or(Value::Null, Value::from)); + + let ints = [Some(34), None, Some(i64::MIN), Some(-1)]; + let unsigned = [Some(34), None, Some(u64::MAX), Some(0)]; + let cases = string_arrays + .into_iter() + .map(|array| (array, string_values.to_vec())) + .chain([ + ( + Arc::new(Int64Array::from(ints.to_vec())) as ArrayRef, + ints.map(|v| v.map_or(Value::Null, Value::Int64)).to_vec(), + ), + ( + Arc::new(UInt64Array::from(unsigned.to_vec())) as ArrayRef, + unsigned + .map(|v| v.map_or(Value::Null, Value::UInt64)) + .to_vec(), + ), + ]); + for (array, values) in cases { + let hashes = IcebergHash::evaluate_array(array.as_ref()).unwrap(); + let hashes = hashes.as_primitive::(); + for (row, value) in values.iter().enumerate() { + let expected = IcebergHash::evaluate_value(value).unwrap(); + let actual = if hashes.is_null(row) { + Value::Null + } else { + Value::Int32(hashes.value(row)) + }; + assert_eq!(actual, expected, "{:?} row {row}", array.data_type()); + } + } + } +} diff --git a/src/partition/src/lib.rs b/src/partition/src/lib.rs index 3bfd7297cc7..f89da2749ac 100644 --- a/src/partition/src/lib.rs +++ b/src/partition/src/lib.rs @@ -19,6 +19,7 @@ pub mod checker; pub mod collider; pub mod error; pub mod expr; +pub mod function; pub mod manager; pub mod multi_dim; pub mod overlap; diff --git a/src/partition/src/multi_dim.rs b/src/partition/src/multi_dim.rs index 25dcc0fb89d..42f7dced14f 100644 --- a/src/partition/src/multi_dim.rs +++ b/src/partition/src/multi_dim.rs @@ -14,17 +14,19 @@ use std::any::Any; use std::cmp::Ordering; -use std::collections::HashMap; use std::collections::hash_map::Entry; +use std::collections::{BTreeSet, HashMap}; use std::sync::{Arc, RwLock}; -use datafusion_expr::ColumnarValue; +use datafusion_common::tree_node::{Transformed, TreeNode}; +use datafusion_expr::{ColumnarValue, Expr}; use datafusion_physical_expr::PhysicalExpr; use datatypes::arrow; use datatypes::arrow::array::{Array, BooleanArray, BooleanBufferBuilder, RecordBatch}; use datatypes::arrow::buffer::BooleanBuffer; use datatypes::arrow::compute::prep_null_mask_filter; -use datatypes::arrow::datatypes::Schema; +use datatypes::arrow::datatypes::{Field, Schema}; +use datatypes::data_type::DataType as _; use datatypes::prelude::Value; use datatypes::vectors::{Helper, VectorRef}; use serde::{Deserialize, Serialize}; @@ -34,13 +36,26 @@ use store_api::storage::RegionNumber; use crate::PartitionRule; use crate::checker::PartitionChecker; use crate::error::{self, Result, UndefinedColumnSnafu}; -use crate::expr::{Operand, PartitionExpr, RestrictedOp}; +use crate::expr::{Operand, PartitionExpr, RestrictedOp, plan_physical_expr}; +use crate::function::IcebergHash; use crate::partition::RegionMask; /// The default region number when no partition exprs are matched. const DEFAULT_REGION: RegionNumber = 0; -type PhysicalExprCache = Option<(Vec>, Arc)>; +type PhysicalExprCache = Option>; + +/// Partition predicates planned for one input schema. +#[derive(Debug)] +struct PhysicalPartitionExprs { + input_schema: Arc, + exprs: Vec>, + /// Input indices of the columns under `iceberg_hash`. Each hash is evaluated + /// once per batch and appended to the input as `eval_schema` describes, and + /// `exprs` read the results instead of hashing per predicate. + hash_inputs: Vec, + eval_schema: Arc, +} /// Multi-Dimiension partition rule. RFC [here](https://github.com/GreptimeTeam/greptimedb/blob/main/docs/rfcs/2024-02-21-multi-dimension-partition-rule/rfc.md) /// @@ -137,6 +152,14 @@ impl MultiDimPartitionRule { let r = &values[*index]; Self::perform_op(l, &expr.op, r) } + (Operand::IcebergHash(hash), Operand::Value(r)) => { + let l = self.evaluate_hash(hash, values)?; + Self::perform_op(&l, &expr.op, r) + } + (Operand::Value(l), Operand::IcebergHash(hash)) => { + let r = self.evaluate_hash(hash, values)?; + Self::perform_op(l, &expr.op, &r) + } (Operand::Expr(lhs), Operand::Expr(rhs)) => { let lhs = self.evaluate_expr(lhs, values)?; let rhs = self.evaluate_expr(rhs, values)?; @@ -150,6 +173,15 @@ impl MultiDimPartitionRule { } } + fn evaluate_hash(&self, hash: &IcebergHash, values: &[Value]) -> Result { + let column = hash.column(); + let index = self + .name_to_index + .get(column) + .context(UndefinedColumnSnafu { column })?; + IcebergHash::evaluate_value(&values[*index]) + } + fn perform_op(lhs: &Value, op: &RestrictedOp, rhs: &Value) -> Result { let result = match op { RestrictedOp::Eq => lhs.eq(rhs), @@ -239,27 +271,38 @@ impl MultiDimPartitionRule { } let physical_exprs = { let cache_read_guard = self.physical_expr_cache.read().unwrap(); - if let Some((cached_exprs, schema)) = cache_read_guard.as_ref() - && schema == record_batch.schema_ref() + if let Some(cached) = cache_read_guard.as_ref() + && &cached.input_schema == record_batch.schema_ref() { - cached_exprs.clone() + cached.clone() } else { drop(cache_read_guard); // Release the read lock before acquiring write lock - let schema = record_batch.schema(); - let new_cache = self - .exprs - .iter() - .map(|e| e.try_as_physical_expr(&schema)) - .collect::>>()?; - + let new_cache = Arc::new(self.plan_physical_exprs(record_batch.schema())?); let mut cache_write_guard = self.physical_expr_cache.write().unwrap(); - cache_write_guard.replace((new_cache.clone(), schema)); + cache_write_guard.replace(new_cache.clone()); new_cache } }; + let eval_batch; + let record_batch = if physical_exprs.hash_inputs.is_empty() { + record_batch + } else { + let mut columns = record_batch.columns().to_vec(); + for index in &physical_exprs.hash_inputs { + columns.push( + IcebergHash::evaluate_array(record_batch.column(*index).as_ref()) + .context(error::EvaluateRecordBatchSnafu)?, + ); + } + eval_batch = RecordBatch::try_new(physical_exprs.eval_schema.clone(), columns) + .context(error::ComputeArrowKernelSnafu)?; + &eval_batch + }; + let mut result: HashMap = physical_exprs + .exprs .iter() .zip(self.regions.iter()) .filter_map(|(expr, region_num)| { @@ -325,6 +368,97 @@ impl MultiDimPartitionRule { } Ok(result) } + + fn plan_physical_exprs(&self, input_schema: Arc) -> Result { + let mut hashes = BTreeSet::new(); + for expr in &self.exprs { + collect_hashes(expr, &mut hashes); + } + if hashes.is_empty() { + let exprs = self + .exprs + .iter() + .map(|e| e.try_as_physical_expr(&input_schema)) + .collect::>>()?; + return Ok(PhysicalPartitionExprs { + eval_schema: input_schema.clone(), + input_schema, + exprs, + hash_inputs: vec![], + }); + } + + let mut fields = input_schema.fields().to_vec(); + let mut hash_columns = HashMap::with_capacity(hashes.len()); + let mut hash_inputs = Vec::with_capacity(hashes.len()); + for hash in hashes { + let index = + input_schema + .index_of(hash.column()) + .ok() + .context(UndefinedColumnSnafu { + column: hash.column(), + })?; + let mut name = format!("__partition_hash_{}", hash_inputs.len()); + while input_schema.field_with_name(&name).is_ok() { + name.push('_'); + } + fields.push(Arc::new(Field::new( + &name, + IcebergHash::output_type().as_arrow_type(), + true, + ))); + hash_columns.insert( + hash.to_logical_expr(), + Expr::Column(datafusion_common::Column::from_name(name)), + ); + hash_inputs.push(index); + } + let eval_schema = Arc::new(Schema::new(fields)); + + let exprs = self + .exprs + .iter() + .map(|expr| { + let expr = expr + .try_as_logical_expr()? + .transform(|expr| { + Ok(match hash_columns.get(&expr) { + Some(column) => Transformed::yes(column.clone()), + None => Transformed::no(expr), + }) + }) + .context(error::CreatePhysicalExprSnafu)? + .data; + debug_assert!( + !expr + .exists(|expr| Ok(matches!(expr, Expr::ScalarFunction(_)))) + .unwrap(), + "iceberg_hash should read the precomputed columns: {expr}" + ); + plan_physical_expr(expr, &eval_schema) + }) + .collect::>>()?; + + Ok(PhysicalPartitionExprs { + input_schema, + exprs, + hash_inputs, + eval_schema, + }) + } +} + +fn collect_hashes(expr: &PartitionExpr, hashes: &mut BTreeSet) { + for operand in [expr.lhs(), expr.rhs()] { + match operand { + Operand::IcebergHash(hash) => { + hashes.insert(hash.clone()); + } + Operand::Expr(expr) => collect_hashes(expr, hashes), + Operand::Column(_) | Operand::Value(_) => {} + } + } } fn columnar_value_to_boolean_array( @@ -858,6 +992,71 @@ mod test_split_record_batch { } } + #[test] + fn test_split_record_batch_widened_time_index_unit_with_hash() { + // The precomputed-hash path must coerce millisecond bounds against a + // nanosecond time index too. + let bound = Value::from(Timestamp::new_millisecond(1_000)); + let hash = Operand::IcebergHash(IcebergHash::new("host")); + let exprs = [ + ( + col("ts").lt(bound.clone()), + hash.clone().lt(Value::Int32(0)), + ), + ( + col("ts").lt(bound.clone()), + hash.clone().gt_eq(Value::Int32(0)), + ), + ( + col("ts").gt_eq(bound.clone()), + hash.clone().lt(Value::Int32(0)), + ), + (col("ts").gt_eq(bound), hash.gt_eq(Value::Int32(0))), + ] + .into_iter() + .map(|(ts, hash)| ts.and(hash)) + .collect(); + let rule = MultiDimPartitionRule::try_new( + vec!["ts".to_string(), "host".to_string()], + vec![1, 2, 3, 4], + exprs, + true, + ) + .unwrap(); + + let hosts = ["a", "b", "c", "server42"]; + let mut ts = vec![]; + let mut host = vec![]; + for nanos in [500_000_000, 2_000_000_000] { + for h in hosts { + ts.push(nanos); + host.push(h); + } + } + let schema = Arc::new(Schema::new(vec![ + Field::new("ts", DataType::Timestamp(TimeUnit::Nanosecond, None), false), + Field::new("host", DataType::Utf8, false), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(TimestampNanosecondArray::from(ts)), + Arc::new(StringArray::from(host)), + ], + ) + .unwrap(); + + let expected = rule.split_record_batch_naive(&batch).unwrap(); + let result = rule.split_record_batch(&batch).unwrap(); + for (region, expected) in expected { + let actual = result + .get(®ion) + .map(|mask| mask.array().clone()) + .unwrap_or_else(|| BooleanArray::from(vec![false; batch.num_rows()])); + assert_eq!(actual, expected, "region {region}"); + } + } + #[test] fn test_split_record_batch_widened_time_index_unit() { // https://github.com/GreptimeTeam/greptimedb/issues/9747 @@ -1021,6 +1220,84 @@ mod test_split_record_batch { assert_eq!(total_rows, batch.num_rows()); } + #[test] + fn test_hash_partition_batch_matches_row() { + use datatypes::arrow::array::{Array, StringDictionaryBuilder, UInt64Array}; + use datatypes::arrow::datatypes::{DataType, Field, UInt32Type}; + + let host_hash = Operand::IcebergHash(IcebergHash::new("host")); + let tenant_hash = Operand::IcebergHash(IcebergHash::new("tenant")); + let exprs = [ + host_hash + .clone() + .lt(Value::Int32(0)) + .and(col("tenant").lt(Value::UInt64(100))), + host_hash + .clone() + .lt(Value::Int32(0)) + .and(col("tenant").gt_eq(Value::UInt64(100))), + host_hash + .clone() + .gt_eq(Value::Int32(0)) + .and(tenant_hash.clone().lt(Value::Int32(0))), + host_hash + .gt_eq(Value::Int32(0)) + .and(tenant_hash.gt_eq(Value::Int32(0))), + ] + .iter() + .map(|expr| { + PartitionExpr::from_json_str(&expr.as_json_str().unwrap()) + .unwrap() + .unwrap() + }) + .collect(); + let rule = MultiDimPartitionRule::try_new( + vec!["host".to_string(), "tenant".to_string()], + vec![1, 2, 3, 4], + exprs, + true, + ) + .unwrap(); + + let hosts = [ + None, + Some(""), + Some("a"), + Some("b"), + Some("中🙂"), + Some("server42"), + ]; + let tenants = [None, Some(0), Some(99), Some(100), Some(u64::MAX)]; + let mut host_builder = StringDictionaryBuilder::::new(); + let mut tenant_values = vec![]; + for host in hosts { + for tenant in tenants { + host_builder.append_option(host); + tenant_values.push(tenant); + } + } + let host_array = Arc::new(host_builder.finish()); + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("host", host_array.data_type().clone(), true), + Field::new("tenant", DataType::UInt64, true), + ])), + vec![host_array, Arc::new(UInt64Array::from(tenant_values))], + ) + .unwrap(); + + let by_row = rule.split_record_batch_naive(&batch).unwrap(); + let by_batch = rule.split_record_batch(&batch).unwrap(); + assert!(!by_batch.contains_key(&DEFAULT_REGION)); + for (region, expected) in by_row { + let actual = by_batch + .get(®ion) + .map(|mask| mask.array().clone()) + .unwrap_or_else(|| BooleanArray::from(vec![false; batch.num_rows()])); + assert_eq!(actual, expected, "region {region}"); + } + } + #[test] fn test_comparison_batch_matches_row() { use datatypes::arrow::array::Array; diff --git a/src/partition/src/simplify.rs b/src/partition/src/simplify.rs index a2bffa51b44..eabbfbb4f98 100644 --- a/src/partition/src/simplify.rs +++ b/src/partition/src/simplify.rs @@ -23,8 +23,8 @@ use std::ops::Bound; use datatypes::value::{OrderedF64, Value}; -use crate::collider::{AtomicExpr, Collider, GluonOp, NucleonExpr}; -use crate::expr::{Operand, PartitionExpr, RestrictedOp, col}; +use crate::collider::{AtomicExpr, Collider, GluonOp, NucleonExpr, PartitionDim}; +use crate::expr::{Operand, PartitionExpr, RestrictedOp}; /// Attempts to simplify a merged partition expression (typically an `OR` of multiple partitions) /// into an equivalent but shorter expression. @@ -36,7 +36,7 @@ pub fn simplify_merged_partition_expr(expr: PartitionExpr) -> PartitionExpr { try_simplify_merged_partition_expr(&expr).unwrap_or(expr) } -type DenormValues = BTreeMap>; +type DenormValues = BTreeMap>; fn try_simplify_merged_partition_expr(expr: &PartitionExpr) -> Option { let collider = Collider::new(std::slice::from_ref(expr)).ok()?; @@ -87,7 +87,7 @@ fn term_from_atomic(atomic: &AtomicExpr, denorm_values: &DenormValues) -> Option let interval = interval_from_nucleons(&atomic.nucleons[start..i])?; if !interval.is_unbounded() { - constraints.insert(column.to_string(), interval); + constraints.insert(column.clone(), interval); } } @@ -105,7 +105,7 @@ fn interval_from_nucleons(nucleons: &[NucleonExpr]) -> Option { #[derive(Debug, Clone, PartialEq, Eq)] struct Term { // Only stores constrained columns. Missing column means unbounded. - constraints: BTreeMap, + constraints: BTreeMap, } impl Term { @@ -358,11 +358,11 @@ fn simplify_terms(mut terms: Vec) -> Option> { fn try_merge_terms(a: &Term, b: &Term) -> Option { // Find the only differing column (treat missing as unbounded). - let mut diff_col: Option<&str> = None; + let mut diff_col: Option<&PartitionDim> = None; let mut cols = BTreeSet::new(); - cols.extend(a.constraints.keys().map(|s| s.as_str())); - cols.extend(b.constraints.keys().map(|s| s.as_str())); + cols.extend(a.constraints.keys()); + cols.extend(b.constraints.keys()); for col in cols { let a_interval = a.constraints.get(col); @@ -387,7 +387,7 @@ fn try_merge_terms(a: &Term, b: &Term) -> Option { constraints.remove(diff_col); } UnionInterval::Interval(interval) => { - constraints.insert(diff_col.to_string(), interval); + constraints.insert(diff_col.clone(), interval); } } @@ -440,7 +440,7 @@ fn term_to_expr(term: &Term, denorm_values: &DenormValues) -> Option Option> { @@ -457,7 +457,7 @@ fn interval_to_exprs( match (lower, upper) { (Included(lv), Included(uv)) if lv == uv => { - return Some(vec![col(column).eq(col_values.get(lv)?.clone())]); + return Some(vec![column.to_operand().eq(col_values.get(lv)?.clone())]); } (Excluded(lv), Excluded(uv)) if lv == uv => return None, (Included(lv), Excluded(uv)) if lv == uv => return None, @@ -468,13 +468,13 @@ fn interval_to_exprs( let mut exprs = Vec::new(); match lower { Unbounded => {} - Included(v) => exprs.push(col(column).gt_eq(col_values.get(v)?.clone())), - Excluded(v) => exprs.push(col(column).gt(col_values.get(v)?.clone())), + Included(v) => exprs.push(column.to_operand().gt_eq(col_values.get(v)?.clone())), + Excluded(v) => exprs.push(column.to_operand().gt(col_values.get(v)?.clone())), } match upper { Unbounded => {} - Included(v) => exprs.push(col(column).lt_eq(col_values.get(v)?.clone())), - Excluded(v) => exprs.push(col(column).lt(col_values.get(v)?.clone())), + Included(v) => exprs.push(column.to_operand().lt_eq(col_values.get(v)?.clone())), + Excluded(v) => exprs.push(column.to_operand().lt(col_values.get(v)?.clone())), } Some(exprs) @@ -487,7 +487,7 @@ mod tests { use datatypes::value::{OrderedFloat, Value}; use super::*; - use crate::expr::Operand; + use crate::expr::col; fn or(lhs: PartitionExpr, rhs: PartitionExpr) -> PartitionExpr { PartitionExpr::new(Operand::Expr(lhs), RestrictedOp::Or, Operand::Expr(rhs)) @@ -508,6 +508,20 @@ mod tests { assert_eq!(simplified.to_string(), "device_id < 100"); } + #[test] + fn simplify_adjacent_hash_ranges() { + use crate::function::IcebergHash; + + let hash = Operand::IcebergHash(IcebergHash::new("host")); + let left = hash.clone().lt(Value::Int32(0)); + let right = hash + .clone() + .gt_eq(Value::Int32(0)) + .and(hash.clone().lt(Value::Int32(100))); + let simplified = simplify_merged_partition_expr(or(left, right)); + assert_eq!(simplified, hash.lt(Value::Int32(100))); + } + #[test] fn simplify_adjacent_ranges() { // host < 'h0' OR (host >= 'h0' AND host < 'h1') -> host < 'h1' diff --git a/src/query/src/dist_plan/predicate_extractor.rs b/src/query/src/dist_plan/predicate_extractor.rs index 020e6d169ef..170526d6c89 100644 --- a/src/query/src/dist_plan/predicate_extractor.rs +++ b/src/query/src/dist_plan/predicate_extractor.rs @@ -227,7 +227,7 @@ impl ExpressionChecker { result: &mut Vec, ) { match operand { - Operand::Column(_) | Operand::Value(_) => { + Operand::Column(_) | Operand::Value(_) | Operand::IcebergHash(_) => { // This shouldn't happen in well-formed expressions } Operand::Expr(expr) => { @@ -252,6 +252,7 @@ impl ExpressionChecker { ) -> bool { match operand { Operand::Column(col) => partition_columns.contains(col), + Operand::IcebergHash(hash) => partition_columns.contains(hash.column()), Operand::Value(_) => true, // Values are always safe Operand::Expr(expr) => { Self::expr_only_involves_partition_columns(expr, partition_columns) diff --git a/src/query/src/dist_plan/region_pruner.rs b/src/query/src/dist_plan/region_pruner.rs index eca6a60ebd5..bab2f4896a3 100644 --- a/src/query/src/dist_plan/region_pruner.rs +++ b/src/query/src/dist_plan/region_pruner.rs @@ -206,6 +206,10 @@ impl ConstraintPruner { true } + // Query predicates never constrain a function dimension, so its values all + // come from partition rules, already typed as the function output. + (Operand::IcebergHash(_), Operand::Value(_)) + | (Operand::Value(_), Operand::IcebergHash(_)) => true, _ => false, } } @@ -270,6 +274,57 @@ mod tests { assert!(pruned.contains(&RegionId::new(1, 2))); } + #[test] + fn test_constraint_pruning_with_hash_dimension() { + use partition::function::IcebergHash; + + let hash = Operand::IcebergHash(IcebergHash::new("host")); + let partitions = [ + col("ts") + .lt(Value::Int64(100)) + .and(hash.clone().lt(Value::Int32(0))), + col("ts") + .lt(Value::Int64(100)) + .and(hash.clone().gt_eq(Value::Int32(0))), + col("ts") + .gt_eq(Value::Int64(100)) + .and(hash.clone().lt(Value::Int32(0))), + col("ts") + .gt_eq(Value::Int64(100)) + .and(hash.gt_eq(Value::Int32(0))), + ] + .into_iter() + .enumerate() + .map(|(i, expr)| create_test_partition_info(i as u64 + 1, Some(expr))) + .collect::>(); + let column_datatypes = HashMap::from_iter([ + ("ts".to_string(), ConcreteDataType::int64_datatype()), + ("host".to_string(), ConcreteDataType::string_datatype()), + ]); + let prune = |query: PartitionExpr| { + let mut regions = + ConstraintPruner::prune_regions(&[query], &partitions, column_datatypes.clone()) + .unwrap(); + regions.sort(); + regions + }; + + // A filter on the hashed column can't exclude any hash range. + assert_eq!( + prune(col("host").eq(Value::String("a".into()))), + (1..=4).map(|i| RegionId::new(1, i)).collect::>() + ); + // Plain columns in the same rule still prune. + assert_eq!( + prune( + col("host") + .eq(Value::String("a".into())) + .and(col("ts").gt_eq(Value::Int64(100))) + ), + vec![RegionId::new(1, 3), RegionId::new(1, 4)] + ); + } + #[test] fn test_constraint_pruning_in_list() { let partitions = vec![ diff --git a/src/sql/src/parsers/create_parser.rs b/src/sql/src/parsers/create_parser.rs index 20eba9f2087..0da6f0a41d8 100644 --- a/src/sql/src/parsers/create_parser.rs +++ b/src/sql/src/parsers/create_parser.rs @@ -28,8 +28,8 @@ pub(crate) use json::parse_json2_type_and_options; pub use json::parse_json2_type_hint_path; use snafu::{OptionExt, ResultExt, ensure}; use sqlparser::ast::{ - ColumnOption, ColumnOptionDef, DataType, Expr, KeyOrIndexDisplay, NullsDistinctOption, - PrimaryKeyConstraint, UniqueConstraint, + ColumnOption, ColumnOptionDef, DataType, Expr, FunctionArg, FunctionArgExpr, FunctionArguments, + KeyOrIndexDisplay, NullsDistinctOption, PrimaryKeyConstraint, UniqueConstraint, }; use sqlparser::dialect::keywords::Keyword; use sqlparser::keywords::ALL_KEYWORDS; @@ -1300,6 +1300,17 @@ fn ensure_one_expr(expr: &Expr, columns: &[&Column]) -> Result<()> { ); Ok(()) } + // The function and its argument shape are validated when the rule is converted. + Expr::Function(function) => { + if let FunctionArguments::List(args) = &function.args { + for arg in &args.args { + if let FunctionArg::Unnamed(FunctionArgExpr::Expr(expr)) = arg { + ensure_one_expr(expr, columns)?; + } + } + } + Ok(()) + } Expr::Value(_) => Ok(()), Expr::UnaryOp { expr, .. } => { ensure_one_expr(expr, columns)?; diff --git a/src/sql/src/statements/create.rs b/src/sql/src/statements/create.rs index 13fb4cb0a3a..54c2bdac19c 100644 --- a/src/sql/src/statements/create.rs +++ b/src/sql/src/statements/create.rs @@ -14,6 +14,7 @@ use std::collections::HashMap; use std::fmt::{Display, Formatter, Write}; +use std::ops::ControlFlow; use common_catalog::consts::FILE_ENGINE; use datatypes::json::{JSON2_DEFAULT_MAX_AUTO_EXPANDED_PATHS, JsonSettings}; @@ -22,7 +23,7 @@ use datatypes::schema::{CompressionOptions, FulltextOptions, SkippingIndexOption use itertools::Itertools; use serde::Serialize; use snafu::ResultExt; -use sqlparser::ast::{ColumnOptionDef, DataType, Expr}; +use sqlparser::ast::{ColumnOptionDef, DataType, Expr, visit_expressions_mut}; use sqlparser_derive::{Visit, VisitMut}; use crate::ast::{ColumnDef, Ident, ObjectName}; @@ -395,11 +396,17 @@ pub struct Partitions { } impl Partitions { - /// set quotes to all [Ident]s from column list + /// Sets quotes to all column [Ident]s, in the column list and in the expressions. pub fn set_quote(&mut self, quote_style: char) { self.column_list .iter_mut() .for_each(|c| c.quote_style = Some(quote_style)); + let _ = visit_expressions_mut(&mut self.exprs, |expr| { + if let Expr::Identifier(ident) = expr { + ident.quote_style = Some(quote_style); + } + ControlFlow::<()>::Continue(()) + }); } } diff --git a/src/table/src/metadata.rs b/src/table/src/metadata.rs index c8941bb1a82..57a4548160a 100644 --- a/src/table/src/metadata.rs +++ b/src/table/src/metadata.rs @@ -1093,6 +1093,16 @@ impl TableMeta { } ); } else { + // Partition rules keep bounds typed as the column, and functions such as + // `iceberg_hash` route by the typed value, so a new type would reroute or + // break routing. Time index widening above keeps timestamps comparable. + ensure!( + !self.partition_key_indices.contains(&index), + error::InvalidAlterRequestSnafu { + table: table_name, + err: format!("Not allowed to change partition column '{}'", column.name) + } + ); ensure!( column .data_type @@ -1911,6 +1921,50 @@ mod tests { } } + #[test] + fn test_modify_partition_column_type() { + // `col2` is nullable and not in the primary key, so only its partition role + // decides whether the change is allowed. + let alter_kind = AlterKind::ModifyColumnTypes { + columns: vec![ModifyColumnTypeRequest { + column_name: "col2".to_string(), + target_type: ConcreteDataType::float64_datatype(), + }], + }; + for (partition_key_indices, allowed) in [(vec![], true), (vec![2], false)] { + let meta = TableMetaBuilder::empty() + .schema(Arc::new(new_test_schema())) + .primary_key_indices(vec![0]) + .partition_key_indices(partition_key_indices) + .engine("engine") + .next_column_id(3) + .build() + .unwrap(); + let result = meta.builder_with_alter_kind("my_table", &alter_kind); + assert_eq!(result.is_ok(), allowed); + } + + // Widening the unit of a time index that is also a partition column is allowed. + let meta = TableMetaBuilder::empty() + .schema(Arc::new(new_test_schema())) + .primary_key_indices(vec![0]) + .partition_key_indices(vec![1]) + .engine("engine") + .next_column_id(3) + .build() + .unwrap(); + let alter_kind = AlterKind::ModifyColumnTypes { + columns: vec![ModifyColumnTypeRequest { + column_name: "ts".to_string(), + target_type: ConcreteDataType::timestamp_microsecond_datatype(), + }], + }; + assert!( + meta.builder_with_alter_kind("my_table", &alter_kind) + .is_ok() + ); + } + #[test] fn test_add_columns() { let schema = Arc::new(new_test_schema()); diff --git a/tests-integration/src/otlp.rs b/tests-integration/src/otlp.rs index 5ab264d23d0..98cc66bcafb 100644 --- a/tests-integration/src/otlp.rs +++ b/tests-integration/src/otlp.rs @@ -747,22 +747,22 @@ mod test { PRIMARY KEY ("service_name") ) PARTITION ON COLUMNS ("trace_id") ( - trace_id < '1', - trace_id >= '1' AND trace_id < '2', - trace_id >= '2' AND trace_id < '3', - trace_id >= '3' AND trace_id < '4', - trace_id >= '4' AND trace_id < '5', - trace_id >= '5' AND trace_id < '6', - trace_id >= '6' AND trace_id < '7', - trace_id >= '7' AND trace_id < '8', - trace_id >= '8' AND trace_id < '9', - trace_id >= '9' AND trace_id < 'a', - trace_id >= 'a' AND trace_id < 'b', - trace_id >= 'b' AND trace_id < 'c', - trace_id >= 'c' AND trace_id < 'd', - trace_id >= 'd' AND trace_id < 'e', - trace_id >= 'e' AND trace_id < 'f', - trace_id >= 'f' + "trace_id" < '1', + "trace_id" >= '1' AND "trace_id" < '2', + "trace_id" >= '2' AND "trace_id" < '3', + "trace_id" >= '3' AND "trace_id" < '4', + "trace_id" >= '4' AND "trace_id" < '5', + "trace_id" >= '5' AND "trace_id" < '6', + "trace_id" >= '6' AND "trace_id" < '7', + "trace_id" >= '7' AND "trace_id" < '8', + "trace_id" >= '8' AND "trace_id" < '9', + "trace_id" >= '9' AND "trace_id" < 'a', + "trace_id" >= 'a' AND "trace_id" < 'b', + "trace_id" >= 'b' AND "trace_id" < 'c', + "trace_id" >= 'c' AND "trace_id" < 'd', + "trace_id" >= 'd' AND "trace_id" < 'e', + "trace_id" >= 'e' AND "trace_id" < 'f', + "trace_id" >= 'f' ) ENGINE=mito WITH( diff --git a/tests-integration/src/tests/instance_test.rs b/tests-integration/src/tests/instance_test.rs index 22a0b993d17..481754e6df5 100644 --- a/tests-integration/src/tests/instance_test.rs +++ b/tests-integration/src/tests/instance_test.rs @@ -989,10 +989,10 @@ PARTITION ON COLUMNS (n) ( | | PRIMARY KEY ("n") | | | ) | | | PARTITION ON COLUMNS ("n") ( | -| | n < 1, | -| | n >= 1 AND n < 10, | -| | n >= 10 AND n < 100, | -| | n >= 100 | +| | "n" < 1, | +| | "n" >= 1 AND "n" < 10, | +| | "n" >= 10 AND "n" < 100, | +| | "n" >= 100 | | | ) | | | ENGINE=mito | | | | diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 11c41431fbc..1827e8b6581 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -10284,7 +10284,7 @@ pub async fn test_otlp_traces_v1(store_type: StorageType) { ) .await; - let expected_ddl = r#"[["mytable","CREATE TABLE IF NOT EXISTS \"mytable\" (\n \"timestamp\" TIMESTAMP(9) NOT NULL,\n \"timestamp_end\" TIMESTAMP(9) NULL COMPRESSION WITH(encoding = 'delta_binary_packed', type = 'zstd'),\n \"duration_nano\" BIGINT NULL,\n \"parent_span_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM') COMPRESSION WITH(encoding = 'plain', page_size = '64KiB', type = 'zstd'),\n \"trace_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_id\" STRING NULL COMPRESSION WITH(encoding = 'plain', page_size = '64KiB', type = 'zstd'),\n \"span_kind\" STRING NULL,\n \"span_name\" STRING NULL,\n \"span_status_code\" STRING NULL,\n \"span_status_message\" STRING NULL,\n \"trace_state\" STRING NULL,\n \"scope_name\" STRING NULL,\n \"scope_version\" STRING NULL,\n \"service_name\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_attributes.net.peer.ip\" STRING NULL,\n \"span_attributes.peer.service\" STRING NULL,\n \"span_events\" JSON NULL,\n \"span_links\" JSON NULL,\n TIME INDEX (\"timestamp\"),\n PRIMARY KEY (\"service_name\")\n)\nPARTITION ON COLUMNS (\"trace_id\") (\n trace_id < '1',\n trace_id >= '1' AND trace_id < '2',\n trace_id >= '2' AND trace_id < '3',\n trace_id >= '3' AND trace_id < '4',\n trace_id >= '4' AND trace_id < '5',\n trace_id >= '5' AND trace_id < '6',\n trace_id >= '6' AND trace_id < '7',\n trace_id >= '7' AND trace_id < '8',\n trace_id >= '8' AND trace_id < '9',\n trace_id >= '9' AND trace_id < 'a',\n trace_id >= 'a' AND trace_id < 'b',\n trace_id >= 'b' AND trace_id < 'c',\n trace_id >= 'c' AND trace_id < 'd',\n trace_id >= 'd' AND trace_id < 'e',\n trace_id >= 'e' AND trace_id < 'f',\n trace_id >= 'f'\n)\nENGINE=mito\nWITH(\n 'comment' = 'Created on insertion',\n append_mode = 'true',\n 'greptime.semantic.entity.service.id' = 'service_name',\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\n 'greptime.semantic.signal_type' = 'trace',\n 'greptime.semantic.source' = 'opentelemetry',\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\n table_data_model = 'greptime_trace_v1'\n)"]]"#; + let expected_ddl = r#"[["mytable","CREATE TABLE IF NOT EXISTS \"mytable\" (\n \"timestamp\" TIMESTAMP(9) NOT NULL,\n \"timestamp_end\" TIMESTAMP(9) NULL COMPRESSION WITH(encoding = 'delta_binary_packed', type = 'zstd'),\n \"duration_nano\" BIGINT NULL,\n \"parent_span_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM') COMPRESSION WITH(encoding = 'plain', page_size = '64KiB', type = 'zstd'),\n \"trace_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_id\" STRING NULL COMPRESSION WITH(encoding = 'plain', page_size = '64KiB', type = 'zstd'),\n \"span_kind\" STRING NULL,\n \"span_name\" STRING NULL,\n \"span_status_code\" STRING NULL,\n \"span_status_message\" STRING NULL,\n \"trace_state\" STRING NULL,\n \"scope_name\" STRING NULL,\n \"scope_version\" STRING NULL,\n \"service_name\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_attributes.net.peer.ip\" STRING NULL,\n \"span_attributes.peer.service\" STRING NULL,\n \"span_events\" JSON NULL,\n \"span_links\" JSON NULL,\n TIME INDEX (\"timestamp\"),\n PRIMARY KEY (\"service_name\")\n)\nPARTITION ON COLUMNS (\"trace_id\") (\n \"trace_id\" < '1',\n \"trace_id\" >= '1' AND \"trace_id\" < '2',\n \"trace_id\" >= '2' AND \"trace_id\" < '3',\n \"trace_id\" >= '3' AND \"trace_id\" < '4',\n \"trace_id\" >= '4' AND \"trace_id\" < '5',\n \"trace_id\" >= '5' AND \"trace_id\" < '6',\n \"trace_id\" >= '6' AND \"trace_id\" < '7',\n \"trace_id\" >= '7' AND \"trace_id\" < '8',\n \"trace_id\" >= '8' AND \"trace_id\" < '9',\n \"trace_id\" >= '9' AND \"trace_id\" < 'a',\n \"trace_id\" >= 'a' AND \"trace_id\" < 'b',\n \"trace_id\" >= 'b' AND \"trace_id\" < 'c',\n \"trace_id\" >= 'c' AND \"trace_id\" < 'd',\n \"trace_id\" >= 'd' AND \"trace_id\" < 'e',\n \"trace_id\" >= 'e' AND \"trace_id\" < 'f',\n \"trace_id\" >= 'f'\n)\nENGINE=mito\nWITH(\n 'comment' = 'Created on insertion',\n append_mode = 'true',\n 'greptime.semantic.entity.service.id' = 'service_name',\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\n 'greptime.semantic.signal_type' = 'trace',\n 'greptime.semantic.source' = 'opentelemetry',\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\n table_data_model = 'greptime_trace_v1'\n)"]]"#; validate_data( "otlp_traces", &client, @@ -10374,7 +10374,7 @@ pub async fn test_otlp_traces_v1(store_type: StorageType) { ) .await; assert_eq!(StatusCode::OK, res.status()); - let expected_ddl = r#"[["trace_table_part4","CREATE TABLE IF NOT EXISTS \"trace_table_part4\" (\n \"timestamp\" TIMESTAMP(9) NOT NULL,\n \"timestamp_end\" TIMESTAMP(9) NULL COMPRESSION WITH(encoding = 'delta_binary_packed', type = 'zstd'),\n \"duration_nano\" BIGINT NULL,\n \"parent_span_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM') COMPRESSION WITH(encoding = 'plain', page_size = '64KiB', type = 'zstd'),\n \"trace_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_id\" STRING NULL COMPRESSION WITH(encoding = 'plain', page_size = '64KiB', type = 'zstd'),\n \"span_kind\" STRING NULL,\n \"span_name\" STRING NULL,\n \"span_status_code\" STRING NULL,\n \"span_status_message\" STRING NULL,\n \"trace_state\" STRING NULL,\n \"scope_name\" STRING NULL,\n \"scope_version\" STRING NULL,\n \"service_name\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_attributes.net.peer.ip\" STRING NULL,\n \"span_attributes.peer.service\" STRING NULL,\n \"span_events\" JSON NULL,\n \"span_links\" JSON NULL,\n TIME INDEX (\"timestamp\"),\n PRIMARY KEY (\"service_name\")\n)\nPARTITION ON COLUMNS (\"trace_id\") (\n trace_id < '4',\n trace_id >= '4' AND trace_id < '8',\n trace_id >= '8' AND trace_id < 'c',\n trace_id >= 'c'\n)\nENGINE=mito\nWITH(\n 'comment' = 'Created on insertion',\n append_mode = 'true',\n 'greptime.semantic.entity.service.id' = 'service_name',\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\n 'greptime.semantic.signal_type' = 'trace',\n 'greptime.semantic.source' = 'opentelemetry',\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\n table_data_model = 'greptime_trace_v1'\n)"]]"#; + let expected_ddl = r#"[["trace_table_part4","CREATE TABLE IF NOT EXISTS \"trace_table_part4\" (\n \"timestamp\" TIMESTAMP(9) NOT NULL,\n \"timestamp_end\" TIMESTAMP(9) NULL COMPRESSION WITH(encoding = 'delta_binary_packed', type = 'zstd'),\n \"duration_nano\" BIGINT NULL,\n \"parent_span_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM') COMPRESSION WITH(encoding = 'plain', page_size = '64KiB', type = 'zstd'),\n \"trace_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_id\" STRING NULL COMPRESSION WITH(encoding = 'plain', page_size = '64KiB', type = 'zstd'),\n \"span_kind\" STRING NULL,\n \"span_name\" STRING NULL,\n \"span_status_code\" STRING NULL,\n \"span_status_message\" STRING NULL,\n \"trace_state\" STRING NULL,\n \"scope_name\" STRING NULL,\n \"scope_version\" STRING NULL,\n \"service_name\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_attributes.net.peer.ip\" STRING NULL,\n \"span_attributes.peer.service\" STRING NULL,\n \"span_events\" JSON NULL,\n \"span_links\" JSON NULL,\n TIME INDEX (\"timestamp\"),\n PRIMARY KEY (\"service_name\")\n)\nPARTITION ON COLUMNS (\"trace_id\") (\n \"trace_id\" < '4',\n \"trace_id\" >= '4' AND \"trace_id\" < '8',\n \"trace_id\" >= '8' AND \"trace_id\" < 'c',\n \"trace_id\" >= 'c'\n)\nENGINE=mito\nWITH(\n 'comment' = 'Created on insertion',\n append_mode = 'true',\n 'greptime.semantic.entity.service.id' = 'service_name',\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\n 'greptime.semantic.signal_type' = 'trace',\n 'greptime.semantic.source' = 'opentelemetry',\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\n table_data_model = 'greptime_trace_v1'\n)"]]"#; validate_data( "otlp_traces", &client, @@ -10411,7 +10411,7 @@ pub async fn test_otlp_traces_v1(store_type: StorageType) { ) .await; assert_eq!(StatusCode::OK, res.status()); - let expected_ddl = r#"[["trace_table_part32","CREATE TABLE IF NOT EXISTS \"trace_table_part32\" (\n \"timestamp\" TIMESTAMP(9) NOT NULL,\n \"timestamp_end\" TIMESTAMP(9) NULL COMPRESSION WITH(encoding = 'delta_binary_packed', type = 'zstd'),\n \"duration_nano\" BIGINT NULL,\n \"parent_span_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM') COMPRESSION WITH(encoding = 'plain', page_size = '64KiB', type = 'zstd'),\n \"trace_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_id\" STRING NULL COMPRESSION WITH(encoding = 'plain', page_size = '64KiB', type = 'zstd'),\n \"span_kind\" STRING NULL,\n \"span_name\" STRING NULL,\n \"span_status_code\" STRING NULL,\n \"span_status_message\" STRING NULL,\n \"trace_state\" STRING NULL,\n \"scope_name\" STRING NULL,\n \"scope_version\" STRING NULL,\n \"service_name\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_attributes.net.peer.ip\" STRING NULL,\n \"span_attributes.peer.service\" STRING NULL,\n \"span_events\" JSON NULL,\n \"span_links\" JSON NULL,\n TIME INDEX (\"timestamp\"),\n PRIMARY KEY (\"service_name\")\n)\nPARTITION ON COLUMNS (\"trace_id\") (\n trace_id < '08',\n trace_id >= '08' AND trace_id < '10',\n trace_id >= '10' AND trace_id < '18',\n trace_id >= '18' AND trace_id < '20',\n trace_id >= '20' AND trace_id < '28',\n trace_id >= '28' AND trace_id < '30',\n trace_id >= '30' AND trace_id < '38',\n trace_id >= '38' AND trace_id < '40',\n trace_id >= '40' AND trace_id < '48',\n trace_id >= '48' AND trace_id < '50',\n trace_id >= '50' AND trace_id < '58',\n trace_id >= '58' AND trace_id < '60',\n trace_id >= '60' AND trace_id < '68',\n trace_id >= '68' AND trace_id < '70',\n trace_id >= '70' AND trace_id < '78',\n trace_id >= '78' AND trace_id < '80',\n trace_id >= '80' AND trace_id < '88',\n trace_id >= '88' AND trace_id < '90',\n trace_id >= '90' AND trace_id < '98',\n trace_id >= '98' AND trace_id < 'a0',\n trace_id >= 'a0' AND trace_id < 'a8',\n trace_id >= 'a8' AND trace_id < 'b0',\n trace_id >= 'b0' AND trace_id < 'b8',\n trace_id >= 'b8' AND trace_id < 'c0',\n trace_id >= 'c0' AND trace_id < 'c8',\n trace_id >= 'c8' AND trace_id < 'd0',\n trace_id >= 'd0' AND trace_id < 'd8',\n trace_id >= 'd8' AND trace_id < 'e0',\n trace_id >= 'e0' AND trace_id < 'e8',\n trace_id >= 'e8' AND trace_id < 'f0',\n trace_id >= 'f0' AND trace_id < 'f8',\n trace_id >= 'f8'\n)\nENGINE=mito\nWITH(\n 'comment' = 'Created on insertion',\n append_mode = 'true',\n 'greptime.semantic.entity.service.id' = 'service_name',\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\n 'greptime.semantic.signal_type' = 'trace',\n 'greptime.semantic.source' = 'opentelemetry',\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\n table_data_model = 'greptime_trace_v1'\n)"]]"#; + let expected_ddl = r#"[["trace_table_part32","CREATE TABLE IF NOT EXISTS \"trace_table_part32\" (\n \"timestamp\" TIMESTAMP(9) NOT NULL,\n \"timestamp_end\" TIMESTAMP(9) NULL COMPRESSION WITH(encoding = 'delta_binary_packed', type = 'zstd'),\n \"duration_nano\" BIGINT NULL,\n \"parent_span_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM') COMPRESSION WITH(encoding = 'plain', page_size = '64KiB', type = 'zstd'),\n \"trace_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_id\" STRING NULL COMPRESSION WITH(encoding = 'plain', page_size = '64KiB', type = 'zstd'),\n \"span_kind\" STRING NULL,\n \"span_name\" STRING NULL,\n \"span_status_code\" STRING NULL,\n \"span_status_message\" STRING NULL,\n \"trace_state\" STRING NULL,\n \"scope_name\" STRING NULL,\n \"scope_version\" STRING NULL,\n \"service_name\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_attributes.net.peer.ip\" STRING NULL,\n \"span_attributes.peer.service\" STRING NULL,\n \"span_events\" JSON NULL,\n \"span_links\" JSON NULL,\n TIME INDEX (\"timestamp\"),\n PRIMARY KEY (\"service_name\")\n)\nPARTITION ON COLUMNS (\"trace_id\") (\n \"trace_id\" < '08',\n \"trace_id\" >= '08' AND \"trace_id\" < '10',\n \"trace_id\" >= '10' AND \"trace_id\" < '18',\n \"trace_id\" >= '18' AND \"trace_id\" < '20',\n \"trace_id\" >= '20' AND \"trace_id\" < '28',\n \"trace_id\" >= '28' AND \"trace_id\" < '30',\n \"trace_id\" >= '30' AND \"trace_id\" < '38',\n \"trace_id\" >= '38' AND \"trace_id\" < '40',\n \"trace_id\" >= '40' AND \"trace_id\" < '48',\n \"trace_id\" >= '48' AND \"trace_id\" < '50',\n \"trace_id\" >= '50' AND \"trace_id\" < '58',\n \"trace_id\" >= '58' AND \"trace_id\" < '60',\n \"trace_id\" >= '60' AND \"trace_id\" < '68',\n \"trace_id\" >= '68' AND \"trace_id\" < '70',\n \"trace_id\" >= '70' AND \"trace_id\" < '78',\n \"trace_id\" >= '78' AND \"trace_id\" < '80',\n \"trace_id\" >= '80' AND \"trace_id\" < '88',\n \"trace_id\" >= '88' AND \"trace_id\" < '90',\n \"trace_id\" >= '90' AND \"trace_id\" < '98',\n \"trace_id\" >= '98' AND \"trace_id\" < 'a0',\n \"trace_id\" >= 'a0' AND \"trace_id\" < 'a8',\n \"trace_id\" >= 'a8' AND \"trace_id\" < 'b0',\n \"trace_id\" >= 'b0' AND \"trace_id\" < 'b8',\n \"trace_id\" >= 'b8' AND \"trace_id\" < 'c0',\n \"trace_id\" >= 'c0' AND \"trace_id\" < 'c8',\n \"trace_id\" >= 'c8' AND \"trace_id\" < 'd0',\n \"trace_id\" >= 'd0' AND \"trace_id\" < 'd8',\n \"trace_id\" >= 'd8' AND \"trace_id\" < 'e0',\n \"trace_id\" >= 'e0' AND \"trace_id\" < 'e8',\n \"trace_id\" >= 'e8' AND \"trace_id\" < 'f0',\n \"trace_id\" >= 'f0' AND \"trace_id\" < 'f8',\n \"trace_id\" >= 'f8'\n)\nENGINE=mito\nWITH(\n 'comment' = 'Created on insertion',\n append_mode = 'true',\n 'greptime.semantic.entity.service.id' = 'service_name',\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\n 'greptime.semantic.signal_type' = 'trace',\n 'greptime.semantic.source' = 'opentelemetry',\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\n table_data_model = 'greptime_trace_v1'\n)"]]"#; validate_data( "otlp_traces", &client, @@ -13576,7 +13576,7 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { .await; assert_eq!(StatusCode::OK, res.status()); - let trace_table_sql = "[[\"mytable\",\"CREATE TABLE IF NOT EXISTS \\\"mytable\\\" (\\n \\\"timestamp\\\" TIMESTAMP(9) NOT NULL,\\n \\\"timestamp_end\\\" TIMESTAMP(9) NULL COMPRESSION WITH(encoding = 'delta_binary_packed', type = 'zstd'),\\n \\\"duration_nano\\\" BIGINT NULL,\\n \\\"parent_span_id\\\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM') COMPRESSION WITH(encoding = 'plain', page_size = '64KiB', type = 'zstd'),\\n \\\"trace_id\\\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\\n \\\"span_id\\\" STRING NULL COMPRESSION WITH(encoding = 'plain', page_size = '64KiB', type = 'zstd'),\\n \\\"span_kind\\\" STRING NULL,\\n \\\"span_name\\\" STRING NULL,\\n \\\"span_status_code\\\" STRING NULL,\\n \\\"span_status_message\\\" STRING NULL,\\n \\\"trace_state\\\" STRING NULL,\\n \\\"scope_name\\\" STRING NULL,\\n \\\"scope_version\\\" STRING NULL,\\n \\\"service_name\\\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\\n \\\"span_attributes.operation.type\\\" STRING NULL,\\n \\\"span_attributes.net.peer.ip\\\" STRING NULL,\\n \\\"span_attributes.peer.service\\\" STRING NULL,\\n \\\"span_events\\\" JSON NULL,\\n \\\"span_links\\\" JSON NULL,\\n TIME INDEX (\\\"timestamp\\\"),\\n PRIMARY KEY (\\\"service_name\\\")\\n)\\nPARTITION ON COLUMNS (\\\"trace_id\\\") (\\n trace_id < '1',\\n trace_id >= '1' AND trace_id < '2',\\n trace_id >= '2' AND trace_id < '3',\\n trace_id >= '3' AND trace_id < '4',\\n trace_id >= '4' AND trace_id < '5',\\n trace_id >= '5' AND trace_id < '6',\\n trace_id >= '6' AND trace_id < '7',\\n trace_id >= '7' AND trace_id < '8',\\n trace_id >= '8' AND trace_id < '9',\\n trace_id >= '9' AND trace_id < 'a',\\n trace_id >= 'a' AND trace_id < 'b',\\n trace_id >= 'b' AND trace_id < 'c',\\n trace_id >= 'c' AND trace_id < 'd',\\n trace_id >= 'd' AND trace_id < 'e',\\n trace_id >= 'e' AND trace_id < 'f',\\n trace_id >= 'f'\\n)\\nENGINE=mito\\nWITH(\\n 'comment' = 'Created on insertion',\\n append_mode = 'true',\\n 'greptime.semantic.entity.service.id' = 'service_name',\\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\\n 'greptime.semantic.signal_type' = 'trace',\\n 'greptime.semantic.source' = 'opentelemetry',\\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\\n table_data_model = 'greptime_trace_v1',\\n ttl = '7days'\\n)\"]]"; + let trace_table_sql = "[[\"mytable\",\"CREATE TABLE IF NOT EXISTS \\\"mytable\\\" (\\n \\\"timestamp\\\" TIMESTAMP(9) NOT NULL,\\n \\\"timestamp_end\\\" TIMESTAMP(9) NULL COMPRESSION WITH(encoding = 'delta_binary_packed', type = 'zstd'),\\n \\\"duration_nano\\\" BIGINT NULL,\\n \\\"parent_span_id\\\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM') COMPRESSION WITH(encoding = 'plain', page_size = '64KiB', type = 'zstd'),\\n \\\"trace_id\\\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\\n \\\"span_id\\\" STRING NULL COMPRESSION WITH(encoding = 'plain', page_size = '64KiB', type = 'zstd'),\\n \\\"span_kind\\\" STRING NULL,\\n \\\"span_name\\\" STRING NULL,\\n \\\"span_status_code\\\" STRING NULL,\\n \\\"span_status_message\\\" STRING NULL,\\n \\\"trace_state\\\" STRING NULL,\\n \\\"scope_name\\\" STRING NULL,\\n \\\"scope_version\\\" STRING NULL,\\n \\\"service_name\\\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\\n \\\"span_attributes.operation.type\\\" STRING NULL,\\n \\\"span_attributes.net.peer.ip\\\" STRING NULL,\\n \\\"span_attributes.peer.service\\\" STRING NULL,\\n \\\"span_events\\\" JSON NULL,\\n \\\"span_links\\\" JSON NULL,\\n TIME INDEX (\\\"timestamp\\\"),\\n PRIMARY KEY (\\\"service_name\\\")\\n)\\nPARTITION ON COLUMNS (\\\"trace_id\\\") (\\n \\\"trace_id\\\" < '1',\\n \\\"trace_id\\\" >= '1' AND \\\"trace_id\\\" < '2',\\n \\\"trace_id\\\" >= '2' AND \\\"trace_id\\\" < '3',\\n \\\"trace_id\\\" >= '3' AND \\\"trace_id\\\" < '4',\\n \\\"trace_id\\\" >= '4' AND \\\"trace_id\\\" < '5',\\n \\\"trace_id\\\" >= '5' AND \\\"trace_id\\\" < '6',\\n \\\"trace_id\\\" >= '6' AND \\\"trace_id\\\" < '7',\\n \\\"trace_id\\\" >= '7' AND \\\"trace_id\\\" < '8',\\n \\\"trace_id\\\" >= '8' AND \\\"trace_id\\\" < '9',\\n \\\"trace_id\\\" >= '9' AND \\\"trace_id\\\" < 'a',\\n \\\"trace_id\\\" >= 'a' AND \\\"trace_id\\\" < 'b',\\n \\\"trace_id\\\" >= 'b' AND \\\"trace_id\\\" < 'c',\\n \\\"trace_id\\\" >= 'c' AND \\\"trace_id\\\" < 'd',\\n \\\"trace_id\\\" >= 'd' AND \\\"trace_id\\\" < 'e',\\n \\\"trace_id\\\" >= 'e' AND \\\"trace_id\\\" < 'f',\\n \\\"trace_id\\\" >= 'f'\\n)\\nENGINE=mito\\nWITH(\\n 'comment' = 'Created on insertion',\\n append_mode = 'true',\\n 'greptime.semantic.entity.service.id' = 'service_name',\\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\\n 'greptime.semantic.signal_type' = 'trace',\\n 'greptime.semantic.source' = 'opentelemetry',\\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\\n table_data_model = 'greptime_trace_v1',\\n ttl = '7days'\\n)\"]]"; validate_data( "trace_v1_create_table", &client, diff --git a/tests/cases/distributed/repartition/repartition.result b/tests/cases/distributed/repartition/repartition.result index 2b03139abd4..5ed205e050e 100644 --- a/tests/cases/distributed/repartition/repartition.result +++ b/tests/cases/distributed/repartition/repartition.result @@ -36,9 +36,9 @@ SHOW CREATE TABLE alter_repartition_table; | | PRIMARY KEY ("device_id") | | | ) | | | PARTITION ON COLUMNS ("device_id", "area") ( | -| | device_id < 100, | -| | device_id >= 100 AND device_id < 200, | -| | device_id >= 200 | +| | "device_id" < 100, | +| | "device_id" >= 100 AND "device_id" < 200, | +| | "device_id" >= 200 | | | ) | | | ENGINE=mito | | | | @@ -66,9 +66,9 @@ SHOW CREATE TABLE alter_repartition_table; | | PRIMARY KEY ("device_id") | | | ) | | | PARTITION ON COLUMNS ("device_id", "area") ( | -| | device_id < 100, | -| | device_id >= 100 AND device_id < 200, | -| | device_id >= 200 | +| | "device_id" < 100, | +| | "device_id" >= 100 AND "device_id" < 200, | +| | "device_id" >= 200 | | | ) | | | ENGINE=mito | | | | @@ -197,9 +197,9 @@ SHOW CREATE TABLE metric_physical_table; | | PRIMARY KEY ("host") | | | ) | | | PARTITION ON COLUMNS ("host") ( | -| | host < 'h1', | -| | host >= 'h1' AND host < 'h2', | -| | host >= 'h2' | +| | "host" < 'h1', | +| | "host" >= 'h1' AND "host" < 'h2', | +| | "host" >= 'h2' | | | ) | | | ENGINE=metric | | | WITH( | @@ -250,9 +250,9 @@ SHOW CREATE TABLE metric_physical_table; | | PRIMARY KEY ("host") | | | ) | | | PARTITION ON COLUMNS ("host") ( | -| | host < 'h1', | -| | host >= 'h1' AND host < 'h2', | -| | host >= 'h2' | +| | "host" < 'h1', | +| | "host" >= 'h1' AND "host" < 'h2', | +| | "host" >= 'h2' | | | ) | | | ENGINE=metric | | | WITH( | diff --git a/tests/cases/standalone/common/create/metric_engine_partition.result b/tests/cases/standalone/common/create/metric_engine_partition.result index d45f6adf326..0497364e2bf 100644 --- a/tests/cases/standalone/common/create/metric_engine_partition.result +++ b/tests/cases/standalone/common/create/metric_engine_partition.result @@ -98,9 +98,9 @@ show create table logical_table_2; | | PRIMARY KEY ("another_partition_key", "host", "one_partition_key") | | | ) | | | PARTITION ON COLUMNS ("host", "one_partition_key", "another_partition_key") ( | -| | host <= 'host1', | -| | host > 'host1' AND host <= 'host2', | -| | host > 'host2' | +| | "host" <= 'host1', | +| | "host" > 'host1' AND "host" <= 'host2', | +| | "host" > 'host2' | | | ) | | | ENGINE=metric | | | WITH( | @@ -227,9 +227,9 @@ show create table logical_table_3; | | PRIMARY KEY ("a", "another_partition_key", "host", "one_partition_key", "z") | | | ) | | | PARTITION ON COLUMNS ("host", "one_partition_key", "another_partition_key") ( | -| | host <= 'host1', | -| | host > 'host1' AND host <= 'host2', | -| | host > 'host2' | +| | "host" <= 'host1', | +| | "host" > 'host1' AND "host" <= 'host2', | +| | "host" > 'host2' | | | ) | | | ENGINE=metric | | | WITH( | @@ -319,9 +319,9 @@ show create table logical_table_4; | | PRIMARY KEY ("another_partition_key", "host", "one_partition_key") | | | ) | | | PARTITION ON COLUMNS ("host", "one_partition_key", "another_partition_key") ( | -| | host <= 'host1', | -| | host > 'host1' AND host <= 'host2', | -| | host > 'host2' | +| | "host" <= 'host1', | +| | "host" > 'host1' AND "host" <= 'host2', | +| | "host" > 'host2' | | | ) | | | ENGINE=metric | | | WITH( | diff --git a/tests/cases/standalone/common/partition_iceberg_hash.result b/tests/cases/standalone/common/partition_iceberg_hash.result new file mode 100644 index 00000000000..18ceee583d2 --- /dev/null +++ b/tests/cases/standalone/common/partition_iceberg_hash.result @@ -0,0 +1,184 @@ +CREATE TABLE hash_routes ( + host STRING, + tenant BIGINT UNSIGNED, + ts TIMESTAMP TIME INDEX, + v INT, + PRIMARY KEY (host, tenant) +) +PARTITION ON COLUMNS (host, tenant) ( + iceberg_hash(host) < 0 AND tenant < 100, + iceberg_hash(host) < 0 AND tenant >= 100, + ICEBERG_HASH(host) >= 0 AND iceberg_hash(tenant) < 0, + 0 <= iceberg_hash(host) AND iceberg_hash(tenant) >= 0 +); + +Affected Rows: 0 + +SHOW CREATE TABLE hash_routes; + ++-------------+-------------------------------------------------------------+ +| Table | Create Table | ++-------------+-------------------------------------------------------------+ +| hash_routes | CREATE TABLE IF NOT EXISTS "hash_routes" ( | +| | "host" STRING NULL, | +| | "tenant" BIGINT UNSIGNED NULL, | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | "v" INT NULL, | +| | TIME INDEX ("ts"), | +| | PRIMARY KEY ("host", "tenant") | +| | ) | +| | PARTITION ON COLUMNS ("host", "tenant") ( | +| | iceberg_hash("host") < 0 AND "tenant" < 100, | +| | iceberg_hash("host") < 0 AND "tenant" >= 100, | +| | iceberg_hash("host") >= 0 AND iceberg_hash("tenant") < 0, | +| | iceberg_hash("host") >= 0 AND iceberg_hash("tenant") >= 0 | +| | ) | +| | ENGINE=mito | +| | | ++-------------+-------------------------------------------------------------+ + +SELECT partition_name, partition_description FROM information_schema.partitions WHERE table_name = 'hash_routes' ORDER BY partition_name; + ++----------------+-------------------------------------------------------+ +| partition_name | partition_description | ++----------------+-------------------------------------------------------+ +| p0 | iceberg_hash(host) < 0 AND tenant < 100 | +| p1 | iceberg_hash(host) < 0 AND tenant >= 100 | +| p2 | iceberg_hash(host) >= 0 AND iceberg_hash(tenant) < 0 | +| p3 | iceberg_hash(host) >= 0 AND iceberg_hash(tenant) >= 0 | ++----------------+-------------------------------------------------------+ + +INSERT INTO hash_routes VALUES + ('a', 1, 1, 1), + ('b', 100, 2, 2), + ('server42', 18446744073709551615, 3, 3), + ('中🙂', 0, 4, 4), + ('', 7, 5, 5), + (NULL, 7, 6, 6), + ('a', NULL, 7, 7); + +Affected Rows: 7 + +SELECT * FROM hash_routes ORDER BY ts; + ++----------+----------------------+-------------------------+---+ +| host | tenant | ts | v | ++----------+----------------------+-------------------------+---+ +| a | 1 | 1970-01-01T00:00:00.001 | 1 | +| b | 100 | 1970-01-01T00:00:00.002 | 2 | +| server42 | 18446744073709551615 | 1970-01-01T00:00:00.003 | 3 | +| 中🙂 | 0 | 1970-01-01T00:00:00.004 | 4 | +| | 7 | 1970-01-01T00:00:00.005 | 5 | +| | 7 | 1970-01-01T00:00:00.006 | 6 | +| a | | 1970-01-01T00:00:00.007 | 7 | ++----------+----------------------+-------------------------+---+ + +-- A filter on the hashed column must not prune the region holding the row. +SELECT * FROM hash_routes WHERE host = 'server42'; + ++----------+----------------------+-------------------------+---+ +| host | tenant | ts | v | ++----------+----------------------+-------------------------+---+ +| server42 | 18446744073709551615 | 1970-01-01T00:00:00.003 | 3 | ++----------+----------------------+-------------------------+---+ + +SELECT * FROM hash_routes WHERE host IS NULL; + ++------+--------+-------------------------+---+ +| host | tenant | ts | v | ++------+--------+-------------------------+---+ +| | 7 | 1970-01-01T00:00:00.006 | 6 | ++------+--------+-------------------------+---+ + +DELETE FROM hash_routes WHERE host = 'a'; + +Affected Rows: 2 + +SELECT * FROM hash_routes ORDER BY ts; + ++----------+----------------------+-------------------------+---+ +| host | tenant | ts | v | ++----------+----------------------+-------------------------+---+ +| b | 100 | 1970-01-01T00:00:00.002 | 2 | +| server42 | 18446744073709551615 | 1970-01-01T00:00:00.003 | 3 | +| 中🙂 | 0 | 1970-01-01T00:00:00.004 | 4 | +| | 7 | 1970-01-01T00:00:00.005 | 5 | +| | 7 | 1970-01-01T00:00:00.006 | 6 | ++----------+----------------------+-------------------------+---+ + +DROP TABLE hash_routes; + +Affected Rows: 0 + +CREATE TABLE hash_float (f DOUBLE, ts TIMESTAMP TIME INDEX, PRIMARY KEY (f)) +PARTITION ON COLUMNS (f) ( + iceberg_hash(f) < 0, + iceberg_hash(f) >= 0 +); + +Error: 1004(InvalidArguments), Invalid partition rule: partition function iceberg_hash does not support column type Float64 + +CREATE TABLE hash_two_args (a STRING, b STRING, ts TIMESTAMP TIME INDEX, PRIMARY KEY (a, b)) +PARTITION ON COLUMNS (a, b) ( + iceberg_hash(a, b) < 0, + iceberg_hash(a, b) >= 0 +); + +Error: 1004(InvalidArguments), Invalid partition rule: unsupported partition function call iceberg_hash(a, b) + +CREATE TABLE hash_gap (a STRING, ts TIMESTAMP TIME INDEX, PRIMARY KEY (a)) +PARTITION ON COLUMNS (a) ( + iceberg_hash(a) < 0, + iceberg_hash(a) > 0 +); + +Error: 1004(InvalidArguments), Checkpoint `iceberg_hash(a)=0` is not covered + +CREATE TABLE quoted_hash ( + "host-name" STRING, + uid BIGINT, + ts TIMESTAMP TIME INDEX, + PRIMARY KEY ("host-name") +) +PARTITION ON COLUMNS ("host-name", uid) ( + iceberg_hash("host-name") < 0, + iceberg_hash("host-name") >= 0 AND iceberg_hash(uid) < 0, + iceberg_hash("host-name") >= 0 AND iceberg_hash(uid) >= 0 +); + +Affected Rows: 0 + +SHOW CREATE TABLE quoted_hash; + ++-------------+---------------------------------------------------------------+ +| Table | Create Table | ++-------------+---------------------------------------------------------------+ +| quoted_hash | CREATE TABLE IF NOT EXISTS "quoted_hash" ( | +| | "host-name" STRING NULL, | +| | "uid" BIGINT NULL, | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | TIME INDEX ("ts"), | +| | PRIMARY KEY ("host-name") | +| | ) | +| | PARTITION ON COLUMNS ("host-name", "uid") ( | +| | iceberg_hash("host-name") < 0, | +| | iceberg_hash("host-name") >= 0 AND iceberg_hash("uid") < 0, | +| | iceberg_hash("host-name") >= 0 AND iceberg_hash("uid") >= 0 | +| | ) | +| | ENGINE=mito | +| | | ++-------------+---------------------------------------------------------------+ + +-- A partition column keeps its type, even when nullable and outside the primary key. +ALTER TABLE quoted_hash MODIFY COLUMN uid DOUBLE; + +Error: 1004(InvalidArguments), Invalid alter table(quoted_hash) request: Not allowed to change partition column 'uid' + +ALTER TABLE quoted_hash MODIFY COLUMN uid STRING; + +Error: 1004(InvalidArguments), Invalid alter table(quoted_hash) request: Not allowed to change partition column 'uid' + +DROP TABLE quoted_hash; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/partition_iceberg_hash.sql b/tests/cases/standalone/common/partition_iceberg_hash.sql new file mode 100644 index 00000000000..2a787b4c8bd --- /dev/null +++ b/tests/cases/standalone/common/partition_iceberg_hash.sql @@ -0,0 +1,78 @@ +CREATE TABLE hash_routes ( + host STRING, + tenant BIGINT UNSIGNED, + ts TIMESTAMP TIME INDEX, + v INT, + PRIMARY KEY (host, tenant) +) +PARTITION ON COLUMNS (host, tenant) ( + iceberg_hash(host) < 0 AND tenant < 100, + iceberg_hash(host) < 0 AND tenant >= 100, + ICEBERG_HASH(host) >= 0 AND iceberg_hash(tenant) < 0, + 0 <= iceberg_hash(host) AND iceberg_hash(tenant) >= 0 +); + +SHOW CREATE TABLE hash_routes; + +SELECT partition_name, partition_description FROM information_schema.partitions WHERE table_name = 'hash_routes' ORDER BY partition_name; + +INSERT INTO hash_routes VALUES + ('a', 1, 1, 1), + ('b', 100, 2, 2), + ('server42', 18446744073709551615, 3, 3), + ('中🙂', 0, 4, 4), + ('', 7, 5, 5), + (NULL, 7, 6, 6), + ('a', NULL, 7, 7); + +SELECT * FROM hash_routes ORDER BY ts; + +-- A filter on the hashed column must not prune the region holding the row. +SELECT * FROM hash_routes WHERE host = 'server42'; + +SELECT * FROM hash_routes WHERE host IS NULL; + +DELETE FROM hash_routes WHERE host = 'a'; + +SELECT * FROM hash_routes ORDER BY ts; + +DROP TABLE hash_routes; + +CREATE TABLE hash_float (f DOUBLE, ts TIMESTAMP TIME INDEX, PRIMARY KEY (f)) +PARTITION ON COLUMNS (f) ( + iceberg_hash(f) < 0, + iceberg_hash(f) >= 0 +); + +CREATE TABLE hash_two_args (a STRING, b STRING, ts TIMESTAMP TIME INDEX, PRIMARY KEY (a, b)) +PARTITION ON COLUMNS (a, b) ( + iceberg_hash(a, b) < 0, + iceberg_hash(a, b) >= 0 +); + +CREATE TABLE hash_gap (a STRING, ts TIMESTAMP TIME INDEX, PRIMARY KEY (a)) +PARTITION ON COLUMNS (a) ( + iceberg_hash(a) < 0, + iceberg_hash(a) > 0 +); + +CREATE TABLE quoted_hash ( + "host-name" STRING, + uid BIGINT, + ts TIMESTAMP TIME INDEX, + PRIMARY KEY ("host-name") +) +PARTITION ON COLUMNS ("host-name", uid) ( + iceberg_hash("host-name") < 0, + iceberg_hash("host-name") >= 0 AND iceberg_hash(uid) < 0, + iceberg_hash("host-name") >= 0 AND iceberg_hash(uid) >= 0 +); + +SHOW CREATE TABLE quoted_hash; + +-- A partition column keeps its type, even when nullable and outside the primary key. +ALTER TABLE quoted_hash MODIFY COLUMN uid DOUBLE; + +ALTER TABLE quoted_hash MODIFY COLUMN uid STRING; + +DROP TABLE quoted_hash; diff --git a/tests/cases/standalone/common/show/show_create.result b/tests/cases/standalone/common/show/show_create.result index ba2c39386d6..8a5010adb03 100644 --- a/tests/cases/standalone/common/show/show_create.result +++ b/tests/cases/standalone/common/show/show_create.result @@ -35,9 +35,9 @@ SHOW CREATE TABLE system_metrics; | | PRIMARY KEY ("id", "host") | | | ) | | | PARTITION ON COLUMNS ("id") ( | -| | id < 5, | -| | id >= 5 AND id < 9, | -| | id >= 9 | +| | "id" < 5, | +| | "id" >= 5 AND "id" < 9, | +| | "id" >= 9 | | | ) | | | ENGINE=mito | | | WITH( | @@ -204,8 +204,8 @@ SHOW CREATE TABLE partition_num_hint_table; | | PRIMARY KEY ("host") | | | ) | | | PARTITION ON COLUMNS ("host") ( | -| | host < 'm', | -| | host >= 'm' | +| | "host" < 'm', | +| | "host" >= 'm' | | | ) | | | ENGINE=mito | | | WITH( | @@ -229,8 +229,8 @@ SHOW CREATE TABLE partition_num_hint_table; | | PRIMARY KEY ("host") | | | ) | | | PARTITION ON COLUMNS ("host") ( | -| | host < 'm', | -| | host >= 'm' | +| | "host" < 'm', | +| | "host" >= 'm' | | | ) | | | ENGINE=mito | | | WITH( | @@ -262,8 +262,8 @@ SHOW CREATE TABLE partition_num_hint_table; | | PRIMARY KEY ("host") | | | ) | | | PARTITION ON COLUMNS ("host") ( | -| | host < 'm', | -| | host >= 'm' | +| | "host" < 'm', | +| | "host" >= 'm' | | | ) | | | ENGINE=mito | | | | diff --git a/tests/compatibility/cases/partition_iceberg_hash/case.toml b/tests/compatibility/cases/partition_iceberg_hash/case.toml new file mode 100644 index 00000000000..6935dccc9eb --- /dev/null +++ b/tests/compatibility/cases/partition_iceberg_hash/case.toml @@ -0,0 +1,9 @@ +name = "partition_iceberg_hash" +reason = "Preserve iceberg_hash partition rules and their routing across restart and upgrade." +introduced_by = "iceberg_hash partition function" +topologies = ["standalone", "distributed"] +# The runner compares core versions only, so this also admits 1.4.0 prereleases. +from_range = [">=v1.4.0"] +to_range = [">=v1.4.0"] +features = ["partition", "table", "metadata"] +owner = "frontend" diff --git a/tests/compatibility/cases/partition_iceberg_hash/setup.sql b/tests/compatibility/cases/partition_iceberg_hash/setup.sql new file mode 100644 index 00000000000..902db740d72 --- /dev/null +++ b/tests/compatibility/cases/partition_iceberg_hash/setup.sql @@ -0,0 +1,22 @@ +CREATE TABLE iceberg_hash_routes ( + host STRING, + tenant BIGINT UNSIGNED, + ts TIMESTAMP TIME INDEX, + v INT, + PRIMARY KEY (host, tenant) +) +PARTITION ON COLUMNS (host, tenant) ( + iceberg_hash(host) < 0 AND tenant < 100, + iceberg_hash(host) < 0 AND tenant >= 100, + iceberg_hash(host) >= 0 AND iceberg_hash(tenant) < 0, + iceberg_hash(host) >= 0 AND iceberg_hash(tenant) >= 0 +); + +INSERT INTO iceberg_hash_routes VALUES + ('a', 1, 1, 1), + ('b', 100, 2, 2), + ('server42', 18446744073709551615, 3, 3), + ('中🙂', 0, 4, 4), + (NULL, 7, 5, 5); + +ADMIN FLUSH_TABLE('iceberg_hash_routes'); diff --git a/tests/compatibility/cases/partition_iceberg_hash/verify.result b/tests/compatibility/cases/partition_iceberg_hash/verify.result new file mode 100644 index 00000000000..c1723e2a4c1 --- /dev/null +++ b/tests/compatibility/cases/partition_iceberg_hash/verify.result @@ -0,0 +1,79 @@ +SHOW CREATE TABLE iceberg_hash_routes; + ++---------------------+-------------------------------------------------------------+ +| Table | Create Table | ++---------------------+-------------------------------------------------------------+ +| iceberg_hash_routes | CREATE TABLE IF NOT EXISTS "iceberg_hash_routes" ( | +| | "host" STRING NULL, | +| | "tenant" BIGINT UNSIGNED NULL, | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | "v" INT NULL, | +| | TIME INDEX ("ts"), | +| | PRIMARY KEY ("host", "tenant") | +| | ) | +| | PARTITION ON COLUMNS ("host", "tenant") ( | +| | iceberg_hash("host") < 0 AND "tenant" < 100, | +| | iceberg_hash("host") < 0 AND "tenant" >= 100, | +| | iceberg_hash("host") >= 0 AND iceberg_hash("tenant") < 0, | +| | iceberg_hash("host") >= 0 AND iceberg_hash("tenant") >= 0 | +| | ) | +| | ENGINE=mito | +| | | ++---------------------+-------------------------------------------------------------+ + +SELECT * FROM iceberg_hash_routes ORDER BY ts; + ++----------+----------------------+-------------------------+---+ +| host | tenant | ts | v | ++----------+----------------------+-------------------------+---+ +| a | 1 | 1970-01-01T00:00:00.001 | 1 | +| b | 100 | 1970-01-01T00:00:00.002 | 2 | +| server42 | 18446744073709551615 | 1970-01-01T00:00:00.003 | 3 | +| 中🙂 | 0 | 1970-01-01T00:00:00.004 | 4 | +| | 7 | 1970-01-01T00:00:00.005 | 5 | ++----------+----------------------+-------------------------+---+ + +-- Overwrites must route to the regions holding the flushed rows. +INSERT INTO iceberg_hash_routes VALUES + ('a', 1, 1, 10), + ('b', 100, 2, 20), + ('server42', 18446744073709551615, 3, 30), + ('中🙂', 0, 4, 40), + (NULL, 7, 5, 50); + +Affected Rows: 5 + +SELECT * FROM iceberg_hash_routes ORDER BY ts; + ++----------+----------------------+-------------------------+----+ +| host | tenant | ts | v | ++----------+----------------------+-------------------------+----+ +| a | 1 | 1970-01-01T00:00:00.001 | 10 | +| b | 100 | 1970-01-01T00:00:00.002 | 20 | +| server42 | 18446744073709551615 | 1970-01-01T00:00:00.003 | 30 | +| 中🙂 | 0 | 1970-01-01T00:00:00.004 | 40 | +| | 7 | 1970-01-01T00:00:00.005 | 50 | ++----------+----------------------+-------------------------+----+ + +SELECT * FROM iceberg_hash_routes WHERE host = 'server42'; + ++----------+----------------------+-------------------------+----+ +| host | tenant | ts | v | ++----------+----------------------+-------------------------+----+ +| server42 | 18446744073709551615 | 1970-01-01T00:00:00.003 | 30 | ++----------+----------------------+-------------------------+----+ + +DELETE FROM iceberg_hash_routes WHERE host = 'b'; + +Affected Rows: 1 + +SELECT * FROM iceberg_hash_routes ORDER BY ts; + ++----------+----------------------+-------------------------+----+ +| host | tenant | ts | v | ++----------+----------------------+-------------------------+----+ +| a | 1 | 1970-01-01T00:00:00.001 | 10 | +| server42 | 18446744073709551615 | 1970-01-01T00:00:00.003 | 30 | +| 中🙂 | 0 | 1970-01-01T00:00:00.004 | 40 | +| | 7 | 1970-01-01T00:00:00.005 | 50 | ++----------+----------------------+-------------------------+----+ diff --git a/tests/compatibility/cases/partition_iceberg_hash/verify.sql b/tests/compatibility/cases/partition_iceberg_hash/verify.sql new file mode 100644 index 00000000000..981022f574a --- /dev/null +++ b/tests/compatibility/cases/partition_iceberg_hash/verify.sql @@ -0,0 +1,19 @@ +SHOW CREATE TABLE iceberg_hash_routes; + +SELECT * FROM iceberg_hash_routes ORDER BY ts; + +-- Overwrites must route to the regions holding the flushed rows. +INSERT INTO iceberg_hash_routes VALUES + ('a', 1, 1, 10), + ('b', 100, 2, 20), + ('server42', 18446744073709551615, 3, 30), + ('中🙂', 0, 4, 40), + (NULL, 7, 5, 50); + +SELECT * FROM iceberg_hash_routes ORDER BY ts; + +SELECT * FROM iceberg_hash_routes WHERE host = 'server42'; + +DELETE FROM iceberg_hash_routes WHERE host = 'b'; + +SELECT * FROM iceberg_hash_routes ORDER BY ts; From 17ed988f49804af40dc79ff81f6c751b18be1dd1 Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Tue, 27 Oct 2026 10:20:45 +0000 Subject: [PATCH 107/239] fix(mito): preserve SST lifetime across index replacement (#9950) * fix(mito): preserve SST lifetime across index replacement Share retirement and final purge across immutable index metadata views. Keep version-specific GC references until readers release them and match Parquet liveness by file identity. Fixes #9944 Signed-off-by: Lei, HUANG * refactor(mito): list-based SST purge and single FilePurger::remove_file The final purge of a retired SST no longer depends on FileMeta: index files of every version are discovered by listing ({region_dir}/index/{file_id}. prefix), so even when the last live view is a stale one, the purge still removes all newer index files. - SstLifetime is slimmed down to { deleted, file_purger }; the purge now runs in the last FileHandleInner::drop (Arc::strong_count == 1, which is race-free here since only inners of the same file hold the lifetime). - FilePurger is merged back into a single remove_file(file_meta, is_delete) method; is_delete means the file was retired AND this is the last live view. - LocalFilePurger deletes the SST and all its index files via the new delete_file_and_all_indexes helper (list before delete, so a listing failure leaves the file untouched). - Remove the dead index_outdated flag and its accessors, which had no callers in the entire repo. Signed-off-by: Lei, HUANG * fix(mito): elect final SST view release atomically Track metadata views with an atomic counter instead of inspecting the lifetime Arc strong count. Concurrent destructors now elect exactly one final releaser while preserving the single remove_file interface. Add a deterministic concurrent-drop regression covering retired and live files. Signed-off-by: Lei, HUANG * refactor(mito): restore Arc-owned SST cleanup and clarify reference release Restore the lifetime design from 835227dcbc, using the final Arc drop and latest metadata instead of manual view counting and per-SST index listing. Rename FilePurger::remove_file to release_file_ref while retaining purge_file for final cleanup. Preserve and adapt the concurrent view-release regression test. Signed-off-by: Lei, HUANG * refactor(mito): track SST references by file lifetime Register one reference guard per logical SST lifetime and release it on final drop, including normal region close. Remove per-view release_file_ref and index_outdated handling. Protect all index generations while their SST is referenced. Preserve the existing GC snapshot format by expanding the recorded index generation upper bound for older version-aware workers. Update coverage for independent references, normal close, cross-region ownership, concurrent final cleanup and both GC listing modes. Signed-off-by: Lei, HUANG * refactor(mito): release SST references directly on lifetime drop Remove the redundant FileReferenceGuard and let SstLifetime register and unregister its file reference directly. Preserve unconditional reference release on normal close and retirement-only purge. Expose the reference manager through FilePurger and adapt reference manager tests without changing GC or snapshot behavior. Signed-off-by: Lei, HUANG * test(mito): verify SST cleanup with real storage and isolate GC policy cases Replace the recording-only concurrent purge test with filesystem assertions using LocalFilePurger and shared fixtures. Keep intermediate-view retention and retired versus live file cleanup coverage without claiming exhaustive thread interleavings. Separate GC reference-policy cases from reader-release ordering and reuse their setup. Production code remains unchanged. Signed-off-by: Lei, HUANG * style(mito): format SST lifetime changes Signed-off-by: Lei, HUANG * fix(mito): reclaim obsolete indexes by reader version Track exact index-version references per metadata view and restore version-specific Puffin liveness in both GC paths. Keep file-level Parquet protection and shared SST lifetime cleanup, without expanding snapshots into unreferenced index generations. Include old readers of imported files even when the source region is not open locally. Add real-storage GC coverage for reader release and lingering while retaining Parquet and the current index. Validated workspace checks with default and all features, all-target all-feature Clippy with warnings denied, and all 1853 mito2 tests. Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- src/cmd/src/datanode/objbench.rs | 9 +- src/mito2/src/engine/file_ref_test.rs | 201 +++++++++++++++ src/mito2/src/gc.rs | 43 +--- src/mito2/src/gc/worker_test.rs | 301 +++++++++++++++++++++- src/mito2/src/region.rs | 6 +- src/mito2/src/sst/file.rs | 98 ++++--- src/mito2/src/sst/file_purger.rs | 354 ++++++++++++++++---------- src/mito2/src/sst/file_ref.rs | 326 +++++++----------------- src/mito2/src/sst/version.rs | 38 ++- 9 files changed, 927 insertions(+), 449 deletions(-) diff --git a/src/cmd/src/datanode/objbench.rs b/src/cmd/src/datanode/objbench.rs index 03696c96fc0..248acc09bf6 100644 --- a/src/cmd/src/datanode/objbench.rs +++ b/src/cmd/src/datanode/objbench.rs @@ -27,7 +27,7 @@ use mito2::config::{FulltextIndexConfig, MitoConfig, Mode}; use mito2::read::FlatSource; use mito2::sst::FormatType; use mito2::sst::file::{FileHandle, FileMeta}; -use mito2::sst::file_purger::{FilePurger, FilePurgerRef}; +use mito2::sst::file_purger::{FilePurgerRef, NoopFilePurger}; use mito2::sst::index::intermediate::IntermediateManager; use mito2::sst::index::puffin_manager::PuffinManagerFactory; use mito2::sst::parquet::WriteOptions; @@ -508,12 +508,7 @@ async fn build_cache_manager( } fn new_noop_file_purger() -> FilePurgerRef { - #[derive(Debug)] - struct Noop; - impl FilePurger for Noop { - fn remove_file(&self, _file_meta: FileMeta, _is_delete: bool, _index_outdated: bool) {} - } - Arc::new(Noop) + Arc::new(NoopFilePurger) } async fn load_parquet_metadata( diff --git a/src/mito2/src/engine/file_ref_test.rs b/src/mito2/src/engine/file_ref_test.rs index 96f5f6661f4..3979c14553c 100644 --- a/src/mito2/src/engine/file_ref_test.rs +++ b/src/mito2/src/engine/file_ref_test.rs @@ -53,3 +53,204 @@ async fn get_snapshot_of_file_refs_skips_missing_file_handle_region() { assert!(manifest.manifest_version.is_empty()); assert!(manifest.cross_region_refs.is_empty()); } + +#[tokio::test] +async fn get_snapshot_preserves_related_region_old_index_readers() { + use store_api::region_engine::{RegionEngine, RegionRole}; + use store_api::region_request::RegionRequest; + use store_api::storage::{FileId, FileRef}; + + use crate::manifest::action::{RegionEdit, RegionMetaAction, RegionMetaActionList}; + use crate::region::RegionLeaderState; + use crate::sst::file::{FileHandle, FileMeta, IndexType}; + use crate::test_util::CreateRequestBuilder; + + let mut env = TestEnv::new().await; + let engine = env + .create_engine(MitoConfig { + gc: crate::gc::GcConfig { + enable: true, + ..Default::default() + }, + ..Default::default() + }) + .await; + let source_region = RegionId::new(1, 1); + let target_region = RegionId::new(1, 2); + engine + .handle_request( + target_region, + RegionRequest::Create(CreateRequestBuilder::new().build()), + ) + .await + .unwrap(); + engine + .set_region_role(target_region, RegionRole::Leader) + .unwrap(); + let meta = FileMeta { + region_id: source_region, + file_id: FileId::random(), + index_version: 2, + available_indexes: smallvec::smallvec![IndexType::InvertedIndex], + ..Default::default() + }; + let target = engine.get_region(target_region).unwrap(); + let old_reader = FileHandle::new( + FileMeta { + index_version: 0, + ..meta.clone() + }, + target.file_purger.clone(), + ); + target + .manifest_ctx + .update_manifest( + RegionLeaderState::Writable, + RegionMetaActionList::with_action(RegionMetaAction::Edit(RegionEdit { + files_to_add: vec![meta.clone()], + files_to_remove: vec![], + timestamp_ms: None, + compaction_time_window: None, + flushed_entry_id: None, + flushed_sequence: None, + committed_sequence: None, + })), + false, + ) + .await + .unwrap(); + + // The source is not open here, but a query on the destination still uses its v0 index. + let related_regions = HashMap::from([(source_region, HashSet::from([target_region]))]); + let snapshot = engine + .get_snapshot_of_file_refs([], related_regions.clone()) + .await + .unwrap(); + assert_eq!( + snapshot.file_refs[&source_region], + [0, 2] + .into_iter() + .map(|version| FileRef::new(source_region, meta.file_id, Some(version))) + .collect() + ); + assert_eq!( + snapshot.cross_region_refs[&source_region], + HashSet::from([target_region]) + ); + drop(old_reader); + let snapshot = engine + .get_snapshot_of_file_refs([], related_regions) + .await + .unwrap(); + assert_eq!( + snapshot.file_refs[&source_region], + HashSet::from([FileRef::new(source_region, meta.file_id, Some(2))]) + ); +} + +#[rstest::rstest] +#[tokio::test] +async fn get_snapshot_preserves_old_index_readers_until_release( + #[values(false, true)] old_has_index: bool, +) { + use store_api::region_engine::{RegionEngine, RegionRole}; + use store_api::region_request::RegionRequest; + use store_api::storage::{FileId, FileRef}; + + use crate::manifest::action::{RegionEdit, RegionMetaAction, RegionMetaActionList}; + use crate::region::RegionLeaderState; + use crate::sst::file::{FileMeta, IndexType}; + use crate::sst::version::SstVersion; + use crate::test_util::CreateRequestBuilder; + + let mut env = TestEnv::new().await; + let engine = env + .create_engine(MitoConfig { + gc: crate::gc::GcConfig { + enable: true, + ..Default::default() + }, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + engine + .handle_request( + region_id, + RegionRequest::Create(CreateRequestBuilder::new().build()), + ) + .await + .unwrap(); + engine + .set_region_role(region_id, RegionRole::Leader) + .unwrap(); + let region = engine.get_region(region_id).unwrap(); + let meta = FileMeta { + region_id, + file_id: FileId::random(), + available_indexes: if old_has_index { + smallvec::smallvec![IndexType::InvertedIndex] + } else { + Default::default() + }, + ..Default::default() + }; + let mut version = SstVersion::new(crate::test_util::memtable_util::metadata_for_test()); + version.add_files(region.file_purger.clone(), [meta.clone()].into_iter()); + let old_reader = version.levels()[0].files[&meta.file_id].clone(); + let current_meta = FileMeta { + index_version: 1, + available_indexes: smallvec::smallvec![IndexType::InvertedIndex], + ..meta.clone() + }; + version.add_files( + region.file_purger.clone(), + [current_meta.clone()].into_iter(), + ); + region + .manifest_ctx + .update_manifest( + RegionLeaderState::Writable, + RegionMetaActionList::with_action(RegionMetaAction::Edit(RegionEdit { + files_to_add: vec![current_meta], + files_to_remove: vec![], + timestamp_ms: None, + compaction_time_window: None, + flushed_entry_id: None, + flushed_sequence: None, + committed_sequence: None, + })), + false, + ) + .await + .unwrap(); + let snapshot = engine + .get_snapshot_of_file_refs([region_id], HashMap::new()) + .await + .unwrap(); + // Only v1 is covered by the manifest; keep the older view's exact reference. + let expected = HashSet::from([FileRef::new( + region_id, + meta.file_id, + old_has_index.then_some(0), + )]); + assert_eq!(snapshot.file_refs[®ion_id], expected); + // Closing the current view without retirement must keep the old reader protected. + drop(version); + let snapshot = engine + .get_snapshot_of_file_refs([region_id], HashMap::new()) + .await + .unwrap(); + assert_eq!(snapshot.file_refs[®ion_id], expected); + drop(old_reader); + let snapshot = engine + .get_snapshot_of_file_refs([region_id], HashMap::new()) + .await + .unwrap(); + assert!( + snapshot + .file_refs + .get(®ion_id) + .is_none_or(HashSet::is_empty) + ); +} diff --git a/src/mito2/src/gc.rs b/src/mito2/src/gc.rs index 541bb4872b4..8d917fbcffb 100644 --- a/src/mito2/src/gc.rs +++ b/src/mito2/src/gc.rs @@ -929,17 +929,8 @@ impl LocalGcWorker { unknown_file_may_linger_until: chrono::DateTime, ) -> Vec { let mut ready_for_delete = vec![]; - // all group by file id for easier checking - let in_tmp_ref: HashMap> = - in_tmp_ref - .iter() - .fold(HashMap::new(), |mut acc, (file, version)| { - let indices = acc.entry(*file).or_default(); - if let Some(version) = version { - indices.insert(*version); - } - acc - }); + // Every index-version reader needs the same Parquet file. + let referenced_ssts: HashSet<_> = in_tmp_ref.iter().map(|(file_id, _)| *file_id).collect(); let may_linger_files: HashMap> = may_linger_files .iter() @@ -976,7 +967,7 @@ impl LocalGcWorker { let should_delete = match file_type { FileType::Parquet => { let is_in_manifest = in_manifest.contains_key(&file_id); - let is_in_tmp_ref = in_tmp_ref.contains_key(&file_id); + let is_in_tmp_ref = referenced_ssts.contains(&file_id); let is_linger = may_linger_files.contains_key(&file_id); let is_eligible_for_delete = eligible_for_delete.contains_key(&file_id); @@ -991,15 +982,8 @@ impl LocalGcWorker { ) } FileType::Puffin(version) => { - // notice need to check both file id and version - let is_in_manifest = in_manifest - .get(&file_id) - .map(|opt_ver| *opt_ver == Some(version)) - .unwrap_or(false); - let is_in_tmp_ref = in_tmp_ref - .get(&file_id) - .map(|versions| versions.contains(&version)) - .unwrap_or(false); + let is_in_manifest = in_manifest.get(&file_id) == Some(&Some(version)); + let is_in_tmp_ref = in_tmp_ref.contains(&(file_id, Some(version))); let is_linger = may_linger_files .get(&file_id) .map(|files| files.contains(&&RemovedFile::Index(file_id, version))) @@ -1100,21 +1084,22 @@ impl LocalGcWorker { // When full_file_listing is false, skip expensive list operations and only delete // files that are tracked in recently_removed_files if !self.full_file_listing { + let referenced_ssts: HashSet<_> = + in_tmp_ref.iter().map(|(file_id, _)| *file_id).collect(); // Only delete files that: // 1. Are in recently_removed_files (tracked in manifest) // 2. Are not in use(in manifest or tmp ref) // 3. Have passed the lingering time let files_to_delete: Vec = eligible_for_removal .iter() - .filter(|file_id| { - let in_use = match file_id { - RemovedFile::File(file_id, index_version) => { - in_manifest.get(file_id) == Some(index_version) - || in_tmp_ref.contains(&(*file_id, *index_version)) + .filter(|file| { + let in_use = match file { + RemovedFile::File(file_id, _) => { + in_manifest.contains_key(file_id) || referenced_ssts.contains(file_id) } - RemovedFile::Index(file_id, index_version) => { - in_manifest.get(file_id) == Some(&Some(*index_version)) - || in_tmp_ref.contains(&(*file_id, Some(*index_version))) + RemovedFile::Index(file_id, version) => { + in_manifest.get(file_id) == Some(&Some(*version)) + || in_tmp_ref.contains(&(*file_id, Some(*version))) } }; !in_use diff --git a/src/mito2/src/gc/worker_test.rs b/src/mito2/src/gc/worker_test.rs index 3d389e7e1d7..f439e01a951 100644 --- a/src/mito2/src/gc/worker_test.rs +++ b/src/mito2/src/gc/worker_test.rs @@ -21,12 +21,13 @@ use async_trait::async_trait; use bytes::Bytes; use common_base::Plugins; use common_telemetry::init_default_ut_logging; +use common_time::Timestamp; use futures::TryStreamExt; use object_store::{Entry, ObjectStore, services}; use store_api::metadata::RegionMetadataRef; -use store_api::region_engine::RegionEngine as _; +use store_api::region_engine::{RegionEngine as _, RegionRole}; use store_api::region_request::{RegionCompactRequest, RegionRequest}; -use store_api::storage::{FileRef, FileRefsManifest, RegionId}; +use store_api::storage::{FileId, FileRef, FileRefsManifest, IndexVersion, RegionId}; use crate::access_layer::AccessLayerRef; use crate::config::MitoConfig; @@ -35,9 +36,11 @@ use crate::engine::compaction_test::{delete_and_flush, put_and_flush}; use crate::engine::region_hook::{RegionGcInfo, RegionHook, RegionHookRef}; use crate::error::UnexpectedSnafu; use crate::gc::{GcConfig, LocalGcWorker, should_delete_file}; -use crate::manifest::action::RemovedFile; +use crate::manifest::action::{RegionEdit, RegionMetaAction, RegionMetaActionList, RemovedFile}; use crate::metrics::GC_SKIPPED_UNPARSABLE_FILES; -use crate::region::MitoRegionRef; +use crate::region::{MitoRegionRef, RegionLeaderState}; +use crate::sst::file::{FileHandle, FileMeta, IndexType, RegionIndexId}; +use crate::sst::location; use crate::test_util::{ CreateRequestBuilder, TestEnv, build_rows, flush_region, put_rows, rows_schema, }; @@ -74,6 +77,296 @@ async fn create_gc_worker( .unwrap() } +struct IndexGcTestEnv { + engine: MitoEngine, + region: MitoRegionRef, + worker: LocalGcWorker, + meta: FileMeta, + entries: Vec, + removed_files: BTreeMap>, + _env: TestEnv, +} + +impl IndexGcTestEnv { + async fn new(full_file_listing: bool) -> Self { + let mut env = TestEnv::new().await; + let engine = env + .create_engine(MitoConfig { + gc: GcConfig { + enable: true, + lingering_time: Some(std::time::Duration::from_secs(60 * 60)), + ..Default::default() + }, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + engine + .handle_request( + region_id, + RegionRequest::Create(CreateRequestBuilder::new().build()), + ) + .await + .unwrap(); + let region = engine.get_region(region_id).unwrap(); + let meta = FileMeta { + region_id, + file_id: FileId::random(), + available_indexes: smallvec::smallvec![IndexType::InvertedIndex], + ..Default::default() + }; + let layer = ®ion.access_layer; + let paths = [ + location::sst_file_path(layer.table_dir(), meta.file_id(), layer.path_type()), + location::index_file_path(layer.table_dir(), meta.index_id(), layer.path_type()), + location::index_file_path( + layer.table_dir(), + RegionIndexId::new(meta.file_id(), 1), + layer.path_type(), + ), + ]; + let mut entries = Vec::new(); + for path in paths { + entries.push(write_and_list_entry(layer.object_store(), &path).await); + } + let removed_files = BTreeMap::from([( + Timestamp::new_millisecond(0), + HashSet::from([ + RemovedFile::File(meta.file_id, Some(1)), + RemovedFile::Index(meta.file_id, 0), + RemovedFile::Index(meta.file_id, 1), + ]), + )]); + let worker = create_gc_worker( + &engine, + BTreeMap::from([(region_id, Some(region.clone()))]), + &FileRefsManifest::default(), + full_file_listing, + ) + .await; + Self { + engine, + region, + worker, + meta, + entries, + removed_files, + _env: env, + } + } +} + +#[rstest::rstest] +#[tokio::test] +async fn test_gc_keeps_sst_across_index_versions( + #[values(false, true)] full_file_listing: bool, + #[values(false, true)] old_reader_last: bool, +) { + use crate::sst::version::SstVersion; + + let env = IndexGcTestEnv::new(full_file_listing).await; + let meta = &env.meta; + let region_id = meta.region_id; + let mut version = SstVersion::new(crate::test_util::memtable_util::metadata_for_test()); + version.add_files(env.region.file_purger.clone(), [meta.clone()].into_iter()); + let old = version.levels()[0].files[&meta.file_id].clone(); + let updated = FileMeta { + index_version: 1, + ..meta.clone() + }; + version.add_files(env.region.file_purger.clone(), [updated].into_iter()); + let new = version.levels()[0].files[&meta.file_id].clone(); + version.mark_all_deleted(); + drop(version); + + let (first, last) = if old_reader_last { + (new, old) + } else { + (old, new) + }; + let released_version = first.index_id().version; + let live_version = last.index_id().version; + let mut handles = [Some(first), Some(last)]; + for released in 0..=2 { + if released > 0 { + drop(handles[released - 1].take()); + } + let refs = env + .engine + .get_snapshot_of_file_refs([region_id], HashMap::new()) + .await + .unwrap(); + let refs = refs.file_refs.get(®ion_id).cloned().unwrap_or_default(); + let expected_versions = match released { + 0 => vec![0, 1], + 1 => vec![live_version], + _ => vec![], + }; + assert_eq!( + refs, + expected_versions + .into_iter() + .map(|v| FileRef::new(region_id, meta.file_id, Some(v))) + .collect() + ); + let tmp_refs = refs + .into_iter() + .map(|f| (f.file_id, f.index_version)) + .collect(); + let deleted: HashSet<_> = env + .worker + .list_to_be_deleted_files( + region_id, + false, + &HashMap::new(), + &tmp_refs, + env.removed_files.clone(), + env.entries.clone(), + ) + .await + .unwrap() + .into_iter() + .collect(); + let expected = match released { + 0 => HashSet::new(), + 1 => HashSet::from([RemovedFile::Index(meta.file_id, released_version)]), + _ => HashSet::from([ + RemovedFile::File(meta.file_id, if full_file_listing { None } else { Some(1) }), + RemovedFile::Index(meta.file_id, 0), + RemovedFile::Index(meta.file_id, 1), + ]), + }; + assert_eq!(deleted, expected); + } +} + +#[rstest::rstest] +#[case::manifest(true, Some(3))] +#[case::no_index_reader(false, None)] +#[case::old_index_reader(false, Some(0))] +#[case::other_index_reader(false, Some(3))] +#[tokio::test] +async fn test_gc_reference_protects_parquet_and_only_its_index( + #[values(false, true)] full_file_listing: bool, + #[case] in_manifest: bool, + #[case] index_version: Option, +) { + let env = IndexGcTestEnv::new(full_file_listing).await; + let file_id = env.meta.file_id; + let (manifest_refs, tmp_refs) = if in_manifest { + (HashMap::from([(file_id, index_version)]), HashSet::new()) + } else { + (HashMap::new(), HashSet::from([(file_id, index_version)])) + }; + let deleted = env + .worker + .list_to_be_deleted_files( + env.meta.region_id, + false, + &manifest_refs, + &tmp_refs, + env.removed_files.clone(), + env.entries.clone(), + ) + .await + .unwrap(); + let expected: HashSet<_> = [0, 1] + .into_iter() + .filter(|version| Some(*version) != index_version) + .map(|version| RemovedFile::Index(file_id, version)) + .collect(); + assert_eq!(deleted.into_iter().collect::>(), expected); +} + +#[rstest::rstest] +#[tokio::test] +async fn test_gc_reclaims_obsolete_index_after_readers_and_linger_release( + #[values(false, true)] full_file_listing: bool, + #[values(false, true)] linger_expired: bool, +) { + let env = IndexGcTestEnv::new(full_file_listing).await; + let region_id = env.meta.region_id; + env.engine + .set_region_role(region_id, RegionRole::Leader) + .unwrap(); + let old_reader = FileHandle::new(env.meta.clone(), env.region.file_purger.clone()); + let current_meta = FileMeta { + index_version: 1, + ..env.meta.clone() + }; + let current_reader = old_reader.with_updated_index(current_meta.clone()); + let removed_at = if linger_expired { + 0 + } else { + chrono::Utc::now().timestamp_millis() + }; + for meta in [env.meta.clone(), current_meta] { + env.region + .manifest_ctx + .update_manifest( + RegionLeaderState::Writable, + RegionMetaActionList::with_action(RegionMetaAction::Edit(RegionEdit { + files_to_add: vec![meta], + files_to_remove: vec![], + timestamp_ms: Some(removed_at), + compaction_time_window: None, + flushed_entry_id: None, + flushed_sequence: None, + committed_sequence: None, + })), + false, + ) + .await + .unwrap(); + } + + let layer = &env.region.access_layer; + let parquet = location::sst_file_path(layer.table_dir(), env.meta.file_id(), layer.path_type()); + let old_index = + location::index_file_path(layer.table_dir(), env.meta.index_id(), layer.path_type()); + let current_index = location::index_file_path( + layer.table_dir(), + current_reader.index_id(), + layer.path_type(), + ); + let mut old_reader = Some(old_reader); + for released in [false, true] { + if released { + drop(old_reader.take()); + } + let refs = env + .engine + .get_snapshot_of_file_refs([region_id], HashMap::new()) + .await + .unwrap(); + let worker = create_gc_worker( + &env.engine, + BTreeMap::from([(region_id, Some(env.region.clone()))]), + &refs, + full_file_listing, + ) + .await; + let report = worker.run().await.unwrap(); + let reclaimed = released && linger_expired; + assert!(report.deleted_files[®ion_id].is_empty()); + assert_eq!( + report.deleted_indexes[®ion_id], + if reclaimed { + vec![(env.meta.file_id, 0)] + } else { + vec![] + } + ); + assert!(layer.object_store().exists(&parquet).await.unwrap()); + assert!(layer.object_store().exists(¤t_index).await.unwrap()); + assert_eq!( + layer.object_store().exists(&old_index).await.unwrap(), + !reclaimed + ); + } + drop(current_reader); +} + #[tokio::test] async fn test_gc_worker_ignores_dropping_marker_file() { init_default_ut_logging(); diff --git a/src/mito2/src/region.rs b/src/mito2/src/region.rs index e6e3de2df1e..ef5b24f9de5 100644 --- a/src/mito2/src/region.rs +++ b/src/mito2/src/region.rs @@ -2363,10 +2363,8 @@ mod tests { struct CountingPurger(AtomicUsize); impl FilePurger for CountingPurger { - fn remove_file(&self, _: FileMeta, is_delete: bool, _: bool) { - if is_delete { - self.0.fetch_add(1, Ordering::SeqCst); - } + fn purge_file(&self, _: FileMeta) { + self.0.fetch_add(1, Ordering::SeqCst); } } diff --git a/src/mito2/src/sst/file.rs b/src/mito2/src/sst/file.rs index 28b6218cd7b..1d4caaf9127 100644 --- a/src/mito2/src/sst/file.rs +++ b/src/mito2/src/sst/file.rs @@ -39,6 +39,7 @@ use crate::cache::CacheManagerRef; use crate::cache::file_cache::{FileType, IndexKey}; use crate::region::options::IndexOptions; use crate::sst::file_purger::FilePurgerRef; +use crate::sst::file_ref::FileReferenceManagerRef; use crate::sst::index::fulltext_index::FULLTEXT_BLOOM_VERSION; use crate::sst::location; use crate::sst::parquet::SstInfo; @@ -551,13 +552,13 @@ impl fmt::Debug for FileHandle { f.debug_struct("FileHandle") .field("meta", self.meta_ref()) .field("compacting", &self.compacting()) - .field("deleted", &self.inner.deleted.load(Ordering::Relaxed)) + .field("deleted", &self.is_deleted()) .finish() } } impl FileHandle { - /// Creates a handle sharing the file's physical state and original statistics. + /// Creates a new logical SST reference. Index replacements must reuse its lifetime. pub fn new(meta: FileMeta, file_purger: FilePurgerRef) -> FileHandle { let pk_range = meta.primary_key_range(); FileHandle { @@ -576,6 +577,23 @@ impl FileHandle { } } + /// Replaces an index metadata view without splitting the underlying SST lifetime. + pub(crate) fn with_updated_index(&self, meta: FileMeta) -> Self { + debug_assert_eq!(self.file_id(), meta.file_id()); + let primary_key_range = meta + .primary_key_range() + .or_else(|| self.raw_primary_key_range()); + let lifetime = self.inner.lifetime.clone(); + *lifetime.meta.lock().unwrap() = meta.clone(); + Self { + inner: Arc::new(FileHandleInner::with_lifetime( + meta, + lifetime, + primary_key_range, + )), + } + } + /// Returns the region id of the file. pub fn region_id(&self) -> RegionId { self.inner.meta.region_id @@ -615,7 +633,7 @@ impl FileHandle { /// Mark the file as deleted and will delete it on drop asynchronously pub fn mark_deleted(&self) { - self.inner.deleted.store(true, Ordering::Relaxed); + self.inner.lifetime.deleted.store(true, Ordering::Relaxed); } pub fn compacting(&self) -> bool { @@ -634,23 +652,13 @@ impl FileHandle { .is_ok() } - pub fn index_outdated(&self) -> bool { - self.inner.index_outdated.load(Ordering::Relaxed) - } - - pub fn set_index_outdated(&self, index_outdated: bool) { - self.inner - .index_outdated - .store(index_outdated, Ordering::Relaxed); - } - /// Returns a reference to the [FileMeta]. pub fn meta_ref(&self) -> &FileMeta { &self.inner.meta } pub fn file_purger(&self) -> FilePurgerRef { - self.inner.file_purger.clone() + self.inner.lifetime.file_purger.clone() } pub fn size(&self) -> u64 { @@ -670,7 +678,7 @@ impl FileHandle { } pub fn is_deleted(&self) -> bool { - self.inner.deleted.load(Ordering::Relaxed) + self.inner.lifetime.deleted.load(Ordering::Relaxed) } /// Returns bounds aligned to the caller's pinned schema, before any comparison @@ -776,39 +784,69 @@ impl PrimaryKeyRange { struct FileHandleInner { meta: FileMeta, compacting: AtomicBool, - deleted: AtomicBool, - index_outdated: AtomicBool, primary_key_range: RwLock, - file_purger: FilePurgerRef, + lifetime: Arc, } impl Drop for FileHandleInner { fn drop(&mut self) { - self.file_purger.remove_file( - self.meta.clone(), - self.deleted.load(Ordering::Acquire), - self.index_outdated.load(Ordering::Acquire), - ); + if let Some(manager) = &self.lifetime.file_ref_manager { + manager.remove_file(&self.meta); + } } } impl FileHandleInner { - /// There should only be one `FileHandleInner` for each file on a datanode fn new( meta: FileMeta, file_purger: FilePurgerRef, primary_key_range: Option<(Bytes, Bytes)>, - ) -> FileHandleInner { - file_purger.new_file(&meta); - FileHandleInner { + ) -> Self { + let file_ref_manager = file_purger.file_ref_manager(); + let lifetime = Arc::new(SstLifetime { + meta: Mutex::new(meta.clone()), + deleted: AtomicBool::new(false), + file_purger, + file_ref_manager, + }); + Self::with_lifetime(meta, lifetime, primary_key_range) + } + + fn with_lifetime( + meta: FileMeta, + lifetime: Arc, + primary_key_range: Option<(Bytes, Bytes)>, + ) -> Self { + if let Some(manager) = &lifetime.file_ref_manager { + manager.register_file(&meta); + } + Self { meta, compacting: AtomicBool::new(false), - deleted: AtomicBool::new(false), - index_outdated: AtomicBool::new(false), primary_key_range: RwLock::new( primary_key_range.map_or(PrimaryKeyRange::Missing, PrimaryKeyRange::Raw), ), - file_purger, + lifetime, + } + } +} + +/// Physical deletion must wait for readers of every index metadata version. +/// Independent logical references (for example, copied regions) do not share this guard. +struct SstLifetime { + /// Latest metadata ensures final purge includes indexes newer than the last reader's view. + meta: Mutex, + deleted: AtomicBool, + file_purger: FilePurgerRef, + /// Tracks each metadata view's readers independently for object-store index GC. + file_ref_manager: Option, +} + +impl Drop for SstLifetime { + fn drop(&mut self) { + if self.deleted.load(Ordering::Acquire) { + self.file_purger + .purge_file(self.meta.get_mut().unwrap().clone()); } } } diff --git a/src/mito2/src/sst/file_purger.rs b/src/mito2/src/sst/file_purger.rs index 2660a19d4af..727c6c1cc9e 100644 --- a/src/mito2/src/sst/file_purger.rs +++ b/src/mito2/src/sst/file_purger.rs @@ -23,23 +23,20 @@ use crate::access_layer::AccessLayerRef; use crate::cache::CacheManagerRef; use crate::error::Result; use crate::schedule::scheduler::SchedulerRef; -use crate::sst::file::{FileMeta, delete_files, delete_index}; +use crate::sst::file::{FileMeta, delete_files}; use crate::sst::file_ref::FileReferenceManagerRef; use crate::sst::range_index::RangeIndexDeleter; /// A worker to delete files in background. pub trait FilePurger: Send + Sync + fmt::Debug { - /// Send a request to remove the file. - /// If `is_delete` is true, the file will be deleted from the storage. - /// Otherwise, only the reference will be removed. - /// If `index_outdated` is true, the index file will be deleted regardless of `is_delete`. - fn remove_file(&self, file_meta: FileMeta, is_delete: bool, index_outdated: bool); + /// Purges a retired SST and its indexes after all metadata views release their readers. + /// Each metadata view has already released its reader reference. + fn purge_file(&self, file_meta: FileMeta); - /// Notify the purger of a new file created. - /// This is useful for object store based storage, where we need to track the file references - /// The default implementation is a no-op. - fn new_file(&self, _: &FileMeta) { - // noop + /// Returns the object-store reference manager used by SST metadata views. + /// Local and no-op purgers do not track GC references. + fn file_ref_manager(&self) -> Option { + None } } @@ -50,9 +47,7 @@ pub type FilePurgerRef = Arc; pub struct NoopFilePurger; impl FilePurger for NoopFilePurger { - fn remove_file(&self, _file_meta: FileMeta, _is_delete: bool, _index_outdated: bool) { - // noop - } + fn purge_file(&self, _file_meta: FileMeta) {} } /// Purger that purges file for current region. @@ -182,33 +177,16 @@ impl LocalFilePurger { error!(e; "Failed to schedule the file purge request"); } } - - fn delete_index(&self, file_meta: FileMeta) { - let sst_layer = self.sst_layer.clone(); - let cache_manager = self.cache_manager.clone(); - if let Err(e) = self.scheduler.schedule(Box::pin(async move { - let index_id = file_meta.index_id(); - if let Err(e) = delete_index(index_id, &sst_layer, &cache_manager).await { - error!(e; "Failed to delete index for file {:?} from storage", file_meta); - } - })) { - error!(e; "Failed to schedule the index purge request"); - } - } } impl FilePurger for LocalFilePurger { - fn remove_file(&self, file_meta: FileMeta, is_delete: bool, index_outdated: bool) { - if is_delete { - schedule_range_index_deletion( - &self.scheduler, - self.range_index_deleter.as_ref(), - file_meta.file_id, - ); - self.delete_file(file_meta); - } else if index_outdated { - self.delete_index(file_meta); - } + fn purge_file(&self, file_meta: FileMeta) { + schedule_range_index_deletion( + &self.scheduler, + self.range_index_deleter.as_ref(), + file_meta.file_id, + ); + self.delete_file(file_meta); } } @@ -246,31 +224,26 @@ fn schedule_range_index_deletion( } impl FilePurger for ObjectStoreFilePurger { - fn remove_file(&self, file_meta: FileMeta, is_delete: bool, _index_outdated: bool) { - // if not on local file system, instead inform the global file purger to remove the file reference. - // notice that no matter whether the file is deleted or not, we need to remove the reference - // because the file is no longer in use nonetheless. - // for same reason, we don't care about index_outdated here. - self.file_ref_manager.remove_file(&file_meta); - if is_delete { - schedule_range_index_deletion( - &self.scheduler, - self.range_index_deleter.as_ref(), - file_meta.file_id, - ); - } + fn purge_file(&self, file_meta: FileMeta) { + schedule_range_index_deletion( + &self.scheduler, + self.range_index_deleter.as_ref(), + file_meta.file_id, + ); } - fn new_file(&self, file_meta: &FileMeta) { - self.file_ref_manager.add_file(file_meta); + fn file_ref_manager(&self) -> Option { + Some(self.file_ref_manager.clone()) } } #[cfg(test)] mod tests { use std::num::NonZeroU64; + use std::sync::Barrier; + use std::thread; - use common_test_util::temp_dir::create_temp_dir; + use common_test_util::temp_dir::{TempDir, create_temp_dir}; use object_store::ObjectStore; use object_store::services::Fs; use smallvec::SmallVec; @@ -288,6 +261,133 @@ mod tests { use crate::sst::index::puffin_manager::PuffinManagerFactory; use crate::sst::location; + struct PurgeTestEnv { + layer: AccessLayerRef, + scheduler: Arc, + file_id: RegionFileId, + _dir: TempDir, + } + + impl PurgeTestEnv { + async fn new() -> Self { + common_telemetry::init_default_ut_logging(); + let dir = create_temp_dir("file-purge"); + let index_aux_path = dir.path().join("index_aux"); + let puffin_mgr = PuffinManagerFactory::new(&index_aux_path, 4096, None, None) + .await + .unwrap(); + let intm_mgr = IntermediateManager::init_fs(index_aux_path.to_str().unwrap()) + .await + .unwrap(); + let store = + ObjectStore::new(Fs::default().root(&dir.path().display().to_string())).unwrap(); + let env = Self { + layer: Arc::new(AccessLayer::new( + "table1", + PathType::Bare, + store, + puffin_mgr, + intm_mgr, + )), + scheduler: Arc::new(LocalScheduler::new(1)), + file_id: RegionFileId::new(RegionId::new(0, 0), FileId::random()), + _dir: dir, + }; + env.layer + .object_store() + .write(&env.sst_path(), vec![0; 4096]) + .await + .unwrap(); + env + } + + fn sst_path(&self) -> String { + location::sst_file_path(self.layer.table_dir(), self.file_id, self.layer.path_type()) + } + + fn index_path(&self, version: u64) -> String { + location::index_file_path( + self.layer.table_dir(), + RegionIndexId::new(self.file_id, version), + self.layer.path_type(), + ) + } + + async fn wait_for_purges(&self) { + // The single worker finishes all earlier purge jobs before this barrier. + let (tx, rx) = tokio::sync::oneshot::channel(); + self.scheduler + .schedule(Box::pin(async move { + tx.send(()).unwrap(); + })) + .unwrap(); + rx.await.unwrap(); + } + } + + #[rstest::rstest] + #[tokio::test] + async fn test_concurrent_index_view_release_keeps_live_files_and_purges_retired_files( + #[values(false, true)] retired: bool, + ) { + let env = PurgeTestEnv::new().await; + let store = env.layer.object_store(); + let purger = Arc::new(LocalFilePurger::new( + env.scheduler.clone(), + env.layer.clone(), + None, + )); + let meta = FileMeta { + region_id: env.file_id.region_id(), + file_id: env.file_id.file_id(), + ..Default::default() + }; + let old_reader = FileHandle::new(meta.clone(), purger); + let mut readers = vec![old_reader]; + for index_version in 1..=2 { + store + .write(&env.index_path(index_version), "index") + .await + .unwrap(); + readers.push(readers.last().unwrap().with_updated_index(FileMeta { + index_version, + available_indexes: smallvec::smallvec![IndexType::InvertedIndex], + ..meta.clone() + })); + } + if retired { + readers.last().unwrap().mark_deleted(); + } + // Dropping one view must not remove files still needed by the other views. + drop(readers.remove(1)); + env.wait_for_purges().await; + assert!(store.exists(&env.sst_path()).await.unwrap()); + for version in 1..=2 { + assert!(store.exists(&env.index_path(version)).await.unwrap()); + } + + // Regression coverage for final releases on different threads. The start barrier + // encourages overlap; it does not exhaustively control destructor interleavings. + let barrier = Barrier::new(readers.len()); + thread::scope(|scope| { + for reader in readers { + let barrier = &barrier; + scope.spawn(move || { + barrier.wait(); + drop(reader); + }); + } + }); + env.scheduler.stop(true).await.unwrap(); + assert_eq!(store.exists(&env.sst_path()).await.unwrap(), !retired); + for version in 1..=2 { + assert_eq!( + store.exists(&env.index_path(version)).await.unwrap(), + !retired + ); + } + } + #[rstest::rstest] #[case(false)] #[case(true)] @@ -295,36 +395,13 @@ mod tests { async fn test_range_index_purge_on_handle_release( #[case] gc_enabled: bool, #[values(false, true)] is_delete: bool, + #[values(false, true)] update_index: bool, ) { - common_telemetry::init_default_ut_logging(); - - let dir = create_temp_dir("file-purge"); - let dir_path = dir.path().display().to_string(); - let builder = Fs::default().root(&dir_path); - let sst_file_id = RegionFileId::new(RegionId::new(0, 0), FileId::random()); - let sst_dir = "table1"; - - let index_aux_path = dir.path().join("index_aux"); - let puffin_mgr = PuffinManagerFactory::new(&index_aux_path, 4096, None, None) - .await - .unwrap(); - let intm_mgr = IntermediateManager::init_fs(index_aux_path.to_str().unwrap()) - .await - .unwrap(); - - let object_store = ObjectStore::new(builder).unwrap(); - - let layer = Arc::new(AccessLayer::new( - sst_dir, - PathType::Bare, - object_store.clone(), - puffin_mgr, - intm_mgr, - )); - let path = location::sst_file_path(sst_dir, sst_file_id, layer.path_type()); - object_store.write(&path, vec![0; 4096]).await.unwrap(); - - let scheduler = Arc::new(LocalScheduler::new(3)); + let env = PurgeTestEnv::new().await; + let object_store = env.layer.object_store(); + let path = env.sst_path(); + let sst_file_id = env.file_id; + let scheduler = &env.scheduler; let index_store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); let owner = RegionId::new(9, 1); @@ -334,40 +411,48 @@ mod tests { gc_enabled, PathType::Bare, scheduler.clone(), - layer, + env.layer.clone(), None, Arc::new(crate::sst::file_ref::FileReferenceManager::new(None)), Some(RangeIndexDeleter::new(index_store.clone(), owner)), ); { - let handle = FileHandle::new( - FileMeta { - region_id: sst_file_id.region_id(), - file_id: sst_file_id.file_id(), - time_range: FileTimeRange::default(), - level: 0, - file_size: 4096, - max_row_group_uncompressed_size: 4096, - available_indexes: Default::default(), - indexes: Default::default(), - index_file_size: 0, - index_version: 0, - num_rows: 0, - num_row_groups: 0, - sequence: None, - partition_expr: None, - num_series: 0, - ..Default::default() - }, - file_purger, + let meta = FileMeta { + region_id: sst_file_id.region_id(), + file_id: sst_file_id.file_id(), + file_size: 4096, + max_row_group_uncompressed_size: 4096, + ..Default::default() + }; + let mut version = crate::sst::version::SstVersion::new( + crate::test_util::memtable_util::metadata_for_test(), ); - if is_delete { - handle.mark_deleted(); + version.add_files(file_purger, [meta].into_iter()); + let reader = version.levels()[0].files[&sst_file_id.file_id()].clone(); + if update_index { + for index_version in 1..=2 { + let mut meta = reader.meta_ref().clone(); + meta.index_version = index_version; + meta.available_indexes.push(IndexType::InvertedIndex); + let path = env.index_path(index_version); + object_store.write(&path, "index").await.unwrap(); + version.add_files(reader.file_purger(), [meta].into_iter()); + } } - let reader = handle.clone(); - drop(handle); + if is_delete { + version.mark_all_deleted(); + } + drop(version); + env.wait_for_purges().await; + assert!(object_store.exists(&path).await.unwrap()); assert!(index_store.exists(&index_path).await.unwrap()); + if update_index { + for index_version in 1..=2 { + let path = env.index_path(index_version); + assert!(object_store.exists(&path).await.unwrap()); + } + } drop(reader); } @@ -378,48 +463,35 @@ mod tests { gc_enabled || !is_delete ); assert_eq!(index_store.exists(&index_path).await.unwrap(), !is_delete); + if update_index { + for index_version in 1..=2 { + let path = env.index_path(index_version); + assert_eq!( + object_store.exists(&path).await.unwrap(), + gc_enabled || !is_delete + ); + } + } } #[tokio::test] async fn test_file_purge_with_index() { - common_telemetry::init_default_ut_logging(); - - let dir = create_temp_dir("file-purge"); - let dir_path = dir.path().display().to_string(); - let builder = Fs::default().root(&dir_path); - let sst_file_id = RegionFileId::new(RegionId::new(0, 0), FileId::random()); - let index_file_id = RegionIndexId::new(sst_file_id, 0); - let sst_dir = "table1"; - - let index_aux_path = dir.path().join("index_aux"); - let puffin_mgr = PuffinManagerFactory::new(&index_aux_path, 4096, None, None) - .await - .unwrap(); - let intm_mgr = IntermediateManager::init_fs(index_aux_path.to_str().unwrap()) - .await - .unwrap(); - - let object_store = ObjectStore::new(builder).unwrap(); - - let layer = Arc::new(AccessLayer::new( - sst_dir, - PathType::Bare, - object_store.clone(), - puffin_mgr, - intm_mgr, - )); - let path = location::sst_file_path(sst_dir, sst_file_id, layer.path_type()); - object_store.write(&path, vec![0; 4096]).await.unwrap(); - - let index_path = location::index_file_path(sst_dir, index_file_id, layer.path_type()); + let env = PurgeTestEnv::new().await; + let object_store = env.layer.object_store(); + let sst_file_id = env.file_id; + let path = env.sst_path(); + let index_path = env.index_path(0); object_store .write(&index_path, vec![0; 4096]) .await .unwrap(); - let scheduler = Arc::new(LocalScheduler::new(3)); - - let file_purger = Arc::new(LocalFilePurger::new(scheduler.clone(), layer, None)); + let scheduler = &env.scheduler; + let file_purger = Arc::new(LocalFilePurger::new( + scheduler.clone(), + env.layer.clone(), + None, + )); { let handle = FileHandle::new( diff --git a/src/mito2/src/sst/file_ref.rs b/src/mito2/src/sst/file_ref.rs index 4dee84ca1fa..68b0facaab4 100644 --- a/src/mito2/src/sst/file_ref.rs +++ b/src/mito2/src/sst/file_ref.rs @@ -24,23 +24,18 @@ use crate::metrics::GC_REF_FILE_CNT; use crate::region::MitoRegionRef; use crate::sst::file::FileMeta; -/// File references for a region. -/// It contains all files referenced by the region. +/// File references grouped by their physical source region. #[derive(Debug, Clone, Default)] -pub struct RegionFileRefs { - /// (FileRef, Ref Count) meaning how many FileHandleInner is opened for this file. - pub files: HashMap, +struct RegionFileRefs { + /// Counts independent metadata views of each physical file and index version. + files: HashMap, } -/// Manages all file references in one datanode. -/// It keeps track of which files are referenced and group by table ids. -/// This is useful for ensuring that files are not deleted while they are still in use by any -/// query. +/// Tracks reader references to SST metadata views in one datanode for object-store GC. #[derive(Debug)] pub struct FileReferenceManager { /// Datanode id. used to determine tmp ref file name. node_id: Option, - /// TODO(discord9): use no hash hasher since table id is sequential. files_per_region: DashMap, /// Whether global GC is enabled. /// This is only meaningful when using object store (not local filesystem). @@ -68,58 +63,49 @@ impl FileReferenceManager { } /// Returns whether global GC is enabled. - /// - /// This is useful for determining the file deletion strategy: - /// - If GC is enabled (and using object store), files are tracked but not immediately deleted - /// - If GC is disabled (or using local filesystem), files are deleted immediately pub fn is_gc_enabled(&self) -> bool { self.gc_enabled } fn ref_file_set(&self, region_id: RegionId) -> Option> { - // region id not found. - let file_refs = self.files_per_region.get(®ion_id)?.clone(); - - if file_refs.files.is_empty() { - // still return an empty manifest to indicate no files are referenced. - // and differentiate from error case where table_id not found. - return Some(HashSet::new()); - } - - let ref_file_set: HashSet = file_refs.files.keys().cloned().collect(); - + let region_refs = self.files_per_region.get(®ion_id)?.clone(); + let refs: HashSet<_> = region_refs.files.keys().cloned().collect(); debug!( "Get file refs for region {}, node {:?}, {} files", region_id, self.node_id, - ref_file_set.len(), + region_refs.files.len(), ); - - Some(ref_file_set) + Some(refs) } - /// Gets all ref files for the given regions, meaning all open FileHandles for those regions - /// and from related regions' manifests. - /// `query_regions_for_mem` queries for in memory file handles. - /// `related_regions_in_manifest` queries for related regions' manifests to get more file refs of given region ids. + /// Gets references from live SST metadata views and related regions' manifests. pub(crate) async fn get_snapshot_of_file_refs( &self, query_regions_for_mem: Vec, dst_region_to_src_regions: Vec<(MitoRegionRef, HashSet)>, ) -> Result { + // Imported SSTs are tracked under their source region even if that region is + // not open on this node. Capture their old readers before reading manifests. + let region_ids: HashSet<_> = query_regions_for_mem + .iter() + .map(|r| r.region_id()) + .chain( + dst_region_to_src_regions + .iter() + .flat_map(|(_, sources)| sources.iter().copied()), + ) + .collect(); let mut ref_files = HashMap::new(); - // get from in memory file handles - for region_id in query_regions_for_mem.iter().map(|r| r.region_id()) { + for region_id in region_ids { if let Some(files) = self.ref_file_set(region_id) { ref_files.insert(region_id, files); } } let mut manifest_version = HashMap::new(); - let mut cross_region_refs = HashMap::new(); - // get file refs from related regions' manifests for (dst_region, src_regions) in &dst_region_to_src_regions { let manifest = dst_region.manifest_ctx.manifest().await; for meta in manifest.files.values() { @@ -128,9 +114,8 @@ impl FileReferenceManager { .entry(meta.region_id) .or_insert_with(HashSet::new) .insert(dst_region.region_id()); - // since gc couldn't happen together with repartition - // (both the queries and related_region acquire region read lock), no need to worry about - // staging manifest in repartition here. + // GC and repartition hold region read locks, so staging manifests + // cannot change ownership while these references are collected. ref_files .entry(meta.region_id) .or_insert_with(HashSet::new) @@ -141,23 +126,24 @@ impl FileReferenceManager { )); } } - // not sure if related region's manifest version is needed, but record it for now. manifest_version.insert(dst_region.region_id(), manifest.manifest_version); } for r in &query_regions_for_mem { let manifest = r.manifest_ctx.manifest().await; - // remove in manifest files for smaller size, since gc worker read from manifest later. + // The manifest only protects its current index generation; older readers + // still need their exact-version temporary references. ref_files.entry(r.region_id()).and_modify(|refs| { - *refs = std::mem::take(refs) - .into_iter() - .filter(|f| !manifest.files.contains_key(&f.file_id)) - .collect(); + refs.retain(|f| { + !manifest + .files + .get(&f.file_id) + .is_some_and(|meta| meta.index_version() == f.index_version) + }); }); manifest_version.insert(r.region_id(), manifest.manifest_version); } - // simply return all ref files, no manifest version filtering for now. Ok(FileRefsManifest { file_refs: ref_files, manifest_version, @@ -165,77 +151,44 @@ impl FileReferenceManager { }) } - /// Adds a new file reference. - /// Also records the access layer for the table if not exists. - /// The access layer will be used to upload ref file to object storage. - pub fn add_file(&self, file_meta: &FileMeta) { - let region_id = file_meta.region_id; - let mut is_new = false; - { - let file_ref = FileRef::new( - file_meta.region_id, - file_meta.file_id, - file_meta.index_version(), - ); - self.files_per_region - .entry(region_id) - .and_modify(|refs| { - refs.files - .entry(file_ref.clone()) - .and_modify(|count| *count += 1) - .or_insert_with(|| { - is_new = true; - 1 - }); - }) - .or_insert_with(|| RegionFileRefs { - files: HashMap::from_iter([(file_ref, 1)]), - }); - } - if is_new { + /// Registers a metadata view, paired with `remove_file` when its last reader leaves. + pub(crate) fn register_file(&self, file_meta: &FileMeta) { + let mut region = self + .files_per_region + .entry(file_meta.region_id) + .or_default(); + let file_ref = FileRef::new( + file_meta.region_id, + file_meta.file_id, + file_meta.index_version(), + ); + let count = region.files.entry(file_ref).or_insert_with(|| { GC_REF_FILE_CNT.inc(); - } + 0 + }); + *count += 1; } - /// Removes a file reference. - /// If the reference count reaches zero, the file reference will be removed from the manager. - pub fn remove_file(&self, file_meta: &FileMeta) { - let region_id = file_meta.region_id; - let file_ref = FileRef::new(region_id, file_meta.file_id, file_meta.index_version()); - - let mut remove_table_entry = false; - let mut remove_file_ref = false; - let mut file_cnt = 0; - - let region_ref = self.files_per_region.entry(region_id).and_modify(|refs| { - let entry = refs.files.entry(file_ref.clone()).and_modify(|count| { - if *count > 0 { - *count -= 1; - } - if *count == 0 { - remove_file_ref = true; - } - }); - if let std::collections::hash_map::Entry::Occupied(o) = entry - && remove_file_ref + /// Releases a metadata view's reference, regardless of whether the SST was retired. + pub(crate) fn remove_file(&self, file_meta: &FileMeta) { + let file_ref = FileRef::new( + file_meta.region_id, + file_meta.file_id, + file_meta.index_version(), + ); + if let Entry::Occupied(mut region) = self.files_per_region.entry(file_meta.region_id) { + if let std::collections::hash_map::Entry::Occupied(mut file) = + region.get_mut().files.entry(file_ref) { - o.remove_entry(); + *file.get_mut() -= 1; + if *file.get() == 0 { + file.remove(); + GC_REF_FILE_CNT.dec(); + } } - - file_cnt = refs.files.len(); - - if refs.files.is_empty() { - remove_table_entry = true; + if region.get().files.is_empty() { + region.remove(); } - }); - - if let Entry::Occupied(o) = region_ref - && remove_table_entry - { - o.remove_entry(); - } - if remove_file_ref { - GC_REF_FILE_CNT.dec(); } } @@ -246,135 +199,42 @@ impl FileReferenceManager { #[cfg(test)] mod tests { - use std::num::NonZeroU64; - - use smallvec::SmallVec; - use store_api::storage::{FileId, RegionId}; + use store_api::storage::FileId; use super::*; - use crate::sst::file::{ColumnIndexMetadata, FileMeta, FileTimeRange, IndexType, RegionFileId}; + use crate::sst::file::IndexType; - #[tokio::test] - async fn test_file_ref_mgr() { - common_telemetry::init_default_ut_logging(); - - let sst_file_id = RegionFileId::new(RegionId::new(0, 0), FileId::random()); - - let file_ref_mgr = FileReferenceManager::new(None); - - let file_meta = FileMeta { - region_id: sst_file_id.region_id(), - file_id: sst_file_id.file_id(), - time_range: FileTimeRange::default(), - level: 0, - file_size: 4096, - max_row_group_uncompressed_size: 4096, - available_indexes: SmallVec::from_iter([IndexType::InvertedIndex]), - indexes: vec![ColumnIndexMetadata { - column_id: 0, - created_indexes: SmallVec::from_iter([IndexType::InvertedIndex]), - ..Default::default() - }], - index_file_size: 4096, - index_version: 0, - num_rows: 1024, - num_row_groups: 1, - sequence: NonZeroU64::new(4096), - partition_expr: None, - num_series: 0, + #[test] + fn test_index_view_references_release_generations_independently() { + let manager = FileReferenceManager::new(None); + let meta = FileMeta { + region_id: RegionId::new(1, 1), + file_id: FileId::random(), + available_indexes: smallvec::smallvec![IndexType::InvertedIndex], ..Default::default() }; + let updated = FileMeta { + index_version: 2, + ..meta.clone() + }; + manager.register_file(&meta); + manager.register_file(&meta); + manager.register_file(&updated); + let expected: HashSet<_> = [0, 2] + .into_iter() + .map(|version| FileRef::new(meta.region_id, meta.file_id, Some(version))) + .collect(); + assert_eq!(manager.ref_file_set(meta.region_id).unwrap(), expected); - file_ref_mgr.add_file(&file_meta); - + // Independent readers of the same index version must both release it. + manager.remove_file(&meta); + assert_eq!(manager.ref_file_set(meta.region_id).unwrap(), expected); + manager.remove_file(&meta); assert_eq!( - file_ref_mgr - .files_per_region - .get(&file_meta.region_id) - .unwrap() - .files, - HashMap::from_iter([( - FileRef::new( - file_meta.region_id, - file_meta.file_id, - file_meta.index_version() - ), - 1 - )]) - ); - - file_ref_mgr.add_file(&file_meta); - - let expected_region_ref_manifest = HashSet::from_iter([FileRef::new( - file_meta.region_id, - file_meta.file_id, - file_meta.index_version(), - )]); - - assert_eq!( - file_ref_mgr.ref_file_set(file_meta.region_id).unwrap(), - expected_region_ref_manifest - ); - - assert_eq!( - file_ref_mgr - .files_per_region - .get(&file_meta.region_id) - .unwrap() - .files, - HashMap::from_iter([( - FileRef::new( - file_meta.region_id, - file_meta.file_id, - file_meta.index_version() - ), - 2 - )]) - ); - - assert_eq!( - file_ref_mgr.ref_file_set(file_meta.region_id).unwrap(), - expected_region_ref_manifest - ); - - file_ref_mgr.remove_file(&file_meta); - - assert_eq!( - file_ref_mgr - .files_per_region - .get(&file_meta.region_id) - .unwrap() - .files, - HashMap::from_iter([( - FileRef::new( - file_meta.region_id, - file_meta.file_id, - file_meta.index_version() - ), - 1 - )]) - ); - - assert_eq!( - file_ref_mgr.ref_file_set(file_meta.region_id).unwrap(), - expected_region_ref_manifest - ); - - file_ref_mgr.remove_file(&file_meta); - - assert!( - file_ref_mgr - .files_per_region - .get(&file_meta.region_id) - .is_none(), - "{:?}", - file_ref_mgr.files_per_region - ); - - assert!( - file_ref_mgr.ref_file_set(file_meta.region_id).is_none(), - "{:?}", - file_ref_mgr.files_per_region + manager.ref_file_set(meta.region_id).unwrap(), + HashSet::from([FileRef::new(meta.region_id, meta.file_id, Some(2))]) ); + manager.remove_file(&updated); + assert!(manager.ref_file_set(meta.region_id).is_none()); } } diff --git a/src/mito2/src/sst/version.rs b/src/mito2/src/sst/version.rs index 1ca0410db7b..5351195eefb 100644 --- a/src/mito2/src/sst/version.rs +++ b/src/mito2/src/sst/version.rs @@ -100,7 +100,7 @@ impl SstVersion { } } else { // include case like old file have no index or index is outdated - *f = FileHandle::new(file.clone(), file_purger.clone()); + *f = f.with_updated_index(file.clone()); } }) .or_insert_with(|| { @@ -352,6 +352,42 @@ mod tests { assert_eq!(version.time_range(), None); } + #[rstest::rstest] + fn test_index_replacement_shares_retirement_with_old_readers( + #[values(false, true)] retire_before_update: bool, + ) { + let purger = new_noop_file_purger(); + let meta = FileMeta { + file_id: FileId::random(), + ..Default::default() + }; + let mut version = SstVersion::new(crate::test_util::memtable_util::metadata_for_test()); + version.add_files(purger.clone(), [meta.clone()].into_iter()); + let old_snapshot = version.clone(); + let old_reader = old_snapshot.levels()[0].files[&meta.file_id].clone(); + let independent = FileHandle::new(meta.clone(), purger.clone()); + if retire_before_update { + old_reader.mark_deleted(); + } + let mut updated = meta.clone(); + updated.index_version = 1; + updated + .available_indexes + .push(crate::sst::file::IndexType::InvertedIndex); + version.add_files(purger, [updated.clone()].into_iter()); + let new_reader = version.levels()[0].files[&meta.file_id].clone(); + + assert_eq!(old_reader.meta_ref(), &meta); + assert_eq!(new_reader.meta_ref(), &updated); + assert_eq!(new_reader.is_deleted(), retire_before_update); + version.remove_files([updated].into_iter()); + assert!(new_reader.is_deleted()); + assert!(old_reader.is_deleted()); + assert!(!independent.is_deleted()); + drop(version); + assert!(old_snapshot.levels()[0].files[&meta.file_id].is_deleted()); + } + #[test] fn test_add_files() { let purger = new_noop_file_purger(); From 5b4efd908ca03c8f5243d308d96146365600c73c Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Tue, 27 Oct 2026 10:49:16 +0000 Subject: [PATCH 108/239] perf(mito2): index encoded bulk parts from memory during flush (#9913) * perf(mito2): build encoded bulk part indexes from in-memory bytes during flush Sync flush uploaded encoded bulk parts without indexes and rebuilt them after the flush, which read every encoded SST back from the object store and needed a second manifest commit. Build the indexes from the encoded bytes still in memory before the flush commits, the same way the regular SST writer does, and drop the post-flush rebuild for these files. Signed-off-by: Dennis Zhuang * fix(mito2): keep flush going when indexing encoded bytes fails Match ParquetWriter: a failure while building indexes from the encoded bytes logs and leaves the SST unindexed instead of failing the flush after the SST was uploaded. Signed-off-by: Dennis Zhuang * fix(mito2): own encoded SSTs for cleanup before uploading them Index creation now runs after an encoded SST is uploaded, so a truncate or drop could cancel the flush after the upload but before UncommittedSsts tracked the file, leaving the SST and its index behind. Track the file before the upload, and fill the index metadata on the caller's SstInfo instead of a partial one. Signed-off-by: Dennis Zhuang * refactor(mito2): name encoded SST index build params explicitly Signed-off-by: Dennis Zhuang * fix(mito2): drop index metadata when the puffin file fails and size index decode batches Indexer::do_finish kept the per-index metadata when finishing the puffin file failed, so callers could publish indexes that were not written. Return an empty output instead. Indexing encoded bulk parts now decodes with the batch size estimated from the largest row group, the same estimate bulk merges use, instead of a fixed 8192 rows. Signed-off-by: Dennis Zhuang * docs(config): document the mito index build_mode option Signed-off-by: Dennis Zhuang * fix(mito2): decode encoded SSTs for indexing on the blocking pool Indexing encoded bulk parts decoded Parquet batches inside the flush task on the global async runtime. Decode each batch with spawn_blocking_global and keep the indexer update on the async task, with one batch in flight. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- config/config.md | 2 + config/datanode.example.toml | 6 + config/standalone.example.toml | 6 + src/mito2/src/access_layer.rs | 69 ++++++- src/mito2/src/batch_size.rs | 19 ++ src/mito2/src/cache/write_cache.rs | 79 ++++++-- src/mito2/src/engine/index_build_test.rs | 214 +++++++++++++++++++--- src/mito2/src/flush.rs | 86 +++------ src/mito2/src/memtable/bulk.rs | 17 +- src/mito2/src/request.rs | 2 - src/mito2/src/sst/file.rs | 5 + src/mito2/src/sst/index.rs | 82 ++++++++- src/mito2/src/sst/index/indexer/finish.rs | 5 + src/mito2/src/worker/handle_flush.rs | 17 +- 14 files changed, 468 insertions(+), 141 deletions(-) diff --git a/config/config.md b/config/config.md index f04689d97f6..15cd4915f9d 100644 --- a/config/config.md +++ b/config/config.md @@ -225,6 +225,7 @@ | `region_engine.mito.index.aux_path` | String | `""` | Auxiliary directory path for the index in filesystem, used to store intermediate files for
creating the index and staging files for searching the index, defaults to `{data_home}/index_intermediate`.
The default name for this directory is `index_intermediate` for backward compatibility.

This path contains two subdirectories:
- `__intm`: for storing intermediate files used during creating index.
- `staging`: for storing staging files used during searching index. | | `region_engine.mito.index.staging_size` | String | `2GB` | The max capacity of the staging directory. | | `region_engine.mito.index.staging_ttl` | String | `7d` | The TTL of the staging directory.
Defaults to 7 days.
Setting it to "0s" to disable TTL. | +| `region_engine.mito.index.build_mode` | String | `sync` | When to build indexes for new SST files.
- `sync`: build indexes while writing SST files during flush and compaction.
- `async`: build indexes in the background after the SST files are written. Queries read
these files without indexes until the build finishes. | | `region_engine.mito.index.metadata_cache_size` | String | `64MiB` | Cache size for inverted index metadata. | | `region_engine.mito.index.content_cache_size` | String | `128MiB` | Cache size for inverted index content. | | `region_engine.mito.index.content_cache_page_size` | String | `64KiB` | Page size for inverted index content cache. | @@ -703,6 +704,7 @@ | `region_engine.mito.index.aux_path` | String | `""` | Auxiliary directory path for the index in filesystem, used to store intermediate files for
creating the index and staging files for searching the index, defaults to `{data_home}/index_intermediate`.
The default name for this directory is `index_intermediate` for backward compatibility.

This path contains two subdirectories:
- `__intm`: for storing intermediate files used during creating index.
- `staging`: for storing staging files used during searching index. | | `region_engine.mito.index.staging_size` | String | `2GB` | The max capacity of the staging directory. | | `region_engine.mito.index.staging_ttl` | String | `7d` | The TTL of the staging directory.
Defaults to 7 days.
Setting it to "0s" to disable TTL. | +| `region_engine.mito.index.build_mode` | String | `sync` | When to build indexes for new SST files.
- `sync`: build indexes while writing SST files during flush and compaction.
- `async`: build indexes in the background after the SST files are written. Queries read
these files without indexes until the build finishes. | | `region_engine.mito.index.metadata_cache_size` | String | `64MiB` | Cache size for inverted index metadata. | | `region_engine.mito.index.content_cache_size` | String | `128MiB` | Cache size for inverted index content. | | `region_engine.mito.index.content_cache_page_size` | String | `64KiB` | Page size for inverted index content cache. | diff --git a/config/datanode.example.toml b/config/datanode.example.toml index f6e9a2571f8..acc65582e81 100644 --- a/config/datanode.example.toml +++ b/config/datanode.example.toml @@ -734,6 +734,12 @@ staging_size = "2GB" ## Setting it to "0s" to disable TTL. staging_ttl = "7d" +## When to build indexes for new SST files. +## - `sync`: build indexes while writing SST files during flush and compaction. +## - `async`: build indexes in the background after the SST files are written. Queries read +## these files without indexes until the build finishes. +build_mode = "sync" + ## Cache size for inverted index metadata. metadata_cache_size = "64MiB" diff --git a/config/standalone.example.toml b/config/standalone.example.toml index f569ff1b952..5cfbb95729e 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -949,6 +949,12 @@ staging_size = "2GB" ## Setting it to "0s" to disable TTL. staging_ttl = "7d" +## When to build indexes for new SST files. +## - `sync`: build indexes while writing SST files during flush and compaction. +## - `async`: build indexes in the background after the SST files are written. Queries read +## these files without indexes until the build finishes. +build_mode = "sync" + ## Cache size for inverted index metadata. metadata_cache_size = "64MiB" diff --git a/src/mito2/src/access_layer.rs b/src/mito2/src/access_layer.rs index b1457935c16..3067c6098a1 100644 --- a/src/mito2/src/access_layer.rs +++ b/src/mito2/src/access_layer.rs @@ -44,9 +44,9 @@ use crate::metrics::{COMPACTION_STAGE_ELAPSED, FLUSH_ELAPSED}; use crate::read::FlatSource; use crate::region::options::IndexOptions; use crate::sst::file::{FileHandle, RegionFileId, RegionIndexId}; -use crate::sst::index::IndexerBuilderImpl; use crate::sst::index::intermediate::IntermediateManager; use crate::sst::index::puffin_manager::{PuffinManagerFactory, SstPuffinManager}; +use crate::sst::index::{IndexerBuilderImpl, build_index_from_parquet_bytes}; use crate::sst::location::{self, region_dir_from_table_dir}; use crate::sst::parquet::reader::ParquetReaderBuilder; use crate::sst::parquet::writer::ParquetWriter; @@ -477,13 +477,19 @@ impl AccessLayer { } /// Puts encoded SST bytes to the write cache (if enabled) and uploads it to the object store. + /// + /// Builds the SST's indexes from `data` into `sst_info.index_metadata` when + /// `index_build_params` is set. pub(crate) async fn put_sst( &self, data: &bytes::Bytes, - region_file_id: RegionFileId, + region_id: RegionId, + sst_info: &mut SstInfo, cache_manager: &CacheManagerRef, write_buffer_size: ReadableSize, + index_build_params: Option, ) -> Result { + let region_file_id = RegionFileId::new(region_id, sst_info.file_id); if let Some(write_cache) = cache_manager.write_cache() { // Write to cache and upload to remote store let upload_request = SstUploadRequest { @@ -494,7 +500,14 @@ impl AccessLayer { remote_store: self.object_store.clone(), }; write_cache - .put_and_upload_sst(data, region_file_id, upload_request, write_buffer_size) + .put_and_upload_sst( + data, + region_id, + sst_info, + upload_request, + write_buffer_size, + index_build_params, + ) .await } else { let start = Instant::now(); @@ -519,6 +532,19 @@ impl AccessLayer { } let mut metrics = Metrics::new(WriteType::Flush); metrics.write_batch = start.elapsed(); + + if let Some(params) = index_build_params { + let start = Instant::now(); + let builder = params.indexer_builder( + self.puffin_manager_factory + .build(self.object_store.clone(), path_provider), + false, + self.intermediate_manager.clone(), + ); + sst_info.index_metadata = + build_index_from_parquet_bytes(&builder, region_file_id, data).await; + metrics.update_index = start.elapsed(); + } Ok(metrics) } } @@ -593,6 +619,36 @@ pub struct SstWriteRequest { pub bloom_filter_index_config: BloomFilterConfig, } +/// Index inputs for SST bytes that were encoded outside [ParquetWriter]. +pub(crate) struct EncodedSstIndexBuildParams { + pub(crate) metadata: RegionMetadataRef, + pub(crate) index_options: IndexOptions, + pub(crate) inverted_index_config: InvertedIndexConfig, + pub(crate) fulltext_index_config: FulltextIndexConfig, + pub(crate) bloom_filter_index_config: BloomFilterConfig, +} + +impl EncodedSstIndexBuildParams { + pub(crate) fn indexer_builder( + self, + puffin_manager: SstPuffinManager, + write_cache_enabled: bool, + intermediate_manager: IntermediateManager, + ) -> IndexerBuilderImpl { + IndexerBuilderImpl { + build_type: OperationType::Flush.into(), + metadata: self.metadata, + puffin_manager, + write_cache_enabled, + intermediate_manager, + index_options: self.index_options, + inverted_index_config: self.inverted_index_config, + fulltext_index_config: self.fulltext_index_config, + bloom_filter_index_config: self.bloom_filter_index_config, + } + } +} + /// Cleaner to remove temp files on the atomic write dir. pub(crate) struct TempFileCleaner { region_id: RegionId, @@ -830,9 +886,14 @@ mod tests { access_layer .put_sst( &encoded, - file_id, + file_id.region_id(), + &mut SstInfo { + file_id: file_id.file_id(), + ..Default::default() + }, &cache_manager, ReadableSize(chunk_size as u64), + None, ) .await .unwrap(); diff --git a/src/mito2/src/batch_size.rs b/src/mito2/src/batch_size.rs index 92fdd75673f..b71265ca7ee 100644 --- a/src/mito2/src/batch_size.rs +++ b/src/mito2/src/batch_size.rs @@ -14,11 +14,30 @@ //! Utilities for choosing an execution batch size from source statistics. +use parquet::file::metadata::ParquetMetaData; + use crate::sst::parquet::DEFAULT_READ_BATCH_SIZE; /// Approximate decoded bytes to produce in one batch. pub(crate) const TARGET_BATCH_BYTES: usize = 64 * 1024 * 1024; +/// Returns `(num_rows, uncompressed_bytes)` of the largest row group in a Parquet file, which +/// bounds the decoded size of a batch read from it. +pub(crate) fn largest_row_group_statistic(metadata: &ParquetMetaData) -> Option<(u64, u64)> { + metadata + .row_groups() + .iter() + .map(|row_group| { + let uncompressed_bytes = row_group + .columns() + .iter() + .map(|column| column.uncompressed_size() as u64) + .sum(); + (row_group.num_rows() as u64, uncompressed_bytes) + }) + .max_by_key(|(_, uncompressed_bytes)| *uncompressed_bytes) +} + /// Estimates a batch size from `(num_rows, estimated_bytes)` pairs. /// /// Uses the widest source with valid statistics and targets approximately 64 MiB of decoded data. diff --git a/src/mito2/src/cache/write_cache.rs b/src/mito2/src/cache/write_cache.rs index afe943eef50..24e2e957219 100644 --- a/src/mito2/src/cache/write_cache.rs +++ b/src/mito2/src/cache/write_cache.rs @@ -26,8 +26,9 @@ use store_api::storage::RegionId; use tokio::sync::mpsc::{UnboundedSender, unbounded_channel}; use crate::access_layer::{ - FilePathProvider, Metrics, OperationType, RegionFilePathFactory, SstInfoArray, SstWriteRequest, - TempFileCleaner, WriteCachePathProvider, WriteType, new_fs_cache_store, + EncodedSstIndexBuildParams, FilePathProvider, Metrics, OperationType, RegionFilePathFactory, + SstInfoArray, SstWriteRequest, TempFileCleaner, WriteCachePathProvider, WriteType, + new_fs_cache_store, }; use crate::cache::file_cache::{FileCache, FileCacheRef, FileType, IndexKey, IndexValue}; use crate::cache::manifest_cache::ManifestCache; @@ -35,11 +36,11 @@ use crate::error::{self, Result}; use crate::metrics::UPLOAD_BYTES_TOTAL; use crate::region::opener::RegionLoadCacheTask; use crate::sst::file::RegionFileId; -use crate::sst::index::IndexerBuilderImpl; use crate::sst::index::intermediate::IntermediateManager; use crate::sst::index::puffin_manager::{PuffinManagerFactory, SstPuffinManager}; -use crate::sst::parquet::WriteOptions; +use crate::sst::index::{IndexerBuilderImpl, build_index_from_parquet_bytes}; use crate::sst::parquet::writer::ParquetWriter; +use crate::sst::parquet::{SstInfo, WriteOptions}; use crate::sst::{DEFAULT_WRITE_BUFFER_SIZE, DEFAULT_WRITE_CONCURRENCY}; /// Wraps the remote object store used by write cache uploads. @@ -173,15 +174,20 @@ impl WriteCache { } /// Put encoded SST data to the cache and upload to the remote object store. + /// + /// Builds the SST's indexes from `data` when `index_build_params` is set and uploads them + /// with the SST. pub(crate) async fn put_and_upload_sst( &self, data: &bytes::Bytes, - region_file_id: RegionFileId, + region_id: RegionId, + sst_info: &mut SstInfo, upload_request: SstUploadRequest, write_buffer_size: ReadableSize, + index_build_params: Option, ) -> Result { - let region_id = region_file_id.region_id(); - let file_id = region_file_id.file_id(); + let file_id = sst_info.file_id; + let region_file_id = RegionFileId::new(region_id, file_id); let mut metrics = Metrics::new(WriteType::Flush); // Create index key for the SST file @@ -204,31 +210,70 @@ impl WriteCache { metrics.write_batch = cache_start.elapsed(); + if let Some(params) = index_build_params { + let start = Instant::now(); + let builder = params.indexer_builder( + self.build_puffin_manager(), + true, + self.intermediate_manager.clone(), + ); + sst_info.index_metadata = + build_index_from_parquet_bytes(&builder, region_file_id, data).await; + metrics.update_index = start.elapsed(); + } + // Upload to remote store + let remote_store = &upload_request.remote_store; + let mut upload_tracker = UploadTracker::new(region_id); let upload_start = Instant::now(); let remote_path = upload_request .dest_path_provider .build_sst_file_path(region_file_id); - - if let Err(e) = self + let mut result = self .upload( parquet_key, &remote_path, - &upload_request.remote_store, + remote_store, UploadOptions { // `put_and_upload_sst` is only used by the flush path. op_type: OperationType::Flush, write_buffer_size, }, ) - .await - { - // Clean up cache on failure - self.remove(parquet_key).await; - return Err(e); + .await; + metrics.upload_parquet = upload_start.elapsed(); + + if result.is_ok() && sst_info.index_metadata.file_size > 0 { + upload_tracker.push_uploaded_file(remote_path); + let puffin_key = IndexKey::new(region_id, file_id, FileType::Puffin(0)); + let puffin_path = upload_request + .dest_path_provider + .build_index_file_path(region_file_id); + let start = Instant::now(); + result = self + .upload( + puffin_key, + &puffin_path, + remote_store, + UploadOptions { + op_type: OperationType::Flush, + write_buffer_size: DEFAULT_WRITE_BUFFER_SIZE, + }, + ) + .await; + metrics.upload_puffin = start.elapsed(); } - metrics.upload_parquet = upload_start.elapsed(); + if let Err(e) = result { + upload_tracker + .clean( + std::slice::from_ref(sst_info), + &self.file_cache, + remote_store, + ) + .await; + return Err(e); + } Ok(metrics) } @@ -537,7 +582,7 @@ impl UploadTracker { /// Cleans uploaded files and files in the file cache at best effort. pub(crate) async fn clean( &self, - sst_info: &SstInfoArray, + sst_info: &[SstInfo], file_cache: &FileCacheRef, remote_store: &ObjectStore, ) { diff --git a/src/mito2/src/engine/index_build_test.rs b/src/mito2/src/engine/index_build_test.rs index 672c5e472d5..2877ac109d9 100644 --- a/src/mito2/src/engine/index_build_test.rs +++ b/src/mito2/src/engine/index_build_test.rs @@ -23,6 +23,8 @@ use api::v1::region::{StrictWindow, compact_request}; use api::v1::{Rows, SemanticType}; use async_trait::async_trait; use common_base::readable_size::ReadableSize; +use common_error::ext::ErrorExt; +use common_error::status_code::StatusCode; use common_recordbatch::RecordBatches; use datafusion_expr::expr::InList; use datafusion_expr::{Expr, col, lit}; @@ -34,7 +36,7 @@ use store_api::metadata::ColumnMetadata; use store_api::region_engine::RegionEngine; use store_api::region_request::{ AddColumn, AlterKind, RegionAlterRequest, RegionBuildIndexRequest, RegionCloseRequest, - RegionCompactRequest, RegionRequest, SetIndexOption, + RegionCompactRequest, RegionFlushRequest, RegionRequest, RegionTruncateRequest, SetIndexOption, }; use store_api::storage::{RegionId, ScanRequest}; use tokio::sync::{Notify, Semaphore}; @@ -2297,8 +2299,9 @@ impl SstReadCounter { #[rstest::rstest] #[tokio::test] -async fn test_sync_flush_indexes_encoded_bulk_parts( - #[values(false, true)] disable_inverted_on_flush: bool, +async fn test_flush_indexes_encoded_bulk_parts( + #[values(false, true)] async_without_flush_indexes: bool, + #[values(false, true)] write_cache: bool, ) { let sst_bytes_read = Arc::new(AtomicUsize::new(0)); let counter = sst_bytes_read.clone(); @@ -2312,19 +2315,27 @@ async fn test_sync_flush_indexes_encoded_bulk_parts( })) .build() .unwrap(); - let mut env = TestEnv::with_prefix("test_sync_flush_indexes_encoded_bulk_parts_") + let mut env = TestEnv::with_prefix("test_flush_indexes_encoded_bulk_parts_") .await .with_mock_layer(layer); - let listener = Arc::new(IndexBuildListener::default()); let mut config = MitoConfig { default_flat_format: true, min_compaction_interval: Duration::from_secs(3600), ..Default::default() }; assert_eq!(config.index.build_mode, IndexBuildMode::Sync); - if disable_inverted_on_flush { + if async_without_flush_indexes { + config.index.build_mode = IndexBuildMode::Async; config.inverted_index.create_on_flush = Mode::Disable; } + if write_cache { + config = config.enable_write_cache( + env.data_home().join("write_cache").display().to_string(), + ReadableSize::mb(32), + None, + ); + } + let listener = Arc::new(IndexBuildListener::default()); let engine = env .create_engine_with(config, None, Some(listener.clone()), None) .await; @@ -2345,7 +2356,7 @@ async fn test_sync_flush_indexes_encoded_bulk_parts( .insert_option("memtable.bulk.merge_threshold", "2") .insert_option("memtable.bulk.encode_row_threshold", "1") .build_with_index(); - if !disable_inverted_on_flush { + if !async_without_flush_indexes { // `build_with_index` only adds inverted indexes on tags; also cover a skipping index. let field = request .column_metadatas @@ -2385,9 +2396,10 @@ async fn test_sync_flush_indexes_encoded_bulk_parts( flush_region(&engine, region_id, None).await; let files = current_file_metas(&engine, region_id).await; - if disable_inverted_on_flush { - // Only inverted indexes are configured and they are disabled on flush, so the builds - // have nothing to create. They must stop before reading the rows of the SSTs. + assert!(!files.is_empty()); + if async_without_flush_indexes { + // Async mode builds indexes after flush. Only inverted indexes are configured and they + // are disabled on flush, so the builds must stop before reading the rows of the SSTs. let total_size: usize = files.iter().map(|f| f.file_size as usize).sum(); // Stopping the engine waits for scheduled index builds. engine.stop().await.unwrap(); @@ -2398,25 +2410,169 @@ async fn test_sync_flush_indexes_encoded_bulk_parts( return; } - let version = engine.get_region(region_id).unwrap().version(); - let indexed = async { - loop { - let files = current_file_metas(&engine, region_id).await; - if files.iter().all(|f| { - f.is_index_consistent_with_region( - &version.metadata.column_metadatas, - &version.options.index_options, - ) - }) { - return files; - } - tokio::time::sleep(Duration::from_millis(20)).await; - } - }; - let files = tokio::time::timeout(Duration::from_secs(10), indexed) - .await - .expect("flushed encoded parts have no index"); - assert!(listener.begin_count() > 0); + // Sync mode builds the indexes from the encoded bytes in memory during flush. + assert_eq!(0, sst_bytes_read.load(Ordering::Relaxed)); + assert_eq!(0, listener.begin_count()); assert!(files.iter().all(|f| f.inverted_index_available())); assert!(files.iter().all(|f| f.bloom_filter_index_available())); + let access_layer = engine.get_region(region_id).unwrap().access_layer.clone(); + for file in &files { + let index_path = location::index_file_path( + access_layer.table_dir(), + RegionIndexId::new(RegionFileId::new(region_id, file.file_id), 0), + access_layer.path_type(), + ); + assert!( + access_layer + .object_store() + .exists(&index_path) + .await + .unwrap() + ); + } +} + +/// Blocks closing a puffin file whose SST is already closed, which only happens for encoded +/// bulk parts: the SST writer closes the puffin file before the SST. +struct BlockingPuffinWriter { + file_name: String, + inner: object_store::layers::mock::Writer, + closed_ssts: Arc>>, + entered: Arc, + release: Arc, +} + +impl object_store::layers::mock::Write for BlockingPuffinWriter { + async fn write( + &mut self, + bs: object_store::layers::mock::Buffer, + ) -> object_store::layers::mock::Result<()> { + self.inner.write(bs).await + } + + async fn close( + &mut self, + ) -> object_store::layers::mock::Result { + if let Some(file_id) = self.file_name.strip_suffix(".puffin") + && self.closed_ssts.lock().unwrap().contains(file_id) + { + self.entered.notify_one(); + self.release.notified().await; + } + let metadata = self.inner.close().await?; + if let Some(file_id) = self.file_name.strip_suffix(".parquet") { + self.closed_ssts.lock().unwrap().insert(file_id.to_string()); + } + Ok(metadata) + } + + async fn abort(&mut self) -> object_store::layers::mock::Result<()> { + self.inner.abort().await + } +} + +#[tokio::test] +async fn test_truncate_while_indexing_encoded_bulk_part_removes_uploaded_sst() { + let entered = Arc::new(Notify::new()); + let release = Arc::new(Notify::new()); + let closed_ssts = Arc::new(std::sync::Mutex::new(HashSet::new())); + let (factory_entered, factory_release) = (entered.clone(), release.clone()); + let layer = object_store::layers::mock::MockLayerBuilder::default() + .writer_factory(Arc::new(move |path, _op, inner| { + Box::new(BlockingPuffinWriter { + file_name: path.rsplit('/').next().unwrap_or_default().to_string(), + inner, + closed_ssts: closed_ssts.clone(), + entered: factory_entered.clone(), + release: factory_release.clone(), + }) as object_store::layers::mock::Writer + })) + .build() + .unwrap(); + let mut env = TestEnv::with_prefix("test_truncate_while_indexing_encoded_bulk_part_") + .await + .with_mock_layer(layer); + let config = MitoConfig { + default_flat_format: true, + ..Default::default() + }; + let engine = env.create_engine_with(config, None, None, None).await; + + let region_id = RegionId::new(1, 1); + env.get_schema_metadata_manager() + .register_region_table_info( + region_id.table_id(), + "test_table", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + let request = CreateRequestBuilder::new() + .insert_option("memtable.type", "bulk") + .insert_option("memtable.bulk.merge_threshold", "2") + .insert_option("memtable.bulk.encode_row_threshold", "1") + .build_with_index(); + let column_schemas = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + for (start, end) in [(0, 100_000), (100_000, 200_000)] { + let rows = Rows { + schema: column_schemas.clone(), + rows: build_rows(start, end), + }; + put_rows(&engine, region_id, rows).await; + } + let version = engine.get_region(region_id).unwrap().version(); + for memtable in version.memtables.list_memtables() { + memtable.compact(false).unwrap(); + } + + let flush_engine = engine.clone(); + let flush_task = tokio::spawn(async move { + flush_engine + .handle_request( + region_id, + RegionRequest::Flush(RegionFlushRequest::default()), + ) + .await + }); + // The encoded SST is uploaded and its index is being written. + tokio::time::timeout(Duration::from_secs(5), entered.notified()) + .await + .expect("flush did not reach index creation"); + + tokio::time::timeout( + Duration::from_secs(5), + engine.handle_request( + region_id, + RegionRequest::Truncate(RegionTruncateRequest::All), + ), + ) + .await + .expect("truncate did not finish after cancelling flush") + .unwrap(); + release.notify_one(); + let flush_err = tokio::time::timeout(Duration::from_secs(5), flush_task) + .await + .expect("cancelled flush did not finish") + .expect("flush task panicked") + .unwrap_err(); + assert_eq!(flush_err.status_code(), StatusCode::Cancelled); + + let access_layer = engine.get_region(region_id).unwrap().access_layer.clone(); + let leftovers = access_layer + .object_store() + .list_with(access_layer.table_dir()) + .recursive(true) + .await + .unwrap() + .into_iter() + .map(|entry| entry.path().to_string()) + .filter(|path| path.ends_with(".parquet") || path.ends_with(".puffin")) + .collect::>(); + assert!(leftovers.is_empty(), "{leftovers:?}"); } diff --git a/src/mito2/src/flush.rs b/src/mito2/src/flush.rs index fba0607235b..0a4f1765f39 100644 --- a/src/mito2/src/flush.rs +++ b/src/mito2/src/flush.rs @@ -29,16 +29,17 @@ use smallvec::{SmallVec, smallvec}; use snafu::{ResultExt, ensure}; use store_api::metadata::RegionMetadataRef; use store_api::region_request::RegionFlushReason; -use store_api::storage::{FileId, RegionId, SequenceNumber}; +use store_api::storage::{RegionId, SequenceNumber}; use strum::IntoStaticStr; use tokio::sync::{Semaphore, mpsc, watch}; use crate::access_layer::{ - AccessLayerRef, Metrics, OperationType, SstInfoArray, SstWriteRequest, WriteType, + AccessLayerRef, EncodedSstIndexBuildParams, Metrics, OperationType, SstInfoArray, + SstWriteRequest, WriteType, }; use crate::cache::CacheManagerRef; use crate::compaction::{collect_json2_rewrite_plans, rewrite_json2_batch, rewrite_json2_schema}; -use crate::config::MitoConfig; +use crate::config::{IndexBuildMode, MitoConfig}; use crate::engine::region_hook::SstFileInfo; use crate::error::{ Error, FlushCancelledSnafu, FlushRegionSnafu, JoinSnafu, RegionBusySnafu, RegionClosedSnafu, @@ -66,7 +67,7 @@ use crate::request::{ }; use crate::schedule::CancellableTaskState; use crate::schedule::scheduler::{Job, SchedulerRef}; -use crate::sst::file::{FileMeta, RegionFileId, UncommittedSsts}; +use crate::sst::file::{FileMeta, UncommittedSsts}; use crate::sst::parquet::metadata::extract_primary_key_range; use crate::sst::parquet::{ DEFAULT_READ_BATCH_SIZE, DEFAULT_ROW_GROUP_SIZE, SstInfo, WriteOptions, flat_format, @@ -394,10 +395,7 @@ impl RegionFlushTask { }; let worker_request = match flush_result { - Ok(FlushedEdit { - edit, - unindexed_files, - }) => { + Ok(edit) => { let memtables_to_remove = version_data .version .memtables @@ -413,7 +411,6 @@ impl RegionFlushTask { senders: std::mem::take(&mut self.senders), _timer: timer, edit, - unindexed_files, memtables_to_remove, is_staging: self.is_staging, }; @@ -451,7 +448,7 @@ impl RegionFlushTask { version_data: &VersionControlData, state: &CancellableTaskState, uncommitted: &UncommittedSsts, - ) -> Result { + ) -> Result { // We must use the immutable memtables list and entry ids from the `version_data` // for consistency as others might already modify the version in the `version_control`. let version = &version_data.version; @@ -470,7 +467,6 @@ impl RegionFlushTask { let DoFlushMemtablesResult { file_metas, - unindexed_files, flushed_bytes, series_count, encoded_part_count, @@ -594,10 +590,7 @@ impl RegionFlushTask { self.reason.as_str() ); - Ok(FlushedEdit { - edit, - unindexed_files, - }) + Ok(edit) } async fn do_flush_memtables( @@ -612,7 +605,6 @@ impl RegionFlushTask { let mut flushed_bytes = 0; let mut series_count = 0; let mut encoded_part_count = 0; - let mut unindexed_files = Vec::new(); let mut flush_metrics = Metrics::new(WriteType::Flush); let partition_expr = parse_partition_expr(self.partition_expr.as_deref())?; let hook = self.manifest_ctx.hook(); @@ -659,12 +651,7 @@ impl RegionFlushTask { .await?; encoded_part_count += num_encoded; for (source_idx, result) in results.into_iter().enumerate() { - let FlushedSsts { - max_sequence, - ssts: ssts_written, - metrics, - indexed, - } = result?; + let (max_sequence, ssts_written, metrics) = result?; if ssts_written.is_empty() { // No data written. continue; @@ -683,9 +670,6 @@ impl RegionFlushTask { for sst_info in &ssts_written { flushed_bytes += sst_info.file_size; - if !indexed { - unindexed_files.push(sst_info.file_id); - } let pk_range = sst_info .file_metadata .as_ref() @@ -719,7 +703,6 @@ impl RegionFlushTask { Ok(DoFlushMemtablesResult { file_metas, - unindexed_files, flushed_bytes, series_count, encoded_part_count, @@ -765,41 +748,43 @@ impl RegionFlushTask { .await?; uncommitted.track(&ssts); FLUSH_FILE_TOTAL.inc_by(ssts.len() as u64); - Ok(FlushedSsts { - max_sequence, - ssts, - metrics, - indexed: true, - }) + Ok((max_sequence, ssts, metrics)) }); tasks.push(task); } - for (encoded, max_sequence) in flat_sources.encoded { + let build_index = self.engine_config.index.build_mode == IndexBuildMode::Sync; + for (mut encoded, max_sequence) in flat_sources.encoded { let access_layer = self.access_layer.clone(); let cache_manager = self.cache_manager.clone(); let region_id = version.metadata.region_id; let write_buffer_size = write_opts.write_buffer_size; let semaphore = self.flush_semaphore.clone(); let uncommitted = uncommitted.clone(); + let index_build_params = build_index.then(|| EncodedSstIndexBuildParams { + metadata: version.metadata.clone(), + index_options: self.index_options.clone(), + inverted_index_config: self.engine_config.inverted_index.clone(), + fulltext_index_config: self.engine_config.fulltext_index.clone(), + bloom_filter_index_config: self.engine_config.bloom_filter_index.clone(), + }); + // Index creation runs after the SST is uploaded, so cancellation can land after the + // upload. Own the files before uploading them. + uncommitted.track_before_upload(encoded.sst_info.file_id, index_build_params.is_some()); let task = common_runtime::spawn_global(async move { let _permit = semaphore.acquire().await.unwrap(); let metrics = access_layer .put_sst( &encoded.data, - RegionFileId::new(region_id, encoded.sst_info.file_id), + region_id, + &mut encoded.sst_info, &cache_manager, write_buffer_size, + index_build_params, ) .await?; uncommitted.track(std::slice::from_ref(&encoded.sst_info)); FLUSH_FILE_TOTAL.inc(); - // Encoded parts are uploaded as-is, without passing through the indexer. - Ok(FlushedSsts { - max_sequence, - ssts: smallvec![encoded.sst_info], - metrics, - indexed: false, - }) + Ok((max_sequence, smallvec![encoded.sst_info], metrics)) }); tasks.push(task); } @@ -935,30 +920,14 @@ impl RegionFlushTask { } } -struct FlushedEdit { - edit: RegionEdit, - /// Files in `edit` written without building their indexes. - unindexed_files: Vec, -} - -struct FlushedSsts { - max_sequence: SequenceNumber, - ssts: SstInfoArray, - metrics: Metrics, - /// Whether the indexer ran while writing `ssts`. - indexed: bool, -} - struct FlushFlatMemResult { num_encoded: usize, num_sources: usize, - results: Vec>, + results: Vec>, } struct DoFlushMemtablesResult { file_metas: Vec, - /// Files written without building their indexes. - unindexed_files: Vec, flushed_bytes: u64, series_count: usize, encoded_part_count: usize, @@ -1970,7 +1939,6 @@ mod tests { flushed_sequence: None, committed_sequence: None, }, - unindexed_files: Vec::new(), memtables_to_remove: smallvec![], is_staging: false, }), diff --git a/src/mito2/src/memtable/bulk.rs b/src/mito2/src/memtable/bulk.rs index 847a12db5ae..e323fb6b340 100644 --- a/src/mito2/src/memtable/bulk.rs +++ b/src/mito2/src/memtable/bulk.rs @@ -1155,20 +1155,9 @@ impl PartToMerge { PartToMerge::Multi { part, .. } => { Some((part.num_rows() as u64, part.estimated_size() as u64)) } - PartToMerge::Encoded { part, .. } => part - .metadata() - .parquet_metadata - .row_groups() - .iter() - .map(|row_group| { - let uncompressed_bytes = row_group - .columns() - .iter() - .map(|column| column.uncompressed_size() as u64) - .sum(); - (row_group.num_rows() as u64, uncompressed_bytes) - }) - .max_by_key(|(_, uncompressed_bytes)| *uncompressed_bytes), + PartToMerge::Encoded { part, .. } => { + crate::batch_size::largest_row_group_statistic(&part.metadata().parquet_metadata) + } } } diff --git a/src/mito2/src/request.rs b/src/mito2/src/request.rs index 2e81fbc4396..d7e12a5da60 100644 --- a/src/mito2/src/request.rs +++ b/src/mito2/src/request.rs @@ -1148,8 +1148,6 @@ pub(crate) struct FlushFinished { pub(crate) _timer: HistogramTimer, /// Region edit to apply. pub(crate) edit: RegionEdit, - /// Files in `edit` written without building their indexes. - pub(crate) unindexed_files: Vec, /// Memtables to remove. pub(crate) memtables_to_remove: SmallVec<[MemtableId; 2]>, /// Whether the region is in staging mode. diff --git a/src/mito2/src/sst/file.rs b/src/mito2/src/sst/file.rs index 1d4caaf9127..4c6613dee40 100644 --- a/src/mito2/src/sst/file.rs +++ b/src/mito2/src/sst/file.rs @@ -944,6 +944,11 @@ impl UncommittedSsts { } } + /// Tracks an SST before it is uploaded, so a cancelled upload is still removed. + pub(crate) fn track_before_upload(&self, file_id: FileId, with_index: bool) { + self.files.lock().unwrap().insert(file_id, (0, with_index)); + } + /// Disarms cleanup after the manifest edit has committed or may have committed. /// /// A manifest update error does not guarantee that the edit was not persisted. Once an edit diff --git a/src/mito2/src/sst/index.rs b/src/mito2/src/sst/index.rs index 4dc7002f147..fe39af344d2 100644 --- a/src/mito2/src/sst/index.rs +++ b/src/mito2/src/sst/index.rs @@ -35,14 +35,16 @@ use std::sync::Arc; use arrow_schema::extension::ExtensionType; use bloom_filter::creator::BloomFilterIndexer; +use bytes::Bytes; use common_telemetry::{debug, error, info, warn}; use datatypes::arrow::record_batch::RecordBatch; use datatypes::extension::json::Json2ExtensionType; use index::target::IndexTarget; use mito_codec::row_converter::DensePrimaryKeyCodec; use object_store::ObjectStore; +use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; use puffin_manager::SstPuffinManager; -use smallvec::{SmallVec, smallvec}; +use smallvec::SmallVec; use snafu::ResultExt; use statistics::{ByteCount, RowCount}; use store_api::codec::PrimaryKeyEncoding; @@ -52,6 +54,7 @@ use strum::IntoStaticStr; use tokio::sync::mpsc::Sender; use crate::access_layer::{AccessLayerRef, FilePathProvider, OperationType, RegionFilePathFactory}; +use crate::batch_size::{estimate_batch_size, largest_row_group_statistic}; use crate::cache::file_cache::{FileCacheRef, FileType, IndexKey}; use crate::cache::write_cache::{UploadOptions, UploadTracker, WriteCacheRef}; use crate::cache::{CacheManagerRef, CacheStrategy}; @@ -442,6 +445,80 @@ impl IndexerBuilder for IndexerBuilderImpl { } } +/// Builds the indexes of an SST from its encoded parquet bytes. +/// +/// Encoded bulk parts are flushed without passing through [ParquetWriter], so their indexes +/// are built from the bytes still in memory instead of reading the uploaded SST back. Like +/// the writer, a failure only drops the index and leaves the SST unindexed. +/// +/// [ParquetWriter]: crate::sst::parquet::writer::ParquetWriter +pub(crate) async fn build_index_from_parquet_bytes( + builder: &IndexerBuilderImpl, + region_file_id: RegionFileId, + data: &Bytes, +) -> IndexOutput { + let reader_builder = match ParquetRecordBatchReaderBuilder::try_new(data.clone()) { + Ok(reader_builder) => reader_builder, + Err(e) => { + warn!(e; "Failed to read encoded SST {:?} for index creation", region_file_id); + return IndexOutput::default(); + } + }; + let row_group_size = reader_builder + .metadata() + .row_groups() + .first() + .map(|row_group| row_group.num_rows() as usize) + .filter(|size| *size > 0); + let mut indexer = builder.build(region_file_id, 0, row_group_size).await; + if indexer.is_empty() { + return IndexOutput::default(); + } + // Follow the batch size the flush reads use, so wide rows are not decoded 8192 at a time. + let batch_size = estimate_batch_size(largest_row_group_statistic(reader_builder.metadata())); + + // TODO(dennis): decode only the columns the enabled indexers need, including the encoded + // primary key for tag indexes. This needs the indexer to stop locating columns by position. + let mut reader = match reader_builder.with_batch_size(batch_size).build() { + Ok(reader) => reader, + Err(e) => { + warn!(e; "Failed to read encoded SST {:?} for index creation", region_file_id); + indexer.abort().await; + return IndexOutput::default(); + } + }; + loop { + // Decoding is CPU-bound, so each batch is decoded on the blocking pool. Only one batch + // is in flight, and an aborted flush stops after the batch being decoded. + let decoded = common_runtime::spawn_blocking_global(move || { + let batch = reader.next(); + (reader, batch) + }) + .await; + let batch = match decoded { + Ok((next_reader, batch)) => { + reader = next_reader; + batch + } + Err(e) => { + warn!(e; "Failed to decode encoded SST {:?} for index creation", region_file_id); + indexer.abort().await; + return IndexOutput::default(); + } + }; + match batch { + Some(Ok(batch)) => indexer.update_flat(&batch).await, + Some(Err(e)) => { + warn!(e; "Failed to decode encoded SST {:?} for index creation", region_file_id); + indexer.abort().await; + return IndexOutput::default(); + } + None => break, + } + } + indexer.finish().await +} + impl IndexerBuilderImpl { fn build_inverted_indexer( &self, @@ -1164,7 +1241,7 @@ impl IndexBuildTask { // Cleans index files on failure. upload_tracker .clean( - &smallvec![SstInfo { + &[SstInfo { file_id, index_metadata: output, ..Default::default() @@ -1593,6 +1670,7 @@ mod tests { use partition::expr::col; use puffin::puffin_manager::{PuffinManager, PuffinReader}; use puffin_manager::PuffinManagerFactory; + use smallvec::smallvec; use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; use tokio::sync::mpsc; diff --git a/src/mito2/src/sst/index/indexer/finish.rs b/src/mito2/src/sst/index/indexer/finish.rs index 4e46f886ebd..1a6d90ac4f5 100644 --- a/src/mito2/src/sst/index/indexer/finish.rs +++ b/src/mito2/src/sst/index/indexer/finish.rs @@ -57,6 +57,11 @@ impl Indexer { self.do_prune_intm_sst_dir().await; output.file_size = self.do_finish_puffin_writer(writer).await; + // A failed puffin file must not advertise the indexes written into it. Not covered by + // tests: `do_finish_puffin_writer` panics on failure in test builds. + if output.file_size == 0 { + return IndexOutput::default(); + } output.version = self.index_version; output } diff --git a/src/mito2/src/worker/handle_flush.rs b/src/mito2/src/worker/handle_flush.rs index f09efe81e47..4f26911077f 100644 --- a/src/mito2/src/worker/handle_flush.rs +++ b/src/mito2/src/worker/handle_flush.rs @@ -439,21 +439,10 @@ impl RegionWorkerLoop { let flush_on_close = request.flush_reason == FlushReason::Closing; - let mut index_build_file_metas = std::mem::take(&mut request.edit.files_to_add); - if self.config.index.build_mode == IndexBuildMode::Sync { - // Sync flush builds indexes while writing SSTs, except for encoded bulk parts, - // which are uploaded as-is. Build the missing indexes for those files here. - let version = region.version(); - index_build_file_metas.retain(|meta| { - request.unindexed_files.contains(&meta.file_id) - && !meta.is_index_consistent_with_region( - &version.metadata.column_metadatas, - &version.options.index_options, - ) - }); - } + let index_build_file_metas = std::mem::take(&mut request.edit.files_to_add); - if !index_build_file_metas.is_empty() { + // In async mode, create indexes after flush. + if self.config.index.build_mode == IndexBuildMode::Async { self.handle_rebuild_index( BuildIndexRequest { region_id, From 984f80e2772c42413f90cd5e6cf13b20acd8d30c Mon Sep 17 00:00:00 2001 From: discord9 Date: Tue, 27 Oct 2026 11:51:43 +0000 Subject: [PATCH 109/239] fix(query): stop windowed sort on precision-lowering time casts (#9524) * test: cover windowed sort cast unit mismatch Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(query): stop windowed sort on unit- or timezone-changing time casts A cast on the time index column keeps the windowed-sort optimization eligible only when it preserves the underlying timestamp values. Casts that change the time unit (e.g. ms -> s) or that add a timezone to a naive timestamp re-interpret the values, so the working ranges derived from the time index no longer describe the sort column and the scan fails with a working-range error or silently wrong order. Reject both cast shapes so such queries fall back to the normal SortExec path. The integration fixture uses two rows 1000 seconds apart so the cast output stays distinct and the ordering assertions are not vacuous; the descending arm gets its own expectation. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(query): stop windowed sort on unit-changing timestamp wrappers to_timestamp_seconds/millis/... cast their argument to their own unit inside the function, so the hidden cast bypassed the explicit-cast guard and kept the windowed-sort optimization eligible on a sort column whose values no longer match the working ranges. The wrapper branch now applies the same unit- and value-preserving predicate as the cast branch, with the predicate extracted into preserves_timestamp_values and shared by both. The integration regression gains to_timestamp_seconds ASC/DESC arms and a to_timestamp_millis value-preserving control; unit tests cover unit-changing, tz-adding, and non-timestamp wrapper inputs. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test: regen windowed sort sqlness after wrapper guard The windowed_sort_nanos and windowed_sort_advance cases sort by to_timestamp_nanos over a millisecond index and to_timestamp_millis over a nanosecond index. Both wrappers change the time unit inside the function, so the windowed-sort rule no longer applies and the plans render a plain SortExec. Query results are unchanged; only the EXPLAIN plan line differs. Regenerated by the sqlness runner. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(query): simplify windowed sort cast regressions Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/query/src/optimizer/windowed_sort.rs | 263 +++++++++++++++--- tests-integration/src/tests/instance_test.rs | 168 +++++++++++ .../optimizer/windowed_sort_advance.result | 2 +- .../optimizer/windowed_sort_nanos.result | 2 +- .../optimizer/windowed_sort_advance.result | 2 +- .../optimizer/windowed_sort_nanos.result | 2 +- 6 files changed, 400 insertions(+), 39 deletions(-) diff --git a/src/query/src/optimizer/windowed_sort.rs b/src/query/src/optimizer/windowed_sort.rs index 62c17c40b7a..d1ec687c6cf 100644 --- a/src/query/src/optimizer/windowed_sort.rs +++ b/src/query/src/optimizer/windowed_sort.rs @@ -222,7 +222,8 @@ fn is_time_index_expr( } if let Some(cast_expr) = expr.downcast_ref::() { - return if matches!(cast_expr.cast_type(), DataType::Timestamp(_, _)) { + let source_type = cast_expr.expr().data_type(plan.schema().as_ref())?; + return if preserves_timestamp_values(cast_expr.cast_type(), &source_type) { is_time_index_expr(plan, cast_expr.expr()) } else { Ok(false) @@ -230,10 +231,22 @@ fn is_time_index_expr( } if let Some(scalar_function_expr) = expr.downcast_ref::() { - return if is_supported_time_index_wrapper(scalar_function_expr) - && scalar_function_expr.args().len() == 1 + if !is_supported_time_index_wrapper(scalar_function_expr) + || scalar_function_expr.args().len() != 1 { - is_time_index_expr(plan, &scalar_function_expr.args()[0]) + return Ok(false); + } + + // Wrappers such as `to_timestamp_seconds` cast their argument to another time unit + // inside the function. That cast does not show up as a `CastExpr` in the plan, but it + // changes the values of the sort column exactly like a cast in the plan does, so the + // wrapper is only still the time index if its result keeps both the unit and the values + // of its argument. + let arg = &scalar_function_expr.args()[0]; + let return_type = scalar_function_expr.data_type(plan.schema().as_ref())?; + let arg_type = arg.data_type(plan.schema().as_ref())?; + return if preserves_timestamp_values(&return_type, &arg_type) { + is_time_index_expr(plan, arg) } else { Ok(false) }; @@ -242,6 +255,24 @@ fn is_time_index_expr( Ok(false) } +/// Returns true if a cast, or a function wrapper, that turns `source` into `target` keeps both the +/// time unit and the underlying values of a timestamp. +/// +/// A cast (or a hidden cast inside a function wrapper) that changes the time unit changes the +/// values of the sort column, so the working ranges derived from the time index, which keep the +/// unit of the time index, no longer apply. The optimizer conservatively rejects every cast that +/// adds a timezone to a naive timestamp, including UTC, even though UTC has a zero offset. Only a +/// target that keeps both the unit and the values of its source can still be the time index. +fn preserves_timestamp_values(target: &DataType, source: &DataType) -> bool { + matches!( + (target, source), + ( + DataType::Timestamp(target_unit, target_tz), + DataType::Timestamp(source_unit, source_tz) + ) if target_unit == source_unit && (source_tz.is_some() || target_tz.is_none()) + ) +} + fn is_time_index_column( plan: &Arc, column_expr: &PhysicalColumn, @@ -344,7 +375,7 @@ mod test { use datafusion::config::ConfigOptions; use datafusion::physical_plan::filter::FilterExecBuilder; use datafusion_common::ScalarValue; - use datafusion_functions::datetime::to_timestamp_millis; + use datafusion_functions::datetime::{to_timestamp, to_timestamp_millis, to_timestamp_seconds}; use datafusion_physical_expr::expressions::{CastExpr, Literal}; use datatypes::data_type::ConcreteDataType; use datatypes::schema::{ColumnSchema, Schema}; @@ -416,37 +447,92 @@ mod test { } #[test] - fn test_is_time_index_expr_tracks_wrapped_aliases_through_projection() { - let scan = new_region_scan(); + fn test_is_time_index_expr_tracks_value_preserving_wrapper_through_projection() { + let scan = new_region_scan_with_unit(TimeUnit::Millisecond); let config = Arc::new(ConfigOptions::default()); - let return_field = Arc::new(Field::new( - "ts", - DataType::Timestamp(TimeUnit::Millisecond, None), - false, - )); + let wrapper = ScalarFunctionExpr::try_new( + to_timestamp_millis(config.as_ref()), + vec![Arc::new(PhysicalColumn::new("ts", 1))], + scan.schema().as_ref(), + config, + ) + .unwrap(); + // `to_timestamp_millis` keeps both the unit and the values of a millisecond timestamp. + assert_eq!( + wrapper.data_type(scan.schema().as_ref()).unwrap(), + DataType::Timestamp(TimeUnit::Millisecond, None) + ); let projection = Arc::new( ProjectionExec::try_new( - vec![( - Arc::new(ScalarFunctionExpr::new( - "to_timestamp_millis", - to_timestamp_millis(config.as_ref()), - vec![Arc::new(PhysicalColumn::new("ts", 1))], - return_field, - config, - )) as Arc, - "ts".to_string(), - )], + vec![(Arc::new(wrapper) as Arc, "ts".to_string())], scan, ) .unwrap(), ) as Arc; let expr = Arc::new(PhysicalColumn::new("ts", 0)) as Arc; + // So the wrapper is still the time index. assert!(is_time_index_expr(&projection, &expr).unwrap()); } #[test] - fn test_is_time_index_expr_tracks_cast_aliases_through_projection() { + fn test_is_time_index_expr_rejects_unit_changing_wrapper() { + let config = Arc::new(ConfigOptions::default()); + + // `to_timestamp_millis` casts its argument to milliseconds, which truncates the values of + // a nanosecond time index, so the wrapper is no longer the time index. + let nanosecond_scan = new_region_scan(); + let truncates_to_millis = Arc::new( + ScalarFunctionExpr::try_new( + to_timestamp_millis(config.as_ref()), + vec![Arc::new(PhysicalColumn::new("ts", 1))], + nanosecond_scan.schema().as_ref(), + config.clone(), + ) + .unwrap(), + ) as Arc; + assert_eq!( + truncates_to_millis + .data_type(nanosecond_scan.schema().as_ref()) + .unwrap(), + DataType::Timestamp(TimeUnit::Millisecond, None) + ); + assert!(!is_time_index_expr(&nanosecond_scan, &truncates_to_millis).unwrap()); + + // Same for a millisecond time index wrapped in `to_timestamp_seconds`. + let millisecond_scan = new_region_scan_with_unit(TimeUnit::Millisecond); + let truncates_to_seconds = Arc::new( + ScalarFunctionExpr::try_new( + to_timestamp_seconds(config.as_ref()), + vec![Arc::new(PhysicalColumn::new("ts", 1))], + millisecond_scan.schema().as_ref(), + config.clone(), + ) + .unwrap(), + ) as Arc; + assert_eq!( + truncates_to_seconds + .data_type(millisecond_scan.schema().as_ref()) + .unwrap(), + DataType::Timestamp(TimeUnit::Second, None) + ); + assert!(!is_time_index_expr(&millisecond_scan, &truncates_to_seconds).unwrap()); + + // A supported wrapper whose argument is not a timestamp is never the time index. + let from_integer = Arc::new( + ScalarFunctionExpr::try_new( + to_timestamp_millis(config.as_ref()), + vec![Arc::new(Literal::new(ScalarValue::Int64(Some(0))))], + nanosecond_scan.schema().as_ref(), + config, + ) + .unwrap(), + ) as Arc; + assert!(!is_time_index_expr(&nanosecond_scan, &from_integer).unwrap()); + } + + #[test] + fn test_is_time_index_expr_rejects_unit_changing_cast_through_projection() { let scan = new_region_scan(); let projection = Arc::new( ProjectionExec::try_new( @@ -464,7 +550,116 @@ mod test { ) as Arc; let expr = Arc::new(PhysicalColumn::new("ts_ms", 0)) as Arc; - assert!(is_time_index_expr(&projection, &expr).unwrap()); + // The cast changes the time unit (ns → ms), so it is no longer the time index. + assert!(!is_time_index_expr(&projection, &expr).unwrap()); + } + + #[test] + fn test_is_time_index_expr_rejects_timezone_adding_cast() { + let scan = new_region_scan(); + + // The optimizer conservatively rejects casts adding any timezone to a naive timestamp. + let adds_timezone = Arc::new(CastExpr::new( + Arc::new(PhysicalColumn::new("ts", 1)), + DataType::Timestamp(TimeUnit::Nanosecond, Some(Arc::from("+07:00"))), + None, + )) as Arc; + assert!(!is_time_index_expr(&scan, &adds_timezone).unwrap()); + + // UTC has a zero offset, but naive-to-aware timestamp casts are all rejected. + let adds_utc_timezone = Arc::new(CastExpr::new( + Arc::new(PhysicalColumn::new("ts", 1)), + DataType::Timestamp(TimeUnit::Nanosecond, Some(Arc::from("UTC"))), + None, + )) as Arc; + assert!(!is_time_index_expr(&scan, &adds_utc_timezone).unwrap()); + + // `to_timestamp` returns an aware timestamp in the session time zone for a naive + // timestamp argument. That re-interprets the underlying values in the same way as a + // timezone-adding cast, so the wrapper is rejected as well. + let mut config = ConfigOptions::default(); + config.execution.time_zone = Some("+08:00".to_string()); + let config = Arc::new(config); + let aware_source = ScalarFunctionExpr::try_new( + to_timestamp(config.as_ref()), + vec![Arc::new(PhysicalColumn::new("ts", 1))], + scan.schema().as_ref(), + config.clone(), + ) + .unwrap(); + let projection = Arc::new( + ProjectionExec::try_new( + vec![( + Arc::new(aware_source) as Arc, + "ts_aware".to_string(), + )], + scan, + ) + .unwrap(), + ) as Arc; + assert_eq!( + projection.schema().field(0).data_type(), + &DataType::Timestamp(TimeUnit::Nanosecond, Some(Arc::from("+08:00"))) + ); + let expr = Arc::new(PhysicalColumn::new("ts_aware", 0)) as Arc; + + assert!(!is_time_index_expr(&projection, &expr).unwrap()); + + // A cast between two aware timestamps of the same unit keeps the underlying values, see + // `test_preserves_timestamp_values`, but it is still rejected here as its input is not + // the time index. + let aware_cast = Arc::new(CastExpr::new( + Arc::new(PhysicalColumn::new("ts_aware", 0)), + DataType::Timestamp(TimeUnit::Nanosecond, Some(Arc::from("+07:00"))), + None, + )) as Arc; + + assert!(!is_time_index_expr(&projection, &aware_cast).unwrap()); + } + + #[test] + fn test_preserves_timestamp_values() { + let naive_nanosecond = DataType::Timestamp(TimeUnit::Nanosecond, None); + let naive_millisecond = DataType::Timestamp(TimeUnit::Millisecond, None); + let aware_nanosecond = DataType::Timestamp(TimeUnit::Nanosecond, Some(Arc::from("+08:00"))); + let other_aware_nanosecond = + DataType::Timestamp(TimeUnit::Nanosecond, Some(Arc::from("+07:00"))); + + // Keeping the unit and the values of a timestamp. + assert!(preserves_timestamp_values( + &naive_nanosecond, + &naive_nanosecond + )); + assert!(preserves_timestamp_values( + &aware_nanosecond, + &other_aware_nanosecond + )); + + // Changing the time unit, adding a timezone to a naive timestamp, or leaving the timestamp + // domain entirely all change the values of the timestamp. + assert!(!preserves_timestamp_values( + &naive_millisecond, + &naive_nanosecond + )); + assert!(!preserves_timestamp_values( + &aware_nanosecond, + &naive_nanosecond + )); + assert!(!preserves_timestamp_values( + &naive_nanosecond, + &DataType::Int64 + )); + assert!(!preserves_timestamp_values( + &DataType::Int64, + &naive_nanosecond + )); + + // Dropping an aware timestamp's timezone keeps the underlying values. Adding one to a + // naive timestamp is conservatively treated as non-preserving by the predicate. + assert!(preserves_timestamp_values( + &naive_nanosecond, + &aware_nanosecond + )); } #[test] @@ -524,7 +719,8 @@ mod test { DataType::Timestamp(TimeUnit::Millisecond, None), None, )) as Arc; - assert!(is_time_index_expr(&scan, &cast_expr).unwrap()); + // The cast changes the time unit (ns → ms), so it is no longer the time index. + assert!(!is_time_index_expr(&scan, &cast_expr).unwrap()); let non_timestamp_cast = Arc::new(CastExpr::new( Arc::new(PhysicalColumn::new("ts", 1)), @@ -626,13 +822,14 @@ mod test { } fn new_region_scan() -> Arc { + new_region_scan_with_unit(TimeUnit::Nanosecond) + } + + fn new_region_scan_with_unit(unit: TimeUnit) -> Arc { + let timestamp_type = ConcreteDataType::timestamp_datatype(unit.into()); let schema = Arc::new(Schema::new(vec![ ColumnSchema::new("value", ConcreteDataType::int32_datatype(), false), - ColumnSchema::new( - "ts", - ConcreteDataType::timestamp_nanosecond_datatype(), - false, - ), + ColumnSchema::new("ts", timestamp_type.clone(), false), ])); let recordbatches = RecordBatches::try_new(schema.clone(), vec![]).unwrap(); let stream = recordbatches.as_stream(); @@ -649,11 +846,7 @@ mod test { column_id: 1, }) .push_column_metadata(ColumnMetadata { - column_schema: ColumnSchema::new( - "ts", - ConcreteDataType::timestamp_nanosecond_datatype(), - false, - ), + column_schema: ColumnSchema::new("ts", timestamp_type, false), semantic_type: SemanticType::Timestamp, column_id: 2, }); diff --git a/tests-integration/src/tests/instance_test.rs b/tests-integration/src/tests/instance_test.rs index 481754e6df5..3b3dca0bbd0 100644 --- a/tests-integration/src/tests/instance_test.rs +++ b/tests-integration/src/tests/instance_test.rs @@ -949,6 +949,174 @@ ORDER BY t.key_id"#; ); } +/// Sorting by `CAST(ts AS TIMESTAMP(0))` must not be optimized as if it were sorting by the +/// (millisecond) time index column. +/// +/// The windowed sort rule still recognizes the cast as the time index column, but the working +/// ranges it derives from the region scan's `PartitionRange`s stay in the unit of the table's time +/// index (millisecond here). `PartSortExec` projects those ranges onto the unit of the sort column +/// while `WindowedSortExec` does not, so sorted runs (second) are checked against working ranges +/// (millisecond) and the range invariant trips with +/// "Current batch have data on the left side of working range, something is very wrong". +/// +/// The two rows are 1000 seconds apart, so casting the millisecond time index to second precision +/// truncates sub-second parts and collapses the first row onto second 0, which is before the start +/// of the region's time range. +/// +/// `to_timestamp_seconds(ts)` casts the time index to second precision *inside* the function, so it +/// has to be rejected exactly like the cast in the plan; the cast is just not visible as a +/// `CastExpr`. `to_timestamp_millis(ts)` keeps the values of the millisecond time index and is used +/// as the value-preserving control. +#[apply(standalone_instance_case)] +async fn test_windowed_sort_sort_column_cast_time_unit_mismatch(instance: Arc) { + let frontend = instance.frontend(); + + execute_sql( + &frontend, + r#"CREATE TABLE f05_mito_wc ( + ts TIMESTAMP(3) TIME INDEX, + host STRING PRIMARY KEY, + val DOUBLE + ) ENGINE = mito"#, + ) + .await; + + // The metric engine flavor only creates its physical table explicitly; rows are written to and + // read from the logical table built on it. + execute_sql( + &frontend, + r#"CREATE TABLE f05_metric_phy ( + ts TIMESTAMP(3) TIME INDEX, + host STRING PRIMARY KEY, + val DOUBLE + ) ENGINE = metric WITH ('physical_metric_table' = '')"#, + ) + .await; + execute_sql( + &frontend, + r#"CREATE TABLE f05_metric_wc ( + ts TIMESTAMP(3) TIME INDEX, + host STRING PRIMARY KEY, + val DOUBLE + ) ENGINE = metric WITH ('on_physical_table' = 'f05_metric_phy')"#, + ) + .await; + + let rows_by_time_index = r#"+-------------------------+------+ +| ts | host | ++-------------------------+------+ +| 1970-01-01T00:00:00.001 | a | +| 1970-01-01T00:16:40.001 | b | ++-------------------------+------+"#; + let rows_by_cast = r#"+---------------------+------+ +| t | host | ++---------------------+------+ +| 1970-01-01T00:00:00 | a | +| 1970-01-01T00:16:40 | b | ++---------------------+------+"#; + // `ORDER BY t DESC` returns the same rows in the opposite order; without distinct cast values + // the two rows tie and this expectation would be vacuous. `rows_by_cast` is an ascending + // fixture, so the descending query needs its own. + let rows_by_cast_desc = r#"+---------------------+------+ +| t | host | ++---------------------+------+ +| 1970-01-01T00:16:40 | b | +| 1970-01-01T00:00:00 | a | ++---------------------+------+"#; + + // `to_timestamp_seconds` casts the millisecond time index to second precision inside the + // function, so it returns the same rows as the direct cast above (the function is registered + // without a session time zone, so its result stays a naive timestamp). + let rows_by_seconds_wrapper = rows_by_cast; + let rows_by_seconds_wrapper_desc = rows_by_cast_desc; + // `to_timestamp_millis` keeps the values of the millisecond time index, so the control below + // returns the same rows as sorting by the time index itself, only under the `t` alias. + let rows_by_millis_wrapper = r#"+-------------------------+------+ +| t | host | ++-------------------------+------+ +| 1970-01-01T00:00:00.001 | a | +| 1970-01-01T00:16:40.001 | b | ++-------------------------+------+"#; + + for table in ["f05_mito_wc", "f05_metric_wc"] { + let output = execute_sql( + &frontend, + &format!("INSERT INTO {table} (ts, host, val) VALUES (1, 'a', 1), (1000001, 'b', 2)"), + ) + .await + .data; + assert!(matches!(output, OutputData::AffectedRows(2)), "{table}"); + + // Controls: sorting by the time index itself, applying a fetch, or sorting by the cast + // together with another column all keep working. + let output = execute_sql( + &frontend, + &format!("SELECT ts, host FROM {table} ORDER BY ts"), + ) + .await + .data; + check_output_stream(output, rows_by_time_index).await; + + let output = execute_sql( + &frontend, + &format!("SELECT to_timestamp_millis(ts) AS t, host FROM {table} ORDER BY t"), + ) + .await + .data; + check_output_stream(output, rows_by_millis_wrapper).await; + + let output = execute_sql( + &frontend, + &format!("SELECT CAST(ts AS TIMESTAMP(0)) AS t, host FROM {table} ORDER BY t LIMIT 10"), + ) + .await + .data; + check_output_stream(output, rows_by_cast).await; + + let output = execute_sql( + &frontend, + &format!("SELECT CAST(ts AS TIMESTAMP(0)) AS t, host FROM {table} ORDER BY t, host"), + ) + .await + .data; + check_output_stream(output, rows_by_cast).await; + } + + // Ordering by the cast or a function wrapper that casts inside the function returns the same + // rows as the controls above. + for table in ["f05_mito_wc", "f05_metric_wc"] { + for parallelism in [1, 4] { + for (sql, expected) in [ + ( + format!("SELECT CAST(ts AS TIMESTAMP(0)) AS t, host FROM {table} ORDER BY t"), + rows_by_cast, + ), + ( + format!( + "SELECT CAST(ts AS TIMESTAMP(0)) AS t, host FROM {table} ORDER BY t DESC" + ), + rows_by_cast_desc, + ), + ( + format!("SELECT to_timestamp_seconds(ts) AS t, host FROM {table} ORDER BY t"), + rows_by_seconds_wrapper, + ), + ( + format!( + "SELECT to_timestamp_seconds(ts) AS t, host FROM {table} ORDER BY t DESC" + ), + rows_by_seconds_wrapper_desc, + ), + ] { + let output = execute_sql_with_query_parallelism(&frontend, &sql, parallelism) + .await + .data; + check_output_stream(output, expected).await; + } + } + } +} + #[apply(both_instances_cases)] async fn test_show_create_table(instance: Arc) { let frontend = instance.frontend(); diff --git a/tests/cases/distributed/optimizer/windowed_sort_advance.result b/tests/cases/distributed/optimizer/windowed_sort_advance.result index 3c695f3879c..618a1547a14 100644 --- a/tests/cases/distributed/optimizer/windowed_sort_advance.result +++ b/tests/cases/distributed/optimizer/windowed_sort_advance.result @@ -107,7 +107,7 @@ EXPLAIN ANALYZE select to_timestamp_millis(ts) as ts, status, value from `a` whe |_|_|_| | 1_| 0_|_ProjectionExec: expr=[to_timestamp_millis(a.ts)@0 as ts, status@1 as status, value@2 as value] REDACTED |_|_|_SortPreservingMergeExec: [to_timestamp_millis(a.ts)@0 ASC NULLS LAST] REDACTED -|_|_|_WindowedSortExec: expr=to_timestamp_millis(a.ts)@0 ASC NULLS LAST num_ranges=REDACTED REDACTED +|_|_|_SortExec: expr=[to_timestamp_millis(a.ts)@0 ASC NULLS LAST], preserve_partitioning=[true] REDACTED |_|_|_ProjectionExec: expr=[to_timestamp_millis(ts@2) as to_timestamp_millis(a.ts), status@1 as status, value@0 as value] REDACTED |_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["value", "status", "ts"], "filters": ["ts >= TimestampNanosecond(1773273600000000000, None)", "ts < TimestampNanosecond(1773273721000000000, None)"], "time_range":REDACTED |_|_|_| diff --git a/tests/cases/distributed/optimizer/windowed_sort_nanos.result b/tests/cases/distributed/optimizer/windowed_sort_nanos.result index 52d264073ae..122c6e64a29 100644 --- a/tests/cases/distributed/optimizer/windowed_sort_nanos.result +++ b/tests/cases/distributed/optimizer/windowed_sort_nanos.result @@ -67,7 +67,7 @@ EXPLAIN ANALYZE select to_timestamp_nanos(ts) as ts, status, value from `a_ms` w |_|_|_| | 1_| 0_|_ProjectionExec: expr=[to_timestamp_nanos(a_ms.ts)@0 as ts, status@1 as status, value@2 as value] REDACTED |_|_|_SortPreservingMergeExec: [to_timestamp_nanos(a_ms.ts)@0 ASC NULLS LAST] REDACTED -|_|_|_WindowedSortExec: expr=to_timestamp_nanos(a_ms.ts)@0 ASC NULLS LAST num_ranges=REDACTED REDACTED +|_|_|_SortExec: expr=[to_timestamp_nanos(a_ms.ts)@0 ASC NULLS LAST], preserve_partitioning=[true] REDACTED |_|_|_ProjectionExec: expr=[to_timestamp_nanos(ts@2) as to_timestamp_nanos(a_ms.ts), status@1 as status, value@0 as value] REDACTED |_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["value", "status", "ts"], "filters": ["ts >= TimestampMillisecond(1773273600000, None)", "ts < TimestampMillisecond(1773273721000, None)"], "time_range":REDACTED |_|_|_| diff --git a/tests/cases/standalone/optimizer/windowed_sort_advance.result b/tests/cases/standalone/optimizer/windowed_sort_advance.result index f8863341471..123635de8c2 100644 --- a/tests/cases/standalone/optimizer/windowed_sort_advance.result +++ b/tests/cases/standalone/optimizer/windowed_sort_advance.result @@ -106,7 +106,7 @@ EXPLAIN ANALYZE select to_timestamp_millis(ts) as ts, status, value from `a` whe |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_SortPreservingMergeExec: [ts@0 ASC NULLS LAST] REDACTED -|_|_|_WindowedSortExec: expr=ts@0 ASC NULLS LAST num_ranges=REDACTED REDACTED +|_|_|_SortExec: expr=[ts@0 ASC NULLS LAST], preserve_partitioning=[true] REDACTED |_|_|_ProjectionExec: expr=[to_timestamp_millis(ts@2) as ts, status@1 as status, value@0 as value] REDACTED |_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["value", "status", "ts"], "filters": ["ts >= TimestampNanosecond(1773273600000000000, None)", "ts < TimestampNanosecond(1773273721000000000, None)"], "time_range":REDACTED |_|_|_| diff --git a/tests/cases/standalone/optimizer/windowed_sort_nanos.result b/tests/cases/standalone/optimizer/windowed_sort_nanos.result index 9577b711704..3c808912b67 100644 --- a/tests/cases/standalone/optimizer/windowed_sort_nanos.result +++ b/tests/cases/standalone/optimizer/windowed_sort_nanos.result @@ -66,7 +66,7 @@ EXPLAIN ANALYZE select to_timestamp_nanos(ts) as ts, status, value from `a_ms` w |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_SortPreservingMergeExec: [ts@0 ASC NULLS LAST] REDACTED -|_|_|_WindowedSortExec: expr=ts@0 ASC NULLS LAST num_ranges=REDACTED REDACTED +|_|_|_SortExec: expr=[ts@0 ASC NULLS LAST], preserve_partitioning=[true] REDACTED |_|_|_ProjectionExec: expr=[to_timestamp_nanos(ts@2) as ts, status@1 as status, value@0 as value] REDACTED |_|_|_SeqScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "projection": ["value", "status", "ts"], "filters": ["ts >= TimestampMillisecond(1773273600000, None)", "ts < TimestampMillisecond(1773273721000, None)"], "time_range":REDACTED |_|_|_| From 44df6e5dedfa6bddfa6c03ad4dd362b3564b38ce Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Tue, 27 Oct 2026 13:11:44 +0000 Subject: [PATCH 110/239] fix(mito2): coordinate range cache invalidation and singleflight reads (#9935) * fix(mito2): invalidate retired SST range results and coalesce loads Coordinate decoded range loads by their full scan key, bound shared work, and prevent late fills after SST retirement. Invalidate range results for local purge and object-store GC while preserving old snapshot reads. Fixes #9934 Signed-off-by: Lei, HUANG * fix(mito): preserve SST lifetime across index replacement Share retirement and final purge across immutable index metadata views. Keep version-specific GC references until readers release them and match Parquet liveness by file identity. Fixes #9944 Signed-off-by: Lei, HUANG * refactor(mito): list-based SST purge and single FilePurger::remove_file The final purge of a retired SST no longer depends on FileMeta: index files of every version are discovered by listing ({region_dir}/index/{file_id}. prefix), so even when the last live view is a stale one, the purge still removes all newer index files. - SstLifetime is slimmed down to { deleted, file_purger }; the purge now runs in the last FileHandleInner::drop (Arc::strong_count == 1, which is race-free here since only inners of the same file hold the lifetime). - FilePurger is merged back into a single remove_file(file_meta, is_delete) method; is_delete means the file was retired AND this is the last live view. - LocalFilePurger deletes the SST and all its index files via the new delete_file_and_all_indexes helper (list before delete, so a listing failure leaves the file untouched). - Remove the dead index_outdated flag and its accessors, which had no callers in the entire repo. Signed-off-by: Lei, HUANG * fix(mito): elect final SST view release atomically Track metadata views with an atomic counter instead of inspecting the lifetime Arc strong count. Concurrent destructors now elect exactly one final releaser while preserving the single remove_file interface. Add a deterministic concurrent-drop regression covering retired and live files. Signed-off-by: Lei, HUANG * refactor(mito): restore Arc-owned SST cleanup and clarify reference release Restore the lifetime design from 835227dcbc, using the final Arc drop and latest metadata instead of manual view counting and per-SST index listing. Rename FilePurger::remove_file to release_file_ref while retaining purge_file for final cleanup. Preserve and adapt the concurrent view-release regression test. Signed-off-by: Lei, HUANG * refactor(mito): track SST references by file lifetime Register one reference guard per logical SST lifetime and release it on final drop, including normal region close. Remove per-view release_file_ref and index_outdated handling. Protect all index generations while their SST is referenced. Preserve the existing GC snapshot format by expanding the recorded index generation upper bound for older version-aware workers. Update coverage for independent references, normal close, cross-region ownership, concurrent final cleanup and both GC listing modes. Signed-off-by: Lei, HUANG * refactor(mito): release SST references directly on lifetime drop Remove the redundant FileReferenceGuard and let SstLifetime register and unregister its file reference directly. Preserve unconditional reference release on normal close and retirement-only purge. Expose the reference manager through FilePurger and adapt reference manager tests without changing GC or snapshot behavior. Signed-off-by: Lei, HUANG * test(mito): verify SST cleanup with real storage and isolate GC policy cases Replace the recording-only concurrent purge test with filesystem assertions using LocalFilePurger and shared fixtures. Keep intermediate-view retention and retired versus live file cleanup coverage without claiming exhaustive thread interleavings. Separate GC reference-policy cases from reader-release ordering and reuse their setup. Production code remains unchanged. Signed-off-by: Lei, HUANG * style(mito): format SST lifetime changes Signed-off-by: Lei, HUANG * fix(mito2): use crate-rooted range cache imports Signed-off-by: Lei, HUANG * fix(mito2): account for range cache delivery lifetime Keep completed range results charged while watch subscribers or replay streams retain them. Cache entries keep weak delivery references, allowing cache-only ownership to release the delivery reservation. Cache hits reuse or nonblockingly acquire the reservation, falling back when the budget is exhausted. Resize input reservations to completed-result estimates and cover successive slow consumers, cache-only handoff, hit sharing, invalidation, admission rejection, and cancellation with regression tests. Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- src/mito2/src/cache.rs | 107 +++- src/mito2/src/cache/range_cache.rs | 733 +++++++++++++++++++++++++ src/mito2/src/cache/test_util.rs | 26 + src/mito2/src/metrics.rs | 8 + src/mito2/src/read/range_cache.rs | 551 ++++++++++++++++--- src/mito2/src/read/seq_scan.rs | 52 +- src/mito2/src/read/series_candidate.rs | 41 +- src/mito2/src/read/series_reader.rs | 47 +- src/mito2/src/sst/file.rs | 8 +- src/mito2/src/sst/file_purger.rs | 34 +- 10 files changed, 1446 insertions(+), 161 deletions(-) create mode 100644 src/mito2/src/cache/range_cache.rs diff --git a/src/mito2/src/cache.rs b/src/mito2/src/cache.rs index cf0e8de3458..c4b8ffa0d73 100644 --- a/src/mito2/src/cache.rs +++ b/src/mito2/src/cache.rs @@ -19,6 +19,7 @@ pub(crate) mod cache_size; pub(crate) mod file_cache; pub(crate) mod index; pub(crate) mod manifest_cache; +pub(crate) mod range_cache; #[cfg(test)] pub(crate) mod test_util; pub(crate) mod write_cache; @@ -145,6 +146,40 @@ impl RangeResultMemoryLimiter { self.semaphore.available_permits() } + /// Never wait while holding a partial load: several partial loads could + /// otherwise exhaust the budget and wait on each other indefinitely. + pub(crate) fn try_acquire(&self, bytes: usize) -> Option { + let permits = bytes.div_ceil(self.permit_bytes).max(1); + if permits > self.total_permits { + return None; + } + self.semaphore + .clone() + .try_acquire_many_owned(permits.try_into().ok()?) + .ok() + } + + /// Transfers a load's reservation to its completed result without a gap in + /// accounting. Compaction may shrink buffers or add metadata; growing the + /// reservation must remain nonblocking just like loading more input. + pub(crate) fn try_resize( + &self, + permit: Option, + bytes: usize, + ) -> Option { + let Some(mut permit) = permit else { + return self.try_acquire(bytes); + }; + let required = bytes.div_ceil(self.permit_bytes).max(1); + let reserved = permit.num_permits(); + if required > reserved { + permit.merge(self.try_acquire((required - reserved).checked_mul(self.permit_bytes)?)?); + } else { + drop(permit.split(reserved - required)); + } + Some(permit) + } + pub(crate) async fn acquire(&self, bytes: usize) -> Result> { let permits = bytes.div_ceil(self.permit_bytes).max(1); if permits > self.total_permits { @@ -898,6 +933,7 @@ impl CacheStrategy { /// Calls [CacheManager::get_range_result()]. /// It returns None if the strategy is [CacheStrategy::Compaction] or [CacheStrategy::Disabled]. + #[cfg(test)] pub(crate) fn get_range_result( &self, key: &RangeScanCacheKey, @@ -910,6 +946,7 @@ impl CacheStrategy { /// Calls [CacheManager::put_range_result()]. /// It does nothing if the strategy isn't [CacheStrategy::EnableAll]. + #[cfg(any(test, feature = "test"))] pub(crate) fn put_range_result( &self, key: RangeScanCacheKey, @@ -920,6 +957,13 @@ impl CacheStrategy { } } + pub(crate) fn range_cache_worker(&self) -> Option<&Arc> { + match self { + Self::EnableAll(manager) => manager.range_result_cache.as_ref(), + Self::Compaction(_) | Self::Disabled => None, + } + } + /// Returns true if the range result cache is enabled. pub(crate) fn has_range_result_cache(&self) -> bool { match self { @@ -1389,7 +1433,15 @@ impl CacheManager { } } + /// Retires decoded range results when an SST leaves the current version. + pub(crate) fn invalidate_range_results(&self, file_id: FileId) { + if let Some(worker) = &self.range_result_cache { + worker.invalidate_file(file_id); + } + } + /// Gets cached result for range scan. + #[cfg(test)] pub(crate) fn get_range_result( &self, key: &RangeScanCacheKey, @@ -1400,15 +1452,13 @@ impl CacheManager { } /// Puts range scan result into cache. + #[cfg(any(test, feature = "test"))] pub(crate) fn put_range_result( &self, key: RangeScanCacheKey, result: Arc, ) { if let Some(cache) = &self.range_result_cache { - CACHE_BYTES - .with_label_values(&[RANGE_RESULT_TYPE]) - .add(range_result_cache_weight(&key, &result).into()); cache.insert(key, result); } } @@ -1649,21 +1699,8 @@ impl CacheManagerBuilder { }) .build() }); - let range_result_cache = (self.range_result_cache_size != 0).then(|| { - Cache::builder() - .max_capacity(self.range_result_cache_size) - .weigher(range_result_cache_weight) - .eviction_listener(move |k, v, cause| { - let size = range_result_cache_weight(&k, &v); - CACHE_BYTES - .with_label_values(&[RANGE_RESULT_TYPE]) - .sub(size.into()); - CACHE_EVICTION - .with_label_values(&[RANGE_RESULT_TYPE, removal_cause_str(cause)]) - .inc(); - }) - .build() - }); + let range_result_cache = (self.range_result_cache_size != 0) + .then(|| range_cache::RangeCacheWorker::new(self.range_result_cache_size)); CacheManager { sst_meta_cache, sst_decoded_meta_cache, @@ -2093,7 +2130,7 @@ type VectorCache = Cache<(ConcreteDataType, Value), VectorRef>; /// Maps (file id, row group id, time series row selector) to [SelectorResultValue]. type SelectorResultCache = Cache>; /// Maps partition-range scan key to cached flat batches. -type RangeResultCache = Cache>; +type RangeResultCache = Arc; #[cfg(test)] mod tests { @@ -2900,6 +2937,38 @@ mod tests { assert_eq!(limiter.available_permits(), 2); } + #[rstest::rstest] + #[case::empty_result(None, 0, Some(1))] + #[case::new_result(None, 1025, Some(2))] + #[case::unchanged(Some(1024), 1024, Some(1))] + #[case::grow(Some(1024), 2048, Some(2))] + #[case::shrink(Some(3072), 1024, Some(1))] + #[case::exhausted(Some(1024), 4096, None)] + #[case::oversized(Some(1024), usize::MAX, None)] + #[test] + fn range_result_memory_reservation_resizes_without_waiting( + #[case] input_size: Option, + #[case] result_size: usize, + #[case] expected_permits: Option, + ) { + let limiter = RangeResultMemoryLimiter::new(4096, 1024); + let blocker = limiter.try_acquire(1024).unwrap(); + let input_permit = input_size.map(|size| limiter.try_acquire(size).unwrap()); + let result_permit = limiter.try_resize(input_permit, result_size); + assert_eq!( + result_permit.as_ref().map(|permit| permit.num_permits()), + expected_permits + ); + assert_eq!( + limiter.available_permits(), + 3 - expected_permits.unwrap_or(0) + ); + drop(result_permit); + assert_eq!(limiter.available_permits(), 3); + drop(blocker); + assert_eq!(limiter.available_permits(), 4); + } + #[tokio::test] async fn test_evict_puffin_cache_clears_all_entries() { use std::collections::{BTreeMap, HashMap}; diff --git a/src/mito2/src/cache/range_cache.rs b/src/mito2/src/cache/range_cache.rs new file mode 100644 index 00000000000..060af974764 --- /dev/null +++ b/src/mito2/src/cache/range_cache.rs @@ -0,0 +1,733 @@ +// 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. + +//! Coordinates range-result cache loads and SST retirement. No I/O or decoding +//! runs under the coordinator lock; loaders use the shared query runtime. + +use std::collections::HashMap; +use std::future::Future; +use std::sync::{Arc, Mutex, Weak, mpsc}; + +use moka::sync::Cache; +use store_api::storage::FileId; +use tokio::sync::{OwnedSemaphorePermit, watch}; + +use crate::cache::{ + RANGE_RESULT_TYPE, RangeResultMemoryLimiter, range_result_cache_weight, removal_cause_str, +}; +use crate::metrics::{CACHE_BYTES, CACHE_EVICTION, RANGE_CACHE_WORKER_EVENTS}; +use crate::read::range_cache::{RangeScanCacheKey, RangeScanCacheValue}; +use crate::sst::file::FileHandle; + +/// Admission, not a queue: excess distinct misses keep using ordinary scans. +const MAX_IN_FLIGHT_RANGES: usize = 32; + +type SharedResult = Option>; + +/// Keeps completed results charged while subscribers or replay streams own them. +/// Cache entries retain only a weak reference: after the last reader leaves, +/// the resident value is covered solely by the cache's capacity accounting. +pub(crate) struct RangeCacheDelivery { + value: Arc, + _permit: OwnedSemaphorePermit, +} + +impl RangeCacheDelivery { + /// Shares a decoded result together with its memory reservation. + pub(crate) fn new(value: Arc, permit: OwnedSemaphorePermit) -> Arc { + Arc::new(Self { + value, + _permit: permit, + }) + } + + /// Borrows the result without detaching it from the delivery reservation. + pub(crate) fn value(&self) -> &RangeScanCacheValue { + &self.value + } +} + +pub(crate) enum RangeCacheLookup { + Hit(Arc), + Loading(watch::Receiver), + Bypass, +} + +struct Entry { + generation: u64, + value: Arc, + delivery: Mutex>, +} + +impl Entry { + fn deliver(&self, memory: &RangeResultMemoryLimiter) -> SharedResult { + // Serialize cache-hit reservations so readers of the same result share + // one charge, including when the last previous reader just left. + let mut shared_delivery = self + .delivery + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(delivery) = shared_delivery.upgrade() { + return Some(delivery); + } + let permit = memory.try_acquire(self.value.estimated_size())?; + let delivery = RangeCacheDelivery::new(self.value.clone(), permit); + *shared_delivery = Arc::downgrade(&delivery); + Some(delivery) + } +} + +struct Flight { + generation: u64, + sender: watch::Sender, + invalidated: bool, +} + +#[derive(Default)] +struct State { + generation: u64, + flights: HashMap, + /// File IDs are globally unique. Index by physical file, not query region: + /// a copied region can read an SST whose original region differs from its own. + files: HashMap, u64>>, +} + +impl State { + fn next_generation(&mut self) -> u64 { + self.generation += 1; + self.generation + } + + fn forget(&mut self, key: &RangeScanCacheKey, generation: u64) { + for (file, _) in &key.row_groups { + if let Some(keys) = self.files.get_mut(file) { + // Eviction callbacks may arrive after the same key is refilled. + if keys.get(key) == Some(&generation) { + keys.remove(key); + } + if keys.is_empty() { + self.files.remove(file); + } + } + } + } +} + +pub(crate) struct RangeCacheWorker { + cache: Cache>, + state: Mutex, + /// Moka can call listeners during a cache operation. Queue only cleanup + /// metadata instead of reentering `state` from its eviction listener. + evictions: Mutex, u64)>>, + /// Covers loads and completed deliveries, separately from cache residency + /// and the temporary concat-memory limiter. Never wait for this budget. + pub(crate) load_memory: RangeResultMemoryLimiter, +} + +impl RangeCacheWorker { + pub(crate) fn new(capacity: u64) -> Arc { + let (tx, rx) = mpsc::channel(); + let cache = Cache::builder() + .max_capacity(capacity) + .weigher(|k: &RangeScanCacheKey, v: &Arc| range_result_cache_weight(k, &v.value)) + .eviction_listener(move |key, entry: Arc, cause| { + CACHE_BYTES + .with_label_values(&[RANGE_RESULT_TYPE]) + .sub(range_result_cache_weight(&key, &entry.value).into()); + CACHE_EVICTION + .with_label_values(&[RANGE_RESULT_TYPE, removal_cause_str(cause)]) + .inc(); + let _ = tx.send((key, entry.generation)); + }) + .build(); + Arc::new(Self { + cache, + state: Mutex::new(State::default()), + evictions: Mutex::new(rx), + load_memory: RangeResultMemoryLimiter::new(capacity as usize, 1024), + }) + } + + fn drain_evictions(&self, state: &mut State) { + for (key, generation) in self.evictions.lock().unwrap().try_iter() { + state.forget(&key, generation); + } + } + + #[cfg(test)] + pub(crate) fn get(&self, key: &RangeScanCacheKey) -> Option> { + self.cache.get(key).map(|entry| entry.value.clone()) + } + + fn get_for_delivery(&self, key: &RangeScanCacheKey) -> Option { + let entry = self.cache.get(key)?; + // A hit without delivery budget must bypass, not start another load. + Some(match entry.deliver(&self.load_memory) { + Some(delivery) => RangeCacheLookup::Hit(delivery), + None => { + RANGE_CACHE_WORKER_EVENTS + .with_label_values(&["bypass_memory"]) + .inc(); + RangeCacheLookup::Bypass + } + }) + } + + /// For test/benchmark producers that do not hold live SST handles. + #[cfg(any(test, feature = "test"))] + pub(crate) fn insert(&self, key: RangeScanCacheKey, value: Arc) { + let Some(permit) = self.load_memory.try_acquire(value.estimated_size()) else { + return; + }; + let delivery = RangeCacheDelivery::new(value, permit); + let mut state = self.state.lock().unwrap(); + self.drain_evictions(&mut state); + self.publish(&mut state, key, &delivery); + } + + fn publish( + &self, + state: &mut State, + key: RangeScanCacheKey, + delivery: &Arc, + ) { + let generation = state.next_generation(); + let indexed_key = Arc::new(key.clone()); + for (file, _) in &key.row_groups { + state + .files + .entry(*file) + .or_default() + .insert(indexed_key.clone(), generation); + } + CACHE_BYTES + .with_label_values(&[RANGE_RESULT_TYPE]) + .add(range_result_cache_weight(&key, &delivery.value).into()); + self.cache.insert( + key, + Arc::new(Entry { + generation, + value: delivery.value.clone(), + delivery: Mutex::new(Arc::downgrade(delivery)), + }), + ); + self.drain_evictions(state); + } + + /// Called at logical retirement, before the last snapshot releases the SST. + /// Existing readers may finish, but their results must not refill this cache. + pub(crate) fn invalidate_file(&self, file: FileId) { + RANGE_CACHE_WORKER_EVENTS + .with_label_values(&["invalidate"]) + .inc(); + let mut state = self.state.lock().unwrap(); + self.drain_evictions(&mut state); + if let Some(keys) = state.files.remove(&file) { + for (key, generation) in keys { + self.cache.invalidate(key.as_ref()); + state.forget(&key, generation); + } + } + for (key, flight) in &mut state.flights { + if key.row_groups.iter().any(|(id, _)| *id == file) { + flight.invalidated = true; + } + } + self.drain_evictions(&mut state); + } + + pub(crate) fn get_or_load( + self: &Arc, + key: RangeScanCacheKey, + files: Vec, + load: F, + ) -> RangeCacheLookup + where + F: Future + Send + 'static, + { + if files.iter().any(FileHandle::is_deleted) { + RANGE_CACHE_WORKER_EVENTS + .with_label_values(&["bypass_retired"]) + .inc(); + return RangeCacheLookup::Bypass; + } + if let Some(lookup) = + crate::cache::update_hit_miss(self.get_for_delivery(&key), RANGE_RESULT_TYPE) + { + return lookup; + } + let mut state = self.state.lock().unwrap(); + self.drain_evictions(&mut state); + // Close the miss/register race with completion and invalidation. + if files.iter().any(FileHandle::is_deleted) { + RANGE_CACHE_WORKER_EVENTS + .with_label_values(&["bypass_retired"]) + .inc(); + return RangeCacheLookup::Bypass; + } + if let Some(lookup) = self.get_for_delivery(&key) { + return lookup; + } + if let Some(flight) = state.flights.get(&key) { + RANGE_CACHE_WORKER_EVENTS + .with_label_values(&["coalesced"]) + .inc(); + return RangeCacheLookup::Loading(flight.sender.subscribe()); + } + if state.flights.len() >= MAX_IN_FLIGHT_RANGES { + RANGE_CACHE_WORKER_EVENTS + .with_label_values(&["bypass_capacity"]) + .inc(); + return RangeCacheLookup::Bypass; + } + let generation = state.next_generation(); + let (sender, receiver) = watch::channel(None); + state.flights.insert( + key.clone(), + Flight { + generation, + sender: sender.clone(), + invalidated: false, + }, + ); + drop(state); + + RANGE_CACHE_WORKER_EVENTS + .with_label_values(&["started"]) + .inc(); + let guard = LoadGuard { + worker: self.clone(), + key, + generation, + }; + common_runtime::spawn_query(async move { + // Losing the first requester does not cancel other subscribers. + // Losing every requester drops the loader and unregisters the flight. + let value = tokio::select! { + biased; + _ = sender.closed() => return, + value = load => value, + }; + guard.complete(value, &files); + }); + RangeCacheLookup::Loading(receiver) + } +} + +/// Also unregisters a cancelled or panicking load. Generation matching keeps a +/// stale task from unregistering a later retry of the same key. +struct LoadGuard { + worker: Arc, + key: RangeScanCacheKey, + generation: u64, +} + +impl LoadGuard { + fn complete(&self, value: SharedResult, files: &[FileHandle]) { + let mut state = self.worker.state.lock().unwrap(); + self.worker.drain_evictions(&mut state); + let Some(flight) = state.flights.get(&self.key) else { + return; + }; + if flight.generation != self.generation { + return; + } + let flight = state.flights.remove(&self.key).unwrap(); + let event = if let Some(value) = &value { + if !flight.invalidated && !files.iter().any(FileHandle::is_deleted) { + self.worker.publish(&mut state, self.key.clone(), value); + "published" + } else { + "late_fill_skipped" + } + } else { + "load_skipped" + }; + RANGE_CACHE_WORKER_EVENTS.with_label_values(&[event]).inc(); + // Fan out independently of Moka admission, including to old snapshots. + // Do not hold the coordinator lock while notifying readers. + drop(state); + let _ = flight.sender.send(value); + } +} + +impl Drop for LoadGuard { + fn drop(&mut self) { + let mut state = self.worker.state.lock().unwrap(); + if state + .flights + .get(&self.key) + .is_some_and(|flight| flight.generation == self.generation) + { + state.flights.remove(&self.key); + RANGE_CACHE_WORKER_EVENTS + .with_label_values(&["abandoned"]) + .inc(); + } + } +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::time::Duration; + + use tokio::sync::oneshot; + + use super::*; + use crate::cache::test_util::range_cache_key; + use crate::sst::file::FileMeta; + use crate::sst::file_purger::NoopFilePurger; + + fn empty_result() -> Arc { + Arc::new(RangeScanCacheValue::new(vec![], 0)) + } + + fn empty_delivery() -> Arc { + let value = empty_result(); + let memory = RangeResultMemoryLimiter::new(1024, 1024); + let permit = memory.try_acquire(value.estimated_size()).unwrap(); + RangeCacheDelivery::new(value, permit) + } + + fn receiver(lookup: RangeCacheLookup) -> watch::Receiver { + match lookup { + RangeCacheLookup::Loading(receiver) => receiver, + _ => panic!("expected a shared load"), + } + } + + async fn loaded(mut receiver: watch::Receiver) -> SharedResult { + tokio::time::timeout(Duration::from_secs(5), async { + receiver + .wait_for(|value| value.is_some()) + .await + .ok() + .and_then(|value| value.clone()) + }) + .await + .expect("load must finish or cancel") + } + + #[tokio::test] + async fn same_key_misses_share_one_load_even_after_first_reader_leaves() { + let worker = RangeCacheWorker::new(1024 * 1024); + let key = range_cache_key(&[FileId::random()]); + let count = Arc::new(AtomicUsize::new(0)); + let (release, gate) = oneshot::channel(); + let first_count = count.clone(); + let first = receiver(worker.get_or_load(key.clone(), vec![], async move { + first_count.fetch_add(1, Ordering::Relaxed); + gate.await.unwrap(); + Some(empty_delivery()) + })); + let mut followers = Vec::new(); + for _ in 0..16 { + let count = count.clone(); + followers.push(receiver(worker.get_or_load( + key.clone(), + vec![], + async move { + count.fetch_add(1, Ordering::Relaxed); + Some(empty_delivery()) + }, + ))); + } + drop(first); + release.send(()).unwrap(); + let mut results = Vec::new(); + for follower in followers { + results.push(loaded(follower).await.unwrap()); + } + assert_eq!(count.load(Ordering::Relaxed), 1); + assert!(results.iter().all(|value| Arc::ptr_eq(value, &results[0]))); + assert!(Arc::ptr_eq(&worker.get(&key).unwrap(), &results[0].value)); + } + + #[test] + fn retirement_removes_all_dependent_keys_but_preserves_other_files() { + let worker = RangeCacheWorker::new(1024 * 1024); + let a = FileId::random(); + let b = FileId::random(); + let combined = range_cache_key(&[a, b, a]); + let mut copied = range_cache_key(&[a]); + copied.region_id = store_api::storage::RegionId::new(2, 1); + let unaffected = range_cache_key(&[b]); + for key in [&combined, &copied, &unaffected] { + worker.insert(key.clone(), empty_result()); + } + worker.invalidate_file(a); + worker.invalidate_file(a); + assert!(worker.get(&combined).is_none()); + assert!(worker.get(&copied).is_none()); + assert!(worker.get(&unaffected).is_some()); + worker.invalidate_file(b); + assert!(worker.get(&unaffected).is_none()); + assert!(worker.state.lock().unwrap().files.is_empty()); + } + + #[test] + fn delayed_replacement_eviction_cannot_hide_a_new_entry_from_invalidation() { + let worker = RangeCacheWorker::new(1024 * 1024); + let file = FileId::random(); + let key = range_cache_key(&[file]); + worker.insert(key.clone(), empty_result()); + worker.cache.run_pending_tasks(); + let replacement = empty_result(); + worker.insert(key.clone(), replacement.clone()); + worker.cache.run_pending_tasks(); + assert!(Arc::ptr_eq(&worker.get(&key).unwrap(), &replacement)); + worker.invalidate_file(file); + assert!(worker.get(&key).is_none()); + } + + #[tokio::test] + async fn retired_inflight_result_serves_old_readers_without_refilling() { + let worker = RangeCacheWorker::new(1024 * 1024); + let file = FileId::random(); + let key = range_cache_key(&[file]); + let (release, gate) = oneshot::channel(); + let pending = receiver(worker.get_or_load(key.clone(), vec![], async move { + gate.await.unwrap(); + Some(empty_delivery()) + })); + worker.invalidate_file(file); + release.send(()).unwrap(); + assert!(loaded(pending).await.is_some()); + assert!(worker.get(&key).is_none()); + } + + #[test] + fn old_snapshot_cannot_start_a_fill_after_retirement() { + let worker = RangeCacheWorker::new(1024 * 1024); + let file = FileHandle::new(FileMeta::default(), Arc::new(NoopFilePurger)); + file.mark_deleted(); + let key = range_cache_key(&[file.file_id().file_id()]); + assert!(matches!( + worker.get_or_load(key, vec![file], async { Some(empty_delivery()) }), + RangeCacheLookup::Bypass + )); + } + + #[rstest::rstest] + #[tokio::test] + async fn index_replacement_retirement_invalidates_and_blocks_old_reader_fills( + #[values(false, true)] gc_enabled: bool, + ) { + use common_base::readable_size::ReadableSize; + use store_api::region_engine::RegionEngine; + use store_api::region_request::RegionRequest; + use store_api::storage::RegionId; + + use crate::config::MitoConfig; + use crate::gc::GcConfig; + use crate::sst::file::IndexType; + use crate::sst::version::SstVersion; + use crate::test_util::{CreateRequestBuilder, TestEnv}; + + let mut env = TestEnv::new().await; + let engine = env + .create_engine(MitoConfig { + range_result_cache_size: ReadableSize::mb(1), + gc: GcConfig { + enable: gc_enabled, + ..Default::default() + }, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + engine + .handle_request( + region_id, + RegionRequest::Create(CreateRequestBuilder::new().build()), + ) + .await + .unwrap(); + let region = engine.get_region(region_id).unwrap(); + let cache = engine.cache_manager(); + let worker = cache.range_result_cache.as_ref().unwrap(); + let meta = FileMeta { + region_id, + file_id: FileId::random(), + ..Default::default() + }; + let mut version = SstVersion::new(crate::test_util::memtable_util::metadata_for_test()); + version.add_files(region.file_purger.clone(), [meta.clone()].into_iter()); + let old_reader = version.levels()[0].files[&meta.file_id].clone(); + let key = range_cache_key(&[meta.file_id]); + worker.insert(key.clone(), empty_result()); + assert!(worker.get(&key).is_some()); + let loading_key = key.clone().for_cross_range_merge(); + let (release, gate) = oneshot::channel(); + let pending = receiver(worker.get_or_load( + loading_key.clone(), + vec![old_reader.clone()], + async move { + gate.await.unwrap(); + Some(empty_delivery()) + }, + )); + let updated = FileMeta { + index_version: 1, + available_indexes: smallvec::smallvec![IndexType::InvertedIndex], + ..meta + }; + version.add_files(region.file_purger.clone(), [updated.clone()].into_iter()); + version.remove_files([updated].into_iter()); + + // R1-C1: neither a pre-existing flight nor a new request through the old + // index handle may refill a range after the current handle is retired. + assert!(old_reader.is_deleted()); + assert!(worker.get(&key).is_none()); + assert!(matches!( + worker.get_or_load(key, vec![old_reader], async { Some(empty_delivery()) }), + RangeCacheLookup::Bypass + )); + release.send(()).unwrap(); + assert!(loaded(pending).await.is_some()); + assert!(worker.get(&loading_key).is_none()); + } + + #[tokio::test] + async fn admission_rejection_does_not_discard_the_shared_result() { + let worker = RangeCacheWorker::new(1); + let key = range_cache_key(&[FileId::random()]); + let value = empty_result(); + let permit = worker + .load_memory + .try_acquire(value.estimated_size()) + .unwrap(); + let delivery = RangeCacheDelivery::new(value, permit); + let pending = + receiver(worker.get_or_load(key.clone(), vec![], async move { Some(delivery) })); + let delivery = loaded(pending).await.unwrap(); + worker.cache.run_pending_tasks(); + assert!(worker.get(&key).is_none()); + { + let mut state = worker.state.lock().unwrap(); + worker.drain_evictions(&mut state); + assert!(state.files.is_empty()); + } + assert_eq!(worker.load_memory.available_permits(), 0); + drop(delivery); + // Publishing wakes readers before the producer drops its watch sender. + // Its LoadGuard also needs the coordinator lock to finish cleaning up. + tokio::time::timeout(Duration::from_secs(5), async { + while worker.load_memory.available_permits() != 1 { + tokio::task::yield_now().await; + } + }) + .await + .expect("a rejected value must release its reservation after all owners leave"); + } + + #[test] + fn cache_hits_share_delivery_budget_until_the_last_reader_leaves() { + let worker = RangeCacheWorker::new(1024); + let file = FileId::random(); + let key = range_cache_key(&[file]); + worker.insert(key.clone(), empty_result()); + assert_eq!(worker.load_memory.available_permits(), 1); + + let blocker = worker.load_memory.try_acquire(1024).unwrap(); + assert!(matches!( + worker.get_or_load(key.clone(), vec![], async { + panic!("must not load on a hit") + }), + RangeCacheLookup::Bypass + )); + assert!(worker.get(&key).is_some()); + drop(blocker); + + let RangeCacheLookup::Hit(first) = + worker.get_or_load(key.clone(), vec![], async { panic!("must hit the cache") }) + else { + panic!("expected an accounted cache hit"); + }; + let RangeCacheLookup::Hit(second) = + worker.get_or_load(key.clone(), vec![], async { panic!("must reuse delivery") }) + else { + panic!("readers of the same value must share its reservation"); + }; + assert!(Arc::ptr_eq(&first, &second)); + assert_eq!(worker.load_memory.available_permits(), 0); + + worker.invalidate_file(file); + worker.cache.run_pending_tasks(); + assert!(worker.get(&key).is_none()); + drop(first); + assert_eq!(worker.load_memory.available_permits(), 0); + drop(second); + assert_eq!(worker.load_memory.available_permits(), 1); + } + + #[tokio::test] + async fn failed_load_releases_followers_and_allows_a_retry() { + let worker = RangeCacheWorker::new(1024 * 1024); + let key = range_cache_key(&[FileId::random()]); + let failed = receiver(worker.get_or_load(key.clone(), vec![], async { None })); + assert!(loaded(failed).await.is_none()); + let retry = receiver(worker.get_or_load(key, vec![], async { Some(empty_delivery()) })); + assert!(loaded(retry).await.is_some()); + } + + #[tokio::test] + async fn abandoning_all_subscribers_cancels_and_unregisters_the_load() { + let worker = RangeCacheWorker::new(1024 * 1024); + let key = range_cache_key(&[FileId::random()]); + let (started, start_rx) = oneshot::channel(); + let marker = Arc::new(()); + let weak_marker = Arc::downgrade(&marker); + let pending = receiver(worker.get_or_load(key.clone(), vec![], async move { + let _marker = marker; + started.send(()).unwrap(); + std::future::pending().await + })); + start_rx.await.unwrap(); + drop(pending); + tokio::time::timeout(Duration::from_secs(5), async { + while weak_marker.upgrade().is_some() + || !worker.state.lock().unwrap().flights.is_empty() + { + tokio::task::yield_now().await; + } + }) + .await + .expect("cancelled loaders must release resources and the flight key"); + assert!(worker.get(&key).is_none()); + let retry = receiver(worker.get_or_load(key, vec![], async { Some(empty_delivery()) })); + assert!(loaded(retry).await.is_some()); + } + + #[tokio::test] + async fn distinct_loads_are_bounded_without_queueing_misses() { + let worker = RangeCacheWorker::new(1024 * 1024); + let mut subscribers = Vec::new(); + for _ in 0..MAX_IN_FLIGHT_RANGES { + subscribers.push(receiver(worker.get_or_load( + range_cache_key(&[FileId::random()]), + vec![], + std::future::pending(), + ))); + } + assert!(matches!( + worker.get_or_load(range_cache_key(&[FileId::random()]), vec![], async { + Some(empty_delivery()) + }), + RangeCacheLookup::Bypass + )); + drop(subscribers); + } +} diff --git a/src/mito2/src/cache/test_util.rs b/src/mito2/src/cache/test_util.rs index 58b8272fbc7..82cabe24297 100644 --- a/src/mito2/src/cache/test_util.rs +++ b/src/mito2/src/cache/test_util.rs @@ -145,3 +145,29 @@ pub(crate) fn assert_parquet_metadata_equal(x: Arc, y: Arc crate::read::range_cache::RangeScanCacheKey { + use crate::read::range_cache::{RangeScanCacheKey, ScanRequestFingerprintBuilder}; + use crate::read::read_columns::ReadColumns; + + RangeScanCacheKey { + region_id: store_api::storage::RegionId::new(1, 1), + row_groups: files.iter().map(|file| (*file, 0)).collect(), + scan: ScanRequestFingerprintBuilder { + read_columns: ReadColumns::new(std::iter::empty()), + read_column_types: vec![], + filters: vec![], + time_filters: vec![], + series_row_selector: None, + append_mode: false, + filter_deleted: true, + merge_mode: crate::region::options::MergeMode::LastRow, + sequence_range: None, + partition_expr_version: 0, + } + .build(), + } +} diff --git a/src/mito2/src/metrics.rs b/src/mito2/src/metrics.rs index e8718f0e758..f7af9d54ffc 100644 --- a/src/mito2/src/metrics.rs +++ b/src/mito2/src/metrics.rs @@ -319,6 +319,14 @@ lazy_static! { "mito upload bytes total", ) .unwrap(); + /// Bounded event labels distinguish lookup misses from actual shared loads. + pub static ref RANGE_CACHE_WORKER_EVENTS: IntCounterVec = register_int_counter_vec!( + "greptime_mito_range_cache_worker_events_total", + "Range cache load coordination and SST retirement events", + &["event"] + ) + .unwrap(); + /// Cache eviction counter, labeled with cache type and eviction reason. pub static ref CACHE_EVICTION: IntCounterVec = register_int_counter_vec!( "greptime_mito_cache_eviction", diff --git a/src/mito2/src/read/range_cache.rs b/src/mito2/src/read/range_cache.rs index 081efa7bec8..a3b857fc24d 100644 --- a/src/mito2/src/read/range_cache.rs +++ b/src/mito2/src/read/range_cache.rs @@ -14,6 +14,7 @@ //! Utilities for the partition range scan result cache. +use std::future::Future; use std::mem; use std::sync::Arc; @@ -36,6 +37,7 @@ use table::predicate::is_string_timestamp_literal; use tokio::sync::{mpsc, oneshot}; use crate::cache::CacheStrategy; +use crate::cache::range_cache::{RangeCacheDelivery, RangeCacheLookup, RangeCacheWorker}; use crate::error::{ComputeArrowSnafu, Result}; use crate::read::BoxedRecordBatchStream; use crate::read::read_columns::ReadColumns; @@ -608,9 +610,14 @@ fn build_range_cache_key_inner( } /// Returns a stream that replays cached record batches. -pub(crate) fn cached_flat_range_stream(value: Arc) -> BoxedRecordBatchStream { +pub(crate) fn cached_flat_range_stream( + delivery: Arc, +) -> BoxedRecordBatchStream { Box::pin(try_stream! { - for cached_batch in &value.cached_batches { + // Borrow through the delivery guard so it survives every yield. Batches + // handed to downstream operators are then owned by the query, not by + // this completed-result backlog. + for cached_batch in &delivery.value().cached_batches { let mut offset = 0; for &len in &cached_batch.slice_lengths { yield cached_batch.batch.slice(offset, len); @@ -620,14 +627,147 @@ pub(crate) fn cached_flat_range_stream(value: Arc) -> Boxed }) } +/// Builds only one decoded result per full cache key. A load must be driven +/// independently of any caller: merge readers can wait for another stream with +/// the same key before polling their first stream to completion. +pub(crate) struct CachedRangeReader<'a> { + pub(crate) stream_ctx: &'a StreamContext, + pub(crate) part_range: &'a PartitionRange, + pub(crate) key: Option, + pub(crate) metrics: &'a PartitionMetrics, +} + +impl CachedRangeReader<'_> { + pub(crate) async fn build(self, build: F) -> Result<(BoxedRecordBatchStream, usize)> + where + F: FnOnce() -> Fut + Clone + Send + 'static, + Fut: Future> + Send + 'static, + { + let Self { + stream_ctx, + part_range, + key, + metrics, + } = self; + let strategy = &stream_ctx.input.cache_strategy; + let (Some(key), Some(worker)) = (key, strategy.range_cache_worker()) else { + return build().await; + }; + let files = stream_ctx.ranges[part_range.identifier] + .row_group_indices + .iter() + .map(|index| stream_ctx.input.file_from_index(*index).clone()) + .collect(); + let loader = build.clone(); + let load_worker = worker.clone(); + let load_strategy = strategy.clone(); + let load_metrics = metrics.clone(); + let key_size = key.estimated_size(); + match worker.get_or_load(key, files, async move { + let result = async { + let (stream, _) = loader().await?; + collect_range_result(stream, &load_worker, &load_strategy).await + } + .await; + match result { + Ok(Some(value)) => { + load_metrics.inc_range_cache_size(key_size + value.value().estimated_size()); + Some(value) + } + Ok(None) => None, + Err(error) => { + warn!(error; "Failed to load range cache; falling back to an uncached scan"); + None + } + } + }) { + RangeCacheLookup::Hit(value) => { + metrics.inc_range_cache_hit(); + Ok((cached_flat_range_stream(value), DEFAULT_READ_BATCH_SIZE)) + } + RangeCacheLookup::Loading(receiver) => { + metrics.inc_range_cache_miss(); + Ok(( + shared_range_stream(receiver, build), + DEFAULT_READ_BATCH_SIZE, + )) + } + RangeCacheLookup::Bypass => { + metrics.inc_range_cache_miss(); + build().await + } + } + } +} + +/// Waiting is lazy, but the producer is independent of every consumer. +fn shared_range_stream( + mut receiver: tokio::sync::watch::Receiver>>, + build: F, +) -> BoxedRecordBatchStream +where + F: FnOnce() -> Fut + Send + 'static, + Fut: Future> + Send + 'static, +{ + Box::pin(try_stream! { + let cached = receiver.wait_for(|cached| cached.is_some()).await + .ok().and_then(|cached| cached.clone()); + drop(receiver); + let mut stream = if let Some(cached) = cached { + cached_flat_range_stream(cached) + } else { + // A failed/oversized/cancelled fill is never a query result. + // Retry once through the ordinary streaming path, not the worker. + build().await?.0 + }; + while let Some(batch) = stream.try_next().await? { + yield batch; + } + }) +} + +/// The load budget covers retained input buffers and completed deliveries. Acquisition +/// must be nonblocking: waiting with a partially accumulated result deadlocks +/// when several loads consume all permits. Oversized ranges retain streaming +/// query semantics by falling back to an uncached scan. +async fn collect_range_result( + mut stream: BoxedRecordBatchStream, + worker: &RangeCacheWorker, + strategy: &CacheStrategy, +) -> Result>> { + let mut permit: Option = None; + let mut buffer = CacheBatchBuffer::new(strategy); + while let Some(batch) = stream.try_next().await? { + let Some(next_permit) = worker + .load_memory + .try_acquire(batch.get_array_memory_size()) + else { + return Ok(None); + }; + if let Some(permit) = &mut permit { + permit.merge(next_permit); + } else { + permit = Some(next_permit); + } + buffer.push(batch)?; + } + let Some(value) = buffer.finish().await.transpose()? else { + return Ok(None); + }; + let Some(permit) = worker + .load_memory + .try_resize(permit, value.estimated_size()) + else { + return Ok(None); + }; + Ok(Some(RangeCacheDelivery::new(value, permit))) +} + enum CacheConcatCommand { Compact(Vec), Finish { pending: Vec, - key: RangeScanCacheKey, - cache_strategy: CacheStrategy, - part_metrics: PartitionMetrics, - result_tx: Option>>>, + result_tx: oneshot::Sender>>, }, } @@ -708,13 +848,7 @@ async fn run_cache_concat_task( return; } } - CacheConcatCommand::Finish { - pending, - key, - cache_strategy, - part_metrics, - result_tx, - } => { + CacheConcatCommand::Finish { pending, result_tx } => { let compact_result = state .compact(pending, &limiter) .await @@ -722,10 +856,6 @@ async fn run_cache_concat_task( let result = match compact_result { Ok(v) => { let value = Arc::new(v); - part_metrics - .inc_range_cache_size(key.estimated_size() + value.estimated_size()); - cache_strategy.put_range_result(key, value.clone()); - Ok(value) } Err(e) => { @@ -735,9 +865,7 @@ async fn run_cache_concat_task( } }; - if let Some(tx) = result_tx { - let _ = tx.send(result); - } + let _ = result_tx.send(result); break; } @@ -750,8 +878,7 @@ struct CacheBatchBuffer { buffered_rows: usize, buffered_size: usize, sender: Option>, - /// Handle of the background concat task. Cleared once the task owns a finish - /// command, so only unfinished scans cancel it. + /// Kept until concat finishes so abandoning the shared load cancels concat too. concat_task: Option>, } @@ -816,39 +943,26 @@ impl CacheBatchBuffer { } } - fn finish( - mut self, - key: RangeScanCacheKey, - cache_strategy: CacheStrategy, - part_metrics: PartitionMetrics, - result_tx: Option>>>, - ) { - let Some(sender) = self.sender.take() else { - return; - }; - - if sender + async fn finish(mut self) -> Option>> { + let sender = self.sender.take()?; + let (result_tx, result_rx) = oneshot::channel(); + sender .send(CacheConcatCommand::Finish { pending: mem::take(&mut self.buffered_batches), - key, - cache_strategy, - part_metrics, result_tx, }) - .is_ok() - { - // The task now owns the finish command, so it may keep populating the - // cache after the scan stream is dropped. - self.concat_task = None; - } + .ok()?; + // Keep ownership until the result arrives so cancellation also aborts + // concat while it waits for its temporary memory permits. + let result = result_rx.await.ok(); + self.concat_task = None; + result } } impl Drop for CacheBatchBuffer { fn drop(&mut self) { - // Still holding the handle means no finish command was sent: the scan was - // cancelled or failed and the queued batches can never reach the cache. - // Aborting releases them even while the task waits for a memory permit. + // Aborting releases queued batches even while concat waits for a permit. if let Some(task) = &self.concat_task { task.abort(); } @@ -856,6 +970,7 @@ impl Drop for CacheBatchBuffer { } /// Wraps a stream to cache its output for future range cache hits. +#[cfg(feature = "test")] pub(crate) fn cache_flat_range_stream( mut stream: BoxedRecordBatchStream, cache_strategy: CacheStrategy, @@ -869,7 +984,10 @@ pub(crate) fn cache_flat_range_stream( yield batch; } - buffer.finish(key, cache_strategy, part_metrics, None); + if let Some(Ok(value)) = buffer.finish().await { + part_metrics.inc_range_cache_size(key.estimated_size() + value.estimated_size()); + cache_strategy.put_range_result(key, value); + } }) } @@ -936,9 +1054,11 @@ mod tests { use datafusion_expr::{Expr, col, lit}; use smallvec::smallvec; use store_api::storage::{FileId, RegionId}; + use tokio::sync::watch; use super::*; use crate::cache::CacheManager; + use crate::cache::test_util::range_cache_key; use crate::read::flat_projection::FlatProjectionMapper; use crate::read::range::{RangeMeta, RowGroupIndex, SourceIndex}; use crate::read::scan_region::{PredicateGroup, ScanInput}; @@ -994,19 +1114,6 @@ mod tests { (key, part_metrics) } - async fn finish_cache_batch_buffer( - buffer: CacheBatchBuffer, - key: RangeScanCacheKey, - cache_strategy: CacheStrategy, - part_metrics: PartitionMetrics, - ) -> Result> { - let (tx, rx) = oneshot::channel(); - common_telemetry::info!("finish start"); - buffer.finish(key, cache_strategy, part_metrics, Some(tx)); - common_telemetry::info!("finish end"); - rx.await.context(crate::error::RecvSnafu)? - } - async fn new_stream_context( filters: Vec, query_time_range: Option, @@ -1050,6 +1157,295 @@ mod tests { (stream_ctx, partition_range) } + #[rstest::rstest] + #[case(1024 * 1024, true)] + #[case(1, false)] + #[tokio::test] + async fn shared_range_read_preserves_batches_or_falls_back_on_budget_exhaustion( + #[case] capacity: u64, + #[case] cached: bool, + ) { + use std::sync::atomic::{AtomicUsize, Ordering}; + + let (mut ctx, range) = + new_stream_context(vec![col("k0").eq(lit("a"))], None, FileTimeRange::default()).await; + ctx.input.cache_strategy = CacheStrategy::EnableAll(Arc::new( + CacheManager::builder() + .range_result_cache_size(capacity) + .build(), + )); + let key = build_range_cache_key(&ctx, &range).unwrap(); + let (_, metrics) = test_cache_context(&ctx.input.cache_strategy); + let calls = Arc::new(AtomicUsize::new(0)); + let expected = vec![make_batch(&[1, 2]), make_batch(&[3, 4, 5])]; + let batches = expected.clone(); + let count = calls.clone(); + let (stream, _) = CachedRangeReader { + stream_ctx: &ctx, + part_range: &range, + key: Some(key.clone()), + metrics: &metrics, + } + .build(move || async move { + count.fetch_add(1, Ordering::Relaxed); + Ok(( + Box::pin(futures::stream::iter(batches.into_iter().map(Ok))) + as BoxedRecordBatchStream, + DEFAULT_READ_BATCH_SIZE, + )) + }) + .await + .unwrap(); + let actual: Vec<_> = stream.try_collect().await.unwrap(); + assert_eq!(actual, expected); + assert_eq!(calls.load(Ordering::Relaxed), if cached { 1 } else { 2 }); + assert_eq!( + ctx.input.cache_strategy.get_range_result(&key).is_some(), + cached + ); + assert_eq!( + ctx.input + .cache_strategy + .range_cache_worker() + .unwrap() + .load_memory + .available_permits(), + (capacity as usize).div_ceil(1024) + ); + } + + fn start_shared_load( + strategy: &CacheStrategy, + key: RangeScanCacheKey, + batches: Vec, + ) -> watch::Receiver>> { + let worker = strategy.range_cache_worker().unwrap(); + let loader_worker = worker.clone(); + let loader_strategy = strategy.clone(); + let lookup = worker.get_or_load(key, vec![], async move { + collect_range_result( + Box::pin(futures::stream::iter(batches.into_iter().map(Ok))), + &loader_worker, + &loader_strategy, + ) + .await + .unwrap() + }); + let RangeCacheLookup::Loading(receiver) = lookup else { + panic!("expected a shared load"); + }; + receiver + } + + async fn wait_for_shared_load( + receiver: &mut watch::Receiver>>, + ) -> Option> { + tokio::time::timeout(std::time::Duration::from_secs(5), async { + receiver + .wait_for(|value| value.is_some()) + .await + .ok() + .and_then(|value| value.clone()) + }) + .await + .expect("a shared load must finish or fall back without waiting for consumers") + } + + #[tokio::test] + async fn completed_load_keeps_memory_until_all_subscribers_release() { + let strategy = test_cache_strategy(); + let worker = strategy.range_cache_worker().unwrap(); + let capacity = worker.load_memory.available_permits(); + let (key, _) = test_cache_context(&strategy); + let mut first = start_shared_load(&strategy, key.clone(), vec![make_batch(&[1, 2, 3])]); + let second = first.clone(); + drop(wait_for_shared_load(&mut first).await.unwrap()); + + // A completed flight no longer counts toward the flight limit, but + // unpolled subscribers must still reserve its decoded result memory. + assert!(worker.load_memory.available_permits() < capacity); + drop(first); + assert!(worker.load_memory.available_permits() < capacity); + drop(second); + tokio::time::timeout(std::time::Duration::from_secs(5), async { + while worker.load_memory.available_permits() != capacity { + tokio::task::yield_now().await; + } + }) + .await + .expect("the last subscriber must release the delivery reservation"); + assert!( + worker.get(&key).is_some(), + "the cache still owns the result" + ); + } + + #[rstest::rstest] + #[tokio::test] + async fn successive_loads_with_slow_consumers_exhaust_delivery_budget( + #[values(false, true)] invalidate: bool, + ) { + let strategy = CacheStrategy::EnableAll(Arc::new( + CacheManager::builder() + .range_result_cache_size(8192) + .build(), + )); + let worker = strategy.range_cache_worker().unwrap(); + let capacity = worker.load_memory.available_permits(); + let mut held_streams = Vec::new(); + let mut fallbacks = 0; + + // Each flight finishes before the next starts, exceeding the 32-flight + // limit cumulatively. Only the delivery budget can bound this backlog. + for _ in 0..40 { + let file = FileId::random(); + let key = range_cache_key(&[file]); + let expected = vec![make_large_binary_batch(1, 2048)]; + let mut receiver = start_shared_load(&strategy, key, expected.clone()); + let shared = wait_for_shared_load(&mut receiver).await.is_some(); + if invalidate { + worker.invalidate_file(file); + } + let fallback_batches = expected.clone(); + let stream = shared_range_stream(receiver, move || async move { + Ok(( + Box::pin(futures::stream::iter(fallback_batches.into_iter().map(Ok))) + as BoxedRecordBatchStream, + DEFAULT_READ_BATCH_SIZE, + )) + }); + if shared { + held_streams.push(stream); + } else { + fallbacks += 1; + let actual: Vec<_> = + tokio::time::timeout(std::time::Duration::from_secs(5), stream.try_collect()) + .await + .expect("fallback must not wait for an unpolled range") + .unwrap(); + assert_eq!(actual, expected); + } + } + + assert!(!held_streams.is_empty()); + assert!(held_streams.len() < 32); + assert!( + fallbacks > 0, + "completed results must not escape the budget" + ); + assert!(worker.load_memory.available_permits() < capacity); + + // Transfer ownership out of watch receivers into actively polled replay + // streams, then leave them suspended at their first yield. + for stream in &mut held_streams { + assert!(stream.try_next().await.unwrap().is_some()); + } + assert!(worker.load_memory.available_permits() < capacity); + drop(held_streams); + tokio::time::timeout(std::time::Duration::from_secs(5), async { + while worker.load_memory.available_permits() != capacity { + tokio::task::yield_now().await; + } + }) + .await + .expect("dropping slow readers must unblock new shared loads"); + + let mut retry = start_shared_load( + &strategy, + range_cache_key(&[FileId::random()]), + vec![make_large_binary_batch(1, 2048)], + ); + assert!(wait_for_shared_load(&mut retry).await.is_some()); + } + + #[tokio::test] + async fn cancelling_a_partial_load_releases_its_reservation() { + let strategy = test_cache_strategy(); + let worker = strategy.range_cache_worker().unwrap().clone(); + let capacity = worker.load_memory.available_permits(); + let (key, _) = test_cache_context(&strategy); + let loader_worker = worker.clone(); + let lookup = worker.get_or_load(key.clone(), vec![], async move { + let stream = Box::pin(try_stream! { + yield make_batch(&[1, 2, 3]); + std::future::pending::<()>().await; + }); + collect_range_result(stream, &loader_worker, &strategy) + .await + .unwrap() + }); + let RangeCacheLookup::Loading(receiver) = lookup else { + panic!("expected a shared load"); + }; + tokio::time::timeout(std::time::Duration::from_secs(5), async { + while worker.load_memory.available_permits() == capacity { + tokio::task::yield_now().await; + } + }) + .await + .expect("the partial result must acquire its input reservation"); + drop(receiver); + tokio::time::timeout(std::time::Duration::from_secs(5), async { + while worker.load_memory.available_permits() != capacity { + tokio::task::yield_now().await; + } + }) + .await + .expect("cancellation must release the partial result reservation"); + assert!(worker.get(&key).is_none()); + } + + #[tokio::test] + async fn follower_can_finish_without_polling_the_first_range_stream() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + let (ctx, range) = + new_stream_context(vec![col("k0").eq(lit("a"))], None, FileTimeRange::default()).await; + let key = build_range_cache_key(&ctx, &range).unwrap(); + let (_, metrics) = test_cache_context(&ctx.input.cache_strategy); + let gate = Arc::new(tokio::sync::Semaphore::new(0)); + let calls = Arc::new(AtomicUsize::new(0)); + let count = calls.clone(); + let loader_gate = gate.clone(); + let build = move || async move { + count.fetch_add(1, Ordering::Relaxed); + let _permit = loader_gate.acquire().await.unwrap(); + Ok(( + Box::pin(futures::stream::iter([Ok(make_batch(&[1, 2, 3]))])) + as BoxedRecordBatchStream, + DEFAULT_READ_BATCH_SIZE, + )) + }; + let (first, _) = CachedRangeReader { + stream_ctx: &ctx, + part_range: &range, + key: Some(key.clone()), + metrics: &metrics, + } + .build(build.clone()) + .await + .unwrap(); + let (second, _) = CachedRangeReader { + stream_ctx: &ctx, + part_range: &range, + key: Some(key), + metrics: &metrics, + } + .build(build) + .await + .unwrap(); + gate.add_permits(1); + let second_batches: Vec<_> = + tokio::time::timeout(std::time::Duration::from_secs(5), second.try_collect()) + .await + .expect("a follower must not depend on polling the first query") + .unwrap(); + let first_batches: Vec<_> = first.try_collect().await.unwrap(); + assert_eq!(first_batches, second_batches); + assert_eq!(second_batches, vec![make_batch(&[1, 2, 3])]); + assert_eq!(calls.load(Ordering::Relaxed), 1); + } + /// Helper to create a timestamp millisecond literal. fn ts_lit(val: i64) -> Expr { lit(ScalarValue::TimestampMillisecond(Some(val), None)) @@ -1600,7 +1996,13 @@ mod tests { make_batch(&[1, 2, 3]).get_array_memory_size(), )); - let replayed = cached_flat_range_stream(value) + let worker = RangeCacheWorker::new(1024 * 1024); + let permit = worker + .load_memory + .try_acquire(value.estimated_size()) + .unwrap(); + let delivery = RangeCacheDelivery::new(value, permit); + let replayed = cached_flat_range_stream(delivery) .try_collect::>() .await .unwrap(); @@ -1652,28 +2054,20 @@ mod tests { let strategy = test_cache_strategy(); let batch = make_batch(&[1, 2, 3]); let expected_size = batch.get_array_memory_size(); - let (key, part_metrics) = test_cache_context(&strategy); let mut buffer = CacheBatchBuffer::new(&strategy); buffer.push(batch).unwrap(); - let value = finish_cache_batch_buffer(buffer, key.clone(), strategy.clone(), part_metrics) - .await - .unwrap(); + let value = buffer.finish().await.unwrap().unwrap(); assert_eq!(value.cached_batches.len(), 1); assert_eq!(value.cached_batches[0].slice_lengths, vec![3]); assert_eq!(value.estimated_batches_size, expected_size); - assert!(Arc::ptr_eq( - &value, - &strategy.get_range_result(&key).unwrap() - )); } #[tokio::test] async fn cache_batch_buffer_compacts_when_rows_exceed_default_batch_size() { let strategy = test_cache_strategy(); let batch = make_batch(&vec![1; DEFAULT_READ_BATCH_SIZE / 2 + 1]); - let (key, part_metrics) = test_cache_context(&strategy); let mut buffer = CacheBatchBuffer::new(&strategy); buffer.push(batch.clone()).unwrap(); @@ -1682,9 +2076,7 @@ mod tests { assert_eq!(buffer.buffered_rows, 0); assert!(buffer.buffered_batches.is_empty()); - let value = finish_cache_batch_buffer(buffer, key, strategy, part_metrics) - .await - .unwrap(); + let value = buffer.finish().await.unwrap().unwrap(); assert_eq!(value.cached_batches.len(), 1); assert_eq!( value.cached_batches[0].slice_lengths, @@ -1703,7 +2095,6 @@ mod tests { .range_result_cache_size((large_batch.get_array_memory_size() * 3) as u64) .build(), )); - let (key, part_metrics) = test_cache_context(&strategy); let mut buffer = CacheBatchBuffer::new(&strategy); buffer.push(large_batch.clone()).unwrap(); @@ -1716,9 +2107,7 @@ mod tests { assert_eq!(buffer.buffered_rows, 0); assert!(buffer.buffered_batches.is_empty()); - let value = finish_cache_batch_buffer(buffer, key, strategy, part_metrics) - .await - .unwrap(); + let value = buffer.finish().await.unwrap().unwrap(); assert_eq!(value.cached_batches.len(), 1); assert_eq!( value.cached_batches[0].slice_lengths, @@ -1736,17 +2125,11 @@ mod tests { .range_result_cache_size(budget) .build(), )); - let (key, part_metrics) = test_cache_context(&strategy); let mut buffer = CacheBatchBuffer::new(&strategy); for _ in 0..4 { buffer.push(large_batch.clone()).unwrap(); } - assert!( - finish_cache_batch_buffer(buffer, key.clone(), strategy.clone(), part_metrics) - .await - .is_err() - ); - assert!(strategy.get_range_result(&key).is_none()); + assert!(!matches!(buffer.finish().await, Some(Ok(_)))); } } diff --git a/src/mito2/src/read/seq_scan.rs b/src/mito2/src/read/seq_scan.rs index 42119135805..4b09533e29a 100644 --- a/src/mito2/src/read/seq_scan.rs +++ b/src/mito2/src/read/seq_scan.rs @@ -42,9 +42,7 @@ use crate::read::flat_merge::FlatMergeReader; use crate::read::last_row::FlatLastRowReader; use crate::read::pruner::{PartitionPruner, Pruner}; use crate::read::range::RangeMeta; -use crate::read::range_cache::{ - build_range_cache_key, cache_flat_range_stream, cached_flat_range_stream, -}; +use crate::read::range_cache::{CachedRangeReader, build_range_cache_key}; use crate::read::scan_region::{ScanInput, StreamContext}; use crate::read::scan_util::{ BATCH_SIZE_THRESHOLD, MIN_COMPACTION_SPLIT_ROWS, PartitionMetrics, PartitionMetricsList, @@ -307,15 +305,39 @@ impl SeqScan { key } }); - - if let Some(key) = cache_key.as_ref() { - if let Some(value) = stream_ctx.input.cache_strategy.get_range_result(key) { - part_metrics.inc_range_cache_hit(); - return Ok((cached_flat_range_stream(value), DEFAULT_READ_BATCH_SIZE)); - } - part_metrics.inc_range_cache_miss(); + let ctx = stream_ctx.clone(); + let range = *part_range; + let metrics = part_metrics.clone(); + CachedRangeReader { + stream_ctx, + part_range, + key: cache_key, + metrics: part_metrics, } + .build(move || async move { + Self::build_uncached_partition_range( + &ctx, + &range, + read_mode, + &metrics, + partition_pruner, + file_scan_semaphore, + merge_semaphore, + ) + .await + }) + .await + } + async fn build_uncached_partition_range( + stream_ctx: &Arc, + part_range: &PartitionRange, + read_mode: FlatReadMode, + part_metrics: &PartitionMetrics, + partition_pruner: Arc, + file_scan_semaphore: Option>, + merge_semaphore: Option>, + ) -> Result<(BoxedRecordBatchStream, usize)> { let mut sources = Vec::new(); let split_batch_size = build_flat_sources( stream_ctx, @@ -339,16 +361,6 @@ impl SeqScan { ) .await?; - let stream = match cache_key { - Some(key) => cache_flat_range_stream( - stream, - stream_ctx.input.cache_strategy.clone(), - key, - part_metrics.clone(), - ), - None => stream, - }; - Ok((stream, estimated_rows_per_batch)) } diff --git a/src/mito2/src/read/series_candidate.rs b/src/mito2/src/read/series_candidate.rs index 3ba5512a7bd..185d3ea0dff 100644 --- a/src/mito2/src/read/series_candidate.rs +++ b/src/mito2/src/read/series_candidate.rs @@ -47,9 +47,7 @@ use crate::error::{ use crate::read::BoxedRecordBatchStream; use crate::read::pruner::{PartitionPruner, Pruner}; use crate::read::range::RowGroupIndex; -use crate::read::range_cache::{ - build_candidate_range_cache_key, cache_flat_range_stream, cached_flat_range_stream, -}; +use crate::read::range_cache::{CachedRangeReader, build_candidate_range_cache_key}; use crate::read::scan_region::StreamContext; use crate::read::scan_util::{PartitionMetrics, new_filter_metrics, scan_flat_mem_ranges}; use crate::series_index::{ @@ -371,14 +369,29 @@ impl SeriesCandidateRangeBuilder { } else { build_candidate_range_cache_key(&self.stream_ctx, &part_range) }; - if let Some(key) = cache_key.as_ref() { - if let Some(value) = self.stream_ctx.input.cache_strategy.get_range_result(key) { - self.part_metrics.inc_range_cache_hit(); - return Ok(cached_flat_range_stream(value)); - } - self.part_metrics.inc_range_cache_miss(); + let builder = self.clone(); + let (stream, _) = CachedRangeReader { + stream_ctx: &self.stream_ctx, + part_range: &part_range, + key: cache_key, + metrics: &self.part_metrics, } + .build(move || async move { + let stream = builder + .build_uncached_range_stream(part_range, merge_partition) + .await?; + Ok((stream, DEFAULT_READ_BATCH_SIZE)) + }) + .await?; + Ok(stream) + } + async fn build_uncached_range_stream( + &self, + part_range: PartitionRange, + merge_partition: usize, + ) -> Result { + let range_meta = &self.stream_ctx.ranges[part_range.identifier]; let mut sources = Vec::with_capacity(range_meta.row_group_indices.len()); for index in &range_meta.row_group_indices { let source = self.build_source(*index, range_meta.time_range).await?; @@ -400,15 +413,7 @@ impl SeriesCandidateRangeBuilder { "SeriesCandidateScanner::range_merge", )?; - Ok(match cache_key { - Some(key) => cache_flat_range_stream( - stream, - self.stream_ctx.input.cache_strategy.clone(), - key, - self.part_metrics.clone(), - ), - None => stream, - }) + Ok(stream) } async fn build_source( diff --git a/src/mito2/src/read/series_reader.rs b/src/mito2/src/read/series_reader.rs index 5ac69d191a6..c030b5b3e7e 100644 --- a/src/mito2/src/read/series_reader.rs +++ b/src/mito2/src/read/series_reader.rs @@ -31,9 +31,7 @@ use crate::error::InvalidRequestSnafu; use crate::error::{JoinSnafu, Result, UnexpectedSnafu}; use crate::read::BoxedRecordBatchStream; use crate::read::pruner::PartitionPruner; -use crate::read::range_cache::{ - build_series_range_cache_key, cache_flat_range_stream, cached_flat_range_stream, -}; +use crate::read::range_cache::{CachedRangeReader, build_series_range_cache_key}; use crate::read::scan_region::StreamContext; use crate::read::scan_util::{ PartitionMetrics, SplitRecordBatchStream, compute_average_batch_size, @@ -377,14 +375,36 @@ async fn build_series_partition_range( .enable_range_cache .then(|| build_series_range_cache_key(&stream_ctx, &part_range, range)) .flatten(); - if let Some(key) = cache_key.as_ref() { - if let Some(value) = stream_ctx.input.cache_strategy.get_range_result(key) { - part_metrics.inc_range_cache_hit(); - return Ok((cached_flat_range_stream(value), DEFAULT_READ_BATCH_SIZE)); - } - part_metrics.inc_range_cache_miss(); + let ctx = stream_ctx.clone(); + let metrics = part_metrics.clone(); + CachedRangeReader { + stream_ctx: &stream_ctx, + part_range: &part_range, + key: cache_key, + metrics: &part_metrics, } + .build(move || async move { + build_uncached_series_partition_range( + ctx, + part_range, + filter, + codec, + partition_pruner, + metrics, + ) + .await + }) + .await +} +async fn build_uncached_series_partition_range( + stream_ctx: Arc, + part_range: PartitionRange, + filter: MetricSeriesFilter, + codec: SparsePrimaryKeyCodec, + partition_pruner: Arc, + part_metrics: PartitionMetrics, +) -> Result<(BoxedRecordBatchStream, usize)> { let range_meta = &stream_ctx.ranges[part_range.identifier]; let split_batch_size = should_split_flat_batches_for_merge(&stream_ctx, range_meta); let mut sources = Vec::with_capacity(range_meta.row_group_indices.len()); @@ -480,15 +500,6 @@ async fn build_series_partition_range( compute_parallel_channel_size(estimated_batch_size), ) .await?; - let stream = match cache_key { - Some(key) => cache_flat_range_stream( - stream, - stream_ctx.input.cache_strategy.clone(), - key, - part_metrics, - ), - None => stream, - }; Ok((stream, estimated_batch_size)) } diff --git a/src/mito2/src/sst/file.rs b/src/mito2/src/sst/file.rs index 4c6613dee40..a627513d310 100644 --- a/src/mito2/src/sst/file.rs +++ b/src/mito2/src/sst/file.rs @@ -633,7 +633,12 @@ impl FileHandle { /// Mark the file as deleted and will delete it on drop asynchronously pub fn mark_deleted(&self) { - self.inner.lifetime.deleted.store(true, Ordering::Relaxed); + if !self.inner.lifetime.deleted.swap(true, Ordering::Relaxed) { + self.inner + .lifetime + .file_purger + .invalidate_file(&self.inner.meta); + } } pub fn compacting(&self) -> bool { @@ -867,6 +872,7 @@ pub async fn delete_files( // Remove meta of the file from cache. if let Some(cache) = &cache_manager { for (file_id, _) in file_ids { + cache.invalidate_range_results(*file_id); cache.remove_parquet_meta_data(RegionFileId::new(region_id, *file_id)); } } diff --git a/src/mito2/src/sst/file_purger.rs b/src/mito2/src/sst/file_purger.rs index 727c6c1cc9e..ad44c5bcc05 100644 --- a/src/mito2/src/sst/file_purger.rs +++ b/src/mito2/src/sst/file_purger.rs @@ -33,6 +33,9 @@ pub trait FilePurger: Send + Sync + fmt::Debug { /// Each metadata view has already released its reader reference. fn purge_file(&self, file_meta: FileMeta); + /// Retire cached scan results immediately, without waiting for physical deletion. + fn invalidate_file(&self, _: &FileMeta) {} + /// Returns the object-store reference manager used by SST metadata views. /// Local and no-op purgers do not track GC references. fn file_ref_manager(&self) -> Option { @@ -105,6 +108,7 @@ pub fn create_file_purger( && matches!(path_type, PathType::Data | PathType::Bare) { Arc::new(ObjectStoreFilePurger { + cache_manager, file_ref_manager, scheduler, range_index_deleter, @@ -180,6 +184,12 @@ impl LocalFilePurger { } impl FilePurger for LocalFilePurger { + fn invalidate_file(&self, file_meta: &FileMeta) { + if let Some(cache) = &self.cache_manager { + cache.invalidate_range_results(file_meta.file_id); + } + } + fn purge_file(&self, file_meta: FileMeta) { schedule_range_index_deletion( &self.scheduler, @@ -191,6 +201,7 @@ impl FilePurger for LocalFilePurger { } pub struct ObjectStoreFilePurger { + cache_manager: Option, file_ref_manager: FileReferenceManagerRef, scheduler: SchedulerRef, range_index_deleter: Option, @@ -224,6 +235,12 @@ fn schedule_range_index_deletion( } impl FilePurger for ObjectStoreFilePurger { + fn invalidate_file(&self, file_meta: &FileMeta) { + if let Some(cache) = &self.cache_manager { + cache.invalidate_range_results(file_meta.file_id); + } + } + fn purge_file(&self, file_meta: FileMeta) { schedule_range_index_deletion( &self.scheduler, @@ -407,12 +424,25 @@ mod tests { let owner = RegionId::new(9, 1); let index_path = crate::sst::range_index::range_index_path(owner, sst_file_id.file_id()); index_store.write(&index_path, "range index").await.unwrap(); + let cache = Arc::new( + crate::cache::CacheManager::builder() + .range_result_cache_size(1024 * 1024) + .build(), + ); + let cache_key = crate::cache::test_util::range_cache_key(&[sst_file_id.file_id()]); + cache.put_range_result( + cache_key.clone(), + Arc::new(crate::read::range_cache::RangeScanCacheValue::new( + vec![], + 0, + )), + ); let file_purger = create_file_purger( gc_enabled, PathType::Bare, scheduler.clone(), env.layer.clone(), - None, + Some(cache.clone()), Arc::new(crate::sst::file_ref::FileReferenceManager::new(None)), Some(RangeIndexDeleter::new(index_store.clone(), owner)), ); @@ -443,6 +473,8 @@ mod tests { if is_delete { version.mark_all_deleted(); } + // Invalidate before any old index-version reader releases the SST. + assert_eq!(cache.get_range_result(&cache_key).is_none(), is_delete); drop(version); env.wait_for_purges().await; assert!(object_store.exists(&path).await.unwrap()); From 62ebf71f90a6167966c1f3e64ed6d01dfc57b14e Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Tue, 27 Oct 2026 13:12:52 +0000 Subject: [PATCH 111/239] fix(meta): exclude route followers from failover targets (#10001) * fix(meta): exclude route followers from failover targets Failover target selection did not exclude the region's route followers. The selector could keep picking a follower, the submit rejected it as a peer conflict, and the region never failed over. Group the failed regions by their followers and exclude them per group. A group without a candidate is skipped instead of failing every region of the datanode. Signed-off-by: jeremyhi * test: check failover continues after a group without candidates Signed-off-by: jeremyhi * fix(meta): keep failover going when follower routes can't be read Fall back to selecting without follower exclusion when the routes can't be read, log skipped groups with their followers at warn level, and test a region without a route. Signed-off-by: jeremyhi * refactor(meta): select failover targets per exclusion set in select_peers generate_failover_tasks now only reads the route followers of the regions. select_peers resolves the affinity and the unschedulable datanodes once, then selects per group of regions that exclude the same datanodes. Signed-off-by: jeremyhi * refactor(meta): drop a comment that repeats the code Signed-off-by: jeremyhi * refactor(meta): tighten failover target selection Name the route followers after the parameter they fill, read the routes with unwrap_or_else, and log a group without a candidate as an error, as the caller did before. Signed-off-by: jeremyhi * fix(meta): let the region-stat-aware selector balance the whole batch The region-stat-aware selector now receives each region's excluded peers along with the whole batch, so it can still balance load across the regions while keeping a region off its own followers. A region without a candidate is left out of its output instead of failing the others. The selectors driven by SelectorOptions keep selecting per group of regions that exclude the same datanodes. Signed-off-by: jeremyhi * test: name the batch type of the stat-aware recorder Signed-off-by: jeremyhi * refactor(meta): treat a selector without candidates the same in both arms Both selector arms now place what they can and treat NoEnoughAvailableNode as nothing placed; the stat-aware selector's contract says so. Rename its parameter to exclude_peer_ids_by_region, check for unplaced regions in linear time, and describe what the code does in the comments. Signed-off-by: jeremyhi * test: keep the stat-aware recorder's placements off excluded peers Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- src/meta-srv/src/region/supervisor.rs | 333 ++++++++++++++++++++++---- src/meta-srv/src/selector.rs | 7 +- 2 files changed, 290 insertions(+), 50 deletions(-) diff --git a/src/meta-srv/src/region/supervisor.rs b/src/meta-srv/src/region/supervisor.rs index c7c5cf7dab7..0e1b02695f4 100644 --- a/src/meta-srv/src/region/supervisor.rs +++ b/src/meta-srv/src/region/supervisor.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::{HashMap, HashSet}; +use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet}; use std::fmt::Debug; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; @@ -23,7 +23,7 @@ use common_meta::affinity::{AffinityFallback, AffinityReference, AffinityResolve use common_meta::datanode::Stat; use common_meta::ddl::{DetectingRegion, RegionFailureDetectorController}; use common_meta::key::runtime_switch::RuntimeSwitchManagerRef; -use common_meta::key::table_route::{TableRouteKey, TableRouteValue}; +use common_meta::key::table_route::{TableRouteKey, TableRouteStorage, TableRouteValue}; use common_meta::key::{MetadataKey, MetadataValue}; use common_meta::kv_backend::KvBackendRef; use common_meta::leadership_notifier::LeadershipChangeListener; @@ -34,6 +34,7 @@ use common_runtime::JoinHandle; use common_telemetry::{debug, error, info, warn}; use common_time::util::current_time_millis; use futures::{StreamExt, TryStreamExt}; +use itertools::Itertools; use snafu::{ResultExt, ensure}; use store_api::storage::RegionId; use tokio::sync::mpsc::error::TrySendError; @@ -44,7 +45,7 @@ use tokio::time::{MissedTickBehavior, interval, interval_at}; use crate::discovery::utils::accept_ingest_workload; use crate::error::{self, Result}; use crate::failure_detector::PhiAccrualFailureDetectorOptions; -use crate::metasrv::{RegionStatAwareSelectorRef, SelectTarget, SelectorContext, SelectorRef}; +use crate::metasrv::{RegionStatAwareSelectorRef, SelectorContext, SelectorRef}; use crate::metrics::METRIC_META_HEARTBEAT_DROPPED; use crate::procedure::region_migration::manager::{ RegionMigrationManagerRef, RegionMigrationTriggerReason, SubmitRegionMigrationTaskResult, @@ -667,11 +668,15 @@ impl RegionSupervisor { .context(error::RuntimeSwitchManagerSnafu) } + /// Selects a target for each of `regions`, excluding `failure_datanodes`, the unschedulable + /// datanodes and the region's datanodes in `exclude_peer_ids_by_region`. Regions without a + /// candidate are left out. async fn select_peers( &self, from_peer_id: DatanodeId, regions: &[RegionId], failure_datanodes: &[DatanodeId], + exclude_peer_ids_by_region: &HashMap>, ) -> Result> { let affinity = if let Some(resolver) = &self.affinity_resolver { match resolver @@ -715,54 +720,66 @@ impl RegionSupervisor { .filter(|datanode| !datanode.is_schedulable()) .map(|datanode| datanode.peer.id); exclude_peer_ids.extend(unschedulable_datanodes); - match &self.selector { - RegionSupervisorSelector::NaiveSelector(selector) => { - let opts = SelectorOptions { - affinity: affinity.as_ref(), - min_required_items: regions.len(), - allow_duplication: true, - exclude_peer_ids, - workload_filter: Some(accept_ingest_workload), - }; - let peers = selector.select(&self.selector_context, opts).await?; - ensure!( - peers.len() == regions.len(), - error::NoEnoughAvailableNodeSnafu { - required: regions.len(), - available: peers.len(), - select_target: SelectTarget::Datanode, - } - ); - let region_peers = regions - .iter() - .zip(peers) - .map(|(region_id, peer)| (*region_id, peer)) - .collect::>(); - Ok(region_peers) + let region_peers = match &self.selector { + RegionSupervisorSelector::NaiveSelector(selector) => { + // `Selector` takes one exclusion set, so select once per distinct set of excluded + // datanodes. + let empty = BTreeSet::new(); + let mut groups = BTreeMap::<_, Vec<_>>::new(); + for region_id in regions { + let excluded = exclude_peer_ids_by_region.get(region_id).unwrap_or(&empty); + groups.entry(excluded).or_default().push(*region_id); + } + let mut region_peers = Vec::with_capacity(regions.len()); + for (excluded, regions) in groups { + let mut exclude_peer_ids = exclude_peer_ids.clone(); + exclude_peer_ids.extend(excluded); + let opts = SelectorOptions { + affinity: affinity.as_ref(), + min_required_items: regions.len(), + allow_duplication: true, + exclude_peer_ids, + workload_filter: Some(accept_ingest_workload), + }; + match selector.select(&self.selector_context, opts).await { + Ok(peers) => region_peers.extend(regions.into_iter().zip(peers)), + // Other groups may still have candidates. + Err(error::Error::NoEnoughAvailableNode { .. }) => {} + Err(err) => return Err(err), + } + } + region_peers } RegionSupervisorSelector::RegionStatAwareSelector(selector) => { - let peers = selector + let selected = selector .select( &self.selector_context, from_peer_id, regions, exclude_peer_ids, + exclude_peer_ids_by_region, affinity.as_ref(), ) - .await?; - ensure!( - peers.len() == regions.len(), - error::NoEnoughAvailableNodeSnafu { - required: regions.len(), - available: peers.len(), - select_target: SelectTarget::Datanode, - } - ); - - Ok(peers) + .await; + match selected { + Err(error::Error::NoEnoughAvailableNode { .. }) => vec![], + selected => selected?, + } } + }; + let placed = region_peers + .iter() + .map(|(region_id, _)| *region_id) + .collect::>(); + let unplaced = regions + .iter() + .filter(|region_id| !placed.contains(region_id)) + .collect::>(); + if !unplaced.is_empty() { + error!("No failover target for regions: {:?}", unplaced); } + Ok(region_peers) } async fn generate_failover_tasks( @@ -780,8 +797,37 @@ impl RegionSupervisor { .flatten() .unwrap_or_else(|| Peer::empty(from_peer_id)); + // A region's route followers can't take it over. + let table_ids = regions + .iter() + .map(|region_id| region_id.table_id()) + .unique() + .collect::>(); + let table_routes = TableRouteStorage::new(self.kv_backend.clone()) + .batch_get(&table_ids) + .await + .unwrap_or_else(|err| { + // Exclude no followers; the submit still rejects a follower as the target. + warn!(err; "Failed to read the routes of failover regions: {:?}", regions); + vec![] + }); + let exclude_peer_ids_by_region = table_routes + .iter() + .flatten() + .filter_map(|table_route| table_route.region_routes().ok()) + .flatten() + .map(|route| { + let peers = route.follower_peers.iter().map(|peer| peer.id); + (route.region.id, peers.collect()) + }) + .collect(); let region_peers = self - .select_peers(from_peer_id, regions, failed_datanodes) + .select_peers( + from_peer_id, + regions, + failed_datanodes, + &exclude_peer_ids_by_region, + ) .await?; for (region_id, peer) in region_peers { @@ -1032,7 +1078,7 @@ pub(crate) mod tests { #[tokio::test] async fn test_failover_affinity_fallback_and_selector_errors() { - use std::collections::HashSet; + use std::collections::{BTreeSet, HashSet}; use std::sync::atomic::{AtomicUsize, Ordering}; use common_meta::affinity::{ @@ -1147,11 +1193,13 @@ pub(crate) mod tests { source: u64, regions: &[RegionId], excluded: HashSet, + exclude_peer_ids_by_region: &HashMap>, affinity: Option<&AffinityRequirements>, ) -> Result { assert_eq!(ctx.affinity_fallback, self.expected_fallback); assert_eq!(source, 7); assert_eq!(excluded, HashSet::from([8])); + assert!(exclude_peer_ids_by_region.is_empty()); Ok(regions .iter() .copied() @@ -1189,21 +1237,25 @@ pub(crate) mod tests { }; let regions = [RegionId::new(1, 1), RegionId::new(1, 2)]; // Preserve the caller's exclusion set without adding the source. - let result = supervisor.select_peers(7, ®ions, &[8]).await; + let result = supervisor + .select_peers(7, ®ions, &[8], &HashMap::new()) + .await; if succeeds { assert_eq!( result.unwrap(), vec![(regions[0], Peer::empty(9)), (regions[1], Peer::empty(9))] ); } else { - let error = result.unwrap_err(); match outcome { - Outcome::Short | Outcome::Exhausted => { - assert!(matches!(error, error::Error::NoEnoughAvailableNode { .. })) - } - Outcome::Error => { - assert!(matches!(error, error::Error::Unexpected { .. })) + // Only the regions with a candidate are placed. + Outcome::Short => { + assert_eq!(result.unwrap(), vec![(regions[0], Peer::empty(9))]) } + Outcome::Exhausted => assert!(result.unwrap().is_empty()), + Outcome::Error => assert!(matches!( + result.unwrap_err(), + error::Error::Unexpected { .. } + )), Outcome::Success => unreachable!(), } } @@ -1276,13 +1328,196 @@ pub(crate) mod tests { ); let region_peers = supervisor - .select_peers(7, &[RegionId::new(1, 1)], &[8]) + .select_peers(7, &[RegionId::new(1, 1)], &[8], &HashMap::new()) .await .unwrap(); assert_eq!(1, region_peers.len()); assert_eq!(Peer::empty(9), region_peers[0].1); } + #[tokio::test] + async fn test_failover_excludes_route_followers() { + use std::collections::{BTreeSet, HashSet as StdHashSet}; + + use common_meta::affinity::AffinityRequirements; + use common_meta::key::MetadataKey; + use common_meta::key::table_route::TableRouteKey; + use common_meta::rpc::store::PutRequest; + + use crate::error::{self, Result}; + use crate::metasrv::SelectTarget; + use crate::selector::{RegionStatAwareSelector, Selector, SelectorOptions}; + + /// Records each selection, and has no candidate once datanode 3 is excluded. + struct RecordingSelector(Mutex, usize)>>); + + #[async_trait::async_trait] + impl Selector for RecordingSelector { + type Context = SelectorContext; + type Output = Vec; + + async fn select( + &self, + _ctx: &Self::Context, + opts: SelectorOptions<'_>, + ) -> Result { + let required = opts.min_required_items; + let no_candidate = opts.exclude_peer_ids.contains(&3); + self.0 + .lock() + .unwrap() + .push((opts.exclude_peer_ids, required)); + if no_candidate { + return error::NoEnoughAvailableNodeSnafu { + required, + available: 0usize, + select_target: SelectTarget::Datanode, + } + .fail(); + } + Ok(vec![Peer::empty(3); required]) + } + } + + /// A batch of regions with their excluded peers. + type Batch = (Vec, HashMap>); + + /// Records each batch, and places every region on datanode 9, which no region excludes. + struct StatAwareRecorder(Mutex>); + + #[async_trait::async_trait] + impl RegionStatAwareSelector for StatAwareRecorder { + type Context = SelectorContext; + type Output = Vec<(RegionId, Peer)>; + + async fn select( + &self, + _ctx: &Self::Context, + _from_peer_id: u64, + region_ids: &[RegionId], + exclude_peer_ids: StdHashSet, + exclude_peer_ids_by_region: &HashMap>, + _affinity: Option<&AffinityRequirements>, + ) -> Result { + assert_eq!(exclude_peer_ids, StdHashSet::from([8])); + self.0 + .lock() + .unwrap() + .push((region_ids.to_vec(), exclude_peer_ids_by_region.clone())); + Ok(region_ids + .iter() + .map(|region_id| (*region_id, Peer::empty(9))) + .collect()) + } + } + + let selector = Arc::new(RecordingSelector(Mutex::new(vec![]))); + let (mut supervisor, _tx) = new_test_supervisor_with( + new_test_selector_context(), + RegionSupervisorSelector::NaiveSelector(selector.clone()), + ); + let table_id = 1024; + let region_route = |region_number, followers: &[u64]| RegionRoute { + region: Region { + id: RegionId::new(table_id, region_number), + ..Default::default() + }, + leader_peer: Some(Peer::empty(8)), + follower_peers: followers.iter().map(|id| Peer::empty(*id)).collect(), + ..Default::default() + }; + let table_route = PhysicalTableRouteValue::new(vec![ + region_route(0, &[]), + region_route(1, &[3]), + region_route(2, &[4]), + region_route(3, &[4]), + ]); + let mut task = test_create_physical_table_task("my_physical_table"); + task.set_table_id(table_id); + TableMetadataManager::new(supervisor.kv_backend.clone()) + .create_table_metadata( + task.table_info, + TableRouteValue::Physical(table_route), + HashMap::new(), + ) + .await + .unwrap(); + + let mut regions = (0..4) + .map(|number| RegionId::new(table_id, number)) + .collect::>(); + // A region of a table without a route has no followers. + regions.push(RegionId::new(table_id + 1, 0)); + let tasks = supervisor + .generate_failover_tasks(8, ®ions, &[8]) + .await + .unwrap(); + + // Regions with the same followers are selected together, excluding those followers. + assert_eq!( + *selector.0.lock().unwrap(), + vec![ + (StdHashSet::from([8]), 2), + (StdHashSet::from([8, 3]), 1), + (StdHashSet::from([8, 4]), 2), + ] + ); + // Only the region whose follower is the last candidate is skipped. + let failed_over = tasks + .iter() + .map(|(task, _)| task.region_id) + .collect::>(); + assert_eq!( + failed_over, + [regions[0], regions[4], regions[2], regions[3]] + ); + + // The stat-aware selector gets the whole batch with each region's followers. + let stat_aware = Arc::new(StatAwareRecorder(Mutex::new(vec![]))); + supervisor.selector = RegionSupervisorSelector::RegionStatAwareSelector(stat_aware.clone()); + let tasks = supervisor + .generate_failover_tasks(8, ®ions, &[8]) + .await + .unwrap(); + let followers = HashMap::from([ + (regions[0], BTreeSet::new()), + (regions[1], BTreeSet::from([3])), + (regions[2], BTreeSet::from([4])), + (regions[3], BTreeSet::from([4])), + ]); + assert_eq!( + *stat_aware.0.lock().unwrap(), + vec![(regions.clone(), followers)] + ); + let failed_over = tasks + .iter() + .map(|(task, _)| task.region_id) + .collect::>(); + assert_eq!(failed_over, regions); + supervisor.selector = RegionSupervisorSelector::NaiveSelector(selector.clone()); + + // A route that can't be read excludes no followers. + supervisor + .kv_backend + .put( + PutRequest::new() + .with_key(TableRouteKey::new(table_id + 1).to_bytes()) + .with_value(b"corrupted".to_vec()), + ) + .await + .unwrap(); + selector.0.lock().unwrap().clear(); + let tasks = supervisor + .generate_failover_tasks(8, ®ions, &[8]) + .await + .unwrap(); + assert_eq!( + *selector.0.lock().unwrap(), + vec![(StdHashSet::from([8]), 5)] + ); + assert_eq!(tasks.len(), 5); + } + #[tokio::test] async fn test_heartbeat() { let (mut supervisor, sender) = new_test_supervisor(); diff --git a/src/meta-srv/src/selector.rs b/src/meta-srv/src/selector.rs index c4d29b8ecea..ca151be0b68 100644 --- a/src/meta-srv/src/selector.rs +++ b/src/meta-srv/src/selector.rs @@ -20,7 +20,7 @@ pub mod round_robin; pub(crate) mod test_utils; pub mod weight_compute; pub mod weighted_choose; -use std::collections::HashSet; +use std::collections::{BTreeSet, HashMap, HashSet}; use api::v1::meta::heartbeat_request::NodeWorkloads; use common_meta::affinity::AffinityRequirements; @@ -43,6 +43,10 @@ pub trait Selector: Send + Sync { /// /// It selects the best destination peer for a list of regions. /// The selection is based on the region statistics, such as the region leader's write throughput. +/// +/// A region is never placed on `exclude_peer_ids`, nor on its own peers in +/// `exclude_peer_ids_by_region`. It returns the regions it placed, possibly none: a region +/// without a remaining candidate is left out. Errors are for other failures. #[async_trait::async_trait] pub trait RegionStatAwareSelector: Send + Sync { type Context; @@ -54,6 +58,7 @@ pub trait RegionStatAwareSelector: Send + Sync { from_peer_id: u64, region_ids: &[RegionId], exclude_peer_ids: HashSet, + exclude_peer_ids_by_region: &HashMap>, affinity: Option<&AffinityRequirements>, ) -> Result; } From 9fc396c148816ff84685e0aee025c5cd4ff80025 Mon Sep 17 00:00:00 2001 From: LFC <990479+MichaelScofield@users.noreply.github.com> Date: Wed, 28 Oct 2026 01:04:44 +0000 Subject: [PATCH 112/239] ci(query-regression): bump RUNNER_IMAGE_EPOCH to 8 (#10019) ci(query-regression): bump RUNNER_IMAGE_EPOCH to 8 for image m-0xihwmu1dxusb9vuf13q Signed-off-by: greptimedb-ci Co-authored-by: greptimedb-ci --- .github/workflows/query-regression.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/query-regression.yml b/.github/workflows/query-regression.yml index 6a368c455d6..6543281409d 100644 --- a/.github/workflows/query-regression.yml +++ b/.github/workflows/query-regression.yml @@ -549,7 +549,7 @@ jobs: readonly EXPECTED_SCCACHE_DIR="/home/runner/.cache/sccache" readonly EXPECTED_RUSTC_WRAPPER="/usr/local/bin/sccache" readonly RUNNER_IMAGE_DIGEST="sha256:e713b294e23b7e15184e558866c90025e59930033e72c97650dbc7f1ca022d11" - readonly RUNNER_IMAGE_EPOCH="7" + readonly RUNNER_IMAGE_EPOCH="8" require_expected_root() { local name="$1" From 180e2716fc5fb3f192a6f6cd57eac308892ccdcc Mon Sep 17 00:00:00 2001 From: Yvan Wang <131545713+BootstrapperSBL@users.noreply.github.com> Date: Wed, 28 Oct 2026 02:34:56 +0000 Subject: [PATCH 113/239] fix(log-store): make raft-engine periodic WAL sync actually fsync (#9715) * fix(log-store): make raft-engine periodic WAL sync actually fsync On raft-engine 0.4.2, Engine::sync() writes an empty LogBatch and write() returns before fsync (tikv/raft-engine#395). The periodic sync task therefore never durable-flushed acknowledged WAL data. Work around by writing a minimal non-empty marker with sync=true so the fsync path is reached until a crates.io release includes tikv/raft-engine#400. Fixes #9712 Signed-off-by: Yvan Wang <131545713+BootstrapperSBL@users.noreply.github.com> * perf(log-store): reuse LogBatch across periodic WAL syncs Keep the periodic sync task's LogBatch across ticks instead of allocating a fresh one (4 KiB encode buffer plus item vec) every period. A successful Engine::write drains the batch but retains its buffers; after a failed write or a panicked blocking task the batch is discarded and recreated on the next tick, since it may be left partially encoded. Keep the marker put rather than deleting an empty key: the put overwrites a single key, whereas every delete would become a recovery tombstone that raft-engine keeps in memory while replaying append files. Also document that the marker value must not decode as a NamespaceImpl (the marker key lies inside the range scanned by list_namespaces), add a test covering batch reuse and namespace listing, and apply rustfmt. Refs #9712 Signed-off-by: Yvan Wang <131545713+BootstrapperSBL@users.noreply.github.com> --------- Signed-off-by: Yvan Wang <131545713+BootstrapperSBL@users.noreply.github.com> --- src/log-store/src/raft_engine/log_store.rs | 122 ++++++++++++++++++++- 1 file changed, 118 insertions(+), 4 deletions(-) diff --git a/src/log-store/src/raft_engine/log_store.rs b/src/log-store/src/raft_engine/log_store.rs index 32a028f1691..dcc1109b569 100644 --- a/src/log-store/src/raft_engine/log_store.rs +++ b/src/log-store/src/raft_engine/log_store.rs @@ -82,17 +82,60 @@ impl TaskFunction for PurgeExpiredFilesFunction { } } +/// Key used by the periodic WAL sync task. Overwritten each period so the +/// write stays non-empty on raft-engine 0.4.2 (see [`durable_sync`]). +const PERIODIC_SYNC_MARKER_KEY: &[u8] = b"$wal_periodic_sync"; + +/// Force an fsync of previously appended WAL data on raft-engine 0.4.2. +/// +/// `Engine::sync()` writes an empty [`LogBatch`] with `sync = true`, but +/// 0.4.2's `Engine::write` returns early for empty batches *before* the +/// fsync (tikv/raft-engine#395). Upstream fixed this in tikv/raft-engine#400, +/// which is not on crates.io yet. Issuing a minimal non-empty put with +/// `sync = true` reaches the fsync path on 0.4.2. +/// +/// `batch` must be empty and can be reused across calls: a successful +/// `Engine::write` drains its items but keeps its encoding buffer. On error the +/// batch may be left partially encoded, so callers must discard it. +fn durable_sync(engine: &Engine, batch: &mut LogBatch) -> raft_engine::Result { + debug_assert!(batch.is_empty()); + // The value is unused, but it must not decode as a `NamespaceImpl`: the + // marker key sorts after `NAMESPACE_PREFIX`, inside the range scanned by + // `list_namespaces`. Overwriting the same key keeps the marker bounded. + batch.put(SYSTEM_NAMESPACE, PERIODIC_SYNC_MARKER_KEY.to_vec(), vec![0])?; + engine.write(batch, true) +} + pub struct SyncWalTaskFunction { engine: Arc, + /// Batch reused across ticks to avoid allocating a new one each period. + /// `None` after a failed sync; a fresh batch is created on the next tick. + batch: Option, } #[async_trait::async_trait] impl TaskFunction for SyncWalTaskFunction { async fn call(&mut self) -> std::result::Result<(), Error> { let engine = self.engine.clone(); - if let Err(e) = tokio::task::spawn_blocking(move || engine.sync()).await { - error!(e; "Failed to sync raft engine log files"); - }; + let mut batch = self + .batch + .take() + .unwrap_or_else(|| LogBatch::with_capacity(1)); + match tokio::task::spawn_blocking(move || { + let res = durable_sync(&engine, &mut batch); + (res, batch) + }) + .await + { + Ok((Ok(_), batch)) => self.batch = Some(batch), + // The batch may be partially encoded after a failed write; drop it. + Ok((Err(e), _)) => { + error!(e; "Failed to sync raft engine log files"); + } + Err(e) => { + error!(e; "Failed to sync raft engine log files"); + } + } Ok(()) } @@ -103,7 +146,10 @@ impl TaskFunction for SyncWalTaskFunction { impl SyncWalTaskFunction { pub fn new(engine: Arc) -> Self { - Self { engine } + Self { + engine, + batch: None, + } } } @@ -555,6 +601,74 @@ mod tests { assert_eq!(0, namespaces.len()); } + /// Regression for #9712: periodic sync must write a non-empty batch so + /// raft-engine 0.4.2 actually fsyncs (empty `Engine::sync()` is a no-op). + #[tokio::test] + async fn test_durable_sync_writes_non_empty_batch() { + let dir = create_temp_dir("raft-engine-durable-sync-test"); + let cfg = Config { + dir: dir.path().to_str().unwrap().to_string(), + recovery_mode: RecoveryMode::TolerateTailCorruption, + ..Default::default() + }; + let engine = Engine::open(cfg).unwrap(); + + // Control: empty sync leaves no marker and returns Ok without fsync + // on 0.4.2 (we cannot count fsyncs here, but we can show the empty + // path writes nothing). + engine.sync().unwrap(); + assert!( + engine + .get(SYSTEM_NAMESPACE, PERIODIC_SYNC_MARKER_KEY) + .is_none(), + "Engine::sync must not create the periodic sync marker" + ); + + let mut batch = LogBatch::with_capacity(1); + let written = durable_sync(&engine, &mut batch).unwrap(); + assert!(written > 0, "durable_sync must write a non-empty batch"); + assert!(batch.is_empty(), "a successful write drains the batch"); + assert_eq!( + engine + .get(SYSTEM_NAMESPACE, PERIODIC_SYNC_MARKER_KEY) + .as_deref(), + Some(&[0][..]), + "durable_sync must persist the periodic sync marker" + ); + + // Reusing the drained batch overwrites the same key and still writes. + let written_again = durable_sync(&engine, &mut batch).unwrap(); + assert!(written_again > 0); + assert!(batch.is_empty()); + } + + #[tokio::test] + async fn test_sync_wal_task_reuses_batch_and_keeps_namespaces() { + let dir = create_temp_dir("raft-engine-sync-wal-task-test"); + let logstore = RaftEngineLogStore::try_new( + dir.path().to_str().unwrap().to_string(), + &RaftEngineConfig::default(), + ) + .await + .unwrap(); + logstore + .create_namespace(&Provider::raft_engine_provider(42)) + .await + .unwrap(); + + let mut task = SyncWalTaskFunction::new(logstore.engine.clone()); + for _ in 0..3 { + task.call().await.unwrap(); + let batch = task.batch.as_ref().expect("batch is kept after a sync"); + assert!(batch.is_empty()); + } + + // The marker lies in the range scanned by `list_namespaces` and must + // not be reported as a namespace. + let namespaces = logstore.list_namespaces().await.unwrap(); + assert_eq!(vec![Provider::raft_engine_provider(42)], namespaces); + } + #[tokio::test] async fn test_manage_namespace() { let dir = create_temp_dir("raft-engine-logstore-test"); From 3c490bd2dca1c7973a876f2fde6c4ce3b614d2ab Mon Sep 17 00:00:00 2001 From: Dhruv Vaishnav Date: Wed, 28 Oct 2026 02:37:49 +0000 Subject: [PATCH 114/239] ci: reuse fuzz corpora across integration runs (#9979) ci: persist fuzz corpora between runs Signed-off-by: dhruvxvaishnav --- .github/actions/fuzz-test/action.yaml | 40 +++++++ .github/scripts/prepare-fuzz-corpus.sh | 38 +++++++ .github/scripts/run-fuzz-targets-test.sh | 128 ++++++++++++++++++++++- .github/scripts/run-fuzz-targets.sh | 10 +- tests-fuzz/README.md | 81 ++++++++++++++ 5 files changed, 292 insertions(+), 5 deletions(-) create mode 100644 .github/scripts/prepare-fuzz-corpus.sh diff --git a/.github/actions/fuzz-test/action.yaml b/.github/actions/fuzz-test/action.yaml index b696d8099dc..b93634b16a2 100644 --- a/.github/actions/fuzz-test/action.yaml +++ b/.github/actions/fuzz-test/action.yaml @@ -28,7 +28,38 @@ inputs: runs: using: composite steps: + - name: Prepare Fuzz Corpus + id: corpus + shell: bash + env: + FUZZ_TARGETS: ${{ inputs.targets }} + run: bash .github/scripts/prepare-fuzz-corpus.sh + + - name: Restore Fuzz Corpus + id: restore-corpus + uses: actions/cache/restore@v4 + continue-on-error: true + with: + path: ${{ steps.corpus.outputs.paths }} + key: ${{ steps.corpus.outputs.key-prefix }}${{ github.run_id }}-${{ github.run_attempt }}-${{ github.job }}-${{ strategy.job-index || 0 }} + restore-keys: ${{ steps.corpus.outputs.key-prefix }} + + - name: Report Fuzz Corpus Cache + shell: bash + env: + CACHE_OUTCOME: ${{ steps.restore-corpus.outcome }} + CACHE_MATCHED_KEY: ${{ steps.restore-corpus.outputs.cache-matched-key }} + run: | + if [[ "${CACHE_OUTCOME}" == failure ]]; then + echo '::warning::Fuzz corpus restore failed; running with available inputs' + elif [[ -n "${CACHE_MATCHED_KEY}" ]]; then + printf 'Restored fuzz corpus: %s\n' "${CACHE_MATCHED_KEY}" + else + echo 'No cached fuzz corpus found; starting with an empty corpus' + fi + - name: Run Fuzz Tests + id: run-fuzz shell: bash env: FUZZ_TARGETS: ${{ inputs.targets }} @@ -39,4 +70,13 @@ runs: FUZZ_COLLECT_CLUSTER_ARTIFACTS: ${{ inputs.collect-cluster-artifacts }} FUZZ_ARTIFACT_ROOT: ${{ inputs.artifact-root }} FUZZ_BIN_DIR: ${{ inputs.prebuilt-binaries-dir }} + FUZZ_CORPUS_ROOT: ${{ github.workspace }}/tests-fuzz/corpus run: .github/scripts/run-fuzz-targets.sh + + - name: Save Fuzz Corpus + if: ${{ !cancelled() && (steps.run-fuzz.outcome == 'success' || steps.run-fuzz.outcome == 'failure') }} + uses: actions/cache/save@v4 + continue-on-error: true + with: + path: ${{ steps.corpus.outputs.paths }} + key: ${{ steps.corpus.outputs.key-prefix }}${{ github.run_id }}-${{ github.run_attempt }}-${{ github.job }}-${{ strategy.job-index || 0 }} diff --git a/.github/scripts/prepare-fuzz-corpus.sh b/.github/scripts/prepare-fuzz-corpus.sh new file mode 100644 index 00000000000..e647cf9ceaf --- /dev/null +++ b/.github/scripts/prepare-fuzz-corpus.sh @@ -0,0 +1,38 @@ +#!/usr/bin/env bash + +set -euo pipefail + +: "${FUZZ_TARGETS:?FUZZ_TARGETS is required}" +: "${GITHUB_WORKSPACE:?GITHUB_WORKSPACE is required}" +: "${GITHUB_OUTPUT:?GITHUB_OUTPUT is required}" + +targets=() +while IFS= read -r target; do + target="${target%$'\r'}" + [[ -z "${target}" ]] && continue + if [[ ! "${target}" =~ ^[a-zA-Z0-9_]+$ ]]; then + printf 'Invalid fuzz target: %s\n' "${target}" >&2 + exit 2 + fi + targets+=("${target}") +done <<<"${FUZZ_TARGETS}" + +if [[ "${#targets[@]}" -eq 0 ]]; then + printf 'FUZZ_TARGETS contains no targets\n' >&2 + exit 2 +fi + +# Corpus identity depends on the targets, not the source revision or their order. +target_set="$(printf '%s\n' "${targets[@]}" | LC_ALL=C sort -u)" +target_hash="$(printf '%s\n' "${target_set}" | sha256sum | cut -d ' ' -f 1)" +corpus_root="${GITHUB_WORKSPACE}/tests-fuzz/corpus" + +{ + printf 'key-prefix=fuzz-corpus-v1-%s-\n' "${target_hash}" + printf 'paths<>"${GITHUB_OUTPUT}" diff --git a/.github/scripts/run-fuzz-targets-test.sh b/.github/scripts/run-fuzz-targets-test.sh index 9705b6bc5af..ddc559c8db2 100755 --- a/.github/scripts/run-fuzz-targets-test.sh +++ b/.github/scripts/run-fuzz-targets-test.sh @@ -5,6 +5,7 @@ set -euo pipefail SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" RUNNER="${SCRIPT_DIR}/run-fuzz-targets.sh" COLLECTOR="${SCRIPT_DIR}/collect-fuzz-target-artifacts.sh" +PREPARE_CORPUS="${SCRIPT_DIR}/prepare-fuzz-corpus.sh" fail() { printf 'FAIL: %s\n' "$*" >&2 @@ -32,6 +33,13 @@ new_fixture() { set -euo pipefail printf '%s\t%s\t%s\n' "${GT_FUZZ_DUMP_DIR}" "$*" "${MOCK_CARGO_MARKER:-}" >>"${MOCK_CARGO_LOG}" target="$4" +corpus="$5" +[[ -d "${corpus}" ]] || exit 18 +printf '%s\t%s\n' "${target}" "${corpus}" >>"${MOCK_CORPUS_LOG}" +if [[ -f "${corpus}/seed" ]]; then + cat "${corpus}/seed" >>"${MOCK_SEED_LOG}" +fi +printf '%s\n' "${target}" >"${corpus}/learned" case " ${MOCK_FAIL_TARGETS:-} " in *" ${target} "*) exit 17 ;; esac @@ -64,12 +72,15 @@ run_fixture() { "MOCK_COLLECTOR_LOG=${fixture}/collector.log" "MOCK_FAIL_TARGETS=${fail_targets}" "MOCK_CARGO_MARKER=fixture" + "MOCK_CORPUS_LOG=${fixture}/corpus.log" + "MOCK_SEED_LOG=${fixture}/seeds.log" "FUZZ_TARGETS=${targets}" "FUZZ_GROUP=test-group" "FUZZ_MAX_TOTAL_TIME=120" "FUZZ_UNSTABLE=${unstable}" "FUZZ_COLLECT_CLUSTER_ARTIFACTS=true" "FUZZ_ARTIFACT_ROOT=${artifact_root}" + "FUZZ_CORPUS_ROOT=${fixture}/corpus with spaces" "FUZZ_ARTIFACT_COLLECTOR=${fixture}/collector" "GT_FUZZ_BINARY_PATH=${GT_FUZZ_BINARY_PATH:-}" "GT_FUZZ_INSTANCE_ROOT_DIR=${GT_FUZZ_INSTANCE_ROOT_DIR:-}" @@ -81,7 +92,7 @@ run_fixture() { fi set +e - env -u FUZZ_FAIL_FAST "${runner_env[@]}" \ + env -u FUZZ_FAIL_FAST -u FUZZ_BIN_DIR "${runner_env[@]}" \ "${RUNNER}" >"${fixture}/stdout.log" 2>&1 fixture_status=$? set -e @@ -89,18 +100,30 @@ run_fixture() { run_prebuilt_fixture() { new_fixture - mkdir -p "${fixture}/prebuilt" + mkdir -p "${fixture}/prebuilt" "${fixture}/corpus with spaces/fuzz_create_table" + printf 'restored prebuilt seed\n' >"${fixture}/corpus with spaces/fuzz_create_table/seed" cat >"${fixture}/prebuilt/fuzz_create_table" <<'EOF' #!/usr/bin/env bash set -euo pipefail printf '%s\t%s\t%s\n' "${GT_FUZZ_DUMP_DIR}" "$*" "${MOCK_BINARY_MARKER:-}" >>"${MOCK_BINARY_LOG}" +[[ -d "$1" ]] || exit 18 +printf '%s\n' "$1" >>"${MOCK_CORPUS_LOG}" +cat "$1/seed" >>"${MOCK_SEED_LOG}" +printf 'new prebuilt input\n' >"$1/learned" +exit "${MOCK_BINARY_EXIT_CODE}" EOF chmod +x "${fixture}/prebuilt/fuzz_create_table" set +e env \ + "PATH=${fixture}/bin:${PATH}" \ + "MOCK_CARGO_LOG=${fixture}/cargo.log" \ "MOCK_BINARY_LOG=${fixture}/binary.log" \ "MOCK_BINARY_MARKER=fixture" \ + "MOCK_BINARY_EXIT_CODE=${1:-0}" \ + "MOCK_CORPUS_LOG=${fixture}/corpus.log" \ + "MOCK_SEED_LOG=${fixture}/seeds.log" \ + "MOCK_COLLECTOR_LOG=${fixture}/collector.log" \ "FUZZ_TARGETS=fuzz_create_table" \ "FUZZ_GROUP=test-group" \ "FUZZ_MAX_TOTAL_TIME=120" \ @@ -108,6 +131,7 @@ EOF "FUZZ_FAIL_FAST=true" \ "FUZZ_COLLECT_CLUSTER_ARTIFACTS=false" \ "FUZZ_ARTIFACT_ROOT=${fixture}/artifacts" \ + "FUZZ_CORPUS_ROOT=${fixture}/corpus with spaces" \ "FUZZ_ARTIFACT_COLLECTOR=${fixture}/collector" \ "FUZZ_BIN_DIR=${fixture}/prebuilt" \ "GITHUB_SHA=deadbeef" \ @@ -143,6 +167,7 @@ test_successful_targets_run_in_order() { "target order" grep -q -- '--features=unstable' "${fixture}/cargo.log" || fail "unstable feature missing" grep -q -- '-max_total_time=120' "${fixture}/cargo.log" || fail "fuzz time missing" + grep -q -- '-print_final_stats=1' "${fixture}/cargo.log" || fail "fuzz statistics missing" grep -q -- '-artifact_prefix=.*/targets/fuzz_create_table/libfuzzer/' "${fixture}/cargo.log" || \ fail "target-scoped libFuzzer prefix missing" assert_eq success "$(jq -r '.[0].status' "${fixture}/artifacts/manifest.json")" "first status" @@ -167,8 +192,49 @@ test_prebuilt_binary_runs_without_cargo() { [[ ! -e "${fixture}/cargo.log" ]] || fail "cargo ran for prebuilt binary" assert_eq 1 "$(wc -l <"${fixture}/binary.log" | tr -d ' ')" "prebuilt invocation count" grep -q -- '-max_total_time=120' "${fixture}/binary.log" || fail "prebuilt fuzz time missing" + grep -q -- '-print_final_stats=1' "${fixture}/binary.log" || fail "prebuilt fuzz statistics missing" grep -q -- '-artifact_prefix=.*/targets/fuzz_create_table/libfuzzer/' "${fixture}/binary.log" || \ fail "prebuilt target-scoped libFuzzer prefix missing" + assert_eq "${fixture}/corpus with spaces/fuzz_create_table" "$(cat "${fixture}/corpus.log")" \ + "prebuilt corpus argument" + assert_eq 'restored prebuilt seed' "$(cat "${fixture}/seeds.log")" "prebuilt restored input" + assert_file "${fixture}/corpus with spaces/fuzz_create_table/learned" + cleanup_fixture +} + +test_cargo_reuses_separate_target_corpora() { + new_fixture + mkdir -p "${fixture}/corpus with spaces/fuzz_create_table" "${fixture}/corpus with spaces/fuzz_insert" + printf 'create seed\n' >"${fixture}/corpus with spaces/fuzz_create_table/seed" + printf 'insert seed\n' >"${fixture}/corpus with spaces/fuzz_insert/seed" + run_fixture $'fuzz_create_table\nfuzz_insert' false false "" + + assert_eq 0 "${fixture_status}" "restored cargo run status" + assert_eq $'create seed\ninsert seed' "$(cat "${fixture}/seeds.log")" "target-scoped restored inputs" + assert_eq \ + "fuzz_create_table"$'\t'"${fixture}/corpus with spaces/fuzz_create_table"$'\n'"fuzz_insert"$'\t'"${fixture}/corpus with spaces/fuzz_insert" \ + "$(cat "${fixture}/corpus.log")" "cargo corpus arguments" + assert_eq fuzz_create_table "$(cat "${fixture}/corpus with spaces/fuzz_create_table/learned")" \ + "first target learned input" + assert_eq fuzz_insert "$(cat "${fixture}/corpus with spaces/fuzz_insert/learned")" \ + "second target learned input" + + run_fixture fuzz_create_table true false "" + assert_eq 0 "${fixture_status}" "repeated cargo run status" + assert_file "${fixture}/corpus with spaces/fuzz_create_table/seed" + assert_file "${fixture}/corpus with spaces/fuzz_insert/learned" + cleanup_fixture +} + +test_prebuilt_failure_preserves_corpus_and_artifacts() { + run_prebuilt_fixture 17 + + assert_eq 1 "${fixture_status}" "prebuilt failure status" + assert_eq 17 "$(jq -r '.[0].exit_code' "${fixture}/artifacts/manifest.json")" \ + "prebuilt failure exit code" + assert_file "${fixture}/corpus with spaces/fuzz_create_table/seed" + assert_file "${fixture}/corpus with spaces/fuzz_create_table/learned" + assert_file "${fixture}/artifacts/targets/fuzz_create_table/mock-artifacts/state.txt" cleanup_fixture } @@ -183,6 +249,8 @@ test_fail_fast_stops_after_first_failure() { assert_eq skipped_after_failure "$(jq -r '.[2].status' "${fixture}/artifacts/manifest.json")" "skipped status" assert_eq true "$(jq -r '.[2].after_prior_failure' "${fixture}/artifacts/manifest.json")" "skipped provenance" assert_file "${fixture}/artifacts/targets/fuzz_insert/mock-artifacts/state.txt" + assert_file "${fixture}/corpus with spaces/fuzz_insert/learned" + [[ ! -e "${fixture}/corpus with spaces/fuzz_alter_table" ]] || fail "skipped target corpus created" grep -q 'title=Fuzz target failed.*target=fuzz_insert' "${fixture}/stdout.log" || \ fail "target failure annotation missing" grep -q 'title=Fuzz target skipped.*target=fuzz_alter_table' "${fixture}/stdout.log" || \ @@ -191,6 +259,8 @@ test_fail_fast_stops_after_first_failure() { fail "reproduction section missing" grep -q 'cargo +nightly fuzz run fuzz_insert' "${fixture}/artifacts/summary.md" || \ fail "reproduction command missing" + grep -q 'fuzz_insert --fuzz-dir tests-fuzz' "${fixture}/artifacts/summary.md" || \ + fail "reproduction command depends on the CI corpus path" cleanup_fixture } @@ -241,6 +311,56 @@ test_invalid_configuration_fails_before_cargo() { cleanup_fixture } +prepare_corpus_fixture() { + local targets="$1" + local output="$2" + set +e + FUZZ_TARGETS="${targets}" \ + GITHUB_WORKSPACE="${fixture}/workspace with spaces" \ + GITHUB_OUTPUT="${output}" \ + bash "${PREPARE_CORPUS}" >"${fixture}/prepare.log" 2>&1 + prepare_status=$? + set -e +} + +test_corpus_cache_identity_and_paths() { + new_fixture + prepare_corpus_fixture $'fuzz_insert\r\n\nfuzz_create_table\r\nfuzz_insert\r' "${fixture}/first-output" + assert_eq 0 "${prepare_status}" "corpus preparation status" + prepare_corpus_fixture $'fuzz_create_table\nfuzz_insert' "${fixture}/second-output" + assert_eq 0 "${prepare_status}" "reordered corpus preparation status" + cmp -s "${fixture}/first-output" "${fixture}/second-output" || \ + fail "cache identity depends on target order, duplicates, or line endings" + assert_eq \ + "${fixture}/workspace with spaces/tests-fuzz/corpus/fuzz_create_table"$'\n'"${fixture}/workspace with spaces/tests-fuzz/corpus/fuzz_insert" \ + "$(sed '1,2d; $d' "${fixture}/first-output")" "selected corpus cache paths" + [[ -d "${fixture}/workspace with spaces/tests-fuzz/corpus/fuzz_insert" ]] || \ + fail "corpus directory not created" + [[ ! -e "${fixture}/workspace with spaces/tests-fuzz/corpus/fuzz_alter_table" ]] || \ + fail "unselected corpus directory created" + + GITHUB_SHA=another-revision prepare_corpus_fixture $'fuzz_create_table\nfuzz_insert' "${fixture}/new-revision-output" + cmp -s "${fixture}/first-output" "${fixture}/new-revision-output" || \ + fail "source revision changed cache identity" + prepare_corpus_fixture fuzz_insert "${fixture}/subset-output" + assert_eq 0 "${prepare_status}" "subset corpus preparation status" + [[ "$(head -n 1 "${fixture}/first-output")" != "$(head -n 1 "${fixture}/subset-output")" ]] || \ + fail "different target sets share cache identity" + cleanup_fixture +} + +test_corpus_preparation_rejects_invalid_targets_before_writing() { + local targets + for targets in $'fuzz_insert\n../escape' 'fuzz target' $'\n\r\n'; do + new_fixture + prepare_corpus_fixture "${targets}" "${fixture}/output" + assert_eq 2 "${prepare_status}" "invalid corpus target status" + [[ ! -e "${fixture}/output" ]] || fail "invalid targets wrote cache outputs" + [[ ! -e "${fixture}/workspace with spaces" ]] || fail "invalid targets created directories" + cleanup_fixture + done +} + test_collector_keeps_target_scopes_separate() { new_fixture service_log="${fixture}/greptime.log" @@ -373,10 +493,14 @@ EOF test_successful_targets_run_in_order test_prebuilt_binary_runs_without_cargo +test_cargo_reuses_separate_target_corpora +test_prebuilt_failure_preserves_corpus_and_artifacts test_fail_fast_stops_after_first_failure test_fail_fast_defaults_to_true test_non_fail_fast_marks_later_targets test_invalid_configuration_fails_before_cargo +test_corpus_cache_identity_and_paths +test_corpus_preparation_rejects_invalid_targets_before_writing test_collector_keeps_target_scopes_separate test_cluster_collector_honors_target_scope_and_namespace test_setup_failure_writes_artifact_contract diff --git a/.github/scripts/run-fuzz-targets.sh b/.github/scripts/run-fuzz-targets.sh index 205ba58c419..1489d20b7f0 100755 --- a/.github/scripts/run-fuzz-targets.sh +++ b/.github/scripts/run-fuzz-targets.sh @@ -5,6 +5,7 @@ set -euo pipefail FUZZ_ARTIFACT_ROOT="${FUZZ_ARTIFACT_ROOT:-/tmp/greptime-fuzz-artifacts}" FUZZ_ARTIFACT_COLLECTOR="${FUZZ_ARTIFACT_COLLECTOR:-.github/scripts/collect-fuzz-target-artifacts.sh}" FUZZ_ARTIFACT_COLLECTION_TIMEOUT_SECS="${FUZZ_ARTIFACT_COLLECTION_TIMEOUT_SECS:-180}" +FUZZ_CORPUS_ROOT="${FUZZ_CORPUS_ROOT:-tests-fuzz/corpus}" log() { printf '[run-fuzz-targets] %s\n' "$*" @@ -187,19 +188,22 @@ for ((index = 0; index < ${#targets[@]}; index++)); do started_epoch="$(date +%s)" echo "::group::Fuzz target: ${target}" log "run target ${target}" + mkdir -p "${FUZZ_CORPUS_ROOT}/${target}" + corpus_dir="$(cd -- "${FUZZ_CORPUS_ROOT}/${target}" && pwd)" + log "corpus directory: ${corpus_dir}" if [[ -n "${FUZZ_BIN_DIR:-}" ]]; then fuzz_binary="${FUZZ_BIN_DIR}/${target}" if [[ ! -x "${fuzz_binary}" ]]; then log "prebuilt fuzz binary does not exist or is not executable: ${fuzz_binary}" exit 2 fi - run_args=("${fuzz_binary}" "-max_total_time=${FUZZ_MAX_TOTAL_TIME}" "-artifact_prefix=${target_dir}/libfuzzer/") + run_args=("${fuzz_binary}" "${corpus_dir}" "-max_total_time=${FUZZ_MAX_TOTAL_TIME}" "-print_final_stats=1" "-artifact_prefix=${target_dir}/libfuzzer/") else - run_args=(cargo +"${RUSTUP_TOOLCHAIN:-nightly}" fuzz run "${target}" --fuzz-dir tests-fuzz -D -s none) + run_args=(cargo +"${RUSTUP_TOOLCHAIN:-nightly}" fuzz run "${target}" "${corpus_dir}" --fuzz-dir tests-fuzz -D -s none) if [[ "${fuzz_unstable}" == true ]]; then run_args+=(--features=unstable) fi - run_args+=(-- "-max_total_time=${FUZZ_MAX_TOTAL_TIME}" "-artifact_prefix=${target_dir}/libfuzzer/") + run_args+=(-- "-max_total_time=${FUZZ_MAX_TOTAL_TIME}" "-print_final_stats=1" "-artifact_prefix=${target_dir}/libfuzzer/") fi if GT_FUZZ_DUMP_DIR="${target_dir}/csv" \ "${run_args[@]}" 2>&1 | tee "${target_dir}/fuzz.log"; then diff --git a/tests-fuzz/README.md b/tests-fuzz/README.md index 7945bae1583..f39a23f2887 100644 --- a/tests-fuzz/README.md +++ b/tests-fuzz/README.md @@ -42,6 +42,87 @@ cargo fuzz list --fuzz-dir tests-fuzz cargo fuzz run fuzz_create_table --fuzz-dir tests-fuzz -D -s none ``` +## CI Corpus Persistence + +The shared `.github/actions/fuzz-test` action restores inputs before running +the selected targets and saves the updated corpus afterward, including when a +target fails. Cache failures are best-effort and do not change the fuzz result. +Crash files, diagnostic artifacts, and fail-fast behavior are unchanged. + +Each target reads and writes its own `tests-fuzz/corpus/` directory. +Both prebuilt executables and the `cargo fuzz` fallback receive that directory +explicitly. For local runs of `.github/scripts/run-fuzz-targets.sh`, set +`FUZZ_CORPUS_ROOT` to use a different corpus root. + +The selected target directories are cached together to preserve the existing +job grouping. The stable key prefix identifies the sorted, deduplicated target +set, not a source hash, so code changes do not invalidate existing inputs. +Each job saves a new immutable entry, distinguished by run ID, attempt, job, +and matrix index. Restore uses the latest accessible entry matching the prefix. +Changing the target set starts a different cache; corpora are not shared between +different target sets. Parallel jobs may start from the same entry, but their +new inputs are not merged automatically. + +The restore step logs the matched key (including prefix matches) or a cache +miss. GitHub's [cache access restrictions](https://docs.github.com/en/actions/reference/workflows-and-actions/dependency-caching#restrictions-for-accessing-a-cache) +still apply: PR caches are scoped to that PR, while a default-branch cache can +seed other PRs. The existing scheduled workflow does not run fuzz jobs; a +maintainer can populate the default-branch corpus through the existing manual +fuzz workflow. This change does not add scheduling, target selection, or corpus +pruning. Rolling entries remain subject to GitHub's cache limits and eviction. + +### Validation + +The orchestration tests use mock fuzz executables and do not build Rust or +start GreptimeDB. They cover both runner paths, restored inputs, separate +target directories, retained inputs after failure, cache identity, and invalid +target names. Run them on Linux with Bash, jq, and GNU coreutils: + +```bash +bash .github/scripts/run-fuzz-targets-test.sh +``` + +These tests do not exercise the hosted GitHub cache service. To verify that +service, inspect two fuzz workflow runs on the same PR: the first should save +an entry, and the second should report its restored key and load a nonempty +corpus. Also verify that cache save still runs after a target failure, without +turning that failure into a successful fuzz result. + +To compare useful coverage, use the same fuzz executable, GreptimeDB revision, +service configuration, and time budget for empty and restored corpora. Both +runner paths enable `-print_final_stats=1`; the target's `fuzz.log` retains the +libFuzzer coverage and final statistics. Seed the restored corpus with a +separate warm-up run (or a downloaded corpus), then compare it with a fresh, +empty directory. Keep warm-up time separate from the comparison and repeat +trials to account for randomness. Report initial corpus size, final covered +edges (`cov:`), executions, and elapsed time, rather than inferring a coverage +improvement merely from a cache hit. + +For example, with a prebuilt `fuzz_create_table` binary and a running standalone +service, the following creates separate warm-up, empty, and seeded corpora. Set +`FUZZ_BIN_DIR` to the absolute binary directory and `GT_MYSQL_ADDR` to the test +service address first. All runs use the same executable and 120-second budget; +the warm-up is not part of the measured comparison. This creates temporary +inputs, logs, and crash files, and executes DDL against the test service. + +```bash +set -euo pipefail +: "${FUZZ_BIN_DIR:?Set the prebuilt fuzz binary directory}" +: "${GT_MYSQL_ADDR:?Set the test service address}" +comparison_dir="$(mktemp -d -t fuzz-corpus-comparison.XXXXXX)" +for trial in seed cold warm; do + mkdir -p "${comparison_dir}/${trial}" "${comparison_dir}/artifacts/${trial}" + if [[ "${trial}" == warm ]]; then + cp -a "${comparison_dir}/seed/." "${comparison_dir}/warm/" + fi + "${FUZZ_BIN_DIR}/fuzz_create_table" "${comparison_dir}/${trial}" \ + -max_total_time=120 -print_final_stats=1 \ + "-artifact_prefix=${comparison_dir}/artifacts/${trial}/" \ + 2>&1 | tee "${comparison_dir}/${trial}.log" +done +printf 'Comparison logs and corpora: %s\n' "${comparison_dir}" +``` + ## Crash Reproduction If you want to reproduce a crash, you first need to obtain the Base64 encoded code, which usually appears at the end of a crash report, and store it in a file. From b369d35d8b8a841bac9fae108508be8c24e354b8 Mon Sep 17 00:00:00 2001 From: LFC <990479+MichaelScofield@users.noreply.github.com> Date: Wed, 28 Oct 2026 02:43:29 +0000 Subject: [PATCH 115/239] feat: enable JSON2 columns on non-append-only tables (#9985) feat(mito2): support JSON2 in non-append-only tables Signed-off-by: luofucong --- src/mito2/src/read/compat.rs | 177 ++++- src/mito2/src/read/flat_projection.rs | 159 ++--- src/mito2/src/read/last_row.rs | 67 +- src/mito2/src/read/scan_util.rs | 215 +++++- src/mito2/src/read/seq_scan.rs | 6 +- src/mito2/src/read/series_reader.rs | 7 +- src/mito2/src/read/unordered_scan.rs | 4 +- src/mito2/src/sst/parquet/file_range.rs | 73 +- src/mito2/src/sst/parquet/json_align.rs | 75 +- src/mito2/src/sst/parquet/reader.rs | 8 + src/mito2/src/test_util/sst_util.rs | 57 ++ src/operator/src/statement/ddl.rs | 75 +- src/query/src/datafusion.rs | 61 +- .../common/types/json/json2_limit.result | 53 +- .../common/types/json/json2_limit.sql | 33 +- .../common/types/json/json2_non_append.result | 675 ++++++++++++++++++ .../common/types/json/json2_non_append.sql | 201 ++++++ 17 files changed, 1655 insertions(+), 291 deletions(-) create mode 100644 tests/cases/standalone/common/types/json/json2_non_append.result create mode 100644 tests/cases/standalone/common/types/json/json2_non_append.sql diff --git a/src/mito2/src/read/compat.rs b/src/mito2/src/read/compat.rs index acc05f0e9d2..6f9a4a36e50 100644 --- a/src/mito2/src/read/compat.rs +++ b/src/mito2/src/read/compat.rs @@ -42,9 +42,9 @@ use crate::error::{ NewRecordBatchSnafu, Result, UnexpectedSnafu, UnsupportedOperationSnafu, }; use crate::read::flat_projection::{FlatProjectionMapper, flat_projected_columns}; +use crate::sst::internal_fields; use crate::sst::parquet::flat_format::{FlatReadFormat, primary_key_column_index}; use crate::sst::parquet::format::{INTERNAL_COLUMN_NUM, PrimaryKeyArray}; -use crate::sst::{internal_fields, tag_maybe_to_dictionary_field, with_field_id}; /// Returns true if the columns in the `projection_mapper` and `read_format` have same data types /// and primary key encodings. @@ -116,8 +116,12 @@ impl FlatCompatBatch { return FlatCompatBatch::try_new_compact_sparse(mapper, actual); } - let (index_or_defaults, fields) = - Self::compute_index_and_fields(&actual_schema, expect_schema, mapper.metadata())?; + let (index_or_defaults, fields) = Self::compute_index_and_fields( + &actual_schema, + expect_schema, + mapper.metadata(), + mapper.input_arrow_schema(compaction).fields(), + )?; let compat_pk = FlatCompatPrimaryKey::new(mapper.metadata(), actual)?; @@ -132,6 +136,7 @@ impl FlatCompatBatch { actual_schema: &[(ColumnId, ConcreteDataType)], expect_schema: &[(ColumnId, ConcreteDataType)], expect_metadata: &RegionMetadata, + target_fields: &[FieldRef], ) -> Result<(Vec, Vec)> { // Maps column id to the index and data type in the actual schema. let actual_schema_index: HashMap<_, _> = actual_schema @@ -142,29 +147,18 @@ impl FlatCompatBatch { let mut index_or_defaults = Vec::with_capacity(expect_schema.len()); let mut fields = Vec::with_capacity(expect_schema.len()); - for (column_id, expect_data_type) in expect_schema { + for (field_index, (column_id, expect_data_type)) in expect_schema.iter().enumerate() { // Safety: expect_schema comes from the same mapper. let column_index = expect_metadata.column_index_by_id(*column_id).unwrap(); let expect_column = &expect_metadata.column_metadatas[column_index]; - let column_field = &expect_metadata.schema.arrow_schema().fields()[column_index]; - // For tag columns, we need to create a dictionary field. - if expect_column.semantic_type == SemanticType::Tag { - let field = tag_maybe_to_dictionary_field( - &expect_column.column_schema.data_type, - column_field, - ); - fields.push(Arc::new(with_field_id( - (*field).clone(), - expect_column.column_id, - ))); - } else { - let field = with_field_id( - Arc::unwrap_or_clone(column_field.clone()), - expect_column.column_id, - ) - .with_data_type(expect_data_type.as_arrow_type()); - fields.push(Arc::new(field)) - }; + // Query target fields follow mapper.batch_schema(); compaction reads all columns in + // PK -> fields -> timestamp order (fields -> timestamp for sparse encoding), before + // trailing internal columns. Keep this order consistent with expect_schema. + // The mapper already applies string-tag dictionary types and JSON2 compaction layouts, + // which can differ from the logical types in expect_schema. + let column_field = &target_fields[field_index]; + debug_assert_eq!(column_field.name(), &expect_column.column_schema.name); + fields.push(column_field.clone()); if let Some((index, actual_data_type)) = actual_schema_index.get(column_id) { let mut cast_type = None; @@ -194,14 +188,13 @@ impl FlatCompatBatch { // Create a default vector with 1 element for that column. let default_vector = if expect_data_type.is_json2() { // JSON2 columns do not support defaults, so missing columns use NULL. - // `expect_data_type` may be a query-concretized JSON2 type, - // while region metadata still carries its unconcretized type. - // The default vector must match the target Arrow field above. - Helper::try_into_vector(new_null_array(&expect_data_type.as_arrow_type(), 1)) - .context(CreateDefaultSnafu { + // A logical JSON2 placeholder cannot describe a compaction layout. + Helper::try_into_vector(new_null_array(column_field.data_type(), 1)).context( + CreateDefaultSnafu { region_id: expect_metadata.region_id, column: &expect_column.column_schema.name, - })? + }, + )? } else { expect_column .column_schema @@ -256,8 +249,12 @@ impl FlatCompatBatch { .map(|col| (col.column_id, col.column_schema.data_type.clone())) .collect(); - let (index_or_defaults, fields) = - Self::compute_index_and_fields(&actual_schema, &expect_schema, mapper.metadata())?; + let (index_or_defaults, fields) = Self::compute_index_and_fields( + &actual_schema, + &expect_schema, + mapper.metadata(), + mapper.input_arrow_schema(true).fields(), + )?; let compat_pk = FlatCompatPrimaryKey::default(); @@ -663,10 +660,15 @@ mod tests { }; use datatypes::arrow::datatypes::UInt32Type; use datatypes::arrow::record_batch::RecordBatch; + use datatypes::extension::json::JsonMetadata; + use datatypes::json::JsonSettings; use datatypes::prelude::ConcreteDataType; - use datatypes::schema::ColumnSchema; + use datatypes::schema::{ColumnDefaultConstraint, ColumnSchema}; use datatypes::types::json_type::JsonNativeType; + use datatypes::types::parse_string_to_jsonb; use datatypes::value::ValueRef; + use datatypes::vectors::json::array::JsonArray; + use datatypes::vectors::json::json2_physical_data_type; use mito_codec::row_converter::{ DensePrimaryKeyCodec, PrimaryKeyCodecExt, SparsePrimaryKeyCodec, }; @@ -678,7 +680,9 @@ mod tests { use crate::read::flat_projection::FlatProjectionMapper; use crate::read::read_columns::ReadColumns; use crate::sst::parquet::flat_format::FlatReadFormat; + use crate::sst::parquet::{Json2RewriteTargets, Json2TargetLayout}; use crate::sst::{FlatSchemaOptions, to_flat_sst_arrow_schema}; + use crate::test_util; /// Creates a new [RegionMetadata]. fn new_metadata( @@ -1484,4 +1488,113 @@ mod tests { assert_eq!(expected_batch, result); } + #[test] + fn test_compaction_json2_uses_target_fields() + -> std::result::Result<(), Box> { + let (metadata, batch) = test_util::sst_util::new_json2_record_batch()?; + let logical_settings = JsonSettings::default(); + let layout = JsonSettings::try_new(vec![], Some(0))?; + let targets: Json2RewriteTargets = Arc::new(BTreeMap::from([( + 0, + Json2TargetLayout { + extension_metadata: serde_json::to_string(&JsonMetadata::new( + logical_settings.clone(), + ))?, + target_layout: layout.clone(), + }, + )])); + // Model the existing Rewrite output; the compatibility layer only selects or fills arrays. + let json = JsonArray::from(batch.column(0)).rewrite_to_v2( + batch.schema_ref().field(0), + &logical_settings, + &layout, + )?; + let mut fields = batch.schema_ref().fields().to_vec(); + fields[0] = Arc::new( + fields[0] + .as_ref() + .clone() + .with_data_type(json.data_type().clone()), + ); + let mut columns = batch.columns().to_vec(); + columns[0] = json; + let batch = RecordBatch::try_new(Arc::new(Schema::new(fields)), columns)?; + for encoding in [PrimaryKeyEncoding::Dense, PrimaryKeyEncoding::Sparse] { + for missing_json in [false, true] { + let source = if missing_json { + new_metadata( + &[( + 1, + SemanticType::Timestamp, + ConcreteDataType::timestamp_millisecond_datatype(), + )], + &[], + ) + } else { + metadata.as_ref().clone() + }; + let mut builder = RegionMetadataBuilder::from_existing(source); + builder.primary_key_encoding(encoding); + let actual = Arc::new(builder.build()?); + let mut builder = RegionMetadataBuilder::from_existing(metadata.as_ref().clone()); + builder.primary_key_encoding(encoding); + builder.push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "added", + ConcreteDataType::int64_datatype(), + true, + ) + .with_default_constraint(Some( + ColumnDefaultConstraint::Value(Value::Int64(42)), + ))?, + semantic_type: SemanticType::Field, + column_id: 2, + }); + let expected = Arc::new(builder.build()?); + let read_columns = ReadColumns::new([0, 1, 2]); + let mapper = FlatProjectionMapper::new_with_json2_rewrite_targets( + &expected, + vec![0, 1, 2], + read_columns.clone(), + &targets, + )?; + let read_format = FlatReadFormat::new(actual, read_columns, None, "test", true)?; + let compat = FlatCompatBatch::try_new(&mapper, &read_format, true)? + .ok_or("missing columns require compatibility")?; + let input = if missing_json { + batch.project(&[1, 2, 3, 4])? + } else { + batch.clone() + }; + let output = compat.compat(input)?; + assert_eq!( + output.column(0).data_type(), + &json2_physical_data_type(&layout) + ); + assert_eq!( + output.column(0).null_count(), + if missing_json { 2 } else { 0 } + ); + let defaults = output + .column(1) + .as_any() + .downcast_ref::() + .ok_or("expected int64 defaults")?; + assert_eq!(defaults.values().as_ref(), &[42, 42]); + if !missing_json { + let jsonb = JsonArray::from(output.column(0)).project_to_v2( + output.schema_ref().field(0), + &datatypes::arrow::datatypes::DataType::Binary, + )?; + let jsonb = jsonb + .as_any() + .downcast_ref::() + .ok_or("expected JSONB")?; + assert_eq!(jsonb.value(0), parse_string_to_jsonb(r#"{"a":1}"#)?); + assert_eq!(jsonb.value(1), parse_string_to_jsonb(r#"{"a":2}"#)?); + } + } + } + Ok(()) + } } diff --git a/src/mito2/src/read/flat_projection.rs b/src/mito2/src/read/flat_projection.rs index 1734ca364af..69df6256762 100644 --- a/src/mito2/src/read/flat_projection.rs +++ b/src/mito2/src/read/flat_projection.rs @@ -17,10 +17,7 @@ use std::sync::Arc; use api::v1::SemanticType; -use common_error::ext::BoxedError; -use common_recordbatch::error::{ - ArrowComputeSnafu, DataTypesSnafu, ExternalSnafu, NewDfRecordBatchSnafu, -}; +use common_recordbatch::error::{ArrowComputeSnafu, NewDfRecordBatchSnafu}; use common_recordbatch::{DfRecordBatch, RecordBatch}; use datatypes::arrow::array::{Array, AsArray}; use datatypes::arrow::datatypes::{DataType as ArrowDataType, Field}; @@ -28,15 +25,13 @@ use datatypes::extension::json::is_json2_extension_type; use datatypes::prelude::{ConcreteDataType, DataType}; use datatypes::schema::{Schema, SchemaRef}; use datatypes::value::Value; -use datatypes::vectors::Helper; -use datatypes::vectors::json::array::JsonArray; use datatypes::vectors::json::json2_physical_data_type; use snafu::{OptionExt, ResultExt}; use store_api::metadata::{RegionMetadata, RegionMetadataRef}; use store_api::storage::ColumnId; use crate::cache::CacheStrategy; -use crate::error::{InvalidRequestSnafu, RecordBatchSnafu, Result}; +use crate::error::{InvalidRequestSnafu, NewRecordBatchSnafu, Result}; use crate::read::projection::{read_column_ids_from_projection, repeated_vector_with_cache}; use crate::read::read_columns::{JsonTargetTypes, ReadColumns}; use crate::sst::parquet::Json2RewriteTargets; @@ -150,7 +145,7 @@ impl FlatProjectionMapper { flat_projected_columns(metadata, &format_projection, read_cols.json_target_types()); let input_arrow_schema = - compute_input_arrow_schema(metadata, &batch_schema, &read_cols, json2_rewrite_targets); + compute_input_arrow_schema(metadata, &batch_schema, json2_rewrite_targets); // If projection is empty, we don't output any column. let output_schema = if is_empty_projection { @@ -307,6 +302,7 @@ impl FlatProjectionMapper { /// Converts a flat format [RecordBatch] to a normal [RecordBatch]. /// /// The batch must match the `projection` using to build the mapper. + /// JSON2 columns must already match query types before merge and deduplication. pub(crate) fn convert( &self, batch: &datatypes::arrow::record_batch::RecordBatch, @@ -365,13 +361,6 @@ impl FlatProjectionMapper { } } - let field = &self.output_schema.arrow_schema().fields()[output_idx]; - if is_json2_extension_type(field) && array.data_type() != field.data_type() { - array = JsonArray::from(&array) - .project_to_v2(batch.schema_ref().field(*index), field.data_type()) - .context(DataTypesSnafu)?; - } - arrays.push(array); } @@ -383,30 +372,6 @@ impl FlatProjectionMapper { df_record_batch, )) } - - /// Projects columns from the input batch and converts them into vectors. - pub(crate) fn project_vectors( - &self, - batch: &datatypes::arrow::record_batch::RecordBatch, - ) -> common_recordbatch::error::Result> { - let mut columns = Vec::with_capacity(self.output_schema.num_columns()); - for index in &self.batch_indices { - let mut array = batch.column(*index).clone(); - // Casts dictionary values to the target type. - if let datatypes::arrow::datatypes::DataType::Dictionary(_key_type, value_type) = - array.data_type() - { - let casted = datatypes::arrow::compute::cast(&array, value_type) - .context(ArrowComputeSnafu)?; - array = casted; - } - let vector = Helper::try_into_vector(array) - .map_err(BoxedError::new) - .context(ExternalSnafu)?; - columns.push(vector); - } - Ok(columns) - } } fn single_value_string_dictionary<'a>( @@ -481,7 +446,6 @@ pub(crate) fn flat_projected_columns( pub(crate) fn compute_input_arrow_schema( metadata: &RegionMetadata, batch_schema: &[(ColumnId, ConcreteDataType)], - read_cols: &ReadColumns, json2_rewrite_targets: &Json2RewriteTargets, ) -> datatypes::arrow::datatypes::SchemaRef { let mut new_fields = Vec::with_capacity(batch_schema.len() + 3); @@ -489,11 +453,6 @@ pub(crate) fn compute_input_arrow_schema( let data_type = json2_rewrite_targets .get(column_id) .map(|x| json2_physical_data_type(&x.target_layout)) - .or_else(|| { - read_cols - .json_target_type(*column_id) - .map(|x| x.as_arrow_type()) - }) .unwrap_or_else(|| data_type.as_arrow_type()); let column_metadata = metadata.column_by_id(*column_id).unwrap(); @@ -521,8 +480,7 @@ pub(crate) fn compute_input_arrow_schema( /// Helper to project compaction batches into flat format columns /// (fields + time index + __primary_key + __sequence + __op_type). pub(crate) struct CompactionProjectionMapper { - mapper: FlatProjectionMapper, - assembler: DfBatchAssembler, + projection: Vec, } impl CompactionProjectionMapper { @@ -544,77 +502,74 @@ impl CompactionProjectionMapper { let read_col_ids = metadata.column_metadatas.iter().map(|col| col.column_id); let read_cols = ReadColumns::new(read_col_ids); let mapper = FlatProjectionMapper::new_with_read_columns(metadata, projection, read_cols)?; - let assembler = DfBatchAssembler::new(mapper.output_schema()); - - Ok(Self { mapper, assembler }) + Ok(Self { + projection: mapper.batch_indices, + }) } /// Projects columns and appends internal columns for compaction output. /// /// The input batch is expected to be in flat format with internal columns appended. pub(crate) fn project(&self, batch: DfRecordBatch) -> Result { - let columns = self - .mapper - .project_vectors(&batch) - .context(RecordBatchSnafu)?; - self.assembler - .build_df_record_batch_with_internal(&batch, columns) - .context(RecordBatchSnafu) - } -} - -/// Builds [DfRecordBatch] with internal columns appended. -pub(crate) struct DfBatchAssembler { - output_arrow_schema_with_internal: datatypes::arrow::datatypes::SchemaRef, -} - -impl DfBatchAssembler { - /// Precomputes the output schema with internal columns. - pub(crate) fn new(output_schema: SchemaRef) -> Self { - let fields = output_schema - .arrow_schema() - .fields() - .into_iter() - .chain(internal_fields().iter()) - .cloned() - .collect::>(); - let output_arrow_schema_with_internal = - Arc::new(datatypes::arrow::datatypes::Schema::new(fields)); - Self { - output_arrow_schema_with_internal, - } - } - - /// Builds a [DfRecordBatch] from projected vectors plus internal columns. - /// - /// Assumes the input batch already contains internal columns as the last three fields - /// ("__primary_key", "__sequence", "__op_type"). - pub(crate) fn build_df_record_batch_with_internal( - &self, - batch: &datatypes::arrow::record_batch::RecordBatch, - mut columns: Vec, - ) -> common_recordbatch::error::Result { - let num_columns = batch.columns().len(); - // The last 3 columns are the internal columns. - let internal_indices = [num_columns - 3, num_columns - 2, num_columns - 1]; - for index in internal_indices.iter() { - let array = batch.column(*index).clone(); - let vector = Helper::try_into_vector(array) - .map_err(BoxedError::new) - .context(ExternalSnafu)?; - columns.push(vector); - } - RecordBatch::to_df_record_batch(self.output_arrow_schema_with_internal.clone(), columns) + let mut projection = self.projection.clone(); + projection.extend(batch.num_columns() - 3..batch.num_columns()); + // Projection must preserve the Arrow layouts already produced by compaction Rewrite. + batch.project(&projection).context(NewRecordBatchSnafu) } } #[cfg(test)] mod tests { use datatypes::schema::ColumnSchema; + use datatypes::vectors::Helper; use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; use store_api::storage::RegionId; use super::*; + use datatypes::types::json_type::JsonNativeType; + + #[test] + fn test_json2_root_hint_uses_output_type_before_merge() + -> std::result::Result<(), Box> { + let mut builder = RegionMetadataBuilder::new(RegionId::new(1024, 0)); + builder + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "j", + ConcreteDataType::json2(JsonNativeType::object()), + true, + ), + semantic_type: SemanticType::Field, + column_id: 0, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: 1, + }); + let metadata = Arc::new(builder.build()?); + let read_cols = ReadColumns::new([0, 1]) + .with_json_target_types([(0, JsonNativeType::String)].into_iter().collect()); + let mapper = FlatProjectionMapper::new_with_read_columns(&metadata, vec![0, 1], read_cols)?; + let input = mapper.input_arrow_schema(false); + assert_eq!( + &ArrowDataType::Binary, + input.field_with_name("j")?.data_type() + ); + assert_eq!( + mapper + .output_schema() + .arrow_schema() + .field_with_name("j")? + .data_type(), + input.field_with_name("j")?.data_type(), + ); + Ok(()) + } fn metadata_with_legacy_json() -> RegionMetadataRef { let mut builder = RegionMetadataBuilder::new(RegionId::new(1024, 0)); diff --git a/src/mito2/src/read/last_row.rs b/src/mito2/src/read/last_row.rs index b7da3e1f70c..6af6a08717a 100644 --- a/src/mito2/src/read/last_row.rs +++ b/src/mito2/src/read/last_row.rs @@ -16,7 +16,8 @@ use std::sync::Arc; -use datatypes::arrow::array::{Array, BinaryArray, UInt32Array}; +use api::v1::OpType; +use datatypes::arrow::array::{Array, BinaryArray, UInt8Array, UInt32Array}; use datatypes::arrow::compute::{concat_batches, take_record_batch}; use datatypes::arrow::record_batch::RecordBatch; use futures::{Stream, TryStreamExt}; @@ -31,7 +32,9 @@ use crate::error::{ComputeArrowSnafu, Result}; use crate::read::read_columns::JsonTargetTypes; use crate::read::{BoxedRecordBatchStream, timestamp_array_to_i64_slice}; use crate::sst::parquet::DEFAULT_READ_BATCH_SIZE; -use crate::sst::parquet::flat_format::{primary_key_column_index, time_index_column_index}; +use crate::sst::parquet::flat_format::{ + op_type_column_index, primary_key_column_index, time_index_column_index, +}; use crate::sst::parquet::format::{PrimaryKeyArray, primary_key_offsets}; use crate::sst::parquet::read_columns::ParquetReadColumns; use crate::sst::parquet::reader::FlatRowGroupReader; @@ -287,6 +290,31 @@ impl LastKeyState { } impl FlatLastTimestampSelector { + /// Keeps the latest timestamp candidates in one sorted PUT batch, including ties. + pub(crate) fn select_batch(batch: RecordBatch) -> Result { + if batch.num_rows() <= 1 { + return Ok(batch); + } + let Some(op_types) = batch + .column(op_type_column_index(batch.num_columns())) + .as_any() + .downcast_ref::() + else { + return Ok(batch); + }; + if op_types.values().iter().any(|op| *op != OpType::Put as u8) { + return Ok(batch); + } + let mut selector = Self::default(); + let mut pending = BatchBuffer::new(); + selector.on_next(batch.clone(), &mut pending)?; + selector.finish(&mut pending)?; + if pending.num_rows == batch.num_rows() { + return Ok(batch); + } + pending.concat() + } + /// Processes the next batch and appends completed-key results into `output_buffer`. pub(crate) fn on_next( &mut self, @@ -471,6 +499,41 @@ mod tests { use super::*; + #[test] + fn test_select_batch_preserves_timestamp_ties_and_deletes() + -> std::result::Result<(), Box> { + let input = new_flat_batch( + &[b"a", b"a", b"a", b"b", b"b", b"c"], + &[1, 2, 2, 1, 2, 1], + &[1, 2, 3, 4, 5, 6], + ); + let expected = take_record_batch(&input, &UInt32Array::from(vec![1, 2, 4, 5]))?; + assert_eq!( + expected, + FlatLastTimestampSelector::select_batch(input.clone())? + ); + let empty = input.slice(0, 0); + assert_eq!( + empty, + FlatLastTimestampSelector::select_batch(empty.clone())? + ); + let mut columns = input.columns().to_vec(); + columns[op_type_column_index(input.num_columns())] = Arc::new(UInt8Array::from(vec![ + OpType::Delete as u8, + OpType::Put as u8, + OpType::Put as u8, + OpType::Put as u8, + OpType::Put as u8, + OpType::Put as u8, + ])); + let with_delete = RecordBatch::try_new(input.schema(), columns)?; + assert_eq!( + with_delete, + FlatLastTimestampSelector::select_batch(with_delete.clone())? + ); + Ok(()) + } + /// Helper to build a flat format RecordBatch for testing. fn new_flat_batch(primary_keys: &[&[u8]], timestamps: &[i64], fields: &[i64]) -> RecordBatch { let num_rows = timestamps.len(); diff --git a/src/mito2/src/read/scan_util.rs b/src/mito2/src/read/scan_util.rs index a67e2449baa..5f82743f92d 100644 --- a/src/mito2/src/read/scan_util.rs +++ b/src/mito2/src/read/scan_util.rs @@ -26,7 +26,9 @@ use common_telemetry::tracing; use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricBuilder, Time}; use datatypes::arrow::record_batch::RecordBatch; use datatypes::timestamp::timestamp_array_to_primitive; -use futures::Stream; +#[cfg(feature = "enterprise")] +use futures::TryStreamExt; +use futures::{Stream, StreamExt}; use prometheus::IntGauge; use smallvec::SmallVec; use snafu::ResultExt; @@ -40,6 +42,7 @@ use crate::metrics::{ }; use crate::read::dedup::{DedupMetrics, DedupMetricsReport}; use crate::read::flat_merge::{MergeMetrics, MergeMetricsReport}; +use crate::read::last_row::FlatLastTimestampSelector; use crate::read::pruner::PartitionPruner; use crate::read::range::{RangeMeta, RowGroupIndex}; use crate::read::scan_region::StreamContext; @@ -50,6 +53,7 @@ use crate::sst::index::fulltext_index::applier::FulltextIndexApplyMetrics; use crate::sst::index::inverted_index::applier::InvertedIndexApplyMetrics; use crate::sst::parquet::file_range::{FileRange, PreFilterMode}; use crate::sst::parquet::flat_format::{sequence_column_index, time_index_column_index}; +use crate::sst::parquet::json_align; use crate::sst::parquet::reader::{MetadataCacheMetrics, ReaderFilterMetrics, ReaderMetrics}; use crate::sst::parquet::row_group::ParquetFetchMetrics; use crate::sst::parquet::{DEFAULT_READ_BATCH_SIZE, DEFAULT_ROW_GROUP_SIZE}; @@ -1150,6 +1154,15 @@ pub(crate) struct SeriesDistributorMetrics { pub(crate) divider_cost: Duration, } +/// Aligns memtable and extension batches; SST readers align before schema compatibility. +pub(crate) fn align_query_source_batch( + context: &StreamContext, + batch: RecordBatch, +) -> Result { + let query_schema = context.input.mapper.input_arrow_schema(false); + json_align::align_query_batch(batch, &query_schema) +} + /// Scans memtable ranges at `index` using flat format that returns RecordBatch. #[tracing::instrument( skip_all, @@ -1187,6 +1200,39 @@ pub(crate) fn scan_flat_mem_ranges( } } +/// Aligns memtable JSON2 columns to query types before merge/dedup. +/// +/// For append-only queries with `TimeSeriesRowSelector::LastRow { after_merge: false }`, +/// first keeps all rows tied at each series' maximum timestamp within each sorted +/// PUT-only batch, avoiding JSON2 conversion for discarded rows. Final cross-source +/// LastRow selection still runs. Batches containing deletes are not reduced. +/// This optimization uses the query's series row selector, not its merge mode; +/// non-append tables retain all input rows for deduplication and last-non-null backfill. +pub(crate) fn prepare_flat_mem_stream( + context: Arc, + stream: impl Stream> + Send + 'static, +) -> BoxedRecordBatchStream { + if context.input.read_cols.json_target_types().is_empty() { + return Box::pin(stream); + } + + let select_last = context.input.append_mode + && matches!( + context.input.series_row_selector, + Some(TimeSeriesRowSelector::LastRow { after_merge: false }) + ); + Box::pin(stream.map(move |batch| { + let batch = batch?; + // Select raw candidates before decoding JSON2; final cross-source selection remains. + let batch = if select_last { + FlatLastTimestampSelector::select_batch(batch)? + } else { + batch + }; + align_query_source_batch(&context, batch) + })) +} + /// Files with row count greater than this threshold can contribute to the estimation. const SPLIT_ROW_THRESHOLD: u64 = DEFAULT_ROW_GROUP_SIZE as u64; /// Number of series threshold for splitting batches. @@ -1562,12 +1608,10 @@ pub fn build_flat_file_range_scan_stream( let mapper = range.compaction_projection_mapper(); while let Some(record_batch) = reader.next_batch().await? { let record_batch = if let Some(mapper) = mapper { - let batch = mapper.project(record_batch)?; - batch + mapper.project(record_batch)? } else { record_batch }; - let Some(record_batch) = filter_flat_batch_by_sequence( record_batch, stream_ctx.input.sequence_range, @@ -1622,10 +1666,15 @@ pub(crate) async fn scan_flat_extension_range( let range = context.input.extension_range(index.index); let reader = range.flat_reader(context.as_ref(), options); let stream = reader - .read(context, partition_metrics, index) + .read(context.clone(), partition_metrics, index) .await .context(crate::error::ScanExternalRangeSnafu)?; - Ok(stream) + if context.input.read_cols.json_target_types().is_empty() { + return Ok(stream); + } + Ok(Box::pin(stream.and_then(move |batch| { + std::future::ready(align_query_source_batch(&context, batch)) + }))) } pub(crate) async fn maybe_scan_flat_other_ranges( @@ -1736,21 +1785,27 @@ pub(crate) fn split_record_batch(record_batch: RecordBatch, batches: &mut VecDeq #[cfg(test)] mod tests { + use arrow_select::concat::concat_batches; + use common_time::Timestamp; + use datatypes::arrow::array::{BinaryArray, TimestampMillisecondArray, UInt64Array}; + use datatypes::types::json_type::JsonNativeType; + use futures::TryStreamExt; + use smallvec::{SmallVec, smallvec}; use std::sync::Arc; use std::time::Instant; - - use common_time::Timestamp; - use smallvec::{SmallVec, smallvec}; use store_api::storage::RegionId; use super::*; use crate::cache::CacheStrategy; + use crate::memtable::bulk::context::BulkIterContext; + use crate::memtable::bulk::part_reader::BulkPartBatchIter; use crate::memtable::{ BoxedBatchIterator, BoxedRecordBatchIterator, IterBuilder, MemtableRange, MemtableRangeContext, MemtableStats, }; use crate::read::flat_projection::FlatProjectionMapper; use crate::read::range::{MemRangeBuilder, SourceIndex}; + use crate::read::read_columns::ReadColumns; use crate::read::scan_region::ScanInput; use crate::sst::file::{FileHandle, FileMeta}; use crate::sst::file_purger::NoopFilePurger; @@ -1853,6 +1908,148 @@ mod tests { } } + #[tokio::test] + async fn test_memtable_last_row_selects_before_json2_alignment() + -> std::result::Result<(), Box> { + struct BatchIter(RecordBatch, Arc, Option); + impl IterBuilder for BatchIter { + fn build(&self, _metrics: Option) -> Result { + Ok(Box::new(std::iter::empty())) + } + fn is_record_batch(&self) -> bool { + true + } + fn build_record_batch( + &self, + _time_range: Option, + metrics: Option, + ) -> Result { + Ok(Box::new(BulkPartBatchIter::from_single( + self.0.clone(), + self.1.clone(), + self.2, + 1, + metrics, + ))) + } + } + let (metadata, batch) = crate::test_util::sst_util::new_json2_record_batch()?; + let batch = concat_batches(&batch.schema(), &[batch.clone(), batch.slice(0, 1)])?; + let mut columns = batch.columns().to_vec(); + columns[1] = Arc::new(TimestampMillisecondArray::from(vec![1, 2, 3])); + let seq_index = sequence_column_index(batch.num_columns()); + columns[seq_index] = Arc::new(UInt64Array::from(vec![1, 2, 3])); + let batch = RecordBatch::try_new(batch.schema(), columns)?; + let reader_context = Arc::new(BulkIterContext::new( + metadata.clone(), + None, + None, + false, + DEFAULT_READ_BATCH_SIZE, + )?); + let time_range = (Timestamp::new_millisecond(1), Timestamp::new_millisecond(3)); + let env = SchedulerEnv::new().await; + for (append, selector, sequence, count) in [ + ( + true, + Some(TimeSeriesRowSelector::LastRow { after_merge: false }), + None, + 1, + ), + ( + true, + Some(TimeSeriesRowSelector::LastRow { after_merge: true }), + None, + 3, + ), + ( + false, + Some(TimeSeriesRowSelector::LastRow { after_merge: false }), + None, + 3, + ), + (true, None, None, 3), + ( + true, + Some(TimeSeriesRowSelector::LastRow { after_merge: false }), + Some(SequenceRange::GtLtEq { min: 0, max: 2 }), + 1, + ), + ] { + let range_context = Arc::new(MemtableRangeContext::new( + 0, + Box::new(BatchIter(batch.clone(), reader_context.clone(), sequence)), + Default::default(), + )); + let stats = MemtableStats { + time_range: Some(time_range), + num_rows: 3, + num_ranges: 1, + series_count: 1, + ..Default::default() + }; + let range = MemtableRange::new(range_context, stats.clone()); + let read_cols = ReadColumns::new([0, 1]) + .with_json_target_types([(0, JsonNativeType::Variant)].into_iter().collect()); + let mapper = + FlatProjectionMapper::new_with_read_columns(&metadata, vec![0, 1], read_cols)?; + let input = ScanInput::builder(env.access_layer.clone(), mapper) + .with_append_mode(append) + .with_series_row_selector(selector) + .with_sequence_range(sequence) + .with_memtables(vec![MemRangeBuilder::new(range, stats)]) + .build(); + let context = Arc::new(StreamContext { + input, + ranges: vec![], + query_start: Instant::now(), + }); + let metrics = PartitionMetrics::new( + metadata.region_id, + 0, + "test_json2", + context.query_start, + false, + &ExecutionPlanMetricsSet::new(), + ); + let mut raw_stream = Box::pin(scan_flat_mem_ranges( + context.clone(), + metrics, + RowGroupIndex { + index: 0, + row_group_index: -1, + }, + time_range, + )); + let raw = raw_stream.try_next().await?.ok_or("expected a raw batch")?; + let expected_raw = if sequence.is_some() { + batch.slice(0, 2) + } else { + batch.clone() + }; + assert_eq!(expected_raw, raw); + assert!(raw_stream.try_next().await?.is_none()); + let mut stream = prepare_flat_mem_stream(context, futures::stream::iter([Ok(raw)])); + let output = stream.try_next().await?.ok_or("expected a JSON2 batch")?; + assert_eq!(count, output.num_rows()); + let values = output + .column(0) + .as_any() + .downcast_ref::() + .ok_or("expected aligned JSONB")?; + assert_eq!( + datatypes::types::parse_string_to_jsonb(if sequence.is_some() { + r#"{"a":2}"# + } else { + r#"{"a":1}"# + })?, + values.value(count - 1) + ); + assert!(stream.try_next().await?.is_none()); + } + Ok(()) + } + #[tokio::test] async fn test_should_split_flat_batches_for_merge_uses_splittable_file_rows_per_series() { let num_rows = SPLIT_ROW_THRESHOLD * 2; diff --git a/src/mito2/src/read/seq_scan.rs b/src/mito2/src/read/seq_scan.rs index 4b09533e29a..b57622e07e7 100644 --- a/src/mito2/src/read/seq_scan.rs +++ b/src/mito2/src/read/seq_scan.rs @@ -46,8 +46,8 @@ use crate::read::range_cache::{CachedRangeReader, build_range_cache_key}; use crate::read::scan_region::{ScanInput, StreamContext}; use crate::read::scan_util::{ BATCH_SIZE_THRESHOLD, MIN_COMPACTION_SPLIT_ROWS, PartitionMetrics, PartitionMetricsList, - SplitRecordBatchStream, compute_parallel_channel_size, scan_flat_file_ranges, - scan_flat_mem_ranges, should_split_flat_batches_for_merge, + SplitRecordBatchStream, compute_parallel_channel_size, prepare_flat_mem_stream, + scan_flat_file_ranges, scan_flat_mem_ranges, should_split_flat_batches_for_merge, }; use crate::read::stream::{ConvertBatchStream, ScanBatch, ScanBatchStream}; use crate::read::{BoxedRecordBatchStream, ScannerMetrics, scan_util}; @@ -725,7 +725,7 @@ pub(crate) async fn build_flat_sources( *index, range_meta.time_range, ); - ordered_sources[position] = Some(Box::pin(stream) as _); + ordered_sources[position] = Some(prepare_flat_mem_stream(stream_ctx.clone(), stream)); } else if stream_ctx.is_file_range_index(*index) { // Common manifest-level fast-skip shared by SeqScan and UnorderedScan. // Compaction should keep reading its selected input ranges completely. diff --git a/src/mito2/src/read/series_reader.rs b/src/mito2/src/read/series_reader.rs index c030b5b3e7e..021928f2c01 100644 --- a/src/mito2/src/read/series_reader.rs +++ b/src/mito2/src/read/series_reader.rs @@ -36,7 +36,7 @@ use crate::read::scan_region::StreamContext; use crate::read::scan_util::{ PartitionMetrics, SplitRecordBatchStream, compute_average_batch_size, compute_parallel_channel_size, filter_flat_batch_by_sequence, new_filter_metrics, - scan_flat_mem_ranges, should_split_flat_batches_for_merge, + prepare_flat_mem_stream, scan_flat_mem_ranges, should_split_flat_batches_for_merge, }; use crate::read::seq_scan::SeqScan; use crate::read::series_candidate::validate_metric_metadata; @@ -411,12 +411,13 @@ async fn build_uncached_series_partition_range( for index in range_meta.row_group_indices.iter().copied() { if stream_ctx.is_mem_range_index(index) { - let stream = Box::pin(scan_flat_mem_ranges( + let stream = scan_flat_mem_ranges( stream_ctx.clone(), part_metrics.clone(), index, range_meta.time_range, - )); + ); + let stream = prepare_flat_mem_stream(stream_ctx.clone(), stream); sources.push(filter_flat_stream_by_series( stream, codec.clone(), diff --git a/src/mito2/src/read/unordered_scan.rs b/src/mito2/src/read/unordered_scan.rs index 9362898023f..05387be0ce1 100644 --- a/src/mito2/src/read/unordered_scan.rs +++ b/src/mito2/src/read/unordered_scan.rs @@ -37,7 +37,8 @@ use crate::error::{PartitionOutOfRangeSnafu, Result}; use crate::read::pruner::{PartitionPruner, Pruner}; use crate::read::scan_region::{ScanInput, StreamContext}; use crate::read::scan_util::{ - PartitionMetrics, PartitionMetricsList, scan_flat_file_ranges, scan_flat_mem_ranges, + PartitionMetrics, PartitionMetricsList, prepare_flat_mem_stream, scan_flat_file_ranges, + scan_flat_mem_ranges, }; use crate::read::stream::{ConvertBatchStream, ScanBatch, ScanBatchStream}; use crate::read::{ScannerMetrics, scan_util}; @@ -134,6 +135,7 @@ impl UnorderedScan { *index, range_meta.time_range, ); + let stream = prepare_flat_mem_stream(stream_ctx.clone(), stream); for await record_batch in stream { yield record_batch?; } diff --git a/src/mito2/src/sst/parquet/file_range.rs b/src/mito2/src/sst/parquet/file_range.rs index 1c15afd386d..f79569d6d4e 100644 --- a/src/mito2/src/sst/parquet/file_range.rs +++ b/src/mito2/src/sst/parquet/file_range.rs @@ -722,10 +722,20 @@ impl RangeBase { tag_decode_state: &mut TagDecodeState, ) -> Result> { let metadata = self.read_format.metadata(); - // A pruned column rejects the batch without decoding any tag payloads. - if self - .filters - .iter() + // Honor skip_fields/skip_tags before inspecting MaybeFilter::Pruned: a predicate + // marked Pruned is still skipped if it must not run before merge/dedup. + // For example, after adding nullable JSON2 column j to a last_non_null table, + // an old SST may contain (key=a, ts=1, f=10) without j, while a newer row has + // (key=a, ts=1, f=NULL, j={"a":1}). WHERE j IS NOT NULL is Pruned for the old + // SST, but discarding that row here would lose f=10 needed for backfill. + // Keep it for merge/dedup and evaluate the field predicate on the merged row. + // This also applies to added non-JSON2 fields, not only JSON2 columns. + let filters = self.filters.iter().filter(|ctx| { + !(skip_fields && ctx.semantic_type() == SemanticType::Field + || skip_tags && ctx.semantic_type() == SemanticType::Tag) + }); + if filters + .clone() .any(|ctx| matches!(ctx.filter(), MaybeFilter::Pruned)) { return Ok(None); @@ -757,7 +767,7 @@ impl RangeBase { let mut mask = BooleanBuffer::new_set(input.num_rows()); // Run filter one by one and combine them result - for filter_ctx in &self.filters { + for filter_ctx in filters { let filter = match filter_ctx.filter() { MaybeFilter::Filter(f) => f, // Column matches. @@ -766,14 +776,6 @@ impl RangeBase { MaybeFilter::Pruned => return Ok(None), }; - // Skip field filters if skip_fields is true - if skip_fields && filter_ctx.semantic_type() == SemanticType::Field { - continue; - } - if skip_tags && filter_ctx.semantic_type() == SemanticType::Tag { - continue; - } - // Get the column directly by its projected index. // If the column is missing and it's not a tag/time column, this filter is skipped. // Assumes the projection indices align with the input batch schema. @@ -998,6 +1000,14 @@ impl RangeBase { mod tests { use std::sync::Arc; + use super::*; + use crate::read::read_columns::ReadColumns; + use crate::sst::parquet::flat_format::FlatReadFormat; + use crate::sst::{FlatSchemaOptions, to_flat_sst_arrow_schema}; + use crate::test_util::sst_util::{ + new_record_batch_with_custom_sequence, sst_region_metadata, + sst_region_metadata_with_encoding, + }; use datafusion_expr::{col, lit}; use datatypes::arrow::array::{ BinaryDictionaryBuilder, TimestampMillisecondArray, UInt8Array, UInt64Array, @@ -1009,15 +1019,7 @@ mod tests { use mito_codec::row_converter::SparsePrimaryKeyCodec; use parquet::arrow::arrow_reader::RowSelector; use partition::expr::col as partition_col; - - use super::*; - use crate::read::read_columns::ReadColumns; - use crate::sst::parquet::flat_format::FlatReadFormat; - use crate::sst::{FlatSchemaOptions, to_flat_sst_arrow_schema}; - use crate::test_util::sst_util::{ - new_record_batch_with_custom_sequence, sst_region_metadata, - sst_region_metadata_with_encoding, - }; + use store_api::metadata::RegionMetadataBuilder; fn new_test_range_base(filters: Vec) -> RangeBase { let metadata: RegionMetadataRef = Arc::new(sst_region_metadata()); @@ -1333,4 +1335,31 @@ mod tests { ]); assert_eq!(actual, expected); } + #[test] + fn test_pruned_missing_field_respects_skip_fields() + -> std::result::Result<(), Box> { + let metadata = Arc::new(sst_region_metadata()); + let mut builder = RegionMetadataBuilder::from_existing(metadata.as_ref().clone()); + builder.push_column_metadata(store_api::metadata::ColumnMetadata { + column_schema: ColumnSchema::new("added", ConcreteDataType::float64_datatype(), true), + semantic_type: SemanticType::Field, + column_id: 100, + }); + let expected = builder.build()?; + let filter = + SimpleFilterContext::new_opt(&metadata, Some(&expected), &col("added").is_not_null()) + .ok_or("expected a missing-column filter")?; + assert!(matches!(filter.filter(), MaybeFilter::Pruned)); + let base = new_test_range_base(vec![filter]); + let batch = new_record_batch_with_custom_sequence(&["b", "x"], 0, 4, 1); + assert!( + base.compute_filter_mask_flat(&batch, false, false, &mut TagDecodeState::new())? + .is_none() + ); + let mask = base + .compute_filter_mask_flat(&batch, true, false, &mut TagDecodeState::new())? + .ok_or("skipped field filter must preserve backfill rows")?; + assert_eq!(batch.num_rows(), mask.count_set_bits()); + Ok(()) + } } diff --git a/src/mito2/src/sst/parquet/json_align.rs b/src/mito2/src/sst/parquet/json_align.rs index 93ecffa072f..c03c7f1d62b 100644 --- a/src/mito2/src/sst/parquet/json_align.rs +++ b/src/mito2/src/sst/parquet/json_align.rs @@ -12,8 +12,12 @@ // See the License for the specific language governing permissions and // limitations under the License. +//! Shared JSON2 conversion for Parquet streams and in-memory query batches. +//! Query types must agree before merge/dedup; compaction Rewrite retains physical layouts. + use std::collections::HashMap; use std::pin::Pin; +use std::sync::Arc; use std::task::{Context, Poll}; use datafusion_common::cast_column; @@ -169,6 +173,45 @@ where } } +/// Aligns existing JSON2 columns to the Arrow types of matching fields in +/// `query_schema`, matching columns by name. +/// Preserves column order and skips columns absent from the batch. +pub(crate) fn align_query_batch( + batch: RecordBatch, + query_schema: &SchemaRef, +) -> Result { + let mut aligned = None; + for target in query_schema + .fields() + .iter() + .filter(|field| is_json2_extension_type(field)) + { + let Ok(index) = batch.schema_ref().index_of(target.name()) else { + continue; + }; + let source = batch.schema_ref().field(index); + if source.data_type() != target.data_type() { + let array = align_array(batch.column(index), source, target)?; + let (fields, columns) = aligned.get_or_insert_with(|| { + ( + batch.schema_ref().fields().to_vec(), + batch.columns().to_vec(), + ) + }); + fields[index] = Arc::new(source.clone().with_data_type(target.data_type().clone())); + columns[index] = array; + } + } + let Some((fields, columns)) = aligned else { + return Ok(batch); + }; + let schema = Arc::new(Schema::new_with_metadata( + fields, + batch.schema().metadata().clone(), + )); + RecordBatch::try_new(schema, columns).context(NewRecordBatchSnafu) +} + fn align_projected_batch( rb: RecordBatch, output_schema: &SchemaRef, @@ -327,6 +370,8 @@ mod tests { use std::collections::HashMap; use std::sync::Arc; + use super::*; + use crate::test_util; use datatypes::arrow::array::{ Array, ArrayRef, BinaryArray, Int64Array, StringArray, StringViewArray, StructArray, }; @@ -337,8 +382,6 @@ mod tests { use datatypes::types::parse_string_to_jsonb; use futures::{StreamExt, stream}; - use super::*; - #[test] fn test_aligner_resolves_json2_rewrite_settings() -> std::result::Result<(), Box> { @@ -955,4 +998,32 @@ mod tests { fn string_array(values: impl IntoIterator) -> ArrayRef { Arc::new(StringArray::from_iter_values(values)) } + #[test] + fn test_query_alignment_preserves_source_columns_until_compatibility() + -> std::result::Result<(), Box> { + let (_, batch) = test_util::sst_util::new_json2_record_batch()?; + let values = int_array([10, 20]); + let mut fields = vec![Arc::new(Field::new("value", DataType::Int64, true))]; + fields.extend(batch.schema_ref().fields().iter().cloned()); + let mut columns = vec![values.clone()]; + columns.extend_from_slice(batch.columns()); + let input = RecordBatch::try_new(Arc::new(Schema::new(fields)), columns)?; + let target = schema([ + Field::new("new_column", DataType::Utf8, true), + Field::new("new_json", DataType::Binary, true) + .with_extension_type(Json2ExtensionType::default()), + Field::new("j", DataType::Binary, true) + .with_extension_type(Json2ExtensionType::default()), + Field::new("value", DataType::Int64, true), + ]); + let aligned = align_query_batch(input, &target)?; + assert_eq!(batch.num_columns() + 1, aligned.num_columns()); + assert_eq!("value", aligned.schema_ref().field(0).name()); + assert!(Arc::ptr_eq(&values, aligned.column(0))); + assert_eq!("j", aligned.schema_ref().field(1).name()); + assert_eq!(&DataType::Binary, aligned.column(1).data_type()); + let repeated = align_query_batch(aligned.clone(), &target)?; + assert!(Arc::ptr_eq(aligned.column(1), repeated.column(1))); + Ok(()) + } } diff --git a/src/mito2/src/sst/parquet/reader.rs b/src/mito2/src/sst/parquet/reader.rs index 56488669f5a..1c98444946c 100644 --- a/src/mito2/src/sst/parquet/reader.rs +++ b/src/mito2/src/sst/parquet/reader.rs @@ -655,6 +655,14 @@ impl ParquetReaderBuilder { let json2_rewrite_targets = self .json2_rewrite_targets .iter() + // Global compaction targets may include JSON2 columns added after this SST was written. + // Rewrite only projected columns present in this SST; schema compatibility fills + // missing columns with NULL later, without looking them up in this SST's metadata. + .filter(|(column_id, _)| { + read_format + .parquet_projected_index_by_id(**column_id) + .is_some() + }) .map(|(column_id, layout)| { let column = region_meta .column_by_id(*column_id) diff --git a/src/mito2/src/test_util/sst_util.rs b/src/mito2/src/test_util/sst_util.rs index fd6321b0895..84933af2946 100644 --- a/src/mito2/src/test_util/sst_util.rs +++ b/src/mito2/src/test_util/sst_util.rs @@ -24,10 +24,16 @@ use datatypes::arrow::array::{ ArrayRef, BinaryDictionaryBuilder, RecordBatch, StringDictionaryBuilder, TimestampMillisecondArray, UInt8Array, UInt64Array, }; +use datatypes::arrow::array::{Int64Array, StructArray}; use datatypes::arrow::datatypes::UInt32Type; +use datatypes::arrow::datatypes::{DataType, Field}; +use datatypes::extension::json::Json2ExtensionType; +use datatypes::json::JsonSettings; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, SkippingIndexOptions}; +use datatypes::types::json_type::JsonNativeType; use datatypes::value::ValueRef; +use datatypes::vectors::json::array::JsonArray; use mito_codec::row_converter::{DensePrimaryKeyCodec, PrimaryKeyCodecExt, SortField}; use object_store::layers::mock::{self, MockLayer, MockLayerBuilder, oio}; use store_api::metadata::{ @@ -44,6 +50,57 @@ use crate::sst::file::{FileHandle, FileMeta}; use crate::sst::{FlatSchemaOptions, to_flat_sst_arrow_schema}; use crate::test_util::{new_batch_builder, new_noop_file_purger}; +/// Builds a current-layout JSON2 flat batch with two timestamps in one series. +pub(crate) fn new_json2_record_batch() +-> Result<(RegionMetadataRef, RecordBatch), Box> { + let legacy: ArrayRef = Arc::new(StructArray::from(vec![( + Arc::new(Field::new("a", DataType::Int64, true)), + Arc::new(Int64Array::from_iter_values([1, 2])) as ArrayRef, + )])); + let field = Field::new("j", legacy.data_type().clone(), true) + .with_extension_type(Json2ExtensionType::default()); + let settings = JsonSettings::default(); + let physical = JsonArray::from(&legacy).rewrite_to_v2(&field, &settings, &settings)?; + let mut column_schema = + ColumnSchema::new("j", ConcreteDataType::json2(JsonNativeType::object()), true); + column_schema.with_extension_type(&Json2ExtensionType::default()); + let mut metadata = RegionMetadataBuilder::new(RegionId::new(1024, 0)); + metadata + .push_column_metadata(ColumnMetadata { + column_schema, + semantic_type: SemanticType::Field, + column_id: 0, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: 1, + }); + let metadata = Arc::new(metadata.build()?); + let mut options = FlatSchemaOptions::default(); + options + .concretized_json_types + .insert("j".to_string(), physical.data_type().clone()); + let mut keys = BinaryDictionaryBuilder::::new(); + keys.append(b"series")?; + keys.append(b"series")?; + let input = RecordBatch::try_new( + to_flat_sst_arrow_schema(&metadata, &options), + vec![ + physical, + Arc::new(TimestampMillisecondArray::from(vec![1, 2])), + Arc::new(keys.finish()), + Arc::new(UInt64Array::from_value(1, 2)), + Arc::new(UInt8Array::from_value(OpType::Put as u8, 2)), + ], + )?; + Ok((metadata, input)) +} + /// Test region id. const REGION_ID: RegionId = RegionId::new(0, 0); diff --git a/src/operator/src/statement/ddl.rs b/src/operator/src/statement/ddl.rs index c19be914a73..40a83613a56 100644 --- a/src/operator/src/statement/ddl.rs +++ b/src/operator/src/statement/ddl.rs @@ -100,7 +100,6 @@ use sqlparser::ast::{ use store_api::metric_engine_consts::{ LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME, PHYSICAL_TABLE_METADATA_KEY, }; -use store_api::mito_engine_options::APPEND_MODE_KEY; use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan}; use table::TableRef; use table::dist_table::DistTable; @@ -2632,7 +2631,7 @@ pub fn verify_alter( .build() .context(error::BuildTableMetaSnafu { table_name })?; - validate_json2_columns_append_mode(&new_meta.schema, &new_meta.options)?; + validate_json2_primary_keys(&new_meta.schema, &new_meta.primary_key_indices)?; Ok(true) } @@ -2682,7 +2681,7 @@ pub fn create_table_info( let mut table_options = TableOptions::try_from_iter(&create_table.table_options) .context(UnrecognizedTableOptionSnafu)?; - validate_json2_columns_append_mode(&schema, &table_options)?; + validate_json2_primary_keys(&schema, &primary_key_indices)?; validate_and_normalize_annotations(&mut table_options, &schema, &partition_key_indices)?; @@ -2721,26 +2720,16 @@ pub fn create_table_info( Ok(table_info) } -fn validate_json2_columns_append_mode(schema: &Schema, table_options: &TableOptions) -> Result<()> { - let append_mode = table_options - .extra_options - .get(APPEND_MODE_KEY) - .is_some_and(|value| value == "true"); - - for column in schema.column_schemas() { - if column.data_type.is_json2() { - ensure!( - append_mode, - InvalidSqlSnafu { - err_msg: format!( - "JSON2 column `{}` requires {}='true'", - column.name, APPEND_MODE_KEY - ), - } - ); - } +fn validate_json2_primary_keys(schema: &Schema, primary_key_indices: &[usize]) -> Result<()> { + for &index in primary_key_indices { + let column = &schema.column_schemas()[index]; + ensure!( + !column.data_type.is_json2(), + InvalidSqlSnafu { + err_msg: format!("JSON2 column `{}` cannot be a primary key", column.name), + } + ); } - Ok(()) } @@ -4147,26 +4136,36 @@ WITH ('repartition.column.hint' = ' host ')", } #[test] - fn test_json2_requires_append_mode() { - let cases = [ - "CREATE TABLE monitor (payload JSON2, ts TIMESTAMP TIME INDEX);", - "CREATE TABLE monitor (payload JSON2, ts TIMESTAMP TIME INDEX) WITH (append_mode='false');", - ]; - - for sql in cases { - let expr = create_expr_from_sql(sql); - let err = create_table_info(&expr, vec![]).unwrap_err(); - assert!( - err.to_string() - .contains("JSON2 column `payload` requires append_mode='true'"), - "{err}" - ); + fn test_json2_supports_non_append_mode() -> std::result::Result<(), Box> + { + for options in [ + "", + "WITH (append_mode='false')", + "WITH (append_mode='true')", + "WITH (merge_mode='last_non_null')", + ] { + let expr = create_expr_from_sql(&format!( + "CREATE TABLE monitor (payload JSON2, ts TIMESTAMP TIME INDEX) {options};" + )); + create_table_info(&expr, vec![])?; } + Ok(()) + } + #[test] + fn test_json2_cannot_be_primary_key() -> std::result::Result<(), Box> { let expr = create_expr_from_sql( - "CREATE TABLE monitor (payload JSON2, ts TIMESTAMP TIME INDEX) WITH (append_mode='true');", + "CREATE TABLE monitor (payload JSON2, ts TIMESTAMP TIME INDEX, PRIMARY KEY(payload));", ); - create_table_info(&expr, vec![]).unwrap(); + let err = create_table_info(&expr, vec![]) + .err() + .ok_or("expected JSON2 primary key rejection")?; + assert!( + err.to_string() + .contains("JSON2 column `payload` cannot be a primary key"), + "{err}" + ); + Ok(()) } #[test] diff --git a/src/query/src/datafusion.rs b/src/query/src/datafusion.rs index 4b23f81eb16..2d9ae7a7baf 100644 --- a/src/query/src/datafusion.rs +++ b/src/query/src/datafusion.rs @@ -233,14 +233,47 @@ impl DatafusionQueryEngine { } } WriteOp::Delete => { + // Read-only system tables may have no time index; reject DELETE before key lookup. + ensure!( + !is_readonly_table(&table_name.schema, &table_name.table), + TableReadOnlySnafu { + table: table_name.table.to_string(), + } + ); + let table_info = table.table_info(); + let table_schema = table.schema(); + let ts_column = table_schema.timestamp_column().with_context(|| { + MissingTimestampColumnSnafu { + table_name: table_info.name.clone(), + } + })?; + let key_indices = stream + .schema() + .column_schemas() + .iter() + .enumerate() + .filter(|(_, column)| { + ts_column.name == column.name + || table_info + .meta + .row_key_column_names() + .any(|name| name == &column.name) + }) + .map(|(index, _)| index) + .collect::>(); while let Some(batch) = stream.next().await { - let batch = batch.context(CreateRecordBatchSnafu)?; + // The input plan already evaluated WHERE; storage deletes need only keys/time. + // Project before vector conversion because JSON2's logical Struct is a placeholder. + let batch = batch + .context(CreateRecordBatchSnafu)? + .try_project(&key_indices) + .context(CreateRecordBatchSnafu)?; let column_vectors = batch .column_vectors(&table_name.to_string(), table.schema()) .map_err(BoxedError::new) .context(QueryExecutionSnafu)?; affected_rows += self - .delete(&table_name, &table, column_vectors, query_ctx.clone()) + .delete(&table_name, column_vectors, query_ctx.clone()) .await?; } } @@ -256,36 +289,12 @@ impl DatafusionQueryEngine { async fn delete( &self, table_name: &ResolvedTableReference, - table: &TableRef, column_vectors: HashMap, query_ctx: QueryContextRef, ) -> Result { let catalog_name = table_name.catalog.to_string(); let schema_name = table_name.schema.to_string(); let table_name = table_name.table.to_string(); - let table_schema = table.schema(); - - ensure!( - !is_readonly_table(&schema_name, &table_name), - TableReadOnlySnafu { table: table_name } - ); - - let ts_column = table_schema - .timestamp_column() - .map(|x| &x.name) - .with_context(|| MissingTimestampColumnSnafu { - table_name: table_name.clone(), - })?; - - let table_info = table.table_info(); - let rowkey_columns = table_info - .meta - .row_key_column_names() - .collect::>(); - let column_vectors = column_vectors - .into_iter() - .filter(|x| &x.0 == ts_column || rowkey_columns.contains(&&x.0)) - .collect::>(); let request = DeleteRequest { catalog_name, diff --git a/tests/cases/standalone/common/types/json/json2_limit.result b/tests/cases/standalone/common/types/json/json2_limit.result index 6c362a5cc62..a3e237bf753 100644 --- a/tests/cases/standalone/common/types/json/json2_limit.result +++ b/tests/cases/standalone/common/types/json/json2_limit.result @@ -65,36 +65,6 @@ drop table json2_whole_and_path_read; Affected Rows: 0 -create table json2_without_append_mode ( - ts timestamp time index, - j json2 -); - -Error: 1004(InvalidArguments), Invalid SQL, error: JSON2 column `j` requires append_mode='true' - -create table json2_append_mode_false ( - ts timestamp time index, - j json2 -) with ( - 'append_mode' = 'false' -); - -Error: 1004(InvalidArguments), Invalid SQL, error: JSON2 column `j` requires append_mode='true' - -create table json2_alter_non_append ( - ts timestamp time index -); - -Affected Rows: 0 - -alter table json2_alter_non_append add column j json2; - -Error: 1004(InvalidArguments), Invalid SQL, error: JSON2 column `j` requires append_mode='true' - -drop table json2_alter_non_append; - -Affected Rows: 0 - create table json2_set_append_mode_false ( ts timestamp time index, j json2 @@ -104,14 +74,35 @@ create table json2_set_append_mode_false ( Affected Rows: 0 +-- SQLNESS REPLACE \d+\(\d+,\s+\d+\) REDACTED alter table json2_set_append_mode_false set 'append_mode' = 'false'; -Error: 1004(InvalidArguments), Invalid SQL, error: JSON2 column `j` requires append_mode='true' +Error: 1004(InvalidArguments), Invalid region request, region_id: REDACTED, err: Only allow changing append_mode from false to true drop table json2_set_append_mode_false; Affected Rows: 0 +create table json2_primary_key ( + ts timestamp time index, + j json2, + primary key (j) +); + +Error: 1004(InvalidArguments), Invalid SQL, error: JSON2 column `j` cannot be a primary key + +create table json2_add_primary_key (ts timestamp time index); + +Affected Rows: 0 + +alter table json2_add_primary_key add column j json2 primary key; + +Error: 1004(InvalidArguments), Invalid SQL, error: JSON2 column `j` cannot be a primary key + +drop table json2_add_primary_key; + +Affected Rows: 0 + -- JSON2 requires flat SST output even when append mode is enabled. create table json2_primary_key_format ( ts timestamp time index, diff --git a/tests/cases/standalone/common/types/json/json2_limit.sql b/tests/cases/standalone/common/types/json/json2_limit.sql index 7b428abb5e0..6a7de6d113b 100644 --- a/tests/cases/standalone/common/types/json/json2_limit.sql +++ b/tests/cases/standalone/common/types/json/json2_limit.sql @@ -38,26 +38,6 @@ order by json_get(j, 'a.b'); drop table json2_whole_and_path_read; -create table json2_without_append_mode ( - ts timestamp time index, - j json2 -); - -create table json2_append_mode_false ( - ts timestamp time index, - j json2 -) with ( - 'append_mode' = 'false' -); - -create table json2_alter_non_append ( - ts timestamp time index -); - -alter table json2_alter_non_append add column j json2; - -drop table json2_alter_non_append; - create table json2_set_append_mode_false ( ts timestamp time index, j json2 @@ -65,10 +45,23 @@ create table json2_set_append_mode_false ( 'append_mode' = 'true' ); +-- SQLNESS REPLACE \d+\(\d+,\s+\d+\) REDACTED alter table json2_set_append_mode_false set 'append_mode' = 'false'; drop table json2_set_append_mode_false; +create table json2_primary_key ( + ts timestamp time index, + j json2, + primary key (j) +); + +create table json2_add_primary_key (ts timestamp time index); + +alter table json2_add_primary_key add column j json2 primary key; + +drop table json2_add_primary_key; + -- JSON2 requires flat SST output even when append mode is enabled. create table json2_primary_key_format ( ts timestamp time index, diff --git a/tests/cases/standalone/common/types/json/json2_non_append.result b/tests/cases/standalone/common/types/json/json2_non_append.result new file mode 100644 index 00000000000..eed66518477 --- /dev/null +++ b/tests/cases/standalone/common/types/json/json2_non_append.result @@ -0,0 +1,675 @@ +create table json2_without_append_mode ( + ts timestamp time index, + j json2 +); + +Affected Rows: 0 + +drop table json2_without_append_mode; + +Affected Rows: 0 + +create table json2_append_mode_false ( + ts timestamp time index, + j json2 +) with ( + 'append_mode' = 'false' +); + +Affected Rows: 0 + +drop table json2_append_mode_false; + +Affected Rows: 0 + +create table json2_alter_non_append ( + ts timestamp time index +); + +Affected Rows: 0 + +insert into json2_alter_non_append values (1); + +Affected Rows: 1 + +alter table json2_alter_non_append add column j json2; + +Affected Rows: 0 + +insert into json2_alter_non_append values (1, '{"a":1}'); + +Affected Rows: 1 + +select ts, j.a from json2_alter_non_append; + ++-------------------------+------------------------------------------------+ +| ts | json_get(json2_alter_non_append.j,Utf8("$.a")) | ++-------------------------+------------------------------------------------+ +| 1970-01-01T00:00:00.001 | 1 | ++-------------------------+------------------------------------------------+ + +drop table json2_alter_non_append; + +Affected Rows: 0 + +-- JSON2 is merged as a whole field, including empty objects and explicit JSON nulls. +CREATE TABLE json2_last_row ( + host STRING, + ts TIMESTAMP TIME INDEX, + j JSON2, + k JSON2(max_auto_expanded_paths = 0), + f DOUBLE, + PRIMARY KEY (host) +) WITH ( + 'append_mode' = 'false', + 'merge_mode' = 'last_row' +); + +Affected Rows: 0 + +INSERT INTO json2_last_row VALUES + ('one', 1, '{"a":1}', '{"a":1}', 10), + ('one', 2, '{"a":2}', '{"a":2}', 10), + ('one', 3, '{"a":3}', '{"a":3}', 10), + ('one', 4, '{"a":4}', '{"a":4}', 10), + ('one', 5, '{"a":{"x":5}}', '{"a":{"x":5}}', 10), + ('one', 6, '{"a":{"x":6}}', '{"a":{"x":6}}', 10), + ('one', 7, '{"a":{"x":8}}', '{"a":{"x":8}}', 10), + ('two', 1, '{"a":99}', '{"a":99}', 99); + +Affected Rows: 8 + +ADMIN FLUSH_TABLE('json2_last_row'); + ++-------------------------------------+ +| ADMIN FLUSH_TABLE('json2_last_row') | ++-------------------------------------+ +| 0 | ++-------------------------------------+ + +-- New objects replace old objects; only SQL NULL can fall back in last_non_null. +INSERT INTO json2_last_row VALUES + ('one', 1, '{"b":2}', '{"b":2}', NULL), + ('one', 2, '{}', '{}', NULL), + ('one', 3, '{"a":null}', '{"a":null}', NULL), + ('one', 4, NULL, NULL, NULL), + ('one', 5, '{"a":42}', '{"a":42}', NULL), + ('one', 6, '{"a":{"x":7}}', '{"a":{"x":7}}', NULL), + ('one', 7, NULL, NULL, NULL); + +Affected Rows: 7 + +-- Path-only reads must keep the new nested value and distinguish SQL NULL fallback. +SELECT host, ts, j, j.a AS a, j.a.x AS ax, k, k.a AS ka, f +FROM json2_last_row ORDER BY host, ts; + ++------+-------------------------+---------------+---------+----+---------------+---------+------+ +| host | ts | j | a | ax | k | ka | f | ++------+-------------------------+---------------+---------+----+---------------+---------+------+ +| one | 1970-01-01T00:00:00.001 | {"b":2} | | | {"b":2} | | | +| one | 1970-01-01T00:00:00.002 | {} | | | {} | | | +| one | 1970-01-01T00:00:00.003 | {"a":null} | | | {"a":null} | | | +| one | 1970-01-01T00:00:00.004 | | | | | | | +| one | 1970-01-01T00:00:00.005 | {"a":42} | 42 | | {"a":42} | 42 | | +| one | 1970-01-01T00:00:00.006 | {"a":{"x":7}} | {"x":7} | 7 | {"a":{"x":7}} | {"x":7} | | +| one | 1970-01-01T00:00:00.007 | | | | | | | +| two | 1970-01-01T00:00:00.001 | {"a":99} | 99 | | {"a":99} | 99 | 99.0 | ++------+-------------------------+---------------+---------+----+---------------+---------+------+ + +SELECT host, ts, j.a.x AS ax, k.a.x AS kax FROM json2_last_row ORDER BY host, ts; + ++------+-------------------------+----+-----+ +| host | ts | ax | kax | ++------+-------------------------+----+-----+ +| one | 1970-01-01T00:00:00.001 | | | +| one | 1970-01-01T00:00:00.002 | | | +| one | 1970-01-01T00:00:00.003 | | | +| one | 1970-01-01T00:00:00.004 | | | +| one | 1970-01-01T00:00:00.005 | | | +| one | 1970-01-01T00:00:00.006 | 7 | 7 | +| one | 1970-01-01T00:00:00.007 | | | +| two | 1970-01-01T00:00:00.001 | | | ++------+-------------------------+----+-----+ + +SELECT host, ts FROM json2_last_row WHERE j.a = '1' ORDER BY host, ts; + ++------+----+ +| host | ts | ++------+----+ ++------+----+ + +ADMIN FLUSH_TABLE('json2_last_row'); + ++-------------------------------------+ +| ADMIN FLUSH_TABLE('json2_last_row') | ++-------------------------------------+ +| 0 | ++-------------------------------------+ + +SELECT host, ts, j, j.a AS a, j.a.x AS ax, k, k.a AS ka, f +FROM json2_last_row ORDER BY host, ts; + ++------+-------------------------+---------------+---------+----+---------------+---------+------+ +| host | ts | j | a | ax | k | ka | f | ++------+-------------------------+---------------+---------+----+---------------+---------+------+ +| one | 1970-01-01T00:00:00.001 | {"b":2} | | | {"b":2} | | | +| one | 1970-01-01T00:00:00.002 | {} | | | {} | | | +| one | 1970-01-01T00:00:00.003 | {"a":null} | | | {"a":null} | | | +| one | 1970-01-01T00:00:00.004 | | | | | | | +| one | 1970-01-01T00:00:00.005 | {"a":42} | 42 | | {"a":42} | 42 | | +| one | 1970-01-01T00:00:00.006 | {"a":{"x":7}} | {"x":7} | 7 | {"a":{"x":7}} | {"x":7} | | +| one | 1970-01-01T00:00:00.007 | | | | | | | +| two | 1970-01-01T00:00:00.001 | {"a":99} | 99 | | {"a":99} | 99 | 99.0 | ++------+-------------------------+---------------+---------+----+---------------+---------+------+ + +SELECT host, ts, j.a.x AS ax, k.a.x AS kax FROM json2_last_row ORDER BY host, ts; + ++------+-------------------------+----+-----+ +| host | ts | ax | kax | ++------+-------------------------+----+-----+ +| one | 1970-01-01T00:00:00.001 | | | +| one | 1970-01-01T00:00:00.002 | | | +| one | 1970-01-01T00:00:00.003 | | | +| one | 1970-01-01T00:00:00.004 | | | +| one | 1970-01-01T00:00:00.005 | | | +| one | 1970-01-01T00:00:00.006 | 7 | 7 | +| one | 1970-01-01T00:00:00.007 | | | +| two | 1970-01-01T00:00:00.001 | | | ++------+-------------------------+----+-----+ + +ADMIN COMPACT_TABLE('json2_last_row'); + ++---------------------------------------+ +| ADMIN COMPACT_TABLE('json2_last_row') | ++---------------------------------------+ +| 0 | ++---------------------------------------+ + +SELECT host, ts, j, j.a AS a, j.a.x AS ax, k, k.a AS ka, f +FROM json2_last_row ORDER BY host, ts; + ++------+-------------------------+---------------+---------+----+---------------+---------+------+ +| host | ts | j | a | ax | k | ka | f | ++------+-------------------------+---------------+---------+----+---------------+---------+------+ +| one | 1970-01-01T00:00:00.001 | {"b":2} | | | {"b":2} | | | +| one | 1970-01-01T00:00:00.002 | {} | | | {} | | | +| one | 1970-01-01T00:00:00.003 | {"a":null} | | | {"a":null} | | | +| one | 1970-01-01T00:00:00.004 | | | | | | | +| one | 1970-01-01T00:00:00.005 | {"a":42} | 42 | | {"a":42} | 42 | | +| one | 1970-01-01T00:00:00.006 | {"a":{"x":7}} | {"x":7} | 7 | {"a":{"x":7}} | {"x":7} | | +| one | 1970-01-01T00:00:00.007 | | | | | | | +| two | 1970-01-01T00:00:00.001 | {"a":99} | 99 | | {"a":99} | 99 | 99.0 | ++------+-------------------------+---------------+---------+----+---------------+---------+------+ + +SELECT host, ts, j.a.x AS ax, k.a.x AS kax FROM json2_last_row ORDER BY host, ts; + ++------+-------------------------+----+-----+ +| host | ts | ax | kax | ++------+-------------------------+----+-----+ +| one | 1970-01-01T00:00:00.001 | | | +| one | 1970-01-01T00:00:00.002 | | | +| one | 1970-01-01T00:00:00.003 | | | +| one | 1970-01-01T00:00:00.004 | | | +| one | 1970-01-01T00:00:00.005 | | | +| one | 1970-01-01T00:00:00.006 | 7 | 7 | +| one | 1970-01-01T00:00:00.007 | | | +| two | 1970-01-01T00:00:00.001 | | | ++------+-------------------------+----+-----+ + +-- SQLNESS ARG restart=true +SELECT host, ts, j, j.a AS a, j.a.x AS ax, k, k.a AS ka, f +FROM json2_last_row ORDER BY host, ts; + ++------+-------------------------+---------------+---------+----+---------------+---------+------+ +| host | ts | j | a | ax | k | ka | f | ++------+-------------------------+---------------+---------+----+---------------+---------+------+ +| one | 1970-01-01T00:00:00.001 | {"b":2} | | | {"b":2} | | | +| one | 1970-01-01T00:00:00.002 | {} | | | {} | | | +| one | 1970-01-01T00:00:00.003 | {"a":null} | | | {"a":null} | | | +| one | 1970-01-01T00:00:00.004 | | | | | | | +| one | 1970-01-01T00:00:00.005 | {"a":42} | 42 | | {"a":42} | 42 | | +| one | 1970-01-01T00:00:00.006 | {"a":{"x":7}} | {"x":7} | 7 | {"a":{"x":7}} | {"x":7} | | +| one | 1970-01-01T00:00:00.007 | | | | | | | +| two | 1970-01-01T00:00:00.001 | {"a":99} | 99 | | {"a":99} | 99 | 99.0 | ++------+-------------------------+---------------+---------+----+---------------+---------+------+ + +SELECT host, ts, j.a.x AS ax, k.a.x AS kax FROM json2_last_row ORDER BY host, ts; + ++------+-------------------------+----+-----+ +| host | ts | ax | kax | ++------+-------------------------+----+-----+ +| one | 1970-01-01T00:00:00.001 | | | +| one | 1970-01-01T00:00:00.002 | | | +| one | 1970-01-01T00:00:00.003 | | | +| one | 1970-01-01T00:00:00.004 | | | +| one | 1970-01-01T00:00:00.005 | | | +| one | 1970-01-01T00:00:00.006 | 7 | 7 | +| one | 1970-01-01T00:00:00.007 | | | +| two | 1970-01-01T00:00:00.001 | | | ++------+-------------------------+----+-----+ + +-- A delete must prevent last_non_null from filling fields from older versions. +DELETE FROM json2_last_row WHERE host = 'one' AND ts = 1; + +Affected Rows: 1 + +INSERT INTO json2_last_row VALUES ('one', 1, '{"c":3}', '{"c":3}', NULL); + +Affected Rows: 1 + +SELECT host, ts, j, j.a AS a, j.a.x AS ax, k, k.a AS ka, f +FROM json2_last_row ORDER BY host, ts; + ++------+-------------------------+---------------+---------+----+---------------+---------+------+ +| host | ts | j | a | ax | k | ka | f | ++------+-------------------------+---------------+---------+----+---------------+---------+------+ +| one | 1970-01-01T00:00:00.001 | {"c":3} | | | {"c":3} | | | +| one | 1970-01-01T00:00:00.002 | {} | | | {} | | | +| one | 1970-01-01T00:00:00.003 | {"a":null} | | | {"a":null} | | | +| one | 1970-01-01T00:00:00.004 | | | | | | | +| one | 1970-01-01T00:00:00.005 | {"a":42} | 42 | | {"a":42} | 42 | | +| one | 1970-01-01T00:00:00.006 | {"a":{"x":7}} | {"x":7} | 7 | {"a":{"x":7}} | {"x":7} | | +| one | 1970-01-01T00:00:00.007 | | | | | | | +| two | 1970-01-01T00:00:00.001 | {"a":99} | 99 | | {"a":99} | 99 | 99.0 | ++------+-------------------------+---------------+---------+----+---------------+---------+------+ + +ADMIN FLUSH_TABLE('json2_last_row'); + ++-------------------------------------+ +| ADMIN FLUSH_TABLE('json2_last_row') | ++-------------------------------------+ +| 0 | ++-------------------------------------+ + +ADMIN COMPACT_TABLE('json2_last_row'); + ++---------------------------------------+ +| ADMIN COMPACT_TABLE('json2_last_row') | ++---------------------------------------+ +| 0 | ++---------------------------------------+ + +SELECT host, ts, j, j.a AS a, j.a.x AS ax, k, k.a AS ka, f +FROM json2_last_row ORDER BY host, ts; + ++------+-------------------------+---------------+---------+----+---------------+---------+------+ +| host | ts | j | a | ax | k | ka | f | ++------+-------------------------+---------------+---------+----+---------------+---------+------+ +| one | 1970-01-01T00:00:00.001 | {"c":3} | | | {"c":3} | | | +| one | 1970-01-01T00:00:00.002 | {} | | | {} | | | +| one | 1970-01-01T00:00:00.003 | {"a":null} | | | {"a":null} | | | +| one | 1970-01-01T00:00:00.004 | | | | | | | +| one | 1970-01-01T00:00:00.005 | {"a":42} | 42 | | {"a":42} | 42 | | +| one | 1970-01-01T00:00:00.006 | {"a":{"x":7}} | {"x":7} | 7 | {"a":{"x":7}} | {"x":7} | | +| one | 1970-01-01T00:00:00.007 | | | | | | | +| two | 1970-01-01T00:00:00.001 | {"a":99} | 99 | | {"a":99} | 99 | 99.0 | ++------+-------------------------+---------------+---------+----+---------------+---------+------+ + +DELETE FROM json2_last_row WHERE j.a = '99'; + +Affected Rows: 1 + +SELECT COUNT(*) FROM json2_last_row WHERE host = 'two'; + ++----------+ +| count(*) | ++----------+ +| 0 | ++----------+ + +DROP TABLE json2_last_row; + +Affected Rows: 0 + +-- JSON2 is merged as a whole field, including empty objects and explicit JSON nulls. +CREATE TABLE json2_last_non_null ( + host STRING, + ts TIMESTAMP TIME INDEX, + j JSON2, + k JSON2(max_auto_expanded_paths = 0), + f DOUBLE, + PRIMARY KEY (host) +) WITH ( + 'append_mode' = 'false', + 'merge_mode' = 'last_non_null' +); + +Affected Rows: 0 + +INSERT INTO json2_last_non_null VALUES + ('one', 1, '{"a":1}', '{"a":1}', 10), + ('one', 2, '{"a":2}', '{"a":2}', 10), + ('one', 3, '{"a":3}', '{"a":3}', 10), + ('one', 4, '{"a":4}', '{"a":4}', 10), + ('one', 5, '{"a":{"x":5}}', '{"a":{"x":5}}', 10), + ('one', 6, '{"a":{"x":6}}', '{"a":{"x":6}}', 10), + ('one', 7, '{"a":{"x":8}}', '{"a":{"x":8}}', 10), + ('two', 1, '{"a":99}', '{"a":99}', 99); + +Affected Rows: 8 + +ADMIN FLUSH_TABLE('json2_last_non_null'); + ++------------------------------------------+ +| ADMIN FLUSH_TABLE('json2_last_non_null') | ++------------------------------------------+ +| 0 | ++------------------------------------------+ + +-- New objects replace old objects; only SQL NULL can fall back in last_non_null. +INSERT INTO json2_last_non_null VALUES + ('one', 1, '{"b":2}', '{"b":2}', NULL), + ('one', 2, '{}', '{}', NULL), + ('one', 3, '{"a":null}', '{"a":null}', NULL), + ('one', 4, NULL, NULL, NULL), + ('one', 5, '{"a":42}', '{"a":42}', NULL), + ('one', 6, '{"a":{"x":7}}', '{"a":{"x":7}}', NULL), + ('one', 7, NULL, NULL, NULL); + +Affected Rows: 7 + +-- Path-only reads must keep the new nested value and distinguish SQL NULL fallback. +SELECT host, ts, j, j.a AS a, j.a.x AS ax, k, k.a AS ka, f +FROM json2_last_non_null ORDER BY host, ts; + ++------+-------------------------+---------------+---------+----+---------------+---------+------+ +| host | ts | j | a | ax | k | ka | f | ++------+-------------------------+---------------+---------+----+---------------+---------+------+ +| one | 1970-01-01T00:00:00.001 | {"b":2} | | | {"b":2} | | 10.0 | +| one | 1970-01-01T00:00:00.002 | {} | | | {} | | 10.0 | +| one | 1970-01-01T00:00:00.003 | {"a":null} | | | {"a":null} | | 10.0 | +| one | 1970-01-01T00:00:00.004 | {"a":4} | 4 | | {"a":4} | 4 | 10.0 | +| one | 1970-01-01T00:00:00.005 | {"a":42} | 42 | | {"a":42} | 42 | 10.0 | +| one | 1970-01-01T00:00:00.006 | {"a":{"x":7}} | {"x":7} | 7 | {"a":{"x":7}} | {"x":7} | 10.0 | +| one | 1970-01-01T00:00:00.007 | {"a":{"x":8}} | {"x":8} | 8 | {"a":{"x":8}} | {"x":8} | 10.0 | +| two | 1970-01-01T00:00:00.001 | {"a":99} | 99 | | {"a":99} | 99 | 99.0 | ++------+-------------------------+---------------+---------+----+---------------+---------+------+ + +SELECT host, ts, j.a.x AS ax, k.a.x AS kax FROM json2_last_non_null ORDER BY host, ts; + ++------+-------------------------+----+-----+ +| host | ts | ax | kax | ++------+-------------------------+----+-----+ +| one | 1970-01-01T00:00:00.001 | | | +| one | 1970-01-01T00:00:00.002 | | | +| one | 1970-01-01T00:00:00.003 | | | +| one | 1970-01-01T00:00:00.004 | | | +| one | 1970-01-01T00:00:00.005 | | | +| one | 1970-01-01T00:00:00.006 | 7 | 7 | +| one | 1970-01-01T00:00:00.007 | 8 | 8 | +| two | 1970-01-01T00:00:00.001 | | | ++------+-------------------------+----+-----+ + +SELECT host, ts FROM json2_last_non_null WHERE j.a = '1' ORDER BY host, ts; + ++------+----+ +| host | ts | ++------+----+ ++------+----+ + +ADMIN FLUSH_TABLE('json2_last_non_null'); + ++------------------------------------------+ +| ADMIN FLUSH_TABLE('json2_last_non_null') | ++------------------------------------------+ +| 0 | ++------------------------------------------+ + +SELECT host, ts, j, j.a AS a, j.a.x AS ax, k, k.a AS ka, f +FROM json2_last_non_null ORDER BY host, ts; + ++------+-------------------------+---------------+---------+----+---------------+---------+------+ +| host | ts | j | a | ax | k | ka | f | ++------+-------------------------+---------------+---------+----+---------------+---------+------+ +| one | 1970-01-01T00:00:00.001 | {"b":2} | | | {"b":2} | | 10.0 | +| one | 1970-01-01T00:00:00.002 | {} | | | {} | | 10.0 | +| one | 1970-01-01T00:00:00.003 | {"a":null} | | | {"a":null} | | 10.0 | +| one | 1970-01-01T00:00:00.004 | {"a":4} | 4 | | {"a":4} | 4 | 10.0 | +| one | 1970-01-01T00:00:00.005 | {"a":42} | 42 | | {"a":42} | 42 | 10.0 | +| one | 1970-01-01T00:00:00.006 | {"a":{"x":7}} | {"x":7} | 7 | {"a":{"x":7}} | {"x":7} | 10.0 | +| one | 1970-01-01T00:00:00.007 | {"a":{"x":8}} | {"x":8} | 8 | {"a":{"x":8}} | {"x":8} | 10.0 | +| two | 1970-01-01T00:00:00.001 | {"a":99} | 99 | | {"a":99} | 99 | 99.0 | ++------+-------------------------+---------------+---------+----+---------------+---------+------+ + +SELECT host, ts, j.a.x AS ax, k.a.x AS kax FROM json2_last_non_null ORDER BY host, ts; + ++------+-------------------------+----+-----+ +| host | ts | ax | kax | ++------+-------------------------+----+-----+ +| one | 1970-01-01T00:00:00.001 | | | +| one | 1970-01-01T00:00:00.002 | | | +| one | 1970-01-01T00:00:00.003 | | | +| one | 1970-01-01T00:00:00.004 | | | +| one | 1970-01-01T00:00:00.005 | | | +| one | 1970-01-01T00:00:00.006 | 7 | 7 | +| one | 1970-01-01T00:00:00.007 | 8 | 8 | +| two | 1970-01-01T00:00:00.001 | | | ++------+-------------------------+----+-----+ + +ADMIN COMPACT_TABLE('json2_last_non_null'); + ++--------------------------------------------+ +| ADMIN COMPACT_TABLE('json2_last_non_null') | ++--------------------------------------------+ +| 0 | ++--------------------------------------------+ + +SELECT host, ts, j, j.a AS a, j.a.x AS ax, k, k.a AS ka, f +FROM json2_last_non_null ORDER BY host, ts; + ++------+-------------------------+---------------+---------+----+---------------+---------+------+ +| host | ts | j | a | ax | k | ka | f | ++------+-------------------------+---------------+---------+----+---------------+---------+------+ +| one | 1970-01-01T00:00:00.001 | {"b":2} | | | {"b":2} | | 10.0 | +| one | 1970-01-01T00:00:00.002 | {} | | | {} | | 10.0 | +| one | 1970-01-01T00:00:00.003 | {"a":null} | | | {"a":null} | | 10.0 | +| one | 1970-01-01T00:00:00.004 | {"a":4} | 4 | | {"a":4} | 4 | 10.0 | +| one | 1970-01-01T00:00:00.005 | {"a":42} | 42 | | {"a":42} | 42 | 10.0 | +| one | 1970-01-01T00:00:00.006 | {"a":{"x":7}} | {"x":7} | 7 | {"a":{"x":7}} | {"x":7} | 10.0 | +| one | 1970-01-01T00:00:00.007 | {"a":{"x":8}} | {"x":8} | 8 | {"a":{"x":8}} | {"x":8} | 10.0 | +| two | 1970-01-01T00:00:00.001 | {"a":99} | 99 | | {"a":99} | 99 | 99.0 | ++------+-------------------------+---------------+---------+----+---------------+---------+------+ + +SELECT host, ts, j.a.x AS ax, k.a.x AS kax FROM json2_last_non_null ORDER BY host, ts; + ++------+-------------------------+----+-----+ +| host | ts | ax | kax | ++------+-------------------------+----+-----+ +| one | 1970-01-01T00:00:00.001 | | | +| one | 1970-01-01T00:00:00.002 | | | +| one | 1970-01-01T00:00:00.003 | | | +| one | 1970-01-01T00:00:00.004 | | | +| one | 1970-01-01T00:00:00.005 | | | +| one | 1970-01-01T00:00:00.006 | 7 | 7 | +| one | 1970-01-01T00:00:00.007 | 8 | 8 | +| two | 1970-01-01T00:00:00.001 | | | ++------+-------------------------+----+-----+ + +-- SQLNESS ARG restart=true +SELECT host, ts, j, j.a AS a, j.a.x AS ax, k, k.a AS ka, f +FROM json2_last_non_null ORDER BY host, ts; + ++------+-------------------------+---------------+---------+----+---------------+---------+------+ +| host | ts | j | a | ax | k | ka | f | ++------+-------------------------+---------------+---------+----+---------------+---------+------+ +| one | 1970-01-01T00:00:00.001 | {"b":2} | | | {"b":2} | | 10.0 | +| one | 1970-01-01T00:00:00.002 | {} | | | {} | | 10.0 | +| one | 1970-01-01T00:00:00.003 | {"a":null} | | | {"a":null} | | 10.0 | +| one | 1970-01-01T00:00:00.004 | {"a":4} | 4 | | {"a":4} | 4 | 10.0 | +| one | 1970-01-01T00:00:00.005 | {"a":42} | 42 | | {"a":42} | 42 | 10.0 | +| one | 1970-01-01T00:00:00.006 | {"a":{"x":7}} | {"x":7} | 7 | {"a":{"x":7}} | {"x":7} | 10.0 | +| one | 1970-01-01T00:00:00.007 | {"a":{"x":8}} | {"x":8} | 8 | {"a":{"x":8}} | {"x":8} | 10.0 | +| two | 1970-01-01T00:00:00.001 | {"a":99} | 99 | | {"a":99} | 99 | 99.0 | ++------+-------------------------+---------------+---------+----+---------------+---------+------+ + +SELECT host, ts, j.a.x AS ax, k.a.x AS kax FROM json2_last_non_null ORDER BY host, ts; + ++------+-------------------------+----+-----+ +| host | ts | ax | kax | ++------+-------------------------+----+-----+ +| one | 1970-01-01T00:00:00.001 | | | +| one | 1970-01-01T00:00:00.002 | | | +| one | 1970-01-01T00:00:00.003 | | | +| one | 1970-01-01T00:00:00.004 | | | +| one | 1970-01-01T00:00:00.005 | | | +| one | 1970-01-01T00:00:00.006 | 7 | 7 | +| one | 1970-01-01T00:00:00.007 | 8 | 8 | +| two | 1970-01-01T00:00:00.001 | | | ++------+-------------------------+----+-----+ + +-- A delete must prevent last_non_null from filling fields from older versions. +DELETE FROM json2_last_non_null WHERE host = 'one' AND ts = 1; + +Affected Rows: 1 + +INSERT INTO json2_last_non_null VALUES ('one', 1, '{"c":3}', '{"c":3}', NULL); + +Affected Rows: 1 + +SELECT host, ts, j, j.a AS a, j.a.x AS ax, k, k.a AS ka, f +FROM json2_last_non_null ORDER BY host, ts; + ++------+-------------------------+---------------+---------+----+---------------+---------+------+ +| host | ts | j | a | ax | k | ka | f | ++------+-------------------------+---------------+---------+----+---------------+---------+------+ +| one | 1970-01-01T00:00:00.001 | {"c":3} | | | {"c":3} | | | +| one | 1970-01-01T00:00:00.002 | {} | | | {} | | 10.0 | +| one | 1970-01-01T00:00:00.003 | {"a":null} | | | {"a":null} | | 10.0 | +| one | 1970-01-01T00:00:00.004 | {"a":4} | 4 | | {"a":4} | 4 | 10.0 | +| one | 1970-01-01T00:00:00.005 | {"a":42} | 42 | | {"a":42} | 42 | 10.0 | +| one | 1970-01-01T00:00:00.006 | {"a":{"x":7}} | {"x":7} | 7 | {"a":{"x":7}} | {"x":7} | 10.0 | +| one | 1970-01-01T00:00:00.007 | {"a":{"x":8}} | {"x":8} | 8 | {"a":{"x":8}} | {"x":8} | 10.0 | +| two | 1970-01-01T00:00:00.001 | {"a":99} | 99 | | {"a":99} | 99 | 99.0 | ++------+-------------------------+---------------+---------+----+---------------+---------+------+ + +ADMIN FLUSH_TABLE('json2_last_non_null'); + ++------------------------------------------+ +| ADMIN FLUSH_TABLE('json2_last_non_null') | ++------------------------------------------+ +| 0 | ++------------------------------------------+ + +ADMIN COMPACT_TABLE('json2_last_non_null'); + ++--------------------------------------------+ +| ADMIN COMPACT_TABLE('json2_last_non_null') | ++--------------------------------------------+ +| 0 | ++--------------------------------------------+ + +SELECT host, ts, j, j.a AS a, j.a.x AS ax, k, k.a AS ka, f +FROM json2_last_non_null ORDER BY host, ts; + ++------+-------------------------+---------------+---------+----+---------------+---------+------+ +| host | ts | j | a | ax | k | ka | f | ++------+-------------------------+---------------+---------+----+---------------+---------+------+ +| one | 1970-01-01T00:00:00.001 | {"c":3} | | | {"c":3} | | | +| one | 1970-01-01T00:00:00.002 | {} | | | {} | | 10.0 | +| one | 1970-01-01T00:00:00.003 | {"a":null} | | | {"a":null} | | 10.0 | +| one | 1970-01-01T00:00:00.004 | {"a":4} | 4 | | {"a":4} | 4 | 10.0 | +| one | 1970-01-01T00:00:00.005 | {"a":42} | 42 | | {"a":42} | 42 | 10.0 | +| one | 1970-01-01T00:00:00.006 | {"a":{"x":7}} | {"x":7} | 7 | {"a":{"x":7}} | {"x":7} | 10.0 | +| one | 1970-01-01T00:00:00.007 | {"a":{"x":8}} | {"x":8} | 8 | {"a":{"x":8}} | {"x":8} | 10.0 | +| two | 1970-01-01T00:00:00.001 | {"a":99} | 99 | | {"a":99} | 99 | 99.0 | ++------+-------------------------+---------------+---------+----+---------------+---------+------+ + +DELETE FROM json2_last_non_null WHERE j.a = '99'; + +Affected Rows: 1 + +SELECT COUNT(*) FROM json2_last_non_null WHERE host = 'two'; + ++----------+ +| count(*) | ++----------+ +| 0 | ++----------+ + +DROP TABLE json2_last_non_null; + +Affected Rows: 0 + +-- A predicate on an added JSON2 field must retain older values of other fields. +CREATE TABLE json2_alter_backfill ( + host STRING, + ts TIMESTAMP TIME INDEX, + f DOUBLE, + PRIMARY KEY (host) +) WITH ('merge_mode' = 'last_non_null'); + +Affected Rows: 0 + +INSERT INTO json2_alter_backfill VALUES ('one', 1, 10); + +Affected Rows: 1 + +ADMIN FLUSH_TABLE('json2_alter_backfill'); + ++-------------------------------------------+ +| ADMIN FLUSH_TABLE('json2_alter_backfill') | ++-------------------------------------------+ +| 0 | ++-------------------------------------------+ + +ALTER TABLE json2_alter_backfill ADD COLUMN j JSON2; + +Affected Rows: 0 + +INSERT INTO json2_alter_backfill VALUES ('one', 1, NULL, '{"a":1}'); + +Affected Rows: 1 + +SELECT host, ts, j, f FROM json2_alter_backfill WHERE j IS NOT NULL; + ++------+-------------------------+---------+------+ +| host | ts | j | f | ++------+-------------------------+---------+------+ +| one | 1970-01-01T00:00:00.001 | {"a":1} | 10.0 | ++------+-------------------------+---------+------+ + +SELECT host, ts, j, f FROM json2_alter_backfill WHERE j.a = '1'; + ++------+-------------------------+---------+------+ +| host | ts | j | f | ++------+-------------------------+---------+------+ +| one | 1970-01-01T00:00:00.001 | {"a":1} | 10.0 | ++------+-------------------------+---------+------+ + +ADMIN FLUSH_TABLE('json2_alter_backfill'); + ++-------------------------------------------+ +| ADMIN FLUSH_TABLE('json2_alter_backfill') | ++-------------------------------------------+ +| 0 | ++-------------------------------------------+ + +SELECT host, ts, j, f FROM json2_alter_backfill WHERE j IS NOT NULL; + ++------+-------------------------+---------+------+ +| host | ts | j | f | ++------+-------------------------+---------+------+ +| one | 1970-01-01T00:00:00.001 | {"a":1} | 10.0 | ++------+-------------------------+---------+------+ + +ADMIN COMPACT_TABLE('json2_alter_backfill'); + ++---------------------------------------------+ +| ADMIN COMPACT_TABLE('json2_alter_backfill') | ++---------------------------------------------+ +| 0 | ++---------------------------------------------+ + +-- SQLNESS ARG restart=true +SELECT host, ts, j, f FROM json2_alter_backfill WHERE j IS NOT NULL; + ++------+-------------------------+---------+------+ +| host | ts | j | f | ++------+-------------------------+---------+------+ +| one | 1970-01-01T00:00:00.001 | {"a":1} | 10.0 | ++------+-------------------------+---------+------+ + +DROP TABLE json2_alter_backfill; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/types/json/json2_non_append.sql b/tests/cases/standalone/common/types/json/json2_non_append.sql new file mode 100644 index 00000000000..7d5082ffbea --- /dev/null +++ b/tests/cases/standalone/common/types/json/json2_non_append.sql @@ -0,0 +1,201 @@ +create table json2_without_append_mode ( + ts timestamp time index, + j json2 +); + +drop table json2_without_append_mode; + +create table json2_append_mode_false ( + ts timestamp time index, + j json2 +) with ( + 'append_mode' = 'false' +); + +drop table json2_append_mode_false; + +create table json2_alter_non_append ( + ts timestamp time index +); + +insert into json2_alter_non_append values (1); + +alter table json2_alter_non_append add column j json2; + +insert into json2_alter_non_append values (1, '{"a":1}'); + +select ts, j.a from json2_alter_non_append; + +drop table json2_alter_non_append; + +-- JSON2 is merged as a whole field, including empty objects and explicit JSON nulls. +CREATE TABLE json2_last_row ( + host STRING, + ts TIMESTAMP TIME INDEX, + j JSON2, + k JSON2(max_auto_expanded_paths = 0), + f DOUBLE, + PRIMARY KEY (host) +) WITH ( + 'append_mode' = 'false', + 'merge_mode' = 'last_row' +); + +INSERT INTO json2_last_row VALUES + ('one', 1, '{"a":1}', '{"a":1}', 10), + ('one', 2, '{"a":2}', '{"a":2}', 10), + ('one', 3, '{"a":3}', '{"a":3}', 10), + ('one', 4, '{"a":4}', '{"a":4}', 10), + ('one', 5, '{"a":{"x":5}}', '{"a":{"x":5}}', 10), + ('one', 6, '{"a":{"x":6}}', '{"a":{"x":6}}', 10), + ('one', 7, '{"a":{"x":8}}', '{"a":{"x":8}}', 10), + ('two', 1, '{"a":99}', '{"a":99}', 99); + +ADMIN FLUSH_TABLE('json2_last_row'); + +-- New objects replace old objects; only SQL NULL can fall back in last_non_null. +INSERT INTO json2_last_row VALUES + ('one', 1, '{"b":2}', '{"b":2}', NULL), + ('one', 2, '{}', '{}', NULL), + ('one', 3, '{"a":null}', '{"a":null}', NULL), + ('one', 4, NULL, NULL, NULL), + ('one', 5, '{"a":42}', '{"a":42}', NULL), + ('one', 6, '{"a":{"x":7}}', '{"a":{"x":7}}', NULL), + ('one', 7, NULL, NULL, NULL); + +-- Path-only reads must keep the new nested value and distinguish SQL NULL fallback. +SELECT host, ts, j, j.a AS a, j.a.x AS ax, k, k.a AS ka, f +FROM json2_last_row ORDER BY host, ts; +SELECT host, ts, j.a.x AS ax, k.a.x AS kax FROM json2_last_row ORDER BY host, ts; +SELECT host, ts FROM json2_last_row WHERE j.a = '1' ORDER BY host, ts; + +ADMIN FLUSH_TABLE('json2_last_row'); + +SELECT host, ts, j, j.a AS a, j.a.x AS ax, k, k.a AS ka, f +FROM json2_last_row ORDER BY host, ts; +SELECT host, ts, j.a.x AS ax, k.a.x AS kax FROM json2_last_row ORDER BY host, ts; + +ADMIN COMPACT_TABLE('json2_last_row'); + +SELECT host, ts, j, j.a AS a, j.a.x AS ax, k, k.a AS ka, f +FROM json2_last_row ORDER BY host, ts; +SELECT host, ts, j.a.x AS ax, k.a.x AS kax FROM json2_last_row ORDER BY host, ts; + +-- SQLNESS ARG restart=true +SELECT host, ts, j, j.a AS a, j.a.x AS ax, k, k.a AS ka, f +FROM json2_last_row ORDER BY host, ts; +SELECT host, ts, j.a.x AS ax, k.a.x AS kax FROM json2_last_row ORDER BY host, ts; + +-- A delete must prevent last_non_null from filling fields from older versions. +DELETE FROM json2_last_row WHERE host = 'one' AND ts = 1; +INSERT INTO json2_last_row VALUES ('one', 1, '{"c":3}', '{"c":3}', NULL); + +SELECT host, ts, j, j.a AS a, j.a.x AS ax, k, k.a AS ka, f +FROM json2_last_row ORDER BY host, ts; + +ADMIN FLUSH_TABLE('json2_last_row'); +ADMIN COMPACT_TABLE('json2_last_row'); + +SELECT host, ts, j, j.a AS a, j.a.x AS ax, k, k.a AS ka, f +FROM json2_last_row ORDER BY host, ts; + +DELETE FROM json2_last_row WHERE j.a = '99'; +SELECT COUNT(*) FROM json2_last_row WHERE host = 'two'; + +DROP TABLE json2_last_row; + +-- JSON2 is merged as a whole field, including empty objects and explicit JSON nulls. +CREATE TABLE json2_last_non_null ( + host STRING, + ts TIMESTAMP TIME INDEX, + j JSON2, + k JSON2(max_auto_expanded_paths = 0), + f DOUBLE, + PRIMARY KEY (host) +) WITH ( + 'append_mode' = 'false', + 'merge_mode' = 'last_non_null' +); + +INSERT INTO json2_last_non_null VALUES + ('one', 1, '{"a":1}', '{"a":1}', 10), + ('one', 2, '{"a":2}', '{"a":2}', 10), + ('one', 3, '{"a":3}', '{"a":3}', 10), + ('one', 4, '{"a":4}', '{"a":4}', 10), + ('one', 5, '{"a":{"x":5}}', '{"a":{"x":5}}', 10), + ('one', 6, '{"a":{"x":6}}', '{"a":{"x":6}}', 10), + ('one', 7, '{"a":{"x":8}}', '{"a":{"x":8}}', 10), + ('two', 1, '{"a":99}', '{"a":99}', 99); + +ADMIN FLUSH_TABLE('json2_last_non_null'); + +-- New objects replace old objects; only SQL NULL can fall back in last_non_null. +INSERT INTO json2_last_non_null VALUES + ('one', 1, '{"b":2}', '{"b":2}', NULL), + ('one', 2, '{}', '{}', NULL), + ('one', 3, '{"a":null}', '{"a":null}', NULL), + ('one', 4, NULL, NULL, NULL), + ('one', 5, '{"a":42}', '{"a":42}', NULL), + ('one', 6, '{"a":{"x":7}}', '{"a":{"x":7}}', NULL), + ('one', 7, NULL, NULL, NULL); + +-- Path-only reads must keep the new nested value and distinguish SQL NULL fallback. +SELECT host, ts, j, j.a AS a, j.a.x AS ax, k, k.a AS ka, f +FROM json2_last_non_null ORDER BY host, ts; +SELECT host, ts, j.a.x AS ax, k.a.x AS kax FROM json2_last_non_null ORDER BY host, ts; +SELECT host, ts FROM json2_last_non_null WHERE j.a = '1' ORDER BY host, ts; + +ADMIN FLUSH_TABLE('json2_last_non_null'); + +SELECT host, ts, j, j.a AS a, j.a.x AS ax, k, k.a AS ka, f +FROM json2_last_non_null ORDER BY host, ts; +SELECT host, ts, j.a.x AS ax, k.a.x AS kax FROM json2_last_non_null ORDER BY host, ts; + +ADMIN COMPACT_TABLE('json2_last_non_null'); + +SELECT host, ts, j, j.a AS a, j.a.x AS ax, k, k.a AS ka, f +FROM json2_last_non_null ORDER BY host, ts; +SELECT host, ts, j.a.x AS ax, k.a.x AS kax FROM json2_last_non_null ORDER BY host, ts; + +-- SQLNESS ARG restart=true +SELECT host, ts, j, j.a AS a, j.a.x AS ax, k, k.a AS ka, f +FROM json2_last_non_null ORDER BY host, ts; +SELECT host, ts, j.a.x AS ax, k.a.x AS kax FROM json2_last_non_null ORDER BY host, ts; + +-- A delete must prevent last_non_null from filling fields from older versions. +DELETE FROM json2_last_non_null WHERE host = 'one' AND ts = 1; +INSERT INTO json2_last_non_null VALUES ('one', 1, '{"c":3}', '{"c":3}', NULL); + +SELECT host, ts, j, j.a AS a, j.a.x AS ax, k, k.a AS ka, f +FROM json2_last_non_null ORDER BY host, ts; + +ADMIN FLUSH_TABLE('json2_last_non_null'); +ADMIN COMPACT_TABLE('json2_last_non_null'); + +SELECT host, ts, j, j.a AS a, j.a.x AS ax, k, k.a AS ka, f +FROM json2_last_non_null ORDER BY host, ts; + +DELETE FROM json2_last_non_null WHERE j.a = '99'; +SELECT COUNT(*) FROM json2_last_non_null WHERE host = 'two'; + +DROP TABLE json2_last_non_null; + +-- A predicate on an added JSON2 field must retain older values of other fields. +CREATE TABLE json2_alter_backfill ( + host STRING, + ts TIMESTAMP TIME INDEX, + f DOUBLE, + PRIMARY KEY (host) +) WITH ('merge_mode' = 'last_non_null'); +INSERT INTO json2_alter_backfill VALUES ('one', 1, 10); +ADMIN FLUSH_TABLE('json2_alter_backfill'); +ALTER TABLE json2_alter_backfill ADD COLUMN j JSON2; +INSERT INTO json2_alter_backfill VALUES ('one', 1, NULL, '{"a":1}'); +SELECT host, ts, j, f FROM json2_alter_backfill WHERE j IS NOT NULL; +SELECT host, ts, j, f FROM json2_alter_backfill WHERE j.a = '1'; +ADMIN FLUSH_TABLE('json2_alter_backfill'); +SELECT host, ts, j, f FROM json2_alter_backfill WHERE j IS NOT NULL; +ADMIN COMPACT_TABLE('json2_alter_backfill'); +-- SQLNESS ARG restart=true +SELECT host, ts, j, f FROM json2_alter_backfill WHERE j IS NOT NULL; +DROP TABLE json2_alter_backfill; From 3d2b73762a6275a4bb44400e750fea4aeda87478 Mon Sep 17 00:00:00 2001 From: "B. Naut" Date: Wed, 28 Oct 2026 03:47:17 +0000 Subject: [PATCH 116/239] fix(catalog): propagate frontend process-list failures (#9999) * fix(catalog): propagate frontend process-list failures Return the existing invocation error when a selected frontend cannot list processes, so a successful response cannot omit that peer. Signed-off-by: Byte-Naut * test(frontend): cover remote process-list errors and aggregation Exercise real metadata discovery and gRPC listeners for a missing management service and successful catalog-scoped aggregation. Signed-off-by: Byte-Naut --------- Signed-off-by: Byte-Naut --- src/catalog/src/process_manager.rs | 18 ++-- tests-integration/tests/grpc.rs | 165 ++++++++++++++++++++++++++++- 2 files changed, 170 insertions(+), 13 deletions(-) diff --git a/src/catalog/src/process_manager.rs b/src/catalog/src/process_manager.rs index 17df06d7648..c0af1f10a2f 100644 --- a/src/catalog/src/process_manager.rs +++ b/src/catalog/src/process_manager.rs @@ -25,7 +25,7 @@ use common_event_recorder::EventRecorderRef; use common_frontend::selector::{FrontendSelector, MetaClientSelector}; use common_frontend::slow_query_event::SlowQueryEvent; use common_telemetry::logging::SlowQueriesRecordType; -use common_telemetry::{debug, info, slow, warn}; +use common_telemetry::{debug, info, slow}; use common_time::util::current_time_millis; use meta_client::MetaClientRef; use promql_parser::parser::EvalStmt; @@ -166,6 +166,9 @@ impl ProcessManager { Ok(result) } + /// Lists running processes from every selected frontend and the local registry. + /// + /// Returns an error if a remote frontend cannot provide its process list. pub async fn list_all_processes( &self, catalog: Option<&str>, @@ -177,20 +180,13 @@ impl ProcessManager { .await .context(error::InvokeFrontendSnafu)?; for mut f in frontends { - let result = f + let resp = f .list_process(ListProcessRequest { catalog: catalog.unwrap_or_default().to_string(), }) .await - .context(error::InvokeFrontendSnafu); - match result { - Ok(resp) => { - processes.extend(resp.processes); - } - Err(e) => { - warn!(e; "Skipping failing node: {:?}", f) - } - } + .context(error::InvokeFrontendSnafu)?; + processes.extend(resp.processes); } } processes.extend(self.local_processes(catalog)?); diff --git a/tests-integration/tests/grpc.rs b/tests-integration/tests/grpc.rs index e9cfee7e9d4..818a940d76b 100644 --- a/tests-integration/tests/grpc.rs +++ b/tests-integration/tests/grpc.rs @@ -15,6 +15,7 @@ // Retain coverage of the legacy constructors, including TLS configuration. #![allow(deprecated)] +use std::net::SocketAddr; use std::sync::Arc; use api::v1::alter_table_expr::Kind; @@ -32,18 +33,23 @@ use api::v1::{ }; use auth::user_provider_from_option; use base64::prelude::{BASE64_STANDARD, Engine as _}; +use catalog::process_manager::{ProcessManager, ProcessManagerRef}; use client::{Client, DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, Database, OutputData}; use common_catalog::consts::{DEFAULT_PRIVATE_SCHEMA_NAME, MITO_ENGINE}; use common_event_recorder::DEFAULT_EVENTS_TABLE_NAME; use common_frontend::slow_query_event::SLOW_QUERY_TABLE_NAME; use common_grpc::channel_manager::ClientTlsOption; use common_memory_manager::OnExhaustedPolicy; +use common_meta::cluster::{FrontendStatus, NodeInfo, NodeInfoKey, NodeStatus, Role}; use common_meta::datanode::REGION_STATS_HISTORY_TABLE_NAME; +use common_meta::peer::Peer; +use common_meta::rpc::store::PutRequest; use common_query::Output; use common_recordbatch::RecordBatches; -use common_runtime::Runtime; use common_runtime::runtime::{BuilderBuild, RuntimeTrait}; +use common_runtime::{Builder as RuntimeBuilder, Runtime}; use common_test_util::find_workspace_path; +use common_time::util; use datatypes::arrow::array::{ Array, ArrayRef, Float64Array, Float64Builder, Int32Array, ListBuilder, StringArray, StructArray, TimestampNanosecondArray, UInt8Array, UInt16Array, UInt32Array, UInt64Array, @@ -52,6 +58,8 @@ use datatypes::arrow::array::{ use datatypes::arrow::datatypes::{DataType, Field}; use datatypes::arrow::ipc::writer::StreamWriter; use datatypes::arrow::record_batch::RecordBatch as ArrowRecordBatch; +use meta_client::client::MetaClientBuilder; +use meta_srv::mocks::{self, MockInfo}; use otel_arrow_rust::otlp::metrics::MetricType as ArrowMetricType; use otel_arrow_rust::proto::opentelemetry::arrow::v1::arrow_metrics_service_client::ArrowMetricsServiceClient; use otel_arrow_rust::proto::opentelemetry::arrow::v1::{ @@ -60,8 +68,9 @@ use otel_arrow_rust::proto::opentelemetry::arrow::v1::{ use otel_arrow_rust::proto::opentelemetry::metrics::v1::AggregationTemporality; use otel_arrow_rust::schema::consts as arrow_consts; use rstest_reuse::apply; -use servers::grpc::GrpcServerConfig; use servers::grpc::builder::GrpcServerBuilder; +use servers::grpc::frontend_grpc_handler::FrontendGrpcHandler; +use servers::grpc::{GrpcServer, GrpcServerConfig}; use servers::http::prometheus::{ PromData, PromQueryResult, PromSampleValue, PromSeriesMatrix, PromSeriesVector, PrometheusJsonResponse, PrometheusResponse, @@ -538,6 +547,158 @@ pub async fn test_grpc_auth(store_type: StorageType) { let _ = fe_grpc_server.shutdown().await; } +async fn start_process_rpc_server(process_manager: Option) -> GrpcServer { + let runtime: Runtime = RuntimeBuilder::default() + .worker_threads(2) + .thread_name("process-rpc-test") + .build() + .unwrap(); + let mut builder = GrpcServerBuilder::new(GrpcServerConfig::default(), runtime); + if let Some(process_manager) = process_manager { + builder = builder.frontend_grpc_handler(FrontendGrpcHandler::new(process_manager)); + } + let mut server = builder.build(); + server + .start("127.0.0.1:0".parse::().unwrap()) + .await + .unwrap(); + server +} + +async fn register_process_rpc_peer(mock: &MockInfo, addr: String) { + let node_info = NodeInfo { + peer: Peer { id: 1, addr }, + last_activity_ts: util::current_time_millis(), + status: NodeStatus::Frontend(FrontendStatus::default()), + version: String::new(), + git_commit: String::new(), + start_time_ms: 0, + total_cpu_millicores: 0, + total_memory_bytes: 0, + cpu_usage_millicores: 0, + memory_usage_bytes: 0, + hostname: String::new(), + env_vars: Default::default(), + schedulable: None, + }; + mock.in_memory + .as_ref() + .unwrap() + .put(PutRequest { + key: Vec::from(&NodeInfoKey { + role: Role::Frontend, + node_id: 1, + }), + value: node_info.try_into().unwrap(), + prev_kv: false, + }) + .await + .unwrap(); +} + +async fn build_process_manager_with_remote_peer( + mock: &MockInfo, + remote_addr: String, +) -> ProcessManagerRef { + register_process_rpc_peer(mock, remote_addr).await; + let mut meta_client = MetaClientBuilder::frontend_default_options() + .channel_manager(mock.channel_manager.clone()) + .build(); + meta_client.start(&[&mock.server_addr]).await.unwrap(); + Arc::new(ProcessManager::new( + "local".to_string(), + Some(Arc::new(meta_client)), + )) +} + +#[tokio::test(flavor = "multi_thread")] +async fn test_list_all_processes_reports_unimplemented_remote_service() { + let mock = mocks::mock_with_memstore().await; + let server = start_process_rpc_server(None).await; + let remote_addr = server.bind_addr().unwrap().to_string(); + let process_manager = build_process_manager_with_remote_peer(&mock, remote_addr).await; + let _local = process_manager.register_query( + "greptime".to_string(), + vec!["public".to_string()], + "local query".to_string(), + String::new(), + None, + None, + ); + + let result = process_manager.list_all_processes(Some("greptime")).await; + server.shutdown().await.unwrap(); + + assert!( + matches!( + &result, + Err(catalog::error::Error::InvokeFrontend { + source: common_frontend::error::Error::InvokeFrontend { error, .. }, + .. + }) if error.code() == tonic::Code::Unimplemented + ), + "a missing remote management service must not return an incomplete list: {result:?}" + ); + assert_eq!( + process_manager + .local_processes(Some("greptime")) + .unwrap() + .len(), + 1 + ); +} + +#[tokio::test(flavor = "multi_thread")] +async fn test_list_all_processes_collects_remote_and_local_catalog() { + let mock = mocks::mock_with_memstore().await; + let remote_manager = Arc::new(ProcessManager::new("remote".to_string(), None)); + let _selected = remote_manager.register_query( + "greptime".to_string(), + vec!["public".to_string()], + "selected remote query".to_string(), + String::new(), + None, + None, + ); + let _other = remote_manager.register_query( + "other".to_string(), + vec!["public".to_string()], + "other remote query".to_string(), + String::new(), + None, + None, + ); + let server = start_process_rpc_server(Some(remote_manager)).await; + let remote_addr = server.bind_addr().unwrap().to_string(); + let process_manager = build_process_manager_with_remote_peer(&mock, remote_addr).await; + let _local = process_manager.register_query( + "greptime".to_string(), + vec!["public".to_string()], + "selected local query".to_string(), + String::new(), + None, + None, + ); + + let processes = process_manager + .list_all_processes(Some("greptime")) + .await + .unwrap(); + server.shutdown().await.unwrap(); + + assert_eq!(processes.len(), 2); + assert!( + processes + .iter() + .any(|process| process.query == "selected local query") + ); + assert!( + processes + .iter() + .any(|process| process.query == "selected remote query") + ); +} + pub async fn test_otel_arrow_auth(store_type: StorageType) { let user_provider = user_provider_from_option("static_user_provider:cmd:greptime_user=greptime_pwd").unwrap(); From 5f0e22171ed54c4bb1245492648e2b5d04704b68 Mon Sep 17 00:00:00 2001 From: Dragon Roar <2181719471@qq.com> Date: Wed, 28 Oct 2026 06:55:55 +0000 Subject: [PATCH 117/239] fix(log-store): skip records with unparsable keys instead of aborting the read (#9749) * fix(log-store): skip records with unparsable keys instead of aborting the read Signed-off-by: lczllx <2181719471@qq.com> * fix(log-store): keep undecodable RecordMeta keys fatal Signed-off-by: lczllx <2181719471@qq.com> --------- Signed-off-by: lczllx <2181719471@qq.com> --- src/log-store/src/kafka/log_store.rs | 511 ++++++++++++++++++++++--- src/log-store/src/kafka/util/record.rs | 52 ++- 2 files changed, 503 insertions(+), 60 deletions(-) diff --git a/src/log-store/src/kafka/log_store.rs b/src/log-store/src/kafka/log_store.rs index daa6dfbcf21..33cdd242564 100644 --- a/src/log-store/src/kafka/log_store.rs +++ b/src/log-store/src/kafka/log_store.rs @@ -40,7 +40,8 @@ use crate::kafka::consumer::{ConsumerBuilder, RecordsBuffer, RegionWalRange}; use crate::kafka::periodic_offset_fetcher::PeriodicOffsetFetcher; use crate::kafka::producer::OrderedBatchProducerRef; use crate::kafka::util::record::{ - ESTIMATED_META_SIZE, Record, convert_to_kafka_records, maybe_emit_entry, remaining_entries, + ESTIMATED_META_SIZE, Record, convert_to_kafka_records, is_record_meta_key, maybe_emit_entry, + remaining_entries, }; use crate::metrics; @@ -408,26 +409,41 @@ impl LogStore for KafkaLogStore { ); // Ignores no-op records. - if kafka_record.value.is_none() { - if check_termination(offset, end_offset) { - if let Some(entries) = remaining_entries(&provider, &mut entry_records) { - yield Ok(entries); + if kafka_record.value.is_some() { + let key = kafka_record.key; + match Record::try_from_parts(key.as_deref(), kafka_record.value) { + Ok(record) => { + // Tries to construct an entry from records consumed so far. + if let Some(mut entry) = + maybe_emit_entry(&provider, record, &mut entry_records)? + { + // We don't rely on the EntryId generated by mito2. + // Instead, we use the offset return from Kafka as EntryId. + // Therefore, we MUST overwrite the EntryId with RecordOffset. + entry.set_entry_id(offset as u64); + yield Ok(vec![entry]); + } + } + Err(err) => { + // A key that cannot be decoded is skipped only when the log store + // never wrote it. A key that is a JSON object with a `version` + // field may be a `RecordMeta` written by an incompatible version, + // for example one with a newer `RecordType` variant read after a + // rollback or during a rolling upgrade; skipping it would lose the + // entry once the WAL is pruned past its offset, so it stays fatal. + if key.as_deref().is_some_and(is_record_meta_key) { + yield Err(err); + return; + } + warn!( + "Skipping unparsable record at offset {} of topic {}, reason: {}", + offset, provider.topic, err + ); } - break; } - continue; - } - - let record = Record::try_from(kafka_record)?; - // Tries to construct an entry from records consumed so far. - if let Some(mut entry) = maybe_emit_entry(&provider, record, &mut entry_records)? { - // We don't rely on the EntryId generated by mito2. - // Instead, we use the offset return from Kafka as EntryId. - // Therefore, we MUST overwrite the EntryId with RecordOffset. - entry.set_entry_id(offset as u64); - yield Ok(vec![entry]); } + // The termination check runs for every consumed record, including skipped ones. if check_termination(offset, end_offset) { if let Some(entries) = remaining_entries(&provider, &mut entry_records) { yield Ok(entries); @@ -529,21 +545,25 @@ mod tests { use common_base::readable_size::ReadableSize; use common_meta::datanode::TopicStatsReporter; use common_telemetry::info; - use common_telemetry::tracing::warn; use common_wal::config::kafka::DatanodeKafkaConfig; use common_wal::config::kafka::common::KafkaConnectionConfig; + use common_wal::maybe_skip_kafka_integration_test; + use common_wal::test_util::get_kafka_endpoints; use dashmap::DashMap; use futures::TryStreamExt; use rand::Rng; use rand::prelude::SliceRandom; - use rskafka::client::partition::OffsetAt; + use rskafka::client::partition::{Compression, OffsetAt}; + use rskafka::record::Record as KafkaRecord; use store_api::logstore::LogStore; use store_api::logstore::entry::{Entry, MultiplePartEntry, MultiplePartHeader, NaiveEntry}; use store_api::logstore::provider::Provider; use store_api::storage::RegionId; use super::build_entry; + use crate::error; use crate::kafka::log_store::{KafkaLogStore, PeriodicTopicStatsReporter, TopicStat}; + use crate::kafka::util::record::{ESTIMATED_META_SIZE, convert_to_kafka_records}; #[test] fn test_build_naive_entry() { @@ -630,26 +650,76 @@ mod tests { .unwrap(); } + async fn new_test_logstore(max_batch_bytes: ReadableSize) -> KafkaLogStore { + let config = DatanodeKafkaConfig { + connection: KafkaConnectionConfig { + broker_endpoints: get_kafka_endpoints(), + ..Default::default() + }, + max_batch_bytes, + ..Default::default() + }; + KafkaLogStore::try_new(&config).await.unwrap() + } + + async fn produce_records( + logstore: &KafkaLogStore, + provider: &Provider, + records: Vec, + ) { + // A freshly created topic can be missing from the broker's metadata for a short + // while; `try_create_client` already retries that case. + let client = logstore + .client_manager + .get_or_insert(provider.as_kafka_provider().unwrap()) + .await + .unwrap() + .client() + .clone(); + client + .produce(records, Compression::NoCompression) + .await + .unwrap(); + } + + /// Builds a record that was not written by this log store. + fn foreign_record(key: Option<&[u8]>) -> KafkaRecord { + KafkaRecord { + key: key.map(|key| key.to_vec()), + value: Some(b"foreign payload".to_vec()), + timestamp: chrono::Utc::now(), + headers: Default::default(), + } + } + + /// Returns the log end offset of the topic, i.e. the exclusive end of the range that + /// `read` consumes. + async fn log_end_offset(logstore: &KafkaLogStore, provider: &Provider) -> i64 { + let client = logstore + .client_manager + .get_or_insert(provider.as_kafka_provider().unwrap()) + .await + .unwrap(); + client.client().get_offset(OffsetAt::Latest).await.unwrap() + } + + /// Returns one record per kind of key that the log store never writes: a missing key, a + /// key that is not JSON, a JSON value that is not an object, and a JSON object without a + /// `version` field. + fn foreign_records() -> Vec { + vec![ + foreign_record(None), + foreign_record(Some(b"not-a-json-key")), + foreign_record(Some(b"[1, 2, 3]")), + foreign_record(Some(br#"{"hello": "world"}"#)), + ] + } + #[tokio::test] async fn test_append_batch_basic() { common_telemetry::init_default_ut_logging(); - let Ok(broker_endpoints) = std::env::var("GT_KAFKA_ENDPOINTS") else { - warn!("The endpoints is empty, skipping the test 'test_append_batch_basic'"); - return; - }; - let broker_endpoints = broker_endpoints - .split(',') - .map(|s| s.trim().to_string()) - .collect::>(); - let config = DatanodeKafkaConfig { - connection: KafkaConnectionConfig { - broker_endpoints, - ..Default::default() - }, - max_batch_bytes: ReadableSize::kb(32), - ..Default::default() - }; - let logstore = KafkaLogStore::try_new(&config).await.unwrap(); + maybe_skip_kafka_integration_test!(); + let logstore = new_test_logstore(ReadableSize::kb(32)).await; let topic_name = uuid::Uuid::new_v4().to_string(); prepare_topic(&logstore, &topic_name).await; let provider = Provider::kafka_provider(topic_name); @@ -715,25 +785,10 @@ mod tests { #[tokio::test] async fn test_append_batch_basic_large() { common_telemetry::init_default_ut_logging(); - let Ok(broker_endpoints) = std::env::var("GT_KAFKA_ENDPOINTS") else { - warn!("The endpoints is empty, skipping the test 'test_append_batch_basic_large'"); - return; - }; + maybe_skip_kafka_integration_test!(); let data_size_kb = rand::rng().random_range(9..31usize); info!("Entry size: {}Ki", data_size_kb); - let broker_endpoints = broker_endpoints - .split(',') - .map(|s| s.trim().to_string()) - .collect::>(); - let config = DatanodeKafkaConfig { - connection: KafkaConnectionConfig { - broker_endpoints, - ..Default::default() - }, - max_batch_bytes: ReadableSize::kb(8), - ..Default::default() - }; - let logstore = KafkaLogStore::try_new(&config).await.unwrap(); + let logstore = new_test_logstore(ReadableSize::kb(8)).await; let topic_name = uuid::Uuid::new_v4().to_string(); prepare_topic(&logstore, &topic_name).await; let provider = Provider::kafka_provider(topic_name); @@ -788,6 +843,356 @@ mod tests { assert_eq!(high_wathermark, (data_size_kb as u64 / 8 + 1) * 20 * 5 - 1); } + #[tokio::test] + async fn test_read_skips_unparsable_record_inside_multipart_entry() { + maybe_skip_kafka_integration_test!(); + let logstore = new_test_logstore(ReadableSize::kb(8)).await; + let topic_name = uuid::Uuid::new_v4().to_string(); + prepare_topic(&logstore, &topic_name).await; + let provider = Provider::kafka_provider(topic_name); + + let region_id = RegionId::new(1, 1); + // Splits the entry into First, Middle(1) and Last. + let max_data_size = 8 * 1024 - ESTIMATED_META_SIZE; + let data = vec![7u8; max_data_size * 2 + 10]; + let entry = logstore + .entry(data.clone(), 0, region_id, &provider) + .unwrap(); + assert_matches!(entry, Entry::MultiplePart(_)); + let mut records = convert_to_kafka_records(entry).unwrap(); + assert_eq!(records.len(), 3); + + let last = records.pop().unwrap(); + // Places an unparsable record before the entry and between every pair of parts. + produce_records( + &logstore, + &provider, + vec![foreign_record(Some(b"not-a-json-key"))], + ) + .await; + produce_records(&logstore, &provider, vec![records[0].clone()]).await; + produce_records( + &logstore, + &provider, + vec![foreign_record(Some(br#"{"hello": "world"}"#))], + ) + .await; + produce_records(&logstore, &provider, vec![records[1].clone()]).await; + produce_records( + &logstore, + &provider, + vec![foreign_record(Some(b"not-a-json-key"))], + ) + .await; + produce_records(&logstore, &provider, vec![last]).await; + + let entries = logstore + .read(&provider, 0) + .await + .unwrap() + .try_collect::>() + .await + .unwrap() + .into_iter() + .flatten() + .collect::>(); + assert_eq!(entries.len(), 1); + let entry = entries + .into_iter() + .next() + .unwrap() + .into_multiple_part_entry() + .unwrap(); + assert_eq!( + entry.headers, + vec![ + MultiplePartHeader::First, + MultiplePartHeader::Middle(1), + MultiplePartHeader::Last + ] + ); + assert_eq!(entry.parts.concat(), data); + } + + #[tokio::test] + async fn test_read_skips_unparsable_record_across_regions() { + maybe_skip_kafka_integration_test!(); + let logstore = new_test_logstore(ReadableSize::kb(8)).await; + let topic_name = uuid::Uuid::new_v4().to_string(); + prepare_topic(&logstore, &topic_name).await; + let provider = Provider::kafka_provider(topic_name); + + let region_a = RegionId::new(1, 1); + let region_b = RegionId::new(1, 2); + let max_data_size = 8 * 1024 - ESTIMATED_META_SIZE; + let data_a = vec![1u8; max_data_size * 2 + 10]; + let data_b = vec![2u8; max_data_size * 2 + 10]; + let records_a = convert_to_kafka_records( + logstore + .entry(data_a.clone(), 0, region_a, &provider) + .unwrap(), + ) + .unwrap(); + let records_b = convert_to_kafka_records( + logstore + .entry(data_b.clone(), 0, region_b, &provider) + .unwrap(), + ) + .unwrap(); + assert_eq!(records_a.len(), 3); + assert_eq!(records_b.len(), 3); + + // Interleaves the parts of both entries with unparsable records. + for i in 0..3 { + produce_records(&logstore, &provider, vec![records_a[i].clone()]).await; + produce_records(&logstore, &provider, vec![records_b[i].clone()]).await; + produce_records( + &logstore, + &provider, + vec![foreign_record(Some(b"not-a-json-key"))], + ) + .await; + } + + let entries = logstore + .read(&provider, 0) + .await + .unwrap() + .try_collect::>() + .await + .unwrap() + .into_iter() + .flatten() + .collect::>(); + assert_eq!(entries.len(), 2); + + let mut entries = entries + .into_iter() + .map(|entry| (entry.region_id(), entry.into_multiple_part_entry().unwrap())) + .collect::>(); + assert_eq!(entries.remove(®ion_a).unwrap().parts.concat(), data_a); + assert_eq!(entries.remove(®ion_b).unwrap().parts.concat(), data_b); + assert!(entries.is_empty()); + } + + #[tokio::test] + async fn test_read_finishes_when_last_record_is_unparsable() { + maybe_skip_kafka_integration_test!(); + let logstore = new_test_logstore(ReadableSize::kb(32)).await; + let topic_name = uuid::Uuid::new_v4().to_string(); + prepare_topic(&logstore, &topic_name).await; + let provider = Provider::kafka_provider(topic_name); + + let region_id = RegionId::new(1, 1); + let mut expected = generate_entries(&logstore, &provider, 3, region_id, 128); + logstore.append_batch(expected.clone()).await.unwrap(); + // Appends an unparsable record as the last record of the read range. + produce_records( + &logstore, + &provider, + vec![foreign_record(Some(b"not-a-json-key"))], + ) + .await; + + let entries = tokio::time::timeout( + Duration::from_secs(30), + logstore + .read(&provider, 0) + .await + .unwrap() + .try_collect::>(), + ) + .await + .expect("read should finish after skipping the unparsable record") + .unwrap() + .into_iter() + .flatten() + .collect::>(); + + expected.iter_mut().for_each(|entry| entry.set_entry_id(0)); + let mut entries = entries; + entries.iter_mut().for_each(|entry| entry.set_entry_id(0)); + assert_eq!(entries, expected); + } + + #[tokio::test] + async fn test_read_skips_foreign_records_of_every_kind() { + maybe_skip_kafka_integration_test!(); + let logstore = new_test_logstore(ReadableSize::kb(32)).await; + let topic_name = uuid::Uuid::new_v4().to_string(); + prepare_topic(&logstore, &topic_name).await; + let provider = Provider::kafka_provider(topic_name); + + // Places every kind of unparsable record before the first entry. + for record in foreign_records() { + produce_records(&logstore, &provider, vec![record]).await; + } + + let region_id = RegionId::new(1, 1); + let mut expected = generate_entries(&logstore, &provider, 1, region_id, 128); + logstore.append_batch(expected.clone()).await.unwrap(); + + let entries = logstore + .read(&provider, 0) + .await + .unwrap() + .try_collect::>() + .await + .unwrap() + .into_iter() + .flatten() + .collect::>(); + + expected.iter_mut().for_each(|entry| entry.set_entry_id(0)); + let mut entries = entries; + entries.iter_mut().for_each(|entry| entry.set_entry_id(0)); + assert_eq!(entries, expected); + } + + #[tokio::test] + async fn test_read_finishes_when_topic_only_has_foreign_records() { + maybe_skip_kafka_integration_test!(); + let logstore = new_test_logstore(ReadableSize::kb(32)).await; + let topic_name = uuid::Uuid::new_v4().to_string(); + prepare_topic(&logstore, &topic_name).await; + let provider = Provider::kafka_provider(topic_name); + + for record in foreign_records() { + produce_records(&logstore, &provider, vec![record]).await; + } + + // Asserts the read range covers the produced records; an empty range would let `read` + // return without exercising the skip path. + assert_eq!(log_end_offset(&logstore, &provider).await, 4); + + let entries = tokio::time::timeout( + Duration::from_secs(30), + logstore + .read(&provider, 0) + .await + .unwrap() + .try_collect::>(), + ) + .await + .expect("read should finish after skipping all unparsable records") + .unwrap(); + + assert!(entries.into_iter().flatten().next().is_none()); + } + + #[tokio::test] + async fn test_read_fails_on_incompatible_record_meta() { + maybe_skip_kafka_integration_test!(); + let logstore = new_test_logstore(ReadableSize::kb(32)).await; + + // Keys that look like a `RecordMeta` but cannot be decoded by this version. They may + // come from an incompatible version of the log store, so the read must not skip them. + let keys: [&[u8]; 2] = [ + // A `RecordType` variant this version does not know. + br#"{"version": 0, "tp": "Compressed", "entry_id": 0, "ns": {"region_id": 1, "topic": ""}}"#, + // A field with the wrong type. + br#"{"version": "0", "tp": "Full", "entry_id": 0, "ns": {"region_id": 1, "topic": ""}}"#, + ]; + for key in keys { + let topic_name = uuid::Uuid::new_v4().to_string(); + prepare_topic(&logstore, &topic_name).await; + let provider = Provider::kafka_provider(topic_name); + produce_records(&logstore, &provider, vec![foreign_record(Some(key))]).await; + + let err = logstore + .read(&provider, 0) + .await + .unwrap() + .try_collect::>>() + .await + .unwrap_err(); + assert_matches!(err, error::Error::DecodeJson { .. }); + } + } + + #[tokio::test] + async fn test_read_fails_on_missing_middle_part() { + maybe_skip_kafka_integration_test!(); + let logstore = new_test_logstore(ReadableSize::kb(8)).await; + let topic_name = uuid::Uuid::new_v4().to_string(); + prepare_topic(&logstore, &topic_name).await; + let provider = Provider::kafka_provider(topic_name); + + let region_id = RegionId::new(1, 1); + // Splits the entry into First, Middle(1), Middle(2) and Last. + let max_data_size = 8 * 1024 - ESTIMATED_META_SIZE; + let data = vec![7u8; max_data_size * 3 + 10]; + let entry = logstore.entry(data, 0, region_id, &provider).unwrap(); + let records = convert_to_kafka_records(entry).unwrap(); + assert_eq!(records.len(), 4); + + // Only a missing part of an entry we wrote is fatal; a foreign record is skipped + // because it cannot be one of our parts. + produce_records( + &logstore, + &provider, + vec![foreign_record(Some(b"not-a-json-key"))], + ) + .await; + produce_records(&logstore, &provider, vec![records[0].clone()]).await; + produce_records(&logstore, &provider, vec![records[2].clone()]).await; + + let err = logstore + .read(&provider, 0) + .await + .unwrap() + .try_collect::>>() + .await + .unwrap_err(); + assert_matches!(err, error::Error::IllegalSequence { .. }); + } + + #[tokio::test] + async fn test_read_succeeds_on_two_part_entry() { + maybe_skip_kafka_integration_test!(); + let logstore = new_test_logstore(ReadableSize::kb(8)).await; + let topic_name = uuid::Uuid::new_v4().to_string(); + prepare_topic(&logstore, &topic_name).await; + let provider = Provider::kafka_provider(topic_name); + + let region_id = RegionId::new(1, 1); + // Splits the entry into First and Last, without any middle part. + let max_data_size = 8 * 1024 - ESTIMATED_META_SIZE; + let data = vec![7u8; max_data_size + 10]; + let entry = logstore + .entry(data.clone(), 0, region_id, &provider) + .unwrap(); + let records = convert_to_kafka_records(entry).unwrap(); + assert_eq!(records.len(), 2); + + for record in records { + produce_records(&logstore, &provider, vec![record]).await; + } + + let entries = logstore + .read(&provider, 0) + .await + .unwrap() + .try_collect::>() + .await + .unwrap() + .into_iter() + .flatten() + .collect::>(); + assert_eq!(entries.len(), 1); + let entry = entries + .into_iter() + .next() + .unwrap() + .into_multiple_part_entry() + .unwrap(); + assert_eq!( + entry.headers, + vec![MultiplePartHeader::First, MultiplePartHeader::Last] + ); + assert_eq!(entry.parts.concat(), data); + } + #[tokio::test] async fn test_topic_stats_reporter() { common_telemetry::init_default_ut_logging(); diff --git a/src/log-store/src/kafka/util/record.rs b/src/log-store/src/kafka/util/record.rs index aefb098e0db..97311bff006 100644 --- a/src/log-store/src/kafka/util/record.rs +++ b/src/log-store/src/kafka/util/record.rs @@ -108,17 +108,32 @@ impl TryFrom for KafkaRecord { } // TODO(niebayes): improve the performance of decoding kafka record. -impl TryFrom for Record { - type Error = crate::error::Error; - - fn try_from(kafka_record: KafkaRecord) -> Result { - let key = kafka_record.key.context(MissingKeySnafu)?; - let meta = serde_json::from_slice(&key).context(DecodeJsonSnafu)?; - let data = kafka_record.value.context(MissingValueSnafu)?; +impl Record { + /// Decodes a record from the key and value of a [`KafkaRecord`]. + /// + /// The key is borrowed so that callers can inspect it if decoding fails, e.g. to + /// tell a foreign key apart from a [`RecordMeta`] written by an incompatible version + /// of the log store. + pub(crate) fn try_from_parts(key: Option<&[u8]>, value: Option>) -> Result { + let key = key.context(MissingKeySnafu)?; + let meta = serde_json::from_slice(key).context(DecodeJsonSnafu)?; + let data = value.context(MissingValueSnafu)?; Ok(Self { meta, data }) } } +/// Returns whether a key is a JSON object with a `version` field, i.e. whether it may have +/// been written as a [`RecordMeta`] key. +/// +/// A `RecordMeta` key that cannot be decoded may come from an incompatible version of the +/// log store, so the reader must not treat such a key as a foreign record. +pub(crate) fn is_record_meta_key(key: &[u8]) -> bool { + matches!( + serde_json::from_slice::(key), + Ok(serde_json::Value::Object(fields)) if fields.contains_key("version") + ) +} + pub(crate) fn convert_to_kafka_records(entry: Entry) -> Result> { match entry { Entry::Naive(entry) => Ok(vec![KafkaRecord::try_from(Record { @@ -559,6 +574,29 @@ mod tests { assert_eq!(full_entry.entry_id(), 3); } + #[test] + fn test_is_record_meta_key() { + let meta = RecordMeta { + version: VERSION, + tp: RecordType::Full, + entry_id: 1, + ns: NamespaceImpl { + region_id: RegionId::new(1, 1).as_u64(), + topic: "greptimedb_wal_topic".to_string(), + }, + }; + assert!(is_record_meta_key(&serde_json::to_vec(&meta).unwrap())); + // A `RecordMeta` key this version cannot decode, e.g. one with a new `RecordType` + // variant. + assert!(is_record_meta_key( + br#"{"version": 0, "tp": "Compressed", "entry_id": 0, "ns": {"region_id": 1, "topic": ""}}"# + )); + // Keys the log store never writes. + assert!(!is_record_meta_key(b"not-a-json-key")); + assert!(!is_record_meta_key(b"[1, 2, 3]")); + assert!(!is_record_meta_key(br#"{"hello": "world"}"#)); + } + #[test] fn test_meta_size() { let meta = RecordMeta { From 724072bb33a442ce8e716f4263c4da4b33704231 Mon Sep 17 00:00:00 2001 From: Chenglong Zheng Date: Wed, 28 Oct 2026 06:56:11 +0000 Subject: [PATCH 118/239] =?UTF-8?q?feat:=20Add=20last=5Fsequence=20for=20m?= =?UTF-8?q?ax=20allocated=20sequence=20and=20add=20tests=20for=20=E2=80=A6?= =?UTF-8?q?=20(#9951)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: Add last_sequence for max allocated sequence and add tests for the reproduced issue cases Signed-off-by: hahaha * test: Add a regression test for Kafka remote catchup Signed-off-by: hahaha --------- Signed-off-by: hahaha --- src/mito2/src/engine.rs | 2 + src/mito2/src/engine/catchup_test.rs | 291 ++++++++++++- src/mito2/src/engine/edit_region_test.rs | 68 ++- .../engine/object_store_wal_recovery_test.rs | 19 +- src/mito2/src/engine/sequence_test.rs | 410 ++++++++++++++++++ src/mito2/src/engine/staging_test.rs | 46 +- src/mito2/src/engine/sync_test.rs | 83 +++- src/mito2/src/region/version.rs | 76 +++- src/mito2/src/region_write_ctx.rs | 200 ++++++++- src/mito2/src/worker/handle_catchup.rs | 4 + src/mito2/src/worker/handle_manifest.rs | 7 +- src/mito2/src/worker/handle_write.rs | 23 +- 12 files changed, 1181 insertions(+), 48 deletions(-) create mode 100644 src/mito2/src/engine/sequence_test.rs diff --git a/src/mito2/src/engine.rs b/src/mito2/src/engine.rs index 0147fb71e2a..3791489d73f 100644 --- a/src/mito2/src/engine.rs +++ b/src/mito2/src/engine.rs @@ -74,6 +74,8 @@ mod scan_corrupt; #[cfg(test)] mod scan_test; #[cfg(test)] +mod sequence_test; +#[cfg(test)] mod set_role_state_test; #[cfg(test)] mod skip_wal_test; diff --git a/src/mito2/src/engine/catchup_test.rs b/src/mito2/src/engine/catchup_test.rs index 5c8e357c056..b470db5e12a 100644 --- a/src/mito2/src/engine/catchup_test.rs +++ b/src/mito2/src/engine/catchup_test.rs @@ -14,6 +14,7 @@ use std::assert_matches; use std::collections::HashMap; +use std::sync::Arc; use api::v1::Rows; use common_error::ext::ErrorExt; @@ -32,6 +33,7 @@ use store_api::storage::{RegionId, ScanRequest}; use crate::config::MitoConfig; use crate::engine::MitoEngine; use crate::error::Error; +use crate::manifest::action::RegionEdit; use crate::test_util::{ CreateRequestBuilder, LogStoreFactory, TestEnv, build_rows, flush_region, kafka_log_store_factory, prepare_test_for_kafka_log_store, put_rows, @@ -168,6 +170,220 @@ async fn test_catchup_with_last_entry_id(factory: Option) { assert!(resp.is_ok()); } +#[apply(single_kafka_log_store_factory)] +async fn test_catchup_preserves_region_edit_sequence( + factory: Option, + #[values(false, true)] open_before_edit: bool, +) { + use store_api::region_engine::SettableRegionRoleState; + + common_telemetry::init_default_ut_logging(); + let Some(factory) = factory else { + return; + }; + + let mut env = TestEnv::with_prefix("catchup_edit_sequence") + .await + .with_log_store_factory(factory.clone()); + let topic = prepare_test_for_kafka_log_store(&factory).await.unwrap(); + let config = MitoConfig { + schedule_compaction_after_edit: false, + ..Default::default() + }; + let leader_engine = env.create_engine(config.clone()).await; + let follower_engine = env.create_follower_engine(config).await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new() + .kafka_topic(Some(topic.clone())) + .build(); + let column_schemas = rows_schema(&request); + let open_request = RegionOpenRequest { + engine: String::new(), + table_dir: request.table_dir.clone(), + path_type: PathType::Bare, + options: HashMap::from([( + WAL_OPTIONS_KEY.to_string(), + serde_json::to_string(&WalOptions::Kafka(KafkaWalOptions::new(topic))).unwrap(), + )]), + skip_wal_replay: true, + checkpoint: None, + requirements: Default::default(), + }; + leader_engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + put_rows( + &leader_engine, + region_id, + Rows { + schema: column_schemas.clone(), + rows: build_rows(0, 3), + }, + ) + .await; + flush_region(&leader_engine, region_id, None).await; + let leader = leader_engine.get_region(region_id).unwrap(); + let flushed_entry_id = leader.version().flushed_entry_id; + let flushed_sequence = leader.version().flushed_sequence; + let flushed_rows = scan_region(&leader_engine, region_id) + .await + .pretty_print() + .unwrap(); + + put_rows( + &leader_engine, + region_id, + Rows { + schema: column_schemas.clone(), + rows: build_rows(3, 5), + }, + ) + .await; + let unflushed = leader.version_control.current(); + let unflushed_entry_id = unflushed.last_entry_id; + assert!(unflushed_entry_id > flushed_entry_id); + assert_eq!(unflushed.version.flushed_entry_id, flushed_entry_id); + + if open_before_edit { + follower_engine + .handle_request(region_id, RegionRequest::Open(open_request.clone())) + .await + .unwrap(); + } + // Reuses the flushed SST without flushing the second batch. + let file = leader.version().ssts.levels()[0] + .files + .values() + .next() + .unwrap() + .meta_ref() + .clone(); + leader_engine + .edit_region( + region_id, + RegionEdit { + files_to_add: vec![file], + files_to_remove: vec![], + timestamp_ms: None, + compaction_time_window: None, + flushed_entry_id: None, + flushed_sequence: None, + committed_sequence: None, + }, + ) + .await + .unwrap(); + let edit_sequence = leader.version_control.committed_sequence(); + assert!(edit_sequence > unflushed.committed_sequence); + let expected_rows = scan_region(&leader_engine, region_id) + .await + .pretty_print() + .unwrap(); + assert_ne!(flushed_rows, expected_rows); + + if !open_before_edit { + follower_engine + .handle_request(region_id, RegionRequest::Open(open_request)) + .await + .unwrap(); + } + let candidate_before_catchup = follower_engine.get_region(region_id).unwrap(); + let before_catchup = candidate_before_catchup.version_control.current(); + assert!(candidate_before_catchup.provider.is_remote_wal()); + assert!(before_catchup.version.memtables.is_empty()); + assert_eq!(before_catchup.last_entry_id, flushed_entry_id); + assert_eq!( + before_catchup.committed_sequence, + if open_before_edit { + flushed_sequence + } else { + edit_sequence + } + ); + assert_eq!( + scan_region(&follower_engine, region_id) + .await + .pretty_print() + .unwrap(), + flushed_rows + ); + + let resp = leader_engine + .set_region_role_state_gracefully(region_id, SettableRegionRoleState::Follower) + .await + .unwrap(); + let last_entry_id = get_last_entry_id(resp); + assert_eq!(last_entry_id, Some(unflushed_entry_id)); + assert_eq!(leader.version().flushed_entry_id, flushed_entry_id); + assert!(!leader.version().memtables.is_empty()); + follower_engine + .handle_request( + region_id, + RegionRequest::Catchup(RegionCatchupRequest { + set_writable: true, + entry_id: last_entry_id, + ..Default::default() + }), + ) + .await + .unwrap(); + let candidate = follower_engine.get_region(region_id).unwrap(); + assert!(candidate.is_writable()); + assert_eq!( + Arc::ptr_eq( + &candidate_before_catchup.version_control, + &candidate.version_control + ), + !open_before_edit + ); + let recovered = candidate.version_control.current(); + assert!(recovered.last_entry_id >= unflushed_entry_id); + assert_eq!(recovered.version.flushed_entry_id, flushed_entry_id); + assert!(recovered.committed_sequence >= edit_sequence); + assert!(recovered.last_sequence >= recovered.committed_sequence); + assert!(!recovered.version.memtables.is_empty()); + assert_eq!( + scan_region(&follower_engine, region_id) + .await + .pretty_print() + .unwrap(), + expected_rows + ); + + put_rows( + &follower_engine, + region_id, + Rows { + schema: column_schemas, + rows: build_rows(5, 6), + }, + ) + .await; + let current = candidate.version_control.current(); + assert_eq!(current.committed_sequence, recovered.last_sequence + 1); + assert!(current.committed_sequence > edit_sequence); + assert!(current.last_sequence >= current.committed_sequence); + let expected = "\ ++-------+---------+---------------------+ +| tag_0 | field_0 | ts | ++-------+---------+---------------------+ +| 0 | 0.0 | 1970-01-01T00:00:00 | +| 1 | 1.0 | 1970-01-01T00:00:01 | +| 2 | 2.0 | 1970-01-01T00:00:02 | +| 3 | 3.0 | 1970-01-01T00:00:03 | +| 4 | 4.0 | 1970-01-01T00:00:04 | +| 5 | 5.0 | 1970-01-01T00:00:05 | ++-------+---------+---------------------+"; + assert_eq!( + scan_region(&follower_engine, region_id) + .await + .pretty_print() + .unwrap(), + expected + ); +} + #[apply(single_kafka_log_store_factory)] async fn test_catchup_with_incorrect_last_entry_id(factory: Option) { use store_api::region_engine::SettableRegionRoleState; @@ -536,7 +752,10 @@ async fn scan_region(engine: &MitoEngine, region_id: RegionId) -> RecordBatches } #[apply(single_raft_engine_log_store_factory)] -async fn test_local_catchup(factory: Option) { +async fn test_local_catchup( + factory: Option, + #[values(false, true)] reopen_before_catchup: bool, +) { use store_api::region_engine::SettableRegionRoleState; use crate::test_util::LogStoreImpl; @@ -599,6 +818,10 @@ async fn test_local_catchup(factory: Option) { rows: build_rows(8, 9), }; put_rows(&leader_engine, region_id, rows).await; + let old_region = leader_engine.get_region(region_id).unwrap(); + if !reopen_before_catchup { + old_region.version_control.update_last_sequence(10); + } let resp = leader_engine .set_region_role_state_gracefully(region_id, SettableRegionRoleState::Follower) .await @@ -606,27 +829,27 @@ async fn test_local_catchup(factory: Option) { let last_entry_id = get_last_entry_id(resp); assert_eq!(last_entry_id.unwrap(), 3); - close_region(&leader_engine, region_id).await; - // Reopen the region, and skip the wal replay. - leader_engine - .handle_request( - region_id, - RegionRequest::Open(RegionOpenRequest { - engine: String::new(), - table_dir: table_dir.clone(), - options: HashMap::new(), - skip_wal_replay: true, - path_type: PathType::Bare, - checkpoint: None, - requirements: Default::default(), - }), - ) - .await - .unwrap(); + assert_eq!(old_region.version_control.committed_sequence(), 7); - // The last entry id should be 1. - let region = leader_engine.get_region(region_id).unwrap(); - assert_eq!(region.version_control.current().last_entry_id, 1); + if reopen_before_catchup { + close_region(&leader_engine, region_id).await; + open_region(&leader_engine, region_id, table_dir, true).await; + } + let region_before_catchup = leader_engine.get_region(region_id).unwrap(); + let current = region_before_catchup.version_control.current(); + if reopen_before_catchup { + assert!(current.version.memtables.is_empty()); + assert_eq!( + ( + current.last_sequence, + current.committed_sequence, + current.last_entry_id + ), + (3, 3, 1) + ); + } else { + assert!(!current.version.memtables.is_empty()); + } // There are 2 entries in the log store. let (start, end) = log_store.span(&RaftEngineProvider::new(region_id.into())); @@ -644,9 +867,24 @@ async fn test_local_catchup(factory: Option) { ) .await; assert!(resp.is_ok()); - // After catchup, the last entry id should be 3. let region = leader_engine.get_region(region_id).unwrap(); - assert_eq!(region.version_control.current().last_entry_id, 3); + assert_eq!( + Arc::ptr_eq( + ®ion_before_catchup.version_control, + ®ion.version_control + ), + reopen_before_catchup + ); + let expected_last_sequence = if reopen_before_catchup { 3 } else { 10 }; + let current = region.version_control.current(); + assert_eq!( + ( + current.last_sequence, + current.committed_sequence, + current.last_entry_id + ), + (expected_last_sequence, 3, 3) + ); // The log store has been obsoleted these 2 entries. let (start, end) = log_store.span(&RaftEngineProvider::new(region_id.into())); @@ -654,7 +892,7 @@ async fn test_local_catchup(factory: Option) { assert_eq!(end, None); // For local WAL, entries are not replayed during catchup. - // Therefore, any rows that were not flushed before closing the region will not be visible. + // Therefore, any rows that were not flushed before reopening the region will not be visible. let batches = scan_region(&leader_engine, region_id).await; let expected = "\ +-------+---------+---------------------+ @@ -672,6 +910,11 @@ async fn test_local_catchup(factory: Option) { rows: build_rows(4, 7), }; put_rows(&leader_engine, region_id, rows).await; + let current = region.version_control.current(); + assert_eq!( + (current.last_sequence, current.committed_sequence), + (expected_last_sequence + 3, expected_last_sequence + 3) + ); let batches = scan_region(&leader_engine, region_id).await; let expected = "\ diff --git a/src/mito2/src/engine/edit_region_test.rs b/src/mito2/src/engine/edit_region_test.rs index ace17905eb4..4ae44196e15 100644 --- a/src/mito2/src/engine/edit_region_test.rs +++ b/src/mito2/src/engine/edit_region_test.rs @@ -26,7 +26,7 @@ use common_time::util::current_time_millis; use datatypes::arrow::array::{ArrayRef, Float64Array, StringArray, TimestampMillisecondArray}; use datatypes::arrow::datatypes::{DataType, Field, Schema, TimeUnit}; use object_store::ObjectStore; -use store_api::region_engine::RegionEngine; +use store_api::region_engine::{RegionEngine, RegionManifestInfo, RegionRole}; use store_api::region_request::{ RegionBulkInsertsRequest, RegionCloseRequest, RegionPutRequest, RegionRequest, }; @@ -38,6 +38,7 @@ use crate::engine::MitoEngine; use crate::engine::listener::EventListener; use crate::manifest::action::RegionEdit; use crate::region::{MitoRegionRef, RegionLeaderState, RegionRoleState}; +use crate::request::{RegionEditRequest, WorkerRequest}; use crate::sst::file::FileMeta; use crate::test_util::{CreateRequestBuilder, TestEnv, build_rows, rows_schema}; use crate::time_provider::mock::MockTimeProvider; @@ -565,6 +566,71 @@ async fn test_region_edit_with_file_sequence_is_not_merged() { assert_eq!(Some(3), region_file_sequence(®ion, third_file_id)); } +#[tokio::test] +async fn test_failed_region_edit_preserves_sequence() { + let mut env = TestEnv::new().await; + let engine = env + .create_engine(MitoConfig { + schedule_compaction_after_edit: false, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + engine + .handle_request( + region_id, + RegionRequest::Create(CreateRequestBuilder::new().build()), + ) + .await + .unwrap(); + let region = engine.get_region(region_id).unwrap(); + region.version_control.update_last_sequence(2); + let manifest_guard = region.manifest_ctx.manifest_manager.write().await; + let (tx, rx) = oneshot::channel(); + engine + .inner + .workers + .submit_to_worker( + region_id, + WorkerRequest::EditRegion(RegionEditRequest::new( + region_id, + test_region_edit(region_id, FileId::random()), + false, + tx, + )), + ) + .await + .unwrap(); + // Sync fails before taking the manifest lock, after the edit starts. + assert!( + engine + .sync_region(region_id, RegionManifestInfo::mito(0, 0, 0).into()) + .await + .is_err() + ); + assert_eq!( + region.state(), + RegionRoleState::Leader(RegionLeaderState::Editing) + ); + assert_eq!(region.version_control.current().last_sequence, 3); + assert_eq!(region.version_control.committed_sequence(), 0); + region.set_role(RegionRole::Follower); + drop(manifest_guard); + assert!(rx.await.unwrap().is_err()); + assert_eq!(region.version_control.current().last_sequence, 3); + assert_eq!(region.version_control.committed_sequence(), 0); + + region.set_role(RegionRole::Leader); + let file_id = FileId::random(); + engine + .edit_region(region_id, test_region_edit(region_id, file_id)) + .await + .unwrap(); + assert_eq!(region_file_sequence(®ion, file_id), Some(4)); + let current = region.version_control.current(); + assert_eq!((current.last_sequence, current.committed_sequence), (4, 4)); +} + #[tokio::test] async fn test_region_edit_clears_preserve_row_sequence() { let mut env = TestEnv::new().await; diff --git a/src/mito2/src/engine/object_store_wal_recovery_test.rs b/src/mito2/src/engine/object_store_wal_recovery_test.rs index d5da18bf521..a9713d02fcc 100644 --- a/src/mito2/src/engine/object_store_wal_recovery_test.rs +++ b/src/mito2/src/engine/object_store_wal_recovery_test.rs @@ -1251,8 +1251,21 @@ async fn test_reopen_after_a_failed_create_shows_only_the_acknowledged_write() { assert_eq!(vec![0, 1], wal_object_seqs(&object_store).await); assert_eq!(0, latest(&store, REGION_A)); - // The same key and timestamp is written again under the same row - // sequence, and acknowledged from object 2. + assert_eq!( + 1, + region(&engine, REGION_A) + .version_control + .current() + .last_sequence + ); + assert_eq!( + 0, + region(&engine, REGION_A) + .version_control + .committed_sequence() + ); + + // Write the same key and timestamp with a higher sequence in object 2. let mut writer = SealedWriter { store: store.clone(), admitted: 1, @@ -1266,7 +1279,7 @@ async fn test_reopen_after_a_failed_create_shows_only_the_acknowledged_write() { assert_eq!(vec![1, 2], wal_object_seqs(&object_store).await); assert_eq!(entry_id(2, 1), latest(&store, REGION_A)); let region_a = region(&engine, REGION_A); - assert_eq!(1, region_a.version_control.committed_sequence()); + assert_eq!(2, region_a.version_control.committed_sequence()); let expected = "\ +-------+---------+---------------------+ | tag_0 | field_0 | ts | diff --git a/src/mito2/src/engine/sequence_test.rs b/src/mito2/src/engine/sequence_test.rs new file mode 100644 index 00000000000..f6c99ebb61b --- /dev/null +++ b/src/mito2/src/engine/sequence_test.rs @@ -0,0 +1,410 @@ +// 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. + +//! Mito sequence recovery tests using a memory WAL with store-assigned entry ids. +//! Uses ObjectStore providers; Kafka producers are not tested. + +use std::collections::HashMap; +use std::sync::{Arc, Mutex}; + +use api::v1::Rows; +use common_recordbatch::RecordBatches; +use common_wal::options::{ObjectStoreWalOptions, WAL_OPTIONS_KEY, WalOptions}; +use futures::stream; +use log_store::error::{Error, IllegalStateSnafu, Result}; +use store_api::logstore::entry::{Entry, NaiveEntry}; +use store_api::logstore::provider::Provider; +use store_api::logstore::{AppendBatchResponse, EntryId, LogStore, SendableEntryStream}; +use store_api::region_engine::{RegionEngine, RegionRole}; +use store_api::region_request::{PathType, RegionOpenRequest, RegionPutRequest, RegionRequest}; +use store_api::storage::{RegionId, ScanRequest}; +use tokio::sync::oneshot; + +use crate::config::MitoConfig; +use crate::engine::MitoEngine; +use crate::test_util::{ + CreateRequestBuilder, TestEnv, build_rows_for_key, flush_region, rows_schema, +}; + +#[derive(Debug)] +enum AppendFault { + Persist, + Partial(RegionId), + Deferred { + release: oneshot::Receiver<()>, + saved: oneshot::Sender<()>, + }, +} + +#[derive(Debug, Default)] +struct LogState { + entries: HashMap>, + fault: Option, + append_calls: usize, +} + +impl LogState { + fn save(&mut self, entries: Vec) -> AppendBatchResponse { + let mut last_entry_ids = HashMap::new(); + for mut entry in entries { + let log = self.entries.entry(entry.provider().clone()).or_default(); + // Entry ids start from 1 (zero means no data). + let offset = log.len() as u64 + 1; + entry.set_entry_id(offset); + last_entry_ids.insert(entry.region_id(), offset); + log.push(entry); + } + AppendBatchResponse { last_entry_ids } + } +} + +#[derive(Debug, Default)] +struct MemoryLogStore { + state: Arc>, +} + +impl MemoryLogStore { + fn fail_next_append(&self, fault: AppendFault) { + let mut state = self.state.lock().unwrap(); + assert!(state.fault.replace(fault).is_none()); + } +} + +#[async_trait::async_trait] +impl LogStore for MemoryLogStore { + type Error = Error; + + async fn stop(&self) -> Result<()> { + Ok(()) + } + + async fn append_batch(&self, entries: Vec) -> Result { + let mut state = self.state.lock().unwrap(); + state.append_calls += 1; + match state.fault.take() { + None => return Ok(state.save(entries)), + Some(AppendFault::Persist) => { + state.save(entries); + } + Some(AppendFault::Partial(region_id)) => { + state.save( + entries + .into_iter() + .filter(|e| e.region_id() == region_id) + .collect(), + ); + } + Some(AppendFault::Deferred { release, saved }) => { + let state = self.state.clone(); + common_runtime::spawn_global(async move { + if release.await.is_ok() { + state.lock().unwrap().save(entries); + let _ = saved.send(()); + } + }); + } + } + IllegalStateSnafu {}.fail() + } + + async fn read( + &self, + provider: &Provider, + id: EntryId, + ) -> Result> { + // Snapshot the log for recovery. + let entries = self + .state + .lock() + .unwrap() + .entries + .get(provider) + .cloned() + .unwrap_or_default(); + Ok(Box::pin(stream::iter([Ok(entries + .into_iter() + .filter(|e| e.entry_id() >= id) + .collect())]))) + } + + async fn create_namespace(&self, _provider: &Provider) -> Result<()> { + Ok(()) + } + async fn delete_namespace(&self, _provider: &Provider) -> Result<()> { + Ok(()) + } + async fn list_namespaces(&self) -> Result> { + Ok(vec![]) + } + async fn obsolete( + &self, + _provider: &Provider, + _region_id: RegionId, + _id: EntryId, + ) -> Result<()> { + Ok(()) + } + async fn obsolete_all(&self, _provider: &Provider, _region_id: RegionId) -> Result<()> { + Ok(()) + } + + fn entry( + &self, + data: Vec, + entry_id: EntryId, + region_id: RegionId, + provider: &Provider, + ) -> Result { + Ok(Entry::Naive(NaiveEntry { + data, + entry_id, + region_id, + provider: provider.clone(), + })) + } + + fn latest_entry_id(&self, provider: &Provider) -> Result { + Ok(self + .state + .lock() + .unwrap() + .entries + .get(provider) + .and_then(|log| log.last()) + .map_or(0, Entry::entry_id)) + } +} + +fn wal_options() -> HashMap { + HashMap::from([( + WAL_OPTIONS_KEY.to_string(), + serde_json::to_string(&WalOptions::ObjectStore(ObjectStoreWalOptions::new( + "sequence-test/".to_string(), + ))) + .unwrap(), + )]) +} + +async fn create_region( + engine: &MitoEngine, + region_id: RegionId, +) -> (String, Vec) { + let request = CreateRequestBuilder::new() + .insert_option(WAL_OPTIONS_KEY, &wal_options()[WAL_OPTIONS_KEY]) + .build(); + let table_dir = request.table_dir.clone(); + let schema = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + (table_dir, schema) +} + +async fn open_region(engine: &MitoEngine, region_id: RegionId, table_dir: String) { + engine + .handle_request( + region_id, + RegionRequest::Open(RegionOpenRequest { + engine: String::new(), + table_dir, + path_type: PathType::Bare, + options: wal_options(), + skip_wal_replay: false, + checkpoint: None, + requirements: Default::default(), + }), + ) + .await + .unwrap(); + engine + .set_region_role(region_id, RegionRole::Leader) + .unwrap(); +} + +fn put_request(schema: &[api::v1::ColumnSchema], keys: &[(&str, usize)]) -> RegionPutRequest { + RegionPutRequest { + rows: Rows { + schema: schema.to_vec(), + rows: keys + .iter() + .flat_map(|(key, value)| build_rows_for_key(key, 0, 1, *value)) + .collect(), + }, + skip_wal: false, + hint: None, + partition_expr_version: None, + } +} + +async fn scan_rows(engine: &MitoEngine, region_id: RegionId) -> String { + let stream = engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(); + RecordBatches::try_collect(stream) + .await + .unwrap() + .pretty_print() + .unwrap() +} + +#[tokio::test] +async fn test_persisted_wal_error_does_not_overwrite_success() { + for flat_format in [false, true] { + for (deferred, flush) in [(false, false), (true, false), (true, true)] { + let mut env = TestEnv::new().await; + let store = Arc::new(MemoryLogStore::default()); + let config = MitoConfig { + default_flat_format: flat_format, + num_workers: 1, + ..Default::default() + }; + let engine = env + .create_engine_with_log_store(config.clone(), store.clone()) + .await; + let region_id = RegionId::new(1, 1); + let (table_dir, schema) = create_region(&engine, region_id).await; + let (release, released) = oneshot::channel(); + let (saved, saved_rx) = oneshot::channel(); + store.fail_next_append(if deferred { + AppendFault::Deferred { + release: released, + saved, + } + } else { + AppendFault::Persist + }); + assert!( + engine + .handle_request( + region_id, + RegionRequest::Put(put_request(&schema, &[("a", 1), ("b", 1)])) + ) + .await + .is_err() + ); + let failed = engine + .get_region(region_id) + .unwrap() + .version_control + .current(); + assert_eq!(failed.committed_sequence, 0); + assert_eq!(failed.last_sequence, 2); + assert_eq!(failed.last_entry_id, 0); + assert_eq!(failed.version.flushed_sequence, 0); + engine + .handle_request( + region_id, + RegionRequest::Put(put_request(&schema, &[("b", 2)])), + ) + .await + .unwrap(); + if flush { + flush_region(&engine, region_id, None).await; + } + if deferred { + release.send(()).unwrap(); + saved_rx.await.unwrap(); + } + assert!( + scan_rows(&engine, region_id) + .await + .contains("| b | 2.0") + ); + engine.stop().await.unwrap(); + drop(engine); + let engine = env.create_engine_with_log_store(config, store).await; + open_region(&engine, region_id, table_dir).await; + let recovered = scan_rows(&engine, region_id).await; + assert!( + recovered.contains("| b | 2.0"), + "deferred={deferred}, flush={flush}, flat={flat_format}: {recovered}" + ); + let current = engine + .get_region(region_id) + .unwrap() + .version_control + .current(); + assert_eq!(current.committed_sequence, 3); + assert_eq!(current.last_sequence, 3); + engine.stop().await.unwrap(); + } + } +} + +#[tokio::test] +async fn test_partial_region_append_preserves_all_allocations() { + let mut env = TestEnv::new().await; + let store = Arc::new(MemoryLogStore::default()); + let engine = env + .create_engine_with_log_store( + MitoConfig { + num_workers: 1, + ..Default::default() + }, + store.clone(), + ) + .await; + let regions = [RegionId::new(1, 1), RegionId::new(1, 2)]; + let mut schemas = Vec::new(); + let mut table_dirs = Vec::new(); + for id in regions { + let (dir, schema) = create_region(&engine, id).await; + table_dirs.push(dir); + schemas.push(schema); + } + store.fail_next_append(AppendFault::Partial(regions[0])); + let requests = regions + .into_iter() + .zip(&schemas) + .map(|(id, schema)| (id, put_request(schema, &[("a", 1), ("b", 1)]))) + .collect(); + for (_, response) in Arc::new(engine.clone()).handle_batch_put_requests(requests) { + assert!(response.await.is_err()); + } + assert_eq!(store.state.lock().unwrap().append_calls, 1); + for (id, schema) in regions.into_iter().zip(&schemas) { + let region = engine.get_region(id).unwrap(); + let current = region.version_control.current(); + assert_eq!( + ( + current.last_sequence, + current.committed_sequence, + current.last_entry_id + ), + (2, 0, 0) + ); + assert!(current.version.memtables.is_empty()); + assert_eq!( + store.latest_entry_id(®ion.provider).unwrap(), + u64::from(id == regions[0]) + ); + engine + .handle_request(id, RegionRequest::Put(put_request(schema, &[("b", 2)]))) + .await + .unwrap(); + assert_eq!(region.version_control.current().committed_sequence, 3); + } + engine.stop().await.unwrap(); + drop(engine); + let engine = env + .create_engine_with_log_store(MitoConfig::default(), store) + .await; + for (id, dir) in regions.into_iter().zip(table_dirs) { + open_region(&engine, id, dir).await; + assert!(scan_rows(&engine, id).await.contains("| b | 2.0")); + } + engine.stop().await.unwrap(); +} diff --git a/src/mito2/src/engine/staging_test.rs b/src/mito2/src/engine/staging_test.rs index 98973de4f30..c29c4bbe3a3 100644 --- a/src/mito2/src/engine/staging_test.rs +++ b/src/mito2/src/engine/staging_test.rs @@ -759,6 +759,39 @@ async fn test_staging_exit_success_with_manifests_with_format(flat_format: bool) .await .unwrap(); + let region = engine.get_region(region_id).unwrap(); + let file = { + let manager = region.manifest_ctx.manifest_manager.read().await; + manager + .staging_manifest() + .unwrap() + .files + .values() + .next() + .unwrap() + .clone() + }; + for expected in [11, 12] { + engine + .edit_region( + region_id, + RegionEdit { + files_to_add: vec![file.clone()], + files_to_remove: vec![], + timestamp_ms: None, + compaction_time_window: None, + flushed_entry_id: None, + flushed_sequence: None, + committed_sequence: None, + }, + ) + .await + .unwrap(); + let current = region.version_control.current(); + assert_eq!(current.last_sequence, expected); + assert_eq!(current.committed_sequence, 10); + } + // Verify we're in staging mode and staging manifests exist let data_home = env.data_home(); let region_dir = format!("{}/data/test/1024_0000000000", data_home.display()); @@ -770,13 +803,11 @@ async fn test_staging_exit_success_with_manifests_with_format(flat_format: bool) .unwrap(); assert_eq!( staging_files_before.len(), - // Two files for flush operation - // One file for entering staging mode - 3, - "Staging manifest directory should contain 3 files before exit, got: {:?}", + // Two flushes, two edits and entering staging mode. + 5, + "Staging manifest directory should contain 5 files before exit, got: {:?}", staging_files_before ); - let region = engine.get_region(region_id).unwrap(); { let manager = region.manifest_ctx.manifest_manager.read().await; let staging_manifest = manager.staging_manifest().unwrap(); @@ -950,6 +981,11 @@ async fn test_staging_exit_success_with_manifests_with_format(flat_format: bool) // Verify we're back in normal mode let workers = &engine.inner.workers; let region = workers.get_region(region_id).unwrap(); + let current = region.version_control.current(); + assert_eq!( + (current.last_sequence, current.committed_sequence), + (12, 12) + ); assert!( !region.is_staging(), "Region should no longer be in staging mode" diff --git a/src/mito2/src/engine/sync_test.rs b/src/mito2/src/engine/sync_test.rs index d75a9ef0313..ee0cd7ebe64 100644 --- a/src/mito2/src/engine/sync_test.rs +++ b/src/mito2/src/engine/sync_test.rs @@ -20,10 +20,10 @@ use common_recordbatch::RecordBatches; use datatypes::prelude::ConcreteDataType; use datatypes::schema::ColumnSchema; use store_api::metadata::ColumnMetadata; -use store_api::region_engine::{RegionEngine, RegionManifestInfo}; +use store_api::region_engine::{RegionEngine, RegionManifestInfo, RegionRole}; use store_api::region_request::{ AddColumn, AddColumnLocation, AlterKind, PathType, RegionAlterRequest, RegionOpenRequest, - RegionRequest, + RegionPutRequest, RegionRequest, }; use store_api::storage::{RegionId, ScanRequest}; @@ -170,6 +170,11 @@ async fn test_sync_after_flush_region_with_format(flat_format: bool) { .await .unwrap(); common_telemetry::info!("Scan the region on the follower engine after sync"); + let follower = follower_engine.get_region(region_id).unwrap(); + assert_eq!(follower.version_control.current().last_sequence, 3); + // Sync does not update committed sequence or entry id. + assert_eq!(follower.version_control.committed_sequence(), 0); + assert_eq!(follower.version_control.current().last_entry_id, 0); // Scan the region on the follower engine let expected = "\ +-------+---------+---------------------+ @@ -180,6 +185,80 @@ async fn test_sync_after_flush_region_with_format(flat_format: bool) { | 2 | 2.0 | 1970-01-01T00:00:02 | +-------+---------+---------------------+"; scan_check(&follower_engine, region_id, expected, 0, 1).await; + + // An edit may have a committed sequence above the flushed sequence. + let file = engine + .get_region(region_id) + .unwrap() + .version() + .ssts + .levels()[0] + .files + .values() + .next() + .unwrap() + .meta_ref() + .clone(); + engine + .edit_region( + region_id, + crate::manifest::action::RegionEdit { + files_to_add: vec![file], + files_to_remove: vec![], + timestamp_ms: None, + compaction_time_window: None, + flushed_entry_id: None, + flushed_sequence: None, + committed_sequence: None, + }, + ) + .await + .unwrap(); + let manifest_version = engine + .get_region(region_id) + .unwrap() + .manifest_ctx + .manifest_version() + .await; + follower_engine + .sync_region( + region_id, + RegionManifestInfo::mito(manifest_version, 0, 0).into(), + ) + .await + .unwrap(); + assert_eq!(follower.version_control.current().last_sequence, 4); + assert_eq!(follower.version_control.committed_sequence(), 0); + + // Writes after promotion must use a higher sequence. + follower_engine + .set_region_role(region_id, RegionRole::Leader) + .unwrap(); + let request = CreateRequestBuilder::new().build(); + // Skip WAL to isolate row sequence allocation. + follower_engine + .handle_request( + region_id, + RegionRequest::Put(RegionPutRequest { + rows: Rows { + schema: rows_schema(&request), + rows: build_rows(0, 1), + }, + skip_wal: true, + hint: None, + partition_expr_version: None, + }), + ) + .await + .unwrap(); + assert_eq!( + 5, + follower_engine + .get_region(region_id) + .unwrap() + .version_control + .committed_sequence() + ); } #[tokio::test] diff --git a/src/mito2/src/region/version.rs b/src/mito2/src/region/version.rs index 1c595fb0b9f..c1940b31ab0 100644 --- a/src/mito2/src/region/version.rs +++ b/src/mito2/src/region/version.rs @@ -59,6 +59,7 @@ impl VersionControl { VersionControl { data: RwLock::new(VersionControlData { version: Arc::new(version), + last_sequence: flushed_sequence, committed_sequence: flushed_sequence, last_entry_id: flushed_entry_id, is_dropped: false, @@ -76,16 +77,24 @@ impl VersionControl { self.data.read().unwrap().clone() } + /// Updates last sequence number to the maximum. + pub(crate) fn update_last_sequence(&self, seq: SequenceNumber) { + let mut data = self.data.write().unwrap(); + data.last_sequence = data.last_sequence.max(seq); + } + /// Updates the `committed_sequence` of version. pub(crate) fn set_committed_sequence(&self, seq: SequenceNumber) { let mut data = self.data.write().unwrap(); - data.committed_sequence = seq; + data.committed_sequence = data.committed_sequence.max(seq); + data.last_sequence = data.last_sequence.max(data.committed_sequence); } /// Updates committed sequence and entry id. pub(crate) fn set_sequence_and_entry_id(&self, seq: SequenceNumber, entry_id: EntryId) { let mut data = self.data.write().unwrap(); - data.committed_sequence = seq; + data.committed_sequence = data.committed_sequence.max(seq); + data.last_sequence = data.last_sequence.max(data.committed_sequence); data.last_entry_id = entry_id; } @@ -162,6 +171,9 @@ impl VersionControl { } else { version_data.committed_sequence }; + version_data.last_sequence = version_data + .last_sequence + .max(version_data.committed_sequence); version_data.version = new_version; } @@ -312,6 +324,7 @@ impl VersionControl { /// Overwrites the current version with a new version. pub(crate) fn overwrite_current(&self, version: VersionRef) { let mut version_data = self.data.write().unwrap(); + version_data.last_sequence = version_data.last_sequence.max(version.flushed_sequence); version_data.version = version; } @@ -336,6 +349,8 @@ pub(crate) type VersionControlRef = Arc; pub(crate) struct VersionControlData { /// Latest version. pub(crate) version: VersionRef, + /// Last allocated sequence number. + pub(crate) last_sequence: SequenceNumber, /// Sequence number of last committed data. /// /// Starts from 1 (zero means no data). @@ -561,3 +576,60 @@ impl VersionBuilder { } } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_util::version_util::VersionControlBuilder; + + #[test] + fn test_sequence_watermarks_are_monotonic() { + let builder = VersionControlBuilder::new(); + let mut version = builder.build_version(); + version.flushed_sequence = 8; + version.flushed_entry_id = 4; + let vc = VersionControl::new(version); + assert_eq!( + (vc.current().last_sequence, vc.committed_sequence()), + (8, 8) + ); + + vc.update_last_sequence(20); + vc.update_last_sequence(9); + vc.set_committed_sequence(10); + vc.set_committed_sequence(9); + assert_eq!( + (vc.current().last_sequence, vc.committed_sequence()), + (20, 10) + ); + assert_eq!(vc.current().last_entry_id, 4); + vc.set_sequence_and_entry_id(12, 5); + vc.set_sequence_and_entry_id(4, 6); + assert_eq!( + (vc.current().last_sequence, vc.committed_sequence()), + (20, 12) + ); + assert_eq!(vc.current().version.flushed_sequence, 8); + + for sequence in [25, 11] { + vc.apply_edit( + Some(RegionEdit { + files_to_add: vec![], + files_to_remove: vec![], + timestamp_ms: None, + compaction_time_window: None, + flushed_entry_id: None, + flushed_sequence: None, + committed_sequence: Some(sequence), + }), + &[], + builder.file_purger(), + ); + assert_eq!( + (vc.current().last_sequence, vc.committed_sequence()), + (25, 25) + ); + assert_eq!(vc.current().last_entry_id, 6); + } + } +} diff --git a/src/mito2/src/region_write_ctx.rs b/src/mito2/src/region_write_ctx.rs index e4e1e505326..d37c42545a0 100644 --- a/src/mito2/src/region_write_ctx.rs +++ b/src/mito2/src/region_write_ctx.rs @@ -100,9 +100,9 @@ pub(crate) struct RegionWriteCtx { /// VersionControl of the region. version_control: VersionControlRef, /// Next sequence number to write. - /// - /// The context assigns a unique sequence number for each row. next_sequence: SequenceNumber, + /// Maximum sequence number allocated by this context. + max_sequence: Option, /// Next entry id of WAL to write. next_entry_id: EntryId, /// Valid WAL entry to write. @@ -148,7 +148,7 @@ impl RegionWriteCtx { ) -> RegionWriteCtx { let VersionControlData { version, - committed_sequence, + last_sequence, last_entry_id, .. } = version_control.current(); @@ -159,7 +159,8 @@ impl RegionWriteCtx { region_id, version, version_control: version_control.clone(), - next_sequence: committed_sequence + 1, + next_sequence: last_sequence + 1, + max_sequence: None, next_entry_id: last_entry_id + 1, wal_entry: WalEntry::default(), memtable_mutations: Vec::new(), @@ -208,8 +209,7 @@ impl RegionWriteCtx { self.wal_notifiers.push(notify); } - // Increase sequence number. - self.next_sequence += num_rows as u64; + self.advance_sequence(num_rows); // Update metrics. match OpType::try_from(op_type) { @@ -224,6 +224,8 @@ impl RegionWriteCtx { &mut self, wal_writer: &mut WalWriter, ) -> Result<()> { + // Reserve sequences before WAL writes to avoid reuse on failure. + self.reserve_sequences(); wal_writer.add_entry( self.region_id, self.next_entry_id, @@ -294,6 +296,10 @@ impl RegionWriteCtx { if self.failed { return; } + // Reserve sequences for WAL replay and skip-WAL writes. + if self.source == WriteSource::WalReplay || self.skip_wal() { + self.reserve_sequences(); + } // A failed installation may leave some rows in memory. if self.has_skip_wal_writes { self.version.memtables.mutable.mark_skip_wal_writes(); @@ -395,7 +401,7 @@ impl RegionWriteCtx { self.bulk_notifiers .push(WriteNotify::new(sender, bulk.num_rows())); - self.next_sequence += bulk.num_rows() as u64; + self.advance_sequence(bulk.num_rows()); self.bulk_parts.push(bulk); true } @@ -453,6 +459,21 @@ impl RegionWriteCtx { } } + /// Advances the sequence number and records the maximum. + fn advance_sequence(&mut self, num_rows: usize) { + self.next_sequence += num_rows as u64; + if num_rows > 0 { + let sequence = self.next_sequence - 1; + self.max_sequence = Some(self.max_sequence.map_or(sequence, |max| max.max(sequence))); + } + } + + fn reserve_sequences(&self) { + if let Some(sequence) = self.max_sequence { + self.version_control.update_last_sequence(sequence); + } + } + /// Publishes the assigned sequences and entry id to the region's committed /// watermark. Only call after both [`write_mutations`](Self::write_mutations) /// and [`write_bulk`](Self::write_bulk) have completed; a failed context @@ -461,8 +482,13 @@ impl RegionWriteCtx { if self.failed { return; } - self.version_control - .set_sequence_and_entry_id(self.next_sequence - 1, self.next_entry_id - 1); + if let Some(sequence) = self.max_sequence { + self.version_control + .set_sequence_and_entry_id(sequence, self.next_entry_id - 1); + } else { + // Empty entries only update the entry id. + self.version_control.set_entry_id(self.next_entry_id - 1); + } } } @@ -1014,4 +1040,160 @@ mod tests { raw_data: None, } } + + fn sequence_rows(num_rows: usize) -> Rows { + use api::v1::helper::{tag_column_schema, time_index_column_schema}; + Rows { + schema: vec![ + time_index_column_schema("ts", api::v1::ColumnDataType::TimestampMillisecond), + tag_column_schema("tag_0", api::v1::ColumnDataType::String), + ], + rows: (0..num_rows) + .map(|i| api::v1::Row { + values: vec![ + crate::test_util::ts_ms_value(i as i64), + api::v1::Value { + value_data: Some(api::v1::value::ValueData::StringValue( + "a".to_string(), + )), + }, + ], + }) + .collect(), + } + } + + #[tokio::test] + async fn test_replay_sequence_maximum_below_allocation_floor() { + let builder = VersionControlBuilder::new(); + let region_id = builder.region_id(); + let vc = Arc::new(builder.build()); + vc.update_last_sequence(100); + let provider = Provider::raft_engine_provider(region_id.as_u64()); + let mut ctx = RegionWriteCtx::new( + region_id, + &vc, + provider.clone(), + None, + WriteSource::WalReplay, + ); + for (sequence, num_rows) in [(40, 2), (3, 1)] { + ctx.push_mutation( + OpType::Put as i32, + Some(sequence_rows(num_rows)), + None, + OptionOutputTx::none(), + Some(sequence), + false, + ); + } + for sequence in [10, 2] { + let mut part = new_bulk_part(); + part.fill_missing_columns(&ctx.version.metadata).unwrap(); + assert!(ctx.push_bulk(OptionOutputTx::none(), part, Some(sequence), false)); + } + assert_eq!(ctx.wal_entry.mutations[0].sequence, 40); + assert_eq!(ctx.bulk_parts[0].sequence, 10); + assert_eq!(ctx.bulk_parts[1].sequence, 2); + assert_eq!(vc.committed_sequence(), 0); + ctx.set_next_entry_id(8); + ctx.write_memtables().await; + assert_eq!(vc.committed_sequence(), 41); + assert_eq!(vc.current().last_sequence, 100); + assert_eq!(vc.current().last_entry_id, 7); + + let mut ctx = RegionWriteCtx::new(region_id, &vc, provider, None, WriteSource::Request); + ctx.push_mutation( + OpType::Put as i32, + Some(sequence_rows(2)), + None, + OptionOutputTx::none(), + None, + true, + ); + assert_eq!(ctx.memtable_mutations[0].0.sequence, 101); + ctx.write_memtables().await; + assert_eq!(vc.committed_sequence(), 102); + assert_eq!(vc.current().last_sequence, 102); + assert_eq!(vc.current().last_entry_id, 7); + } + + #[tokio::test] + async fn test_empty_context_does_not_publish_allocation_floor() { + let builder = VersionControlBuilder::new(); + let region_id = builder.region_id(); + let vc = Arc::new(builder.build()); + vc.update_last_sequence(100); + for empty_objects in [false, true] { + let mut ctx = + RegionWriteCtx::new(region_id, &vc, Provider::Noop, None, WriteSource::WalReplay); + if empty_objects { + ctx.push_mutation( + OpType::Put as i32, + Some(sequence_rows(0)), + None, + OptionOutputTx::none(), + Some(500), + false, + ); + ctx.push_mutation( + OpType::Put as i32, + None, + None, + OptionOutputTx::none(), + Some(600), + false, + ); + let mut part = new_bulk_part(); + part.fill_missing_columns(&ctx.version.metadata).unwrap(); + part.batch = part.batch.slice(0, 0); + assert!(ctx.push_bulk(OptionOutputTx::none(), part, Some(700), true)); + } + ctx.set_next_entry_id(4); + ctx.write_memtables().await; + let current = vc.current(); + assert_eq!(current.last_sequence, 100); + assert_eq!(current.committed_sequence, 0); + assert_eq!(current.last_entry_id, 3); + } + } + + #[tokio::test] + async fn test_rejected_bulk_does_not_publish_sequence() { + use datatypes::arrow::array::StringDictionaryBuilder; + use datatypes::arrow::datatypes::Int32Type; + + let builder = VersionControlBuilder::new(); + let vc = Arc::new(builder.build()); + vc.update_last_sequence(100); + let mut ctx = RegionWriteCtx::new( + builder.region_id(), + &vc, + Provider::Noop, + None, + WriteSource::Request, + ); + let mut part = new_bulk_part(); + let mut dictionary = StringDictionaryBuilder::::new(); + dictionary.append("a").unwrap(); + dictionary.append("b").unwrap(); + let dictionary = Arc::new(dictionary.finish()) as ArrayRef; + let mut fields = part.batch.schema().fields().to_vec(); + fields.push(Arc::new(Field::new( + "tag_0", + dictionary.data_type().clone(), + true, + ))); + part.batch = DfRecordBatch::try_new( + Arc::new(Schema::new(fields)), + vec![part.batch.column(0).clone(), dictionary], + ) + .unwrap(); + let (tx, rx) = oneshot::channel(); + assert!(!ctx.push_bulk(OptionOutputTx::from(tx), part, Some(500), false)); + assert!(rx.await.unwrap().is_err()); + ctx.write_memtables().await; + assert_eq!(vc.committed_sequence(), 0); + assert_eq!(vc.current().last_sequence, 100); + } } diff --git a/src/mito2/src/worker/handle_catchup.rs b/src/mito2/src/worker/handle_catchup.rs index 92874672789..65f20c5837a 100644 --- a/src/mito2/src/worker/handle_catchup.rs +++ b/src/mito2/src/worker/handle_catchup.rs @@ -153,6 +153,10 @@ impl RegionWorkerLoop { .open(&self.config, &self.wal) .await?; debug_assert!(!reopened_region.is_writable()); + // Preserve allocated sequences without publishing the old memtables. + reopened_region + .version_control + .update_last_sequence(region.version_control.current().last_sequence); self.register_region(reopened_region.clone()); if let Some(state) = &self.series_index_task_state { state.wake(); diff --git a/src/mito2/src/worker/handle_manifest.rs b/src/mito2/src/worker/handle_manifest.rs index 373080aee90..9510b484d84 100644 --- a/src/mito2/src/worker/handle_manifest.rs +++ b/src/mito2/src/worker/handle_manifest.rs @@ -274,6 +274,10 @@ impl RegionWorkerLoop { ); let version = version_builder.build(); region.version_control.overwrite_current(Arc::new(version)); + // Restore last sequence from the manifest. + if let Some(sequence) = manifest.committed_sequence { + region.version_control.update_last_sequence(sequence); + } let updated = manifest.manifest_version > original_manifest_version; let _ = sender.send(Ok((manifest.manifest_version, updated))); @@ -311,7 +315,7 @@ impl RegionWorkerLoop { waiters, preload_sst_cache, } = request; - let file_sequence = region.version_control.committed_sequence() + 1; + let file_sequence = region.version_control.current().last_sequence + 1; edit.committed_sequence = Some(file_sequence); // Generic region edits (direct/import/staging) assign a new destination @@ -337,6 +341,7 @@ impl RegionWorkerLoop { waiters.reply_with(|| Err(e.clone()).context(EditRegionSnafu { region_id })); return; } + region.version_control.update_last_sequence(file_sequence); let request_sender = self.sender.clone(); let cache_manager = self.cache_manager.clone(); diff --git a/src/mito2/src/worker/handle_write.rs b/src/mito2/src/worker/handle_write.rs index 7e290a768fb..6b8e5be4276 100644 --- a/src/mito2/src/worker/handle_write.rs +++ b/src/mito2/src/worker/handle_write.rs @@ -1045,6 +1045,13 @@ mod tests { assert!(region_ctxs[&failing_region].is_failed()); assert!(!region_ctxs[&ok_region].is_failed()); assert_eq!(entry_id + 1, region_ctxs[&ok_region].next_entry_id()); + assert_eq!( + region_ctxs[&failing_region] + .version_control() + .current() + .last_sequence, + 1 + ); for region_ctx in region_ctxs.values_mut() { region_ctx.write_memtables().await; @@ -1077,6 +1084,7 @@ mod tests { async fn check_bulk_write_sequence_not_committed_before_install_worker_level(skip_wal: bool) { let region_id = RegionId::new(1, 1); let version_control = Arc::new(VersionControlBuilder::new().build()); + version_control.update_last_sequence(2); let mut region_ctxs = HashMap::new(); let mut ctx = RegionWriteCtx::new( @@ -1111,6 +1119,9 @@ mod tests { .await .expect("bulk write never reached the install barrier"); + assert_eq!(version_control.current().last_sequence, 5); + assert_eq!(version_control.current().last_entry_id, 0); + assert_eq!( 0, version_control.committed_sequence(), @@ -1120,11 +1131,14 @@ mod tests { barrier.release(); write_handle.await.expect("bulk write should complete"); assert_eq!( - 3, + 5, version_control.committed_sequence(), "committed sequence must cover the installed bulk rows" ); + let current = version_control.current(); + assert_eq!(current.last_sequence, 5); + assert_eq!(current.last_entry_id, u64::from(!skip_wal)); assert_eq!(3, rx.await.unwrap().unwrap()); } @@ -1159,6 +1173,7 @@ mod tests { let mut region_ctxs = HashMap::new(); let (ctx, rx) = new_region_ctx(region_id, false); + let version_control = ctx.version_control().clone(); region_ctxs.insert(region_id, ctx); assert!(!write_wal(&wal, &mut region_ctxs).await); @@ -1166,5 +1181,11 @@ mod tests { // All contexts are consumed and waiters are notified with the error. assert!(region_ctxs.is_empty()); assert!(rx.await.unwrap().is_err()); + let current = version_control.current(); + assert_eq!(current.last_sequence, 1); + assert_eq!(current.committed_sequence, 0); + assert_eq!(current.last_entry_id, 0); + assert_eq!(current.version.flushed_sequence, 0); + assert_eq!(current.version.flushed_entry_id, 0); } } From cdc0d3a055ecf71ba32278a2951881f88de1a1dc Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Wed, 28 Oct 2026 07:26:22 +0000 Subject: [PATCH 119/239] feat(procedure): support detached subprocedure recovery (#10009) Select Discard or Detach when submitting children via Status::Suspended. Persist the policy from the child's first normal or rollback checkpoint. Recover detached children as independent roots with their original id and context, using existing execution and cleanup without rebuilding the graph. Keep fresh root admission serialized while allowing accepted procedures to recover above the submission limit. Cover first persistence, repeated recovery, rollback, cleanup, defaults, and concurrent admission. Signed-off-by: WenyXu --- src/common/procedure/src/lib.rs | 1 + src/common/procedure/src/local.rs | 179 ++++--- src/common/procedure/src/local/recovery.rs | 582 +++++++++++++++++++++ src/common/procedure/src/local/runner.rs | 49 +- src/common/procedure/src/procedure.rs | 33 ++ src/common/procedure/src/store.rs | 35 +- 6 files changed, 793 insertions(+), 86 deletions(-) create mode 100644 src/common/procedure/src/local/recovery.rs diff --git a/src/common/procedure/src/lib.rs b/src/common/procedure/src/lib.rs index a21e5d1c47a..419fb41ca30 100644 --- a/src/common/procedure/src/lib.rs +++ b/src/common/procedure/src/lib.rs @@ -33,5 +33,6 @@ pub use crate::procedure::{ ContextProviderRef, EventContext, EventTrigger, LockKey, Output, ParseIdError, PoisonKey, PoisonKeys, Procedure, ProcedureContext, ProcedureId, ProcedureInfo, ProcedureManager, ProcedureManagerRef, ProcedureState, ProcedureWithId, RetryPhase, Status, StringKey, + SubprocedureRecovery, }; pub use crate::watcher::Watcher; diff --git a/src/common/procedure/src/local.rs b/src/common/procedure/src/local.rs index 2ae7f5e3e08..8801028b198 100644 --- a/src/common/procedure/src/local.rs +++ b/src/common/procedure/src/local.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +mod recovery; mod runner; use std::collections::hash_map::Entry; @@ -32,9 +33,9 @@ use tokio::sync::watch::{self, Receiver, Sender}; use tokio::sync::{Mutex as TokioMutex, Notify}; use crate::error::{ - self, CheckStatusSnafu, DuplicateProcedureSnafu, Error, LoaderConflictSnafu, - ManagerNotStartSnafu, ManagerPasuedSnafu, PoisonKeyNotDefinedSnafu, ProcedureNotFoundSnafu, - Result, StartRemoveOutdatedMetaTaskSnafu, StopRemoveOutdatedMetaTaskSnafu, + CheckStatusSnafu, DuplicateProcedureSnafu, Error, LoaderConflictSnafu, ManagerNotStartSnafu, + ManagerPasuedSnafu, PoisonKeyNotDefinedSnafu, ProcedureNotFoundSnafu, Result, + StartRemoveOutdatedMetaTaskSnafu, StopRemoveOutdatedMetaTaskSnafu, TooManyRunningProceduresSnafu, }; use crate::local::runner::Runner; @@ -44,7 +45,8 @@ use crate::store::poison_store::PoisonStoreRef; use crate::store::{ProcedureMessage, ProcedureMessages, ProcedureStore, StateStoreRef}; use crate::{ BoxedProcedure, ContextProvider, EventTrigger, LockKey, PoisonKey, ProcedureContext, - ProcedureId, ProcedureManager, ProcedureState, ProcedureWithId, StringKey, Watcher, + ProcedureId, ProcedureManager, ProcedureState, ProcedureWithId, StringKey, + SubprocedureRecovery, Watcher, }; /// The expired time of a procedure's metadata. @@ -81,6 +83,8 @@ pub(crate) struct ProcedureMeta { type_name: String, /// Parent procedure id. parent_id: Option, + /// Recovery policy captured when the procedure is submitted as a child. + subprocedure_recovery: SubprocedureRecovery, /// Context associated with the root submission. context: ProcedureContext, /// Notify to wait for subprocedures. @@ -116,6 +120,7 @@ impl ProcedureMeta { ProcedureMeta { id, parent_id, + subprocedure_recovery: SubprocedureRecovery::Discard, context, child_notify: Notify::new(), lock_key, @@ -129,6 +134,12 @@ impl ProcedureMeta { } } + /// Sets the recovery policy for this subprocedure. + fn with_subprocedure_recovery(mut self, recovery: SubprocedureRecovery) -> Self { + self.subprocedure_recovery = recovery; + self + } + /// Returns current [ProcedureState]. fn state(&self) -> ProcedureState { self.state_receiver.borrow().clone() @@ -724,7 +735,6 @@ impl LocalManager { origin: RootSubmissionOrigin, ) -> Result { ensure!(self.manager_ctx.running(), ManagerNotStartSnafu); - let meta = Arc::new(ProcedureMeta::new( procedure_id, procedure_state, @@ -735,6 +745,16 @@ impl LocalManager { procedure.type_name(), )); let event_recorder = self.event_recorder.0.lock().unwrap(); + // Serialize admission and registration so concurrent fresh submissions cannot + // exceed the limit. Recovery resumes already accepted work and bypasses it. + if matches!(origin, RootSubmissionOrigin::Fresh) { + ensure!( + self.manager_ctx.num_running_procedures() < self.config.max_running_procedures, + TooManyRunningProceduresSnafu { + max_running_procedures: self.config.max_running_procedures, + } + ); + } let runner = Runner { meta: meta.clone(), procedure, @@ -750,21 +770,13 @@ impl LocalManager { rollback_retry_attempt: 0, }; - let watcher = meta.state_receiver.clone(); - - ensure!( - self.manager_ctx.num_running_procedures() < self.config.max_running_procedures, - TooManyRunningProceduresSnafu { - max_running_procedures: self.config.max_running_procedures, - } - ); - // Inserts meta into the manager before actually spawnd the runner. ensure!( self.manager_ctx.try_insert_procedure(meta), DuplicateProcedureSnafu { procedure_id }, ); + let watcher = runner.meta.state_receiver.clone(); let initial_event = runner.build_event(origin.event_trigger()); let tracing_context = TracingContext::from_current_span(); @@ -773,7 +785,7 @@ impl LocalManager { self.manager_ctx.spawn_runner_task(procedure_id, || { common_runtime::spawn_global(async move { let span = tracing_context.attach(tracing::info_span!( - "LocalManager::submit_root_procedure", + "LocalManager::submit_root_procedure", procedure_name = %runner.meta.type_name, procedure_id = %runner.meta.id, )); @@ -798,60 +810,6 @@ impl LocalManager { Ok(watcher) } - fn submit_recovered_messages( - &self, - messages: HashMap, - init_state: InitProcedureState, - ) { - for (procedure_id, message) in &messages { - if message.parent_id.is_none() { - // This is the root procedure. We only submit the root procedure as it will - // submit sub-procedures to the manager. - let Some(mut loaded_procedure) = self - .manager_ctx - .load_one_procedure_from_message(*procedure_id, message) - else { - // Try to load other procedures. - continue; - }; - - info!( - "Recover root procedure {}-{}, step: {}", - loaded_procedure.procedure.type_name(), - procedure_id, - loaded_procedure.step - ); - - let procedure_state = match init_state { - InitProcedureState::RollingBack => ProcedureState::RollingBack { - error: Arc::new( - error::RollbackProcedureRecoveredSnafu { - error: message.error.clone().unwrap_or("Unknown error".to_string()), - } - .build(), - ), - }, - InitProcedureState::Running => ProcedureState::Running, - }; - - if let Err(e) = loaded_procedure.procedure.recover() { - error!(e; "Failed to recover procedure {}", procedure_id); - } - - if let Err(e) = self.submit_root( - *procedure_id, - procedure_state, - loaded_procedure.step, - loaded_procedure.procedure, - message.context.clone(), - RootSubmissionOrigin::Recovery, - ) { - error!(e; "Failed to recover procedure {}", procedure_id); - } - } - } - } - /// Recovers unfinished procedures and reruns them. async fn recover(&self) -> Result<()> { info!("LocalManager start to recover"); @@ -862,9 +820,13 @@ impl LocalManager { rollback_messages, finished_ids, } = self.procedure_store.load_messages().await?; - // Submits recovered messages first. - self.submit_recovered_messages(rollback_messages, InitProcedureState::RollingBack); - self.submit_recovered_messages(messages, InitProcedureState::Running); + let rollback_messages = rollback_messages + .into_iter() + .map(|(id, message)| (id, message, InitProcedureState::RollingBack)); + let messages = messages + .into_iter() + .map(|(id, message)| (id, message, InitProcedureState::Running)); + self.submit_recovered_messages(rollback_messages.chain(messages)); if !finished_ids.is_empty() { info!( @@ -2352,6 +2314,79 @@ mod tests { ); } + #[tokio::test] + async fn test_concurrent_root_submission_respects_limit() { + struct ConcurrentProcedure(Arc); + + #[async_trait::async_trait] + impl Procedure for ConcurrentProcedure { + fn type_name(&self) -> &str { + "ConcurrentProcedure" + } + + async fn execute(&mut self, _ctx: &Context) -> Result { + std::future::pending().await + } + + fn dump(&self) -> Result { + Ok(String::new()) + } + + fn lock_key(&self) -> LockKey { + // Both submissions reach metadata construction before either registers. + self.0.wait(); + LockKey::default() + } + } + + let dir = create_temp_dir("concurrent_root_submission"); + let state_store = Arc::new(ObjectStateStore::new(test_util::new_object_store(&dir))); + let manager = LocalManager::new( + ManagerConfig { + max_running_procedures: 1, + ..Default::default() + }, + state_store, + Arc::new(InMemoryPoisonStore::new()), + None, + None, + ); + manager.manager_ctx.start(); + let barrier = Arc::new(std::sync::Barrier::new(2)); + let results = std::thread::scope(|scope| { + let handles: Vec<_> = (0..2) + .map(|_| { + let barrier = barrier.clone(); + let manager = &manager; + scope.spawn(move || { + manager.submit_root( + ProcedureId::random(), + ProcedureState::Running, + 0, + Box::new(ConcurrentProcedure(barrier)), + ProcedureContext::default(), + RootSubmissionOrigin::Fresh, + ) + }) + }) + .collect(); + handles + .into_iter() + .map(|h| h.join().unwrap()) + .collect::>() + }); + manager.stop().await.unwrap(); + + assert_eq!(1, results.iter().filter(|result| result.is_ok()).count()); + assert_eq!( + 1, + results + .iter() + .filter(|result| matches!(result, Err(Error::TooManyRunningProcedures { .. }))) + .count() + ); + } + #[tokio::test] async fn test_too_many_running_procedures() { let dir = create_temp_dir("too_many_running_procedures"); diff --git a/src/common/procedure/src/local/recovery.rs b/src/common/procedure/src/local/recovery.rs new file mode 100644 index 00000000000..10de0890d12 --- /dev/null +++ b/src/common/procedure/src/local/recovery.rs @@ -0,0 +1,582 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use common_telemetry::{error, info}; + +use crate::error; +use crate::local::{LocalManager, RootSubmissionOrigin}; +use crate::procedure::InitProcedureState; +use crate::store::ProcedureMessage; +use crate::{ProcedureId, ProcedureState, ProcedureWithId}; + +/// A loaded procedure ready to resume from its checkpoint. +struct RecoveredProcedure { + procedure: ProcedureWithId, + state: ProcedureState, + step: u32, +} + +impl LocalManager { + /// Loads all eligible procedures before starting any recovered runner. + pub(crate) fn submit_recovered_messages( + &self, + messages: impl IntoIterator, + ) { + // Load both rollback and execution checkpoints before spawning. Loaders + // rebuild trackers that parents consult before submitting new work. + let procedures: Vec<_> = messages + .into_iter() + .filter_map(|(id, message, state)| self.load_recovered_procedure(id, message, state)) + .collect(); + + for RecoveredProcedure { + procedure, + state, + step, + } in procedures + { + let procedure_id = procedure.id; + // Detached children use the existing root lifecycle and cleanup. + if let Err(e) = self.submit_root( + procedure_id, + state, + step, + procedure.procedure, + procedure.context, + RootSubmissionOrigin::Recovery, + ) { + error!(e; "Failed to recover procedure {}", procedure_id); + } + } + } + + fn load_recovered_procedure( + &self, + procedure_id: ProcedureId, + message: ProcedureMessage, + init_state: InitProcedureState, + ) -> Option { + if message.parent_id.is_some() && message.subprocedure_recovery.is_discard() { + return None; + } + // The loader logs failures; skip them without adding a duplicate error. + let mut loaded = self + .manager_ctx + .load_one_procedure_from_message(procedure_id, &message)?; + match message.parent_id { + Some(parent_id) => info!( + "Recover detached subprocedure {}-{}, former parent: {}, step: {}", + loaded.procedure.type_name(), + procedure_id, + parent_id, + loaded.step + ), + None => info!( + "Recover root procedure {}-{}, step: {}", + loaded.procedure.type_name(), + procedure_id, + loaded.step + ), + } + let state = match init_state { + InitProcedureState::Running => ProcedureState::Running, + InitProcedureState::RollingBack => ProcedureState::RollingBack { + error: Arc::new( + error::RollbackProcedureRecoveredSnafu { + error: message.error.unwrap_or_else(|| "Unknown error".to_string()), + } + .build(), + ), + }, + }; + if let Err(e) = loaded.procedure.recover() { + error!(e; "Failed to recover procedure {}", procedure_id); + } + Some(RecoveredProcedure { + procedure: ProcedureWithId { + id: procedure_id, + procedure: loaded.procedure, + context: message.context, + }, + state, + step: loaded.step, + }) + } +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::time::Duration; + + use async_trait::async_trait; + use common_test_util::temp_dir::create_temp_dir; + use tokio::sync::Notify; + use tokio::time::timeout; + + use super::*; + use crate::error::Result; + use crate::local::{ManagerConfig, test_util}; + use crate::store::state_store::ObjectStateStore; + use crate::test_util::InMemoryPoisonStore; + use crate::{ + Context, LockKey, Procedure, ProcedureContext, ProcedureManager, Status, + SubprocedureRecovery, + }; + + struct MockParentProcedure { + child: Option, + recovery: SubprocedureRecovery, + } + + #[async_trait] + impl Procedure for MockParentProcedure { + fn type_name(&self) -> &str { + "MockParentProcedure" + } + + async fn execute(&mut self, _: &Context) -> Result { + let Some(child) = self.child.take() else { + return Ok(Status::done()); + }; + Ok(match self.recovery { + SubprocedureRecovery::Discard => Status::suspended(vec![child], false), + SubprocedureRecovery::Detach => Status::Suspended { + subprocedures: vec![child], + persist: false, + recovery: SubprocedureRecovery::Detach, + }, + }) + } + + fn dump(&self) -> Result { + unreachable!("The fixture deliberately does not checkpoint the parent") + } + + fn lock_key(&self) -> LockKey { + LockKey::default() + } + } + + #[derive(Clone)] + struct MockChildProcedure { + checkpointed: bool, + init_state: InitProcedureState, + entered: Arc, + release: Arc, + recoveries: Arc, + rollbacks: Arc, + } + + #[async_trait] + impl Procedure for MockChildProcedure { + fn type_name(&self) -> &str { + "MockChildProcedure" + } + + async fn execute(&mut self, _: &Context) -> Result { + if matches!(self.init_state, InitProcedureState::RollingBack) { + return error::UnexpectedSnafu { + err_msg: "Child needs rollback", + } + .fail(); + } + if !self.checkpointed { + self.checkpointed = true; + return Ok(Status::executing(true)); + } + self.entered.notify_one(); + self.release.notified().await; + Ok(Status::done()) + } + + fn recover(&mut self) -> Result<()> { + self.recoveries.fetch_add(1, Ordering::Relaxed); + // The next execution checkpoint must persist the detached root identity. + self.checkpointed = false; + Ok(()) + } + + fn dump(&self) -> Result { + Ok(self.checkpointed.to_string()) + } + + fn lock_key(&self) -> LockKey { + LockKey::default() + } + + fn rollback_supported(&self) -> bool { + true + } + + async fn rollback(&mut self, _: &Context) -> Result<()> { + self.rollbacks.fetch_add(1, Ordering::Relaxed); + self.entered.notify_one(); + self.release.notified().await; + Ok(()) + } + } + + #[tokio::test] + async fn test_detach_child_after_restart() { + check_child_recovery(SubprocedureRecovery::Detach, InitProcedureState::Running).await; + } + + #[tokio::test] + async fn test_detach_rollback_child_after_restart() { + check_child_recovery( + SubprocedureRecovery::Detach, + InitProcedureState::RollingBack, + ) + .await; + } + + #[tokio::test] + async fn test_discard_child_after_restart() { + check_child_recovery(SubprocedureRecovery::Discard, InitProcedureState::Running).await; + } + + #[tokio::test] + async fn test_discard_rollback_child_after_restart() { + check_child_recovery( + SubprocedureRecovery::Discard, + InitProcedureState::RollingBack, + ) + .await; + } + + #[tokio::test] + async fn test_load_all_procedures_before_spawning() { + // Cover both checkpoint kinds and both cross-kind parent/child orderings. + for (parent_rollback, child_rollback) in + [(false, false), (false, true), (true, false), (true, true)] + { + let dir = create_temp_dir("load-before-spawn"); + let manager = LocalManager::new( + ManagerConfig::default(), + Arc::new(ObjectStateStore::new(test_util::new_object_store(&dir))), + Arc::new(InMemoryPoisonStore::new()), + None, + None, + ); + let manager_ctx = manager.manager_ctx.clone(); + let loads = Arc::new(AtomicUsize::new(0)); + let recoveries = Arc::new(AtomicUsize::new(0)); + let loaded_count = loads.clone(); + let recovered_count = recoveries.clone(); + manager + .register_loader( + "MockChildProcedure", + Box::new(move |data| { + // A loader must never observe an already spawned parent. + // This check is deterministic even if its runner has not run yet. + assert!(manager_ctx.runner_tasks.lock().unwrap().is_empty()); + loaded_count.fetch_add(1, Ordering::Relaxed); + Ok(Box::new(MockChildProcedure { + checkpointed: data.parse().unwrap(), + init_state: InitProcedureState::Running, + entered: Arc::new(Notify::new()), + release: Arc::new(Notify::new()), + recoveries: recovered_count.clone(), + rollbacks: Arc::new(AtomicUsize::new(0)), + })) + }), + ) + .unwrap(); + let parent_id = ProcedureId::random(); + let child_id = ProcedureId::random(); + for (id, parent, rollback) in [ + (parent_id, None, parent_rollback), + (child_id, Some(parent_id), child_rollback), + ] { + let message = ProcedureMessage { + type_name: "MockChildProcedure".to_string(), + data: "true".to_string(), + parent_id: parent, + subprocedure_recovery: SubprocedureRecovery::Detach, + step: 7, + error: rollback.then(|| "Needs rollback".to_string()), + context: ProcedureContext::default(), + }; + if rollback { + manager + .procedure_store + .rollback_procedure(id, message) + .await + .unwrap(); + } else { + manager + .procedure_store + .store_procedure_message(id, &message) + .await + .unwrap(); + } + } + + manager.start().await.unwrap(); + assert_eq!(loads.load(Ordering::Relaxed), 2); + assert_eq!(recoveries.load(Ordering::Relaxed), 2); + assert_eq!(manager.manager_ctx.runner_tasks.lock().unwrap().len(), 2); + manager.stop().await.unwrap(); + } + } + + #[tokio::test] + async fn test_recover_roots_and_detached_children_over_limit() { + let dir = create_temp_dir("detached-recovery-limit"); + let manager = LocalManager::new( + ManagerConfig { + max_running_procedures: 1, + ..Default::default() + }, + Arc::new(ObjectStateStore::new(test_util::new_object_store(&dir))), + Arc::new(InMemoryPoisonStore::new()), + None, + None, + ); + manager + .register_loader( + "MockChildProcedure", + Box::new(|data| { + Ok(Box::new(MockChildProcedure { + checkpointed: data.parse().unwrap(), + init_state: InitProcedureState::Running, + entered: Arc::new(Notify::new()), + release: Arc::new(Notify::new()), + recoveries: Arc::new(AtomicUsize::new(0)), + rollbacks: Arc::new(AtomicUsize::new(0)), + })) + }), + ) + .unwrap(); + let root_id = ProcedureId::random(); + let child_ids = [ProcedureId::random(), ProcedureId::random()]; + let root = ProcedureMessage { + type_name: "MockChildProcedure".to_string(), + data: "true".to_string(), + parent_id: None, + subprocedure_recovery: SubprocedureRecovery::Discard, + step: 7, + error: None, + context: ProcedureContext::default(), + }; + manager + .procedure_store + .store_procedure_message(root_id, &root) + .await + .unwrap(); + for child_id in child_ids { + let child = ProcedureMessage { + parent_id: Some(root_id), + subprocedure_recovery: SubprocedureRecovery::Detach, + ..root.clone() + }; + manager + .procedure_store + .store_procedure_message(child_id, &child) + .await + .unwrap(); + } + + manager.start().await.unwrap(); + { + let procedures = manager.manager_ctx.procedures.read().unwrap(); + assert_eq!(procedures.len(), 3); + for id in [root_id, child_ids[0], child_ids[1]] { + let meta = &procedures[&id]; + assert_eq!(meta.parent_id, None); + assert_eq!(meta.num_children(), 0); + } + } + // Wait for detached children to checkpoint as roots, then verify that + // recovery retained their nonzero step rather than restarting from zero. + timeout(Duration::from_secs(5), async { + loop { + let stored = manager.procedure_store.load_messages().await.unwrap(); + if child_ids.iter().all(|id| { + stored + .messages + .get(id) + .is_some_and(|message| message.parent_id.is_none()) + }) { + for child_id in child_ids { + assert_eq!(stored.messages[&child_id].step, 7); + } + break; + } + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + manager.stop().await.unwrap(); + } + + async fn check_child_recovery(policy: SubprocedureRecovery, init_state: InitProcedureState) { + let dir = create_temp_dir("child-recovery"); + let state_store = Arc::new(ObjectStateStore::new(test_util::new_object_store(&dir))); + let manager = LocalManager::new( + ManagerConfig { + max_running_procedures: 1, + ..Default::default() + }, + state_store, + Arc::new(InMemoryPoisonStore::new()), + None, + None, + ); + let child_id = ProcedureId::random(); + let entered = Arc::new(Notify::new()); + let release = Arc::new(Notify::new()); + let recoveries = Arc::new(AtomicUsize::new(0)); + let rollbacks = Arc::new(AtomicUsize::new(0)); + let child = MockChildProcedure { + checkpointed: false, + init_state: init_state.clone(), + entered: entered.clone(), + release: release.clone(), + recoveries: recoveries.clone(), + rollbacks: rollbacks.clone(), + }; + let loaded_child = child.clone(); + manager + .register_loader( + "MockChildProcedure", + Box::new(move |data| { + assert_eq!( + policy, + SubprocedureRecovery::Detach, + "Discarded children must not be loaded" + ); + let mut child = loaded_child.clone(); + child.checkpointed = data.parse().unwrap(); + Ok(Box::new(child)) + }), + ) + .unwrap(); + manager.start().await.unwrap(); + let parent = ProcedureWithId::with_random_id(Box::new(MockParentProcedure { + child: Some(ProcedureWithId { + id: child_id, + procedure: Box::new(child), + context: ProcedureContext::default(), + }), + recovery: policy, + })) + .with_context(ProcedureContext { + actor: Some("alice".to_string()), + ..Default::default() + }); + let parent_id = parent.id; + manager.submit(parent).await.unwrap(); + timeout(Duration::from_secs(5), entered.notified()) + .await + .unwrap(); + + // The first child checkpoint carries the submission policy, even though + // Status::Suspended.persist was false and the parent has no checkpoint. + let stored = manager.procedure_store.load_messages().await.unwrap(); + let checkpoint = match init_state { + InitProcedureState::Running => &stored.messages[&child_id], + InitProcedureState::RollingBack => &stored.rollback_messages[&child_id], + } + .clone(); + assert_eq!(checkpoint.parent_id, Some(parent_id)); + assert_eq!(checkpoint.subprocedure_recovery, policy); + assert_eq!(checkpoint.context.actor.as_deref(), Some("alice")); + assert_eq!(checkpoint.step, 0); + if matches!(init_state, InitProcedureState::RollingBack) { + assert!(!stored.messages.contains_key(&child_id)); + assert_eq!(checkpoint.data, "false"); + } + assert!(!stored.messages.contains_key(&parent_id)); + manager.stop().await.unwrap(); + manager.start().await.unwrap(); + + if policy == SubprocedureRecovery::Discard { + assert!(manager.procedure_state(child_id).await.unwrap().is_none()); + assert_eq!(recoveries.load(Ordering::Relaxed), 0); + let stored = manager.procedure_store.load_messages().await.unwrap(); + let retained = match init_state { + InitProcedureState::Running => &stored.messages[&child_id], + InitProcedureState::RollingBack => &stored.rollback_messages[&child_id], + }; + assert_eq!(retained, &checkpoint); + manager.stop().await.unwrap(); + return; + } + + // The original parent is absent. The child owns its cleanup as a root, + // and remains recoverable across a second restart. + for restart in 1..=2 { + timeout(Duration::from_secs(5), entered.notified()) + .await + .unwrap(); + assert!(manager.procedure_state(parent_id).await.unwrap().is_none()); + let meta = manager.manager_ctx.procedures.read().unwrap()[&child_id].clone(); + assert_eq!(meta.parent_id, None); + assert_eq!(meta.context, checkpoint.context); + assert_eq!(recoveries.load(Ordering::Relaxed), restart); + let stored = manager.procedure_store.load_messages().await.unwrap(); + match init_state { + InitProcedureState::Running => { + let recovered = &stored.messages[&child_id]; + assert_eq!(recovered.parent_id, None); + assert_eq!(recovered.step, checkpoint.step); + assert_eq!(recovered.context, checkpoint.context); + } + InitProcedureState::RollingBack => { + assert!(meta.state().is_rolling_back()); + assert_eq!(stored.rollback_messages[&child_id], checkpoint); + } + } + if restart == 1 { + manager.stop().await.unwrap(); + manager.start().await.unwrap(); + } + } + release.notify_one(); + timeout(Duration::from_secs(5), async { + while manager + .manager_ctx + .runner_tasks + .lock() + .unwrap() + .contains_key(&child_id) + { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + let stored = manager.procedure_store.load_messages().await.unwrap(); + assert!(stored.messages.is_empty()); + assert!(stored.rollback_messages.is_empty()); + assert!(stored.finished_ids.is_empty()); + let state = manager.procedure_state(child_id).await.unwrap().unwrap(); + match init_state { + InitProcedureState::Running => { + assert!(state.is_done()); + assert_eq!(rollbacks.load(Ordering::Relaxed), 0); + } + InitProcedureState::RollingBack => { + assert!(state.is_failed()); + assert_eq!(rollbacks.load(Ordering::Relaxed), 3); + } + } + manager.stop().await.unwrap(); + } +} diff --git a/src/common/procedure/src/local/runner.rs b/src/common/procedure/src/local/runner.rs index 30b775dbd7d..082ff7e10b7 100644 --- a/src/common/procedure/src/local/runner.rs +++ b/src/common/procedure/src/local/runner.rs @@ -35,7 +35,7 @@ use crate::rwlock::OwnedKeyRwLockGuard; use crate::store::{ProcedureMessage, ProcedureStore}; use crate::{ BoxedProcedure, ChildSubmissionOutcome, Context, Error, EventContext, EventTrigger, Procedure, - ProcedureId, ProcedureState, ProcedureWithId, RetryPhase, Status, + ProcedureId, ProcedureState, ProcedureWithId, RetryPhase, Status, SubprocedureRecovery, }; /// A guard to cleanup procedure state. @@ -390,7 +390,11 @@ impl Runner { self.meta.set_state(ProcedureState::Running); } } - Status::Suspended { subprocedures, .. } => { + Status::Suspended { + subprocedures, + recovery, + .. + } => { let prev_state = self.meta.state(); if !matches!(prev_state, ProcedureState::Running) { info!( @@ -401,7 +405,7 @@ impl Runner { ); self.meta.set_state(ProcedureState::Running); } - self.on_suspended(subprocedures).await; + self.on_suspended(subprocedures, recovery).await; } Status::Done { output } => { if let Err(e) = self.commit_procedure().await { @@ -488,6 +492,7 @@ impl Runner { procedure_id: ProcedureId, procedure_state: ProcedureState, procedure: BoxedProcedure, + recovery: SubprocedureRecovery, ) -> (ChildSubmissionOutcome, Option) { if !self.running() { warn!( @@ -505,7 +510,7 @@ impl Runner { let step = 0; - let meta = Arc::new(ProcedureMeta::new( + let meta = ProcedureMeta::new( procedure_id, procedure_state, Some(self.meta.id), @@ -513,7 +518,9 @@ impl Runner { procedure.lock_key(), procedure.poison_keys(), procedure.type_name(), - )); + ) + .with_subprocedure_recovery(recovery); + let meta = Arc::new(meta); let runner = Runner { meta: meta.clone(), procedure, @@ -578,7 +585,11 @@ impl Runner { time::sleep(d).await; } - async fn on_suspended(&mut self, subprocedures: Vec) { + async fn on_suspended( + &mut self, + subprocedures: Vec, + recovery: SubprocedureRecovery, + ) { let has_child = !subprocedures.is_empty(); // Pre-check: detect potential deadlocks BEFORE submitting any subprocedure. @@ -629,6 +640,7 @@ impl Runner { subprocedure.id, ProcedureState::Running, subprocedure.procedure, + recovery, ); if let Some(event) = submitted_event && let Some(recorder) = self.event_recorder.as_ref() @@ -664,13 +676,17 @@ impl Runner { let data = self.procedure.dump()?; self.store - .store_procedure_with_context( + .store_procedure_message( self.meta.id, - self.step, - type_name, - data, - self.meta.parent_id, - self.meta.context.clone(), + &ProcedureMessage { + type_name, + data, + parent_id: self.meta.parent_id, + subprocedure_recovery: self.meta.subprocedure_recovery, + step: self.step, + error: None, + context: self.meta.context.clone(), + }, ) .await .map_err(|e| { @@ -709,6 +725,7 @@ impl Runner { type_name, data, parent_id: self.meta.parent_id, + subprocedure_recovery: self.meta.subprocedure_recovery, step: self.step, error: Some(error), context: self.meta.context.clone(), @@ -1061,6 +1078,7 @@ mod tests { let exec_fn = move |_| { async move { Ok(Status::Suspended { + recovery: SubprocedureRecovery::Discard, subprocedures: Vec::new(), persist: false, }) @@ -1138,6 +1156,7 @@ mod tests { if times == 1 { // Submit subprocedures. Ok(Status::Suspended { + recovery: SubprocedureRecovery::Discard, subprocedures: children_ids .into_iter() .zip(keys) @@ -1165,6 +1184,7 @@ mod tests { } else { // Return suspended to wait for notify. Ok(Status::Suspended { + recovery: SubprocedureRecovery::Discard, subprocedures: Vec::new(), persist: false, }) @@ -1491,6 +1511,7 @@ mod tests { }; Ok(Status::Suspended { + recovery: SubprocedureRecovery::Discard, subprocedures: vec![ProcedureWithId { id: child_id, procedure: Box::new(fail), @@ -1684,6 +1705,7 @@ mod tests { }; Ok(Status::Suspended { + recovery: SubprocedureRecovery::Discard, subprocedures: vec![ProcedureWithId { id: child_id, procedure: Box::new(fail), @@ -1704,6 +1726,7 @@ mod tests { } else { // Return suspended to wait for notify. Ok(Status::Suspended { + recovery: SubprocedureRecovery::Discard, subprocedures: Vec::new(), persist: false, }) @@ -2205,6 +2228,7 @@ mod tests { rollback_fn: None, }; Ok(Status::Suspended { + recovery: SubprocedureRecovery::Discard, subprocedures: vec![ProcedureWithId { id: child_id, procedure: Box::new(child), @@ -2262,6 +2286,7 @@ mod tests { rollback_fn: None, }; Ok(Status::Suspended { + recovery: SubprocedureRecovery::Discard, subprocedures: vec![ProcedureWithId { id: child_id, procedure: Box::new(child), diff --git a/src/common/procedure/src/procedure.rs b/src/common/procedure/src/procedure.rs index ba6a5d275cb..b43f275b8a2 100644 --- a/src/common/procedure/src/procedure.rs +++ b/src/common/procedure/src/procedure.rs @@ -56,6 +56,33 @@ impl ProcedureContext { } } +/// Controls how children submitted by [`Status::Suspended`] are handled after a restart. +/// +/// The policy is captured at submission and persisted with each child's checkpoints. +/// It does not change parent-child ownership during normal execution. A child without +/// a persisted checkpoint cannot be recovered, regardless of its policy. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +pub enum SubprocedureRecovery { + /// Does not recover the child. The parent decides whether to resubmit its work. + /// This does not roll back side effects or delete the child's persisted records. + #[default] + Discard, + /// Resumes the child's persisted execution or rollback as an independent root. + /// Its original id and context are retained, but its parent link is removed. + /// Completion no longer notifies the former parent. + /// All eligible procedures are loaded and their recovery hooks are called before + /// any recovered runner starts. Runner execution is not ordered relative to the + /// former parent; callers must prevent conflicting work when the parent recovers. + Detach, +} + +impl SubprocedureRecovery { + /// Returns whether the child is excluded from recovery. + pub(crate) fn is_discard(&self) -> bool { + matches!(self, Self::Discard) + } +} + /// Procedure execution status. #[derive(Debug)] pub enum Status { @@ -69,6 +96,8 @@ pub enum Status { /// The procedure has suspended itself and is waiting for subprocedures. Suspended { subprocedures: Vec, + /// Recovery policy for every child in this submission, recorded at its first checkpoint. + recovery: SubprocedureRecovery, /// Whether the framework needs to persist the procedure. persist: bool, }, @@ -85,9 +114,11 @@ pub enum Status { impl Status { /// Returns a [Status::Suspended] with given `subprocedures` and `persist` flag. + /// Children default to [`SubprocedureRecovery::Discard`]. pub fn suspended(subprocedures: Vec, persist: bool) -> Status { Status::Suspended { subprocedures, + recovery: SubprocedureRecovery::Discard, persist, } } @@ -785,12 +816,14 @@ mod tests { assert!(status.need_clean_poisons()); let status = Status::Suspended { + recovery: SubprocedureRecovery::Discard, subprocedures: Vec::new(), persist: false, }; assert!(!status.need_persist()); let status = Status::Suspended { + recovery: SubprocedureRecovery::Discard, subprocedures: Vec::new(), persist: true, }; diff --git a/src/common/procedure/src/store.rs b/src/common/procedure/src/store.rs index b97dab00d00..dd41c13e072 100644 --- a/src/common/procedure/src/store.rs +++ b/src/common/procedure/src/store.rs @@ -22,7 +22,7 @@ use snafu::ResultExt; use crate::error::{Result, ToJsonSnafu}; pub(crate) use crate::store::state_store::StateStoreRef; -use crate::{ProcedureContext, ProcedureId}; +use crate::{ProcedureContext, ProcedureId, SubprocedureRecovery}; pub mod poison_store; pub mod state_store; @@ -50,6 +50,9 @@ pub struct ProcedureMessage { pub data: String, /// Parent procedure id. pub parent_id: Option, + /// Recovery policy for a child submission, ignored for root procedures. + #[serde(default, skip_serializing_if = "SubprocedureRecovery::is_discard")] + pub subprocedure_recovery: SubprocedureRecovery, /// Current step. pub step: u32, /// Errors raised during the procedure. @@ -111,6 +114,7 @@ impl ProcedureStore { .await } + #[cfg(test)] pub(crate) async fn store_procedure_with_context( &self, procedure_id: ProcedureId, @@ -124,14 +128,24 @@ impl ProcedureStore { type_name, data, parent_id, + subprocedure_recovery: SubprocedureRecovery::Discard, step, error: None, context, }; + self.store_procedure_message(procedure_id, &message).await + } + + /// Persists a procedure checkpoint with its recovery policy and submission context. + pub(crate) async fn store_procedure_message( + &self, + procedure_id: ProcedureId, + message: &ProcedureMessage, + ) -> Result<()> { let key = ParsedKey { prefix: &self.proc_path, procedure_id, - step, + step: message.step, key_type: KeyType::Step, } .to_string(); @@ -486,6 +500,7 @@ mod tests { type_name: "TestMessage".to_string(), data: "no parent id".to_string(), parent_id: None, + subprocedure_recovery: SubprocedureRecovery::Discard, step: 4, error: None, context: ProcedureContext::default(), @@ -529,6 +544,20 @@ mod tests { assert_eq!(legacy.context, ProcedureContext::default()); } + #[test] + fn test_subprocedure_recovery_serialization() { + let legacy = r#"{"type_name":"TestMessage","data":"checkpoint","parent_id":null,"step":4}"#; + let mut message: ProcedureMessage = serde_json::from_str(legacy).unwrap(); + assert_eq!(message.subprocedure_recovery, SubprocedureRecovery::Discard); + assert_eq!(serde_json::to_string(&message).unwrap(), legacy); + message.subprocedure_recovery = SubprocedureRecovery::Detach; + let json = serde_json::to_string(&message).unwrap(); + assert_eq!( + serde_json::from_str::(&json).unwrap(), + message + ); + } + struct MockProcedure { data: String, } @@ -589,6 +618,7 @@ mod tests { type_name: "MockProcedure".to_string(), data: "test store procedure".to_string(), parent_id: None, + subprocedure_recovery: SubprocedureRecovery::Discard, step: 0, error: None, context: ProcedureContext::default(), @@ -638,6 +668,7 @@ mod tests { type_name, data, parent_id: None, + subprocedure_recovery: SubprocedureRecovery::Discard, step: 1, error: None, context: ProcedureContext::default(), From bd77b599cc25caf7980d3a1627cdd7727c888778 Mon Sep 17 00:00:00 2001 From: discord9 Date: Wed, 28 Oct 2026 07:27:04 +0000 Subject: [PATCH 120/239] refactor(flow): remove tracked-write protocol (#9992) Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- Cargo.lock | 1 - licenserc-enterprise.toml | 5 - licenserc.toml | 5 - src/common/grpc/Cargo.toml | 4 - src/common/grpc/src/flow_write.rs | 477 --------- src/common/grpc/src/lib.rs | 2 - src/datanode/Cargo.toml | 2 +- src/datanode/src/region_server.rs | 19 - src/datanode/src/region_server/flow_write.rs | 969 ------------------ src/flow/Cargo.toml | 2 +- src/flow/src/batching_mode/frontend_client.rs | 98 -- .../frontend_client/flow_write.rs | 758 -------------- src/frontend/Cargo.toml | 2 - src/frontend/src/instance.rs | 31 - src/frontend/src/instance/builder.rs | 2 - src/frontend/src/instance/flow_write.rs | 227 ---- src/frontend/src/instance/grpc.rs | 14 - src/operator/Cargo.toml | 1 - src/operator/src/insert.rs | 43 - src/operator/src/insert/flow_write.rs | 391 ------- src/session/Cargo.toml | 3 - src/session/src/context.rs | 5 - 22 files changed, 2 insertions(+), 3059 deletions(-) delete mode 100644 src/common/grpc/src/flow_write.rs delete mode 100644 src/datanode/src/region_server/flow_write.rs delete mode 100644 src/flow/src/batching_mode/frontend_client/flow_write.rs delete mode 100644 src/frontend/src/instance/flow_write.rs delete mode 100644 src/operator/src/insert/flow_write.rs diff --git a/Cargo.lock b/Cargo.lock index 639a771816a..6ada3b0a295 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2588,7 +2588,6 @@ dependencies = [ "tokio-util", "tonic 0.14.2", "tower 0.5.2", - "uuid", "vec1", ] diff --git a/licenserc-enterprise.toml b/licenserc-enterprise.toml index 805358f1f1e..d11d344b64d 100644 --- a/licenserc-enterprise.toml +++ b/licenserc-enterprise.toml @@ -15,11 +15,6 @@ headerPath = "licenses/enterprise-header.txt" # with the "# enterprise" block in `licenserc.toml`'s `excludes` — both are # enforced by `make check-enterprise-license`. includes = [ - "src/common/grpc/src/flow_write.rs", - "src/frontend/src/instance/flow_write.rs", - "src/datanode/src/region_server/flow_write.rs", - "src/operator/src/insert/flow_write.rs", - "src/flow/src/batching_mode/frontend_client/flow_write.rs", "src/common/meta/src/rpc/ddl/trigger.rs", "src/operator/src/expr_helper/trigger.rs", "src/sql/src/statements/alter/trigger.rs", diff --git a/licenserc.toml b/licenserc.toml index e84d3a1f8ff..42e6464017c 100644 --- a/licenserc.toml +++ b/licenserc.toml @@ -31,11 +31,6 @@ excludes = [ # Keep this list in sync with that file's `includes`; enforced by # `make check-enterprise-license`. # enterprise:start - "src/common/grpc/src/flow_write.rs", - "src/frontend/src/instance/flow_write.rs", - "src/datanode/src/region_server/flow_write.rs", - "src/operator/src/insert/flow_write.rs", - "src/flow/src/batching_mode/frontend_client/flow_write.rs", "src/common/meta/src/rpc/ddl/trigger.rs", "src/operator/src/expr_helper/trigger.rs", "src/sql/src/statements/alter/trigger.rs", diff --git a/src/common/grpc/Cargo.toml b/src/common/grpc/Cargo.toml index fbfab1b6318..58bd04db093 100644 --- a/src/common/grpc/Cargo.toml +++ b/src/common/grpc/Cargo.toml @@ -7,9 +7,6 @@ license.workspace = true [lints] workspace = true -[features] -enterprise = [] - [dependencies] api.workspace = true arrow-flight.workspace = true @@ -35,7 +32,6 @@ tokio.workspace = true tokio-util.workspace = true tonic.workspace = true tower.workspace = true -uuid.workspace = true vec1 = "1.12" [dev-dependencies] diff --git a/src/common/grpc/src/flow_write.rs b/src/common/grpc/src/flow_write.rs deleted file mode 100644 index 6d085a500bd..00000000000 --- a/src/common/grpc/src/flow_write.rs +++ /dev/null @@ -1,477 +0,0 @@ -// Copyright 2023-2026 GrepTime Inc. -// -// This file is part of the GreptimeDB Enterprise Edition and is licensed under -// the GreptimeDB Enterprise License. You may not use this file except in -// compliance with that license. A copy of the license is available at the root -// of this repository in the file LICENSE-ENTERPRISE. -// -// Unless required by applicable law or agreed to in writing, this software is -// distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, -// either express or implied. - -use std::collections::HashMap; -use std::sync::{Arc, Mutex}; - -use common_error::ext::{ErrorExt, StackError}; -use common_error::status_code::StatusCode; - -use common_recordbatch::{RecordBatch, RecordBatches}; -use datatypes::arrow::array::{Array, StringArray}; -use datatypes::prelude::{ConcreteDataType, VectorRef}; -use datatypes::schema::{ColumnSchema, Schema}; -use datatypes::vectors::StringVector; -use serde::{Deserialize, Serialize}; -use tokio::sync::watch; -use uuid::Uuid; - -pub const COMMAND_KEY: &str = "flow_write_command"; -pub const CHILD_KEY: &str = "flow_write_child"; - -/// The control command carried by a Flow write request. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -pub enum Command { - Open { nonce: String }, - Write { capability: String }, - Close { capability: String }, -} - -impl Command { - /// Serializes this command as JSON. - pub fn encode(&self) -> serde_json::Result { - serde_json::to_string(self) - } - - /// Parses a serialized command. - pub fn decode(value: &str) -> serde_json::Result { - serde_json::from_str(value) - } -} - -/// A safe terminal application error. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -pub struct TerminalError { - pub code: u32, - pub message: String, -} - -impl std::fmt::Display for TerminalError { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "{}", self.message) - } -} - -impl std::error::Error for TerminalError {} - -impl StackError for TerminalError { - fn debug_fmt(&self, layer: usize, buf: &mut Vec) { - buf.push(format!("{layer}: {self:?}")); - } - - fn next(&self) -> Option<&dyn StackError> { - None - } -} - -impl ErrorExt for TerminalError { - fn status_code(&self) -> StatusCode { - StatusCode::from_u32(self.code).unwrap_or(StatusCode::Internal) - } - - fn as_any(&self) -> &dyn std::any::Any { - self - } -} - -/// A finite response for a tracked Flow write. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub enum Reply { - Opened { - capability: String, - }, - Closed, - Complete { - rows: usize, - cost: usize, - metrics: Option, - error: Option, - }, -} - -fn envelope_schema() -> Arc { - Arc::new(Schema::new(vec![ColumnSchema::new( - "envelope", - ConcreteDataType::string_datatype(), - false, - )])) -} - -fn response_error(reason: impl Into) -> common_recordbatch::error::Error { - common_recordbatch::error::Error::CreateRecordBatches { - reason: reason.into(), - location: snafu::Location::new(file!(), line!(), column!()), - } -} - -/// Encodes a reply in the finite one-string-column record-batch envelope. -pub fn reply_batches(reply: &Reply) -> common_recordbatch::error::Result { - let value = serde_json::to_string(reply).map_err(|e| response_error(e.to_string()))?; - let batch = RecordBatch::new( - envelope_schema(), - vec![Arc::new(StringVector::from(vec![value])) as VectorRef], - )?; - RecordBatches::try_new(batch.schema.clone(), vec![batch]) -} - -/// Decodes a reply from its finite one-string-column record-batch envelope. -pub fn decode_reply(batches: &RecordBatches) -> common_recordbatch::error::Result { - let batch = batches - .iter() - .next() - .ok_or_else(|| response_error("missing reply batch"))?; - if batches.iter().count() != 1 || batch.num_columns() != 1 || batch.num_rows() != 1 { - return Err(response_error("reply must be exactly one string value")); - } - let values = batch.columns()[0] - .as_any() - .downcast_ref::() - .ok_or_else(|| response_error("reply value is not a string"))?; - if values.is_null(0) { - return Err(response_error("reply value is null")); - } - let value = values.value(0); - serde_json::from_str(value).map_err(|e| response_error(e.to_string())) -} - -/// Errors that reject admission or indicate that an admitted writer was lost. -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum RegistryError { - UnknownCapability, - Closing, - WriterLost, - PayloadMismatch, -} - -impl std::fmt::Display for RegistryError { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "{self:?}") - } -} - -impl std::error::Error for RegistryError {} - -impl StackError for RegistryError { - fn debug_fmt(&self, layer: usize, buf: &mut Vec) { - buf.push(format!("{layer}: {self:?}")); - } - - fn next(&self) -> Option<&dyn StackError> { - None - } -} - -impl ErrorExt for RegistryError { - fn status_code(&self) -> StatusCode { - StatusCode::InvalidArguments - } - - fn as_any(&self) -> &dyn std::any::Any { - self - } -} - -/// Result of atomically admitting or joining a payload. -#[derive(Debug)] -pub enum Admission { - Start(watch::Sender>>), - Join(watch::Receiver>>), -} - -#[derive(Debug)] -struct Entry { - payload: Option

, - sender: Option>>>, - receiver: watch::Receiver>>, - closing: bool, -} - -#[derive(Debug)] -struct State { - openers: HashMap, - entries: HashMap>, -} - -/// Receiver-issued capabilities for single-write admission and completion sharing. -#[derive(Debug)] -pub struct FlowWriteRegistry { - state: Mutex>, -} - -impl Default for FlowWriteRegistry { - fn default() -> Self { - Self { - state: Mutex::new(State { - openers: HashMap::new(), - entries: HashMap::new(), - }), - } - } -} - -impl FlowWriteRegistry { - /// Returns the active capability for this opener nonce, creating one if needed. - pub fn open(&self, nonce: &str) -> Result { - let mut state = self - .state - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - if let Some(capability) = state.openers.get(nonce) { - if state - .entries - .get(capability) - .is_some_and(|entry| entry.closing) - { - return Err(RegistryError::Closing); - } - return Ok(capability.clone()); - } - let capability = Uuid::new_v4().to_string(); - let (sender, receiver) = watch::channel(None); - state.openers.insert(nonce.to_string(), capability.clone()); - state.entries.insert( - capability.clone(), - Entry { - payload: None, - sender: Some(sender), - receiver, - closing: false, - }, - ); - Ok(capability) - } - - /// Atomically starts one payload or joins its identical in-flight completion. - pub fn admit(&self, capability: &str, payload: P) -> Result, RegistryError> { - let mut state = self - .state - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - let entry = state - .entries - .get_mut(capability) - .ok_or(RegistryError::UnknownCapability)?; - if entry.closing { - return Err(RegistryError::Closing); - } - match &entry.payload { - Some(existing) if existing != &payload => Err(RegistryError::PayloadMismatch), - Some(_) => Ok(Admission::Join(entry.receiver.clone())), - None => { - entry.payload = Some(payload); - let sender = entry.sender.take().ok_or(RegistryError::Closing)?; - Ok(Admission::Start(sender)) - } - } - } - - /// Closes admission, drains an admitted writer, then deletes the capability. - pub async fn close(&self, capability: &str) -> Result<(), RegistryError> { - let receiver = { - let mut state = self - .state - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - let Some(entry) = state.entries.get_mut(capability) else { - return Ok(()); - }; - entry.closing = true; - if entry.payload.is_none() { - state.entries.remove(capability); - state.openers.retain(|_, value| value != capability); - return Ok(()); - } - entry.receiver.clone() - }; - wait_completion(receiver).await?; - let mut state = self - .state - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - state.entries.remove(capability); - state.openers.retain(|_, value| value != capability); - Ok(()) - } -} - -/// Waits for an admitted writer's positive terminal completion. -pub async fn wait_completion( - mut receiver: watch::Receiver>>, -) -> Result, RegistryError> { - loop { - if let Some(value) = receiver.borrow().clone() { - return Ok(value); - } - if receiver.changed().await.is_err() { - return Err(RegistryError::WriterLost); - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use std::sync::Barrier; - use std::thread; - - #[tokio::test] - async fn concurrent_admission_has_one_owner_and_same_payload_joins() { - let registry = Arc::new(FlowWriteRegistry::::default()); - let cap = registry.open("nonce").unwrap(); - let barrier = Arc::new(Barrier::new(8)); - let mut workers = Vec::new(); - for _ in 0..8 { - let registry = registry.clone(); - let cap = cap.clone(); - let barrier = barrier.clone(); - workers.push(thread::spawn(move || { - barrier.wait(); - registry.admit(&cap, "payload".into()).unwrap() - })); - } - let admissions: Vec<_> = workers.into_iter().map(|w| w.join().unwrap()).collect(); - assert_eq!( - admissions - .iter() - .filter(|a| matches!(a, Admission::Start(_))) - .count(), - 1 - ); - let sender = admissions - .into_iter() - .find_map(|a| match a { - Admission::Start(s) => Some(s), - _ => None, - }) - .unwrap(); - sender.send(Some(Arc::new(7))).unwrap(); - registry.close(&cap).await.unwrap(); - } - - #[tokio::test] - async fn close_drains_active_writer_and_deletes_old_capability() { - let registry = Arc::new(FlowWriteRegistry::::default()); - let cap = registry.open("first").unwrap(); - let Admission::Start(sender) = registry.admit(&cap, "payload".into()).unwrap() else { - panic!() - }; - let close_registry = registry.clone(); - let close_cap = cap.clone(); - let close = tokio::spawn(async move { close_registry.close(&close_cap).await }); - tokio::task::yield_now().await; - assert!(matches!( - registry.admit(&cap, "payload".into()), - Err(RegistryError::Closing) - )); - assert!(!close.is_finished()); - sender.send(Some(Arc::new(9))).unwrap(); - close.await.unwrap().unwrap(); - assert!(matches!( - registry.admit(&cap, "payload".into()), - Err(RegistryError::UnknownCapability) - )); - registry.close(&cap).await.unwrap(); - let new_cap = registry.open("first").unwrap(); - assert_ne!(cap, new_cap); - assert!(matches!( - registry.admit(&cap, "late".into()), - Err(RegistryError::UnknownCapability) - )); - assert!(matches!( - registry.admit(&new_cap, "new".into()).unwrap(), - Admission::Start(_) - )); - } - - #[tokio::test] - async fn writer_loss_is_not_completion() { - let registry = FlowWriteRegistry::::default(); - let cap = registry.open("nonce").unwrap(); - let Admission::Start(sender) = registry.admit(&cap, "payload".into()).unwrap() else { - panic!() - }; - let Admission::Join(receiver) = registry.admit(&cap, "payload".into()).unwrap() else { - panic!() - }; - drop(sender); - assert_eq!( - wait_completion(receiver).await.unwrap_err(), - RegistryError::WriterLost - ); - assert_eq!( - registry.close(&cap).await.unwrap_err(), - RegistryError::WriterLost - ); - } - - #[test] - fn command_and_reply_envelope_round_trip_and_malformed_envelopes() { - let command = Command::Open { nonce: "n".into() }; - assert_eq!( - Command::decode(&command.encode().unwrap()).unwrap(), - command - ); - let reply = Reply::Complete { - rows: 2, - cost: 3, - metrics: Some(common_recordbatch::adapter::RecordBatchMetrics::default()), - error: Some(TerminalError { - code: 4, - message: "terminal".into(), - }), - }; - let batches = reply_batches(&reply).unwrap(); - let decoded = decode_reply(&batches).unwrap(); - assert!(matches!( - decoded, - Reply::Complete { - rows: 2, - cost: 3, - error: Some(_), - .. - } - )); - - let schema = envelope_schema(); - let two_rows = RecordBatch::new( - schema.clone(), - vec![Arc::new(StringVector::from(vec!["a", "b"])) as VectorRef], - ) - .unwrap(); - assert!( - decode_reply(&RecordBatches::try_new(schema.clone(), vec![two_rows]).unwrap()).is_err() - ); - let int_schema = Arc::new(Schema::new(vec![ColumnSchema::new( - "envelope", - ConcreteDataType::int32_datatype(), - false, - )])); - let non_string = RecordBatch::new( - int_schema.clone(), - vec![Arc::new(datatypes::vectors::Int32Vector::from(vec![Some(1)])) as VectorRef], - ) - .unwrap(); - assert!( - decode_reply(&RecordBatches::try_new(int_schema, vec![non_string]).unwrap()).is_err() - ); - let invalid_json = RecordBatch::new( - envelope_schema(), - vec![Arc::new(StringVector::from(vec!["not json"])) as VectorRef], - ) - .unwrap(); - assert!( - decode_reply( - &RecordBatches::try_new(invalid_json.schema.clone(), vec![invalid_json]).unwrap() - ) - .is_err() - ); - } -} diff --git a/src/common/grpc/src/lib.rs b/src/common/grpc/src/lib.rs index 0ac05863f3e..8527dd079b8 100644 --- a/src/common/grpc/src/lib.rs +++ b/src/common/grpc/src/lib.rs @@ -15,8 +15,6 @@ pub mod channel_manager; pub mod error; pub mod flight; -#[cfg(feature = "enterprise")] -pub mod flow_write; pub mod precision; pub mod reloadable_tls; pub mod select; diff --git a/src/datanode/Cargo.toml b/src/datanode/Cargo.toml index bef2641415e..0720a576b9c 100644 --- a/src/datanode/Cargo.toml +++ b/src/datanode/Cargo.toml @@ -6,7 +6,7 @@ license.workspace = true [features] testing = [] -enterprise = ["mito2/enterprise", "common-grpc/enterprise"] +enterprise = ["mito2/enterprise"] [lints] workspace = true diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index 06b9d42fd00..6055b07dbf6 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -13,10 +13,6 @@ // limitations under the License. mod catalog; -#[cfg(feature = "enterprise")] -mod flow_write; -#[cfg(feature = "enterprise")] -use flow_write::TrackedPayload; mod registrations; mod remote_dyn_filter; @@ -42,8 +38,6 @@ use async_trait::async_trait; use bytes::Bytes; use common_error::ext::{BoxedError, ErrorExt}; use common_error::status_code::StatusCode; -#[cfg(feature = "enterprise")] -use common_grpc::flow_write::{COMMAND_KEY, FlowWriteRegistry, Reply}; use common_meta::datanode::TopicStatsReporter; use common_query::OutputData; use common_query::request::QueryRequest; @@ -122,8 +116,6 @@ pub struct RegionServer { inner: Arc, flight_compression: FlightCompression, suspend: Arc, - #[cfg(feature = "enterprise")] - flow_write_registry: Arc>, } pub struct RegionStat { @@ -172,8 +164,6 @@ impl RegionServer { )), flight_compression, suspend: Arc::new(AtomicBool::new(false)), - #[cfg(feature = "enterprise")] - flow_write_registry: Arc::new(FlowWriteRegistry::default()), } } @@ -1069,15 +1059,6 @@ impl RegionServerHandler for RegionServer { &self, request: RegionRequestV1, ) -> ServerResult { - #[cfg(feature = "enterprise")] - if request - .header - .as_ref() - .and_then(|header| header.query_context.as_ref()) - .is_some_and(|context| context.extensions.contains_key(COMMAND_KEY)) - { - return self.handle_flow_write_request(request).await; - } let accept_partial_result = request .header .as_ref() diff --git a/src/datanode/src/region_server/flow_write.rs b/src/datanode/src/region_server/flow_write.rs deleted file mode 100644 index bc9bf65d69d..00000000000 --- a/src/datanode/src/region_server/flow_write.rs +++ /dev/null @@ -1,969 +0,0 @@ -// Copyright 2023-2026 GrepTime Inc. -// -// This file is part of the GreptimeDB Enterprise Edition and is licensed under -// the GreptimeDB Enterprise License. You may not use this file except in -// compliance with that license. A copy of the license is available at the root -// of this repository in the file LICENSE-ENTERPRISE. -// -// Unless required by applicable law or agreed to in writing, this software is -// distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, -// either express or implied. - -//! Datanode receiver for tracked Flow writes. -//! -//! A single immutable operation per receiver-issued capability: the first -//! admitted `Write` becomes the owner and executes the Inserts exactly once on -//! the region engines; concurrent or retried identical writers join the -//! owner's completion. Dropping a waiter (for example a lost RPC response) -//! never cancels the engine work, and closing a capability rejects later -//! writes with that capability. If the owner's outcome becomes unknown (a -//! child ingest worker was lost), no positive terminal completion is -//! published: waiters observe `WriterLost`; the capability admits no second -//! owner, so no overlapping engine write can be produced. - -use std::sync::Arc; - -use api::v1::region::RegionRequest as RegionRequestV1; -use api::v1::region::region_request; -use common_error::ext::ErrorExt; -use common_grpc::flow_write::{Admission, COMMAND_KEY, Command, Reply, TerminalError}; -use common_runtime::runtime::RuntimeTrait; -use common_telemetry::tracing_context::{FutureExt, TracingContext}; -use futures_util::future::join_all; -use prost::Message; -use servers::error::{InvalidQuerySnafu, Result as ServerResult}; -use snafu::OptionExt; - -use crate::region_server::RegionServer; - -/// Payload of a tracked write: encoded Inserts rows (exact bytes, so floats -/// keep bit equality even for NaN), schemas compared by protobuf equality, and -/// semantic header identity (dbname and command-free query context). Retrying -/// with different business data or context is rejected instead of joining. -pub(crate) type TrackedPayload = ( - Vec, - Vec>>, - String, - Option, -); - -fn tracked_payload(request: &RegionRequestV1) -> ServerResult { - let body = request.body.as_ref().context(InvalidQuerySnafu { - reason: "Expecting non-empty region request body.", - })?; - if !matches!(body, region_request::Body::Inserts(_)) { - return Err(servers::error::Error::InvalidQuery { - reason: "Tracked Flow write supports Inserts requests only.".to_string(), - location: snafu::Location::new(file!(), line!(), column!()), - }); - } - // Oneof is not `prost::Message`; the variant message itself is. Schema - // options contain unordered maps, so compare schemas separately rather than - // letting their wire ordering affect retries; retain encoded rows for NaN. - let (bytes, schemas) = match body { - region_request::Body::Inserts(inserts) => { - let mut inserts = inserts.clone(); - let schemas = inserts - .requests - .iter_mut() - .map(|request| { - request - .rows - .as_mut() - .map(|rows| std::mem::take(&mut rows.schema)) - }) - .collect(); - (inserts.encode_to_vec(), schemas) - } - _ => unreachable!("checked above"), - }; - let header = request.header.as_ref(); - let mut query_context = header.and_then(|header| header.query_context.clone()); - if let Some(context) = query_context.as_mut() { - // The command extension differs per control request, not per business. - context.extensions.remove(COMMAND_KEY); - } - Ok(( - bytes, - schemas, - header - .map(|header| header.dbname.clone()) - .unwrap_or_default(), - query_context, - )) -} - -fn ensure_empty_control_body(request: &RegionRequestV1) -> ServerResult<()> { - match request.body.as_ref() { - Some(region_request::Body::Inserts(inserts)) if inserts.requests.is_empty() => Ok(()), - _ => Err(servers::error::Error::InvalidQuery { - reason: "Tracked Flow write control requests must carry an empty Inserts body." - .to_string(), - location: snafu::Location::new(file!(), line!(), column!()), - }), - } -} - -fn registry_error(error: common_grpc::flow_write::RegistryError) -> servers::error::Error { - servers::error::Error::UnexpectedResult { - reason: format!("Tracked Flow write registry rejected request: {error:?}"), - location: snafu::Location::new(file!(), line!(), column!()), - } -} - -fn response(reply: Reply) -> ServerResult { - let encoded = - serde_json::to_vec(&reply).map_err(|error| servers::error::Error::UnexpectedResult { - reason: format!("Failed to encode tracked Flow write reply: {error}"), - location: snafu::Location::new(file!(), line!(), column!()), - })?; - let mut extensions = std::collections::HashMap::new(); - extensions.insert(COMMAND_KEY.to_string(), encoded); - Ok(api::v1::region::RegionResponse { - header: Some(api::v1::ResponseHeader { - status: Some(api::v1::Status { - status_code: common_error::status_code::StatusCode::Success as _, - ..Default::default() - }), - }), - affected_rows: 0, - extensions, - metadata: Vec::new(), - failures: Vec::new(), - }) -} - -impl RegionServer { - pub(crate) async fn handle_flow_write_request( - &self, - request: RegionRequestV1, - ) -> ServerResult { - let command = request - .header - .as_ref() - .and_then(|header| header.query_context.as_ref()) - .and_then(|context| context.extensions.get(COMMAND_KEY)) - .and_then(|value| Command::decode(value).ok()); - let command = match command { - Some(command) => command, - None => { - // The extension exists but is not a decodable command: reject - // before any ordinary DML dispatch. - return Err(InvalidQuerySnafu { - reason: "Malformed tracked Flow write command.", - } - .build()); - } - }; - - match command { - Command::Open { nonce } => { - ensure_empty_control_body(&request)?; - let capability = self - .flow_write_registry - .open(&nonce) - .map_err(registry_error)?; - response(Reply::Opened { capability }) - } - Command::Close { capability } => { - ensure_empty_control_body(&request)?; - self.flow_write_registry - .close(&capability) - .await - .map_err(registry_error)?; - response(Reply::Closed) - } - Command::Write { capability } => { - let payload = tracked_payload(&request)?; - let admission = self - .flow_write_registry - .admit(&capability, payload) - .map_err(registry_error)?; - let reply = match admission { - Admission::Start(owner) => { - // Subscribe before the owner sender moves into the - // spawned task: this caller stays a plain waiter whose - // drop cannot cancel the engine work. - let receiver = owner.subscribe(); - let server = self.clone(); - let body = request.body; - let header = request.header; - self.inner.runtime.spawn(async move { - let completion = server - .run_tracked_write(header, body) - .trace(TracingContext::from_current_span().attach( - common_telemetry::tracing::info_span!( - "RegionServer::tracked_flow_write" - ), - )) - .await; - match completion { - Some(reply) => { - let _ = owner.send(Some(Arc::new(reply))); - } - None => { - // Unknown worker outcome: no positive terminal - // result is published. Dropping the sender lets - // waiters observe `WriterLost`. - drop(owner); - } - } - }); - common_grpc::flow_write::wait_completion(receiver) - .await - .map_err(registry_error)? - } - Admission::Join(waiter) => common_grpc::flow_write::wait_completion(waiter) - .await - .map_err(registry_error)?, - }; - response((*reply).clone()) - } - } - } - - /// Executes one admitted tracked write. Returns `None` only when the - /// outcome is unknown (a child worker was lost); malformed request bodies - /// are terminal because no child was ever dispatched. - async fn run_tracked_write( - &self, - header: Option, - body: Option, - ) -> Option { - let requests = - match body.map(store_api::region_request::RegionRequest::try_from_request_body) { - Some(Ok(requests)) => requests, - Some(Err(error)) => { - return Some(Reply::Complete { - rows: 0, - cost: 0, - metrics: None, - error: Some(TerminalError { - code: error.status_code() as u32, - message: error.output_msg(), - }), - }); - } - // `None` is unreachable here: the Write path already rejected an - // empty body before admission. Treat it as unknown rather than - // inventing a terminal result. - None => return None, - }; - let tracing = header - .as_ref() - .map(|header| TracingContext::from_w3c(&header.tracing_context)) - .unwrap_or_default(); - - let mut unknown_worker = false; - let mut first_error: Option = None; - let mut affected_rows = 0usize; - - let tasks = requests.into_iter().map(|(region_id, request)| { - let span = tracing.attach(common_telemetry::tracing::info_span!( - "RegionServer::tracked_region_write", - region_id = region_id.to_string() - )); - async move { self.handle_request(region_id, request).trace(span).await } - }); - // Drain every per-region child before any terminal decision. A lost - // ingest worker (`RuntimeJoin`) is unknown and outranks any ordinary - // application error observed during the same drain. - for result in join_all(tasks).await { - match result { - Ok(response) => affected_rows += response.affected_rows, - Err(crate::error::Error::RuntimeJoin { .. }) => { - unknown_worker = true; - } - Err(error) if first_error.is_none() => first_error = Some(error), - Err(_) => {} - } - } - - if unknown_worker { - return None; - } - - Some(match first_error { - None => Reply::Complete { - rows: affected_rows, - cost: affected_rows, - metrics: None, - error: None, - }, - Some(error) => Reply::Complete { - rows: 0, - cost: 0, - metrics: None, - error: Some(TerminalError { - code: error.status_code() as u32, - message: error.output_msg(), - }), - }, - }) - } -} - -#[cfg(test)] -mod tests { - use std::collections::HashMap; - use std::sync::Arc; - use std::sync::atomic::{AtomicUsize, Ordering}; - - use api::region::RegionResponse; - use api::v1::region::region_request::Body; - use api::v1::region::{ - InsertRequest, InsertRequests, RegionRequest as Request, RegionRequestHeader, - }; - use api::v1::{ColumnDataType, ColumnSchema, QueryContext, Row, Rows, SemanticType, Value}; - use common_grpc::flow_write::{self, Command, Reply}; - use servers::grpc::region_server::RegionServerHandler; - use store_api::region_engine::{RegionEngine, RegionEngineRef}; - use store_api::storage::RegionId; - use tokio::sync::Notify; - use tokio::time::{Duration, timeout}; - - use super::*; - use crate::tests::{MockRegionEngine, mock_region_server}; - - #[derive(Clone, Default)] - struct Gate { - entered: Arc, - notified: Arc, - open: Arc, - fail: bool, - } - - impl Gate { - async fn wait_entered(&self, count: usize) { - timeout(Duration::from_secs(5), async { - loop { - let notified = self.notified.notified(); - tokio::pin!(notified); - notified.as_mut().enable(); - if self.entered.load(Ordering::SeqCst) >= count { - break; - } - notified.await; - } - }) - .await - .unwrap(); - } - fn release(&self) { - self.open.store(1, Ordering::SeqCst); - self.notified.notify_waiters(); - } - async fn wait_open(&self) { - loop { - let notified = self.notified.notified(); - tokio::pin!(notified); - notified.as_mut().enable(); - if self.open.load(Ordering::SeqCst) != 0 { - break; - } - notified.await; - } - } - } - - struct GatedEngine { - gate: Gate, - inner: Arc, - panic_id: Option, - } - #[async_trait::async_trait] - impl RegionEngine for GatedEngine { - fn name(&self) -> &str { - self.inner.name() - } - async fn handle_request( - &self, - id: RegionId, - _request: store_api::region_request::RegionRequest, - ) -> Result { - self.gate.entered.fetch_add(1, Ordering::SeqCst); - self.gate.notified.notify_waiters(); - if self.panic_id == Some(id.as_u64()) { - panic!("test ingest panic"); - } - if self.gate.fail && id.as_u64() == 1 { - return Err(common_error::ext::BoxedError::new( - crate::error::RegionNotFoundSnafu { region_id: id }.build(), - )); - } - self.gate.wait_open().await; - Ok(RegionResponse::new(1)) - } - async fn handle_query( - &self, - id: RegionId, - req: store_api::storage::ScanRequest, - ) -> Result - { - self.inner.handle_query(id, req).await - } - async fn get_committed_sequence( - &self, - id: RegionId, - ) -> Result { - self.inner.get_committed_sequence(id).await - } - async fn get_metadata( - &self, - id: RegionId, - ) -> Result { - self.inner.get_metadata(id).await - } - fn region_statistic( - &self, - id: RegionId, - ) -> Option { - self.inner.region_statistic(id) - } - async fn stop(&self) -> Result<(), common_error::ext::BoxedError> { - self.inner.stop().await - } - fn set_region_role( - &self, - id: RegionId, - role: store_api::region_engine::RegionRole, - ) -> Result<(), common_error::ext::BoxedError> { - self.inner.set_region_role(id, role) - } - - fn set_region_role_with_lease( - &self, - id: RegionId, - role: store_api::region_engine::RegionRole, - deadline: std::time::Instant, - ) -> Result<(), common_error::ext::BoxedError> { - self.inner.set_region_role_with_lease(id, role, deadline) - } - async fn sync_region( - &self, - id: RegionId, - req: store_api::region_engine::SyncRegionFromRequest, - ) -> Result - { - self.inner.sync_region(id, req).await - } - async fn remap_manifests( - &self, - req: store_api::region_engine::RemapManifestsRequest, - ) -> Result - { - self.inner.remap_manifests(req).await - } - async fn set_region_role_state_gracefully( - &self, - id: RegionId, - state: store_api::region_engine::SettableRegionRoleState, - ) -> Result< - store_api::region_engine::SetRegionRoleStateResponse, - common_error::ext::BoxedError, - > { - self.inner.set_region_role_state_gracefully(id, state).await - } - fn role(&self, id: RegionId) -> Option { - self.inner.role(id) - } - fn as_any(&self) -> &dyn std::any::Any { - self - } - } - - fn body(ids: &[u64], nan: bool) -> Body { - Body::Inserts(InsertRequests { - requests: ids - .iter() - .map(|id| InsertRequest { - region_id: *id, - rows: Some(Rows { - schema: vec![ - ColumnSchema { - column_name: "ts".into(), - datatype: ColumnDataType::TimestampMillisecond as i32, - semantic_type: SemanticType::Timestamp as i32, - datatype_extension: None, - options: None, - }, - ColumnSchema { - column_name: "v".into(), - datatype: ColumnDataType::Float64 as i32, - semantic_type: SemanticType::Field as i32, - datatype_extension: None, - options: None, - }, - ], - rows: vec![Row { - values: vec![ - Value { - value_data: Some( - api::v1::value::ValueData::TimestampMillisecondValue(1), - ), - }, - Value { - value_data: Some(api::v1::value::ValueData::F64Value(if nan { - f64::NAN - } else { - 7. - })), - }, - ], - }], - }), - partition_expr_version: None, - skip_wal: false, - }) - .collect(), - }) - } - - fn request(cmd: Command, body: Option) -> Request { - let mut extensions = HashMap::new(); - extensions.insert(flow_write::COMMAND_KEY.into(), cmd.encode().unwrap()); - Request { - header: Some(RegionRequestHeader { - dbname: "db".into(), - tracing_context: Default::default(), - query_context: Some(QueryContext { - extensions, - ..Default::default() - }), - }), - body, - } - } - fn decode(resp: &api::v1::region::RegionResponse) -> Reply { - assert_eq!( - resp.header - .as_ref() - .unwrap() - .status - .as_ref() - .unwrap() - .status_code, - common_error::status_code::StatusCode::Success as u32 - ); - serde_json::from_slice(resp.extensions.get(flow_write::COMMAND_KEY).unwrap()).unwrap() - } - async fn open(s: &RegionServer) -> String { - match decode( - &s.handle_region_request(request( - Command::Open { nonce: "n".into() }, - Some(Body::Inserts(InsertRequests { requests: vec![] })), - )) - .await - .unwrap(), - ) { - Reply::Opened { capability } => capability, - x => panic!("{x:?}"), - } - } - fn setup(gate: Gate, panic_id: Option, ids: &[u64]) -> RegionServer { - let s = mock_region_server(); - let (inner, _) = MockRegionEngine::new("mock"); - let engine: RegionEngineRef = Arc::new(GatedEngine { - gate, - inner, - panic_id, - }); - for id in ids { - s.inner.region_map.insert( - RegionId::from_u64(*id), - super::super::RegionEngineWithStatus::Ready(engine.clone()), - ); - } - s - } - fn write(cap: &str, ids: &[u64], nan: bool) -> Request { - request( - Command::Write { - capability: cap.into(), - }, - Some(body(ids, nan)), - ) - } - - #[tokio::test] - async fn dropped_rpc_retry_and_close_drain_exactly_once() { - let gate = Gate::default(); - let s = setup(gate.clone(), None, &[1, 2]); - let cap = open(&s).await; - let w = write(&cap, &[1, 2], false); - let first = tokio::spawn({ - let s = s.clone(); - let w = w.clone(); - async move { s.handle_region_request(w).await } - }); - gate.wait_entered(2).await; - first.abort(); - // Exercise the same actual registry admission used by retry, before release. - let payload = tracked_payload(&w).unwrap(); - let Admission::Join(waiter) = s.flow_write_registry.admit(&cap, payload).unwrap() else { - panic!("retry did not join") - }; - let mut retry_waiter = tokio::spawn(common_grpc::flow_write::wait_completion(waiter)); - assert!( - timeout(Duration::from_millis(20), &mut retry_waiter) - .await - .is_err() - ); - let mut closer = tokio::spawn({ - let s = s.clone(); - let cap = cap.clone(); - async move { - s.handle_region_request(request( - Command::Close { capability: cap }, - Some(Body::Inserts(InsertRequests { requests: vec![] })), - )) - .await - } - }); - assert!( - timeout(Duration::from_millis(20), &mut closer) - .await - .is_err() - ); - gate.release(); - let joined = timeout(Duration::from_secs(5), &mut retry_waiter) - .await - .unwrap() - .unwrap() - .unwrap(); - assert!(matches!( - *joined, - Reply::Complete { - rows: 2, - error: None, - .. - } - )); - assert!(matches!( - decode( - &timeout(Duration::from_secs(5), &mut closer) - .await - .unwrap() - .unwrap() - .unwrap() - ), - Reply::Closed - )); - assert!( - s.handle_region_request(write(&cap, &[1], false)) - .await - .is_err() - ); - let cap2 = open(&s).await; - gate.release(); - let r = s - .handle_region_request(write(&cap2, &[1], false)) - .await - .unwrap(); - assert!(matches!(decode(&r), Reply::Complete { rows: 1, .. })); - assert!( - s.handle_region_request(write(&cap, &[1], false)) - .await - .is_err() - ); - assert_eq!(gate.entered.load(Ordering::SeqCst), 3); - let close2 = request( - Command::Close { capability: cap2 }, - Some(Body::Inserts(InsertRequests { requests: vec![] })), - ); - assert!(matches!( - decode(&s.handle_region_request(close2).await.unwrap()), - Reply::Closed - )); - } - - #[tokio::test] - async fn application_error_waits_for_sibling_and_panic_overrides_it() { - let gate = Gate { - fail: true, - ..Default::default() - }; - let s = setup(gate.clone(), None, &[1, 2]); - let cap = open(&s).await; - let task = tokio::spawn({ - let s = s.clone(); - async move { s.handle_region_request(write(&cap, &[1, 2], false)).await } - }); - gate.wait_entered(2).await; - // The ordinary error is immediate for region 1, but region 2 remains gated. - assert!(!task.is_finished()); - gate.release(); - let r = timeout(Duration::from_secs(5), task) - .await - .unwrap() - .unwrap() - .unwrap(); - assert!(matches!(decode(&r), Reply::Complete { error: Some(_), .. })); - - let gate = Gate { - fail: true, - ..Default::default() - }; - let s = setup(gate.clone(), Some(2), &[1, 2, 3]); - let cap = open(&s).await; - let task = tokio::spawn({ - let s = s.clone(); - let cap = cap.clone(); - async move { - s.handle_region_request(write(&cap, &[1, 2, 3], false)) - .await - } - }); - gate.wait_entered(3).await; - assert!(!task.is_finished()); - gate.release(); - assert!( - timeout(Duration::from_secs(5), task) - .await - .unwrap() - .unwrap() - .is_err() - ); - assert!( - s.handle_region_request(write(&cap, &[1, 2, 3], false)) - .await - .is_err() - ); - assert_eq!(gate.entered.load(Ordering::SeqCst), 3); - } - - #[tokio::test] - async fn decoded_schema_option_maps_join_live_nan_write() { - let gate = Gate::default(); - let s = setup(gate.clone(), None, &[1]); - let cap = open(&s).await; - let mut original = write(&cap, &[1], true); - let Body::Inserts(inserts) = original.body.as_mut().unwrap() else { - unreachable!() - }; - inserts.requests[0].rows.as_mut().unwrap().schema[1].options = - Some(api::v1::ColumnOptions { - options: [ - ("alpha".to_string(), "one".to_string()), - ("beta".to_string(), "two".to_string()), - ("gamma".to_string(), "three".to_string()), - ] - .into_iter() - .collect(), - }); - let encoded = original.encode_to_vec(); - let Body::Inserts(original_inserts) = original.body.as_ref().unwrap() else { - unreachable!() - }; - let original_rows = original_inserts.requests[0].rows.clone(); - let original_body_bytes = original_inserts.encode_to_vec(); - let mut retries = Vec::new(); - let mut found_different_wire_order = false; - for _ in 0..128 { - let decoded = Request::decode(encoded.as_slice()).unwrap(); - let Body::Inserts(decoded_inserts) = decoded.body.as_ref().unwrap() else { - unreachable!() - }; - assert_eq!( - decoded_inserts.requests[0].rows.as_ref().unwrap().schema, - original_rows.as_ref().unwrap().schema - ); - assert_eq!( - tracked_payload(&decoded).unwrap(), - tracked_payload(&original).unwrap() - ); - found_different_wire_order |= decoded_inserts.encode_to_vec() != original_body_bytes; - retries.push(decoded); - if found_different_wire_order && retries.len() >= 16 { - break; - } - } - assert!(found_different_wire_order); - - let owner = tokio::spawn({ - let s = s.clone(); - async move { s.handle_region_request(original).await } - }); - gate.wait_entered(1).await; - let retry_tasks: Vec<_> = retries - .into_iter() - .map(|retry| { - let s = s.clone(); - tokio::spawn(async move { s.handle_region_request(retry).await }) - }) - .collect(); - tokio::task::yield_now().await; - assert!(retry_tasks.iter().all(|retry| !retry.is_finished())); - assert_eq!(gate.entered.load(Ordering::SeqCst), 1); - gate.release(); - let owner_reply = decode( - &timeout(Duration::from_secs(5), owner) - .await - .unwrap() - .unwrap() - .unwrap(), - ); - assert!(matches!( - &owner_reply, - Reply::Complete { - rows: 1, - error: None, - .. - } - )); - let owner_reply_json = serde_json::to_string(&owner_reply).unwrap(); - for retry in retry_tasks { - let retry_reply = decode( - &timeout(Duration::from_secs(5), retry) - .await - .unwrap() - .unwrap() - .unwrap(), - ); - assert!(matches!( - &retry_reply, - Reply::Complete { - rows: 1, - error: None, - .. - } - )); - assert_eq!( - serde_json::to_string(&retry_reply).unwrap(), - owner_reply_json - ); - } - assert_eq!(gate.entered.load(Ordering::SeqCst), 1); - - let replay = Request::decode(encoded.as_slice()).unwrap(); - assert!(matches!( - decode(&s.handle_region_request(replay.clone()).await.unwrap()), - Reply::Complete { rows: 1, .. } - )); - let mut changed_schema = replay.clone(); - let Body::Inserts(inserts) = changed_schema.body.as_mut().unwrap() else { - unreachable!() - }; - inserts.requests[0].rows.as_mut().unwrap().schema[1] - .options - .as_mut() - .unwrap() - .options - .insert("gamma".into(), "changed".into()); - assert!(s.handle_region_request(changed_schema).await.is_err()); - let close = request( - Command::Close { - capability: cap.clone(), - }, - Some(Body::Inserts(InsertRequests { requests: vec![] })), - ); - assert!(matches!( - decode(&s.handle_region_request(close).await.unwrap()), - Reply::Closed - )); - assert!(s.handle_region_request(replay).await.is_err()); - assert_eq!(gate.entered.load(Ordering::SeqCst), 1); - } - - #[tokio::test] - async fn payload_identity_nan_semantics_and_pre_admission_rejection() { - let s = mock_region_server(); - let cap = open(&s).await; - let a = write(&cap, &[1], true); - let b = write(&cap, &[1], true); - assert_eq!(tracked_payload(&a).unwrap(), tracked_payload(&b).unwrap()); - let owner = match s - .flow_write_registry - .admit(&cap, tracked_payload(&a).unwrap()) - .unwrap() - { - Admission::Start(owner) => owner, - _ => panic!(), - }; - let mut semantic_mismatch = b.clone(); - semantic_mismatch.header.as_mut().unwrap().dbname = "other_db".into(); - assert_eq!( - s.flow_write_registry - .admit(&cap, tracked_payload(&semantic_mismatch).unwrap()) - .unwrap_err(), - flow_write::RegistryError::PayloadMismatch - ); - semantic_mismatch = b.clone(); - semantic_mismatch - .header - .as_mut() - .unwrap() - .query_context - .as_mut() - .unwrap() - .extensions - .insert("flow_lane".into(), "changed".into()); - assert_eq!( - s.flow_write_registry - .admit(&cap, tracked_payload(&semantic_mismatch).unwrap()) - .unwrap_err(), - flow_write::RegistryError::PayloadMismatch - ); - let Admission::Join(waiter) = s - .flow_write_registry - .admit(&cap, tracked_payload(&b).unwrap()) - .unwrap() - else { - panic!() - }; - let mut joined = tokio::spawn(common_grpc::flow_write::wait_completion(waiter)); - assert!( - timeout(Duration::from_millis(20), &mut joined) - .await - .is_err() - ); - owner - .send(Some(Arc::new(Reply::Complete { - rows: 1, - cost: 1, - metrics: None, - error: None, - }))) - .unwrap(); - assert!(matches!( - *timeout(Duration::from_secs(5), &mut joined) - .await - .unwrap() - .unwrap() - .unwrap(), - Reply::Complete { rows: 1, .. } - )); - s.flow_write_registry.close(&cap).await.unwrap(); - assert_ne!( - tracked_payload(&a).unwrap(), - tracked_payload(&write(&cap, &[1], false)).unwrap() - ); - assert!( - s.handle_region_request(request( - Command::Write { - capability: cap.clone() - }, - Some(Body::Deletes(Default::default())) - )) - .await - .is_err() - ); - let fresh = open(&s).await; - let mut bad = request( - Command::Write { - capability: fresh.clone(), - }, - Some(body(&[1], false)), - ); - bad.header - .as_mut() - .unwrap() - .query_context - .as_mut() - .unwrap() - .extensions - .insert(flow_write::COMMAND_KEY.into(), "not-json".into()); - assert!(s.handle_region_request(bad).await.is_err()); - s.flow_write_registry.close(&fresh).await.unwrap(); - } -} diff --git a/src/flow/Cargo.toml b/src/flow/Cargo.toml index 2cf316ad245..2d9ff283aa2 100644 --- a/src/flow/Cargo.toml +++ b/src/flow/Cargo.toml @@ -5,7 +5,7 @@ edition.workspace = true license.workspace = true [features] -enterprise = ["common-grpc/enterprise", "session/enterprise"] +enterprise = [] [lints] workspace = true diff --git a/src/flow/src/batching_mode/frontend_client.rs b/src/flow/src/batching_mode/frontend_client.rs index 8c74c6897f0..1a354c55e42 100644 --- a/src/flow/src/batching_mode/frontend_client.rs +++ b/src/flow/src/batching_mode/frontend_client.rs @@ -95,12 +95,8 @@ pub enum FrontendClient { control_channel_manager: ChannelManager, query: QueryOptions, batch_opts: BatchingModeOptions, - #[cfg(feature = "enterprise")] - tracked_writes: bool, }, Standalone { - #[cfg(feature = "enterprise")] - tracked_writes: bool, /// for the sake of simplicity still use grpc even in standalone mode /// notice the client here should all be lazy, so that can wait after frontend is booted then make conn database_client: HandlerMutable, @@ -109,26 +105,6 @@ pub enum FrontendClient { } impl FrontendClient { - #[cfg(feature = "enterprise")] - pub fn with_tracked_writes(&self) -> Self { - let mut client = self.clone(); - match &mut client { - Self::Distributed { tracked_writes, .. } | Self::Standalone { tracked_writes, .. } => { - *tracked_writes = true; - } - } - client - } - - #[cfg(feature = "enterprise")] - fn tracked_writes_enabled(&self) -> bool { - match self { - Self::Distributed { tracked_writes, .. } | Self::Standalone { tracked_writes, .. } => { - *tracked_writes - } - } - } - /// Create a new empty frontend client, with a `HandlerMutable` to set the grpc handler later pub fn from_empty_grpc_handler(query: QueryOptions) -> (Self, HandlerMutable) { let is_initialized = Arc::new(SetOnce::new()); @@ -140,8 +116,6 @@ impl FrontendClient { Self::Standalone { database_client: handler.clone(), query, - #[cfg(feature = "enterprise")] - tracked_writes: false, }, handler, ) @@ -178,8 +152,6 @@ impl FrontendClient { control_channel_manager, query, batch_opts, - #[cfg(feature = "enterprise")] - tracked_writes: false, }) } @@ -196,8 +168,6 @@ impl FrontendClient { Self::Standalone { database_client: handler, query, - #[cfg(feature = "enterprise")] - tracked_writes: false, } } } @@ -235,39 +205,7 @@ impl DatabaseWithPeer { } } -#[cfg(feature = "enterprise")] -pub(crate) mod flow_write; - impl FrontendClient { - /// Preflight tracked-write support by opening and closing a capability without writing. - #[cfg(feature = "enterprise")] - pub async fn check_tracked_writes(&self, catalog: &str, schema: &str) -> Result<(), Error> { - if !self.tracked_writes_enabled() { - return UnexpectedSnafu { - reason: "Tracked writes are not enabled on this frontend client", - } - .fail(); - } - flow_write::check_tracked_writes(self, catalog, schema).await - } - - /// Execute a finite SQL write with receiver-issued completion tracking. - #[cfg(feature = "enterprise")] - pub async fn tracked_sql_write( - &self, - catalog: &str, - schema: &str, - sql: &str, - ) -> Result { - if !self.tracked_writes_enabled() { - return UnexpectedSnafu { - reason: "Tracked writes are not enabled on this frontend client", - } - .fail(); - } - flow_write::tracked_sql_write(self, catalog, schema, sql).await - } - /// scan for available frontend from metadata pub(crate) async fn scan_for_frontend(&self) -> Result, Error> { let Self::Distributed { meta_client, .. } = self else { @@ -558,33 +496,6 @@ impl FrontendClient { peer_desc: &mut Option, ) -> Result { let flow_extensions = build_flow_extensions(extensions)?; - #[cfg(feature = "enterprise")] - if self.tracked_writes_enabled() && is_tracked_insert_request(&request) { - let effective_extensions: HashMap = extensions - .iter() - .map(|(key, value)| ((*key).to_string(), (*value).to_string())) - .collect(); - if effective_extensions.contains_key(common_grpc::flow_write::COMMAND_KEY) { - return UnexpectedSnafu { - reason: "Client-supplied tracked Flow command is forbidden", - } - .fail(); - } - let effective_extension_refs: Vec<(&str, &str)> = effective_extensions - .iter() - .map(|(key, value)| (key.as_str(), value.as_str())) - .collect(); - return flow_write::tracked_query( - self, - catalog, - schema, - Some(request), - &effective_extension_refs, - snapshot_seqs, - peer_desc, - ) - .await; - } match self { FrontendClient::Distributed { query, batch_opts, .. @@ -835,15 +746,6 @@ impl FrontendClient { } } -#[cfg(feature = "enterprise")] -fn is_tracked_insert_request(request: &QueryRequest) -> bool { - match request.query.as_ref() { - Some(Query::InsertIntoPlan(_)) => true, - Some(Query::Sql(sql)) => sql.trim_start().to_ascii_uppercase().starts_with("INSERT "), - _ => false, - } -} - fn build_flow_extensions(extensions: &[(&str, &str)]) -> Result { let flow_extensions = HashMap::from_iter( extensions diff --git a/src/flow/src/batching_mode/frontend_client/flow_write.rs b/src/flow/src/batching_mode/frontend_client/flow_write.rs deleted file mode 100644 index 0d7663ccd4d..00000000000 --- a/src/flow/src/batching_mode/frontend_client/flow_write.rs +++ /dev/null @@ -1,758 +0,0 @@ -// Copyright 2023-2026 GrepTime Inc. -// -// This file is part of the GreptimeDB Enterprise Edition and is licensed under -// the GreptimeDB Enterprise License. You may not use this file except in -// compliance with that license. A copy of the license is available at the root -// of this repository in the file LICENSE-ENTERPRISE. -// -// Unless required by applicable law or agreed to in writing, this software is -// distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, -// either express or implied. - -use std::collections::HashMap; -use std::sync::Arc; -use std::time::Duration; - -use api::v1::query_request::Query; -use api::v1::{QueryRequest, greptime_request::Request}; -use client::OutputWithMetrics; -use common_error::ext::{BoxedError, ErrorExt}; -use common_error::status_code::StatusCode; -use common_grpc::flow_write::{self, COMMAND_KEY, Command, Reply, TerminalError}; -use common_query::{Output, OutputData, OutputMeta}; -use common_telemetry::warn; -use session::context::{Channel, QueryContextBuilder, generate_remote_query_id}; -use snafu::OptionExt; -use tokio::time::sleep; - -use crate::Error; -use crate::batching_mode::frontend_client::{DatabaseWithPeer, FrontendClient, PeerDesc}; -use crate::error::UnexpectedSnafu; - -const RETRY_PACE: Duration = Duration::from_secs(1); - -pub(crate) async fn tracked_sql_write( - frontend: &FrontendClient, - catalog: &str, - schema: &str, - sql: &str, -) -> Result { - tracked_query( - frontend, - catalog, - schema, - Some(QueryRequest { - query: Some(Query::Sql(sql.to_string())), - }), - &[], - &HashMap::new(), - &mut None, - ) - .await -} - -pub(crate) async fn check_tracked_writes( - frontend: &FrontendClient, - catalog: &str, - schema: &str, -) -> Result<(), Error> { - tracked_query( - frontend, - catalog, - schema, - None, - &[], - &HashMap::new(), - &mut None, - ) - .await - .map(|_| ()) -} - -pub(crate) async fn tracked_query( - frontend: &FrontendClient, - catalog: &str, - schema: &str, - request: Option, - extensions: &[(&str, &str)], - snapshot_seqs: &HashMap, - peer_desc: &mut Option, -) -> Result { - let (selected, handler, parallelism, read_preference, timeout) = match frontend { - FrontendClient::Distributed { - query, batch_opts, .. - } => ( - Some( - frontend - .get_random_active_frontend(catalog, schema) - .await - .map_err(prewrite_error)?, - ), - None, - query.parallelism.to_string(), - Some(batch_opts.read_preference.as_ref().to_string()), - batch_opts.experimental_flight_do_get_timeout, - ), - FrontendClient::Standalone { - database_client, - query, - .. - } => { - let handler = database_client - .handler - .lock() - .map_err(|e| { - prewrite_reason( - format!("Failed to lock database client: {e}"), - StatusCode::Internal, - ) - })? - .as_ref() - .context(UnexpectedSnafu { - reason: "Standalone's frontend instance is not set", - }) - .map_err(prewrite_error)? - .upgrade() - .context(UnexpectedSnafu { - reason: "Failed to upgrade database client", - }) - .map_err(prewrite_error)?; - ( - None, - Some(handler), - query.parallelism.to_string(), - None, - RETRY_PACE, - ) - } - }; - *peer_desc = Some(match &selected { - Some(DatabaseWithPeer { peer, .. }) => PeerDesc::Dist { peer: peer.clone() }, - None => PeerDesc::Standalone, - }); - - let nonce = generate_remote_query_id(); - let mut command = Command::Open { nonce }; - let mut completion: Option<( - usize, - usize, - Option, - Option, - )> = None; - - loop { - let encoded = command.encode().map_err(|error| { - UnexpectedSnafu { - reason: format!("Failed to encode tracked Flow command: {error}"), - } - .build() - })?; - let mut operation_extensions: Vec<(String, String)> = extensions - .iter() - .filter(|(key, _)| *key != COMMAND_KEY) - .map(|(key, value)| ((*key).to_string(), (*value).to_string())) - .collect(); - operation_extensions.push((COMMAND_KEY.to_string(), encoded)); - let extension_refs: Vec<(&str, &str)> = operation_extensions - .iter() - .map(|(k, v)| (k.as_str(), v.as_str())) - .collect(); - let phase_request = match &command { - Command::Write { .. } => request - .as_ref() - .context(UnexpectedSnafu { - reason: "Tracked Flow write command has no request", - })? - .clone(), - Command::Open { .. } | Command::Close { .. } => QueryRequest { - query: Some(Query::Sql(String::new())), - }, - }; - let response = if let Some(DatabaseWithPeer { database, .. }) = &selected { - let mut hints = vec![( - query::datafusion::QUERY_PARALLELISM_HINT, - parallelism.as_str(), - )]; - if let Some(preference) = read_preference.as_deref() { - hints.push((session::hints::READ_PREFERENCE_HINT, preference)); - } - database - .flight_request() - .with_hints(&hints) - .with_flow_extensions(&extension_refs) - .with_snapshot_seqs(snapshot_seqs) - .with_timeout(timeout) - .query_with_terminal_metrics(phase_request) - .await - .map(|output| output.output) - .map_err(|e| (e.status_code(), e.to_string())) - } else { - let mut ctx_extensions: HashMap = - operation_extensions.into_iter().collect(); - ctx_extensions.insert( - query::datafusion::QUERY_PARALLELISM_HINT.to_string(), - parallelism.clone(), - ); - let ctx = QueryContextBuilder::default() - .current_catalog(catalog.to_string()) - .current_schema(schema.to_string()) - .channel(Channel::Internal) - .extensions(ctx_extensions) - .snapshot_seqs(Arc::new(std::sync::RwLock::new(snapshot_seqs.clone()))) - .build(); - let handler = handler.as_ref().ok_or_else(|| { - UnexpectedSnafu { - reason: "Standalone handler was not pinned", - } - .build() - })?; - handler - .do_query(Request::Query(phase_request), Arc::new(ctx)) - .await - .map_err(|e| (e.status_code(), e.to_string())) - }; - let reply = match response { - Ok(output) => match decode_output(output).await { - Ok(reply) => Some(reply), - Err(e) if matches!(&command, Command::Open { .. }) => { - return Err(prewrite_reason( - format!("Frontend does not support tracked Flow writes: {e}"), - StatusCode::Unsupported, - )); - } - Err(e) => { - warn!("Malformed tracked Flow {:?} reply: {}", command, e); - None - } - }, - Err((code, msg)) => { - // Only Open is pre-admission; definitive rejection there is safe to surface. - if matches!(&command, Command::Open { .. }) && is_definitive_rejection(code) { - return Err(prewrite_reason( - format!("Frontend rejected tracked Flow open: {msg}"), - code, - )); - } - warn!( - "Tracked Flow {:?} transport failure (code={:?}): {}", - command, code, msg - ); - None - } - }; - match (command.clone(), reply) { - (Command::Open { .. }, Some(Reply::Opened { capability: cap })) if !cap.is_empty() => { - command = if request.is_some() { - Command::Write { capability: cap } - } else { - Command::Close { capability: cap } - }; - continue; - } - (Command::Open { .. }, Some(other)) => { - return Err(prewrite_reason( - format!("Invalid tracked Flow open reply: {other:?}"), - StatusCode::InvalidArguments, - )); - } - ( - Command::Write { capability }, - Some(Reply::Complete { - rows, - cost, - metrics, - error, - }), - ) => { - completion = Some((rows, cost, metrics, error)); - command = Command::Close { capability }; - continue; - } - (Command::Write { .. }, Some(other)) => { - warn!("Unexpected tracked Flow write reply: {other:?}") - } - (Command::Close { .. }, Some(Reply::Closed)) => { - let Some((rows, cost, metrics, error)) = completion.take() else { - if request.is_none() { - return Ok(OutputWithMetrics::from_output(Output::new( - OutputData::AffectedRows(0), - OutputMeta::default(), - ))); - } - return UnexpectedSnafu { - reason: "Close completed without positive write result", - } - .fail(); - }; - if let Some(error) = error { - return Err(Error::External { - source: BoxedError::new(error), - location: snafu::Location::new(file!(), line!(), column!()), - }); - } - let output = OutputWithMetrics::from_output(Output::new( - OutputData::AffectedRows(rows), - OutputMeta { - cost, - ..Default::default() - }, - )); - output.metrics.update(metrics); - return Ok(output); - } - (Command::Close { .. }, Some(other)) => { - warn!("Unexpected tracked Flow close reply: {other:?}") - } - (_, None) => {} - } - sleep(RETRY_PACE).await; - } -} - -fn prewrite_error(error: Error) -> Error { - let status = error.status_code(); - prewrite_reason(error.to_string(), status) -} - -/// Returns true when a tracked Flow command error is a definitive, non-retryable -/// rejection from the frontend rather than an ambiguous transport failure. These -/// codes indicate the frontend itself refused the command (unsupported protocol, -/// malformed/empty command, or wrong channel), so retrying the same pinned -/// frontend can never succeed; the caller must fail fast instead of looping. -fn is_definitive_rejection(code: StatusCode) -> bool { - matches!( - code, - StatusCode::Unsupported | StatusCode::InvalidArguments | StatusCode::PermissionDenied - ) -} - -fn prewrite_reason(message: String, status: StatusCode) -> Error { - Error::External { - source: BoxedError::new(TerminalError { - code: status as u32, - message, - }), - location: snafu::Location::new(file!(), line!(), column!()), - } -} - -async fn decode_output(output: Output) -> Result { - let batches = match output.data { - OutputData::RecordBatches(batches) => batches, - OutputData::Stream(stream) => common_recordbatch::util::collect_batches(stream) - .await - .map_err(|e| e.to_string())?, - _ => return Err("tracked control request returned a non-record-batch result".to_string()), - }; - flow_write::decode_reply(&batches).map_err(|e| e.to_string()) -} - -#[cfg(test)] -mod tests { - use super::*; - use std::sync::atomic::{AtomicUsize, Ordering}; - - use common_error::ext::{BoxedError, PlainError}; - use common_error::status_code::StatusCode; - use session::context::QueryContextRef; - - #[derive(Debug, Default)] - struct LostReplyHandler { - opens: AtomicUsize, - writes: AtomicUsize, - closes: AtomicUsize, - executions: AtomicUsize, - registry: common_grpc::flow_write::FlowWriteRegistry, - capabilities: std::sync::Mutex>, - } - - #[async_trait::async_trait] - impl super::super::GrpcQueryHandlerWithBoxedError for LostReplyHandler { - async fn do_query( - &self, - request: Request, - ctx: QueryContextRef, - ) -> Result { - let encoded = ctx.extensions().get(COMMAND_KEY).cloned().ok_or_else(|| { - BoxedError::new(PlainError::new( - "missing command".into(), - StatusCode::Unexpected, - )) - })?; - match Command::decode(&encoded).map_err(|error| { - BoxedError::new(PlainError::new( - error.to_string(), - StatusCode::InvalidArguments, - )) - })? { - Command::Open { nonce } => { - let call = self.opens.fetch_add(1, Ordering::SeqCst); - let capability = self.registry.open(&nonce).unwrap(); - let previous = self - .capabilities - .lock() - .unwrap() - .insert(nonce, capability.clone()); - if let Some(previous) = previous { - assert_eq!(previous, capability); - } - if call == 0 { - return Err(transport_error("lost Open reply")); - } - assert_empty_sql(&request); - response(Reply::Opened { capability }) - } - Command::Write { capability } => { - let sql = request_sql(&request); - let call = self.writes.fetch_add(1, Ordering::SeqCst); - let admission = self.registry.admit(&capability, sql.clone()).unwrap(); - let result = match admission { - common_grpc::flow_write::Admission::Start(sender) => { - self.executions.fetch_add(1, Ordering::SeqCst); - let reply = Reply::Complete { - rows: 1, - cost: 9, - metrics: Some(common_recordbatch::adapter::RecordBatchMetrics { - region_watermarks: vec![ - common_recordbatch::adapter::RegionWatermarkEntry { - region_id: 41, - watermark: Some(17), - }, - ], - ..Default::default() - }), - error: None, - }; - sender.send(Some(Arc::new(reply.clone()))).unwrap(); - reply - } - common_grpc::flow_write::Admission::Join(receiver) => { - common_grpc::flow_write::wait_completion(receiver) - .await - .unwrap() - .as_ref() - .clone() - } - }; - if call == 0 { - return Err(transport_error("lost Write reply")); - } - response(result) - } - Command::Close { capability } => { - let call = self.closes.fetch_add(1, Ordering::SeqCst); - self.registry.close(&capability).await.unwrap(); - if call == 0 { - return Err(transport_error("lost Close reply")); - } - response(Reply::Closed) - } - } - } - } - - fn request_sql(request: &Request) -> String { - match request { - Request::Query(query) => match query.query.as_ref() { - Some(Query::Sql(sql)) => sql.clone(), - other => panic!("expected SQL INSERT, got {other:?}"), - }, - other => panic!("expected Query request, got {other:?}"), - } - } - - fn assert_empty_sql(request: &Request) { - assert!( - matches!(request, Request::Query(QueryRequest { query: Some(Query::Sql(sql)) }) if sql.is_empty()) - ); - } - - fn transport_error(message: &str) -> BoxedError { - BoxedError::new(PlainError::new( - message.to_string(), - StatusCode::StorageUnavailable, - )) - } - - fn response(reply: Reply) -> Result { - let batches = flow_write::reply_batches(&reply).map_err(BoxedError::new)?; - Ok(Output::new_with_record_batches(batches)) - } - - #[tokio::test] - async fn retries_same_open_write_and_close_capability_and_executes_once() { - let handler = Arc::new(LostReplyHandler::default()); - let trait_handler: Arc = handler.clone(); - let weak = Arc::downgrade(&trait_handler); - let frontend = - FrontendClient::from_grpc_handler(weak, query::options::QueryOptions::default()) - .with_tracked_writes(); - let output = frontend - .tracked_sql_write("greptime", "public", "INSERT INTO t VALUES (7)") - .await - .unwrap(); - assert!(matches!(output.output.data, OutputData::AffectedRows(1))); - assert_eq!(output.output.meta.cost, 9); - assert_eq!( - output.metrics.region_watermark_map(), - Some(HashMap::from([(41, 17)])) - ); - assert_eq!(handler.opens.load(Ordering::SeqCst), 2); - assert_eq!(handler.writes.load(Ordering::SeqCst), 2); - assert_eq!(handler.closes.load(Ordering::SeqCst), 2); - assert_eq!(handler.executions.load(Ordering::SeqCst), 1); - let second = frontend - .tracked_sql_write("greptime", "public", "INSERT INTO t VALUES (8)") - .await - .unwrap(); - assert!(matches!(second.output.data, OutputData::AffectedRows(1))); - assert_eq!(handler.executions.load(Ordering::SeqCst), 2); - } - - /// A legacy frontend that does not understand the tracked-write extension - /// processes `Open` as an ordinary empty-SQL query and rejects it with - /// `InvalidArguments`. That is a definitive application rejection, not an - /// ambiguous transport failure: the client must fail fast instead of - /// retrying the same pinned frontend forever. - #[derive(Debug, Default)] - struct LegacyRejectHandler { - opens: AtomicUsize, - } - - #[async_trait::async_trait] - impl super::super::GrpcQueryHandlerWithBoxedError for LegacyRejectHandler { - async fn do_query( - &self, - _request: Request, - _ctx: QueryContextRef, - ) -> Result { - self.opens.fetch_add(1, Ordering::SeqCst); - Err(BoxedError::new(PlainError::new( - "empty statements".to_string(), - StatusCode::InvalidArguments, - ))) - } - } - - #[derive(Debug)] - struct PreflightHandler { - opens: AtomicUsize, - closes: AtomicUsize, - reject_open: bool, - } - - #[async_trait::async_trait] - impl super::super::GrpcQueryHandlerWithBoxedError for PreflightHandler { - async fn do_query( - &self, - request: Request, - ctx: QueryContextRef, - ) -> Result { - let encoded = ctx.extensions().get(COMMAND_KEY).cloned().unwrap(); - assert_empty_sql(&request); - match Command::decode(&encoded).unwrap() { - Command::Open { .. } => { - self.opens.fetch_add(1, Ordering::SeqCst); - if self.reject_open { - Err(BoxedError::new(PlainError::new( - "preflight rejected".into(), - StatusCode::PermissionDenied, - ))) - } else { - response(Reply::Opened { - capability: "preflight-capability".into(), - }) - } - } - Command::Close { capability } => { - assert_eq!(capability, "preflight-capability"); - self.closes.fetch_add(1, Ordering::SeqCst); - response(Reply::Closed) - } - Command::Write { .. } => panic!("preflight must not write"), - } - } - } - - #[tokio::test] - async fn preflight_opens_and_closes_without_writing() { - let handler = Arc::new(PreflightHandler { - opens: AtomicUsize::default(), - closes: AtomicUsize::default(), - reject_open: false, - }); - let trait_handler: Arc = handler.clone(); - let frontend = FrontendClient::from_grpc_handler( - Arc::downgrade(&trait_handler), - query::options::QueryOptions::default(), - ) - .with_tracked_writes(); - tokio::time::timeout( - Duration::from_secs(10), - frontend.check_tracked_writes("greptime", "public"), - ) - .await - .expect("tracked-write preflight should finish within 10 seconds") - .unwrap(); - assert_eq!(handler.opens.load(Ordering::SeqCst), 1); - assert_eq!(handler.closes.load(Ordering::SeqCst), 1); - } - - #[tokio::test] - async fn preflight_preserves_definitive_open_error() { - let handler = Arc::new(PreflightHandler { - opens: AtomicUsize::default(), - closes: AtomicUsize::default(), - reject_open: true, - }); - let trait_handler: Arc = handler.clone(); - let frontend = FrontendClient::from_grpc_handler( - Arc::downgrade(&trait_handler), - query::options::QueryOptions::default(), - ) - .with_tracked_writes(); - let error = tokio::time::timeout( - Duration::from_secs(10), - frontend.check_tracked_writes("greptime", "public"), - ) - .await - .expect("tracked-write preflight rejection should finish within 10 seconds") - .unwrap_err(); - assert_eq!(error.status_code(), StatusCode::PermissionDenied); - assert!(matches!( - &error, - Error::External { source, .. } - if source.as_any().downcast_ref::().is_some_and(|terminal| { - terminal.code == StatusCode::PermissionDenied as u32 - && terminal.message.contains("preflight rejected") - }) - )); - assert_eq!(handler.opens.load(Ordering::SeqCst), 1); - assert_eq!(handler.closes.load(Ordering::SeqCst), 0); - } - - #[derive(Debug)] - struct RejectionRetryHandler { - status: StatusCode, - opens: AtomicUsize, - writes: AtomicUsize, - closes: AtomicUsize, - capability: std::sync::Mutex>, - } - - #[async_trait::async_trait] - impl super::super::GrpcQueryHandlerWithBoxedError for RejectionRetryHandler { - async fn do_query( - &self, - request: Request, - ctx: QueryContextRef, - ) -> Result { - let encoded = ctx.extensions().get(COMMAND_KEY).cloned().unwrap(); - match Command::decode(&encoded).unwrap() { - Command::Open { nonce } => { - assert_empty_sql(&request); - self.opens.fetch_add(1, Ordering::SeqCst); - let mut capability = self.capability.lock().unwrap(); - let value = capability.get_or_insert_with(|| nonce); - response(Reply::Opened { - capability: value.clone(), - }) - } - Command::Write { capability } => { - assert_eq!(request_sql(&request), "INSERT INTO t VALUES (11)"); - assert_eq!(self.capability.lock().unwrap().as_ref(), Some(&capability)); - if self.writes.fetch_add(1, Ordering::SeqCst) == 0 { - Err(BoxedError::new(PlainError::new( - "write rejection".into(), - self.status, - ))) - } else { - response(Reply::Complete { - rows: 2, - cost: 0, - metrics: None, - error: None, - }) - } - } - Command::Close { capability } => { - assert_empty_sql(&request); - assert_eq!(self.capability.lock().unwrap().as_ref(), Some(&capability)); - if self.closes.fetch_add(1, Ordering::SeqCst) == 0 { - Err(BoxedError::new(PlainError::new( - "close rejection".into(), - self.status, - ))) - } else { - response(Reply::Closed) - } - } - } - } - } - - #[tokio::test] - async fn post_admission_rejections_retry_same_capability_until_close() { - for status in [ - StatusCode::Unsupported, - StatusCode::InvalidArguments, - StatusCode::PermissionDenied, - ] { - let handler = Arc::new(RejectionRetryHandler { - status, - opens: AtomicUsize::default(), - writes: AtomicUsize::default(), - closes: AtomicUsize::default(), - capability: std::sync::Mutex::new(None), - }); - let trait_handler: Arc = - handler.clone(); - let frontend = FrontendClient::from_grpc_handler( - Arc::downgrade(&trait_handler), - query::options::QueryOptions::default(), - ) - .with_tracked_writes(); - let output = tokio::time::timeout( - Duration::from_secs(10), - frontend.tracked_sql_write("greptime", "public", "INSERT INTO t VALUES (11)"), - ) - .await - .expect("post-admission retries should finish within 10 seconds") - .unwrap(); - assert!(matches!(output.output.data, OutputData::AffectedRows(2))); - assert_eq!(handler.opens.load(Ordering::SeqCst), 1); - assert_eq!(handler.writes.load(Ordering::SeqCst), 2); - assert_eq!(handler.closes.load(Ordering::SeqCst), 2); - } - } - - #[tokio::test] - async fn fails_fast_on_definitive_open_rejection() { - let handler = Arc::new(LegacyRejectHandler::default()); - let trait_handler: Arc = handler.clone(); - let weak = Arc::downgrade(&trait_handler); - let frontend = - FrontendClient::from_grpc_handler(weak, query::options::QueryOptions::default()) - .with_tracked_writes(); - let result = tokio::time::timeout( - Duration::from_secs(10), - frontend.tracked_sql_write("greptime", "public", "INSERT INTO t VALUES (7)"), - ) - .await - .expect("definitive Open rejection should finish within 10 seconds"); - let err = result.expect_err("definitive Open rejection must not be retried forever"); - assert_eq!(err.status_code(), StatusCode::InvalidArguments); - // The durable-execution layer keys repair off a TerminalError inside - // Error::External; pin that exact shape so a refactor cannot silently - // downgrade the rejection to an unresolved (permanently stalled) error. - assert!( - matches!(&err, Error::External { source, .. } if source.as_any().is::()), - "expected Error::External carrying TerminalError, got {err:?}" - ); - assert_eq!( - handler.opens.load(Ordering::SeqCst), - 1, - "Open must not be retried after a definitive rejection" - ); - } -} diff --git a/src/frontend/Cargo.toml b/src/frontend/Cargo.toml index 1fd347c68e9..0e4a15e5d97 100644 --- a/src/frontend/Cargo.toml +++ b/src/frontend/Cargo.toml @@ -7,8 +7,6 @@ license.workspace = true [features] testing = [] enterprise = [ - "common-grpc/enterprise", - "session/enterprise", "common-meta/enterprise", "operator/enterprise", "sql/enterprise", diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index 67f1e6003a0..554a9677af7 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -17,8 +17,6 @@ mod copy_in; mod dashboard; mod entity_graph; mod export_database; -#[cfg(feature = "enterprise")] -mod flow_write; mod grpc; mod import_packed; mod influxdb; @@ -136,13 +134,6 @@ pub struct Instance { trace_aux_limiter: Option, otlp_resource_info: bool, suspend: Arc, - #[cfg(feature = "enterprise")] - flow_write_registry: Arc< - common_grpc::flow_write::FlowWriteRegistry< - flow_write::WritePayload, - common_grpc::flow_write::Reply, - >, - >, } impl Instance { @@ -332,17 +323,6 @@ impl Instance { timeout_recorder, ); - #[cfg(feature = "enterprise")] - let result = if query_ctx.channel() == Channel::Internal - && query_ctx.extension(common_grpc::flow_write::CHILD_KEY) == Some("1") - { - query_fut.await - } else { - CancellableFuture::new(query_fut, ticket.cancellation_handle.clone()) - .await - .map_err(|_| error::CancelledSnafu.build())? - }; - #[cfg(not(feature = "enterprise"))] let result = CancellableFuture::new(query_fut, ticket.cancellation_handle.clone()) .await .map_err(|_| error::CancelledSnafu.build())?; @@ -913,17 +893,6 @@ impl Instance { let query_fut = self.exec_plan_with_timeout(plan, query_ctx.clone(), timeout_recorder); - #[cfg(feature = "enterprise")] - let result = if query_ctx.channel() == Channel::Internal - && query_ctx.extension(common_grpc::flow_write::CHILD_KEY) == Some("1") - { - query_fut.await - } else { - CancellableFuture::new(query_fut, ticket.cancellation_handle.clone()) - .await - .map_err(|_| error::CancelledSnafu.build())? - }; - #[cfg(not(feature = "enterprise"))] let result = CancellableFuture::new(query_fut, ticket.cancellation_handle.clone()) .await .map_err(|_| error::CancelledSnafu.build())?; diff --git a/src/frontend/src/instance/builder.rs b/src/frontend/src/instance/builder.rs index 732c587e07b..05efd38d28a 100644 --- a/src/frontend/src/instance/builder.rs +++ b/src/frontend/src/instance/builder.rs @@ -438,8 +438,6 @@ impl FrontendBuilder { trace_aux_limiter, otlp_resource_info: self.options.otlp.experimental_enable_resource_info, suspend: Arc::new(AtomicBool::new(false)), - #[cfg(feature = "enterprise")] - flow_write_registry: Default::default(), }) } } diff --git a/src/frontend/src/instance/flow_write.rs b/src/frontend/src/instance/flow_write.rs deleted file mode 100644 index f908d438d4e..00000000000 --- a/src/frontend/src/instance/flow_write.rs +++ /dev/null @@ -1,227 +0,0 @@ -// Copyright 2023-2026 GrepTime Inc. -// -// This file is part of the GreptimeDB Enterprise Edition and is licensed under -// the GreptimeDB Enterprise License. You may not use this file except in -// compliance with that license. A copy of the license is available at the root -// of this repository in the file LICENSE-ENTERPRISE. -// -// Unless required by applicable law or agreed to in writing, this software is -// distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, -// either express or implied. - -use std::sync::Arc; - -use api::v1::{greptime_request::Request, query_request::Query}; -use common_error::ext::ErrorExt; -use common_grpc::flow_write::{ - Admission, CHILD_KEY, COMMAND_KEY, Command, Reply, TerminalError, wait_completion, -}; -use common_query::{Output, OutputData}; -use servers::error::{self as server_error, Error as ServerError}; -use servers::interceptor::GrpcQueryInterceptor; -use servers::query_handler::grpc::GrpcQueryHandler; -use session::context::QueryContextRef; - -use crate::instance::Instance; - -#[derive(Clone, Debug, PartialEq)] -pub(crate) struct WritePayload { - request: Request, - catalog: String, - schema: String, - snapshots: Vec<(u64, u64)>, - sst_min_sequences: Vec<(u64, u64)>, - extensions: Vec<(String, String)>, -} - -impl WritePayload { - fn new(request: &Request, ctx: &QueryContextRef) -> Result { - if !is_finite_insert(request) { - return Err("tracked flow writes require INSERT or InsertIntoPlan".to_string()); - } - let mut snapshots: Vec<_> = ctx.snapshots().into_iter().collect(); - snapshots.sort_unstable(); - let mut sst_min_sequences: Vec<_> = ctx.sst_min_sequences().into_iter().collect(); - sst_min_sequences.sort_unstable(); - let mut extensions: Vec<_> = ctx - .extensions() - .into_iter() - .filter(|(key, _)| { - key != COMMAND_KEY - && key != CHILD_KEY - && key != session::context::REMOTE_QUERY_ID_EXTENSION_KEY - }) - .collect(); - extensions.sort_unstable(); - Ok(Self { - request: request.clone(), - catalog: ctx.current_catalog().to_string(), - schema: ctx.current_schema(), - snapshots, - sst_min_sequences, - extensions, - }) - } -} - -fn is_finite_insert(request: &Request) -> bool { - match request { - Request::Query(query) => match query.query.as_ref() { - Some(Query::InsertIntoPlan(_)) => true, - Some(Query::Sql(sql)) => sql.split_whitespace().next().is_some_and(|statement| { - statement - .trim_end_matches(';') - .eq_ignore_ascii_case("INSERT") - }), - _ => false, - }, - _ => false, - } -} - -impl Instance { - fn intercept_control( - &self, - request: &Request, - ctx: &QueryContextRef, - ) -> server_error::Result<()> { - let interceptor = self - .plugins - .get::>(); - interceptor - .as_ref() - .pre_execute(request, ctx.clone()) - .map_err(|error| ServerError::ExecuteQuery { - source: common_error::ext::BoxedError::new(error), - location: snafu::Location::new(file!(), line!(), column!()), - }) - } - - pub(crate) async fn handle_flow_write_request( - &self, - request: Request, - ctx: QueryContextRef, - encoded_command: String, - ) -> server_error::Result { - let command = Command::decode(&encoded_command).map_err(|e| ServerError::InvalidQuery { - reason: e.to_string(), - location: snafu::Location::new(file!(), line!(), column!()), - })?; - let reply = match command { - Command::Open { nonce } => { - self.intercept_control(&request, &ctx)?; - if !is_empty_control(&request) { - return Err(ServerError::InvalidQuery { - reason: "Open requires an empty SQL request".into(), - location: snafu::Location::new(file!(), line!(), column!()), - }); - } - Reply::Opened { - capability: self - .flow_write_registry - .open(&nonce) - .map_err(registry_error)?, - } - } - Command::Close { capability } => { - self.intercept_control(&request, &ctx)?; - if !is_empty_control(&request) { - return Err(ServerError::InvalidQuery { - reason: "Close requires an empty SQL request".into(), - location: snafu::Location::new(file!(), line!(), column!()), - }); - } - self.flow_write_registry - .close(&capability) - .await - .map_err(registry_error)?; - Reply::Closed - } - Command::Write { capability } => { - let payload = WritePayload::new(&request, &ctx).map_err(|reason| { - ServerError::InvalidQuery { - reason, - location: snafu::Location::new(file!(), line!(), column!()), - } - })?; - match self - .flow_write_registry - .admit(&capability, payload) - .map_err(registry_error)? - { - Admission::Join(receiver) => { - (*wait_completion(receiver).await.map_err(registry_error)?).clone() - } - Admission::Start(sender) => { - let receiver = sender.subscribe(); - let instance = self.clone(); - let mut child_ctx = (*ctx).clone(); - child_ctx.remove_extension(COMMAND_KEY); - child_ctx.set_extension(CHILD_KEY, "1"); - let child_ctx = Arc::new(child_ctx); - common_runtime::spawn_global(async move { - let terminal = match instance.do_query(request, child_ctx).await { - Ok(output) => terminal_reply(output), - Err(error) => Reply::Complete { - rows: 0, - cost: 0, - metrics: None, - error: Some(TerminalError { - code: error.status_code() as u32, - message: error.output_msg(), - }), - }, - }; - let _ = sender.send(Some(Arc::new(terminal))); - }); - (*wait_completion(receiver).await.map_err(registry_error)?).clone() - } - } - } - }; - let batches = common_grpc::flow_write::reply_batches(&reply).map_err(|e| { - ServerError::InvalidQuery { - reason: e.to_string(), - location: snafu::Location::new(file!(), line!(), column!()), - } - })?; - Ok(Output::new_with_record_batches(batches)) - } -} - -fn is_empty_control(request: &Request) -> bool { - matches!(request, Request::Query(query) if matches!(&query.query, Some(Query::Sql(sql)) if sql.trim().is_empty())) -} - -fn terminal_reply(output: Output) -> Reply { - let (rows, cost) = output.extract_rows_and_cost(); - let metrics = output - .meta - .plan - .clone() - .and_then(query::metrics::terminal_recordbatch_metrics_from_plan); - match output.data { - OutputData::AffectedRows(_) => Reply::Complete { - rows, - cost, - metrics, - error: None, - }, - _ => Reply::Complete { - rows: 0, - cost, - metrics, - error: Some(TerminalError { - code: 1001, - message: "tracked INSERT did not produce finite affected rows".into(), - }), - }, - } -} - -fn registry_error(error: common_grpc::flow_write::RegistryError) -> ServerError { - ServerError::InvalidQuery { - reason: format!("flow write capability rejected: {error:?}"), - location: snafu::Location::new(file!(), line!(), column!()), - } -} diff --git a/src/frontend/src/instance/grpc.rs b/src/frontend/src/instance/grpc.rs index 3fb8a6feef8..7dfd97cd69d 100644 --- a/src/frontend/src/instance/grpc.rs +++ b/src/frontend/src/instance/grpc.rs @@ -66,20 +66,6 @@ impl GrpcQueryHandler for Instance { request: Request, ctx: QueryContextRef, ) -> server_error::Result { - #[cfg(feature = "enterprise")] - if let Some(command) = ctx - .extension(common_grpc::flow_write::COMMAND_KEY) - .map(str::to_string) - { - if ctx.channel() != session::context::Channel::Internal { - return Err(server_error::Error::InvalidQuery { - reason: "flow write commands are only accepted on the internal channel".into(), - location: snafu::Location::new(file!(), line!(), column!()), - }); - } - return self.handle_flow_write_request(request, ctx, command).await; - } - let result: Result = async { let interceptor_ref = self.plugins.get::>(); let interceptor = interceptor_ref.as_ref(); diff --git a/src/operator/Cargo.toml b/src/operator/Cargo.toml index 225f1bb6dec..b89bf4b13d3 100644 --- a/src/operator/Cargo.toml +++ b/src/operator/Cargo.toml @@ -7,7 +7,6 @@ license.workspace = true [features] testing = [] enterprise = [ - "common-grpc/enterprise", "common-meta/enterprise", "sql/enterprise", "query/enterprise", diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index 9a0cd64a828..88cb565c779 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -#[cfg(feature = "enterprise")] -mod flow_write; mod otlp; pub use otlp::MetricInsertPreflight; @@ -857,47 +855,6 @@ impl Inserter { has_instant_rows, )?; - // Tracked Flow children use a same-capability write protocol. Keep the - // ordinary dispatch path unchanged for every other request. - #[cfg(feature = "enterprise")] - if flow_write::is_child(ctx) { - let dispatcher = self.write_dispatcher.clone(); - let request = request_factory.build_insert(normal_requests); - let writes = dispatcher.prepare(vec![request]).await?; - let tracked_tasks = writes.into_iter().map(|(peer, mut request)| { - let mut query_context = crate::utils::to_meta_query_context(ctx.clone()); - query_context.extensions.insert( - common_grpc::flow_write::CHILD_KEY.to_string(), - "1".to_string(), - ); - if let Some(header) = request.header.as_mut() { - header.query_context = Some(query_context.into()); - } - let node_manager = self.node_manager.clone(); - async move { flow_write::dispatch(&node_manager, &peer, request).await } - }); - let results = future::join_all(tracked_tasks).await; - let mut affected_rows = 0; - let mut first_error = None; - for result in results { - match result { - Ok(response) => affected_rows += response.affected_rows, - Err(error) if first_error.is_none() => first_error = Some(error), - Err(_) => {} - } - } - if let Some(error) = first_error { - return Err(error); - } - crate::metrics::DIST_INGEST_ROW_COUNT - .with_label_values(&[ctx.get_db_string().as_str()]) - .inc_by(affected_rows as u64); - return Ok(Output::new( - OutputData::AffectedRows(affected_rows), - OutputMeta::new_with_cost(write_cost as _), - )); - } - let dispatcher = self.write_dispatcher.clone(); let request = request_factory.build_insert(normal_requests); let writes = dispatcher.prepare(vec![request]).await?; diff --git a/src/operator/src/insert/flow_write.rs b/src/operator/src/insert/flow_write.rs deleted file mode 100644 index 62935e46981..00000000000 --- a/src/operator/src/insert/flow_write.rs +++ /dev/null @@ -1,391 +0,0 @@ -// Copyright 2023-2026 GrepTime Inc. -// -// This file is part of the GreptimeDB Enterprise Edition and is licensed under -// the GreptimeDB Enterprise License. You may not use this file except in -// compliance with that license. A copy of the license is available at the root -// of this repository in the file LICENSE-ENTERPRISE. -// -// Unless required by applicable law or agreed to in writing, this software is -// distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, -// either express or implied. - -use std::time::Duration; - -use api::v1::region::RegionRequest; -use common_error::ext::BoxedError; -use common_grpc::flow_write::{self, Command, Reply}; -use common_meta::node_manager::{DatanodeRef, NodeManagerRef}; -use common_meta::peer::Peer; -use session::context::generate_remote_query_id; - -use crate::error::{Error, Result}; - -const RETRY_DELAY: Duration = Duration::from_secs(1); - -fn error(message: impl Into) -> Error { - Error::Unexpected { - violated: message.into(), - location: snafu::Location::new(file!(), line!(), column!()), - } -} - -fn command_request( - request: &RegionRequest, - command: &Command, - control: bool, -) -> Result { - let mut request = request.clone(); - if control { - request.body = Some(api::v1::region::region_request::Body::Inserts( - Default::default(), - )); - } - let header = request.header.get_or_insert_with(Default::default); - let query_context = header - .query_context - .get_or_insert_with(api::v1::QueryContext::default); - query_context.extensions.insert( - flow_write::COMMAND_KEY.to_string(), - command.encode().map_err(|e| error(e.to_string()))?, - ); - Ok(request) -} - -fn reply(response: &api::region::RegionResponse) -> Option { - response - .extensions - .get(flow_write::COMMAND_KEY) - .and_then(|value| std::str::from_utf8(value).ok()) - .and_then(|value| serde_json::from_str(value).ok()) -} - -pub(crate) async fn dispatch( - node_manager: &NodeManagerRef, - peer: &Peer, - request: RegionRequest, -) -> Result { - dispatch_datanode(node_manager.datanode(peer).await, request).await -} - -async fn dispatch_datanode( - datanode: DatanodeRef, - request: RegionRequest, -) -> Result { - let nonce = generate_remote_query_id(); - let open = command_request(&request, &Command::Open { nonce }, true)?; - let capability = loop { - match datanode.handle(open.clone()).await { - Ok(response) => match reply(&response) { - Some(Reply::Opened { capability }) if !capability.is_empty() => break capability, - _ => return Err(error("Datanode does not support tracked Flow writes")), - }, - Err(_) => tokio::time::sleep(RETRY_DELAY).await, - } - }; - - let write = command_request( - &request, - &Command::Write { - capability: capability.clone(), - }, - false, - )?; - let complete = loop { - match datanode.handle(write.clone()).await { - Ok(response) => match reply(&response) { - Some(Reply::Complete { - rows, - cost: _, - metrics: _, - error, - }) => break (response, rows, error), - _ => tokio::time::sleep(RETRY_DELAY).await, - }, - Err(_) => tokio::time::sleep(RETRY_DELAY).await, - } - }; - - let close = command_request( - &request, - &Command::Close { - capability: capability.clone(), - }, - true, - )?; - loop { - match datanode.handle(close.clone()).await { - Ok(response) if matches!(reply(&response), Some(Reply::Closed)) => break, - _ => tokio::time::sleep(RETRY_DELAY).await, - } - } - - let (mut response, rows, terminal_error) = complete; - if let Some(terminal_error) = terminal_error { - return Err(Error::External { - source: BoxedError::new(terminal_error), - location: snafu::location!(), - }); - } - response.affected_rows = rows as _; - Ok(response) -} - -pub(crate) fn is_child(ctx: &session::context::QueryContextRef) -> bool { - ctx.channel() == common_base::protocol::Channel::Internal - && ctx.extension(flow_write::CHILD_KEY) == Some("1") -} - -#[cfg(test)] -mod tests { - use std::collections::HashMap; - use std::sync::atomic::{AtomicUsize, Ordering}; - use std::sync::{Arc, Mutex}; - - use api::region::RegionResponse; - use api::v1::region::{RegionRequest, region_request::Body}; - use common_error::ext::BoxedError; - use common_error::ext::ErrorExt; - use common_grpc::flow_write::{self, Command, Reply}; - use common_meta::node_manager::{Datanode, DatanodeRef}; - use common_query::request::QueryRequest; - use common_recordbatch::SendableRecordBatchStream; - use futures_util::future::join_all; - use tokio::sync::Notify; - - use super::dispatch_datanode; - - #[derive(Default)] - struct MockDatanode { - capabilities: Mutex>, - business_write_attempts: AtomicUsize, - write_commands: Mutex>, - requests: Mutex>, - drop_open_reply: AtomicUsize, - drop_write_reply: AtomicUsize, - drop_close_reply: AtomicUsize, - latch_write: bool, - terminal_error: bool, - write_started: Notify, - release_write: Notify, - close_started: Notify, - } - - impl MockDatanode { - fn response(reply: Reply) -> RegionResponse { - let mut response = RegionResponse::new(0); - response.extensions.insert( - flow_write::COMMAND_KEY.to_string(), - serde_json::to_vec(&reply).unwrap(), - ); - response - } - - fn command(request: &RegionRequest) -> Command { - let serialized = request - .header - .as_ref() - .and_then(|header| header.query_context.as_ref()) - .and_then(|context| context.extensions.get(flow_write::COMMAND_KEY)) - .unwrap(); - Command::decode(serialized).unwrap() - } - } - - #[async_trait::async_trait] - impl Datanode for MockDatanode { - async fn handle_request( - &self, - request: Arc, - ) -> common_meta::error::Result { - self.requests.lock().unwrap().push(request.as_ref().clone()); - match Self::command(&request) { - Command::Open { nonce } => { - if !matches!(request.body, Some(Body::Inserts(ref inserts)) if inserts.requests.is_empty()) - { - return Err(common_meta::error::Error::External { - source: BoxedError::new(common_error::mock::MockError::new( - common_error::status_code::StatusCode::InvalidArguments, - )), - location: snafu::location!(), - }); - } - let mut capabilities = self.capabilities.lock().unwrap(); - let capability = capabilities - .entry(nonce) - .or_insert_with(|| "cap-1".into()) - .clone(); - if self - .drop_open_reply - .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |n| n.checked_sub(1)) - .is_ok() - { - return Err(common_meta::error::Error::External { - source: BoxedError::new(common_error::mock::MockError::new( - common_error::status_code::StatusCode::TableUnavailable, - )), - location: snafu::location!(), - }); - } - Ok(Self::response(Reply::Opened { capability })) - } - Command::Write { capability } => { - self.write_commands.lock().unwrap().push(capability.clone()); - if matches!(request.body, Some(Body::Inserts(ref inserts)) if !inserts.requests.is_empty()) - && self.business_write_attempts.fetch_add(1, Ordering::SeqCst) == 0 - && self.latch_write - { - self.write_started.notify_one(); - self.release_write.notified().await; - } - if self - .drop_write_reply - .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |n| n.checked_sub(1)) - .is_ok() - { - return Err(common_meta::error::Error::External { - source: BoxedError::new(common_error::mock::MockError::new( - common_error::status_code::StatusCode::TableUnavailable, - )), - location: snafu::location!(), - }); - } - Ok(Self::response(Reply::Complete { - rows: 7, - cost: 0, - metrics: None, - error: self.terminal_error.then(|| flow_write::TerminalError { - code: common_error::status_code::StatusCode::TableNotFound as u32, - message: "terminal child error".to_string(), - }), - })) - } - Command::Close { .. } => { - self.close_started.notify_one(); - if self - .drop_close_reply - .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |n| n.checked_sub(1)) - .is_ok() - { - return Err(common_meta::error::Error::External { - source: BoxedError::new(common_error::mock::MockError::new( - common_error::status_code::StatusCode::TableUnavailable, - )), - location: snafu::location!(), - }); - } - Ok(Self::response(Reply::Closed)) - } - } - } - - async fn handle_query( - &self, - _: QueryRequest, - ) -> common_meta::error::Result { - unreachable!() - } - } - - fn request() -> RegionRequest { - RegionRequest { - header: None, - body: Some(Body::Inserts(api::v1::region::InsertRequests { - requests: vec![Default::default()], - })), - } - } - - #[tokio::test] - async fn lost_control_and_write_replies_reuse_capability_and_exact_request() { - let datanode = Arc::new(MockDatanode::default()); - datanode.drop_open_reply.store(1, Ordering::SeqCst); - datanode.drop_write_reply.store(1, Ordering::SeqCst); - datanode.drop_close_reply.store(1, Ordering::SeqCst); - let response = dispatch_datanode(datanode.clone(), request()) - .await - .unwrap(); - assert_eq!(response.affected_rows, 7); - assert_eq!(datanode.business_write_attempts.load(Ordering::SeqCst), 2); - let writes = datanode.write_commands.lock().unwrap(); - assert_eq!(writes.len(), 2); - assert!(writes.iter().all(|capability| capability == "cap-1")); - let requests = datanode.requests.lock().unwrap(); - let writes: Vec<_> = requests - .iter() - .filter(|request| matches!(MockDatanode::command(request), Command::Write { .. })) - .collect(); - assert_eq!(writes.len(), 2); - assert_eq!(writes[0].body, writes[1].body); - assert_eq!(writes[0].header, writes[1].header); - assert_eq!( - requests - .iter() - .filter(|request| matches!( - MockDatanode::command(request), - Command::Open { .. } | Command::Close { .. } - )) - .count(), - 4 - ); - assert!(requests - .iter() - .filter(|request| matches!(MockDatanode::command(request), Command::Open { .. } | Command::Close { .. })) - .all(|request| matches!(request.body, Some(Body::Inserts(ref inserts)) if inserts.requests.is_empty()))); - } - - #[tokio::test] - async fn unsupported_open_never_submits_business_request() { - struct Unsupported; - #[async_trait::async_trait] - impl Datanode for Unsupported { - async fn handle_request( - &self, - _: Arc, - ) -> common_meta::error::Result { - Ok(RegionResponse::new(0)) - } - async fn handle_query( - &self, - _: QueryRequest, - ) -> common_meta::error::Result { - unreachable!() - } - } - let datanode: DatanodeRef = Arc::new(Unsupported); - let error = dispatch_datanode(datanode, request()).await.unwrap_err(); - assert_eq!( - error.status_code(), - common_error::status_code::StatusCode::Unexpected - ); - } - - #[tokio::test] - async fn all_peer_writes_drain_before_terminal_error_returns() { - let failing = Arc::new(MockDatanode { - terminal_error: true, - ..Default::default() - }); - let latched = Arc::new(MockDatanode { - latch_write: true, - ..Default::default() - }); - let first_datanode: DatanodeRef = failing.clone(); - let second_datanode: DatanodeRef = latched.clone(); - let aggregate = tokio::spawn(async move { - join_all([ - dispatch_datanode(first_datanode, request()), - dispatch_datanode(second_datanode, request()), - ]) - .await - }); - latched.write_started.notified().await; - failing.close_started.notified().await; - assert!(!aggregate.is_finished()); - latched.release_write.notify_one(); - let results = aggregate.await.unwrap(); - assert_eq!(results.len(), 2); - assert_eq!(failing.business_write_attempts.load(Ordering::SeqCst), 1); - assert_eq!(latched.business_write_attempts.load(Ordering::SeqCst), 1); - assert!(results.into_iter().any(|result| result.is_err())); - } -} diff --git a/src/session/Cargo.toml b/src/session/Cargo.toml index 630d0713a62..93c46ff7ff7 100644 --- a/src/session/Cargo.toml +++ b/src/session/Cargo.toml @@ -4,9 +4,6 @@ version.workspace = true edition.workspace = true license.workspace = true -[features] -enterprise = [] - [lints] workspace = true diff --git a/src/session/src/context.rs b/src/session/src/context.rs index f46a2a590b2..02ee621e0e4 100644 --- a/src/session/src/context.rs +++ b/src/session/src/context.rs @@ -467,11 +467,6 @@ impl QueryContext { self.extensions.insert(key.into(), value.into()); } - #[cfg(feature = "enterprise")] - pub fn remove_extension(&mut self, key: &str) { - self.extensions.remove(key); - } - pub fn extension>(&self, key: S) -> Option<&str> { self.extensions.get(key.as_ref()).map(|v| v.as_str()) } From 4135e65a76cf7e9d832cd00b0ff7415e971ef478 Mon Sep 17 00:00:00 2001 From: discord9 Date: Wed, 28 Oct 2026 10:40:43 +0000 Subject: [PATCH 121/239] fix(query): keep requested sort when filtering on a lossy cast (#10026) fix(query): pin DataFusion cast constraint propagation fix Update the pinned GreptimeTeam/datafusion revision from e3221b8 to cc55344 (the 55.1 fork carrying apache/datafusion#25531 plus a stronger regression test). CastExpr::propagate_constraints used to cast a filter's output interval straight back to the input type. For CAST(k AS BIGINT) = 1 the singleton output [1, 1] was inverted to [1.0, 1.0] although k = 1.5 matches too, so FilterExec advertised k as constant, the equivalence classes dropped the ordering required by ORDER BY k, and rows returned in scan order. The sqlness case keeps every #9600 assertion and the explicit ORDER BY, and adds ascending and descending checks for the partitioned table, the control table and a table whose matching rows are stored in neither ascending nor descending k order. Cargo.lock only rewrites the 35 datafusion source lines. Signed-off-by: discord9 --- Cargo.lock | 70 +++++++++--------- Cargo.toml | 32 ++++---- .../optimizer/partition_pruning_cast.result | 74 +++++++++++++++++++ .../optimizer/partition_pruning_cast.sql | 36 +++++++++ 4 files changed, 161 insertions(+), 51 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 6ada3b0a295..69996d9cbb2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3756,7 +3756,7 @@ checksum = "a4ae5f15dda3c708c0ade84bfee31ccab44a3da4f88015ed22f63732abe300c8" [[package]] name = "datafusion" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=cc5534492d5408831395609cdafb6520c7d514b7#cc5534492d5408831395609cdafb6520c7d514b7" dependencies = [ "arrow 59.2.0", "arrow-schema 59.2.0", @@ -3808,7 +3808,7 @@ dependencies = [ [[package]] name = "datafusion-catalog" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=cc5534492d5408831395609cdafb6520c7d514b7#cc5534492d5408831395609cdafb6520c7d514b7" dependencies = [ "arrow 59.2.0", "async-trait", @@ -3832,7 +3832,7 @@ dependencies = [ [[package]] name = "datafusion-catalog-listing" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=cc5534492d5408831395609cdafb6520c7d514b7#cc5534492d5408831395609cdafb6520c7d514b7" dependencies = [ "arrow 59.2.0", "async-trait", @@ -3855,7 +3855,7 @@ dependencies = [ [[package]] name = "datafusion-common" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=cc5534492d5408831395609cdafb6520c7d514b7#cc5534492d5408831395609cdafb6520c7d514b7" dependencies = [ "arrow 59.2.0", "arrow-ipc 59.2.0", @@ -3881,7 +3881,7 @@ dependencies = [ [[package]] name = "datafusion-common-runtime" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=cc5534492d5408831395609cdafb6520c7d514b7#cc5534492d5408831395609cdafb6520c7d514b7" dependencies = [ "futures", "log", @@ -3891,7 +3891,7 @@ dependencies = [ [[package]] name = "datafusion-datasource" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=cc5534492d5408831395609cdafb6520c7d514b7#cc5534492d5408831395609cdafb6520c7d514b7" dependencies = [ "arrow 59.2.0", "async-compression", @@ -3927,7 +3927,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-arrow" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=cc5534492d5408831395609cdafb6520c7d514b7#cc5534492d5408831395609cdafb6520c7d514b7" dependencies = [ "arrow 59.2.0", "arrow-ipc 59.2.0", @@ -3951,7 +3951,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-csv" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=cc5534492d5408831395609cdafb6520c7d514b7#cc5534492d5408831395609cdafb6520c7d514b7" dependencies = [ "arrow 59.2.0", "async-trait", @@ -3974,7 +3974,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-json" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=cc5534492d5408831395609cdafb6520c7d514b7#cc5534492d5408831395609cdafb6520c7d514b7" dependencies = [ "arrow 59.2.0", "async-trait", @@ -3997,7 +3997,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-parquet" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=cc5534492d5408831395609cdafb6520c7d514b7#cc5534492d5408831395609cdafb6520c7d514b7" dependencies = [ "arrow 59.2.0", "arrow-schema 59.2.0", @@ -4029,12 +4029,12 @@ dependencies = [ [[package]] name = "datafusion-doc" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=cc5534492d5408831395609cdafb6520c7d514b7#cc5534492d5408831395609cdafb6520c7d514b7" [[package]] name = "datafusion-execution" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=cc5534492d5408831395609cdafb6520c7d514b7#cc5534492d5408831395609cdafb6520c7d514b7" dependencies = [ "arrow 59.2.0", "arrow-buffer 59.2.0", @@ -4059,7 +4059,7 @@ dependencies = [ [[package]] name = "datafusion-expr" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=cc5534492d5408831395609cdafb6520c7d514b7#cc5534492d5408831395609cdafb6520c7d514b7" dependencies = [ "arrow 59.2.0", "arrow-schema 59.2.0", @@ -4083,7 +4083,7 @@ dependencies = [ [[package]] name = "datafusion-expr-common" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=cc5534492d5408831395609cdafb6520c7d514b7#cc5534492d5408831395609cdafb6520c7d514b7" dependencies = [ "arrow 59.2.0", "datafusion-common", @@ -4094,7 +4094,7 @@ dependencies = [ [[package]] name = "datafusion-functions" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=cc5534492d5408831395609cdafb6520c7d514b7#cc5534492d5408831395609cdafb6520c7d514b7" dependencies = [ "arrow 59.2.0", "arrow-buffer 59.2.0", @@ -4125,7 +4125,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=cc5534492d5408831395609cdafb6520c7d514b7#cc5534492d5408831395609cdafb6520c7d514b7" dependencies = [ "arrow 59.2.0", "datafusion-common", @@ -4145,7 +4145,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate-common" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=cc5534492d5408831395609cdafb6520c7d514b7#cc5534492d5408831395609cdafb6520c7d514b7" dependencies = [ "arrow 59.2.0", "datafusion-common", @@ -4156,7 +4156,7 @@ dependencies = [ [[package]] name = "datafusion-functions-nested" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=cc5534492d5408831395609cdafb6520c7d514b7#cc5534492d5408831395609cdafb6520c7d514b7" dependencies = [ "arrow 59.2.0", "arrow-ord 59.2.0", @@ -4180,7 +4180,7 @@ dependencies = [ [[package]] name = "datafusion-functions-table" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=cc5534492d5408831395609cdafb6520c7d514b7#cc5534492d5408831395609cdafb6520c7d514b7" dependencies = [ "arrow 59.2.0", "async-trait", @@ -4195,7 +4195,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=cc5534492d5408831395609cdafb6520c7d514b7#cc5534492d5408831395609cdafb6520c7d514b7" dependencies = [ "arrow 59.2.0", "datafusion-common", @@ -4211,7 +4211,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window-common" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=cc5534492d5408831395609cdafb6520c7d514b7#cc5534492d5408831395609cdafb6520c7d514b7" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -4220,7 +4220,7 @@ dependencies = [ [[package]] name = "datafusion-macros" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=cc5534492d5408831395609cdafb6520c7d514b7#cc5534492d5408831395609cdafb6520c7d514b7" dependencies = [ "datafusion-doc", "quote", @@ -4230,7 +4230,7 @@ dependencies = [ [[package]] name = "datafusion-optimizer" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=cc5534492d5408831395609cdafb6520c7d514b7#cc5534492d5408831395609cdafb6520c7d514b7" dependencies = [ "arrow 59.2.0", "chrono", @@ -4278,7 +4278,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=cc5534492d5408831395609cdafb6520c7d514b7#cc5534492d5408831395609cdafb6520c7d514b7" dependencies = [ "arrow 59.2.0", "arrow-schema 59.2.0", @@ -4301,7 +4301,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-adapter" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=cc5534492d5408831395609cdafb6520c7d514b7#cc5534492d5408831395609cdafb6520c7d514b7" dependencies = [ "arrow 59.2.0", "datafusion-common", @@ -4315,7 +4315,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-common" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=cc5534492d5408831395609cdafb6520c7d514b7#cc5534492d5408831395609cdafb6520c7d514b7" dependencies = [ "arrow 59.2.0", "chrono", @@ -4332,7 +4332,7 @@ dependencies = [ [[package]] name = "datafusion-physical-optimizer" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=cc5534492d5408831395609cdafb6520c7d514b7#cc5534492d5408831395609cdafb6520c7d514b7" dependencies = [ "arrow 59.2.0", "datafusion-common", @@ -4351,7 +4351,7 @@ dependencies = [ [[package]] name = "datafusion-physical-plan" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=cc5534492d5408831395609cdafb6520c7d514b7#cc5534492d5408831395609cdafb6520c7d514b7" dependencies = [ "arrow 59.2.0", "arrow-data 59.2.0", @@ -4388,7 +4388,7 @@ dependencies = [ [[package]] name = "datafusion-proto" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=cc5534492d5408831395609cdafb6520c7d514b7#cc5534492d5408831395609cdafb6520c7d514b7" dependencies = [ "arrow 59.2.0", "datafusion-catalog", @@ -4414,7 +4414,7 @@ dependencies = [ [[package]] name = "datafusion-proto-common" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=cc5534492d5408831395609cdafb6520c7d514b7#cc5534492d5408831395609cdafb6520c7d514b7" dependencies = [ "arrow 59.2.0", "datafusion-common", @@ -4424,7 +4424,7 @@ dependencies = [ [[package]] name = "datafusion-proto-models" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=cc5534492d5408831395609cdafb6520c7d514b7#cc5534492d5408831395609cdafb6520c7d514b7" dependencies = [ "datafusion-common", "datafusion-proto-common", @@ -4434,7 +4434,7 @@ dependencies = [ [[package]] name = "datafusion-pruning" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=cc5534492d5408831395609cdafb6520c7d514b7#cc5534492d5408831395609cdafb6520c7d514b7" dependencies = [ "arrow 59.2.0", "datafusion-common", @@ -4449,7 +4449,7 @@ dependencies = [ [[package]] name = "datafusion-session" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=cc5534492d5408831395609cdafb6520c7d514b7#cc5534492d5408831395609cdafb6520c7d514b7" dependencies = [ "arrow-schema 59.2.0", "async-trait", @@ -4463,7 +4463,7 @@ dependencies = [ [[package]] name = "datafusion-sql" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=cc5534492d5408831395609cdafb6520c7d514b7#cc5534492d5408831395609cdafb6520c7d514b7" dependencies = [ "arrow 59.2.0", "bigdecimal 0.4.8", @@ -4482,7 +4482,7 @@ dependencies = [ [[package]] name = "datafusion-substrait" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=e3221b8f74b15c7192621576e5d1e354a6cf097a#e3221b8f74b15c7192621576e5d1e354a6cf097a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=cc5534492d5408831395609cdafb6520c7d514b7#cc5534492d5408831395609cdafb6520c7d514b7" dependencies = [ "async-recursion", "async-trait", diff --git a/Cargo.toml b/Cargo.toml index 0e75ddbd912..15539abb2c7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -367,22 +367,22 @@ rev = "f743c3a5fe8c1f57363eac5866c356b7d91f9b39" [patch.crates-io] substrait = { git = "https://github.com/GreptimeTeam/substrait-rs.git", rev = "91ec978b0649417ad3da8390e7a515baec723b1b" } -datafusion = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } -datafusion-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } -datafusion-datasource = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } -datafusion-execution = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } -datafusion-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } -datafusion-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } -datafusion-functions = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } -datafusion-functions-aggregate-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } -datafusion-functions-window-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } -datafusion-optimizer = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } -datafusion-physical-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } -datafusion-physical-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } -datafusion-physical-plan = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } -datafusion-proto = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } -datafusion-sql = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } -datafusion-substrait = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" } +datafusion = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "cc5534492d5408831395609cdafb6520c7d514b7" } +datafusion-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "cc5534492d5408831395609cdafb6520c7d514b7" } +datafusion-datasource = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "cc5534492d5408831395609cdafb6520c7d514b7" } +datafusion-execution = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "cc5534492d5408831395609cdafb6520c7d514b7" } +datafusion-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "cc5534492d5408831395609cdafb6520c7d514b7" } +datafusion-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "cc5534492d5408831395609cdafb6520c7d514b7" } +datafusion-functions = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "cc5534492d5408831395609cdafb6520c7d514b7" } +datafusion-functions-aggregate-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "cc5534492d5408831395609cdafb6520c7d514b7" } +datafusion-functions-window-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "cc5534492d5408831395609cdafb6520c7d514b7" } +datafusion-optimizer = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "cc5534492d5408831395609cdafb6520c7d514b7" } +datafusion-physical-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "cc5534492d5408831395609cdafb6520c7d514b7" } +datafusion-physical-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "cc5534492d5408831395609cdafb6520c7d514b7" } +datafusion-physical-plan = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "cc5534492d5408831395609cdafb6520c7d514b7" } +datafusion-proto = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "cc5534492d5408831395609cdafb6520c7d514b7" } +datafusion-sql = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "cc5534492d5408831395609cdafb6520c7d514b7" } +datafusion-substrait = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "cc5534492d5408831395609cdafb6520c7d514b7" } datafusion-pg-catalog = { git = "https://github.com/datafusion-contrib/datafusion-postgres.git", rev = "3ae0a6cb30b1c5131687f4637247ff4f84ac11e4" } arrow-pg = { git = "https://github.com/datafusion-contrib/datafusion-postgres.git", rev = "3ae0a6cb30b1c5131687f4637247ff4f84ac11e4" } sqlparser = { git = "https://github.com/GreptimeTeam/sqlparser-rs.git", rev = "9e9019bb1c7040ed956f654e39378dd42ab17884" } diff --git a/tests/cases/distributed/optimizer/partition_pruning_cast.result b/tests/cases/distributed/optimizer/partition_pruning_cast.result index 2b18475fda8..2fda5e7a12a 100644 --- a/tests/cases/distributed/optimizer/partition_pruning_cast.result +++ b/tests/cases/distributed/optimizer/partition_pruning_cast.result @@ -97,6 +97,80 @@ ORDER BY k; | 1970-01-01T00:00:01 | 1.5 | 15 | 1 | +---------------------+-----+----+--------+ +-- The same rows must also come back in descending order: a filter on a lossy +-- cast must not turn `k` into a constant and drop the requested sort. +SELECT ts, k, v, CAST(k AS BIGINT) AS cast_k +FROM partition_pruning_cast_partitioned +WHERE CAST(k AS BIGINT) = 1 +ORDER BY k DESC; + ++---------------------+-----+----+--------+ +| ts | k | v | cast_k | ++---------------------+-----+----+--------+ +| 1970-01-01T00:00:01 | 1.5 | 15 | 1 | +| 1970-01-01T00:00:01 | 1.0 | 10 | 1 | ++---------------------+-----+----+--------+ + +SELECT ts, k, v, CAST(k AS BIGINT) AS cast_k +FROM partition_pruning_cast_control +WHERE CAST(k AS BIGINT) = 1 +ORDER BY k DESC; + ++---------------------+-----+----+--------+ +| ts | k | v | cast_k | ++---------------------+-----+----+--------+ +| 1970-01-01T00:00:01 | 1.5 | 15 | 1 | +| 1970-01-01T00:00:01 | 1.0 | 10 | 1 | ++---------------------+-----+----+--------+ + +-- A table whose matching rows are stored in neither ascending nor descending +-- order, so a dropped sort cannot accidentally satisfy the ASC assertion. +CREATE TABLE partition_pruning_cast_unsorted ( + ts TIMESTAMP TIME INDEX, + k FLOAT PRIMARY KEY, + v BIGINT, +); + +Affected Rows: 0 + +INSERT INTO partition_pruning_cast_unsorted VALUES + (1000, 1.5, 15), + (1000, 2.5, 25), + (1000, 1.0, 10), + (1000, 1.25, 12); + +Affected Rows: 4 + +SELECT ts, k, v, CAST(k AS BIGINT) AS cast_k +FROM partition_pruning_cast_unsorted +WHERE CAST(k AS BIGINT) = 1 +ORDER BY k; + ++---------------------+------+----+--------+ +| ts | k | v | cast_k | ++---------------------+------+----+--------+ +| 1970-01-01T00:00:01 | 1.0 | 10 | 1 | +| 1970-01-01T00:00:01 | 1.25 | 12 | 1 | +| 1970-01-01T00:00:01 | 1.5 | 15 | 1 | ++---------------------+------+----+--------+ + +SELECT ts, k, v, CAST(k AS BIGINT) AS cast_k +FROM partition_pruning_cast_unsorted +WHERE CAST(k AS BIGINT) = 1 +ORDER BY k DESC; + ++---------------------+------+----+--------+ +| ts | k | v | cast_k | ++---------------------+------+----+--------+ +| 1970-01-01T00:00:01 | 1.5 | 15 | 1 | +| 1970-01-01T00:00:01 | 1.25 | 12 | 1 | +| 1970-01-01T00:00:01 | 1.0 | 10 | 1 | ++---------------------+------+----+--------+ + +DROP TABLE partition_pruning_cast_unsorted; + +Affected Rows: 0 + -- SQLNESS REPLACE (-+) - -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (elapsed_compute.*) REDACTED diff --git a/tests/cases/distributed/optimizer/partition_pruning_cast.sql b/tests/cases/distributed/optimizer/partition_pruning_cast.sql index 058045f733a..f72355b892d 100644 --- a/tests/cases/distributed/optimizer/partition_pruning_cast.sql +++ b/tests/cases/distributed/optimizer/partition_pruning_cast.sql @@ -43,6 +43,42 @@ FROM partition_pruning_cast_control WHERE CAST(k AS BIGINT) = 1 ORDER BY k; +-- The same rows must also come back in descending order: a filter on a lossy +-- cast must not turn `k` into a constant and drop the requested sort. +SELECT ts, k, v, CAST(k AS BIGINT) AS cast_k +FROM partition_pruning_cast_partitioned +WHERE CAST(k AS BIGINT) = 1 +ORDER BY k DESC; +SELECT ts, k, v, CAST(k AS BIGINT) AS cast_k +FROM partition_pruning_cast_control +WHERE CAST(k AS BIGINT) = 1 +ORDER BY k DESC; + +-- A table whose matching rows are stored in neither ascending nor descending +-- order, so a dropped sort cannot accidentally satisfy the ASC assertion. +CREATE TABLE partition_pruning_cast_unsorted ( + ts TIMESTAMP TIME INDEX, + k FLOAT PRIMARY KEY, + v BIGINT, +); + +INSERT INTO partition_pruning_cast_unsorted VALUES + (1000, 1.5, 15), + (1000, 2.5, 25), + (1000, 1.0, 10), + (1000, 1.25, 12); + +SELECT ts, k, v, CAST(k AS BIGINT) AS cast_k +FROM partition_pruning_cast_unsorted +WHERE CAST(k AS BIGINT) = 1 +ORDER BY k; +SELECT ts, k, v, CAST(k AS BIGINT) AS cast_k +FROM partition_pruning_cast_unsorted +WHERE CAST(k AS BIGINT) = 1 +ORDER BY k DESC; + +DROP TABLE partition_pruning_cast_unsorted; + -- SQLNESS REPLACE (-+) - -- SQLNESS REPLACE (\s\s+) _ -- SQLNESS REPLACE (elapsed_compute.*) REDACTED From c818b6ac2a290a4b40802b959d8d687fdfa6efed Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Wed, 28 Oct 2026 08:00:41 +0000 Subject: [PATCH 122/239] feat(postgres): influxdb line protocol ingest via COPY FROM STDIN (#9982) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(postgres): influxdb line protocol ingest via COPY FROM STDIN COPY FROM STDIN WITH (FORMAT influxdb [, precision 'ns']) ingests InfluxDB line protocol over the postgres wire protocol, reusing the HTTP ingestion pipeline wholesale for correctness and performance parity: the whole payload is buffered and completed once at CopyDone through the frontend's InfluxdbLineProtocolHandler::exec — conversion, permissions, line-protocol interceptors, timestamp alignment to table time-index units, merge mode and inserts are the same path an HTTP write of the same bytes takes, under the connection's regular postgres query context. Formats are an explicit, closed enum (CopyFormat::Csv | CopyFormat::Influxdb(options)); each variant states its own rules in the copy lifecycle's match arms — csv requires a target table and parses incrementally, influxdb forbids one (each line names its own measurement) and buffers one bounded request body. The COPY layer owns only the buffering: record boundaries, including newlines inside quoted field values and after escapes, belong to the line-protocol parser alone — one COPY statement is one bounded ingestion request (64 MiB; larger imports are multiple COPY statements). - Wiring: PostgresServer::with_influxdb_handler attaches the shared line-protocol handler; unwired servers reject the format when a copy starts. - Grammar: COPY FROM STDIN without a table is scanned (an unquoted FROM cannot be a table name); table-bound formats reject it with guidance, influxdb rejects a named table. - Options: precision accepts the HTTP endpoint's values via the shared parse_time_precision (moved to servers::influxdb), defaulting to nanoseconds. Review findings (fengjiachun) are fixed at the root by removing the pre-splitter rather than repairing it: the retired naive line splitter broke quoted multiline records at chunk/flush boundaries (found by removing record splitting entirely), drained the buffer quadratically (found by append-only buffering), and capped nothing (found by an explicit payload limit). e2e: multi-measurement ingest with chunked feeds; a 32,768-physical-line flush-threshold regression with a multiline quoted record split across CopyData chunks; precision 'ms' epoch verification; invalid UTF-8 and trailing-record-without-newline handling; payload limit; table-named rejection; malformed-line abort with connection recovery. Signed-off-by: Ning Sun * chore: typo --------- Signed-off-by: Ning Sun --- src/frontend/src/server.rs | 1 + src/servers/src/http/influxdb.rs | 23 +- src/servers/src/influxdb.rs | 23 +- src/servers/src/postgres.rs | 10 +- src/servers/src/postgres/copy_in/codec.rs | 118 +++---- src/servers/src/postgres/copy_in/csv.rs | 52 ++- src/servers/src/postgres/copy_in/influxdb.rs | 95 ++++++ src/servers/src/postgres/copy_in/mod.rs | 305 +++++++++++++----- src/servers/src/postgres/copy_in/parse.rs | 73 ++++- src/servers/src/postgres/server.rs | 27 +- tests-integration/src/test_util.rs | 69 ++-- .../src/tests/postgres_copy_in_test.rs | 300 +++++++++++++++++ 12 files changed, 851 insertions(+), 245 deletions(-) create mode 100644 src/servers/src/postgres/copy_in/influxdb.rs diff --git a/src/frontend/src/server.rs b/src/frontend/src/server.rs index 0b1e2b9713e..d84f01fce95 100644 --- a/src/frontend/src/server.rs +++ b/src/frontend/src/server.rs @@ -429,6 +429,7 @@ where user_provider.clone(), Some(self.instance.process_manager().clone()), ) + .with_influxdb_handler(instance.clone()) .with_batching_enabled(postgres_batching), ) as Box; diff --git a/src/servers/src/http/influxdb.rs b/src/servers/src/http/influxdb.rs index 671beb7bae2..44815cae3a1 100644 --- a/src/servers/src/http/influxdb.rs +++ b/src/servers/src/http/influxdb.rs @@ -24,7 +24,7 @@ use common_grpc::precision::Precision; use common_telemetry::tracing; use session::context::{Channel, QueryContext, QueryContextRef}; -use crate::error::{Result, TimePrecisionSnafu}; +use crate::error::Result; use crate::http::header::write_cost_header_map; use crate::influxdb::InfluxdbRequest; use crate::query_handler::InfluxdbLineProtocolHandlerRef; @@ -107,30 +107,13 @@ pub async fn influxdb_write( )) } -fn parse_time_precision(value: &str) -> Result { - // Precision conversion needs to be compatible with influxdb v1 v2 api. - // For details, see the Influxdb documents. - // https://docs.influxdata.com/influxdb/v1.8/tools/api/#apiv2write-http-endpoint - // https://docs.influxdata.com/influxdb/v1.8/tools/api/#write-http-endpoint - match value { - "n" | "ns" => Ok(Precision::Nanosecond), - "u" | "us" => Ok(Precision::Microsecond), - "ms" => Ok(Precision::Millisecond), - "s" => Ok(Precision::Second), - "m" => Ok(Precision::Minute), - "h" => Ok(Precision::Hour), - unknown => TimePrecisionSnafu { - name: unknown.to_string(), - } - .fail(), - } -} +pub(crate) use crate::influxdb::parse_time_precision; #[cfg(test)] mod tests { use common_grpc::precision::Precision; - use crate::http::influxdb::parse_time_precision; + use super::parse_time_precision; #[test] fn test_parse_time_precision() { diff --git a/src/servers/src/influxdb.rs b/src/servers/src/influxdb.rs index 9bff0bbc6e4..9dd26d9cda7 100644 --- a/src/servers/src/influxdb.rs +++ b/src/servers/src/influxdb.rs @@ -20,7 +20,7 @@ use hyper::Request; use influxdb_line_protocol::{FieldValue, parse_lines}; use snafu::ResultExt; -use crate::error::{Error, InfluxdbLineProtocolSnafu}; +use crate::error::{Error, InfluxdbLineProtocolSnafu, TimePrecisionSnafu}; use crate::row_writer::{self, MultiTableData}; const INFLUXDB_API_PATH_NAME: &str = "influxdb"; @@ -113,6 +113,27 @@ fn unwrap_or_default_precision(precision: Option) -> Precision { } } +/// Parses an InfluxDB time-precision name shared by the HTTP write +/// endpoints (v1/v2) and the `precision` option of COPY FROM STDIN. +pub(crate) fn parse_time_precision(value: &str) -> crate::error::Result { + // Precision conversion needs to be compatible with influxdb v1 v2 api. + // For details, see the Influxdb documents. + // https://docs.influxdata.com/influxdb/v1.8/tools/api/#apiv2write-http-endpoint + // https://docs.influxdata.com/influxdb/v1.8/tools/api/#write-http-endpoint + match value { + "n" | "ns" => Ok(Precision::Nanosecond), + "u" | "us" => Ok(Precision::Microsecond), + "ms" => Ok(Precision::Millisecond), + "s" => Ok(Precision::Second), + "m" => Ok(Precision::Minute), + "h" => Ok(Precision::Hour), + unknown => TimePrecisionSnafu { + name: unknown.to_string(), + } + .fail(), + } +} + #[cfg(test)] mod tests { use api::v1::value::ValueData; diff --git a/src/servers/src/postgres.rs b/src/servers/src/postgres.rs index e13ec7392d5..fbef2e9ac4c 100644 --- a/src/servers/src/postgres.rs +++ b/src/servers/src/postgres.rs @@ -43,8 +43,8 @@ use session::context::Channel; use self::auth_handler::{PgUserProviderSource, ZeroPidSecretKeyGenerator}; use self::handler::DefaultQueryParser; -use crate::query_handler::CopyInHandlerRef; use crate::query_handler::sql::ServerSqlQueryHandlerRef; +use crate::query_handler::{CopyInHandlerRef, InfluxdbLineProtocolHandlerRef}; pub(crate) struct GreptimeDBStartupParameters { version: String, @@ -81,6 +81,11 @@ pub struct PostgresServerHandlerInner { query_handler: ServerSqlQueryHandlerRef, auth_source: Option>, copy_in_handler: CopyInHandlerRef, + /// Executes influxdb line-protocol batches flushed by a table-less + /// COPY FROM STDIN — the same handler the HTTP write endpoints use. + /// Absent on servers not wired for it; such connections reject the + /// influxdb format when a copy starts. + influxdb_handler: Option, force_tls: bool, auth_driver: AuthDriver, UserInfoRef>, /// SCRAM selected; failures get the uniform metric label. @@ -99,6 +104,8 @@ pub struct PostgresServerHandlerInner { pub(crate) struct MakePostgresServerHandler { query_handler: ServerSqlQueryHandlerRef, copy_in_handler: CopyInHandlerRef, + #[builder(default)] + influxdb_handler: Option, user_provider: Option, #[builder(default = "Arc::new(GreptimeDBStartupParameters::new())")] param_provider: Arc, @@ -149,6 +156,7 @@ impl MakePostgresServerHandler { .clone() .map(|p| Arc::new(PgUserProviderSource::new(p))), copy_in_handler: self.copy_in_handler.clone(), + influxdb_handler: self.influxdb_handler.clone(), force_tls: self.force_tls, auth_driver: AuthDriver::new(self.param_provider.clone()) .with_pid_secret_key_generator(Arc::new(ZeroPidSecretKeyGenerator)), diff --git a/src/servers/src/postgres/copy_in/codec.rs b/src/servers/src/postgres/copy_in/codec.rs index 37a94811309..a7672c619ef 100644 --- a/src/servers/src/postgres/copy_in/codec.rs +++ b/src/servers/src/postgres/copy_in/codec.rs @@ -12,58 +12,63 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! The COPY FROM STDIN codec framework. +//! The COPY FROM STDIN formats. //! -//! A *codec* binds one COPY data format to its implementation. Codecs are -//! stateless strategies: they validate the options they accept at -//! statement-parse time, describe which column types they can ingest, and -//! create the stateful [`RecordParser`] per stream. -//! -//! Formats are not limited to PostgreSQL's own — GreptimeDB-specific -//! formats (a line protocol, a JSON-lines dialect, ...) register exactly -//! like `csv` does, and may define their own options through the open -//! `WITH (...)` clause. +//! [`CopyFormat`] is the closed, explicit set of supported formats — a +//! small enum over the real implementations rather than a generic +//! format framework. Each variant carries its validated options and +//! its rules are stated where they apply: a variant's match arms in +//! the copy lifecycle (see the parent module) decide whether a target +//! table is required, how the `CopyData` bytes are consumed, and which +//! existing ingestion pipeline completes the copy. Adding a format is +//! a normal extension: a new variant plus its arms. //! //! # Parsing and conversion //! -//! A [`RecordParser`] consumes raw `CopyData` bytes and emits complete, -//! insertable rows: one [`ParsedRow`] per record, in the order of the -//! targeted columns. Parsers validate record shape (field counts) and -//! convert field values themselves — `csv` delegates both splitting and -//! typed conversion to arrow-csv's incremental decoder (the engine -//! behind `COPY ... FROM 'file.csv'`) and turns decoded batches into -//! rows with the shared `api::helper::vectors_to_rows`. A future format -//! that carries values as text (e.g. a line protocol) would reintroduce -//! shared field-conversion helpers here; see the removed `text` codec in -//! git history for the shape they had. +//! Table-bound formats parse incrementally through a [`RecordParser`]: +//! it consumes raw `CopyData` bytes and emits complete, insertable +//! rows — one [`ParsedRow`] per record, in the order of the targeted +//! columns. `csv` delegates splitting and typed conversion to +//! arrow-csv's incremental decoder (the engine behind +//! `COPY ... FROM 'file.csv'`) and turns decoded batches into rows +//! with the shared `api::helper::vectors_to_rows`. +//! +//! Formats that delegate to an existing ingestion pipeline (the +//! influxdb line protocol) do not parse here at all: the whole payload +//! is buffered as one bounded request body and completed through that +//! pipeline, so record boundaries and values stay with the pipeline's +//! own parser. //! //! # Options //! //! The `WITH (...)` clause is open: every `KEY value` entry the client -//! writes is collected into a [`CopyOptionSet`] (see `options`) regardless -//! of whether any format understands it — the statement is parsed -//! textually (see `parse`) precisely so options beyond sqlparser's COPY -//! grammar can reach the codecs. A codec builder takes the options it -//! understands (`take_string`, `take_single_byte`, `take_flag`, or raw -//! `take_value` for codec-specific shapes) and rejects the rest with -//! `reject_leftovers`, so a future format like -//! `COPY t FROM STDIN WITH (FORMAT influxdb_line, precision 'ns')` only -//! needs its own codec and registry entry. -//! -//! To add a format: -//! -//! 1. implement [`CopyInCodec`] for a config struct holding the format's -//! validated options (see `csv::CsvCodec` for the example), -//! 2. write a builder `fn(CopyOptionSet) -> PgWireResult>` that applies defaults and rejects options the format -//! does not accept, -//! 3. register the builder under every accepted `FORMAT` name (including -//! aliases) in `REGISTERED_FORMATS` in the parent module. +//! writes is collected into a [`CopyOptionSet`] (see `options`) +//! regardless of whether any format understands it — the statement is +//! parsed textually (see `parse`) precisely so options beyond +//! sqlparser's COPY grammar can reach the formats. A format builder +//! takes the options it understands (`take_string`, +//! `take_single_byte`, `take_flag`, or raw `take_value` for +//! format-specific shapes) and rejects the rest with +//! `reject_leftovers`. -use datatypes::prelude::ConcreteDataType; -use datatypes::schema::ColumnSchema; use pgwire::error::PgWireResult; +use crate::postgres::copy_in::csv::CsvCodec; +use crate::postgres::copy_in::influxdb::InfluxdbOptions; + +/// A resolved COPY FROM STDIN format with its validated options. +#[derive(Debug, Clone)] +pub(super) enum CopyFormat { + /// Table-bound CSV: records are parsed incrementally against the + /// targeted table's columns and inserted in batches. + Csv(CsvCodec), + /// Table-less influxdb line protocol: the whole payload is one + /// bounded request body, completed through the line-protocol + /// pipeline the HTTP write endpoints use (one COPY statement is one + /// ingestion request; streaming ingest is not supported). + Influxdb(InfluxdbOptions), +} + /// A parsed, insertable row: one gRPC value per targeted column, in COPY /// field order. Every parser emits these directly, whatever mechanism it /// used to produce them. @@ -83,34 +88,3 @@ pub(super) trait RecordParser: Send { /// stream ended exactly at a record boundary. fn finish(&mut self) -> PgWireResult>; } - -/// A COPY FROM STDIN data format implementation. -/// -/// Implementations are shared per statement (hence `Send + Sync`) and must -/// be stateless; per-stream state belongs to the [`RecordParser`] they -/// create. Every method is tailored to the format: `check_column_type` -/// declares the types the format can decode (csv is bounded by what -/// arrow-csv supports). -pub(super) trait CopyInCodec: Send + Sync + std::fmt::Debug { - /// The canonical format name, used in error messages. - fn name(&self) -> &'static str; - /// The format code advertised in the `CopyInResponse` header: 0 when - /// the client sends textual data (the default), 1 for binary. - fn format_code(&self) -> i8 { - 0 - } - /// Validates that a column of `data_type` can be ingested by this - /// format. - fn check_column_type(&self, data_type: &ConcreteDataType) -> PgWireResult<()>; - /// Creates the stateful parser for one copy stream over the targeted - /// columns (in COPY field order). - /// - /// `timezone` is the session timezone (an IANA name or a UTC offset): - /// formats that parse offset-free timestamp strings from text should - /// interpret them in this zone, matching SQL INSERT. - fn create_parser( - &self, - columns: &[ColumnSchema], - timezone: &str, - ) -> PgWireResult>; -} diff --git a/src/servers/src/postgres/copy_in/csv.rs b/src/servers/src/postgres/copy_in/csv.rs index b2919210e3f..ab8f521098d 100644 --- a/src/servers/src/postgres/copy_in/csv.rs +++ b/src/servers/src/postgres/copy_in/csv.rs @@ -40,12 +40,12 @@ use datatypes::schema::{ColumnSchema, Schema as GreptimeSchema}; use datatypes::vectors::{BooleanVector, Helper, StringVector, Vector}; use pgwire::error::{PgWireError, PgWireResult}; -use crate::postgres::copy_in::codec::{CopyInCodec, ParsedRow, RecordParser}; +use crate::postgres::copy_in::codec::{CopyFormat, ParsedRow, RecordParser}; use crate::postgres::copy_in::options::CopyOptionSet; use crate::postgres::copy_in::{COPY_IN_FLUSH_ROWS, bad_copy_data, invalid_text_repr, unsupported}; /// The validated `csv` format configuration. -#[derive(Debug)] +#[derive(Debug, Clone)] pub(super) struct CsvCodec { delimiter: u8, quote: u8, @@ -54,7 +54,7 @@ pub(super) struct CsvCodec { } /// Builds the CSV codec from the parsed options. -pub(super) fn build_csv_codec(mut options: CopyOptionSet) -> PgWireResult> { +pub(super) fn build_csv_codec(mut options: CopyOptionSet) -> PgWireResult { let delimiter = options.take_single_byte("DELIMITER")?.unwrap_or(b','); let quote = options.take_single_byte("QUOTE")?.unwrap_or(b'"'); // In PostgreSQL the default escape character for CSV is the quote @@ -75,7 +75,7 @@ pub(super) fn build_csv_codec(mut options: CopyOptionSet) -> PgWireResult PgWireResult &'static str { - "csv" - } - +impl CsvCodec { /// The column types arrow-csv can decode against. - fn check_column_type(&self, data_type: &ConcreteDataType) -> PgWireResult<()> { + pub(super) fn check_column_type(&self, data_type: &ConcreteDataType) -> PgWireResult<()> { let supported = matches!( data_type, ConcreteDataType::Boolean(_) @@ -117,7 +113,7 @@ impl CopyInCodec for CsvCodec { Ok(()) } - fn create_parser( + pub(super) fn create_parser( &self, columns: &[ColumnSchema], timezone: &str, @@ -292,9 +288,15 @@ mod tests { ] } + fn build_codec(set: CopyOptionSet) -> CsvCodec { + match build_csv_codec(set).unwrap() { + CopyFormat::Csv(codec) => codec, + other => panic!("expected csv format, got {other:?}"), + } + } + fn parser() -> Box { - build_csv_codec(CopyOptionSet::default()) - .unwrap() + build_codec(CopyOptionSet::default()) .create_parser(&test_columns(), "UTC") .unwrap() } @@ -308,8 +310,7 @@ mod tests { fn rows_with(mut set: CopyOptionSet) -> Box { set.take_format().unwrap(); - build_csv_codec(set) - .unwrap() + build_codec(set) .create_parser(&test_columns(), "UTC") .unwrap() } @@ -329,8 +330,8 @@ mod tests { #[test] fn test_build_defaults() { - let codec = build_csv_codec(CopyOptionSet::default()).unwrap(); - assert_eq!(codec.name(), "csv"); + let codec = build_codec(CopyOptionSet::default()); + assert!(matches!(codec, CsvCodec { .. })); } #[test] @@ -368,7 +369,7 @@ mod tests { #[test] fn test_check_column_type() { - let codec = build_csv_codec(CopyOptionSet::default()).unwrap(); + let codec = build_codec(CopyOptionSet::default()); assert!( codec .check_column_type(&ConcreteDataType::json_datatype()) @@ -524,8 +525,7 @@ mod tests { .into_bytes(); let check = |chunk: usize| { - let mut parser = build_csv_codec(CopyOptionSet::default()) - .unwrap() + let mut parser = build_codec(CopyOptionSet::default()) .create_parser(&sweep_columns, "UTC") .unwrap(); let mut ids: Vec = Vec::new(); @@ -560,8 +560,7 @@ mod tests { check(chunk); let mut shifted = payload.clone(); shifted.drain(..37); - let mut parser = build_csv_codec(CopyOptionSet::default()) - .unwrap() + let mut parser = build_codec(CopyOptionSet::default()) .create_parser(&sweep_columns, "UTC") .unwrap(); let mut count = 0usize; @@ -585,8 +584,7 @@ mod tests { false, ), ]; - let mut parser = build_csv_codec(CopyOptionSet::default()) - .unwrap() + let mut parser = build_codec(CopyOptionSet::default()) .create_parser(&columns, "UTC") .unwrap(); let rows = feed_all( @@ -628,8 +626,7 @@ mod tests { false, ), ]; - let mut parser = build_csv_codec(CopyOptionSet::default()) - .unwrap() + let mut parser = build_codec(CopyOptionSet::default()) .create_parser(&columns, "UTC") .unwrap(); parser.feed(b"maybe,2024-01-01 00:00:00\n").unwrap(); @@ -656,8 +653,7 @@ mod tests { false, ), ]; - let mut parser = build_csv_codec(CopyOptionSet::default()) - .unwrap() + let mut parser = build_codec(CopyOptionSet::default()) .create_parser(&columns, "Asia/Shanghai") .unwrap(); let rows = feed_all( diff --git a/src/servers/src/postgres/copy_in/influxdb.rs b/src/servers/src/postgres/copy_in/influxdb.rs new file mode 100644 index 00000000000..b4e2a9add30 --- /dev/null +++ b/src/servers/src/postgres/copy_in/influxdb.rs @@ -0,0 +1,95 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! The `influxdb` line protocol format for COPY FROM STDIN. +//! +//! `COPY FROM STDIN WITH (FORMAT influxdb, precision 'ns')` ingests +//! [InfluxDB line protocol](https://docs.influxdata.com/influxdb/v1.8/write_protocols/line_protocol_reference/) +//! text. The format is table-less — each line names its own measurement +//! — and does not parse here at all: the COPY layer buffers the whole +//! payload (bounded by [`crate::postgres::copy_in::COPY_IN_MAX_PAYLOAD_BYTES`]) +//! and completes the copy by handing it to the frontend's line-protocol +//! write pipeline as one body, exactly like the HTTP +//! `/v1/influxdb/v1/write` endpoint receives one request body: +//! conversion, permissions, interceptors, timestamp alignment, merge +//! mode and inserts are the same path, and record boundaries (including +//! newlines inside quoted field values and after escapes) are decided +//! by the line-protocol parser alone — the COPY layer never splits +//! records. +//! +//! One COPY statement is therefore one bounded ingestion request; a +//! larger import is multiple COPY statements. +//! +//! The `precision` option carries the same values as the HTTP +//! `precision` query parameter (`ns`, `us`, `ms`, `s`, ...), defaults to +//! nanoseconds, and is resolved by the shared `parse_time_precision`. + +use common_grpc::precision::Precision; +use pgwire::error::PgWireResult; + +use crate::postgres::copy_in::bad_copy_data; +use crate::postgres::copy_in::codec::CopyFormat; +use crate::postgres::copy_in::options::CopyOptionSet; + +/// The format's validated options. A `None` precision carries the +/// pipeline's default (nanoseconds), exactly like an HTTP write without +/// a `precision` parameter. +#[derive(Debug, Clone)] +pub(super) struct InfluxdbOptions { + pub(super) precision: Option, +} + +pub(super) fn build_influxdb_codec(mut options: CopyOptionSet) -> PgWireResult { + let precision = match options.take_string("precision")? { + Some(value) => Some( + crate::influxdb::parse_time_precision(&value) + .map_err(|e| bad_copy_data(format!("invalid precision option: {e}")))?, + ), + None => None, + }; + options.reject_leftovers("influxdb")?; + Ok(CopyFormat::Influxdb(InfluxdbOptions { precision })) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::postgres::copy_in::options::CopyOptionValue; + + #[test] + fn test_builder_options() { + let format = build_influxdb_codec(CopyOptionSet::default()).unwrap(); + assert!(matches!(format, CopyFormat::Influxdb(_))); + assert!(matches!(format, CopyFormat::Influxdb(_))); + + let mut set = CopyOptionSet::default(); + set.insert("precision".into(), CopyOptionValue::String("ms".into())) + .unwrap(); + assert!(matches!( + build_influxdb_codec(set).unwrap(), + CopyFormat::Influxdb(InfluxdbOptions { + precision: Some(Precision::Millisecond) + }) + )); + + let mut set = CopyOptionSet::default(); + set.insert("precision".into(), CopyOptionValue::String("bogus".into())) + .unwrap(); + assert!(build_influxdb_codec(set).is_err()); + + let mut set = CopyOptionSet::default(); + set.insert("unknown".into(), CopyOptionValue::Flag).unwrap(); + assert!(build_influxdb_codec(set).is_err()); + } +} diff --git a/src/servers/src/postgres/copy_in/mod.rs b/src/servers/src/postgres/copy_in/mod.rs index a3959a23830..92528d4ab2f 100644 --- a/src/servers/src/postgres/copy_in/mod.rs +++ b/src/servers/src/postgres/copy_in/mod.rs @@ -21,38 +21,42 @@ //! the data stream incrementally and flushes it into the query handler as //! [`RowInsertRequests`] batches. //! -//! # Format framework +//! # Formats //! -//! Each COPY data format is implemented as a [`codec::CopyInCodec`]: a -//! stateless strategy that validates its options at statement-parse time, -//! describes the column types it can ingest, and creates the stateful -//! record parser for one stream. Formats register a builder under every -//! accepted `FORMAT` name in [`REGISTERED_FORMATS`] — `csv` is the only -//! implemented format today; PostgreSQL's `text`/`binary` and -//! GreptimeDB-specific formats alike can be added the same way. See the -//! `codec` module docs. +//! The supported formats are the explicit [`codec::CopyFormat`] enum — +//! `csv` (table-bound, incremental row parsing) and `influxdb` +//! (table-less, one bounded whole-payload request). Each format states +//! its own target rule and ingestion shape in the match arms of the +//! copy lifecycle below; adding a format is a new variant plus its +//! arms. Builders resolve a `FORMAT` name (and its aliases) through +//! [`REGISTERED_FORMATS`]. See the `codec` module docs. mod codec; mod csv; +mod influxdb; mod options; mod parse; use std::fmt::Debug; use std::sync::Arc; +use crate::influxdb::InfluxdbRequest; use api::helper::ColumnDataTypeWrapper; use api::v1::column_def::options_from_column_schema; use api::v1::{ Row as GrpcRow, RowInsertRequest, RowInsertRequests, Rows as GrpcRows, SemanticType, }; use async_trait::async_trait; -use codec::{CopyInCodec, RecordParser}; +use codec::{CopyFormat, RecordParser}; +use common_grpc::precision::Precision; use common_query::OutputData; use common_telemetry::{info, warn}; +use csv::CsvCodec; use datafusion::sql::sqlparser::ast::{CopyTarget, Statement as SqlParserStatement}; use datatypes::data_type::DataType; use datatypes::schema::{ColumnSchema, SchemaRef}; use futures::{Sink, SinkExt, stream}; +use influxdb::InfluxdbOptions; use pgwire::api::copy::CopyHandler; use pgwire::api::results::{CopyResponse, Response, Tag}; use pgwire::api::{ClientInfo, PgWireConnectionState}; @@ -80,38 +84,43 @@ use crate::postgres::utils::convert_err; /// the per-batch fixed costs negligible. pub(crate) const COPY_IN_FLUSH_ROWS: usize = 32768; -/// Builds a codec from the format-neutral option set. -type CodecBuilder = fn(options::CopyOptionSet) -> PgWireResult>; +/// Upper bound of an influxdb COPY FROM STDIN payload. One COPY +/// statement is one bounded ingestion request (the whole payload is +/// buffered like one HTTP write body); a larger import is multiple COPY +/// statements. Streaming ingest may arrive later, with framing owned by +/// the shared line-protocol parser rather than the COPY layer. +pub(crate) const COPY_IN_MAX_PAYLOAD_BYTES: usize = 64 * 1024 * 1024; + +/// Builds a format from the format-neutral option set. +type FormatBuilder = fn(options::CopyOptionSet) -> PgWireResult; /// The registered COPY FROM STDIN formats: accepted `FORMAT` name to /// builder. `csv` is the only implemented format; the others are /// registered so they report a clear "not supported yet" instead of an /// unknown-format error. -const REGISTERED_FORMATS: &[(&str, CodecBuilder)] = &[ +const REGISTERED_FORMATS: &[(&str, FormatBuilder)] = &[ ("csv", csv::build_csv_codec), + ("influxdb", influxdb::build_influxdb_codec), + ("influxdb_line", influxdb::build_influxdb_codec), ("text", build_unimplemented_text_codec), ("txt", build_unimplemented_text_codec), ("binary", build_unimplemented_binary_codec), ]; -fn build_unimplemented_text_codec( - _options: options::CopyOptionSet, -) -> PgWireResult> { +fn build_unimplemented_text_codec(_options: options::CopyOptionSet) -> PgWireResult { Err(unsupported( "the text format is not supported yet; use WITH (FORMAT csv)", )) } -fn build_unimplemented_binary_codec( - _options: options::CopyOptionSet, -) -> PgWireResult> { +fn build_unimplemented_binary_codec(_options: options::CopyOptionSet) -> PgWireResult { Err(unsupported( "COPY FROM STDIN WITH (FORMAT binary) is not supported yet", )) } -/// Resolves the format codec from the parsed options. -fn resolve_codec(options: &mut options::CopyOptionSet) -> PgWireResult> { +/// Resolves the format from the parsed options. +fn resolve_format(options: &mut options::CopyOptionSet) -> PgWireResult { let name = options.take_format()?.ok_or_else(|| { unsupported( "COPY FROM STDIN requires an explicit format in this version; use WITH (FORMAT csv)", @@ -133,15 +142,15 @@ pub(crate) struct CopyFromStdin { table: Vec, /// Requested column list; empty means all columns in table order. columns: Vec, - /// The resolved format codec, validated and ready to run. - codec: Arc, + /// The resolved format with its validated options. + format: CopyFormat, } fn copy_in_error(code: PgErrorCode, message: impl Into) -> PgWireError { PgWireError::UserError(Box::new(code.to_err_info(message.into()))) } -fn unsupported(message: impl Into) -> PgWireError { +pub(super) fn unsupported(message: impl Into) -> PgWireError { copy_in_error(PgErrorCode::Ec0A000, message) } @@ -174,11 +183,11 @@ pub(crate) fn copy_from_stdin_from_sql(sql: &str) -> PgWireResult raw, }; let mut option_set = raw.options; - let codec = resolve_codec(&mut option_set)?; + let format = resolve_format(&mut option_set)?; Ok(Some(CopyFromStdin { table: raw.table, columns: raw.columns, - codec, + format, })) } @@ -221,26 +230,49 @@ fn resolve_table_name( } /// Per-connection state of an in-progress COPY FROM STDIN. +/// +/// The payload shape follows the codec: table-bound formats (csv) parse +/// records into rows against the targeted table's columns; table-less +/// formats (influxdb line protocol) split the stream into complete +/// lines that each route to their own measurement, and flush through +/// the handler's line-protocol pipeline — the same one the HTTP write +/// endpoint uses. pub(crate) struct CopyInState { query_ctx: QueryContextRef, - /// The gRPC column schemas of the targeted columns. Cloned into each - /// flushed batch: `Rows.schema` is an owned `Vec` in the proto - /// request, so one copy per request is inherent. - grpc_schema: Vec, - table_name: String, - parser: Box, + payload: Payload, /// The resolved codec's name, for logging. codec_name: &'static str, /// The wire format code of the resolved codec. codec_format_code: i8, - /// Parsed rows not yet flushed to the handler. - rows: Vec, rows_written: u64, } +enum Payload { + /// Table-bound rows: parsed against `grpc_schema` and flushed to + /// `table_name` via `copy_in_insert`. The schema is cloned into each + /// flushed batch: `Rows.schema` is an owned `Vec` in the proto + /// request, so one copy per request is inherent. + Table { + grpc_schema: Vec, + table_name: String, + parser: Box, + /// Parsed rows not yet flushed to the handler. + rows: Vec, + }, + /// The whole buffered influxdb line-protocol payload with its + /// request precision. One COPY statement is one bounded ingestion + /// request, buffered until `CopyDone` — exactly like one HTTP write + /// body — and record boundaries belong to the line-protocol parser, + /// never to the COPY layer. + Influxdb { + data: Vec, + precision: Option, + }, +} + impl CopyInState { fn new( - codec: Arc, + codec: &CsvCodec, columns: Vec, table_info: Arc, query_ctx: QueryContextRef, @@ -249,7 +281,7 @@ impl CopyInState { let schema = table_info.meta.schema.clone(); let timestamp_index = schema.timestamp_index(); - let columns = resolve_columns(codec.as_ref(), columns, &schema, &table_info)?; + let columns = resolve_columns(codec, columns, &schema, &table_info)?; if let Some(ts_index) = timestamp_index { let ts_column = &schema.column_schemas()[ts_index]; if !columns.iter().any(|column| column.name == ts_column.name) { @@ -268,58 +300,122 @@ impl CopyInState { Ok(Self { query_ctx, - grpc_schema, - table_name, - parser, - codec_name: codec.name(), - codec_format_code: codec.format_code(), - rows: Vec::with_capacity(COPY_IN_FLUSH_ROWS), + payload: Payload::Table { + grpc_schema, + table_name, + parser, + rows: Vec::with_capacity(COPY_IN_FLUSH_ROWS), + }, + codec_name: "csv", + codec_format_code: 0, rows_written: 0, }) } + /// State for the influxdb format: the whole payload is buffered and + /// completed once at `CopyDone` through the line-protocol pipeline. + fn new_influxdb(options: &InfluxdbOptions, query_ctx: QueryContextRef) -> Self { + Self { + query_ctx, + payload: Payload::Influxdb { + data: Vec::new(), + precision: options.precision, + }, + codec_name: "influxdb", + codec_format_code: 0, + rows_written: 0, + } + } + /// The wire format code advertised to the client. pub(crate) fn format_code(&self) -> i8 { self.codec_format_code } - /// Feeds a `CopyData` chunk into the parser and buffers parsed rows. + /// Feeds a `CopyData` chunk into the parser and buffers parsed rows + /// (csv) or payload bytes (influxdb). fn feed(&mut self, data: &[u8]) -> PgWireResult<()> { - self.rows.extend(self.parser.feed(data)?); + match &mut self.payload { + Payload::Table { parser, rows, .. } => rows.extend(parser.feed(data)?), + Payload::Influxdb { data: buffered, .. } => { + if buffered.len() + data.len() > COPY_IN_MAX_PAYLOAD_BYTES { + return Err(bad_copy_data(format!( + "COPY FROM STDIN payload exceeds the {} limit; split the import into multiple COPY statements", + COPY_IN_MAX_PAYLOAD_BYTES + ))); + } + buffered.extend_from_slice(data); + } + } Ok(()) } /// Parses the buffered tail at `CopyDone`; a final record without a /// trailing newline is accepted. fn finish(&mut self) -> PgWireResult<()> { - self.rows.extend(self.parser.finish()?); + match &mut self.payload { + Payload::Table { parser, rows, .. } => rows.extend(parser.finish()?), + // The whole payload is one request body; nothing to finalize. + Payload::Influxdb { .. } => {} + } Ok(()) } fn should_flush(&self) -> bool { - self.rows.len() >= COPY_IN_FLUSH_ROWS + match &self.payload { + Payload::Table { rows, .. } => rows.len() >= COPY_IN_FLUSH_ROWS, + // Completed once at CopyDone, not mid-stream. + Payload::Influxdb { .. } => false, + } } /// Takes the buffered rows as an insert request, if any. fn take_requests(&mut self) -> Option { - if self.rows.is_empty() { + let Payload::Table { + grpc_schema, + table_name, + rows, + .. + } = &mut self.payload + else { + return None; + }; + if rows.is_empty() { return None; } let rows = GrpcRows { - schema: self.grpc_schema.clone(), - rows: std::mem::take(&mut self.rows), + schema: grpc_schema.clone(), + rows: std::mem::take(rows), }; Some(RowInsertRequests { inserts: vec![RowInsertRequest { - table_name: self.table_name.clone(), + table_name: table_name.clone(), rows: Some(rows), }], }) } + + /// Takes the buffered payload as one line-protocol request body, if + /// any — the same bytes an HTTP write of the same records would + /// carry. + fn take_lines(&mut self) -> PgWireResult> { + let Payload::Influxdb { data, precision } = &mut self.payload else { + return Ok(None); + }; + if data.is_empty() { + return Ok(None); + } + let lines = String::from_utf8(std::mem::take(data)) + .map_err(|e| bad_copy_data(format!("COPY FROM STDIN data is not valid UTF-8: {e}")))?; + Ok(Some(InfluxdbRequest { + precision: *precision, + lines, + })) + } } fn resolve_columns( - codec: &dyn CopyInCodec, + codec: &CsvCodec, columns: Vec, schema: &SchemaRef, table_info: &Arc, @@ -411,27 +507,63 @@ impl PostgresServerHandlerInner { // pending-rows batcher like SQL INSERT when it is enabled for // this connection. let query_ctx = self.new_query_context(); - let (catalog, schema, table) = resolve_table_name(&query_ctx, &stmt.table)?; - let Some(table_info) = self - .copy_in_handler - .copy_in_table(&catalog, &schema, &table, query_ctx.clone()) - .await - .map_err(convert_err)? - else { - return Err(copy_in_error( - PgErrorCode::Ec42P01, - format!("relation \"{}\" does not exist", stmt.table.join(".")), - )); + // Each format states its own target rule and ingestion shape + // explicitly. + let state = match &stmt.format { + // csv requires a target table: records are parsed against + // its columns and inserted in batches. + CopyFormat::Csv(codec) => { + if stmt.table.is_empty() { + return Err(unsupported( + "format csv requires a table: COPY tbl [(cols)] FROM STDIN WITH (FORMAT csv)", + )); + } + let (catalog, schema, table) = resolve_table_name(&query_ctx, &stmt.table)?; + + let Some(table_info) = self + .copy_in_handler + .copy_in_table(&catalog, &schema, &table, query_ctx.clone()) + .await + .map_err(convert_err)? + else { + return Err(copy_in_error( + PgErrorCode::Ec42P01, + format!("relation \"{}\" does not exist", stmt.table.join(".")), + )); + }; + + // The insert path resolves the bare table name against the + // query context, so point the context at the resolved + // catalog and schema when the COPY target is qualified. + let query_ctx = adjust_query_context(query_ctx, &catalog, &schema); + + CopyInState::new(codec, stmt.columns, table_info, query_ctx, table)? + } + // influxdb forbids a target table — each line names its own + // measurement — and completes through the line-protocol + // pipeline the HTTP write endpoints use; the server must be + // wired for it. + CopyFormat::Influxdb(options) => { + if !stmt.table.is_empty() { + return Err(unsupported(format!( + "format influxdb routes each record to its own table; \"COPY {} FROM STDIN\" must not name one", + stmt.table.join(".") + ))); + } + if self.influxdb_handler.is_none() { + return Err(unsupported( + "this server is not wired for the influxdb line protocol format", + )); + } + CopyInState::new_influxdb(options, query_ctx) + } + }; + let num_columns = match &state.payload { + Payload::Table { grpc_schema, .. } => grpc_schema.len(), + // No fixed column set: the format routes per record. + Payload::Influxdb { .. } => 0, }; - - // The insert path resolves the bare table name against the query - // context, so point the context at the resolved catalog and schema - // when the COPY target is qualified. - let query_ctx = adjust_query_context(query_ctx, &catalog, &schema); - - let state = CopyInState::new(stmt.codec, stmt.columns, table_info, query_ctx, table)?; - let num_columns = state.grpc_schema.len(); let format_code = state.format_code(); *self.copy_in_state.lock().await = Some(state); @@ -443,13 +575,34 @@ impl PostgresServerHandlerInner { ))) } - /// Flushes buffered rows of an in-progress copy. Returns the number of - /// rows written by the flushed batch. + /// Flushes the buffered payload of an in-progress copy. Returns the + /// number of rows written by the flushed batch. async fn flush_copy_in(&self, state: &mut CopyInState) -> PgWireResult { + let query_ctx = state.query_ctx.clone(); + // The buffered influxdb payload is completed through the + // line-protocol pipeline — the same handler the HTTP write + // endpoints use, so conversion, permissions, interceptors, + // timestamp alignment and inserts are identical to an HTTP + // write of the same bytes. The context stays this connection's + // postgres context. + if let Some(request) = state.take_lines()? { + let Some(influxdb_handler) = &self.influxdb_handler else { + return Err(unsupported( + "this server is not wired for the influxdb line protocol format", + )); + }; + let output = influxdb_handler + .exec(request, query_ctx) + .await + .map_err(convert_err)?; + return match output.data { + OutputData::AffectedRows(rows) => Ok(rows as u64), + OutputData::RecordBatches(_) | OutputData::Stream(_) => Ok(0), + }; + } let Some(requests) = state.take_requests() else { return Ok(0); }; - let query_ctx = state.query_ctx.clone(); let output = self .copy_in_handler .copy_in_insert(requests, query_ctx) @@ -587,8 +740,8 @@ mod tests { .unwrap(); assert_eq!(copy.table, vec!["t"]); assert!(copy.columns.is_empty()); - assert_eq!(copy.codec.name(), "csv"); - assert_eq!(copy.codec.format_code(), 0); + assert!(matches!(copy.format, CopyFormat::Csv(_))); + // Both formats are textual: format code 0. } #[test] @@ -599,7 +752,7 @@ mod tests { .unwrap() .unwrap(); assert_eq!(copy.columns, vec!["a", "b"]); - assert_eq!(copy.codec.name(), "csv"); + assert!(matches!(copy.format, CopyFormat::Csv(_))); } #[test] diff --git a/src/servers/src/postgres/copy_in/parse.rs b/src/servers/src/postgres/copy_in/parse.rs index 8139960bfa4..e9be443eea5 100644 --- a/src/servers/src/postgres/copy_in/parse.rs +++ b/src/servers/src/postgres/copy_in/parse.rs @@ -32,7 +32,7 @@ use pgwire::error::PgWireResult; use crate::postgres::copy_in::options::{CopyOptionSet, CopyOptionValue}; use crate::postgres::copy_in::{bad_copy_data, copy_in_must_run_alone_error}; -/// A scanned `COPY tbl [(cols)] FROM STDIN [WITH (...)]` statement. +/// A scanned `COPY [tbl [(cols)]] FROM STDIN [WITH (...)]` statement. /// /// Unquoted identifiers are folded to lowercase and quoted ones keep their /// case (`COPY Metrics` addresses the same table as `CREATE TABLE Metrics`, @@ -41,7 +41,9 @@ use crate::postgres::copy_in::{bad_copy_data, copy_in_must_run_alone_error}; #[derive(Debug)] pub(super) struct RawCopyFromStdin { /// Table name parts, as written (`t`, `schema.t`, - /// `catalog.schema.t`); quoted names keep their case. + /// `catalog.schema.t`); quoted names keep their case. Empty for + /// table-less formats whose records carry their own routing (e.g. + /// the influxdb line protocol names its measurement per line). pub(super) table: Vec, /// Requested column list; empty means all columns in table order. pub(super) columns: Vec, @@ -89,13 +91,27 @@ pub(super) fn scan_copy_from_stdin(sql: &str) -> PgWireResult { return Ok(CopyInScan::NotCopyIn); } - let table = match parse_object_name(&mut cursor) { - Some(table) => table, - None => return Ok(CopyInScan::NotCopyIn), + // Table-less statements (`COPY FROM STDIN WITH (FORMAT influxdb) ...`) + // are table-format specific: the codec decides whether a missing + // table is legal. An unquoted `FROM` cannot be a table name (SQL + // would require quoting it), so peek instead of letting the + // identifier parser swallow the keyword. A column list without a + // table is not. + let table = if cursor.peek_keyword("from") { + Vec::new() + } else { + match parse_object_name(&mut cursor) { + Some(table) => table, + None => return Ok(CopyInScan::NotCopyIn), + } }; - let columns = match parse_column_list(&mut cursor) { - Some(columns) => columns, - None => return Ok(CopyInScan::NotCopyIn), + let columns = if table.is_empty() { + Vec::new() + } else { + match parse_column_list(&mut cursor) { + Some(columns) => columns, + None => return Ok(CopyInScan::NotCopyIn), + } }; if !cursor.take_keyword("from") || !cursor.take_keyword("stdin") { @@ -151,16 +167,20 @@ impl<'a> Tokens<'a> { /// Consumes the next token if it is the unquoted keyword `keyword`. fn take_keyword(&mut self, keyword: &str) -> bool { - match self.peek() { - Some(Token::Word(w)) - if w.quote_style.is_none() && w.value.eq_ignore_ascii_case(keyword) => - { - self.pos += 1; - true - } - _ => false, + if self.peek_keyword(keyword) { + self.pos += 1; + true + } else { + false } } + + /// Whether the next token is the unquoted keyword `keyword`, without + /// consuming it. + fn peek_keyword(&self, keyword: &str) -> bool { + matches!(self.tokens.get(self.pos), Some(Token::Word(w)) + if w.quote_style.is_none() && w.value.eq_ignore_ascii_case(keyword)) + } } /// Returns the identifier with the SQL parser's canonicalization: quoted @@ -275,6 +295,27 @@ mod tests { }) } + #[test] + fn test_scan_tableless() { + // Table-less statements: the codec decides whether a missing + // table is legal (the influxdb line protocol names its + // measurement per line). + let raw = scan("COPY FROM STDIN WITH (FORMAT influxdb, precision 'ns')") + .unwrap() + .unwrap(); + assert!(raw.table.is_empty()); + assert!(raw.columns.is_empty()); + + let raw = scan("copy from stdin (format influxdb)").unwrap().unwrap(); + assert!(raw.table.is_empty()); + + // A column list without a table is not a copy-in statement. + assert!(scan("COPY (a, b) FROM STDIN").unwrap().is_none()); + // COPY OUT and file COPY are unaffected. + assert!(scan("COPY (SELECT 1) TO STDOUT").unwrap().is_none()); + assert!(scan("COPY t FROM '/tmp/x.csv'").unwrap().is_none()); + } + #[test] fn test_scan_basic() { let raw = scan("COPY t FROM STDIN").unwrap().unwrap(); diff --git a/src/servers/src/postgres/server.rs b/src/servers/src/postgres/server.rs index 7e3b3e2b7bd..3919081ba62 100644 --- a/src/servers/src/postgres/server.rs +++ b/src/servers/src/postgres/server.rs @@ -30,6 +30,7 @@ use tokio_rustls::TlsAcceptor; use crate::error::Result; use crate::postgres::{MakePostgresServerHandler, MakePostgresServerHandlerBuilder}; use crate::query_handler::CopyInHandlerRef; +use crate::query_handler::InfluxdbLineProtocolHandlerRef; use crate::query_handler::sql::ServerSqlQueryHandlerRef; use crate::server::{AbortableStream, BaseTcpServer, Server}; use crate::tls::ReloadableTlsServerConfig; @@ -40,6 +41,7 @@ pub struct PostgresServer { tls_server_config: Arc, keep_alive_secs: u64, batching_enabled: bool, + influxdb_handler: Option, bind_addr: Option, process_manager: Option, } @@ -77,6 +79,7 @@ impl PostgresServer { tls_server_config, keep_alive_secs, batching_enabled: false, + influxdb_handler: None, bind_addr: None, process_manager, } @@ -88,12 +91,34 @@ impl PostgresServer { self } + /// Wires the influxdb line-protocol handler used by table-less + /// `COPY FROM STDIN WITH (FORMAT influxdb)` flushes — the same + /// handler the HTTP influxdb write endpoints use, so a flushed batch + /// runs through the identical pipeline. + pub fn with_influxdb_handler(mut self, handler: InfluxdbLineProtocolHandlerRef) -> Self { + self.influxdb_handler = Some(handler); + self + } + fn accept( &self, io_runtime: Runtime, accepting_stream: AbortableStream, ) -> impl Future + use<> { - let handler_maker = self.make_handler.clone(); + // Assemble the per-connection handler directly, including any + // `with_*` wiring applied after construction (the influxdb + // line-protocol handler). Constructing the struct — instead of + // going through the builder — keeps this infallible: no + // `.build().unwrap()` on the hot accept path. + let make_handler = &self.make_handler; + let handler_maker = Arc::new(MakePostgresServerHandler { + query_handler: make_handler.query_handler.clone(), + copy_in_handler: make_handler.copy_in_handler.clone(), + influxdb_handler: self.influxdb_handler.clone(), + user_provider: make_handler.user_provider.clone(), + param_provider: make_handler.param_provider.clone(), + force_tls: make_handler.force_tls, + }); let batching_enabled = self.batching_enabled; let tls_server_config = self.tls_server_config.clone(); let process_manager = self.process_manager.clone(); diff --git a/tests-integration/src/test_util.rs b/tests-integration/src/test_util.rs index 3cc820619c6..3dd451b4716 100644 --- a/tests-integration/src/test_util.rs +++ b/tests-integration/src/test_util.rs @@ -1204,16 +1204,19 @@ pub async fn setup_pg_server_with_slow_query_threshold( .expect("Failed to load certificates and keys"), ); - let mut pg_server = Box::new(PostgresServer::new( - fe_instance_ref.clone(), - fe_instance_ref.clone(), - opts.tls.should_force_tls(), - tls_server_config, - 0, - runtime, - None, - None, - )); + let mut pg_server = Box::new( + PostgresServer::new( + fe_instance_ref.clone(), + fe_instance_ref.clone(), + opts.tls.should_force_tls(), + tls_server_config, + 0, + runtime, + None, + None, + ) + .with_influxdb_handler(fe_instance_ref.clone()), + ); pg_server .start(fe_pg_addr.parse::().unwrap()) @@ -1248,16 +1251,19 @@ pub async fn setup_pg_server_with_user_provider( .expect("Failed to load certificates and keys"), ); - let mut pg_server = Box::new(PostgresServer::new( - fe_instance_ref.clone(), - fe_instance_ref.clone(), - opts.tls.should_force_tls(), - tls_server_config, - 0, - runtime, - user_provider, - None, - )); + let mut pg_server = Box::new( + PostgresServer::new( + fe_instance_ref.clone(), + fe_instance_ref.clone(), + opts.tls.should_force_tls(), + tls_server_config, + 0, + runtime, + user_provider, + None, + ) + .with_influxdb_handler(fe_instance_ref.clone()), + ); pg_server .start(fe_pg_addr.parse::().unwrap()) @@ -1304,16 +1310,19 @@ pub async fn setup_pg_server_with_prom_native_histogram( .expect("Failed to load certificates and keys"), ); - let mut pg_server = Box::new(PostgresServer::new( - instance.fe_instance().clone(), - instance.fe_instance().clone(), - opts.tls.should_force_tls(), - tls_server_config, - 0, - runtime, - None, - None, - )); + let mut pg_server = Box::new( + PostgresServer::new( + instance.fe_instance().clone(), + instance.fe_instance().clone(), + opts.tls.should_force_tls(), + tls_server_config, + 0, + runtime, + None, + None, + ) + .with_influxdb_handler(instance.fe_instance().clone()), + ); pg_server .start(fe_pg_addr.parse::().unwrap()) diff --git a/tests-integration/src/tests/postgres_copy_in_test.rs b/tests-integration/src/tests/postgres_copy_in_test.rs index 82e076e7a65..94b19c1315d 100644 --- a/tests-integration/src/tests/postgres_copy_in_test.rs +++ b/tests-integration/src/tests/postgres_copy_in_test.rs @@ -316,6 +316,7 @@ async fn test_postgres_copy_in_failed_copy_then_new_copy() { Err(e) => e, }; assert!(err.as_db_error().is_some(), "{err}"); + tokio::time::sleep(std::time::Duration::from_millis(200)).await; // The connection is usable and a NEW copy succeeds on it. let rows = feed_copy_in( @@ -519,3 +520,302 @@ async fn test_postgres_copy_in_empty_field_is_null() { server.shutdown().await.unwrap(); } + +/// `COPY FROM STDIN WITH (FORMAT influxdb)` ingests line protocol through +/// the same pipeline as the HTTP influxdb write endpoint: lines route to +/// their own measurements, tables are auto-created, and a batch behaves +/// like an HTTP write of the same bytes. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_postgres_copy_in_influxdb_line_protocol() { + init_default_ut_logging(); + let (_guard, server) = setup_pg_server(StorageType::File, "pg_copy_in_influxdb").await; + let port = server.bind_addr().unwrap().port(); + let client = connect(port).await; + + let payload = b"monitor,host=hosta cpu=66.6,memory=1024 1663840496100023100\n\ + monitor,host=hostb memory=1027 1663840496400340001\n\ + weather,location=us-midwest temperature=82 1663840496100023102\n"; + + // Feed byte-by-byte-ish: lines and UTF-8 may span CopyData chunks. + let rows = feed_copy_in( + &client, + "COPY FROM STDIN WITH (FORMAT influxdb)", + payload.chunks(7).map(Bytes::copy_from_slice).collect(), + ) + .await + .unwrap(); + assert_eq!(rows, 3); + + let messages = client + .simple_query("SELECT host, cpu, memory FROM monitor ORDER BY host") + .await + .unwrap(); + let rows: Vec<(String, String, String)> = messages + .iter() + .filter_map(|m| { + let tokio_postgres::SimpleQueryMessage::Row(row) = m else { + return None; + }; + Some(( + row.get(0).unwrap_or("").to_string(), + row.get(1).unwrap_or("").to_string(), + row.get(2).unwrap_or("").to_string(), + )) + }) + .collect(); + assert_eq!( + rows, + vec![ + ("hosta".into(), "66.6".into(), "1024.0".into()), + ("hostb".into(), "".into(), "1027.0".into()), + ] + ); + + let messages = client + .simple_query("SELECT temperature FROM weather") + .await + .unwrap(); + let temps: Vec = messages + .iter() + .filter_map(|m| match m { + tokio_postgres::SimpleQueryMessage::Row(row) => row.get(0).map(str::to_string), + _ => None, + }) + .collect(); + assert_eq!(temps, vec!["82.0".to_string()]); + + server.shutdown().await.unwrap(); +} + +/// The `precision` option carries the HTTP endpoint's semantics: +/// nanosecond timestamps by default, other units on request. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_postgres_copy_in_influxdb_precision() { + init_default_ut_logging(); + let (_guard, server) = setup_pg_server(StorageType::File, "pg_copy_in_influxdb_prec").await; + let port = server.bind_addr().unwrap().port(); + let client = connect(port).await; + + let rows = feed_copy_in( + &client, + "COPY FROM STDIN WITH (FORMAT influxdb, precision 'ms')", + vec![Bytes::from_static(b"prec,host=h v=1 1465839830100\n")], + ) + .await + .unwrap(); + assert_eq!(rows, 1); + + let messages = client + .simple_query("SELECT CAST(greptime_timestamp AS BIGINT) FROM prec") + .await + .unwrap(); + let ts: Vec = messages + .iter() + .filter_map(|m| match m { + tokio_postgres::SimpleQueryMessage::Row(row) => row.get(0).map(str::to_string), + _ => None, + }) + .collect(); + // 1465839830100 ms == 1465839830100000000 ns + assert_eq!(ts.len(), 1); + let ts_ns: i64 = ts[0].parse().unwrap(); + assert_eq!(ts_ns, 1465839830100000000); + + server.shutdown().await.unwrap(); +} + +/// Records with newlines INSIDE quoted field values are one line-protocol +/// record: the parser owns record boundaries, so a CopyData chunk boundary +/// inside the quoted string — the case a naive line splitter breaks on — +/// must not split it. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_postgres_copy_in_influxdb_multiline_quoted_field() { + init_default_ut_logging(); + let (_guard, server) = setup_pg_server(StorageType::File, "pg_copy_in_influxdb_ml").await; + let port = server.bind_addr().unwrap().port(); + let client = connect(port).await; + + // A record whose string field value contains a physical newline, and + // one with a backslash right before the newline (the parser skips the + // escaped newline; also not a record boundary). + let payload: Vec = b"mlog,host=a line=\"first\nsecond\" 100\n\ + mlog,host=b line=\"tail\\\nmore\" 200\n" + .to_vec(); + + // Split exactly on the internal newline of the first record: chunk 1 + // ends after `first`, chunk 2 starts with the newline. + let split_at = payload.windows(6).position(|w| w == b"first\n").unwrap() + b"first".len(); + let (head, tail) = payload.split_at(split_at); + let rows = feed_copy_in( + &client, + "COPY FROM STDIN WITH (FORMAT influxdb)", + vec![Bytes::copy_from_slice(head), Bytes::copy_from_slice(tail)], + ) + .await + .unwrap(); + assert_eq!(rows, 2); + + let messages = client + .simple_query("SELECT line FROM mlog ORDER BY host") + .await + .unwrap(); + let lines: Vec = messages + .iter() + .filter_map(|m| match m { + tokio_postgres::SimpleQueryMessage::Row(row) => row.get(0).map(str::to_string), + _ => None, + }) + .collect(); + // The quoted values keep their physical newlines: `first`, newline, + // `second`; and `tail`, backslash, newline, `more` (a backslash only + // escapes `\` and `"` in line-protocol strings, so it stays literal + // before the newline). + assert_eq!( + lines, + vec!["first\nsecond".to_string(), "tail\\\nmore".to_string()] + ); + + server.shutdown().await.unwrap(); +} + +/// The payload is bounded by the HTTP write endpoints' own body limit. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_postgres_copy_in_influxdb_payload_limit() { + init_default_ut_logging(); + let (_guard, server) = setup_pg_server(StorageType::File, "pg_copy_in_influxdb_cap").await; + let port = server.bind_addr().unwrap().port(); + let client = connect(port).await; + + // The HTTP write endpoints' default body limit (64 MiB) — kept in + // sync with servers::http::DEFAULT_BODY_LIMIT. + let limit = 64 * 1024 * 1024; + // Several chunks over the limit, none containing a record boundary. + let chunk = vec![b'x'; 1024 * 1024]; + let mut sent = 0usize; + let mut sink = Box::pin( + client + .copy_in("COPY FROM STDIN WITH (FORMAT influxdb)") + .await + .unwrap(), + ); + while sent <= limit { + sink.as_mut() + .send(Bytes::from(chunk.clone())) + .await + .unwrap(); + sent += chunk.len(); + } + // The cap fires mid-stream; the error surfaces when the copy + // completes (the sink buffers sends locally). + let error = sink + .as_mut() + .finish() + .await + .expect_err("expected the oversized payload to fail"); + assert!( + error + .as_db_error() + .map(|d| d.message().contains("exceeds")) + .unwrap_or(false), + "unexpected error: {error}" + ); + + // The connection stays usable. + let messages = client.simple_query("SELECT 1").await.unwrap(); + assert!(!messages.is_empty()); + + server.shutdown().await.unwrap(); +} + +/// The regression for the retired line-splitter's failure mode: with +/// 32,767 ordinary records plus a multiline quoted record whose first +/// physical line lands exactly on the old 32,768-line flush threshold — +/// and a CopyData chunk boundary at the quoted value's internal newline — +/// the whole payload must still ingest as one request. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_postgres_copy_in_influxdb_flush_threshold_multiline() { + init_default_ut_logging(); + let (_guard, server) = + setup_pg_server(StorageType::File, "pg_copy_in_influxdb_threshold").await; + let port = server.bind_addr().unwrap().port(); + let client = connect(port).await; + + // 32,767 ordinary records, then a record whose quoted value spans + // three physical lines. + let mut head = String::new(); + for i in 0..32_767 { + head.push_str(&format!( + "m,host=h{i} v={i} {} +", + 1_000 + i + )); + } + head.push_str("mlog,host=z line=\"first\nsecond\nthird\" 99999"); + let tail = " +" + .to_string(); + + let rows = feed_copy_in( + &client, + "COPY FROM STDIN WITH (FORMAT influxdb)", + vec![ + Bytes::from(head.into_bytes()), + Bytes::from(tail.into_bytes()), + ], + ) + .await + .unwrap(); + assert_eq!(rows, 32_768); + + let messages = client.simple_query("SELECT line FROM mlog").await.unwrap(); + let lines: Vec = messages + .iter() + .filter_map(|m| match m { + tokio_postgres::SimpleQueryMessage::Row(row) => row.get(0).map(str::to_string), + _ => None, + }) + .collect(); + assert_eq!(lines, vec!["first\nsecond\nthird".to_string()]); + + server.shutdown().await.unwrap(); +} + +/// Invalid UTF-8 fails the copy with a clear error and the connection +/// stays usable; a trailing record without a newline is still ingested. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_postgres_copy_in_influxdb_invalid_utf8_and_trailing_record() { + init_default_ut_logging(); + let (_guard, server) = setup_pg_server(StorageType::File, "pg_copy_in_influxdb_utf8").await; + let port = server.bind_addr().unwrap().port(); + let client = connect(port).await; + + // Trailing record without a newline: accepted. + let rows = feed_copy_in( + &client, + "COPY FROM STDIN WITH (FORMAT influxdb)", + vec![Bytes::from_static(b"m2,host=a v=7 300")], + ) + .await + .unwrap(); + assert_eq!(rows, 1); + + // Invalid UTF-8: the copy fails, the connection stays usable. + let error = feed_copy_in( + &client, + "COPY FROM STDIN WITH (FORMAT influxdb)", + vec![Bytes::from_static(b"m3 v=1 100\n\xff\xfe")], + ) + .await + .unwrap_err(); + assert!( + error + .as_db_error() + .map(|d| d.message().contains("not valid UTF-8")) + .unwrap_or(false), + "unexpected error: {error}" + ); + let messages = client.simple_query("SELECT 1").await.unwrap(); + assert!(!messages.is_empty()); + + server.shutdown().await.unwrap(); +} From df33737e9d093d2415df363f2942ff62ef904597 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Wed, 28 Oct 2026 08:45:10 +0000 Subject: [PATCH 123/239] feat: add shared asynchronous cache primitives (#9790) Signed-off-by: WenyXu --- Cargo.lock | 10 + Cargo.toml | 1 + src/common/cache/Cargo.toml | 16 ++ src/common/cache/src/cache.rs | 473 ++++++++++++++++++++++++++++++++ src/common/cache/src/lib.rs | 21 ++ src/common/cache/src/metrics.rs | 194 +++++++++++++ 6 files changed, 715 insertions(+) create mode 100644 src/common/cache/Cargo.toml create mode 100644 src/common/cache/src/cache.rs create mode 100644 src/common/cache/src/lib.rs create mode 100644 src/common/cache/src/metrics.rs diff --git a/Cargo.lock b/Cargo.lock index 69996d9cbb2..8e4183994d4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2340,6 +2340,16 @@ dependencies = [ "tokio", ] +[[package]] +name = "common-cache" +version = "1.4.0-alpha.0" +dependencies = [ + "futures", + "moka", + "prometheus 0.14.0", + "tokio", +] + [[package]] name = "common-catalog" version = "1.4.0-alpha.0" diff --git a/Cargo.toml b/Cargo.toml index 15539abb2c7..a5e3db37960 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,6 +9,7 @@ members = [ "src/cmd", "src/common/base", "src/common/batcher", + "src/common/cache", "src/common/catalog", "src/common/config", "src/common/datasource", diff --git a/src/common/cache/Cargo.toml b/src/common/cache/Cargo.toml new file mode 100644 index 00000000000..1e2f9d83a5a --- /dev/null +++ b/src/common/cache/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "common-cache" +version.workspace = true +edition.workspace = true +license.workspace = true + +[lints] +workspace = true + +[dependencies] +moka = { workspace = true, features = ["future"] } +prometheus.workspace = true + +[dev-dependencies] +futures.workspace = true +tokio = { workspace = true, features = ["macros", "rt", "sync"] } diff --git a/src/common/cache/src/cache.rs b/src/common/cache/src/cache.rs new file mode 100644 index 00000000000..98efb890f29 --- /dev/null +++ b/src/common/cache/src/cache.rs @@ -0,0 +1,473 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::borrow::Borrow; +use std::hash::Hash; +use std::sync::Arc; + +use crate::metrics::{CacheMetrics, CacheSizeMetric, CacheSizeRegistration}; + +/// Reason for releasing an entry from the cache. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum RemovalCause { + /// The entry reached its expiration time. + Expired, + /// The entry was explicitly invalidated. + Explicit, + /// A new value replaced the entry. + Replaced, + /// The entry was removed to satisfy the capacity policy. + Size, +} + +impl RemovalCause { + /// Returns the established eviction metric label. + pub fn as_str(self) -> &'static str { + match self { + Self::Expired => "expired", + Self::Explicit => "explicit", + Self::Replaced => "replaced", + Self::Size => "size", + } + } +} + +impl From for RemovalCause { + fn from(value: moka::notification::RemovalCause) -> Self { + match value { + moka::notification::RemovalCause::Expired => Self::Expired, + moka::notification::RemovalCause::Explicit => Self::Explicit, + moka::notification::RemovalCause::Replaced => Self::Replaced, + moka::notification::RemovalCause::Size => Self::Size, + } + } +} + +/// A weighted memory cache with backend-provided asynchronous load deduplication. +/// Values should be inexpensive to clone, such as `Arc` or `Bytes`. +/// Capacity is maintained asynchronously; insertion does not guarantee retention. +#[derive(Clone)] +pub struct Cache { + inner: moka::future::Cache, + capacity: u64, + size_registration: Option, + metrics: Option>, +} + +impl Cache +where + K: Eq + Hash + Send + Sync + 'static, + V: Clone + Send + Sync + 'static, +{ + /// Creates a cache with the caller's weight and eviction accounting. + pub fn new(name: &str, capacity: u64, weigher: W, listener: L) -> Self + where + W: Fn(&K, &V) -> u32 + Send + Sync + 'static, + L: Fn(Arc, V, RemovalCause) + Send + Sync + 'static, + { + Self { + inner: moka::future::Cache::builder() + .name(name) + .max_capacity(capacity) + .support_invalidation_closures() + .weigher(weigher) + .eviction_listener(move |key, value, cause| listener(key, value, cause.into())) + .build(), + capacity, + size_registration: None, + metrics: None, + } + } + + /// Attaches caller-bound metrics without changing existing hit/miss accounting. + pub fn with_metrics(mut self, metrics: Arc) -> Self { + self.metrics = Some(metrics); + self + } + + /// Samples this cache's retained weight when the caller collects byte metrics. + pub fn with_size_metric(mut self, metric: CacheSizeMetric) -> Self { + let cache = self.inner.clone(); + self.size_registration = Some(metric.register(move || cache.weighted_size())); + self + } + + /// Gets a completed value without starting a load. + pub fn get<'a, Q>(&'a self, key: &'a Q) -> impl Future> + Send + 'a + where + K: Borrow, + Q: Hash + Eq + Sync + ?Sized, + { + self.inner.get(key) + } + + /// Shares a successful load among concurrent callers with the same key and error type. + /// Errors are shared with current waiters but are not retained as cached values. + /// Cancelling the active loader lets a remaining caller execute its own loader. + pub async fn try_get_with(&self, key: &Q, load: F) -> Result> + where + K: Borrow, + Q: ToOwned + Hash + Eq + ?Sized, + F: Future>, + E: Send + Sync + 'static, + { + let request = self + .metrics + .as_ref() + .map(|metrics| metrics.requests.start()); + let result = self + .inner + .try_get_with_by_ref(key, async { + // Only the backend-selected loader contributes physical load metrics. + let operation = self.metrics.as_ref().map(|metrics| metrics.loads.start()); + let result = load.await; + if let Some(operation) = operation { + operation.finish(&result); + } + result + }) + .await; + if let Some(request) = request { + request.finish(&result); + } + result + } + + /// Submits a completed value to the cache's admission and eviction policy. + pub async fn insert(&self, key: K, value: V) { + self.inner.insert(key, value).await; + } + + /// Invalidates a completed entry. This does not cancel a concurrent loader. + pub async fn invalidate(&self, key: &Q) + where + K: Borrow, + Q: Hash + Eq + ?Sized, + { + self.inner.invalidate(key).await; + } + + /// Invalidates completed entries matching a predicate. Pending loaders are not fenced. + /// Physical removal and eviction accounting run asynchronously. + pub fn invalidate_if(&self, predicate: F) -> std::io::Result<()> + where + F: Fn(&K, &V) -> bool + Send + Sync + 'static, + { + self.inner + .invalidate_entries_if(predicate) + .map(|_| ()) + .map_err(std::io::Error::other) + } + + /// Returns the configured weight capacity. + pub fn capacity(&self) -> u64 { + self.capacity + } + + /// Returns the approximate number of retained entries. + pub fn entry_count(&self) -> u64 { + self.inner.entry_count() + } + + /// Returns the approximate total retained weight. + pub fn weighted_size(&self) -> u64 { + self.inner.weighted_size() + } + + /// Applies pending admission, eviction, and accounting work. + pub async fn run_pending_tasks(&self) { + self.inner.run_pending_tasks().await; + } +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use tokio::sync::Notify; + + use super::*; + + // This is also a compile-time contract for generic callers crossing async-trait boundaries. + fn send_lookup<'a, K, V, Q>( + cache: &'a Cache, + key: &'a Q, + ) -> impl Future> + Send + 'a + where + K: Borrow + Eq + Hash + Send + Sync + 'static, + V: Clone + Send + Sync + 'static, + Q: Hash + Eq + Sync + ?Sized, + { + cache.get(key) + } + + #[tokio::test] + async fn test_borrowed_lookup_has_a_send_future() { + let cache = cache(1024); + cache + .insert("key".to_string(), Arc::new("value".to_string())) + .await; + let future: futures::future::BoxFuture<'_, Option>> = + Box::pin(send_lookup(&cache, "key")); + assert_eq!(future.await.as_deref().map(String::as_str), Some("value")); + } + + #[tokio::test] + async fn test_predicate_invalidation_preserves_other_entries() { + let cache = cache(1024); + cache + .insert("old".to_string(), Arc::new("a".to_string())) + .await; + cache + .insert("keep".to_string(), Arc::new("b".to_string())) + .await; + cache.run_pending_tasks().await; + cache.invalidate_if(|key, _| key == "old").unwrap(); + assert!(cache.get("old").await.is_none()); + assert_eq!( + cache.get("keep").await.as_deref().map(String::as_str), + Some("b") + ); + } + + fn cache(capacity: u64) -> Cache> { + Cache::new( + "test", + capacity, + |_, value: &Arc| value.len() as u32, + |_, _, _| {}, + ) + } + + fn operation_metrics(name: &str) -> crate::metrics::OperationMetrics { + use prometheus::{Histogram, HistogramOpts, IntCounter, IntGauge}; + + crate::metrics::OperationMetrics { + started: IntCounter::new(format!("{name}_started"), "test").unwrap(), + succeeded: IntCounter::new(format!("{name}_succeeded"), "test").unwrap(), + failed: IntCounter::new(format!("{name}_failed"), "test").unwrap(), + cancelled: IntCounter::new(format!("{name}_cancelled"), "test").unwrap(), + active: IntGauge::new(format!("{name}_active"), "test").unwrap(), + elapsed: Histogram::with_opts(HistogramOpts::new(format!("{name}_seconds"), "test")) + .unwrap(), + } + } + + fn instrumented_cache() -> (Cache>, Arc) { + let metrics = Arc::new(CacheMetrics { + requests: operation_metrics("requests"), + loads: operation_metrics("loads"), + }); + (cache(1024).with_metrics(metrics.clone()), metrics) + } + + #[tokio::test] + async fn test_metrics_distinguish_requests_from_executed_loads() { + let (cache, metrics) = instrumented_cache(); + // Merely constructing an async call does not start an operation. + drop(cache.try_get_with("unused", async { Ok::<_, ()>(Arc::new(String::new())) })); + assert_eq!(0, metrics.requests.started.get()); + let gate = Notify::new(); + let mut first = Box::pin(cache.try_get_with("key", async { + gate.notified().await; + Ok::<_, ()>(Arc::new("value".to_string())) + })); + let mut second = Box::pin(cache.try_get_with("key", async { Err::, _>(()) })); + assert!(futures::poll!(&mut first).is_pending()); + assert!(futures::poll!(&mut second).is_pending()); + assert_eq!(2, metrics.requests.active.get()); + assert_eq!(1, metrics.loads.active.get()); + gate.notify_waiters(); + let (first, second) = tokio::join!(first, second); + assert!(Arc::ptr_eq(&first.unwrap(), &second.unwrap())); + cache + .try_get_with("key", async { Err::, _>(()) }) + .await + .unwrap(); + assert_eq!(3, metrics.requests.started.get()); + assert_eq!(3, metrics.requests.succeeded.get()); + assert_eq!(1, metrics.loads.started.get()); + assert_eq!(1, metrics.loads.succeeded.get()); + assert_eq!(0, metrics.requests.active.get()); + assert_eq!(0, metrics.loads.active.get()); + assert_eq!(3, metrics.requests.elapsed.get_sample_count()); + assert_eq!(1, metrics.loads.elapsed.get_sample_count()); + assert_eq!(0, metrics.requests.failed.get()); + assert_eq!(0, metrics.requests.cancelled.get()); + } + + #[tokio::test] + async fn test_metrics_release_active_work_on_cancellation_and_error() { + let (cache, metrics) = instrumented_cache(); + let mut producer = + Box::pin(cache.try_get_with("key", std::future::pending::, ()>>())); + let mut waiter = + Box::pin(cache.try_get_with("key", std::future::pending::, ()>>())); + assert!(futures::poll!(&mut producer).is_pending()); + assert!(futures::poll!(&mut waiter).is_pending()); + drop(waiter); + assert_eq!(1, metrics.requests.cancelled.get()); + assert_eq!(0, metrics.loads.cancelled.get()); + drop(producer); + assert_eq!(2, metrics.requests.cancelled.get()); + assert_eq!(1, metrics.loads.cancelled.get()); + let result = cache + .try_get_with("key", async { Err::, _>(()) }) + .await; + assert!(result.is_err()); + assert_eq!(1, metrics.requests.failed.get()); + assert_eq!(1, metrics.loads.failed.get()); + assert_eq!(0, metrics.requests.active.get()); + assert_eq!(0, metrics.loads.active.get()); + assert_eq!(3, metrics.requests.elapsed.get_sample_count()); + assert_eq!(2, metrics.loads.elapsed.get_sample_count()); + } + + #[tokio::test] + async fn test_instrumented_loading_future_is_send() { + fn assert_send(value: T) -> T { + value + } + + let (cache, _) = instrumented_cache(); + assert_send( + cache.try_get_with("key", async { Ok::<_, ()>(Arc::new("value".to_string())) }), + ) + .await + .unwrap(); + } + + #[tokio::test] + async fn test_shares_load_and_bypasses_loader_on_hit() { + let cache = cache(1024); + let calls = AtomicUsize::new(0); + let gate = Notify::new(); + let load = || async { + calls.fetch_add(1, Ordering::Relaxed); + gate.notified().await; + Ok::<_, ()>(Arc::new("value".to_string())) + }; + let mut first = Box::pin(cache.try_get_with("key", load())); + let mut second = Box::pin(cache.try_get_with("key", load())); + assert!(futures::poll!(&mut first).is_pending()); + assert!(futures::poll!(&mut second).is_pending()); + assert_eq!(1, calls.load(Ordering::Relaxed)); + gate.notify_waiters(); + let (first, second) = tokio::join!(first, second); + let first = first.unwrap(); + assert!(Arc::ptr_eq(&first, &second.unwrap())); + let hit = cache + .try_get_with("key", async { Err::, _>(()) }) + .await + .unwrap(); + assert!(Arc::ptr_eq(&first, &hit)); + } + + #[tokio::test] + async fn test_remaining_waiter_loads_after_producer_cancellation() { + let (cache, metrics) = instrumented_cache(); + let mut first = + Box::pin(cache.try_get_with("key", std::future::pending::, ()>>())); + assert!(futures::poll!(&mut first).is_pending()); + let mut second = Box::pin( + cache.try_get_with("key", async { Ok::<_, ()>(Arc::new("retry".to_string())) }), + ); + assert!(futures::poll!(&mut second).is_pending()); + drop(first); + assert_eq!("retry", second.await.unwrap().as_str()); + assert_eq!(2, metrics.requests.started.get()); + assert_eq!(1, metrics.requests.cancelled.get()); + assert_eq!(1, metrics.requests.succeeded.get()); + assert_eq!(2, metrics.loads.started.get()); + assert_eq!(1, metrics.loads.cancelled.get()); + assert_eq!(1, metrics.loads.succeeded.get()); + assert_eq!(0, metrics.requests.active.get()); + assert_eq!(0, metrics.loads.active.get()); + } + + #[tokio::test] + async fn test_errors_are_shared_but_not_cached() { + let (cache, metrics) = instrumented_cache(); + let calls = AtomicUsize::new(0); + let gate = Notify::new(); + let load = || async { + calls.fetch_add(1, Ordering::Relaxed); + gate.notified().await; + Err::, _>("read failed") + }; + let mut first = Box::pin(cache.try_get_with("key", load())); + let mut second = Box::pin(cache.try_get_with("key", load())); + assert!(futures::poll!(&mut first).is_pending()); + assert!(futures::poll!(&mut second).is_pending()); + gate.notify_waiters(); + let (first, second) = tokio::join!(first, second); + assert!(Arc::ptr_eq(&first.unwrap_err(), &second.unwrap_err())); + assert_eq!(2, metrics.requests.failed.get()); + assert_eq!(1, metrics.loads.failed.get()); + assert_eq!(0, metrics.requests.active.get()); + assert_eq!(0, metrics.loads.active.get()); + assert_eq!(1, calls.load(Ordering::Relaxed)); + assert!(cache.get("key").await.is_none()); + assert_eq!( + "retry", + cache + .try_get_with("key", async { + Ok::<_, &'static str>(Arc::new("retry".to_string())) + }) + .await + .unwrap() + .as_str() + ); + } + + #[tokio::test] + async fn test_sampled_size_sums_instances_and_preserves_manual_bytes() { + use crate::metrics::CacheSizeMetrics; + use prometheus::Opts; + use prometheus::core::Collector; + + let metrics = + CacheSizeMetrics::new(Opts::new("test_cache_bytes", "test"), &["type"]).unwrap(); + let manual = metrics.with_label_values(&["shared"]); + manual.set(3); + let make_cache = || { + Cache::new("sampled", 100, |_: &u64, _: &Arc| 7, |_, _, _| {}) + .with_size_metric(metrics.bind(&["shared"])) + }; + let a = make_cache(); + let b = make_cache(); + let sample = || { + metrics.collect()[0].get_metric()[0] + .get_gauge() + .as_ref() + .map_or(0.0, |gauge| gauge.value()) + }; + a.insert(1, Arc::new("one".into())).await; + b.insert(2, Arc::new("two".into())).await; + a.run_pending_tasks().await; + b.run_pending_tasks().await; + assert_eq!(17.0, sample()); + assert_eq!(3, manual.get()); + let consumer = a.get(&1).await.unwrap(); + a.invalidate(&1).await; + a.run_pending_tasks().await; + assert_eq!(10.0, sample()); + drop(consumer); + let cloned = b.clone(); + drop(b); + assert_eq!(10.0, sample()); + drop(cloned); + assert_eq!(3.0, sample()); + } +} diff --git a/src/common/cache/src/lib.rs b/src/common/cache/src/lib.rs new file mode 100644 index 00000000000..58d54adee68 --- /dev/null +++ b/src/common/cache/src/lib.rs @@ -0,0 +1,21 @@ +// 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. + +//! Shared memory cache primitives with native asynchronous load deduplication. + +mod cache; + +pub mod metrics; + +pub use cache::{Cache, RemovalCause}; diff --git a/src/common/cache/src/metrics.rs b/src/common/cache/src/metrics.rs new file mode 100644 index 00000000000..fd02c32b689 --- /dev/null +++ b/src/common/cache/src/metrics.rs @@ -0,0 +1,194 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::{Arc, Mutex, Weak}; +use std::time::Instant; + +use prometheus::core::{Collector, Desc}; +use prometheus::proto::{Gauge, MetricFamily}; +use prometheus::{Histogram, IntCounter, IntGauge, IntGaugeVec, Opts}; + +/// Caller-bound metrics for native cache loading. No collectors are registered here. +/// Requests count only `try_get_with` calls, including its cache hits and callers sharing +/// another request's load. Plain `get` calls are not counted. Loads count +/// only loader futures actually polled by the backend; their difference is not a waiter count. +#[derive(Clone)] +pub struct CacheMetrics { + pub requests: OperationMetrics, + pub loads: OperationMetrics, +} + +/// Metrics for one operation, with labels bound once when constructing a cache. +/// Completion counters are mutually exclusive. Dropping an unfinished future is cancellation. +#[derive(Clone)] +pub struct OperationMetrics { + pub started: IntCounter, + pub succeeded: IntCounter, + pub failed: IntCounter, + pub cancelled: IntCounter, + pub active: IntGauge, + /// Wall-clock lifetime, including suspended time, for every outcome. + pub elapsed: Histogram, +} + +impl OperationMetrics { + /// Starts tracking an operation and its elapsed time. + pub(crate) fn start(&self) -> OperationGuard<'_> { + self.started.inc(); + self.active.inc(); + OperationGuard { + metrics: self, + started: Instant::now(), + completed: false, + } + } +} + +/// Tracks an operation, recording cancellation if dropped before completion. +pub(crate) struct OperationGuard<'a> { + metrics: &'a OperationMetrics, + started: Instant, + completed: bool, +} + +impl OperationGuard<'_> { + /// Records success or failure and completes the operation. + pub(crate) fn finish(mut self, result: &Result) { + if result.is_ok() { + self.metrics.succeeded.inc(); + } else { + self.metrics.failed.inc(); + } + self.completed = true; + } +} + +impl Drop for OperationGuard<'_> { + fn drop(&mut self) { + if !self.completed { + self.metrics.cancelled.inc(); + } + self.metrics.active.dec(); + self.metrics + .elapsed + .observe(self.started.elapsed().as_secs_f64()); + } +} + +/// Cache byte metrics combining manually maintained gauges and sampled cache weights. +/// Sampling does not force backend maintenance or include pending loader values. +#[derive(Clone)] +pub struct CacheSizeMetrics { + gauges: IntGaugeVec, + sources: Arc>>, +} + +/// A byte metric with labels bound for one or more cache instances. +#[derive(Clone)] +pub struct CacheSizeMetric { + metrics: CacheSizeMetrics, + gauge: IntGauge, +} + +/// Keeps a cache instance registered for size sampling while any owner remains. +#[derive(Clone)] +pub struct CacheSizeRegistration { + _reader: Arc u64 + Send + Sync>, +} + +struct CacheSizeSource { + gauge: IntGauge, + read: Weak u64 + Send + Sync>, +} + +impl CacheSizeMetrics { + /// Creates an unregistered collector. Register a clone with the caller's registry. + pub fn new(opts: Opts, labels: &[&str]) -> prometheus::Result { + Ok(Self { + gauges: IntGaugeVec::new(opts, labels)?, + sources: Arc::new(Mutex::new(Vec::new())), + }) + } + + /// Returns the backing gauges for callers retaining manual accounting. + pub fn gauges(&self) -> &IntGaugeVec { + &self.gauges + } + + /// Returns a manually maintained gauge for existing non-cache accounting. + pub fn with_label_values(&self, labels: &[&str]) -> IntGauge { + self.gauges.with_label_values(labels) + } + + /// Binds labels for sampled cache weights. Instances sharing labels are summed. + pub fn bind(&self, labels: &[&str]) -> CacheSizeMetric { + CacheSizeMetric { + metrics: self.clone(), + gauge: self.gauges.with_label_values(labels), + } + } +} + +impl CacheSizeMetric { + /// Registers a cache-level weight reader without retaining it in the collector. + /// Keep the returned registration alongside the cache, including in its clones. + pub fn register( + &self, + read: impl Fn() -> u64 + Send + Sync + 'static, + ) -> CacheSizeRegistration { + let read: Arc u64 + Send + Sync> = Arc::new(read); + self.metrics + .sources + .lock() + .unwrap_or_else(|e| e.into_inner()) + .push(CacheSizeSource { + gauge: self.gauge.clone(), + read: Arc::downgrade(&read), + }); + CacheSizeRegistration { _reader: read } + } +} + +impl Collector for CacheSizeMetrics { + fn desc(&self) -> Vec<&Desc> { + self.gauges.desc() + } + + fn collect(&self) -> Vec { + let mut families = self.gauges.collect(); + let mut sources = self.sources.lock().unwrap_or_else(|e| e.into_inner()); + sources.retain(|source| { + let Some(read) = source.read.upgrade() else { + return false; + }; + let labels = source.gauge.collect(); + for family in &mut families { + for metric in family.mut_metric() { + if metric.get_label() == labels[0].get_metric()[0].get_label() { + let bytes = metric + .get_gauge() + .as_ref() + .map_or(0.0, |gauge| gauge.value()) + + read() as f64; + let mut gauge = Gauge::default(); + gauge.set_value(bytes); + metric.set_gauge(gauge); + } + } + } + true + }); + families + } +} From 05536cd7f52a3b5fbe96f3d3f672b6e9c3855b3e Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Wed, 28 Oct 2026 09:39:45 +0000 Subject: [PATCH 124/239] fix(meta): register failure detectors for regions switched before a rollback (#9989) * fix(meta): register failure detectors for regions switched before a rollback A region migration switches the routes of its tables one table at a time. If it aborts after some tables were switched, those regions are already served by the target, but the rollback deregistered the target's failure detectors for every region and registered only the source's. The target's next heartbeat registers them again, so if the target stopped before sending one, nothing detected its failure and the switched regions stayed unavailable until the metasrv leader changed or the target came back. The rollback now collects, from the routes it already reads, the regions whose leader is the target and resets the target's failure detectors for them at the end. It also skips a table whose route no longer exists, so a table dropped meanwhile no longer fails the rollback before the detectors are handled. Signed-off-by: jeremyhi * test(meta): roll back a downgrading leader in the rollback detector test The table that was not switched yet keeps a downgrading leader in a real abort, so the test now builds it that way and checks that the rollback clears it. Signed-off-by: jeremyhi * fix(meta): monitor regions switched before a rollback only on the target The rollback registered the source's failure detectors for every region of the migration, including the regions already switched to the target. Those regions are served by the target after the rollback, so their source detectors watched the wrong peer: when one fired, the failover found the leader had changed and only removed it. The rollback now registers the source's detectors only for the regions that were not switched, and keeps resetting the target's detectors for the others. Signed-off-by: jeremyhi * refactor(meta): rename reset_failure_detectors_for_candidate_regions to reset_failure_detectors Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- .../src/procedure/region_migration.rs | 110 ++++++++++++++++-- .../rollback_downgraded_region.rs | 3 +- .../upgrade_candidate_region.rs | 3 +- 3 files changed, 103 insertions(+), 13 deletions(-) diff --git a/src/meta-srv/src/procedure/region_migration.rs b/src/meta-srv/src/procedure/region_migration.rs index 75ab5c79397..4f7967dd3c4 100644 --- a/src/meta-srv/src/procedure/region_migration.rs +++ b/src/meta-srv/src/procedure/region_migration.rs @@ -630,13 +630,13 @@ impl Context { Ok(datanode_table_value) } - /// Notifies the RegionSupervisor to register failure detectors of failed region. + /// Notifies the RegionSupervisor to register failure detectors of the given regions on the + /// source peer. /// /// The original failure detector was removed once the procedure was triggered. /// Now, we need to register the failure detector for the failed region again. - pub async fn register_failure_detectors(&self) { + pub async fn register_failure_detectors(&self, region_ids: &[RegionId]) { let datanode_id = self.persistent_ctx.from_peer.id; - let region_ids = &self.persistent_ctx.region_ids; let detecting_regions = region_ids .iter() .map(|region_id| (datanode_id, *region_id)) @@ -650,10 +650,10 @@ impl Context { ); } - /// Notifies the RegionSupervisor to reset failure detectors of candidate regions. - pub async fn reset_failure_detectors_for_candidate_regions(&self) { + /// Notifies the RegionSupervisor to reset failure detectors of the given regions on the target + /// peer. + pub async fn reset_failure_detectors(&self, region_ids: &[RegionId]) { let datanode_id = self.persistent_ctx.to_peer.id; - let region_ids = &self.persistent_ctx.region_ids; let detecting_regions = region_ids .iter() .map(|region_id| (datanode_id, *region_id)) @@ -662,7 +662,7 @@ impl Context { .reset_failure_detectors(detecting_regions) .await; info!( - "Reset failure detectors after migration success for datanode {}, regions {:?}", + "Reset failure detectors of regions for datanode {}, regions {:?}", datanode_id, region_ids ); } @@ -878,12 +878,29 @@ impl RegionMigrationProcedure { .with_label_values(&["rollback"]) .start_timer(); let ctx = &self.context; + let to_peer_id = ctx.persistent_ctx.to_peer.id; let table_regions = ctx.persistent_ctx.table_regions(); + // Regions whose route was switched to the target before the procedure aborted. + let mut switched_regions = Vec::new(); for (table_id, regions) in table_regions { let table_lock = TableLock::Write(table_id).into(); let _guard = procedure_ctx.provider.acquire_lock(&table_lock).await; - let table_route = ctx.get_table_route_value(table_id).await?; + let table_route = match ctx.get_table_route_value(table_id).await { + Ok(table_route) => table_route, + // The table has been dropped meanwhile. + Err(error::Error::TableRouteNotFound { .. }) => continue, + Err(err) => return Err(err), + }; let region_routes = table_route.region_routes().unwrap(); + switched_regions.extend( + region_routes + .iter() + .filter(|route| regions.contains(&route.region.id)) + .filter(|route| { + route.leader_peer.as_ref().map(|peer| peer.id) == Some(to_peer_id) + }) + .map(|route| route.region.id), + ); let downgraded = region_routes .iter() .filter(|route| regions.contains(&route.region.id)) @@ -912,7 +929,24 @@ impl RegionMigrationProcedure { self.context .deregister_failure_detectors_for_candidate_regions() .await; - self.context.register_failure_detectors().await; + // The switched regions are served by the target, so they are monitored there instead. + let unswitched_regions = ctx + .persistent_ctx + .region_ids + .iter() + .filter(|region_id| !switched_regions.contains(region_id)) + .copied() + .collect::>(); + self.context + .register_failure_detectors(&unswitched_regions) + .await; + // The target's next heartbeat would register these detectors again, but nothing would if + // the target stopped before sending one. + if !switched_regions.is_empty() { + self.context + .reset_failure_detectors(&switched_regions) + .await; + } Ok(()) } @@ -1014,8 +1048,8 @@ mod tests { use common_meta::distributed_time_constants::default_distributed_time_constants; use common_meta::instruction::Instruction; - use common_meta::key::test_utils::new_test_table_info; - use common_meta::rpc::router::{Region, RegionRoute}; + use common_meta::key::test_utils::{new_test_table_info, new_test_table_info_with_name}; + use common_meta::rpc::router::{LeaderState, Region, RegionRoute}; use super::*; use crate::handler::HeartbeatMailbox; @@ -1025,12 +1059,66 @@ mod tests { new_downgrade_region_reply, new_flush_region_reply_for_region, new_open_region_reply, new_upgrade_region_reply, }; + use crate::region::supervisor::{Event, RegionFailureDetectorControl}; use crate::service::mailbox::Channel; fn new_persistent_context() -> PersistentContext { test_util::new_persistent_context(1, 2, RegionId::new(1024, 1)) } + #[tokio::test] + async fn test_rollback_registers_detectors_for_switched_regions() { + let env = TestingEnv::new(); + // The route switch aborted after the table of `switched` was switched to the target and + // before the table of `unswitched`, whose leader is still downgrading; the table of + // `dropped` has been dropped. + let switched = RegionId::new(1024, 1); + let unswitched = RegionId::new(1025, 1); + let dropped = RegionId::new(1026, 1); + for (region_id, leader) in [(switched, 2), (unswitched, 1)] { + let table_id = region_id.table_id(); + let table_info = new_test_table_info_with_name(table_id, &format!("table_{table_id}")); + let region_routes = vec![RegionRoute { + region: Region::new_test(region_id), + leader_peer: Some(Peer::empty(leader)), + leader_state: (region_id == unswitched).then_some(LeaderState::Downgrading), + ..Default::default() + }]; + env.create_physical_table_metadata(table_info, region_routes) + .await; + } + let mut persistent_ctx = new_persistent_context(); + persistent_ctx.region_ids = vec![switched, unswitched, dropped]; + let (tx, mut rx) = tokio::sync::mpsc::channel(8); + let mut context_factory = env.context_factory(); + context_factory.region_failure_detector_controller = + Arc::new(RegionFailureDetectorControl::new(tx)); + let mut procedure = RegionMigrationProcedure::new(persistent_ctx, context_factory, vec![]); + + procedure + .rollback_inner(&new_procedure_context()) + .await + .unwrap(); + + // The target's detectors are replaced by the source's, except for `switched`, which stays + // monitored on the target. + let target = [switched, unswitched, dropped].map(|region_id| (2, region_id)); + let source = [unswitched, dropped].map(|region_id| (1, region_id)); + assert_matches!(rx.try_recv(), Ok(Event::DeregisterFailureDetectors(r)) if r == target); + assert_matches!(rx.try_recv(), Ok(Event::RegisterFailureDetectors(r)) if r == source); + assert_matches!(rx.try_recv(), Ok(Event::ResetFailureDetectors(r)) if r == [(2, switched)]); + assert!(rx.try_recv().is_err()); + let table_route = env + .table_metadata_manager() + .table_route_manager() + .table_route_storage() + .get(unswitched.table_id()) + .await + .unwrap() + .unwrap(); + assert_eq!(table_route.region_routes().unwrap()[0].leader_state, None); + } + #[test] fn test_lock_key() { let persistent_context = new_persistent_context(); diff --git a/src/meta-srv/src/procedure/region_migration/update_metadata/rollback_downgraded_region.rs b/src/meta-srv/src/procedure/region_migration/update_metadata/rollback_downgraded_region.rs index 2f0ed0fbe9a..1c9c12a9bc1 100644 --- a/src/meta-srv/src/procedure/region_migration/update_metadata/rollback_downgraded_region.rs +++ b/src/meta-srv/src/procedure/region_migration/update_metadata/rollback_downgraded_region.rs @@ -64,7 +64,8 @@ impl UpdateMetadata { "Rolling back downgraded leader region table route success, table_id: {table_id}, regions: {regions:?}" ); } - ctx.register_failure_detectors().await; + ctx.register_failure_detectors(&ctx.persistent_ctx.region_ids) + .await; Ok(()) } diff --git a/src/meta-srv/src/procedure/region_migration/update_metadata/upgrade_candidate_region.rs b/src/meta-srv/src/procedure/region_migration/update_metadata/upgrade_candidate_region.rs index 5c55891e590..68931e21761 100644 --- a/src/meta-srv/src/procedure/region_migration/update_metadata/upgrade_candidate_region.rs +++ b/src/meta-srv/src/procedure/region_migration/update_metadata/upgrade_candidate_region.rs @@ -215,7 +215,8 @@ impl UpdateMetadata { } ctx.deregister_failure_detectors().await; - ctx.reset_failure_detectors_for_candidate_regions().await; + ctx.reset_failure_detectors(&ctx.persistent_ctx.region_ids) + .await; // Consumes the guard. ctx.volatile_ctx.opening_region_guards.clear(); From 89a5df31418add4a168c7318d4cae55b90ddce00 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Wed, 28 Oct 2026 11:18:25 +0000 Subject: [PATCH 125/239] fix: ignore stale leader grants after a region is downgraded (#9987) * fix: ignore stale leader grants after a region is downgraded Signed-off-by: jeremyhi * fix: compare stale leader grants with the longest lease seen Signed-off-by: jeremyhi * fix: treat a lease too long to store as unknown A lease longer than a u64 of nanoseconds saturated to a shorter one, so the derived send time could be later than the real one. Now such a lease makes the send time unknown, which counts as earlier than any gate. Signed-off-by: jeremyhi * fix: address review comments on stale leader grants Scope the early send time claim to after a lease has been seen, drop the unused role parameter of demote, clarify a test comment and check that a fresh grant still makes a reopened region writable. Signed-off-by: jeremyhi * refactor: compare stale leader grants with the heartbeat send time Pass the send time of the heartbeat behind a lease to set_region_role_with_lease, and return whether the role was applied. Mito compares the send time with the grant gate directly, so it no longer derives it from the longest lease seen. The datanode finalizes leadership only when the grant was applied. Signed-off-by: jeremyhi * test: cover stale grants through the datanode and repeated downgrades Signed-off-by: jeremyhi * refactor: keep leader grant rules in one region method Apply a role grant in MitoRegion::apply_role_grant, which renews the lease and decides whether the role changes; the engine only looks up the region. Rename the gate to downgrade_started_at and demote to begin_downgrade, keep the field private, and copy it on reopen through inherit_downgrade_state_from. Signed-off-by: jeremyhi * refactor: drop a needless borrow Signed-off-by: jeremyhi * refactor: keep lease concepts in the engine and role ordering in the region ManifestContext, which owns the role, records when the role was last demoted and promotes a region only for a promotion issued after that. The time starts when the context is created, so a reopened region ignores grants for heartbeats sent before the reopen, without copying state. The engine translates leases: it renews the deadline and promotes with the heartbeat send time. MitoRegion keeps no grant state. Signed-off-by: jeremyhi * fix: keep the data region of a metric region off a grant its metadata region ignores The metric engine applies a grant to the metadata region first and stops if it is ignored, so the data region, whose role the engine reports, never becomes a leader that skips finalization. Also describe what promoting returns. Signed-off-by: jeremyhi * docs: describe when a lease grant is ignored and flatten the finalize check Signed-off-by: jeremyhi * refactor: move lease types out of the region and the Leader deadline into ManifestContext Signed-off-by: jeremyhi * refactor: name the promotion fence after the downgrade it records Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- src/datanode/src/alive_keeper.rs | 129 ++++++++++++++++++---- src/datanode/src/region_server.rs | 9 +- src/datanode/src/tests.rs | 8 +- src/file-engine/src/engine.rs | 5 +- src/metric-engine/src/engine.rs | 54 ++++++++-- src/metric-engine/src/engine/put.rs | 2 +- src/mito2/src/engine.rs | 12 ++- src/mito2/src/engine/lease_test.rs | 112 +++++++++++++++++++- src/mito2/src/lease.rs | 91 ++++++++++++++++ src/mito2/src/lib.rs | 1 + src/mito2/src/region.rs | 141 +++++++++++-------------- src/mito2/src/region/opener.rs | 2 - src/mito2/src/worker.rs | 12 ++- src/mito2/src/worker/handle_catchup.rs | 2 +- src/mito2/src/worker/handle_create.rs | 2 +- src/mito2/src/worker/handle_write.rs | 6 +- src/query/src/optimizer/test_util.rs | 3 +- src/store-api/src/region_engine.rs | 7 +- 18 files changed, 457 insertions(+), 141 deletions(-) create mode 100644 src/mito2/src/lease.rs diff --git a/src/datanode/src/alive_keeper.rs b/src/datanode/src/alive_keeper.rs index bef59d715c5..13af76202d5 100644 --- a/src/datanode/src/alive_keeper.rs +++ b/src/datanode/src/alive_keeper.rs @@ -134,13 +134,18 @@ impl RegionAliveKeeper { } } - /// Renews the lease of regions to `deadline`. - async fn renew_region_leases(&self, regions: &[GrantedRegion], deadline: Instant) { + /// Renews the lease of regions to `deadline`, granted for a heartbeat sent at `sent_at`. + async fn renew_region_leases( + &self, + regions: &[GrantedRegion], + sent_at: Instant, + deadline: Instant, + ) { for region in regions { let (role, region_id) = (region.role().into(), RegionId::from(region.region_id)); if let Some(handle) = self.find_handle(region_id).await { handle - .reset_deadline(role, deadline, region.extensions.clone()) + .reset_deadline(role, sent_at, deadline, region.extensions.clone()) .await; } else { warn!( @@ -281,7 +286,7 @@ impl HeartbeatResponseHandler for RegionAliveKeeper { if let Some(mito) = self.region_server.mito_engine() { mito.renew_node_lease(deadline.into_std()); } - self.renew_region_leases(®ion_lease.regions, deadline) + self.renew_region_leases(®ion_lease.regions, start_instant, deadline) .await; self.close_staled_regions(®ion_lease.closeable_region_ids) .await; @@ -296,8 +301,8 @@ enum CountdownCommand { /// 4 * `heartbeat_interval_millis` Start(u64), /// Reset countdown deadline to the given instance. - /// (NextRole, Deadline, ExtensionInfo) - Reset((RegionRole, Instant, HashMap>)), + /// (NextRole, HeartbeatSentAt, Deadline, ExtensionInfo) + Reset((RegionRole, Instant, Instant, HashMap>)), /// Returns the current deadline of the countdown task. #[cfg(test)] Deadline(oneshot::Sender), @@ -377,12 +382,18 @@ impl CountdownTaskHandle { async fn reset_deadline( &self, role: RegionRole, + sent_at: Instant, deadline: Instant, extension_info: HashMap>, ) { if let Err(e) = self .tx - .send(CountdownCommand::Reset((role, deadline, extension_info))) + .send(CountdownCommand::Reset(( + role, + sent_at, + deadline, + extension_info, + ))) .await { warn!( @@ -436,20 +447,25 @@ impl CountdownTask { started = true; } }, - Some(CountdownCommand::Reset((role, deadline, extension_info))) => { + Some(CountdownCommand::Reset((role, sent_at, deadline, extension_info))) => { let prev_role_leader = self.region_server.is_region_leader(self.region_id).unwrap_or(false); - if let Err(err) = self.region_server.set_region_role_with_lease(self.region_id, role, deadline.into_std()) { - if err.status_code() == StatusCode::RegionNotFound { - // Table metadata in metasrv is deleted after its regions are dropped. - // The datanode may still receive lease renewal responses that depend on the metadata - // during the short period before it is removed. - warn!(err; "Failed to set region role to {role} for region {region_id}"); - } else { - error!(err; "Failed to set region role to {role} for region {region_id}"); + let applied = match self.region_server.set_region_role_with_lease(self.region_id, role, sent_at.into_std(), deadline.into_std()) { + Ok(applied) => applied, + Err(err) => { + if err.status_code() == StatusCode::RegionNotFound { + // Table metadata in metasrv is deleted after its regions are dropped. + // The datanode may still receive lease renewal responses that depend on the metadata + // during the short period before it is removed. + warn!(err; "Failed to set region role to {role} for region {region_id}"); + } else { + error!(err; "Failed to set region role to {role} for region {region_id}"); + } + false } + }; - // Finalize leadership: persist backfilled metadata. - } else if !prev_role_leader && role == RegionRole::Leader + // Finalize leadership if the grant was applied: persist backfilled metadata. + if applied && !prev_role_leader && role == RegionRole::Leader && let Err(err) = self .region_server .set_region_role_state_gracefully(self.region_id, SettableRegionRoleState::Leader) @@ -503,26 +519,42 @@ impl CountdownTask { #[cfg(test)] mod test { + use std::sync::atomic::AtomicUsize; + use api::v1::Rows; use api::v1::meta::{HeartbeatResponse, RegionLease}; use common_meta::heartbeat::mailbox::HeartbeatMailbox; use mito2::config::MitoConfig; + use mito2::engine::MITO_ENGINE_NAME; use mito2::test_util::{CreateRequestBuilder, TestEnv, build_rows, rows_schema}; - use store_api::region_engine::RegionEngine; + use store_api::region_engine::{ + RegionEngine, SetRegionRoleStateResponse, SetRegionRoleStateSuccess, + }; use store_api::region_request::{ EnterStagingRequest, RegionPutRequest, StagingPartitionDirective, }; use super::*; - use crate::tests::mock_region_server; + use crate::tests::{MockRegionEngine, mock_region_server}; /// Handles a heartbeat response that grants `regions` for a heartbeat sent now. async fn handle_lease_response(alive_keeper: &RegionAliveKeeper, regions: Vec) { + let sent_at = alive_keeper.epoch().elapsed(); + handle_lease_response_sent_at(alive_keeper, sent_at, regions).await; + } + + /// Handles a heartbeat response that grants `regions` for a heartbeat sent at `sent_at` + /// after the epoch. + async fn handle_lease_response_sent_at( + alive_keeper: &RegionAliveKeeper, + sent_at: Duration, + regions: Vec, + ) { let (tx, _rx) = mpsc::channel(1); let response = HeartbeatResponse { region_lease: Some(RegionLease { regions, - duration_since_epoch: alive_keeper.epoch().elapsed().as_millis() as u64, + duration_since_epoch: sent_at.as_millis() as u64, // The countdown of a granted region does not expire during the test. lease_seconds: 86400, closeable_region_ids: vec![], @@ -534,6 +566,38 @@ mod test { alive_keeper.handle(&mut ctx).await.unwrap(); } + #[tokio::test(flavor = "multi_thread")] + async fn ignored_leader_grant_does_not_finalize_leadership() { + common_telemetry::init_default_ut_logging(); + let region_server = mock_region_server(); + let finalized = Arc::new(AtomicUsize::new(0)); + let (engine, _) = MockRegionEngine::with_custom_apply_fn(MITO_ENGINE_NAME, |engine| { + // The engine ignores the grant, so the region stays a follower. + engine.mock_role = Some(Some(RegionRole::Follower)); + let finalized = finalized.clone(); + engine.handle_set_readonly_gracefully_mock_fn = Some(Box::new(move |_| { + finalized.fetch_add(1, Ordering::Relaxed); + Ok(SetRegionRoleStateResponse::success( + SetRegionRoleStateSuccess::mito(0), + )) + })); + }); + let region_id = RegionId::new(1024, 1); + region_server.register_test_region(region_id, engine); + let alive_keeper = + RegionAliveKeeper::new(region_server.clone(), None, Duration::from_millis(100)); + alive_keeper.register_region(region_id).await; + let grant = GrantedRegion { + region_id: region_id.as_u64(), + role: api::v1::meta::RegionRole::Leader.into(), + extensions: HashMap::new(), + }; + handle_lease_response(&alive_keeper, vec![grant]).await; + // The countdown task has handled the grant once it answers. + alive_keeper.deadline(region_id).await.unwrap(); + assert_eq!(0, finalized.load(Ordering::Relaxed)); + } + #[tokio::test(flavor = "multi_thread")] async fn lease_responses_let_regions_write() { common_telemetry::init_default_ut_logging(); @@ -583,10 +647,26 @@ mod test { role: api::v1::meta::RegionRole::Leader.into(), extensions: HashMap::new(), }; - handle_lease_response(&alive_keeper, vec![grant]).await; + handle_lease_response(&alive_keeper, vec![grant.clone()]).await; // The countdown task has applied the grant once it answers. alive_keeper.deadline(region_id).await.unwrap(); put(region_id).await.unwrap(); + + // A grant for a heartbeat sent before a downgrade does not make the region a leader again. + let sent_before = alive_keeper.epoch().elapsed(); + region_server + .set_region_role_state_gracefully(region_id, SettableRegionRoleState::Follower) + .await + .unwrap(); + handle_lease_response_sent_at(&alive_keeper, sent_before, vec![grant.clone()]).await; + alive_keeper.deadline(region_id).await.unwrap(); + assert_eq!(region_server.is_region_leader(region_id), Some(false)); + + // Send times have millisecond precision, so the next heartbeat is sent after the downgrade. + tokio::time::sleep(Duration::from_millis(2)).await; + handle_lease_response(&alive_keeper, vec![grant]).await; + alive_keeper.deadline(region_id).await.unwrap(); + put(region_id).await.unwrap(); } #[tokio::test(flavor = "multi_thread")] @@ -641,6 +721,7 @@ mod test { role: api::v1::meta::RegionRole::Leader.into(), extensions: HashMap::new(), }], + Instant::now(), Instant::now() + Duration::from_millis(200), ) .await; @@ -694,6 +775,7 @@ mod test { countdown_handle .reset_deadline( RegionRole::Leader, + Instant::now(), Instant::now() + Duration::from_millis(heartbeat_interval_millis * 5), HashMap::new(), ) @@ -744,6 +826,7 @@ mod test { role: api::v1::meta::RegionRole::StagingLeader.into(), extensions: HashMap::new(), }], + Instant::now(), Instant::now() + Duration::from_millis(3000), ) .await; @@ -792,6 +875,7 @@ mod test { role: api::v1::meta::RegionRole::Leader.into(), extensions: HashMap::new(), }], + Instant::now(), Instant::now() + Duration::from_millis(3000), ) .await; @@ -834,6 +918,7 @@ mod test { role: api::v1::meta::RegionRole::StagingLeader.into(), extensions: HashMap::new(), }], + Instant::now(), Instant::now() + Duration::from_millis(3000), ) .await; diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index 6055b07dbf6..9d4c4b375fb 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -461,20 +461,23 @@ impl RegionServer { .with_context(|_| HandleRegionRequestSnafu { region_id }) } - /// Sets the role of a region from a lease the metasrv grants until `deadline`. + /// Sets the role of a region from a lease the metasrv grants until `deadline`, in reply to a + /// heartbeat sent at `sent_at`. Returns false if the engine ignores a Leader lease because the + /// region was downgraded or (re)opened after `sent_at`. pub fn set_region_role_with_lease( &self, region_id: RegionId, role: RegionRole, + sent_at: std::time::Instant, deadline: std::time::Instant, - ) -> Result<()> { + ) -> Result { let engine = self .inner .region_map .get(®ion_id) .with_context(|| RegionNotFoundSnafu { region_id })?; engine - .set_region_role_with_lease(region_id, role, deadline) + .set_region_role_with_lease(region_id, role, sent_at, deadline) .with_context(|_| HandleRegionRequestSnafu { region_id }) } diff --git a/src/datanode/src/tests.rs b/src/datanode/src/tests.rs index 4fa247d07ca..a126d3429fd 100644 --- a/src/datanode/src/tests.rs +++ b/src/datanode/src/tests.rs @@ -292,10 +292,12 @@ impl RegionEngine for MockRegionEngine { fn set_region_role_with_lease( &self, _region_id: RegionId, - _role: RegionRole, + role: RegionRole, + _sent_at: std::time::Instant, _deadline: std::time::Instant, - ) -> Result<(), BoxedError> { - Ok(()) + ) -> Result { + // A mocked role stays, so a grant of another role is not applied. + Ok(self.mock_role.flatten().is_none_or(|mocked| mocked == role)) } async fn set_region_role_state_gracefully( diff --git a/src/file-engine/src/engine.rs b/src/file-engine/src/engine.rs index be1d141bcbd..98ad7409f43 100644 --- a/src/file-engine/src/engine.rs +++ b/src/file-engine/src/engine.rs @@ -137,9 +137,10 @@ impl RegionEngine for FileRegionEngine { &self, region_id: RegionId, role: RegionRole, + _sent_at: Instant, _deadline: Instant, - ) -> Result<(), BoxedError> { - self.set_region_role(region_id, role) + ) -> Result { + self.set_region_role(region_id, role).map(|_| true) } async fn set_region_role_state_gracefully( diff --git a/src/metric-engine/src/engine.rs b/src/metric-engine/src/engine.rs index 0c1b1114448..319ddc539d2 100644 --- a/src/metric-engine/src/engine.rs +++ b/src/metric-engine/src/engine.rs @@ -388,23 +388,27 @@ impl RegionEngine for MetricEngine { &self, region_id: RegionId, role: RegionRole, + sent_at: Instant, deadline: Instant, - ) -> Result<(), BoxedError> { - // ignore the region not found error + ) -> Result { + // The metadata region goes first. If it ignores the grant, the data region, whose role the + // engine reports, must not take it either. A region that is not found is ignored. for x in [ utils::to_metadata_region_id(region_id), utils::to_data_region_id(region_id), ] { - if let Err(e) = self + match self .inner .mito - .set_region_role_with_lease(x, role, deadline) - && e.status_code() != StatusCode::RegionNotFound + .set_region_role_with_lease(x, role, sent_at, deadline) { - return Err(e); + Ok(true) => {} + Ok(false) => return Ok(false), + Err(e) if e.status_code() != StatusCode::RegionNotFound => return Err(e), + Err(_) => {} } } - Ok(()) + Ok(true) } async fn sync_region( @@ -887,6 +891,42 @@ mod test { assert!(!object_store.exists(&data_region_dir).await.unwrap()); } + #[tokio::test] + async fn test_leader_grant_ignored_by_metadata_region_leaves_data_region() { + let env = TestEnv::new().await; + env.init_metric_region().await; + let (engine, mito) = (env.metric(), env.mito()); + let physical_region_id = env.default_physical_region_id(); + let lease = std::time::Duration::from_secs(10); + // The data region is downgraded before the heartbeat is sent, the metadata region after it. + let downgrade = |region_id| { + mito.set_region_role_state_gracefully(region_id, SettableRegionRoleState::Follower) + }; + downgrade(utils::to_data_region_id(physical_region_id)) + .await + .unwrap(); + let sent_at = Instant::now(); + downgrade(utils::to_metadata_region_id(physical_region_id)) + .await + .unwrap(); + + let grant = |sent_at: Instant| { + engine + .set_region_role_with_lease( + physical_region_id, + RegionRole::Leader, + sent_at, + sent_at + lease, + ) + .unwrap() + }; + assert!(!grant(sent_at)); + assert_eq!(engine.role(physical_region_id), Some(RegionRole::Follower)); + + assert!(grant(Instant::now())); + assert_eq!(engine.role(physical_region_id), Some(RegionRole::Leader)); + } + #[tokio::test] async fn test_role() { let env = TestEnv::new().await; diff --git a/src/metric-engine/src/engine/put.rs b/src/metric-engine/src/engine/put.rs index e8098131612..eac38d0ab39 100644 --- a/src/metric-engine/src/engine/put.rs +++ b/src/metric-engine/src/engine/put.rs @@ -2379,7 +2379,7 @@ mod tests { let at = |n| start + n * Duration::from_secs(10); let grant = |n| { engine - .set_region_role_with_lease(physical_region_id, RegionRole::Leader, at(n)) + .set_region_role_with_lease(physical_region_id, RegionRole::Leader, start, at(n)) .unwrap() }; let expire = |n| mito.freeze_lease_clock(at(n)); diff --git a/src/mito2/src/engine.rs b/src/mito2/src/engine.rs index 3791489d73f..c628b968b2b 100644 --- a/src/mito2/src/engine.rs +++ b/src/mito2/src/engine.rs @@ -1605,15 +1605,21 @@ impl RegionEngine for MitoEngine { &self, region_id: RegionId, role: RegionRole, + sent_at: Instant, deadline: Instant, - ) -> Result<(), BoxedError> { + ) -> Result { let region = self.inner.find_region(region_id).map_err(BoxedError::new)?; // Before the role, and even if the role leaves the region's state unchanged. if role.writable() { - region.leader_deadline.extend(deadline); + region.manifest_ctx.leader_deadline.extend(deadline); + } + if role == RegionRole::Leader { + return Ok(region + .manifest_ctx + .promote_unless_downgrade_started_after(sent_at, region_id)); } region.set_role(role); - Ok(()) + Ok(true) } async fn set_region_role_state_gracefully( diff --git a/src/mito2/src/engine/lease_test.rs b/src/mito2/src/engine/lease_test.rs index a0e68f7e94b..5fa0a897d63 100644 --- a/src/mito2/src/engine/lease_test.rs +++ b/src/mito2/src/engine/lease_test.rs @@ -13,8 +13,10 @@ // limitations under the License. //! Tests that a region acknowledges only writes whose WAL append completed within its leader -//! lease. +//! lease, and that Leader grants for heartbeats sent before a downgrade do not make it writable +//! again. +use std::collections::HashMap; use std::sync::Arc; use std::time::{Duration, Instant}; @@ -24,8 +26,11 @@ use common_error::ext::{BoxedError, ErrorExt}; use common_recordbatch::RecordBatches; use rstest::rstest; use rstest_reuse::apply; -use store_api::region_engine::{RegionEngine, RegionRole}; -use store_api::region_request::{RegionCatchupRequest, RegionPutRequest, RegionRequest}; +use store_api::region_engine::{RegionEngine, RegionRole, SettableRegionRoleState}; +use store_api::region_request::{ + EnterStagingRequest, RegionCatchupRequest, RegionPutRequest, RegionRequest, + StagingPartitionDirective, +}; use store_api::storage::{RegionId, ScanRequest}; use tokio::task::JoinHandle; @@ -34,10 +39,11 @@ use crate::engine::MitoEngine; use crate::engine::basic_test::build_bulk_insert_request; use crate::engine::listener::WriteWalListener; use crate::error::Error; +use crate::region::{RegionLeaderState, RegionRoleState}; use crate::test_util::{ CreateRequestBuilder, LogStoreFactory, TestEnv, build_rows, kafka_log_store_factory, multiple_log_store_factories, prepare_test_for_kafka_log_store, raft_engine_log_store_factory, - rows_schema, + reopen_region, rows_schema, }; const LEASE: Duration = Duration::from_secs(10); @@ -66,10 +72,18 @@ async fn create_region( fn grant(engine: &MitoEngine, role: RegionRole, deadline: Instant) { engine - .set_region_role_with_lease(REGION_ID, role, deadline) + .set_region_role_with_lease(REGION_ID, role, Instant::now(), deadline) .unwrap(); } +/// Grants Leader for a heartbeat sent at `sent_at`, and returns whether the role was applied. +fn grant_leader_at(engine: &MitoEngine, sent_at: Instant) -> bool { + let deadline = Instant::now() + LEASE; + engine + .set_region_role_with_lease(REGION_ID, RegionRole::Leader, sent_at, deadline) + .unwrap() +} + /// Returns a request that writes rows `start..start + 3`. fn put_request(schema: &[ColumnSchema], start: usize) -> RegionRequest { let rows = Rows { @@ -282,3 +296,91 @@ async fn test_upgraded_region_writes_before_leader_grant() { .unwrap(); put(&engine, REGION_ID, &schema, 0).await.unwrap().unwrap(); } + +#[rstest] +#[case::downgrading(SettableRegionRoleState::DowngradingLeader)] +#[case::follower(SettableRegionRoleState::Follower)] +#[tokio::test] +async fn test_stale_leader_grant_after_downgrade(#[case] downgraded_to: SettableRegionRoleState) { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::with_prefix("stale-leader-grant").await; + let (_listener, engine, _schema) = create_region(&mut env, None).await; + let sent_before = Instant::now() - Duration::from_millis(1); + engine + .set_region_role_state_gracefully(REGION_ID, downgraded_to) + .await + .unwrap(); + let downgraded = engine.get_region(REGION_ID).unwrap().state(); + // Another downgrade raises the gate again, though the role does not change. + let sent_between = Instant::now(); + engine + .set_region_role_state_gracefully(REGION_ID, downgraded_to) + .await + .unwrap(); + + // Grants for heartbeats sent before either downgrade arrive late. + assert!(!grant_leader_at(&engine, sent_before)); + assert!(!grant_leader_at(&engine, sent_between)); + assert_eq!(downgraded, engine.get_region(REGION_ID).unwrap().state()); + + assert!(grant_leader_at(&engine, Instant::now())); + assert!(engine.get_region(REGION_ID).unwrap().is_writable()); +} + +#[rstest] +#[case::catchup(false)] +#[case::close_and_open(true)] +#[tokio::test] +async fn test_reopened_region_ignores_earlier_grants(#[case] close: bool) { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::with_prefix("reopened-region-grant").await; + let (_listener, engine, schema) = create_region(&mut env, None).await; + grant(&engine, RegionRole::Leader, Instant::now() + LEASE); + // The memtable is not empty, so the catchup below reopens the region. + put(&engine, REGION_ID, &schema, 0).await.unwrap().unwrap(); + let sent_before = Instant::now() - Duration::from_millis(1); + let region = engine.get_region(REGION_ID).unwrap(); + if close { + let table_dir = CreateRequestBuilder::new().build().table_dir; + reopen_region(&engine, REGION_ID, table_dir, false, HashMap::new()).await; + } else { + engine + .set_region_role_state_gracefully(REGION_ID, SettableRegionRoleState::Follower) + .await + .unwrap(); + let request = RegionRequest::Catchup(RegionCatchupRequest::default()); + engine.handle_request(REGION_ID, request).await.unwrap(); + } + assert!(!Arc::ptr_eq( + ®ion, + &engine.get_region(REGION_ID).unwrap() + )); + + assert!(!grant_leader_at(&engine, sent_before)); + assert!(!engine.get_region(REGION_ID).unwrap().is_writable()); + + grant(&engine, RegionRole::Leader, Instant::now() + LEASE); + assert!(engine.get_region(REGION_ID).unwrap().is_writable()); +} + +#[tokio::test] +async fn test_leader_grant_exits_staging() { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::with_prefix("staging-leader-grant").await; + let (_listener, engine, _schema) = create_region(&mut env, None).await; + let request = EnterStagingRequest { + partition_directive: StagingPartitionDirective::RejectAllWrites, + }; + engine + .handle_request(REGION_ID, RegionRequest::EnterStaging(request)) + .await + .unwrap(); + let state = || engine.get_region(REGION_ID).unwrap().state(); + assert_eq!(RegionRoleState::Leader(RegionLeaderState::Staging), state()); + + grant(&engine, RegionRole::Leader, Instant::now() + LEASE); + assert_eq!( + RegionRoleState::Leader(RegionLeaderState::Writable), + state() + ); +} diff --git a/src/mito2/src/lease.rs b/src/mito2/src/lease.rs new file mode 100644 index 00000000000..95e076825eb --- /dev/null +++ b/src/mito2/src/lease.rs @@ -0,0 +1,91 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::LazyLock; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::time::Instant; + +/// The instant lease deadlines count from, so that a deadline fits in an [AtomicU64]. +static LEASE_EPOCH: LazyLock = LazyLock::new(Instant::now); + +fn lease_nanos(instant: Instant) -> u64 { + instant.saturating_duration_since(*LEASE_EPOCH).as_nanos() as u64 +} + +/// A lease deadline in nanoseconds since [LEASE_EPOCH]. It only moves forward; 0 means none. +#[derive(Debug, Default)] +pub(crate) struct Deadline(AtomicU64); + +impl Deadline { + pub(crate) fn extend(&self, deadline: Instant) { + self.0.fetch_max(lease_nanos(deadline), Ordering::Relaxed); + } + + pub(crate) fn get(&self) -> u64 { + self.0.load(Ordering::Relaxed) + } +} + +/// Lease of this node: the latest deadline of the heartbeat responses that carry region leases. +#[derive(Debug, Default)] +pub(crate) struct NodeLease { + /// Whether writes need a lease. Only a datanode the metasrv controls sets it. + enabled: AtomicBool, + deadline: Deadline, + /// The time of the lease checks while a test freezes it, in the unit of [Deadline]; 0 when + /// it is not frozen. + #[cfg(any(test, feature = "test"))] + frozen_now: AtomicU64, +} + +impl NodeLease { + pub(crate) fn enable(&self) { + self.enabled.store(true, Ordering::Relaxed); + } + + pub(crate) fn is_enabled(&self) -> bool { + self.enabled.load(Ordering::Relaxed) + } + + pub(crate) fn renew(&self, deadline: Instant) { + self.deadline.extend(deadline); + } + + /// Returns the current time in the unit of [Deadline]. + pub(crate) fn now(&self) -> u64 { + #[cfg(any(test, feature = "test"))] + { + let frozen_now = self.frozen_now.load(Ordering::Relaxed); + if frozen_now != 0 { + return frozen_now; + } + } + lease_nanos(Instant::now()) + } + + #[cfg(any(test, feature = "test"))] + pub(crate) fn freeze_clock(&self, at: Instant) { + self.frozen_now + .store(lease_nanos(at).max(1), Ordering::Relaxed); + } + + /// Lets a region created or upgraded on this node write until the node's deadline. The + /// metasrv granted that deadline before the creation or upgrade, so a later migration of the + /// region waits past it. Later responses do not extend it; the region's Leader grants do. + pub(crate) fn seed(&self, leader_deadline: &Deadline) { + leader_deadline + .0 + .fetch_max(self.deadline.get(), Ordering::Relaxed); + } +} diff --git a/src/mito2/src/lib.rs b/src/mito2/src/lib.rs index a01bca47dd7..116e8ed5833 100644 --- a/src/mito2/src/lib.rs +++ b/src/mito2/src/lib.rs @@ -31,6 +31,7 @@ pub mod error; pub mod extension; pub mod flush; pub mod gc; +mod lease; pub mod manifest; pub mod memtable; mod metrics; diff --git a/src/mito2/src/region.rs b/src/mito2/src/region.rs index ef5b24f9de5..923f223c8c7 100644 --- a/src/mito2/src/region.rs +++ b/src/mito2/src/region.rs @@ -23,8 +23,8 @@ pub(crate) mod version; use std::cell::OnceCell; use std::collections::hash_map::Entry; use std::collections::{HashMap, HashSet}; -use std::sync::atomic::{AtomicBool, AtomicI64, AtomicU64, Ordering}; -use std::sync::{Arc, LazyLock, Mutex, RwLock}; +use std::sync::atomic::{AtomicI64, AtomicU64, Ordering}; +use std::sync::{Arc, Mutex, RwLock}; use std::time::Instant; use common_base::hash::partition_expr_version; @@ -59,6 +59,7 @@ use crate::error::{ FlushableRegionStateSnafu, InvalidPartitionExprSnafu, RegionNotFoundSnafu, RegionStateSnafu, RegionTruncatedSnafu, Result, UnexpectedSnafu, UpdateManifestSnafu, }; +use crate::lease::Deadline; use crate::manifest::action::{ RegionChange, RegionManifest, RegionMetaAction, RegionMetaActionList, }; @@ -186,8 +187,6 @@ pub struct MitoRegion { /// There are no WAL entries in range [flushed_entry_id, topic_latest_entry_id] for current region, /// which means these WAL entries maybe able to be pruned up to `topic_latest_entry_id`. pub(crate) topic_latest_entry_id: AtomicU64, - /// Deadline of the region's latest Leader or StagingLeader lease from the metasrv. - pub(crate) leader_deadline: Deadline, /// Region stats. pub(crate) region_stats: RegionStats, /// manifest stats @@ -224,80 +223,6 @@ impl RegionStats { pub type MitoRegionRef = Arc; -/// The instant lease deadlines count from, so that a deadline fits in an [AtomicU64]. -static LEASE_EPOCH: LazyLock = LazyLock::new(Instant::now); - -fn lease_nanos(instant: Instant) -> u64 { - instant.saturating_duration_since(*LEASE_EPOCH).as_nanos() as u64 -} - -/// A lease deadline in nanoseconds since [LEASE_EPOCH]. It only moves forward; 0 means none. -#[derive(Debug, Default)] -pub(crate) struct Deadline(AtomicU64); - -impl Deadline { - pub(crate) fn extend(&self, deadline: Instant) { - self.0.fetch_max(lease_nanos(deadline), Ordering::Relaxed); - } - - pub(crate) fn get(&self) -> u64 { - self.0.load(Ordering::Relaxed) - } -} - -/// Lease of this node: the latest deadline of the heartbeat responses that carry region leases. -#[derive(Debug, Default)] -pub(crate) struct NodeLease { - /// Whether writes need a lease. Only a datanode the metasrv controls sets it. - enabled: AtomicBool, - deadline: Deadline, - /// The time of the lease checks while a test freezes it, in the unit of [Deadline]; 0 when - /// it is not frozen. - #[cfg(any(test, feature = "test"))] - frozen_now: AtomicU64, -} - -impl NodeLease { - pub(crate) fn enable(&self) { - self.enabled.store(true, Ordering::Relaxed); - } - - pub(crate) fn is_enabled(&self) -> bool { - self.enabled.load(Ordering::Relaxed) - } - - pub(crate) fn renew(&self, deadline: Instant) { - self.deadline.extend(deadline); - } - - /// Returns the current time in the unit of [Deadline]. - pub(crate) fn now(&self) -> u64 { - #[cfg(any(test, feature = "test"))] - { - let frozen_now = self.frozen_now.load(Ordering::Relaxed); - if frozen_now != 0 { - return frozen_now; - } - } - lease_nanos(Instant::now()) - } - - #[cfg(any(test, feature = "test"))] - pub(crate) fn freeze_clock(&self, at: Instant) { - self.frozen_now - .store(lease_nanos(at).max(1), Ordering::Relaxed); - } - - /// Lets a region created or upgraded on this node write until the node's deadline. The - /// metasrv granted that deadline before the creation or upgrade, so a later migration of the - /// region waits past it. Later responses do not extend it; the region's Leader grants do. - pub(crate) fn seed(&self, leader_deadline: &Deadline) { - leader_deadline - .0 - .fetch_max(self.deadline.get(), Ordering::Relaxed); - } -} - #[derive(Debug, Clone)] pub(crate) struct StagingPartitionInfo { pub(crate) partition_directive: StagingPartitionDirective, @@ -1290,6 +1215,17 @@ pub(crate) struct ManifestContext { /// The state of the region. The region checks the state before updating /// manifest. state: AtomicCell, + /// When the region last began a downgrade. A promotion issued before it does not apply. The + /// lock covers comparing or raising it and the role change. + /// + /// It starts at creation time, so promotions issued before the region was opened or reopened + /// do not apply. A context that replaces this one is created after the worker loop closes or + /// reopens the region, so after its last downgrade began. + downgrade_started_at: Mutex, + /// Deadline of the region's latest Leader or StagingLeader lease from the metasrv, which is how + /// long the Leader role stays valid. The engine extends it from those leases, and writes check + /// it. + pub(crate) leader_deadline: Deadline, /// Partition info of the region in staging mode. /// /// During the staging mode, the region metadata in [`VersionControlRef`] is not updated, @@ -1308,6 +1244,8 @@ impl ManifestContext { ManifestContext { manifest_manager: tokio::sync::RwLock::new(manager), state: AtomicCell::new(state), + downgrade_started_at: Mutex::new(Instant::now()), + leader_deadline: Deadline::default(), staging_partition_info: Mutex::new(None), hook, } @@ -1666,6 +1604,33 @@ impl ManifestContext { self.update_and_fire(manager, action_list, is_staging).await } + /// Raises the downgrade start time to now, then marks the region downgrading. + pub(crate) fn begin_downgrade(&self, region_id: RegionId) { + let mut downgrade_started_at = self.downgrade_started_at.lock().unwrap(); + *downgrade_started_at = Instant::now(); + self.set_role(RegionRole::DowngradingLeader, region_id); + } + + /// Promotes the region to a writable leader unless a downgrade started after `issued_at`. + /// Returns false if the promotion is ignored for that reason; otherwise the role changes as + /// [`Self::set_role`] allows. + pub(crate) fn promote_unless_downgrade_started_after( + &self, + issued_at: Instant, + region_id: RegionId, + ) -> bool { + // A writable leader needs no change; a staging leader is promoted below. + if self.state.load() == RegionRoleState::Leader(RegionLeaderState::Writable) { + return true; + } + let downgrade_started_at = self.downgrade_started_at.lock().unwrap(); + if issued_at < *downgrade_started_at { + return false; + } + self.set_role(RegionRole::Leader, region_id); + true + } + /// Sets the [`RegionRole`]. /// /// ```text @@ -2341,7 +2306,6 @@ mod tests { last_schedule_compaction_millis: Default::default(), time_provider: Arc::new(StdTimeProvider), topic_latest_entry_id: Default::default(), - leader_deadline: Default::default(), region_stats: RegionStats::new(), stats: ManifestStats::default(), } @@ -2457,6 +2421,24 @@ mod tests { } } + #[tokio::test] + async fn test_promoting_writable_leader_takes_no_lock() { + let env = SchedulerEnv::new().await; + let region = Arc::new(build_test_region(&env).await); + let _downgrade_started_at = region.manifest_ctx.downgrade_started_at.lock().unwrap(); + let (tx, rx) = std::sync::mpsc::channel(); + let promoted = region.clone(); + std::thread::spawn(move || { + let manifest_ctx = &promoted.manifest_ctx; + let issued_at = std::time::Instant::now(); + tx.send( + manifest_ctx.promote_unless_downgrade_started_after(issued_at, promoted.region_id), + ) + .unwrap(); + }); + assert!(rx.recv_timeout(std::time::Duration::from_secs(10)).unwrap()); + } + #[tokio::test] async fn test_compaction_update_manifest_allows_editing_state() { let env = SchedulerEnv::new().await; @@ -2945,7 +2927,6 @@ mod tests { last_schedule_compaction_millis: Default::default(), time_provider: Arc::new(StdTimeProvider), topic_latest_entry_id: Default::default(), - leader_deadline: Default::default(), region_stats: RegionStats::new(), stats: ManifestStats::default(), }; diff --git a/src/mito2/src/region/opener.rs b/src/mito2/src/region/opener.rs index 947844c8f20..d10545198f7 100644 --- a/src/mito2/src/region/opener.rs +++ b/src/mito2/src/region/opener.rs @@ -459,7 +459,6 @@ impl RegionOpener { last_schedule_compaction_millis: AtomicI64::new(now), time_provider: self.time_provider.clone(), topic_latest_entry_id: AtomicU64::new(flushed_entry_id), - leader_deadline: Default::default(), region_stats: RegionStats::new(), stats: self.stats, })) @@ -698,7 +697,6 @@ impl RegionOpener { last_schedule_compaction_millis: AtomicI64::new(now), time_provider: self.time_provider.clone(), topic_latest_entry_id: AtomicU64::new(topic_latest_entry_id), - leader_deadline: Default::default(), region_stats: RegionStats::new(), stats: self.stats.clone(), }; diff --git a/src/mito2/src/worker.rs b/src/mito2/src/worker.rs index 50fb5bc0eb1..b93331d99dc 100644 --- a/src/mito2/src/worker.rs +++ b/src/mito2/src/worker.rs @@ -52,7 +52,7 @@ use rand::{Rng, rng}; use snafu::ResultExt; use store_api::logstore::LogStore; use store_api::region_engine::{ - RegionRole, SetRegionRoleStateResponse, SetRegionRoleStateSuccess, SettableRegionRoleState, + SetRegionRoleStateResponse, SetRegionRoleStateSuccess, SettableRegionRoleState, }; use store_api::storage::{FileId, RegionId}; use tokio::sync::mpsc::{Receiver, Sender}; @@ -67,12 +67,13 @@ use crate::config::MitoConfig; use crate::error::{self, CreateDirSnafu, JoinSnafu, Result, WorkerStoppedSnafu}; use crate::flush::{FlushScheduler, WriteBufferManagerImpl, WriteBufferManagerRef}; use crate::gc::{GcLimiter, GcLimiterRef}; +use crate::lease::NodeLease; use crate::memtable::MemtableBuilderProvider; use crate::metrics::{REGION_COUNT, REQUEST_WAIT_TIME, WRITE_STALLING}; use crate::region::opener::PartitionExprFetcherRef; use crate::region::{ - CatchupRegions, CatchupRegionsRef, MitoRegionRef, NodeLease, OpeningRegions, OpeningRegionsRef, - RegionMap, RegionMapRef, + CatchupRegions, CatchupRegionsRef, MitoRegionRef, OpeningRegions, OpeningRegionsRef, RegionMap, + RegionMapRef, }; use crate::request::{ BackgroundNotify, BulkInsertRequest, DdlRequest, OptionOutputTx, SenderBulkRequest, @@ -1516,8 +1517,9 @@ impl RegionWorkerLoop { ) { // Fence writes here: writes the worker handles later see // Downgrading, and earlier ones have finished. This only - // changes a writable or staging leader. - region.set_role(RegionRole::DowngradingLeader); + // changes a writable or staging leader, but always fences + // promotions issued before. + region.manifest_ctx.begin_downgrade(region_id); } #[cfg(feature = "hdfs-object-store")] let sweep_older_than = self.config.gc.unknown_file_lingering_time; diff --git a/src/mito2/src/worker/handle_catchup.rs b/src/mito2/src/worker/handle_catchup.rs index 65f20c5837a..0ca51b98a11 100644 --- a/src/mito2/src/worker/handle_catchup.rs +++ b/src/mito2/src/worker/handle_catchup.rs @@ -87,7 +87,7 @@ impl RegionWorkerLoop { if request.set_writable { // The region takes writes under the node lease until its next Leader // grant. - node_lease.seed(®ion.leader_deadline); + node_lease.seed(®ion.manifest_ctx.leader_deadline); region.set_role(RegionRole::Leader); // Finalize leadership: persist backfilled metadata. match region diff --git a/src/mito2/src/worker/handle_create.rs b/src/mito2/src/worker/handle_create.rs index 15641385e59..d2eab4a3f59 100644 --- a/src/mito2/src/worker/handle_create.rs +++ b/src/mito2/src/worker/handle_create.rs @@ -107,7 +107,7 @@ impl RegionWorkerLoop { } // The region takes writes under the node lease until its first Leader grant. - self.node_lease.seed(®ion.leader_deadline); + self.node_lease.seed(®ion.manifest_ctx.leader_deadline); // Insert the MitoRegion into the RegionMap. self.register_region(region); if let Some(state) = &self.series_index_task_state { diff --git a/src/mito2/src/worker/handle_write.rs b/src/mito2/src/worker/handle_write.rs index 6b8e5be4276..269859cb7e8 100644 --- a/src/mito2/src/worker/handle_write.rs +++ b/src/mito2/src/worker/handle_write.rs @@ -340,7 +340,7 @@ impl RegionWorkerLoop { match region.state() { RegionRoleState::Leader(RegionLeaderState::Writable) | RegionRoleState::Leader(RegionLeaderState::Staging) => { - let leader_deadline = region.leader_deadline.get(); + let leader_deadline = region.manifest_ctx.leader_deadline.get(); if lease_now.is_some_and(|now| now >= leader_deadline) { sender_req .sender @@ -497,7 +497,7 @@ impl RegionWorkerLoop { match region.state() { RegionRoleState::Leader(RegionLeaderState::Writable) | RegionRoleState::Leader(RegionLeaderState::Staging) => { - let leader_deadline = region.leader_deadline.get(); + let leader_deadline = region.manifest_ctx.leader_deadline.get(); if lease_now.is_some_and(|now| now >= leader_deadline) { bulk_req .sender @@ -606,7 +606,7 @@ impl RegionWorkerLoop { for (region_id, region_ctx) in region_ctxs.iter_mut() { // A Leader grant may have extended the deadline during the append. let expired = region_ctx.lease.as_ref().is_some_and(|(region, deadline)| { - now >= *deadline && now >= region.leader_deadline.get() + now >= *deadline && now >= region.manifest_ctx.leader_deadline.get() }); if !expired || region_ctx.is_failed() { continue; diff --git a/src/query/src/optimizer/test_util.rs b/src/query/src/optimizer/test_util.rs index a6b95ffdeba..09f72bb08b7 100644 --- a/src/query/src/optimizer/test_util.rs +++ b/src/query/src/optimizer/test_util.rs @@ -107,8 +107,9 @@ impl RegionEngine for MetaRegionEngine { &self, _region_id: RegionId, _role: RegionRole, + _sent_at: std::time::Instant, _deadline: std::time::Instant, - ) -> Result<(), BoxedError> { + ) -> Result { unimplemented!() } diff --git a/src/store-api/src/region_engine.rs b/src/store-api/src/region_engine.rs index 6dffdd03c3b..2a2bc7b7f1f 100644 --- a/src/store-api/src/region_engine.rs +++ b/src/store-api/src/region_engine.rs @@ -1082,7 +1082,9 @@ pub trait RegionEngine: Send + Sync { /// take effect. fn set_region_role(&self, region_id: RegionId, role: RegionRole) -> Result<(), BoxedError>; - /// Sets [RegionRole] for a region from a lease the metasrv grants until `deadline`. + /// Sets [RegionRole] for a region from a lease the metasrv grants until `deadline`, in reply + /// to a heartbeat sent at `sent_at`. Returns false if the engine ignores a Leader lease because + /// the region was downgraded or (re)opened after `sent_at`. /// /// The region must not acknowledge a write whose WAL append completes at or after the /// deadline of its latest Leader or StagingLeader lease. Leases of other roles do not extend @@ -1091,8 +1093,9 @@ pub trait RegionEngine: Send + Sync { &self, region_id: RegionId, role: RegionRole, + sent_at: Instant, deadline: Instant, - ) -> Result<(), BoxedError>; + ) -> Result; /// Syncs the region manifest to the given manifest version. async fn sync_region( From ebb7e22566f8cb6246f39eab01e32488c682a2c8 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Wed, 28 Oct 2026 12:51:52 +0000 Subject: [PATCH 126/239] chore(ci): gate fuzz, E2E and KV backends by diff scope (#9927) * ci: gate the integration CI long tail by diff scope `integration.yml` starts its long tail on every pull request: the three fuzz families, the chaos legs, the export/import E2E job and the two sqlness KvBackend legs. Most diffs cannot affect them, so a `scope` job now classifies the pull request and tells those jobs whether to start. The classifier is deterministic and fail-open. Forced modes (merge queue, `schedule`, `workflow_dispatch`, `ci:full`, `CI_SCOPE_FORCE_ALL`) and three floors -- always-run paths, an uncovered `src/` change, and a PromQL/flow filter -- can only widen the result; everything else is a glob match against the suite catalog. Every step of the job is `continue-on-error` and every gate reads `needs.scope.outputs. != 'false'`, so a classifier crash runs the jobs instead of skipping them. The rules are checked out from the pull request base, so a pull request cannot relax its own gate. The policy is data: `.github/ci/scope-rules.toml` holds the path groups, the suite catalog and the always-run list, with a `reason` on every exclusion. `ci_scope_rules.py` loads it with the stdlib `tomllib` and validates the shape (unknown or cyclic `@group` references, a `family` that disagrees with the suite id, an empty reason) and that every glob's literal prefix still exists, so a renamed crate cannot leave a rule that matches nothing. `make check-ci-scope` runs that check locally; the `CI Scope Rules` job runs it on `ubuntu-slim`, the same image the classifier uses. `.github/ci/scope-cases.toml` holds golden cases that `test_ci_scope.py` replays through the real classifier. The PostgreSQL and MySQL sqlness legs move into their own `sqlness-kvbackend` job so they can be gated: a job-level `if:` cannot read the matrix, and gating them inside `sqlness` would skip required checks and block the merge. They run the same cases as the required Basic leg, so only the backend itself matters. The verdict is visible where triage looks: one log line per suite, the exact gate values, `::warning::` when a lookup fails and the gate widens, and a job summary that lists each suite with the rule that decided it and the changed files that matched. Measured on 200 merged and open pull requests: 27.5% always run, 41% decided by path rules, 31.5% skip every gated job; 995 of 2600 suite selections are skipped, with no selection that a path rule makes. The classifier itself runs in about ten seconds on a single-vCPU runner. Signed-off-by: WenyXu * fix(ci): let the scope job read pull requests The `scope` job declared `permissions: contents: read`, and an explicit block sets every unlisted permission to `none`. `GET /pulls/{n}` therefore answered 403, `build_plan` widened to a full run as designed, and every gate read an absent output as "run": the gate could never skip anything. - `pull-requests: read` on the `scope` job. - `tests/runner/**` joins `@kv-backend`: the runner implements `--setup-pg`/`--setup-mysql`, which only the KV-backend legs pass, so a change there must not skip the only job that exercises it. The `UNMAPPED_CRATES` excuse is gone with it. - The README no longer claims a queue failure blocks the merge: the gated jobs are not required checks. Signed-off-by: WenyXu * fix(ci): run every suite when the kind registry script changes The distributed and chaos fuzz jobs start their cluster through `.github/actions/setup-kind`, which runs `.github/scripts/kind-with-registry.sh`. `.github/actions/**` is in `ALWAYS_RUN` but the script is not, so a change confined to it selected no suite at all and every gate stayed `false`: the jobs that exercise the changed cluster setup were skipped. The script joins `ALWAYS_RUN`, consistent with the action that calls it, and a golden case pins it. Signed-off-by: WenyXu * fix(ci): widen the gate for the fuzz corpus seeding script `fuzz-test/action.yaml` runs `.github/scripts/prepare-fuzz-corpus.sh` before the cache restore in every fuzz job, but the script was not in `ALWAYS_RUN`. A change confined to it therefore selected no suite: `any_fuzz` stayed `false` and every fuzz job that runs it was skipped. The script joins `ALWAYS_RUN` next to the other scripts the fuzz jobs execute, with a golden case for it. Both misses so far (this one and `kind-with-registry.sh`) came through a composite action, which the workflow never names. `WorkflowContractTests` now resolves `uses: ./.github/actions/` transitively for the gated jobs and fails when one of the scripts they run is not in `ALWAYS_RUN`. Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- .github/ci/README.md | 168 ++++ .github/ci/ci_scope.py | 623 ++++++++++++ .github/ci/ci_scope_rules.py | 277 ++++++ .github/ci/scope-rules.toml | 284 ++++++ .github/ci/tests/fixtures/scope-cases.toml | 158 +++ .github/ci/tests/test_ci_scope.py | 1018 ++++++++++++++++++++ .github/scripts/tests/test_ci_slash.py | 2 +- .github/workflows/checks.yml | 19 + .github/workflows/integration.yml | 161 +++- Makefile | 4 + 10 files changed, 2675 insertions(+), 39 deletions(-) create mode 100644 .github/ci/README.md create mode 100755 .github/ci/ci_scope.py create mode 100755 .github/ci/ci_scope_rules.py create mode 100644 .github/ci/scope-rules.toml create mode 100644 .github/ci/tests/fixtures/scope-cases.toml create mode 100755 .github/ci/tests/test_ci_scope.py diff --git a/.github/ci/README.md b/.github/ci/README.md new file mode 100644 index 00000000000..07de1fcae7b --- /dev/null +++ b/.github/ci/README.md @@ -0,0 +1,168 @@ +# Integration CI admission gate + +`integration.yml` runs its expensive jobs — fuzz, export/import E2E, the sqlness +KvBackend legs — on every non-draft pull request. Most pull requests cannot +affect any of them. This directory decides which of those jobs start, using only +the diff. + +The gate can only *skip* work, never add it, and it fails open: when the verdict +is missing or uncertain, every gated job runs. + +| File | Role | +|---|---| +| `ci_scope.py` | the classifier: reads the diff, writes `$GITHUB_OUTPUT` | +| `scope-rules.toml` | the policy: which paths select which suite, as data | +| `ci_scope_rules.py` | loads and validates the policy; `--check-paths` | +| `tests/` | unit tests, and `fixtures/scope-cases.toml` golden cases | + +Two words appear everywhere below. A **suite** is one matrix leg of a gated job, +named `/`. A **gate** is one `$GITHUB_OUTPUT` key that switches a +whole job. + +## What is gated + +Only `integration.yml`, and only these jobs: + +| Gated job | Runs | +|---|---| +| `prepare-fuzz-targets` | the fuzz targets only | +| `fuzztest` | standalone fuzz, 2 legs | +| `unstable-fuzztest` | unstable standalone create-table fuzz, 1 leg | +| `distributed-fuzztest` | distributed fuzz with remote WAL, 4 legs | +| `distributed-fuzztest-with-chaos` | Chaos Mesh, 3 legs | +| `build-greptime-ci` | the profile-CI binary only | +| `export-import-v2-e2e` | CLI `COPY DATABASE` export/import against S3 | +| `sqlness-kvbackend` | sqlness with PostgreSQL and MySQL as the metadata store | + +The legs themselves are the `matrix` entries in `integration.yml`; the rules +file repeats their ids so a verdict can name them. + +Everything else always runs: `rust.yml`, `checks.yml`, `docs.yml`, the debug +`build`, the required `sqlness` job (Basic, Remote WAL, Flat format), `compat` +and `compat-updater-check`. None of the gated jobs is in the required-check +matrix in `docs.yml`, so a skipped job can never block a merge. + +Two choices follow from that rule. The two KvBackend legs sit in their own +`sqlness-kvbackend` job because a job-level `if:` cannot read `matrix`, so +gating them inside `sqlness` would skip required checks and make the pull +request unmergeable. The multi-language SDK job stays ungated: it takes about +three minutes, less than the cost of getting its verdict wrong. + +## How a verdict is made + +First match wins: + +1. **Forced** — `merge_group`, `schedule`, `workflow_dispatch`, the `ci:full` + label or `CI_SCOPE_FORCE_ALL`. Everything runs. +2. **Too large to judge** — more than 300 changed files. Everything runs. +3. **Always-run file changed** — `Cargo.toml`, `Cargo.lock`, `tests-fuzz/**`, + this directory, or anything else that changes what the suites build or how + they start. Everything runs. +4. **Uncovered change** — a path under `src/` that no suite claims. This is the + net for a crate added, split or renamed after the rules were written. + Everything runs. +5. **Path rules** — the suites whose globs match the changed files, and only + those. + +The code is the source of truth for the details; its module docstring is worth +reading before changing it. + +## Why a check was skipped + +The run page's summary pane answers this without opening a log: one row per +suite, with its verdict, the rule that decided it, and up to three of the +changed files that matched. + +The step log of the `CI Scope` job has the same information in text form. It +starts with `CI scope: mode=… full=…` (`grep '^CI scope'` finds it), then one +line per suite — `run` or `skip`, source, reason — then `gates:` with the exact +values the jobs read. + +A degraded decision is never silent. A failed label or file lookup emits +`::warning::` and widens the gate to every suite. + +## Forcing a full run + +Add the `ci:full` label to the pull request and click "Re-run all jobs". The +label is read through the API at run time, not from the event payload, so one +added after the first run still counts. `/ci` dispatch +(`.github/scripts/ci-slash.py`), the `CI_SCOPE_FORCE_ALL` repository variable +and every merge-queue run also run everything: the classifier returns a forced +plan there, so a wrong skip is still exercised before the change lands. None of +the gated jobs is a required check, so a failure in the queue does not block the +merge by itself. + +`CI_SCOPE_FORCE_ALL` is the switch to flip while debugging the rules themselves. + +## Changing the rules + +`scope-rules.toml` is data, not code: `ci_scope_rules.py` loads it with the +standard library's `tomllib`. Two things are validated: + +- the shape — an unknown or cyclic `@group` reference, a `gate` that disagrees + with the suite id, a gate missing from `[gates]`, or an exemption without a + `reason` all raise; +- that every glob still points at something that exists, so a renamed or + deleted crate fails the check instead of silently matching nothing. + +Run both with `make check-ci-scope`, or the path half alone with +`python3 .github/ci/ci_scope_rules.py --check-paths`. The same command runs in +the `CI Scope Rules` job of `checks.yml`, on `ubuntu-slim` — the image the +classifier itself uses, so a path that exists there also exists when the verdict +is computed. + +A malformed file fails the classify step, and because every step of that job is +`continue-on-error`, every gated job runs instead. + +`[[matching.exempt]]` lists paths outside the gate's remit: they select no suite +and never widen the gate. Each one carries a `reason`, and the loader rejects an +empty one — what CI never runs is the part a reader must be able to justify. + +`tests/fixtures/scope-cases.toml` is test data, not configuration: a diff shape +and the suites it must start. `ScopeCaseTests` replays each case through the +classifier, so a rule edit that changes behaviour fails with the case name +instead of silently rerouting CI. The cases carry no pull request numbers; +`git log` has those. + +## Adding a suite + +1. Add a `[suites."/"]` table to `scope-rules.toml` with `gate`, + `paths` (globs or `@group` references) and a one-line `description`. The id + is `/`, and `gate` is that id with `-` replaced by `_`; + the loader and `test_suite_metadata` both enforce it. +2. Add the matching matrix leg to the job in `integration.yml`. +3. Gate the job's `if:` with the matching output, and list `scope` in its + `needs`. A new gate also needs a key in `derive_outputs` and in the `scope` + job's `outputs`; `test_every_gate_routes_to_an_output` fails when that + routing is missing. +4. Add a case to `tests/fixtures/scope-cases.toml` and run + `python3 .github/ci/tests/test_ci_scope.py`. `CrateCoverageTests` separately + fails when a workspace crate matches no suite and is not excused in + `UNMAPPED_CRATES`: the runtime net would run everything, but the rule table + is what should be updated. + +## Ground rules for the workflow + +These six rules keep the gate failing open. Break one and either CI runs jobs it +should have skipped, or a classifier failure takes CI down with it. + +- **Read the output as `!= 'false'`.** An absent output means "run"; writing + `== 'true'` inverts the failure direction. +- **Start every gate with `!cancelled()`.** Without a status function the + condition carries an implicit `success()` over `needs`, so a `scope` job that + times out would skip the gates. Other dependencies are therefore required + explicitly with `needs..result == 'success'`. +- **Every fallible step of `scope` is `continue-on-error: true`.** GitHub skips + the dependents of a *failed* job, so a classifier crash must leave the job + green with no outputs rather than take CI down with it. +- **Rules come from the base revision.** The `scope` job checks out + `pull_request.base.sha` (or `merge_group.base_sha`), so a pull request cannot + weaken its own gate. On the first pull request that introduces this directory + the base has no `ci_scope.py`: the step fails, no outputs are written, and + every gate runs. +- **One gate per job, never per leg.** GitHub does not expose `matrix` to + `jobs..if`. Per-suite verdicts stay visible in the log. +- **Pass the run's `head`/`base` SHAs and let the classifier compare them with + the pull request's live revisions.** A re-run keeps the original SHAs while + GitHub serves the *current* file list; without the comparison the new diff + would be judged against the old tree. With it, a mismatch runs everything. diff --git a/.github/ci/ci_scope.py b/.github/ci/ci_scope.py new file mode 100755 index 00000000000..571e975325a --- /dev/null +++ b/.github/ci/ci_scope.py @@ -0,0 +1,623 @@ +#!/usr/bin/env python3 +# 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. + +"""Decide which integration CI suites must run for one pull request. + +Pipeline, in priority order: + +1. forced mode -- `merge_group`, `schedule`, `workflow_dispatch`, + `CI_SCOPE_FORCE_ALL`, or the `ci:full` label. Every suite runs. +2. always-run floor -- a changed file matches `ALWAYS_RUN`. Every suite runs. +3. uncovered build change -- a path under `SRC_TREE` that no suite claims, so a + crate added or renamed later still runs everything. +4. static path matches -- glob match of the changed files against the suite + catalog in `scope-rules.toml`. Paths under `EXEMPT_TREES` are outside the + gate's remit and are dropped first: `@server` claims `src/query/**`, which + would otherwise put a PromQL-only change through the whole fleet. + +The verdict is deterministic: the same diff always produces the same plan, and +a change that no rule explains falls through to "run everything". + +The plan maps to `$GITHUB_OUTPUT` keys that the `integration.yml` gate jobs +consume; `--dry-run` prints them as JSON. +""" + +import argparse +import json +import os +import re +import sys +import urllib.error +import urllib.request + +from ci_scope_rules import ( + ALWAYS_RUN, + EXEMPT_TREES, + SRC_TREE, + SUITES, + suite_group, + suite_paths, +) + +MAX_FILES = 300 +HTTP_TIMEOUT = 45 +FORCED_EVENTS = {"merge_group", "schedule", "workflow_dispatch"} +# A diff with no code, config or workflow file never needs a suite. +CODE_FILE = re.compile( + r"\.(rs|toml|lock|sh|ya?ml)$|(^|/)Makefile$|(^|/)Cargo\.toml$|\.sql$" +) +FUZZ_GATES = ("fuzz_standalone", "fuzz_distributed", "fuzz_chaos") + + +class ScopeError(Exception): + """A handled failure: the caller falls back instead of crashing.""" + + +# --------------------------------------------------------------------------- +# HTTP seam (tests replace `_request`) +# --------------------------------------------------------------------------- + + +def _request(method, url, headers, body=None, timeout=HTTP_TIMEOUT): + request = urllib.request.Request(url, data=body, headers=headers, method=method) + try: + with urllib.request.urlopen(request, timeout=timeout) as response: + return response.status, response.read() + except urllib.error.HTTPError as err: + return err.code, err.read() + except (urllib.error.URLError, TimeoutError, OSError) as err: + raise ScopeError(f"{method} {url}: {err}") from err + + +def _get(url, headers, timeout=HTTP_TIMEOUT): + return _request("GET", url, headers, None, timeout) + + +def _post(url, headers, body, timeout=HTTP_TIMEOUT): + return _request("POST", url, headers, body, timeout) + + +def get_json(url, headers, timeout=HTTP_TIMEOUT): + """GET returning (status, parsed json); transport failures raise `ScopeError`.""" + status, raw = _get(url, headers, timeout) + try: + payload = json.loads(raw.decode("utf-8", "replace")) if raw else None + except json.JSONDecodeError: + payload = None + return status, payload + + +# --------------------------------------------------------------------------- +# Path matching +# --------------------------------------------------------------------------- + + +def glob_to_regex(pattern): + """Translate a repo glob into an anchored regexp. + + `**/` also matches zero directories, `**` spans `/`, `*` and `?` do not. + """ + out = ["^"] + index = 0 + while index < len(pattern): + char = pattern[index] + if char == "*": + end = index + while end < len(pattern) and pattern[end] == "*": + end += 1 + if end - index >= 2: + if end < len(pattern) and pattern[end] == "/": + out.append("(?:.*/)?") + index = end + 1 + continue + out.append(".*") + else: + out.append("[^/]*") + index = end + elif char == "?": + out.append("[^/]") + index += 1 + else: + out.append(re.escape(char)) + index += 1 + out.append("$") + return re.compile("".join(out)) + + +def compile_patterns(patterns): + return [glob_to_regex(pattern) for pattern in patterns] + + +def matches_any(filename, compiled): + return any(regex.match(filename) for regex in compiled) + + +# --------------------------------------------------------------------------- +# GitHub API +# --------------------------------------------------------------------------- + + +def github_api(method, path, token, payload=None): + base = os.environ.get("GITHUB_API_URL", "https://api.github.com").rstrip("/") + url = base + path + headers = { + "Accept": "application/vnd.github+json", + "User-Agent": "greptimedb-ci-scope", + } + if token: + headers["Authorization"] = "Bearer " + token + body = None + if payload is not None: + body = json.dumps(payload).encode() + headers["Content-Type"] = "application/json" + if method == "GET": + return get_json(url, headers) + status, raw = _request(method, url, headers, body, HTTP_TIMEOUT) + try: + parsed = json.loads(raw.decode("utf-8", "replace")) if raw else None + except json.JSONDecodeError: + parsed = None + return status, parsed + + +def fetch_changed_files(repo, number, token): + """Return (files, truncated). `truncated` means the PR has more than MAX_FILES.""" + files = [] + page = 1 + while True: + status, payload = github_api( + "GET", f"/repos/{repo}/pulls/{number}/files?per_page=100&page={page}", token + ) + if status != 200 or not isinstance(payload, list): + raise ScopeError(f"cannot list pull request files (HTTP {status})") + for item in payload: + if not isinstance(item, dict): + continue + patch = item.get("patch") + files.append( + { + "filename": item.get("filename", ""), + "previous_filename": item.get("previous_filename") or "", + "status": item.get("status", "modified"), + "additions": item.get("additions", 0), + "deletions": item.get("deletions", 0), + "patch": patch if isinstance(patch, str) else "", + "patch_included": isinstance(patch, str) and bool(patch), + } + ) + if len(payload) < 100: + return files, False + if len(files) >= MAX_FILES: + # Keep `per_page=100`: with `per_page=1`, page N is item N, so the + # probe would answer "is there a file 4?" instead of "file 301?". + status, extra = github_api( + "GET", + f"/repos/{repo}/pulls/{number}/files?per_page=100&page={page + 1}", + token, + ) + if status != 200 or not isinstance(extra, list): + raise ScopeError( + "cannot confirm the changed file list is complete " + f"(HTTP {status})" + ) + del files[MAX_FILES:] + return files, bool(extra) + page += 1 + + +def fetch_pull(repo, number, token): + """Return (labels, head sha, base sha) for one pull request.""" + status, payload = github_api("GET", f"/repos/{repo}/pulls/{number}", token) + if status != 200 or not isinstance(payload, dict): + raise ScopeError(f"cannot read pull request (HTTP {status})") + labels = payload.get("labels") or [] + head = ((payload.get("head") or {}).get("sha") or "").strip() + base = ((payload.get("base") or {}).get("sha") or "").strip() + names = {label.get("name", "") for label in labels if isinstance(label, dict)} + return names, head, base + + +def revision_drift(args, head, base): + """Return a note when the tested revisions are no longer the live ones. + + A re-run keeps the run's original SHAs while the file list served by GitHub + is the pull request's current one. After a force-push the classifier would + judge the new diff for the old tree, so a mismatch runs every suite instead. + """ + for supplied, live, name in ( + (args.head_sha, head, "head"), + (args.base_sha, base, "base"), + ): + if supplied and live and supplied.lower() != live.lower(): + return ( + f"pull request {name} moved from {supplied[:12]} to " + f"{live[:12]} after this run started" + ) + return "" + + +# --------------------------------------------------------------------------- +# Decisions +# --------------------------------------------------------------------------- + + +def file_paths(item): + """Every path the change touched; a rename keeps the deleted side too.""" + names = [item["filename"]] + previous = item.get("previous_filename") + if previous and previous != item["filename"]: + names.append(previous) + return names + + +def gated_paths(item): + """The paths of one changed file that the gate judges. + + A path under `EXEMPT_TREES` is outside the remit: it selects no suite, and + it is not a change that "no rule covers", so it cannot widen the gate. + """ + return [name for name in file_paths(item) if not name.startswith(EXEMPT_TREES)] + + +def path_matches(files): + """Return (always_run_hits, {suite_id: [matching file paths]}).""" + always_compiled = compile_patterns(ALWAYS_RUN) + suite_matchers = { + suite_id: compile_patterns(suite_paths(suite)) + for suite_id, suite in SUITES.items() + } + always_hits = set() + matches = {suite_id: set() for suite_id in SUITES} + for item in files: + for name in file_paths(item): + if matches_any(name, always_compiled): + always_hits.add(name) + for name in gated_paths(item): + for suite_id, compiled in suite_matchers.items(): + if matches_any(name, compiled): + matches[suite_id].add(name) + return sorted(always_hits), { + suite_id: sorted(matched) for suite_id, matched in matches.items() + } + + +def uncovered_changes(files, matches): + """True when a change can reach code that no suite path claims. + + This is the net for a crate that was added, split or renamed after the rules + were written: without it such a change matches no suite and would skip every + gated job. A single unmatched path is enough, even when other paths in the + same diff match a suite. + """ + return first_uncovered(files, matches) is not None + + +def first_uncovered(files, matches): + """The first changed path that no suite claims, or None. + + Naming it is what makes the widening explainable in the log. + """ + matched = {name for names in matches.values() for name in names} + return next( + ( + name + for item in files + for name in gated_paths(item) + if name.startswith(SRC_TREE) and name not in matched + ), + None, + ) + + +def has_meaningful_changes(files): + return any(CODE_FILE.search(item["filename"]) for item in files) + + +def finalize(plan): + """`full` means every gated suite runs, whatever produced the decision.""" + plan["full"] = all(entry["run"] for entry in plan["suites"].values()) + return plan + + +def all_suites(run, reason, source): + return { + suite_id: {"run": run, "source": source, "reason": reason} + for suite_id in SUITES + } + + +def full_plan(plan, reason, source, mode="static"): + """Every suite runs; `finalize` derives `full` from the verdicts.""" + plan.update(mode=mode, suites=all_suites(True, reason, source)) + return finalize(plan) + + +def forced_reason(args, labels): + """Return why every suite must run, or an empty string.""" + event = args.event or os.environ.get("GITHUB_EVENT_NAME") or "pull_request" + if event in FORCED_EVENTS: + return f"event={event}" + if os.environ.get("CI_SCOPE_FORCE_ALL", "").strip().lower() in { + "1", + "true", + "yes", + "on", + }: + return "CI_SCOPE_FORCE_ALL" + if "ci:full" in labels: + return "label ci:full" + return "" + + +def build_plan(args, token): + notes = [] + plan = { + "mode": "static", + "head_sha": args.head_sha or "", + "base_sha": args.base_sha or "", + "full": False, + "suites": {}, + "matching": {}, + "notes": notes, + } + + # Events and variables that force a run need no GitHub call at all; the + # `ci:full` label is the only reason to read the pull request early. + forced = forced_reason(args, set()) + labels, head_live, base_live = set(), "", "" + if args.pr and not forced: + try: + labels, head_live, base_live = fetch_pull( + args.repo, args.pr, token + ) + except ScopeError as err: + # A failed lookup must widen the gate, never narrow it. + notes.append(f"pull request lookup failed: {err}") + return full_plan(plan, "pull request lookup failed", "always-run") + forced = forced_reason(args, labels) + + if forced: + notes.append(f"forced: {forced}") + return full_plan(plan, f"forced by {forced}", "forced", mode="forced") + + if not args.pr: + raise ScopeError("--pr is required unless the event forces every suite") + + drift = revision_drift(args, head_live, base_live) + if drift: + notes.append(drift) + return full_plan(plan, drift, "always-run") + + files, truncated = fetch_changed_files(args.repo, args.pr, token) + if truncated: + notes.append(f"changed file list truncated at {MAX_FILES}") + return full_plan( + plan, + f"pull request has more than {MAX_FILES} changed files", + "always-run", + ) + + always_hits, matches = path_matches(files) + # Keep the evidence: the step summary explains a verdict by naming the files + # that matched, which is what triage asks for first. + plan["matching"] = {suite_id: sorted(names) for suite_id, names in matches.items()} + if always_hits: + notes.append(f"always-run: {', '.join(sorted(always_hits)[:5])}") + return full_plan( + plan, + "always-run path: " + ", ".join(sorted(always_hits)[:3]), + "always-run", + ) + + first = first_uncovered(files, matches) + if first is not None: + notes.append(f"uncovered build change: {first}") + return full_plan(plan, f"uncovered build change: {first}", "uncovered") + + static_run = {suite_id: bool(matched) for suite_id, matched in matches.items()} + + if not has_meaningful_changes(files) and not any(static_run.values()): + plan.update( + mode="static", + full=False, + suites={ + suite_id: { + "run": False, + "source": "static", + "reason": "no code, config or workflow file changed", + } + for suite_id in SUITES + }, + ) + notes.append("no meaningful changes: skip every suite") + return finalize(plan) + + plan.update( + mode="static", + suites={ + suite_id: { + "run": static_run[suite_id], + "source": "static", + "reason": ("path rule: " + ", ".join(matches[suite_id][:3])) + if matches[suite_id] + else "no path rule matched", + } + for suite_id in SUITES + }, + ) + return finalize(plan) + + +def derive_outputs(plan): + """Flatten the plan into the `$GITHUB_OUTPUT` keys the workflow gates on. + + + Every gate in `SUITES` must route to a key here, otherwise its suites are + classified but never switch a job; `test_every_gate_routes_to_an_output` + asks this function directly, so a new gate fails there until it is routed. + + + Job-level `if:` cannot read the `matrix` context, so the workflow gates a + whole fuzz gate instead of individual matrix legs; the per-suite verdicts + stay visible in the step log. + """ + + def suites_of(gate): + return [ + suite_id + for suite_id, entry in plan["suites"].items() + if entry["run"] and SUITES[suite_id]["gate"] == gate + ] + + outputs = {} + legs = {} + for gate in FUZZ_GATES: + running = suites_of(gate) + legs[gate] = [suite_group(suite_id) for suite_id in running] + legs[gate] = [leg for leg in legs[gate] if leg] + outputs[gate] = "true" if running else "false" + unstable = bool(suites_of("fuzz_unstable")) + # Every gate outside the fuzz group switches exactly one job of its own. + for gate in sorted({suite["gate"] for suite in SUITES.values()}): + if gate not in FUZZ_GATES: + outputs[gate] = "true" if suites_of(gate) else "false" + outputs["any_fuzz"] = ( + "true" + if unstable or any(outputs[gate] == "true" for gate in FUZZ_GATES) + else "false" + ) + distributed = outputs["fuzz_distributed"] == "true" or outputs["fuzz_chaos"] == "true" + outputs["build_ci_needed"] = "true" if unstable or distributed else "false" + return outputs, legs + + +def render_log(plan): + """One line per suite plus the emitted gates, for the step log. + + Triage starts in the run log, so the verdict has to be readable there + without any other artifact. + """ + head = f"CI scope: mode={plan['mode']} full={'true' if plan['full'] else 'false'}" + lines = [head] + width = max((len(suite_id) for suite_id in plan["suites"]), default=0) + for suite_id, entry in plan["suites"].items(): + lines.append( + f" {'run ' if entry['run'] else 'skip'} {suite_id:<{width}} " + f"{entry['source']:<12} {entry['reason']}" + ) + outputs, _ = derive_outputs(plan) + lines.append( + " gates: " + " ".join(f"{key}={value}" for key, value in outputs.items()) + ) + for note in plan["notes"]: + lines.append(f" note: {note}") + return "\n".join(lines) + + +SUMMARY_MAX_FILES = 3 + + +def render_summary(plan): + """Markdown for the job summary: the matching evidence behind each verdict. + + The step log answers "what ran"; this answers "why", which is the question + when a pull request unexpectedly skips or starts a suite. + """ + lines = [ + f"### CI scope: mode={plan['mode']}, full={'true' if plan['full'] else 'false'}", + "", + "| suite | verdict | rule | matched files |", + "| --- | --- | --- | --- |", + ] + matching = plan.get("matching") or {} + for suite_id, entry in plan["suites"].items(): + matched = matching.get(suite_id) or [] + shown = ", ".join(f"`{name}`" for name in matched[:SUMMARY_MAX_FILES]) + if len(matched) > SUMMARY_MAX_FILES: + shown += f" +{len(matched) - SUMMARY_MAX_FILES} more" + lines.append( + f"| `{suite_id}` | {'run' if entry['run'] else 'skip'} | " + f"{entry['source']} | {shown or '—'} |" + ) + if plan["notes"]: + lines += ["", "

notes", ""] + lines += [f"- {note}" for note in plan["notes"]] + lines += ["", "
"] + return "\n".join(lines) + + +def log_annotations(plan): + """Workflow-run annotations for every degraded decision. + + A lookup failure widens the gate; it must be visible in the run list without + opening the job log, otherwise it looks like a confident verdict. + """ + annotations = [] + for note in plan["notes"]: + if note.startswith("pull request lookup failed"): + suffix = "every suite runs" + level = "warning" + else: + continue + message = f"CI scope: {note}; {suffix}"[:300] + annotations.append(f"::{level}::{message}") + return annotations + + +def emit(plan, args): + outputs, _ = derive_outputs(plan) + if args.dry_run: + print(json.dumps(outputs, indent=2)) + elif args.github_output: + with open(args.github_output, "a", encoding="utf-8") as handle: + for key, value in outputs.items(): + handle.write(f"{key}={value}\n") + if args.summary: + with open(args.summary, "a", encoding="utf-8") as handle: + handle.write(render_summary(plan) + "\n") + + +def parse_args(argv): + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("--repo", default=os.environ.get("GITHUB_REPOSITORY", "")) + parser.add_argument("--pr", type=int) + parser.add_argument("--head-sha", default="") + parser.add_argument("--base-sha", default="") + parser.add_argument( + "--event", + choices=["pull_request", "merge_group", "schedule", "workflow_dispatch"], + ) + parser.add_argument("--github-output") + parser.add_argument( + "--summary", + help="append the markdown matching table to this file ($GITHUB_STEP_SUMMARY)", + ) + parser.add_argument("--dry-run", action="store_true") + return parser.parse_args(argv) + + +def main(argv=None): + args = parse_args(argv) + token = os.environ.get("GITHUB_TOKEN") or os.environ.get("GH_TOKEN") or "" + plan = build_plan(args, token) + for annotation in log_annotations(plan): + print(annotation) + print(render_log(plan)) + emit(plan, args) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/.github/ci/ci_scope_rules.py b/.github/ci/ci_scope_rules.py new file mode 100755 index 00000000000..c7453424848 --- /dev/null +++ b/.github/ci/ci_scope_rules.py @@ -0,0 +1,277 @@ +#!/usr/bin/env python3 +# 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. + +"""Policy for the integration CI admission gate. + +The rules are data and live in `.github/ci/scope-rules.toml` so they can be +reviewed and changed without touching the classifier. This module parses that +file once at import, validates it, and exposes the names `ci_scope.py` uses. + +A malformed file raises at import: the classifier job's steps are +`continue-on-error`, so a broken rules file leaves the gate outputs empty and +every gated job runs. +""" + +import argparse +import pathlib +import sys +import tomllib + +RULES_PATH = pathlib.Path(__file__).resolve().with_name("scope-rules.toml") +# `.github/ci/scope-rules.toml`: parents[0] is `.github/ci`, parents[2] the repo. +REPO_ROOT = RULES_PATH.parents[2] + + + +class RulesError(ValueError): + """The rules file is missing, malformed or internally inconsistent.""" + + +def _require(table, key, where): + if key not in table: + raise RulesError(f"{where} is missing `{key}`") + return table[key] + + +def _strings(value, where): + if not isinstance(value, list) or not all(isinstance(item, str) for item in value): + raise RulesError(f"{where} must be an array of strings") + return list(value) + + +def _trees(entries, where): + """Normalise `[[matching.*_tree]]` tables into bare path prefixes. + + Each entry carries a `reason`: the exclusions decide what CI never runs, so + the justification is part of the contract, not a comment. + """ + if not isinstance(entries, list): + raise RulesError(f"{where} must be an array of tables") + paths = [] + for index, entry in enumerate(entries): + if not isinstance(entry, dict): + raise RulesError(f"{where}[{index}] is not a table") + path = _require(entry, "path", f"{where}[{index}]") + if not isinstance(path, str) or not path: + raise RulesError(f"{where}[{index}].path must be a non-empty string") + reason = _require(entry, "reason", f"{where}[{index}]") + if not isinstance(reason, str) or not reason.strip(): + raise RulesError(f"{where}[{index}].reason must explain the exclusion") + paths.append(path) + return paths + + +def validate(data): + """Check the shape of a parsed rules document and return it normalised.""" + if not isinstance(data, dict): + raise RulesError("rules document is not a table") + matching = _require(data, "matching", "the rules file") + gates_in = _require(data, "gates", "the rules file") + groups_in = _require(data, "groups", "the rules file") + suites_in = _require(data, "suites", "the rules file") + + groups = {} + for name, paths in groups_in.items(): + groups[name] = _strings(paths, f"groups.{name}") + + def expand(entries, stack=()): + out = [] + for entry in entries: + if not entry.startswith("@"): + out.append(entry) + continue + name = entry[1:] + if name not in groups: + raise RulesError(f"unknown path group `{entry}`") + if name in stack: + raise RulesError(f"path group cycle: {' -> '.join(stack + (name,))}") + out.extend(expand(groups[name], stack + (name,))) + return out + + for name in groups: + expand([f"@{name}"]) # every group must resolve on its own + + suites = {} + for suite_id, suite in suites_in.items(): + if not isinstance(suite, dict): + raise RulesError(f"suites.{suite_id} is not a table") + gate = _require(suite, "gate", f"suites.{suite_id}") + paths = _strings( + _require(suite, "paths", f"suites.{suite_id}"), f"suites.{suite_id}.paths" + ) + description = _require(suite, "description", f"suites.{suite_id}") + if not isinstance(description, str) or not description.strip(): + raise RulesError(f"suites.{suite_id}.description is empty") + # `gate/leg` for a matrix job, or a bare gate for a single-leg job. + prefix = suite_id.split("/", 1)[0].replace("-", "_") + if prefix != gate: + raise RulesError( + f"suites.{suite_id}.gate is `{gate}`, expected `{prefix}`" + ) + expand(paths) + suites[suite_id] = { + "gate": gate, + "paths": paths, + "description": description, + } + + gate_jobs = {} + for gate, job in gates_in.items(): + if not isinstance(job, str) or not job: + raise RulesError(f"gates.{gate} must name a job id") + gate_jobs[gate] = job + used = {suite["gate"] for suite in suites.values()} + unknown = sorted(used - set(gate_jobs)) + if unknown: + raise RulesError(f"suites use gates missing from [gates]: {', '.join(unknown)}") + dead = sorted(set(gate_jobs) - used) + if dead: + raise RulesError(f"[gates] names gates no suite uses: {', '.join(dead)}") + + return { + "always_run": _strings( + _require(matching, "always_run", "matching"), "matching.always_run" + ), + "src_tree": _require(matching, "src_tree", "matching"), + "exempt_trees": _trees( + _require(matching, "exempt", "matching"), "matching.exempt" + ), + "gates": gate_jobs, + "groups": groups, + "suites": suites, + } + + +def load(path=None): + """Parse and validate the rules file.""" + path = RULES_PATH if path is None else path + try: + raw = pathlib.Path(path).read_bytes() + except OSError as err: + raise RulesError(f"cannot read {path}: {err}") from err + try: + data = tomllib.loads(raw.decode("utf-8")) + except (tomllib.TOMLDecodeError, UnicodeDecodeError) as err: + raise RulesError(f"{path} is not valid TOML: {err}") from err + return validate(data) + + +_RULES = load() + +ALWAYS_RUN = _RULES["always_run"] +SRC_TREE = _RULES["src_tree"] +EXEMPT_TREES = tuple(_RULES["exempt_trees"]) +GATE_JOBS = _RULES["gates"] +PATH_GROUPS = _RULES["groups"] +SUITES = _RULES["suites"] + + +def main(argv=None): + """Validate the rules file and print what it holds; non-zero on failure.""" + parser = argparse.ArgumentParser(description="Validate the CI scope rules.") + parser.add_argument( + "--check-paths", + action="store_true", + help="also fail when a pattern's literal prefix does not exist", + ) + args = parser.parse_args(argv) + try: + rules = load() + except RulesError as err: + print(f"scope rules invalid: {err}", file=sys.stderr) + return 1 + if args.check_paths: + missing = missing_paths(rules) + for pattern, prefix in missing: + print(f"missing path: {pattern} (nothing under {prefix})", file=sys.stderr) + return 1 if missing else 0 + print( + f"scope rules ok: {len(rules['suites'])} suites, " + f"{len(rules['groups'])} path groups, " + f"{len(rules['always_run'])} always-run entries" + ) + return 0 + + +def suite_group(suite_id): + """Return the matrix group id encoded in a suite id, or None.""" + _, _, group = suite_id.partition("/") + return group or None + + +def expand_paths(entries, groups=None): + """Expand `@group` references, including groups nested in other groups.""" + groups = PATH_GROUPS if groups is None else groups + patterns = [] + for entry in entries: + if entry.startswith("@"): + patterns.extend(expand_paths(groups[entry[1:]], groups)) + else: + patterns.append(entry) + return patterns + + +def suite_paths(suite): + """Expand `@group` references in one suite definition into globs.""" + return expand_paths(suite["paths"]) + + +def _literal_prefix(pattern): + """Return the longest leading path that the glob cannot vary. + + Everything from the first `*`, `?` or `[` onwards is variable. When that + wildcard sits inside a segment the incomplete segment is dropped, so + `src/a/foo*.rs` yields `src/a/`; a wildcard in the very first segment (for + example `*.md`) yields `""` because the whole path is variable. + """ + cuts = [pattern.find(ch) for ch in "*?[" if ch in pattern] + if not cuts: + return pattern + prefix = pattern[: min(cuts)] + if "/" not in prefix: + return "" + return prefix.rsplit("/", 1)[0] + "/" + + +def _all_patterns(rules): + """Every glob the rules mention, with `@group` references expanded.""" + patterns = [] + for key in ("always_run", "exempt_trees"): + patterns.extend(expand_paths(rules[key], rules["groups"])) + for suite in rules["suites"].values(): + patterns.extend(expand_paths(suite["paths"], rules["groups"])) + return patterns + + +def missing_paths(rules=None, root=None): + """Return `(pattern, prefix)` for patterns whose literal prefix is absent. + + A rule that names a directory nobody renamed would otherwise match nothing + forever, silently narrowing the gate. A pattern with no fixed prefix is + skipped: there is nothing on disk to check. + """ + rules = _RULES if rules is None else rules + root = pathlib.Path(REPO_ROOT if root is None else root) + missing = set() + for pattern in _all_patterns(rules): + prefix = _literal_prefix(pattern) + if prefix and not (root / prefix).exists(): + missing.add((pattern, prefix)) + return sorted(missing) + + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/.github/ci/scope-rules.toml b/.github/ci/scope-rules.toml new file mode 100644 index 00000000000..ba2b92deda2 --- /dev/null +++ b/.github/ci/scope-rules.toml @@ -0,0 +1,284 @@ +# Rules for the integration CI admission gate. +# +# Data only: `ci_scope_rules.py` loads this file with the stdlib `tomllib`, and +# `ci_scope.py` matches diffs against it. Edit this file, not the code, when a +# suite's paths change. +# +# Conventions +# ----------- +# - `**` crosses directories, `*` does not; a trailing `/**` covers everything +# under a directory. +# - `@name` inside `suites.*.paths` expands `[groups].name`, and groups may +# reference other groups. Expansion is recursive. +# - A suite id is `/`; `gate` is the `$GITHUB_OUTPUT` +# key that switches the job, and it must equal the id prefix with `-` +# replaced by `_`. +# - Everything here is a floor, never a filter: an unmatched `src/` path runs +# every suite, and a malformed file fails the classifier job, whose steps are +# `continue-on-error`, so the gated jobs run. + +[matching] +# A changed file matching one of these runs every suite. +# Put here anything that changes what the suites build or how they start. +always_run = [ + "Cargo.toml", + "Cargo.lock", + "rust-toolchain.toml", + ".cargo/**", + "tests-fuzz/**", + "tests-integration/**", + "config/**", + "docker/**", + ".github/ci/**", + ".github/workflows/integration.yml", + ".github/actions/**", + # Both run through a composite action, so neither appears in the workflow: + # `setup-kind` starts the cluster for the distributed and chaos jobs, and + # `fuzz-test` seeds the corpus every fuzz job starts from. + ".github/scripts/kind-with-registry.sh", + ".github/scripts/prepare-fuzz-corpus.sh", + ".github/scripts/run-fuzz-targets.sh", + ".github/scripts/collect-fuzz-target-artifacts.sh", + ".github/scripts/wait-kafka-ready.sh", + ".github/scripts/pull-test-deps-images.sh", +] + +# A change under `src/` that no suite claims runs every suite: that is the net +# for a crate added, split or renamed after these rules were written. +src_tree = "src/" + +# Paths outside the gate's remit. They select no suite, and they do not count as +# "no rule covers this" either, so they never widen the gate. Each needs a +# `reason`: what CI covers them instead. +[[matching.exempt]] +path = "src/promql/" +reason = "PromQL crate: rust.yml and the promql sqlness cases own it; the fuzz suites speak SQL." + +[[matching.exempt]] +path = "src/query/src/promql/" +reason = "PromQL module inside `src/query`, which `@server` claims as `src/query/**`." + +[[matching.exempt]] +path = "src/flow/" +reason = "Flow engine: rust.yml and the flow sqlness cases own it; no fuzz suite loads it." + +# Which job each gate switches. The keys are `$GITHUB_OUTPUT` names; the values +# are job ids in `integration.yml`. The classifier cannot read YAML, so the +# mapping is data here, and a test proves both directions line up. +[gates] +fuzz_standalone = "fuzztest" +fuzz_unstable = "unstable-fuzztest" +fuzz_distributed = "distributed-fuzztest" +fuzz_chaos = "distributed-fuzztest-with-chaos" +export_import_v2 = "export-import-v2-e2e" +sqlness_kvbackend = "sqlness-kvbackend" + +[groups] +server = [ + "src/operator/**", + "src/sql/**", + "src/catalog/**", + "src/datatypes/**", + "src/table/**", + "src/store-api/**", + "src/session/**", + "src/api/**", + "src/servers/**", + "src/frontend/**", + # Also covers `src/query/src/promql/`, which is exempt (see + # `[[matching.exempt]]`): a PromQL-only change must not run the fuzz fleet. + "src/query/**", + "src/partition/**", + "src/index/**", + "src/object-store/**", + "src/cache/**", + "src/puffin/**", + "src/auth/**", + "src/client/**", + "src/log-query/**", + "src/pipeline/**", + "src/plugins/**", + "src/cmd/**", + "src/standalone/**", + "src/datanode/**", + "src/common/base/**", + "src/common/batcher/**", + "src/common/catalog/**", + "src/common/config/**", + "src/common/datasource/**", + "src/common/decimal/**", + "src/common/error/**", + "src/common/event-recorder/**", + "src/common/frontend/**", + "src/common/function/**", + "src/common/grpc/**", + "src/common/grpc-expr/**", + "src/common/greptimedb-telemetry/**", + "src/common/macro/**", + "src/common/mem-prof/**", + "src/common/memory-manager/**", + "src/common/options/**", + "src/common/plugins/**", + "src/common/pprof/**", + "src/common/query/**", + "src/common/recordbatch/**", + "src/common/runtime/**", + "src/common/session/**", + "src/common/sql/**", + "src/common/stat/**", + "src/common/substrait/**", + "src/common/telemetry/**", + "src/common/test-util/**", + "src/common/time/**", + "src/common/version/**", + "src/common/workload/**", +] + +# Storage engines and WAL. +engine = [ + "src/mito2/**", + "src/mito-codec/**", + "src/metric-engine/**", + "src/file-engine/**", + "src/log-store/**", + "src/common/wal/**", +] + +# DDL and procedure code standalone executes in-process: it builds a +# `DdlManager` and runs `LocalProcedureExecutor` without any metasrv. +meta-core = [ + "src/common/meta/**", + "src/common/procedure/**", + "src/common/procedure-test/**", +] + +# Everything that sits on the KV backend the two sqlness KvBackend legs +# exercise: the backend implementations, the metadata service and the startup +# wiring that picks a backend. +kv-backend = [ + # The runner implements `--setup-pg`/`--setup-mysql`, which only these legs pass. + "tests/runner/**", + "src/common/meta/**", + "src/meta-srv/**", + "src/meta-client/**", + "src/common/procedure/**", + "src/cmd/**", + "src/common/config/**", + "src/common/options/**", +] + +# Adds the metasrv/meta-client crates on top of the in-process DDL code. +meta = [ + "src/meta-srv/**", + "src/meta-client/**", + "@meta-core", +] + +[suites."fuzz-standalone/database-and-regular-table"] +gate = "fuzz_standalone" +paths = [ + "@server", + "@engine", + "@meta-core", +] +description = "Standalone-server black-box fuzz of CREATE DATABASE/TABLE, ALTER, INSERT over MySQL." + +[suites."fuzz-standalone/logical-table"] +gate = "fuzz_standalone" +paths = [ + "@server", + "@engine", + "@meta-core", +] +description = "Standalone logical-table DDL/INSERT fuzz (mito2 logical tables)." + +[suites."fuzz-unstable/standalone"] +gate = "fuzz_unstable" +paths = [ + "@server", + "@engine", + "@meta-core", +] +description = "Unstable standalone create-table fuzz (nix build feature)." + +[suites."fuzz-distributed/database-and-regular-table"] +gate = "fuzz_distributed" +paths = [ + "@server", + "@engine", + "@meta", +] +description = "Distributed (frontend/datanode/metasrv, remote WAL) DDL/DML fuzz." + +[suites."fuzz-distributed/logical-table"] +gate = "fuzz_distributed" +paths = [ + "@server", + "@engine", + "@meta", +] +description = "Distributed logical-table fuzz." + +[suites."fuzz-distributed/repartition"] +gate = "fuzz_distributed" +paths = [ + "@server", + "@engine", + "@meta", +] +description = "Distributed repartition ALTER TABLE + GC fuzz." + +[suites."fuzz-distributed/logical-pruning"] +gate = "fuzz_distributed" +paths = [ + "@engine", + "@meta", +] +description = "Remote-WAL logical table pruning of stale manifests." + +[suites."fuzz-chaos/migration"] +gate = "fuzz_chaos" +paths = [ + "@server", + "@engine", + "@meta", +] +description = "Region migration under Chaos Mesh, mito and metric engines." + +[suites."fuzz-chaos/failover"] +gate = "fuzz_chaos" +paths = [ + "@server", + "@engine", + "@meta", +] +description = "Region failover under Chaos Mesh, mito and metric engines." + +[suites."fuzz-chaos/repartition"] +gate = "fuzz_chaos" +paths = [ + "@server", + "@engine", + "@meta", +] +description = "Repartition + GC under Chaos Mesh." + +[suites."export-import-v2"] +gate = "export_import_v2" +paths = [ + "src/cli/**", + "@server", + "@engine", + "@meta-core", +] +description = "CLI COPY DATABASE export/import v2 end-to-end against S3." + +[suites."sqlness-kvbackend/metasrv-postgres"] +gate = "sqlness_kvbackend" +paths = ["@kv-backend"] +description = "Sqlness with PostgreSQL as the metadata KV backend." + +[suites."sqlness-kvbackend/metasrv-mysql"] +gate = "sqlness_kvbackend" +paths = ["@kv-backend"] +description = "Sqlness with MySQL as the metadata KV backend." diff --git a/.github/ci/tests/fixtures/scope-cases.toml b/.github/ci/tests/fixtures/scope-cases.toml new file mode 100644 index 00000000000..d159bd3e32d --- /dev/null +++ b/.github/ci/tests/fixtures/scope-cases.toml @@ -0,0 +1,158 @@ +# Golden cases for the admission gate. +# +# `tests/test_ci_scope.py` replays every case through the real classifier, +# so an edit to `scope-rules.toml` that changes behaviour fails with the case +# name here. The cases encode intent, not implementation: each one says which +# suites a kind of change must start. +# +# `suites` is a list of globs over suite ids; `["*"]` means every suite and +# `[]` means none. + +[[case]] +name = "server crate: the whole fuzz fleet except the pruning suite" +note = "A change to `src/api` reaches every suite that builds the server binary." +paths = [ + "src/api/src/lib.rs", +] +suites = [ + "fuzz-standalone/database-and-regular-table", + "fuzz-standalone/logical-table", + "fuzz-unstable/standalone", + "fuzz-distributed/database-and-regular-table", + "fuzz-distributed/logical-table", + "fuzz-distributed/repartition", + "fuzz-chaos/migration", + "fuzz-chaos/failover", + "fuzz-chaos/repartition", + "export-import-v2", +] + +[[case]] +name = "storage engine: every fuzz suite, including remote-WAL pruning" +note = "logical-pruning needs engine and DDL code, so an engine change starts it." +paths = [ + "src/mito2/src/lib.rs", +] +suites = [ + "fuzz-standalone/*", + "fuzz-unstable/*", + "fuzz-distributed/*", + "fuzz-chaos/*", + "export-import-v2", +] + +[[case]] +name = "in-process DDL code registers the procedures standalone executes" +paths = [ + "src/common/meta/src/lib.rs", +] +suites = [ + "fuzz-standalone/*", + "fuzz-unstable/*", + "fuzz-distributed/*", + "fuzz-chaos/*", + "export-import-v2", + "sqlness-kvbackend/*", +] + +[[case]] +name = "metasrv only: distributed and chaos families" +note = "Standalone builds a DdlManager in-process and never links a metasrv." +paths = [ + "src/meta-srv/src/discovery/lease.rs", +] +suites = [ + "fuzz-distributed/database-and-regular-table", + "fuzz-distributed/logical-table", + "fuzz-distributed/repartition", + "fuzz-distributed/logical-pruning", + "fuzz-chaos/migration", + "fuzz-chaos/failover", + "fuzz-chaos/repartition", + "sqlness-kvbackend/*", +] + +[[case]] +name = "CLI only touches the export/import path" +paths = [ + "src/cli/src/lib.rs", +] +suites = [ + "export-import-v2", +] + +[[case]] +name = "PromQL crate: no suite" +note = "rust.yml and the promql sqlness cases own it; the fuzz suites speak SQL." +paths = [ + "src/promql/src/lib.rs", +] +suites = [] + +[[case]] +name = "PromQL module inside the query crate is not fuzzable either" +note = "`@server` claims `src/query/**`, so the module must be exempt." +paths = [ + "src/query/src/promql/src/lib.rs", +] +suites = [] + +[[case]] +name = "flow engine: no suite" +paths = [ + "src/flow/src/lib.rs", +] +suites = [] + +[[case]] +name = "mixed diff: the PromQL file subtracts nothing" +note = "The engine change still starts every suite." +paths = [ + "src/mito2/src/lib.rs", + "src/promql/src/lib.rs", +] +suites = [ + "fuzz-standalone/*", + "fuzz-unstable/*", + "fuzz-distributed/*", + "fuzz-chaos/*", + "export-import-v2", +] + +[[case]] +name = "crate added after the rules were written" +note = "No suite claims it, so the gate widens instead of skipping everything." +paths = [ + "src/brand-new/src/lib.rs", +] +suites = ["*"] + +[[case]] +name = "integration harness change" +paths = [ + "tests-integration/tests/fuzz.rs", +] +suites = ["*"] + +[[case]] +name = "documentation only" +paths = [ + "docs/readme.md", +] +suites = [] + +[[case]] +name = "kind registry setup: every suite" +note = "The distributed and chaos fuzz jobs start their cluster with this script." +paths = [ + ".github/scripts/kind-with-registry.sh", +] +suites = ["*"] + +[[case]] +name = "fuzz corpus seeding: every suite" +note = "Every fuzz job seeds its corpus through the `fuzz-test` action first." +paths = [ + ".github/scripts/prepare-fuzz-corpus.sh", +] +suites = ["*"] diff --git a/.github/ci/tests/test_ci_scope.py b/.github/ci/tests/test_ci_scope.py new file mode 100755 index 00000000000..77839aaaffa --- /dev/null +++ b/.github/ci/tests/test_ci_scope.py @@ -0,0 +1,1018 @@ +#!/usr/bin/env python3 +# 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. + +"""Offline tests for the integration CI admission gate.""" + +import contextlib +import fnmatch +import io +import json +import os +import re +import pathlib +import sys +import tempfile +import tomllib +import unittest +from urllib.parse import parse_qs, urlparse + +GATE_DIR = pathlib.Path(__file__).resolve().parents[1] +WORKFLOWS = GATE_DIR.parent / "workflows" +REPO_ROOT = GATE_DIR.parents[1] +sys.path.insert(0, str(GATE_DIR)) + +import ci_scope # noqa: E402 +import ci_scope_rules # noqa: E402 + +SCOPE_CASES = pathlib.Path(__file__).resolve().parent / "fixtures" / "scope-cases.toml" +from ci_scope_rules import ( # noqa: E402 + ALWAYS_RUN, + PATH_GROUPS, + SUITES, + suite_paths, +) + +def file_entry(name, patch="@@ -1 +1 @@\n- old\n+ new\n", status="modified"): + entry = { + "filename": name, + "status": status, + "additions": 1, + "deletions": 1, + } + if patch is not None: + entry["patch"] = patch + return entry + + +class FakeHTTP: + """Records requests and serves canned GitHub API responses.""" + + def __init__(self): + self.calls = [] + self.file_pages = {1: []} + self.labels = [] + self.pull_status = 200 + self.draft = False + self.head_sha = "" + self.base_sha = "" + self.file_queries = [] + + def __call__(self, method, url, headers, body=None, timeout=None): + parsed = urlparse(url) + query = parse_qs(parsed.query) + self.calls.append((method, parsed.path, query.get("page", ["1"])[0])) + if parsed.path.endswith("/files"): + self.file_queries.append(dict(query)) + page = int(query.get("page", ["1"])[0]) + per_page = int(query.get("per_page", ["100"])[0]) + items = [] + for pageno in sorted(self.file_pages): + canned = self.file_pages[pageno] + if isinstance(canned, int): + # An int stands for an HTTP error status for that page. + if pageno == page: + return canned, b"{}" + continue + items.extend(canned) + # `file_pages` is stored at per_page=100; honour the requested + # window so a probe at per_page=1 cannot read past item 1. + start = (page - 1) * per_page + return 200, json.dumps(items[start : start + per_page]).encode() + if parsed.path.endswith("/pulls/1"): + if self.pull_status != 200: + return self.pull_status, b"{}" + return 200, json.dumps( + { + "labels": [{"name": name} for name in self.labels], + "head": {"sha": self.head_sha}, + "base": {"sha": self.base_sha}, + } + ).encode() + return 404, b"{}" + + +class ScopeTestCase(unittest.TestCase): + def setUp(self): + self.fake = FakeHTTP() + self._original = ci_scope._request + ci_scope._request = self.fake + self._environ = dict(os.environ) + for key in ["CI_SCOPE_FORCE_ALL"]: + os.environ.pop(key, None) + os.environ["GITHUB_API_URL"] = "https://api.github.com" + + def tearDown(self): + ci_scope._request = self._original + os.environ.clear() + os.environ.update(self._environ) + + def plan(self, files, pages=None, event="pull_request", extra_args=(), pr=1): + if pages is None: + pages = {1: files} + self.fake.file_pages = dict(pages) + args = ci_scope.parse_args( + ["--repo", "GreptimeTeam/greptimedb", "--pr", str(pr), "--event", event] + + list(extra_args) + ) + return ci_scope.build_plan(args, "token") + + def assert_all_run(self, plan): + for suite_id, entry in plan["suites"].items(): + self.assertTrue(entry["run"], suite_id) + + +class GlobTests(unittest.TestCase): + def test_double_star_crosses_directories(self): + regex = ci_scope.glob_to_regex("src/**/*.rs") + self.assertTrue(regex.match("src/a.rs")) + self.assertTrue(regex.match("src/a/b/c.rs")) + self.assertFalse(regex.match("tests/a.rs")) + + def test_single_star_stays_in_one_segment(self): + regex = ci_scope.glob_to_regex("src/*.rs") + self.assertTrue(regex.match("src/main.rs")) + self.assertFalse(regex.match("src/cmd/main.rs")) + self.assertFalse(regex.match("src/main.toml")) + + def test_question_mark_matches_one_character(self): + regex = ci_scope.glob_to_regex("a?.rs") + self.assertTrue(regex.match("ab.rs")) + self.assertFalse(regex.match("abc.rs")) + self.assertFalse(regex.match("a/.rs")) + + def test_suffix_glob(self): + regex = ci_scope.glob_to_regex(".cargo/**") + self.assertTrue(regex.match(".cargo/config.toml")) + self.assertFalse(regex.match("cargo/config.toml")) + + def test_suite_groups_expand(self): + for group in PATH_GROUPS: + self.assertTrue(PATH_GROUPS[group]) + self.assertTrue(any( + entry == "@" + group + for suite in SUITES.values() + for entry in suite["paths"] + ), group) + paths = suite_paths(SUITES["fuzz-distributed/repartition"]) + self.assertIn("src/mito2/**", paths) + self.assertIn("src/meta-srv/**", paths) + self.assertNotIn("@server", paths) + + def test_every_gate_routes_to_an_output(self): + # Ask the producer, not a copy of its key list: a new gate must fail + # here until `derive_outputs` routes it to a gate key. + plan = { + "mode": "static", + "full": True, + "notes": [], + "suites": { + suite_id: {"run": True, "source": "static", "reason": ""} + for suite_id in SUITES + }, + } + outputs, _ = ci_scope.derive_outputs(plan) + gates = {suite["gate"] for suite in SUITES.values()} + self.assertLessEqual(gates, set(outputs)) + + def test_suite_metadata(self): + for suite_id, suite in SUITES.items(): + self.assertEqual(suite_id, suite_id.strip()) + self.assertTrue(suite["description"]) + self.assertEqual(suite_id.split("/")[0].replace("-", "_"), suite["gate"]) + if suite["gate"] in ci_scope.FUZZ_GATES: + self.assertEqual(ci_scope.suite_group(suite_id), suite_id.split("/", 1)[1]) + + +class ForcedModeTests(ScopeTestCase): + def test_merge_group_forces_without_http(self): + plan = self.plan([], event="merge_group") + self.assertEqual(plan["mode"], "forced") + self.assert_all_run(plan) + self.assertEqual(self.fake.calls, []) + + def test_schedule_and_workflow_dispatch_force(self): + for event in ("schedule", "workflow_dispatch"): + plan = self.plan([], event=event) + self.assertEqual(plan["mode"], "forced", event) + + def test_force_all_env_forces(self): + os.environ["CI_SCOPE_FORCE_ALL"] = "true" + plan = self.plan([], event="pull_request") + self.assertEqual(plan["mode"], "forced") + + def test_ci_full_label_forces(self): + self.fake.labels = ["ci:full"] + plan = self.plan([file_entry("src/mito2/src/lib.rs")]) + self.assertEqual(plan["mode"], "forced") + self.assert_all_run(plan) + + +class UncoveredChangeTests(ScopeTestCase): + def test_uncovered_crate_runs_everything(self): + # `src/common/telemetry` is mapped by the server group; a crate added + # later is not, and must widen the gate. + plan = self.plan([file_entry("src/common/brand-new/src/lib.rs")]) + self.assertTrue(plan["full"], plan["notes"]) + self.assertEqual(plan["suites"]["export-import-v2"]["source"], "uncovered") + self.assertIn("uncovered build change", " ".join(plan["notes"])) + + def test_mapped_crate_keeps_the_path_decision(self): + plan = self.plan([file_entry("src/query/src/lib.rs")]) + self.assertFalse(plan["full"]) + self.assertTrue(plan["suites"]["fuzz-standalone/logical-table"]["run"]) + # `src/query` reaches `COPY DATABASE` and the chaos scenarios too. + self.assertTrue(plan["suites"]["export-import-v2"]["run"]) + self.assertTrue(plan["suites"]["fuzz-chaos/migration"]["run"]) + + def test_exempt_trees_do_not_widen(self): + plan = self.plan( + [file_entry("src/promql/src/lib.rs"), file_entry("src/flow/src/lib.rs")] + ) + self.assertFalse(plan["full"]) + for entry in plan["suites"].values(): + self.assertFalse(entry["run"], entry) + + def test_exempt_tree_is_not_claimed_by_a_broader_glob(self): + # `@server` claims `src/query/**`; the PromQL module under it must be + # dropped before matching, and must not count as uncovered either. + plan = self.plan([file_entry("src/query/src/promql/src/lib.rs")]) + self.assertFalse(plan["full"]) + for entry in plan["suites"].values(): + self.assertFalse(entry["run"], entry) + + +class StaticDecisionTests(ScopeTestCase): + def outputs(self, plan): + return ci_scope.derive_outputs(plan)[0] + + def test_always_run_forces_everything(self): + plan = self.plan([file_entry("Cargo.lock")]) + self.assertEqual(plan["mode"], "static") + self.assertTrue(plan["full"]) + self.assert_all_run(plan) + + def test_meta_only_skips_standalone_fuzz(self): + plan = self.plan([file_entry("src/meta-srv/src/lib.rs")]) + self.assertEqual(plan["mode"], "static") + outputs, legs = ci_scope.derive_outputs(plan) + self.assertEqual(outputs["fuzz_standalone"], "false") + self.assertEqual(outputs["fuzz_unstable"], "false") + self.assertIn("migration", legs["fuzz_chaos"]) + self.assertIn("logical-pruning", legs["fuzz_distributed"]) + self.assertEqual(outputs["export_import_v2"], "false") + self.assertEqual(outputs["any_fuzz"], "true") + self.assertEqual(outputs["build_ci_needed"], "true") + + def test_flow_and_promql_skip_every_gated_suite(self): + plan = self.plan( + [file_entry("src/flow/src/lib.rs"), file_entry("src/promql/src/lib.rs")] + ) + self.assertEqual(plan["mode"], "static") + for entry in plan["suites"].values(): + self.assertFalse(entry["run"], entry) + outputs, legs = ci_scope.derive_outputs(plan) + self.assertEqual(outputs["any_fuzz"], "false") + self.assertEqual(outputs["build_ci_needed"], "false") + self.assertEqual(outputs["export_import_v2"], "false") + for gate in ci_scope.FUZZ_GATES: + self.assertEqual(outputs[gate], "false") + self.assertTrue(all(not values for values in legs.values())) + + def test_docs_only_short_circuits(self): + plan = self.plan([file_entry("docs/readme.md")]) + self.assertEqual(plan["mode"], "static") + self.assertIn("no meaningful changes", " ".join(plan["notes"])) + + def test_mapped_change_records_the_path_rule_source(self): + plan = self.plan([file_entry("src/mito2/src/lib.rs")]) + self.assertEqual(plan["mode"], "static") + self.assertEqual(plan["suites"]["fuzz-standalone/logical-table"]["source"], "static") + + def test_file_cap_truncates(self): + files = [file_entry("src/mito2/f%d.rs" % index) for index in range(300)] + extra = [file_entry("src/mito2/overflow.rs")] + plan = self.plan(files, pages={1: files[:100], 2: files[100:200], 3: files[200:], 4: extra}) + self.assertTrue(plan["full"]) + self.assert_all_run(plan) + self.assertTrue(all(entry["source"] == "always-run" for entry in plan["suites"].values())) + + def test_exactly_max_files_is_not_truncated(self): + files = [file_entry("src/promql/f%d.rs" % index) for index in range(300)] + plan = self.plan([], pages={1: files[:100], 2: files[100:200], 3: files[200:], 4: []}) + self.assertFalse(plan["full"], plan["notes"]) + self.assertNotIn("truncated", " ".join(plan["notes"])) + self.assertEqual(len(self.fake.file_queries), 4) + self.assertEqual( + self.fake.file_queries[-1], {"per_page": ["100"], "page": ["4"]} + ) + + def test_force_push_after_the_run_widens_the_gate(self): + self.fake.head_sha = "b" * 40 + plan = self.plan( + [file_entry("src/promql/src/lib.rs")], + extra_args=("--head-sha", "a" * 40, "--base-sha", "c" * 40), + ) + self.assertTrue(plan["full"]) + self.assertIn("head moved from aaaaaaaaaaaa", " ".join(plan["notes"])) + + def test_matching_revisions_keep_the_static_verdict(self): + self.fake.head_sha = "a" * 40 + self.fake.base_sha = "c" * 40 + plan = self.plan( + [file_entry("src/promql/src/lib.rs")], + extra_args=("--head-sha", "a" * 40, "--base-sha", "c" * 40), + ) + self.assertFalse(plan["full"]) + self.assertNotIn("moved", " ".join(plan["notes"])) + + def test_pull_request_is_read_once_per_run(self): + self.plan([file_entry("src/mito2/src/lib.rs")]) + pulls = [ + path + for _, path, _ in self.fake.calls + if path.endswith("/pulls/1") + ] + self.assertEqual(len(pulls), 1) + + def test_failed_completeness_probe_fails_instead_of_narrowing(self): + # Raising leaves the gate outputs empty, which the workflow treats as + # "run everything"; returning False here would silently classify a + # truncated file list as complete. + files = [file_entry("src/mito2/f%d.rs" % index) for index in range(300)] + with self.assertRaises(ci_scope.ScopeError): + self.plan( + files, + pages={1: files[:100], 2: files[100:200], 3: files[200:], 4: 500}, + ) + + def test_file_list_pagination_stops_at_short_page(self): + files = [file_entry("src/mito2/f%d.rs" % index) for index in range(120)] + self.plan(files, pages={1: files[:100], 2: files[100:]}) + pages = [ + page for _, path, page in self.fake.calls if path.endswith("/files") + ] + self.assertEqual(pages, ["1", "2"]) + + +# Crates no gated suite exercises. Changing one either runs everything through +# the always-run floor, or reaches CI that is not gated here. +UNMAPPED_CRATES = { + "src/flow": "flow execution is covered by rust.yml and sqlness", + "src/promql": "PromQL execution is covered by rust.yml and sqlness", + "tests-fuzz": "in ALWAYS_RUN", + "tests-integration": "in ALWAYS_RUN", +} + + +def workspace_crates(): + """Every workspace member directory, so nested crates are not collapsed.""" + data = tomllib.loads((REPO_ROOT / "Cargo.toml").read_text(encoding="utf-8")) + return sorted(member.rstrip("/") for member in data["workspace"]["members"]) + + +def rules_doc(): + return { + "matching": { + "always_run": ["Cargo.toml"], + "src_tree": "src/", + "exempt": [{"path": "src/flow/", "reason": "test fixture"}], + }, + "gates": {"fuzz_x": "fuzztest"}, + "groups": {"server": ["src/a/**"]}, + "suites": { + "fuzz-x/g": { + "gate": "fuzz_x", + "paths": ["@server"], + "description": "a suite", + } + }, + } + + +class RulesFileTests(unittest.TestCase): + def test_repository_rules_are_valid(self): + rules = ci_scope_rules.load() + self.assertEqual(len(rules["suites"]), 13) + + def test_unknown_group_is_rejected(self): + doc = rules_doc() + doc["suites"]["fuzz-x/g"]["paths"] = ["@nope"] + with self.assertRaises(ci_scope_rules.RulesError): + ci_scope_rules.validate(doc) + + def test_group_cycle_is_rejected(self): + doc = rules_doc() + doc["groups"]["a"] = ["@b"] + doc["groups"]["b"] = ["@a"] + with self.assertRaises(ci_scope_rules.RulesError): + ci_scope_rules.validate(doc) + + def test_gate_must_match_the_suite_prefix(self): + doc = rules_doc() + doc["suites"]["fuzz-x/g"]["gate"] = "fuzz_y" + with self.assertRaises(ci_scope_rules.RulesError): + ci_scope_rules.validate(doc) + + def test_unknown_key_in_matching_is_ignored_but_missing_is_not(self): + doc = rules_doc() + del doc["matching"]["src_tree"] + with self.assertRaises(ci_scope_rules.RulesError): + ci_scope_rules.validate(doc) + + def test_validator_entry_point_succeeds_on_the_repository_rules(self): + out = io.StringIO() + with contextlib.redirect_stdout(out): + self.assertEqual(ci_scope_rules.main([]), 0) + self.assertIn("scope rules ok", out.getvalue()) + + def test_malformed_toml_names_the_file(self): + with tempfile.TemporaryDirectory() as tmp: + path = pathlib.Path(tmp) / "rules.toml" + path.write_text("this is not toml [[[") + with self.assertRaises(ci_scope_rules.RulesError) as ctx: + ci_scope_rules.load(path) + self.assertIn("not valid TOML", str(ctx.exception)) + + +def stub_rules(**overrides): + """A rules dict in the shape `validate` returns, naming only real paths.""" + rules = { + "always_run": ["Cargo.toml"], + "src_tree": "src/", + "exempt_trees": [], + "gates": {"fuzz_x": "fuzztest"}, + "groups": {"server": ["Cargo.toml"]}, + "suites": { + "fuzz-x": { + "gate": "fuzz_x", + "paths": ["@server"], + "description": "a suite", + } + }, + } + rules.update(overrides) + return rules + + +class PathCheckTests(unittest.TestCase): + def test_repository_rules_name_only_existing_paths(self): + self.assertEqual(ci_scope_rules.missing_paths(), []) + + def test_a_vanished_path_is_reported_with_its_prefix(self): + doc = stub_rules() + doc["suites"]["fuzz-x"]["paths"] = ["@server", "src/does-not-exist/**"] + self.assertEqual( + ci_scope_rules.missing_paths(doc), + [("src/does-not-exist/**", "src/does-not-exist/")], + ) + + def test_patterns_without_a_fixed_prefix_are_skipped(self): + doc = stub_rules(groups={"server": ["*.md"]}, always_run=["*.md"]) + self.assertEqual(ci_scope_rules.missing_paths(doc), []) + + + def test_cli_reports_a_missing_path(self): + text = """\ +[matching] +always_run = ["Cargo.toml"] +src_tree = "src/" + +[[matching.exempt]] +path = "src/flow/" +reason = "test fixture" + +[[matching.exempt]] +path = "src/promql/" +reason = "test fixture" + +[gates] +fuzz_x = "fuzztest" + +[groups] +server = ["src/does-not-exist/**"] + +[suites."fuzz-x"] +gate = "fuzz_x" +paths = ["@server"] +description = "a suite" +""" + with tempfile.TemporaryDirectory() as tmp: + path = pathlib.Path(tmp) / "scope-rules.toml" + path.write_text(text) + original = ci_scope_rules.RULES_PATH + ci_scope_rules.RULES_PATH = path + err = io.StringIO() + try: + with contextlib.redirect_stderr(err): + code = ci_scope_rules.main(["--check-paths"]) + finally: + ci_scope_rules.RULES_PATH = original + self.assertEqual(code, 1) + self.assertIn("missing path: src/does-not-exist/**", err.getvalue()) + + def test_cli_succeeds_on_the_repository_rules(self): + out, err = io.StringIO(), io.StringIO() + with contextlib.redirect_stdout(out), contextlib.redirect_stderr(err): + code = ci_scope_rules.main(["--check-paths"]) + self.assertEqual(code, 0) + self.assertEqual(err.getvalue(), "") + + + +class RulesReasonTests(unittest.TestCase): + def test_exclusion_without_a_reason_is_rejected(self): + doc = rules_doc() + del doc["matching"]["exempt"][0]["reason"] + with self.assertRaises(ci_scope_rules.RulesError): + ci_scope_rules.validate(doc) + + def test_blank_reason_is_rejected(self): + doc = rules_doc() + doc["matching"]["exempt"][0]["reason"] = " " + with self.assertRaises(ci_scope_rules.RulesError): + ci_scope_rules.validate(doc) + + def test_repository_rules_explain_every_exclusion(self): + raw = tomllib.loads(ci_scope_rules.RULES_PATH.read_text()) + matching = raw["matching"] + self.assertTrue(matching["exempt"]) + for entry in matching["exempt"]: + self.assertTrue(entry["reason"].strip(), entry) + + +class ScopeCaseTests(ScopeTestCase): + """Replay `fixtures/scope-cases.toml` through the real classifier.""" + + @staticmethod + def cases(): + return tomllib.loads(SCOPE_CASES.read_text())["case"] + + def test_every_case_matches_the_classifier(self): + for case in self.cases(): + with self.subTest(case=case["name"]): + plan = self.plan([file_entry(path) for path in case["paths"]]) + expected = { + suite_id + for suite_id in ci_scope.SUITES + if any( + fnmatch.fnmatchcase(suite_id, pattern) + for pattern in case["suites"] + ) + } + actual = { + suite_id + for suite_id, entry in plan["suites"].items() + if entry["run"] + } + self.assertEqual(expected, actual, plan["notes"]) + + def test_cases_cover_every_suite(self): + covered = set() + for case in self.cases(): + covered |= { + suite_id + for suite_id in ci_scope.SUITES + if any( + fnmatch.fnmatchcase(suite_id, pattern) + for pattern in case["suites"] + ) + } + self.assertEqual(covered, set(ci_scope.SUITES)) + + +class SummaryTests(ScopeTestCase): + def summary(self, plan, path): + args = ci_scope.parse_args( + ["--repo", "GreptimeTeam/greptimedb", "--pr", "1", "--summary", path] + ) + ci_scope.emit(plan, args) + return pathlib.Path(path).read_text() + + def test_summary_names_the_files_that_matched(self): + plan = self.plan([file_entry("src/mito2/src/lib.rs")]) + with tempfile.TemporaryDirectory() as tmp: + text = self.summary(plan, str(pathlib.Path(tmp) / "summary.md")) + self.assertIn("### CI scope: mode=static", text) + row = next( + line + for line in text.splitlines() + if line.startswith("| `fuzz-standalone/logical-table`") + ) + self.assertIn("| run |", row) + self.assertIn("`src/mito2/src/lib.rs`", row) + + def test_summary_marks_a_skipped_suite_without_files(self): + plan = self.plan([file_entry("src/cli/src/lib.rs")]) + with tempfile.TemporaryDirectory() as tmp: + text = self.summary(plan, str(pathlib.Path(tmp) / "summary.md")) + row = next( + line + for line in text.splitlines() + if line.startswith("| `fuzz-standalone/database-and-regular-table`") + ) + self.assertIn("| skip |", row) + self.assertTrue(row.endswith("| — |"), row) + + def test_summary_truncates_long_file_lists(self): + files = [file_entry(f"src/mito2/src/f{index}.rs") for index in range(5)] + plan = self.plan(files) + with tempfile.TemporaryDirectory() as tmp: + text = self.summary(plan, str(pathlib.Path(tmp) / "summary.md")) + self.assertIn("+2 more", text) + + def test_summary_appends_instead_of_overwriting(self): + plan = self.plan([file_entry("src/mito2/src/lib.rs")]) + with tempfile.TemporaryDirectory() as tmp: + path = pathlib.Path(tmp) / "summary.md" + path.write_text("existing\n") + self.summary(plan, str(path)) + text = path.read_text() + self.assertTrue(text.startswith("existing\n")) + self.assertEqual(text.count("### CI scope"), 1) + + def test_summary_carries_the_notes(self): + plan = self.plan([file_entry("Cargo.lock")]) + with tempfile.TemporaryDirectory() as tmp: + text = self.summary(plan, str(pathlib.Path(tmp) / "summary.md")) + self.assertIn("
notes", text) + + +class CrateCoverageTests(unittest.TestCase): + """A new crate must not silently match no suite: the runtime net would + catch it by running everything, but the rule table should be updated.""" + + def test_every_workspace_crate_is_mapped_or_excused(self): + patterns = [ + ci_scope.glob_to_regex(path) + for suite in SUITES.values() + for path in suite_paths(suite) + ] + unmapped = [ + crate + for crate in workspace_crates() + if crate not in UNMAPPED_CRATES + and not any(p.match(crate + "/src/lib.rs") for p in patterns) + ] + self.assertEqual(unmapped, [], "crate matches no suite and no excuse") + + def test_excuses_reference_existing_crates(self): + stale = sorted(set(UNMAPPED_CRATES) - set(workspace_crates())) + self.assertEqual(stale, [], "excuse points at a crate that is gone") + + +class PathRuleCoverageTests(ScopeTestCase): + """A rule must select every suite the change can actually break.""" + + def test_copy_parser_selects_export_import(self): + plan = self.plan([file_entry("src/sql/src/parsers/copy_parser.rs")]) + self.assertTrue(plan["suites"]["export-import-v2"]["run"]) + + def test_shared_ddl_selects_standalone_fuzz(self): + plan = self.plan([file_entry("src/common/meta/src/ddl/create_table.rs")]) + for suite in ( + "fuzz-standalone/database-and-regular-table", + "fuzz-standalone/logical-table", + "fuzz-unstable/standalone", + ): + self.assertTrue(plan["suites"][suite]["run"], suite) + + def test_sqlness_runner_selects_the_kvbackend_legs(self): + # `--setup-pg` / `--setup-mysql` live in the runner and only these legs + # pass them, so a change there must not skip the only job that uses it. + plan = self.plan([file_entry("tests/runner/src/server_mode.rs")]) + for suite in ("sqlness-kvbackend/metasrv-postgres", "sqlness-kvbackend/metasrv-mysql"): + self.assertTrue(plan["suites"][suite]["run"], suite) + + def test_shared_ddl_selects_export_import(self): + # The E2E creates databases, tables and views and recreates them on + # import, all through the in-process DDL manager. + plan = self.plan([file_entry("src/common/procedure/src/lib.rs")]) + self.assertTrue(plan["suites"]["export-import-v2"]["run"]) + + def test_group_references_resolve_recursively(self): + self.assertIn("src/common/meta/**", ci_scope.suite_paths(SUITES["fuzz-chaos/migration"])) + + +class OutputTests(ScopeTestCase): + def test_flat_outputs(self): + plan = self.plan([file_entry("src/mito2/src/lib.rs")]) + outputs, legs = ci_scope.derive_outputs(plan) + self.assertEqual( + set(outputs), + { + "any_fuzz", + "build_ci_needed", + "fuzz_standalone", + "fuzz_distributed", + "fuzz_chaos", + "fuzz_unstable", + "export_import_v2", + "sqlness_kvbackend", + }, + ) + for key in outputs: + self.assertIn(outputs[key], ("true", "false"), key) + self.assertEqual( + sorted(legs["fuzz_standalone"]), + ["database-and-regular-table", "logical-table"], + ) + + def test_emit_writes_the_gate_outputs(self): + import tempfile + + with tempfile.TemporaryDirectory() as tmp: + out_path = os.path.join(tmp, "github_output") + args = ci_scope.parse_args( + [ + "--repo", + "GreptimeTeam/greptimedb", + "--pr", + "1", + "--event", + "merge_group", + "--github-output", + out_path, + ] + ) + ci_scope.emit(ci_scope.build_plan(args, "token"), args) + with open(out_path, encoding="utf-8") as handle: + text = handle.read() + self.assertIn("any_fuzz=true" + "\n", text) + self.assertIn("fuzz_unstable=true", text) + self.assertIn("fuzz_chaos=true", text) + self.assertEqual(text.count("\n"), 8) + + + +class GateJobContractTests(unittest.TestCase): + """The rules say which job each gate switches; prove the YAML agrees.""" + + WORKFLOW = WORKFLOWS / "integration.yml" + + def job_blocks(self): + blocks = {} + current = None + for line in self.WORKFLOW.read_text().splitlines(): + marker = re.match(r"^ ([A-Za-z0-9_-]+):$", line) + if marker: + current = marker.group(1) + blocks[current] = [] + elif current: + blocks[current].append(line) + return {name: "\n".join(lines) for name, lines in blocks.items()} + + def test_every_gate_names_a_job_that_reads_it(self): + blocks = self.job_blocks() + for gate, job in ci_scope_rules.GATE_JOBS.items(): + self.assertIn(job, blocks, gate) + self.assertIn(f"needs.scope.outputs.{gate} != 'false'", blocks[job], gate) + + def test_every_gated_job_is_named_by_a_gate(self): + derived = {"any_fuzz", "build_ci_needed"} + declared = set(ci_scope_rules.GATE_JOBS) + for name, block in self.job_blocks().items(): + for gate in re.findall(r"needs\.scope\.outputs\.([a-z0-9_]+)", block): + self.assertIn(gate, declared | derived, name) + + def test_a_suite_gate_missing_from_the_table_is_rejected(self): + doc = rules_doc() + del doc["gates"]["fuzz_x"] + with self.assertRaises(ci_scope_rules.RulesError) as ctx: + ci_scope_rules.validate(doc) + self.assertIn("missing from [gates]", str(ctx.exception)) + + def test_gates_no_suite_uses_are_rejected(self): + doc = rules_doc() + doc["gates"]["fuzz_y"] = "fuzztest" + with self.assertRaises(ci_scope_rules.RulesError) as ctx: + ci_scope_rules.validate(doc) + self.assertIn("no suite uses", str(ctx.exception)) + + +class WorkflowContractTests(unittest.TestCase): + """Keep the YAML gate contract in sync with the classifier outputs.""" + + WORKFLOW = WORKFLOWS / "integration.yml" + + def job_blocks(self): + text = self.WORKFLOW.read_text() + blocks = {} + current = None + for line in text.splitlines(): + marker = re.match(r"^ ([A-Za-z0-9_-]+):$", line) + if marker: + current = marker.group(1) + blocks[current] = [] + elif current: + blocks[current].append(line) + return {name: "\n".join(lines) for name, lines in blocks.items()} + + def test_every_script_a_gated_job_runs_is_always_run(self): + """A script a gated job executes must widen the gate. + + A change confined to such a script would otherwise select no suite, so + the jobs that run it are skipped exactly when they are worth running. + The reference is often not in the workflow: a composite action the job + `uses` can run the script, so the search follows those actions too. + """ + action_use = re.compile(r"uses:\s*\./(\.github/actions/[A-Za-z0-9_.-]+)") + script_ref = re.compile( + r"(?:\./)?(\.github/[A-Za-z0-9_./-]+\.(?:sh|py|js|cjs|ts))\b" + ) + always = [ci_scope.glob_to_regex(pattern) for pattern in ALWAYS_RUN] + + gated = [b for b in self.job_blocks().values() if "needs.scope.outputs" in b] + pending = {ref for block in gated for ref in script_ref.findall(block)} + pending |= {a for block in gated for a in action_use.findall(block)} + seen = set() + while pending: + reference = pending.pop() + if reference in seen: + continue + seen.add(reference) + if not reference.startswith(".github/actions/"): + continue + for suffix in ("action.yml", "action.yaml"): + path = REPO_ROOT / reference / suffix + if path.exists(): + text = path.read_text() + pending |= set(script_ref.findall(text)) + pending |= set(action_use.findall(text)) + break + uncovered = sorted( + ref for ref in seen if not any(pattern.match(ref) for pattern in always) + ) + self.assertEqual(uncovered, [], "a gated job runs a script ALWAYS_RUN misses") + + def test_gates_treat_a_missing_output_as_run(self): + blocks = self.job_blocks() + gated = [ + "export-import-v2-e2e", + "prepare-fuzz-targets", + "fuzztest", + "unstable-fuzztest", + "build-greptime-ci", + "distributed-fuzztest", + "distributed-fuzztest-with-chaos", + "sqlness-kvbackend", + ] + declared = set( + re.findall( + r"^ ([a-z0-9_]+): \$\{\{ steps\.classify\.outputs\.\1 \}\}$", + blocks["scope"], + re.MULTILINE, + ) + ) + self.assertIn("any_fuzz", declared) + seen = set() + for name in gated: + self.assertIn("scope", blocks[name], name) + match = re.search( + r"needs\.scope\.outputs\.([a-z0-9_]+) != 'false'", blocks[name] + ) + self.assertTrue(match, name) + seen.add(match.group(1)) + self.assertEqual(seen, declared) + + def test_gates_override_the_implicit_success_check(self): + # A condition without a status function carries an implicit success() + # over `needs`, so a scope job that times out or dies on the runner + # would skip the gates despite the fail-open output comparison. + blocks = self.job_blocks() + for name in ( + "export-import-v2-e2e", + "prepare-fuzz-targets", + "fuzztest", + "unstable-fuzztest", + "build-greptime-ci", + "distributed-fuzztest", + "distributed-fuzztest-with-chaos", + ): + block = blocks[name] + condition = re.search(r"^ if: \$\{\{ (.*?) \}\}$", block, re.M) + self.assertTrue(condition, name) + self.assertTrue(condition.group(1).startswith("!cancelled() && "), name) + needs = re.search(r"^ needs: \[(.*?)\]$", block, re.M) + self.assertTrue(needs, name) + for dep in [item.strip() for item in needs.group(1).split(",")]: + if not dep or dep == "scope": + continue + key = f"needs['{dep}'].result" if "-" in dep else f"needs.{dep}.result" + self.assertIn(key + " == 'success'", condition.group(1), name) + + def test_scope_job_cannot_fail_and_skip_the_gates(self): + block = self.job_blocks()["scope"] + self.assertNotIn("Fail-open defaults", block) + self.assertIn("github.event.merge_group.base_sha", block) + self.assertIn("CI_SCOPE_FORCE_ALL", block) + # Every step that can fail must be continue-on-error: a failed scope job + # would make GitHub skip the gated jobs instead of running them. + steps = block.split("\n - ") + guarded = [ + step + for step in steps + if step.startswith("uses: actions/checkout") + or step.startswith("name: Classify") + ] + self.assertEqual(len(guarded), 2) + for step in guarded: + self.assertIn("continue-on-error: true", step, step.splitlines()[0]) + + + +class LogRenderingTests(ScopeTestCase): + """The step log is where a triage starts; the verdict must be there.""" + + def sample_plan(self): + return { + "mode": "static", + "head_sha": "a" * 40, + "base_sha": "b" * 40, + "full": False, + "notes": ["always-run: Cargo.lock"], + "suites": { + "fuzz-standalone/logical-table": { + "run": True, + "source": "static", + "reason": "touches mito2", + }, + "fuzz-chaos/migration": { + "run": False, + "source": "static", + "reason": "meta docs only", + }, + }, + } + + def test_header_names_mode_and_full(self): + self.assertEqual( + ci_scope.render_log(self.sample_plan()).splitlines()[0], + "CI scope: mode=static full=false", + ) + + def test_one_line_per_suite_with_source_and_reason(self): + lines = ci_scope.render_log(self.sample_plan()).splitlines() + self.assertIn(" run fuzz-standalone/logical-table", lines[1]) + self.assertIn("static", lines[1]) + self.assertIn("touches mito2", lines[1]) + self.assertIn(" skip fuzz-chaos/migration", lines[2]) + self.assertIn("meta docs only", lines[2]) + + def test_gates_and_notes_are_printed(self): + log = ci_scope.render_log(self.sample_plan()) + self.assertIn(" gates: fuzz_standalone=true", log) + self.assertIn("any_fuzz=true", log) + self.assertIn("export_import_v2=false", log) + self.assertIn(" note: always-run: Cargo.lock", log) + + def test_a_failed_pull_request_lookup_is_annotated_and_runs_everything(self): + # Drive the real failure path: the label lookup and the revision + # comparison share one request, so a 500 here must widen the gate and + # still produce the warning the annotation promises. + self.fake.pull_status = 500 + plan = self.plan([file_entry("src/promql/src/lib.rs")]) + self.assertTrue(plan["full"]) + self.assertIn("pull request lookup failed", " ".join(plan["notes"])) + annotations = ci_scope.log_annotations(plan) + self.assertEqual(len(annotations), 1) + self.assertTrue(annotations[0].startswith("::warning::"), annotations) + self.assertIn("every suite runs", annotations[0]) + + def test_main_prints_the_log(self): + self.fake.file_pages = {1: [file_entry("src/mito2/src/lib.rs")]} + stdout = io.StringIO() + with contextlib.redirect_stdout(stdout): + self.assertEqual( + ci_scope.main( + [ + "--repo", + "GreptimeTeam/greptimedb", + "--pr", + "1", + "--event", + "pull_request", + "--dry-run", + ] + ), + 0, + ) + text = stdout.getvalue() + self.assertIn("CI scope: mode=", text) + self.assertIn(" run fuzz-standalone/logical-table", text) + self.assertIn("any_fuzz=true", text) + + +if __name__ == "__main__": + unittest.main() diff --git a/.github/scripts/tests/test_ci_slash.py b/.github/scripts/tests/test_ci_slash.py index c8c01ceda15..e551a4f5297 100644 --- a/.github/scripts/tests/test_ci_slash.py +++ b/.github/scripts/tests/test_ci_slash.py @@ -136,7 +136,7 @@ class CiSlashTest(unittest.TestCase): def test_fork_rerun_gates_cover_standard_jobs(self): import re - for name, count in [("rust", 7), ("integration", 5), ("checks", 5), ("docs", 3)]: + for name, count in [("rust", 7), ("integration", 7), ("checks", 6), ("docs", 3)]: workflow = (SCRIPT.parents[1] / "workflows" / (name + ".yml")).read_text() gates = re.findall(r"^ if:.*github.run_attempt > 1.*$", workflow, re.MULTILINE) self.assertEqual(len(gates), count, name) diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml index 3bc54372d06..b3ff212ade4 100644 --- a/.github/workflows/checks.yml +++ b/.github/workflows/checks.yml @@ -92,6 +92,8 @@ jobs: run: .github/scripts/run-fuzz-targets-test.sh - name: Run CI comment-command tests run: python3 .github/scripts/tests/test_ci_slash.py + - name: Run CI scope gate tests + run: python3 .github/ci/tests/test_ci_scope.py - name: Run metric snapshot benchmark tests run: python3 -m unittest discover -s .github/scripts/tests -p 'test_metric_snapshot*.py' - name: Run query-regression Python tooling tests @@ -107,6 +109,23 @@ jobs: python3 tests/perf/test_agent_observability_summary.py python3 tests/perf/test_tracesbench_summary.py + scope-rules: + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} + name: CI Scope Rules + # Same runner image as the `scope` job in integration.yml: if the rules file + # stops parsing there, the gate silently runs every job, so this check exists + # to fail the pull request instead. + runs-on: ubuntu-slim + timeout-minutes: 5 + steps: + - uses: actions/checkout@v4 + with: + persist-credentials: false + - name: Show the interpreter + run: python3 -VV + - name: Validate .github/ci/scope-rules.toml + run: python3 .github/ci/ci_scope_rules.py --check-paths + check: if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} name: Check diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml index afce4a8c796..6106e5a2860 100644 --- a/.github/workflows/integration.yml +++ b/.github/workflows/integration.yml @@ -42,6 +42,57 @@ concurrency: cancel-in-progress: true jobs: + # Which jobs this gate skips, and why: .github/ci/README.md. + scope: + name: CI Scope + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} + # ubuntu-slim ships python3, git and curl; this job only classifies. + runs-on: ubuntu-slim + timeout-minutes: 10 + permissions: + contents: read + # `/pulls/{n}` and `/pulls/{n}/files`. An explicit block sets every + # unlisted permission to none, and a 403 here would silently widen every + # gate instead of classifying. + pull-requests: read + outputs: + any_fuzz: ${{ steps.classify.outputs.any_fuzz }} + build_ci_needed: ${{ steps.classify.outputs.build_ci_needed }} + fuzz_standalone: ${{ steps.classify.outputs.fuzz_standalone }} + fuzz_distributed: ${{ steps.classify.outputs.fuzz_distributed }} + fuzz_chaos: ${{ steps.classify.outputs.fuzz_chaos }} + fuzz_unstable: ${{ steps.classify.outputs.fuzz_unstable }} + export_import_v2: ${{ steps.classify.outputs.export_import_v2 }} + sqlness_kvbackend: ${{ steps.classify.outputs.sqlness_kvbackend }} + steps: + # Rules come from the pull request base (the queue base on merge_group), + # so a pull request cannot weaken its own gate. If this job or the Classify + # step fails, the outputs stay empty and every gate below treats an absent + # value as "run", which keeps CI fail-open. + - uses: actions/checkout@v4 + # A checkout failure leaves the outputs empty, which every gate treats + # as "run"; letting the step fail would skip the gates instead. + continue-on-error: true + with: + ref: ${{ github.event.pull_request.base.sha || github.event.merge_group.base_sha || github.sha }} + persist-credentials: false + - name: Classify + id: classify + continue-on-error: true + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + CI_SCOPE_FORCE_ALL: ${{ vars.CI_SCOPE_FORCE_ALL || 'false' }} + run: | + args=(--repo "$GITHUB_REPOSITORY" --event "$GITHUB_EVENT_NAME" + --github-output "$GITHUB_OUTPUT" + --summary "$GITHUB_STEP_SUMMARY") + if [[ "${GITHUB_EVENT_NAME}" == "pull_request" ]]; then + args+=(--pr "${{ github.event.pull_request.number }}" + --head-sha "${{ github.event.pull_request.head.sha }}" + --base-sha "${{ github.event.pull_request.base.sha }}") + fi + python3 .github/ci/ci_scope.py "${args[@]}" + build: if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} name: Build GreptimeDB binaries @@ -116,28 +167,12 @@ jobs: - name: "Basic" opts: "" kafka: false - postgres: false - mysql: false - name: "Remote WAL" opts: "-w kafka -k 127.0.0.1:9092" kafka: true - postgres: false - mysql: false - - name: "PostgreSQL KvBackend" - opts: "--setup-pg postgresql://greptimedb:admin@127.0.0.1:5432/postgres" - kafka: false - postgres: true - mysql: false - - name: "MySQL KvBackend" - opts: "--setup-mysql mysql://greptimedb:admin@127.0.0.1:3306/mysql" - kafka: false - postgres: false - mysql: true - name: "Flat format" opts: "--enable-flat-format" kafka: false - postgres: false - mysql: false timeout-minutes: 60 steps: - uses: actions/checkout@v4 @@ -149,16 +184,6 @@ jobs: working-directory: tests-integration/fixtures run: ../../.github/scripts/pull-test-deps-images.sh && docker compose up -d --wait kafka - - if: matrix.mode.postgres - name: Setup PostgreSQL - working-directory: tests-integration/fixtures - run: ../../.github/scripts/pull-test-deps-images.sh && docker compose up -d --wait postgres - - - if: matrix.mode.mysql - name: Setup MySQL - working-directory: tests-integration/fixtures - run: ../../.github/scripts/pull-test-deps-images.sh && docker compose up -d --wait mysql - - name: Download pre-built binaries uses: actions/download-artifact@v4 with: @@ -183,10 +208,68 @@ jobs: /tmp/sqlness* retention-days: 3 + sqlness-kvbackend: + # The two foreign KV backends are the only sqlness legs that are not + # required checks, so they are the only ones this gate may skip. They run + # the same cases as the `Basic` leg, so the diff only matters when it can + # reach the backend itself; the merge queue runs them regardless. + if: ${{ !cancelled() && github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) && needs.build.result == 'success' && needs.scope.outputs.sqlness_kvbackend != 'false' }} + name: Sqlness Test (${{ matrix.mode.name }}) + needs: [scope, build] + runs-on: ${{ matrix.os }} + strategy: + fail-fast: false + matrix: + os: [ ubuntu-latest ] + mode: + - name: "Metasrv on PostgreSQL" + id: "metasrv-postgres" + opts: "--setup-pg postgresql://greptimedb:admin@127.0.0.1:5432/postgres" + postgres: true + mysql: false + - name: "Metasrv on MySQL" + id: "metasrv-mysql" + opts: "--setup-mysql mysql://greptimedb:admin@127.0.0.1:3306/mysql" + postgres: false + mysql: true + timeout-minutes: 60 + steps: + - uses: actions/checkout@v4 + with: + persist-credentials: false + + - if: matrix.mode.postgres + name: Setup PostgreSQL + working-directory: tests-integration/fixtures + run: ../../.github/scripts/pull-test-deps-images.sh && docker compose up -d --wait postgres + + - if: matrix.mode.mysql + name: Setup MySQL + working-directory: tests-integration/fixtures + run: ../../.github/scripts/pull-test-deps-images.sh && docker compose up -d --wait mysql + + - name: Download pre-built binaries + uses: actions/download-artifact@v4 + with: + name: bins + path: . + - name: Unzip binaries + run: tar -xvf ./bins.tar.gz + - name: Run sqlness + run: RUST_BACKTRACE=1 ./bins/sqlness-runner bare ${{ matrix.mode.opts }} -c ./tests/cases --bins-dir ./bins --preserve-state + - name: Upload sqlness logs + if: failure() + uses: actions/upload-artifact@v4 + with: + name: sqlness-logs-${{ matrix.mode.name }} + path: | + /tmp/sqlness* + retention-days: 3 + export-import-v2-e2e: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} + if: ${{ !cancelled() && github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) && needs.build.result == 'success' && needs.scope.outputs.export_import_v2 != 'false' }} name: Export/Import V2 E2E Test - needs: build + needs: [scope, build] runs-on: ubuntu-latest timeout-minutes: 60 steps: @@ -319,8 +402,9 @@ jobs: retention-days: 3 prepare-fuzz-targets: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && ((github.event_name == 'pull_request' && github.event.pull_request.draft == false) || (github.event_name == 'workflow_dispatch' && inputs.ci_command)) && (github.event_name == 'pull_request' || inputs.fuzz_profile != 'none') }} + if: ${{ !cancelled() && github.repository == 'GreptimeTeam/greptimedb' && ((github.event_name == 'pull_request' && github.event.pull_request.draft == false) || (github.event_name == 'workflow_dispatch' && inputs.ci_command) || github.event_name == 'merge_group') && (github.event_name == 'pull_request' || github.event_name == 'merge_group' || inputs.fuzz_profile != 'none') && needs.scope.outputs.any_fuzz != 'false' }} name: Prepare Fuzz Targets + needs: [scope] runs-on: ubuntu-latest timeout-minutes: 60 steps: @@ -433,9 +517,9 @@ jobs: if-no-files-found: error fuzztest: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && ((github.event_name == 'pull_request' && github.event.pull_request.draft == false) || (github.event_name == 'workflow_dispatch' && inputs.ci_command)) && (github.event_name == 'pull_request' || inputs.fuzz_profile == 'standalone' || inputs.fuzz_profile == 'all') }} + if: ${{ !cancelled() && github.repository == 'GreptimeTeam/greptimedb' && ((github.event_name == 'pull_request' && github.event.pull_request.draft == false) || (github.event_name == 'workflow_dispatch' && inputs.ci_command) || github.event_name == 'merge_group') && (github.event_name == 'pull_request' || github.event_name == 'merge_group' || inputs.fuzz_profile == 'standalone' || inputs.fuzz_profile == 'all') && needs.build.result == 'success' && needs['prepare-fuzz-targets'].result == 'success' && needs.scope.outputs.fuzz_standalone != 'false' }} name: Fuzz Test (Standalone, ${{ matrix.group.name }}) - needs: [build, prepare-fuzz-targets] + needs: [build, prepare-fuzz-targets, scope] runs-on: ubuntu-latest timeout-minutes: 60 strategy: @@ -514,9 +598,9 @@ jobs: artifact-name: fuzz-standalone-${{ matrix.group.id }} unstable-fuzztest: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && ((github.event_name == 'pull_request' && github.event.pull_request.draft == false) || (github.event_name == 'workflow_dispatch' && inputs.ci_command)) && (github.event_name == 'pull_request' || inputs.fuzz_profile == 'standalone' || inputs.fuzz_profile == 'all') }} + if: ${{ !cancelled() && github.repository == 'GreptimeTeam/greptimedb' && ((github.event_name == 'pull_request' && github.event.pull_request.draft == false) || (github.event_name == 'workflow_dispatch' && inputs.ci_command) || github.event_name == 'merge_group') && (github.event_name == 'pull_request' || github.event_name == 'merge_group' || inputs.fuzz_profile == 'standalone' || inputs.fuzz_profile == 'all') && needs['build-greptime-ci'].result == 'success' && needs['prepare-fuzz-targets'].result == 'success' && needs.scope.outputs.fuzz_unstable != 'false' }} name: Unstable Fuzz Test - needs: [build-greptime-ci, prepare-fuzz-targets] + needs: [build-greptime-ci, prepare-fuzz-targets, scope] runs-on: ubuntu-latest timeout-minutes: 60 steps: @@ -574,8 +658,9 @@ jobs: artifact-name: fuzz-unstable-standalone build-greptime-ci: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && ((github.event_name == 'pull_request' && github.event.pull_request.draft == false) || (github.event_name == 'workflow_dispatch' && inputs.ci_command && inputs.fuzz_profile != 'none')) }} + if: ${{ !cancelled() && github.repository == 'GreptimeTeam/greptimedb' && ((github.event_name == 'pull_request' && github.event.pull_request.draft == false) || (github.event_name == 'workflow_dispatch' && inputs.ci_command && inputs.fuzz_profile != 'none') || github.event_name == 'merge_group') && needs.scope.outputs.build_ci_needed != 'false' }} name: Build GreptimeDB binary (profile-CI) + needs: [scope] runs-on: ${{ matrix.os }} strategy: matrix: @@ -634,10 +719,10 @@ jobs: version: current distributed-fuzztest: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && ((github.event_name == 'pull_request' && github.event.pull_request.draft == false) || (github.event_name == 'workflow_dispatch' && inputs.ci_command)) && (github.event_name == 'pull_request' || inputs.fuzz_profile == 'distributed' || inputs.fuzz_profile == 'all') }} + if: ${{ !cancelled() && github.repository == 'GreptimeTeam/greptimedb' && ((github.event_name == 'pull_request' && github.event.pull_request.draft == false) || (github.event_name == 'workflow_dispatch' && inputs.ci_command) || github.event_name == 'merge_group') && (github.event_name == 'pull_request' || github.event_name == 'merge_group' || inputs.fuzz_profile == 'distributed' || inputs.fuzz_profile == 'all') && needs['build-greptime-ci'].result == 'success' && needs['prepare-fuzz-targets'].result == 'success' && needs.scope.outputs.fuzz_distributed != 'false' }} name: Fuzz Test (Distributed, ${{ matrix.mode.name }}, ${{ matrix.group.name }}) runs-on: ubuntu-latest - needs: [build-greptime-ci, prepare-fuzz-targets] + needs: [build-greptime-ci, prepare-fuzz-targets, scope] timeout-minutes: 60 strategy: fail-fast: false @@ -873,10 +958,10 @@ jobs: docker system prune -f distributed-fuzztest-with-chaos: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && ((github.event_name == 'pull_request' && github.event.pull_request.draft == false) || (github.event_name == 'workflow_dispatch' && inputs.ci_command)) && (github.event_name == 'pull_request' || inputs.fuzz_profile == 'chaos' || inputs.fuzz_profile == 'all') }} + if: ${{ !cancelled() && github.repository == 'GreptimeTeam/greptimedb' && ((github.event_name == 'pull_request' && github.event.pull_request.draft == false) || (github.event_name == 'workflow_dispatch' && inputs.ci_command) || github.event_name == 'merge_group') && (github.event_name == 'pull_request' || github.event_name == 'merge_group' || inputs.fuzz_profile == 'chaos' || inputs.fuzz_profile == 'all') && needs['build-greptime-ci'].result == 'success' && needs['prepare-fuzz-targets'].result == 'success' && needs.scope.outputs.fuzz_chaos != 'false' }} name: Fuzz Test with Chaos (Distributed, ${{ matrix.mode.name }}, ${{ matrix.group.name }}) runs-on: ubuntu-latest - needs: [build-greptime-ci, prepare-fuzz-targets] + needs: [build-greptime-ci, prepare-fuzz-targets, scope] timeout-minutes: 60 strategy: fail-fast: false diff --git a/Makefile b/Makefile index 21c87cc3707..073a95cb1fa 100644 --- a/Makefile +++ b/Makefile @@ -249,6 +249,10 @@ fmt-check: ## Check code format. check-enterprise-license: ## Check enterprise-gated files are listed in both license configs. python3 scripts/check-enterprise-license.py +.PHONY: check-ci-scope +check-ci-scope: ## Check the CI scope rules still match existing paths. + python3 .github/ci/ci_scope_rules.py --check-paths + .PHONY: start-etcd start-etcd: ## Start single node etcd for testing purpose. docker run --rm -d --network=host -p 2379-2380:2379-2380 ${ETCD_IMAGE} From f646f06430516aa14a9dd78236c4cae4c27f9857 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Wed, 28 Oct 2026 13:32:23 +0000 Subject: [PATCH 127/239] ci: enqueue eligible pull requests during the nightly window (#10040) * ci: enqueue eligible pull requests during the nightly window Signed-off-by: WenyXu * docs: hide nightly merge instructions in PR template Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- .github/nightly-merge.md | 43 ++++++ .github/pull_request_template.md | 11 ++ .github/scripts/tests/nightly-merge.test.cjs | 108 +++++++++++++++ .github/workflows/checks.yml | 2 +- .github/workflows/nightly-merge.yml | 138 +++++++++++++++++++ 5 files changed, 301 insertions(+), 1 deletion(-) create mode 100644 .github/nightly-merge.md create mode 100644 .github/scripts/tests/nightly-merge.test.cjs create mode 100644 .github/workflows/nightly-merge.yml diff --git a/.github/nightly-merge.md b/.github/nightly-merge.md new file mode 100644 index 00000000000..a8699a9b159 --- /dev/null +++ b/.github/nightly-merge.md @@ -0,0 +1,43 @@ +# Nightly merge + +The PR template opts new PRs into the `merge nightly` label. The Nightly Merge +workflow reconciles the label on PR creation, edits, and reopening. Existing PRs +without the checkbox are not opted in. Unchecking removes the label; it does not +request an immediate merge or cancel an existing queue entry/auto-merge request. + +Every day at 02:10, 03:10, 04:10, 05:10, 06:10, 07:10, and 08:10 +Asia/Shanghai (UTC+8), the workflow enqueues up to five eligible PRs targeting +`main` per run, oldest first. Already queued PRs do not count toward this limit. +GitHub may delay scheduled runs. +PRs must still have the checkbox and label, be approved and clean, and have all +review threads resolved. Drafts are skipped. +Already queued PRs and manually enabled auto-merge requests are left alone; +Maintainers can use **Merge when ready** to proceed earlier. Unready PRs wait for a later +scheduled run; this workflow never enables auto-merge on their behalf. + +## Repository setup + +- Create the `merge nightly` label before enabling the workflow. +- Require the merge queue, reviews, status checks, and conversation resolution + on `main`. Conversation resolution must also be enforced by GitHub to cover + new comments arriving between the script's check and enqueueing. +- Ensure the existing `GH_PERSONAL_ACCESS_TOKEN` automation secret can read PRs + and review threads and enqueue PRs in this repository. Do not grant it a + branch-rule bypass. This credential is already referenced by other workflows; + its actual permissions have not been validated by the local tests. +- Run **Nightly Merge** manually with the default `dry_run: true` first and + inspect its job summary. Setting it to false explicitly enqueues outside the + scheduled time as well. Merge queue build concurrency is configured separately + in GitHub; this workflow does not change it. + +The workflow never checks out or executes PR code. Label reconciliation uses +`GITHUB_TOKEN`; enqueueing uses the automation credential so downstream CI can +run. No runner remains allocated while waiting for the nightly schedule. + +## Local validation + +Run `actionlint .github/workflows/nightly-merge.yml .github/workflows/checks.yml` +and `node --test .github/scripts/tests/nightly-merge.test.cjs`. +The tests execute the actual inline workflow scripts with mocked GitHub APIs, +including one-factor comparisons against an eligible PR. They do not validate +live GitHub permissions, branch rules, or runner savings. diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index a33531c2733..5b26f64ceb3 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -25,3 +25,14 @@ Please convert it to a draft if some of the following conditions are not met. - [ ] API changes are backward compatible. - [ ] Schema or data changes are backward compatible. - [ ] This PR needs to be backported to release branches, and I have added the `backport-` labels (e.g. `backport-v1.3` targets `release/v1.3`). + +## Merge timing + +- [x] Merge during the nightly window + + diff --git a/.github/scripts/tests/nightly-merge.test.cjs b/.github/scripts/tests/nightly-merge.test.cjs new file mode 100644 index 00000000000..885876f5610 --- /dev/null +++ b/.github/scripts/tests/nightly-merge.test.cjs @@ -0,0 +1,108 @@ +// 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. + +// Run with: node --test .github/scripts/tests/nightly-merge.test.cjs +const { readFileSync } = require('node:fs'); +const { resolve } = require('node:path'); +const assert = require('node:assert/strict'); +const { test } = require('node:test'); +const AsyncFunction = Object.getPrototypeOf(async function () {}).constructor; +const workflow = readFileSync(resolve(__dirname, '../../workflows/nightly-merge.yml'), 'utf8'); +const scripts = [...workflow.matchAll(/ script: \|\n((?: .*\n|\n)+)/g)] + .map(match => new AsyncFunction('github', 'context', 'core', 'process', match[1].replace(/^ /gm, ''))); +assert.equal(scripts.length, 2); +const checked = '- [x] Merge during the nightly window '; + +async function run({ sync = false, body = checked, labels = ['merge nightly'], + detail = {}, current = {}, threads = [[true]], count = 1, dryRun = false, + failFirst = false } = {}) { + const calls = []; + let page = 0; + const github = { + rest: { + pulls: { get: async () => ({ data: { state: 'open', draft: false, base: { ref: 'main' }, + body, labels: labels.map(name => ({ name })), ...current } }) }, + issues: { + listForRepo: () => {}, + addLabels: async args => calls.push(['add', args.labels]), + removeLabel: async args => calls.push(['remove', args.name]) + } + }, + paginate: async () => Array.from({ length: count }, (_, i) => ({ number: i + 1, pull_request: {} })), + graphql: async (query, args) => { + if (query.includes('mutation')) { + calls.push(['enqueue', args]); + if (failFirst && calls.filter(call => call[0] === 'enqueue').length === 1) throw Error('head changed'); + return {}; + } + if (query.includes('reviewThreads')) { + const index = page++ % threads.length; + calls.push(['threads', args.cursor]); + return { repository: { pullRequest: { reviewThreads: { + nodes: threads[index].map(isResolved => ({ isResolved })), + pageInfo: { hasNextPage: index + 1 < threads.length, endCursor: `page-${index}` } + } } } }; + } + return { repository: { pullRequest: { id: `PR_${args.number}`, headRefOid: 'abc', + reviewDecision: 'APPROVED', mergeStateStatus: 'CLEAN', ...detail } } }; + } + }; + const summary = { addHeading() { return this; }, addTable(rows) { calls.push(['summary', rows]); return this; }, async write() {} }; + await scripts[sync ? 0 : 1](github, + { repo: { owner: 'GreptimeTeam', repo: 'greptimedb' }, payload: { pull_request: { number: 1 } } }, + { summary, warning() {}, setFailed: message => calls.push(['failed', message]) }, + { env: { DRY_RUN: String(dryRun) } }); + return calls; +} + +test('template opts in and checkbox edits reconcile only the nightly label', async () => { + assert.ok(readFileSync(resolve(__dirname, '../../pull_request_template.md'), 'utf8').includes(checked)); + assert.deepEqual(await run({ sync: true, labels: ['bug'] }), [['add', ['merge nightly']]]); + assert.deepEqual(await run({ sync: true }), []); + for (const body of [checked.replace('[x]', '[ ]'), '', null, `> ${checked}`]) { + assert.deepEqual(await run({ sync: true, body }), [['remove', 'merge nightly']]); + } + assert.deepEqual(await run({ sync: true, body: checked.replace('[x]', '[X]'), labels: [] }), [['add', ['merge nightly']]]); +}); + +test('eligible PR is enqueued with its checked head; dry run performs no mutation', async () => { + assert.deepEqual((await run()).filter(c => c[0] === 'enqueue'), [['enqueue', { id: 'PR_1', sha: 'abc' }]]); + assert.equal((await run({ dryRun: true })).some(c => c[0] === 'enqueue'), false); +}); + +test('each eligibility gate independently prevents enqueueing', async () => { + for (const options of [ + { body: checked.replace('[x]', '[ ]') }, { labels: [] }, + { current: { draft: true } }, { current: { state: 'closed' } }, { current: { base: { ref: 'release/v1.3' } } }, + { detail: { reviewDecision: 'REVIEW_REQUIRED' } }, { detail: { reviewDecision: 'CHANGES_REQUESTED' } }, + { detail: { mergeStateStatus: 'BLOCKED' } }, { detail: { mergeStateStatus: 'UNKNOWN' } }, + { detail: { mergeQueueEntry: { id: 'entry' } } }, { detail: { autoMergeRequest: { enabledAt: 'now' } } }, + { threads: [[false]] }, { threads: [[true], [false]] } + ]) { + assert.equal((await run(options)).some(c => c[0] === 'enqueue'), false, JSON.stringify(options)); + } +}); + +test('all thread pages are checked before enqueueing', async () => { + const calls = await run({ threads: [[true], [true]] }); + assert.deepEqual(calls.filter(c => c[0] === 'threads'), [['threads', null], ['threads', 'page-0']]); + assert.equal(calls.filter(c => c[0] === 'enqueue').length, 1); +}); + +test('per-run cap and individual failures do not prevent later candidates', async () => { + assert.equal((await run({ count: 8 })).filter(c => c[0] === 'enqueue').length, 5); + const calls = await run({ count: 2, failFirst: true }); + assert.equal(calls.filter(c => c[0] === 'enqueue').length, 2); + assert.equal(calls.filter(c => c[0] === 'failed').length, 1); +}); diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml index b3ff212ade4..941f36d9726 100644 --- a/.github/workflows/checks.yml +++ b/.github/workflows/checks.yml @@ -75,7 +75,7 @@ jobs: with: node-version: 24 - name: Run Node GitHub script tests - run: node --test .github/scripts/query-regression-comment.test.cjs + run: node --test .github/scripts/query-regression-comment.test.cjs .github/scripts/tests/nightly-merge.test.cjs - name: Setup uv uses: astral-sh/setup-uv@v3 with: diff --git a/.github/workflows/nightly-merge.yml b/.github/workflows/nightly-merge.yml new file mode 100644 index 00000000000..ff6d5874ff1 --- /dev/null +++ b/.github/workflows/nightly-merge.yml @@ -0,0 +1,138 @@ +# Setup and behavior: .github/nightly-merge.md +name: Nightly Merge + +on: + pull_request_target: + types: [opened, edited, reopened] + schedule: + # Hourly from 02:10 through 08:10 Asia/Shanghai (UTC+8), inclusive. + # GitHub may delay scheduled runs. + - cron: '10 0,18-23 * * *' + workflow_dispatch: + inputs: + dry_run: + description: List eligible PRs without adding them to the merge queue + type: boolean + default: true + +permissions: {} + +jobs: + sync-label: + if: github.repository == 'GreptimeTeam/greptimedb' && github.event_name == 'pull_request_target' + runs-on: ubuntu-latest + timeout-minutes: 5 + permissions: + pull-requests: read + issues: write + concurrency: + group: nightly-merge-label-${{ github.event.pull_request.number }} + cancel-in-progress: false + steps: + # No checkout: PR contents are untrusted data, never executable code. + - uses: actions/github-script@v7 + with: + script: | + const { owner, repo } = context.repo; + const issue_number = context.payload.pull_request.number; + const { data: pr } = await github.rest.pulls.get({ owner, repo, pull_number: issue_number }); + const selected = /^- \[[xX]\] Merge during the nightly window \s*$/m.test(pr.body || ''); + const label = 'merge nightly'; + const labeled = pr.labels.some(item => item.name === label); + if (selected && !labeled) { + await github.rest.issues.addLabels({ owner, repo, issue_number, labels: [label] }); + } else if (!selected && labeled) { + await github.rest.issues.removeLabel({ owner, repo, issue_number, name: label }); + } + + enqueue: + if: github.repository == 'GreptimeTeam/greptimedb' && github.event_name != 'pull_request_target' + runs-on: ubuntu-latest + timeout-minutes: 10 + concurrency: + group: nightly-merge-enqueue + cancel-in-progress: false + steps: + # Reuses the automation credential used by backport and compat updates. + # A PAT (or App token) lets subsequent merge_group workflows run. + - uses: actions/github-script@v7 + env: + DRY_RUN: ${{ github.event_name == 'workflow_dispatch' && inputs.dry_run }} + with: + github-token: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + script: | + const { owner, repo } = context.repo; + const dryRun = process.env.DRY_RUN === 'true'; + const candidates = await github.paginate(github.rest.issues.listForRepo, { + owner, repo, state: 'open', labels: 'merge nightly', + sort: 'created', direction: 'asc', per_page: 100 + }); + const rows = []; + let enqueued = 0; + for (const candidate of candidates) { + if (!candidate.pull_request || enqueued >= 5) continue; + const number = candidate.number; + try { + // Fetch current body/labels, not the earlier list snapshot. + const { data: current } = await github.rest.pulls.get({ owner, repo, pull_number: number }); + if (current.state !== 'open' || current.draft || current.base.ref !== 'main' || + !current.labels.some(label => label.name === 'merge nightly') || + !/^- \[[xX]\] Merge during the nightly window \s*$/m.test(current.body || '')) { + rows.push([String(number), 'Skipped: not opted in, draft, closed, or not targeting main']); + continue; + } + const { repository: { pullRequest: pr } } = await github.graphql(` + query($owner: String!, $repo: String!, $number: Int!) { + repository(owner: $owner, name: $repo) { + pullRequest(number: $number) { + id headRefOid reviewDecision mergeStateStatus + mergeQueueEntry { id } + autoMergeRequest { enabledAt } + } + } + }`, { owner, repo, number }); + if (pr.mergeQueueEntry || pr.autoMergeRequest || + pr.reviewDecision !== 'APPROVED' || pr.mergeStateStatus !== 'CLEAN') { + rows.push([String(number), 'Skipped: already queued/auto-merge enabled, or not approved and clean']); + continue; + } + let cursor = null; + let unresolved = false; + do { + const { repository: { pullRequest: { reviewThreads } } } = await github.graphql(` + query($owner: String!, $repo: String!, $number: Int!, $cursor: String) { + repository(owner: $owner, name: $repo) { + pullRequest(number: $number) { + reviewThreads(first: 100, after: $cursor) { + nodes { isResolved } + pageInfo { hasNextPage endCursor } + } + } + } + }`, { owner, repo, number, cursor }); + unresolved = reviewThreads.nodes.some(thread => !thread.isResolved); + cursor = reviewThreads.pageInfo.hasNextPage ? reviewThreads.pageInfo.endCursor : null; + } while (cursor && !unresolved); + if (unresolved) { + rows.push([String(number), 'Skipped: unresolved review threads']); + continue; + } + if (!dryRun) { + // GitHub enforces branch rules; never bypass or enable auto-merge. + await github.graphql(` + mutation($id: ID!, $sha: GitObjectID!) { + enqueuePullRequest(input: {pullRequestId: $id, expectedHeadOid: $sha}) { + mergeQueueEntry { id } + } + }`, { id: pr.id, sha: pr.headRefOid }); + } + enqueued++; + rows.push([String(number), dryRun ? 'Would enqueue' : 'Enqueued']); + } catch (error) { + core.warning(`PR #${number}: ${error.message}`); + rows.push([String(number), 'Failed: see workflow logs']); + core.setFailed('Some PRs could not be processed; see workflow logs.'); + } + } + await core.summary.addHeading(dryRun ? 'Nightly merge (dry run)' : 'Nightly merge') + .addTable([['PR', 'Result'], ...rows]).write(); From 8e16710425da8fde4be4b87614864e8367005348 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Wed, 28 Oct 2026 13:33:39 +0000 Subject: [PATCH 128/239] fix(log-store): accept a retried first create of the conditional create probe (#10034) * fix(log-store): accept a retried first create of the conditional create probe The startup probe of the object store WAL wrote its first create by hand and failed on any error. When that create is stored but its response is lost, the retry layer repeats it, the repeat meets the object, and the probe reported a failed write, so the WAL refused to open. Move the create of put_if_absent, which reads the object back and accepts identical content, into a path-based core and use it for the probe's first create too. The second create, which the store must refuse, is unchanged. Signed-off-by: jeremyhi * test(log-store): reuse LosingWriter for the retried probe create test Move LosingWriter from the fault matrix into the test module and use it, armed with the lose_next pattern, in place of RetriedWriter. Say why the probe's second create bypasses put_path_if_absent: a read-back would find the first create's object and let an overwriting store pass. Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- src/log-store/src/object_store_wal/io.rs | 35 ++++++----- src/log-store/src/object_store_wal/store.rs | 69 +++++++++++++++------ 2 files changed, 70 insertions(+), 34 deletions(-) diff --git a/src/log-store/src/object_store_wal/io.rs b/src/log-store/src/object_store_wal/io.rs index 0c087141c52..c14a078f833 100644 --- a/src/log-store/src/object_store_wal/io.rs +++ b/src/log-store/src/object_store_wal/io.rs @@ -96,16 +96,15 @@ impl ObjectStoreIo { /// call, twice, and then deletes it. pub(crate) async fn probe_conditional_create(&self) -> Result<()> { let path = format!("{}{}", self.probe_prefix, Uuid::new_v4()); - let create = || { - self.store - .write_with(&path, PROBE_CONTENT) - .if_not_exists(true) - }; - create().await.context(WalObjectStoreSnafu { - operation: "write", - path: path.clone(), - })?; - let recreated = create().await; + self.put_path_if_absent(path.clone(), Bytes::from_static(PROBE_CONTENT)) + .await?; + // Sent directly: only the store's own refusal shows that it honours + // the condition, and a read-back would find the first create's object. + let recreated = self + .store + .write_with(&path, PROBE_CONTENT) + .if_not_exists(true) + .await; if let Err(error) = self.store.delete(&path).await { warn!(error; "Failed to delete the conditional create probe {path} of the object store WAL"); } @@ -134,15 +133,23 @@ impl ObjectStoreIo { } } - /// Creates the object `object_seq` unless it already exists. A retry with - /// the same content is accepted, different content is rejected. + /// Creates the object `object_seq` unless it already exists, see + /// [`Self::put_path_if_absent`]. + pub(crate) async fn put_if_absent(&self, object_seq: u64, content: Bytes) -> Result { + self.put_path_if_absent(self.object_path(object_seq), content) + .await + } + + /// Creates `path` unless it already exists. A retry with the same content + /// is accepted, different content is rejected: a create the store stored + /// but whose response was lost can be retried, by the caller or by the + /// retry layer, and then finds its own object. /// /// When the store reports that the object exists but the read that /// compares its content fails, the read failure is returned, so that its /// retry hint says whether comparing again may settle the outcome. Any /// other create failure is returned as the write failure it is. - pub(crate) async fn put_if_absent(&self, object_seq: u64, content: Bytes) -> Result { - let path = self.object_path(object_seq); + async fn put_path_if_absent(&self, path: String, content: Bytes) -> Result { let write_result = self .store .write_with(&path, content.clone()) diff --git a/src/log-store/src/object_store_wal/store.rs b/src/log-store/src/object_store_wal/store.rs index a741c337d88..354521f38d2 100644 --- a/src/log-store/src/object_store_wal/store.rs +++ b/src/log-store/src/object_store_wal/store.rs @@ -4038,6 +4038,25 @@ mod tests { } } + /// A writer that writes the object and then reports a failure, as a + /// create whose response is lost. + struct LosingWriter(oio::Writer); + + impl oio::Write for LosingWriter { + async fn write(&mut self, buffer: mock::Buffer) -> mock::Result<()> { + self.0.write(buffer).await + } + + async fn close(&mut self) -> mock::Result { + self.0.close().await?; + Err(mock::Error::new(mock::ErrorKind::Unexpected, "lost response").set_temporary()) + } + + async fn abort(&mut self) -> mock::Result<()> { + self.0.abort().await + } + } + async fn keys(object_store: &ObjectStore) -> Vec { let mut keys = object_store .list_with(&format!("{PREFIX}/")) @@ -4096,6 +4115,35 @@ mod tests { assert!(keys(&object_store).await.is_empty()); } + #[tokio::test] + async fn test_store_open_accepts_a_retried_first_probe_create() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + let lose_next = Arc::new(AtomicBool::new(true)); + let losing = lose_next.clone(); + let retrying = object_store.clone().layer( + MockLayerBuilder::default() + .writer_factory(Arc::new(move |_, _, inner| { + if losing.swap(false, Ordering::SeqCst) { + Box::new(LosingWriter(inner)) + } else { + inner + } + })) + .build() + .unwrap(), + ); + + // The first write of the open is the first create of the probe. + let store = open(retrying, &eager()).await; + assert!(!lose_next.load(Ordering::SeqCst)); + store.stop().await.unwrap(); + assert_eq!( + vec![object_path(&object_store, 0)], + keys(&object_store).await + ); + } + #[tokio::test] async fn test_store_recovery_ignores_a_probe_left_behind() { let _serialized = SERIALIZED.lock().await; @@ -10500,7 +10548,7 @@ mod tests { /// object store access through the store, one case per fault, and checks /// that every row of the *Failure matrix* has one. mod fault_matrix { - use object_store::layers::mock::{self, MockLayerBuilder, oio}; + use object_store::layers::mock::MockLayerBuilder; use prometheus::Histogram; use super::*; @@ -11459,25 +11507,6 @@ mod tests { assert!(!lose_next.load(Ordering::SeqCst)); } - /// A writer that writes the object and then reports a failure, as a - /// create whose response is lost. - struct LosingWriter(oio::Writer); - - impl oio::Write for LosingWriter { - async fn write(&mut self, buffer: mock::Buffer) -> mock::Result<()> { - self.0.write(buffer).await - } - - async fn close(&mut self) -> mock::Result { - self.0.close().await?; - Err(mock::Error::new(mock::ErrorKind::Unexpected, "lost response").set_temporary()) - } - - async fn abort(&mut self) -> mock::Result<()> { - self.0.abort().await - } - } - /// Object 2 is created and object 3 in flight when object 1 fails. async fn create_fails_transiently_durable() { let (store, io, mut parked) = open_parking_creates(memory_store(), &eager()).await; From 4c37ed4a0e44a1cb968ba9dc1e3a88c7c97cfa2b Mon Sep 17 00:00:00 2001 From: Yvan Wang <131545713+BootstrapperSBL@users.noreply.github.com> Date: Wed, 28 Oct 2026 17:09:21 +0000 Subject: [PATCH 129/239] feat(pipeline): add logfmt processor (#9586) Add a built-in pipeline processor that parses logfmt lines (https://brandur.org/logfmt) into structured objects, with support for quoted values, escapes, bare-key flags, and ignore_missing. Closes #5049 Signed-off-by: Yvan Wang <131545713+BootstrapperSBL@users.noreply.github.com> --- src/pipeline/src/error.rs | 8 + src/pipeline/src/etl/processor.rs | 4 + src/pipeline/src/etl/processor/logfmt.rs | 328 +++++++++++++++++++++++ src/pipeline/tests/logfmt.rs | 122 +++++++++ 4 files changed, 462 insertions(+) create mode 100644 src/pipeline/src/etl/processor/logfmt.rs create mode 100644 src/pipeline/tests/logfmt.rs diff --git a/src/pipeline/src/error.rs b/src/pipeline/src/error.rs index 0da65d25436..534dcc14512 100644 --- a/src/pipeline/src/error.rs +++ b/src/pipeline/src/error.rs @@ -619,6 +619,13 @@ pub enum Error { #[snafu(implicit)] location: Location, }, + #[snafu(display("Failed to parse logfmt input: {input}, reason: {reason}"))] + LogfmtParse { + input: String, + reason: String, + #[snafu(implicit)] + location: Location, + }, #[snafu(display("Failed to parse json"))] JsonParse { #[snafu(source)] @@ -1010,6 +1017,7 @@ impl ErrorExt for Error { | UnsupportedIndexType { .. } | IdentifyPipelineColumnTypeMismatch { .. } | JsonParse { .. } + | LogfmtParse { .. } | JsonPathParse { .. } | JsonPathParseResultIndex { .. } | FieldRequiredForDispatcher diff --git a/src/pipeline/src/etl/processor.rs b/src/pipeline/src/etl/processor.rs index 8c999cd13ee..911bd5b9136 100644 --- a/src/pipeline/src/etl/processor.rs +++ b/src/pipeline/src/etl/processor.rs @@ -25,6 +25,7 @@ pub mod join; pub mod json_parse; pub mod json_path; pub mod letter; +pub mod logfmt; pub mod regex; pub mod select; pub mod simple_extract; @@ -45,6 +46,7 @@ use gsub::GsubProcessor; use join::JoinProcessor; use json_path::JsonPathProcessor; use letter::LetterProcessor; +use logfmt::LogfmtProcessor; use regex::RegexProcessor; use snafu::{OptionExt, ResultExt}; use urlencoding::UrlEncodingProcessor; @@ -143,6 +145,7 @@ pub enum ProcessorKind { Date(DateProcessor), JsonPath(JsonPathProcessor), JsonParse(JsonParseProcessor), + Logfmt(LogfmtProcessor), SimpleJsonPath(SimpleExtractProcessor), Decolorize(DecolorizeProcessor), Digest(DigestProcessor), @@ -226,6 +229,7 @@ fn parse_processor(doc: &yaml_rust::Yaml, allow_vrl_host_access: bool) -> Result json_parse::PROCESSOR_JSON_PARSE => { ProcessorKind::JsonParse(JsonParseProcessor::try_from(value)?) } + logfmt::PROCESSOR_LOGFMT => ProcessorKind::Logfmt(LogfmtProcessor::try_from(value)?), vrl_processor::PROCESSOR_VRL => { ProcessorKind::Vrl(VrlProcessor::from_yaml(value, allow_vrl_host_access)?) } diff --git a/src/pipeline/src/etl/processor/logfmt.rs b/src/pipeline/src/etl/processor/logfmt.rs new file mode 100644 index 00000000000..efe804fe869 --- /dev/null +++ b/src/pipeline/src/etl/processor/logfmt.rs @@ -0,0 +1,328 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! The `logfmt` processor parses a [logfmt](https://brandur.org/logfmt) line, +//! such as `level=info msg="hello world" duration=12ms`, into an object. + +use std::collections::BTreeMap; + +use snafu::OptionExt as _; +use vrl::prelude::Bytes; +use vrl::value::{KeyString, Value as VrlValue}; + +use crate::Processor; +use crate::error::{ + Error, FieldMustBeTypeSnafu, KeyMustBeStringSnafu, LogfmtParseSnafu, + ProcessorMissingFieldSnafu, Result, ValueMustBeMapSnafu, +}; +use crate::etl::field::Fields; +use crate::etl::processor::{ + FIELD_NAME, FIELDS_NAME, IGNORE_MISSING_NAME, yaml_bool, yaml_new_field, yaml_new_fields, +}; + +pub(crate) const PROCESSOR_LOGFMT: &str = "logfmt"; + +/// Parses logfmt strings into objects. +/// +/// Every `key=value` pair becomes an entry of the output object with a string +/// value. A bare key without `=` (e.g. `debug` in `debug msg=hi`) is treated +/// as a flag and set to `true`. When a key appears more than once, the last +/// value wins. The output is written to the target field, or replaces the +/// input field in place when no target is given. +#[derive(Debug, Default)] +pub struct LogfmtProcessor { + fields: Fields, + ignore_missing: bool, +} + +impl TryFrom<&yaml_rust::yaml::Hash> for LogfmtProcessor { + type Error = Error; + + fn try_from(value: &yaml_rust::yaml::Hash) -> std::result::Result { + let mut fields = Fields::default(); + let mut ignore_missing = false; + + for (k, v) in value.iter() { + let key = k + .as_str() + .with_context(|| KeyMustBeStringSnafu { k: k.clone() })?; + match key { + FIELD_NAME => { + fields = Fields::one(yaml_new_field(v, FIELD_NAME)?); + } + FIELDS_NAME => { + fields = yaml_new_fields(v, FIELDS_NAME)?; + } + IGNORE_MISSING_NAME => { + ignore_missing = yaml_bool(v, IGNORE_MISSING_NAME)?; + } + _ => {} + } + } + + Ok(LogfmtProcessor { + fields, + ignore_missing, + }) + } +} + +impl LogfmtProcessor { + fn process_field(&self, val: &VrlValue) -> Result { + let Some(line) = val.as_str() else { + return FieldMustBeTypeSnafu { + field: val.to_string(), + ty: "string", + } + .fail(); + }; + let pairs = parse_logfmt(&line).map_err(|reason| { + LogfmtParseSnafu { + input: line.to_string(), + reason, + } + .build() + })?; + Ok(VrlValue::Object(pairs)) + } +} + +impl Processor for LogfmtProcessor { + fn kind(&self) -> &str { + PROCESSOR_LOGFMT + } + + fn ignore_missing(&self) -> bool { + self.ignore_missing + } + + fn exec_mut(&self, mut val: VrlValue) -> Result { + for field in self.fields.iter() { + let index = field.input_field(); + let val = val.as_object_mut().context(ValueMustBeMapSnafu)?; + match val.get(index) { + Some(VrlValue::Null) | None => { + if !self.ignore_missing { + return ProcessorMissingFieldSnafu { + processor: self.kind(), + field: field.input_field(), + } + .fail(); + } + } + Some(v) => { + let processed = self.process_field(v)?; + let output_index = field.target_or_input_field(); + val.insert(KeyString::from(output_index.to_string()), processed); + } + } + } + Ok(val) + } +} + +/// Parses a logfmt line into key-value pairs. +/// +/// Follows the grammar used by common logfmt implementations (e.g. go-logfmt): +/// - pairs are separated by whitespace; +/// - a key is a run of printable characters other than `=`, `"` and whitespace; +/// - a value is either unquoted (ends at the next whitespace) or double-quoted, +/// where `\"`, `\\`, `\n`, `\r`, `\t`, `\/`, `\b`, `\f` and `\uXXXX` escapes are supported; +/// - `key=` yields an empty string and a bare `key` yields `true`. +fn parse_logfmt(line: &str) -> std::result::Result, String> { + let mut output = BTreeMap::new(); + let mut chars = line.char_indices().peekable(); + + loop { + // Skip separators. + while chars.next_if(|(_, c)| c.is_whitespace()).is_some() {} + let Some(&(key_start, first)) = chars.peek() else { + break; + }; + if first == '=' || first == '"' { + return Err(format!("unexpected '{first}' at position {key_start}")); + } + + let mut key_end = line.len(); + while let Some(&(i, c)) = chars.peek() { + if c == '=' || c.is_whitespace() { + key_end = i; + break; + } + if c == '"' || c.is_control() { + return Err(format!("invalid character {c:?} in key at position {i}")); + } + chars.next(); + } + let key = KeyString::from(&line[key_start..key_end]); + + if chars.next_if(|(_, c)| *c == '=').is_none() { + // Bare key, e.g. `debug`. + output.insert(key, VrlValue::Boolean(true)); + continue; + } + + let value = match chars.peek() { + Some(&(_, '"')) => { + chars.next(); + parse_quoted_value(&mut chars)? + } + Some(&(value_start, _)) => { + let mut value_end = line.len(); + while let Some(&(i, c)) = chars.peek() { + if c.is_whitespace() { + value_end = i; + break; + } + if c == '"' || c == '=' { + return Err(format!( + "unexpected '{c}' in unquoted value at position {i}" + )); + } + chars.next(); + } + line[value_start..value_end].to_string() + } + None => String::new(), + }; + output.insert(key, VrlValue::Bytes(Bytes::from(value))); + } + + Ok(output) +} + +/// Parses the rest of a double-quoted value; the opening quote is already consumed. +fn parse_quoted_value( + chars: &mut std::iter::Peekable>, +) -> std::result::Result { + let mut value = String::new(); + loop { + let Some((i, c)) = chars.next() else { + return Err("unterminated quoted value".to_string()); + }; + match c { + '"' => { + if let Some(&(j, next)) = chars.peek() + && !next.is_whitespace() + { + return Err(format!( + "expect whitespace after quoted value at position {j}, found '{next}'" + )); + } + return Ok(value); + } + '\\' => { + let Some((_, escaped)) = chars.next() else { + return Err("unterminated escape sequence".to_string()); + }; + match escaped { + '"' => value.push('"'), + '\\' => value.push('\\'), + '/' => value.push('/'), + 'n' => value.push('\n'), + 'r' => value.push('\r'), + 't' => value.push('\t'), + 'b' => value.push('\u{8}'), + 'f' => value.push('\u{c}'), + 'u' => { + let hex: String = chars.by_ref().take(4).map(|(_, c)| c).collect(); + let decoded = (hex.len() == 4) + .then(|| u32::from_str_radix(&hex, 16).ok()) + .flatten() + .and_then(char::from_u32) + .ok_or_else(|| format!("invalid unicode escape '\\u{hex}'"))?; + value.push(decoded); + } + other => { + return Err(format!("invalid escape '\\{other}' at position {i}")); + } + } + } + c => value.push(c), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn obj(pairs: &[(&str, VrlValue)]) -> BTreeMap { + pairs + .iter() + .map(|(k, v)| (KeyString::from(*k), v.clone())) + .collect() + } + + fn s(v: &str) -> VrlValue { + VrlValue::Bytes(Bytes::from(v.to_string())) + } + + #[test] + fn test_parse_logfmt() { + let parsed = parse_logfmt( + r#"ts=2026-10-14T10:00:00Z level=info msg="hello \"world\"" path=/api/v1 empty= debug"#, + ) + .unwrap(); + assert_eq!( + parsed, + obj(&[ + ("ts", s("2026-10-14T10:00:00Z")), + ("level", s("info")), + ("msg", s(r#"hello "world""#)), + ("path", s("/api/v1")), + ("empty", s("")), + ("debug", VrlValue::Boolean(true)), + ]) + ); + } + + #[test] + fn test_parse_logfmt_whitespace_escapes_and_duplicates() { + let parsed = parse_logfmt(" a=1\tb=\"x\\ny\\u00e9\" a=2 ").unwrap(); + assert_eq!(parsed, obj(&[("a", s("2")), ("b", s("x\nyé"))])); + assert!(parse_logfmt("").unwrap().is_empty()); + assert!(parse_logfmt(" ").unwrap().is_empty()); + assert_eq!(parse_logfmt("k=\"\"").unwrap(), obj(&[("k", s(""))])); + assert_eq!(parse_logfmt("k=").unwrap(), obj(&[("k", s(""))])); + } + + #[test] + fn test_parse_logfmt_invalid() { + for input in [ + "=value", + "\"key\"=value", + "msg=\"unterminated", + "msg=\"bad\\q\"", + "msg=\"a\"b", + "k=a\"b", + "k=a=b", + "msg=\"\\u12\"", + ] { + assert!(parse_logfmt(input).is_err(), "expect error for {input:?}"); + } + } + + #[test] + fn test_logfmt_processor() { + let processor = LogfmtProcessor::default(); + let result = processor.process_field(&s("level=warn code=42")).unwrap(); + assert_eq!( + result, + VrlValue::Object(obj(&[("code", s("42")), ("level", s("warn"))])) + ); + assert!(processor.process_field(&VrlValue::Integer(1)).is_err()); + assert!(processor.process_field(&s("msg=\"oops")).is_err()); + } +} diff --git a/src/pipeline/tests/logfmt.rs b/src/pipeline/tests/logfmt.rs new file mode 100644 index 00000000000..891721c8402 --- /dev/null +++ b/src/pipeline/tests/logfmt.rs @@ -0,0 +1,122 @@ +// 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. + +mod common; + +use common_query::prelude::greptime_timestamp; +use greptime_proto::v1::value::ValueData; +use greptime_proto::v1::{ColumnDataType, SemanticType}; + +#[test] +fn test_logfmt() { + let input_value_str = r#" +[ + { + "message": "level=info msg=\"user logged in\" user=alice duration=12ms cached" + } +] +"#; + + let pipeline_yaml = r#" +--- +processors: + - logfmt: + field: message, parsed + - simple_extract: + field: parsed, level + key: "level" + - simple_extract: + field: parsed, msg + key: "msg" + - simple_extract: + field: parsed, user + key: "user" + - simple_extract: + field: parsed, cached + key: "cached" + +transform: + - fields: + - level + - msg + - user + type: string + - field: cached + type: boolean +"#; + + let output = common::parse_and_exec(input_value_str, pipeline_yaml); + + let expected_schema = vec![ + common::make_column_schema( + "level".to_string(), + ColumnDataType::String, + SemanticType::Field, + ), + common::make_column_schema( + "msg".to_string(), + ColumnDataType::String, + SemanticType::Field, + ), + common::make_column_schema( + "user".to_string(), + ColumnDataType::String, + SemanticType::Field, + ), + common::make_column_schema( + "cached".to_string(), + ColumnDataType::Boolean, + SemanticType::Field, + ), + common::make_column_schema( + greptime_timestamp().to_string(), + ColumnDataType::TimestampNanosecond, + SemanticType::Timestamp, + ), + ]; + assert_eq!(output.schema, expected_schema); + + let values = &output.rows[0].values; + assert_eq!( + values[0].value_data, + Some(ValueData::StringValue("info".to_string())) + ); + assert_eq!( + values[1].value_data, + Some(ValueData::StringValue("user logged in".to_string())) + ); + assert_eq!( + values[2].value_data, + Some(ValueData::StringValue("alice".to_string())) + ); + assert_eq!(values[3].value_data, Some(ValueData::BoolValue(true))); +} + +#[test] +fn test_logfmt_ignore_missing() { + let pipeline_yaml = r#" +--- +processors: + - logfmt: + field: message + ignore_missing: true + +transform: + - field: message + type: string +"#; + + let output = common::parse_and_exec(r#"[{"other": "x"}]"#, pipeline_yaml); + assert_eq!(output.rows[0].values[0].value_data, None); +} From b83ceb5e7bf260dd9cbd1e11b63ac8187ba0ff0a Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Wed, 28 Oct 2026 17:09:23 +0000 Subject: [PATCH 130/239] feat: add reusable paged byte cache (#9791) * feat: add reusable paged byte cache Signed-off-by: WenyXu * fix(ci): include common cache in scope rules Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- .github/ci/scope-rules.toml | 1 + Cargo.lock | 3 + src/common/cache/Cargo.toml | 5 +- src/common/cache/src/lib.rs | 2 + src/common/cache/src/page.rs | 156 +++++++++ src/common/cache/src/paged.rs | 384 ++++++++++++++++++++++ src/common/cache/src/paged/tests.rs | 492 ++++++++++++++++++++++++++++ 7 files changed, 1042 insertions(+), 1 deletion(-) create mode 100644 src/common/cache/src/page.rs create mode 100644 src/common/cache/src/paged.rs create mode 100644 src/common/cache/src/paged/tests.rs diff --git a/.github/ci/scope-rules.toml b/.github/ci/scope-rules.toml index ba2b92deda2..0c9920d5711 100644 --- a/.github/ci/scope-rules.toml +++ b/.github/ci/scope-rules.toml @@ -103,6 +103,7 @@ server = [ "src/datanode/**", "src/common/base/**", "src/common/batcher/**", + "src/common/cache/**", "src/common/catalog/**", "src/common/config/**", "src/common/datasource/**", diff --git a/Cargo.lock b/Cargo.lock index 8e4183994d4..f68dfbe2334 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2344,6 +2344,9 @@ dependencies = [ name = "common-cache" version = "1.4.0-alpha.0" dependencies = [ + "async-trait", + "bytes", + "common-base", "futures", "moka", "prometheus 0.14.0", diff --git a/src/common/cache/Cargo.toml b/src/common/cache/Cargo.toml index 1e2f9d83a5a..afa9a8f3623 100644 --- a/src/common/cache/Cargo.toml +++ b/src/common/cache/Cargo.toml @@ -8,9 +8,12 @@ license.workspace = true workspace = true [dependencies] +bytes.workspace = true +common-base.workspace = true +futures.workspace = true moka = { workspace = true, features = ["future"] } prometheus.workspace = true [dev-dependencies] -futures.workspace = true +async-trait.workspace = true tokio = { workspace = true, features = ["macros", "rt", "sync"] } diff --git a/src/common/cache/src/lib.rs b/src/common/cache/src/lib.rs index 58d54adee68..228f0ca3cf9 100644 --- a/src/common/cache/src/lib.rs +++ b/src/common/cache/src/lib.rs @@ -17,5 +17,7 @@ mod cache; pub mod metrics; +pub mod page; +pub mod paged; pub use cache::{Cache, RemovalCause}; diff --git a/src/common/cache/src/page.rs b/src/common/cache/src/page.rs new file mode 100644 index 00000000000..19568debc5e --- /dev/null +++ b/src/common/cache/src/page.rs @@ -0,0 +1,156 @@ +// 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. + +//! Fixed-size byte page layout for immutable cached content. + +use std::io; +use std::num::NonZeroU64; +use std::ops::Range; + +/// Fixed-size logical byte pages within one immutable file. +#[derive(Debug, Clone, Copy)] +pub struct PageLayout { + page_size: NonZeroU64, + file_size: u64, +} + +impl PageLayout { + /// Creates a layout. The final page is clipped to the file size. + pub fn new(page_size: NonZeroU64, file_size: u64) -> Self { + Self { + page_size, + file_size, + } + } + + /// Returns the half-open page-ID range intersecting a byte request. + /// + /// Requests are clipped at EOF. Empty or wholly out-of-file requests return + /// an empty range; reversed ranges are rejected. This does not allocate one + /// entry per page, even for large files. + pub fn page_ids(&self, request: Range) -> io::Result> { + validate_range(&request)?; + let start = request.start.min(self.file_size); + let end = request.end.min(self.file_size); + let first = start / self.page_size.get(); + if start == end { + return Ok(first..first); + } + // end > 0 here. Unlike (end + page_size - 1) / page_size, + // this expression cannot overflow at the end of a large file. + let last = (end - 1) / self.page_size.get(); + Ok(first..last + 1) + } + + /// Returns a page's file byte range, or `None` for a page outside the file. + pub fn page_range(&self, page_id: u64) -> Option> { + let start = page_id.checked_mul(self.page_size.get())?; + if start >= self.file_size { + return None; + } + let len = self.page_size.get().min(self.file_size - start); + Some(start..start + len) + } + + /// Returns the requested slice relative to a page's buffer. + /// + /// Returns `None` if the page and request do not intersect. The caller can + /// borrow this slice for an immediate response; caching a slice of a much + /// larger I/O buffer may retain that entire allocation and needs a copy. + pub fn request_slice( + &self, + page_id: u64, + request: Range, + ) -> io::Result>> { + validate_range(&request)?; + let Some(page) = self.page_range(page_id) else { + return Ok(None); + }; + let start = request.start.max(page.start); + let end = request.end.min(page.end); + if start >= end { + return Ok(None); + } + Ok(Some( + to_usize(start - page.start)?..to_usize(end - page.start)?, + )) + } +} + +fn validate_range(range: &Range) -> io::Result<()> { + if range.start > range.end { + return Err(invalid_input("Byte range start exceeds end")); + } + Ok(()) +} + +fn to_usize(value: u64) -> io::Result { + usize::try_from(value).map_err(|_| invalid_input("Byte offset exceeds addressable memory")) +} + +fn invalid_input(message: &'static str) -> io::Error { + io::Error::new(io::ErrorKind::InvalidInput, message) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn nonzero(value: u64) -> NonZeroU64 { + NonZeroU64::new(value).unwrap() + } + + #[test] + fn test_page_alignment_and_eof() { + let layout = PageLayout::new(nonzero(64), 150); + for (request, expected) in [ + (0..0, 0..0), + (0..64, 0..1), + (32..96, 0..2), + (64..128, 1..2), + (128..200, 2..3), + (150..200, 2..2), + ] { + assert_eq!(layout.page_ids(request).unwrap(), expected); + } + assert_eq!(layout.page_range(2), Some(128..150)); + assert_eq!(layout.page_range(3), None); + assert_eq!(layout.request_slice(0, 32..96).unwrap(), Some(32..64)); + assert_eq!(layout.request_slice(1, 32..96).unwrap(), Some(0..32)); + assert_eq!(layout.request_slice(2, 128..200).unwrap(), Some(0..22)); + assert_eq!(layout.request_slice(1, 0..64).unwrap(), None); + assert_eq!(PageLayout::new(nonzero(64), 0).page_range(0), None); + assert!(layout.page_ids(Range { start: 2, end: 1 }).is_err()); + assert!(layout.request_slice(0, Range { start: 2, end: 1 }).is_err()); + } + + #[test] + fn test_large_file_boundaries() { + let layout = PageLayout::new(nonzero(1), u64::MAX); + assert_eq!(layout.page_ids(0..u64::MAX).unwrap(), 0..u64::MAX); + assert_eq!( + layout.page_range(u64::MAX - 1), + Some(u64::MAX - 1..u64::MAX) + ); + assert_eq!(layout.page_range(u64::MAX), None); + let layout = PageLayout::new(nonzero(64), u64::MAX); + assert_eq!(layout.page_range(u64::MAX), None); + let last = u64::MAX / 64; + assert_eq!(layout.page_range(last), Some(last * 64..u64::MAX)); + assert_eq!( + layout.page_ids(u64::MAX - 1..u64::MAX).unwrap(), + last..last + 1 + ); + } +} diff --git a/src/common/cache/src/paged.rs b/src/common/cache/src/paged.rs new file mode 100644 index 00000000000..f2e8a659ddf --- /dev/null +++ b/src/common/cache/src/paged.rs @@ -0,0 +1,384 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Immutable byte pages backed by native cache loading and direct range reads. + +#[cfg(test)] +mod tests; + +use std::collections::{BTreeMap, BTreeSet}; +use std::future::Future; +use std::hash::Hash; +use std::io; +use std::num::NonZeroU64; +use std::ops::Range; +use std::sync::Arc; +use std::task::Poll; + +use bytes::{Bytes, BytesMut}; +use common_base::range_read::RangeReader; +use futures::channel::{mpsc, oneshot}; +use futures::{StreamExt, future}; + +use crate::metrics::{CacheMetrics, CacheSizeMetric}; +use crate::page::PageLayout; +use crate::{Cache, RemovalCause}; + +/// Business-level events; callers retain their existing hit/miss metric units. +pub enum PageEvent { + /// Unique logical page lookups in one request, before any missing loads execute. + Lookup { hits: usize, misses: usize }, + /// One successful native page loader, excluding waiters and hits. + Loaded { bytes: u64 }, + /// Backend weight, not process RSS; sampled after fetch/maintenance. + Residency { usage: u64, capacity: u64 }, + /// Backend eviction, excluding unadmitted values released by their owners. + Evicted { cause: RemovalCause }, +} + +/// Caller-bound metrics without storage-engine dependencies or global registrations. +#[derive(Clone)] +pub struct PageCacheObserver { + pub bytes: CacheSizeMetric, + pub loads: Option>, + pub events: Arc, +} + +/// Byte accounting for one successful fetch. Hits are initial lookups, not waiters. +#[derive(Debug, Default)] +pub struct PageFetchStats { + pub cache_hit: usize, + pub cache_miss: usize, + /// Requested intersections already cached; overlapping requests count repeatedly. + pub cached_bytes: u64, + pub cached_fragments: usize, + /// Original request lengths, including duplicates and overlaps. + pub requested_bytes: u64, + /// Sum of unique logical page lengths, clipped at EOF, including cache hits. + pub aligned_bytes: u64, + /// Union of original request byte ranges, excluding overlaps. + pub unique_requested_bytes: u64, +} + +/// Caches fixed-size byte pages of immutable content identified by `K`. +/// File identity must include any version needed to distinguish different contents. +/// Missing pages elected by this fetch are submitted through [`RangeReader::read_vec`]. +pub struct PagedCache { + cache: Cache<(K, u64), Arc>, + page_size: u64, + observer: PageCacheObserver, +} + +impl PagedCache { + /// Creates an independent cache budget. Zero page size is rejected by fetch. + pub fn new( + name: &str, + capacity: u64, + page_size: u64, + weigher: impl Fn(&(K, u64), &Bytes) -> u32 + Send + Sync + 'static, + observer: PageCacheObserver, + ) -> Self { + let events = observer.events.clone(); + let mut cache = Cache::new( + name, + capacity, + move |key, value: &Arc| weigher(key, &value.bytes), + move |_, _, cause| { + events(PageEvent::Evicted { cause }); + }, + ) + .with_size_metric(observer.bytes.clone()); + if let Some(metrics) = &observer.loads { + cache = cache.with_metrics(metrics.clone()); + } + Self { + cache, + page_size, + observer, + } + } + + /// Invalidates completed pages. Concurrent loads are not fenced. + pub fn invalidate_if( + &self, + predicate: impl Fn(&K) -> bool + Send + Sync + 'static, + ) -> io::Result<()> { + self.cache + .invalidate_if(move |(file, _), _| predicate(file)) + } + + /// Looks up a completed page without updating business lookup counters. + pub async fn contains(&self, file: &K, page_id: u64) -> bool { + self.cache.get(&(file.clone(), page_id)).await.is_some() + } + + /// Recovers immutable file size from a requested page already in memory. + pub async fn known_file_size(&self, file: &K, ranges: &[Range]) -> Option { + let layout = PageLayout::new(NonZeroU64::new(self.page_size)?, u64::MAX); + for range in ranges { + for id in layout.page_ids(range.clone()).ok()? { + if let Some(page) = self.cache.get(&(file.clone(), id)).await { + return Some(page.file_size); + } + } + } + None + } + + /// Applies pending cache admission and eviction work. + pub async fn run_pending_tasks(&self) { + self.cache.run_pending_tasks().await; + self.observe_residency(); + } + + /// Fetches bytes in input order. The reader must address the same immutable + /// content as `file`. Each successful loader reads one complete missing page. + pub async fn fetch( + &self, + file: K, + file_size: u64, + ranges: &[Range], + reader: &R, + ) -> Result<(Vec, PageFetchStats), Arc> { + if ranges.iter().any(|r| r.start > r.end || r.end > file_size) { + return Err(invalid_request()); + } + let page_size = NonZeroU64::new(self.page_size).ok_or_else(invalid_request)?; + let layout = PageLayout::new(page_size, file_size); + let mut stats = PageFetchStats::default(); + let mut ids = BTreeSet::new(); + for range in ranges { + stats.requested_bytes = stats + .requested_bytes + .checked_add(range.end - range.start) + .ok_or_else(invalid_request)?; + ids.extend(layout.page_ids(range.clone()).map_err(Arc::new)?); + } + stats.unique_requested_bytes = unique_requested_bytes(ranges); + let mut pages = BTreeMap::new(); + let mut missing = Vec::new(); + for id in ids { + let range = layout.page_range(id).ok_or_else(invalid_request)?; + stats.aligned_bytes += range.end - range.start; + if let Some(page) = self.cache.get(&(file.clone(), id)).await { + check_file_size(&page, file_size)?; + pages.insert(id, page.bytes.clone()); + } else { + missing.push(id); + } + } + for range in ranges { + for id in layout.page_ids(range.clone()).map_err(Arc::new)? { + if pages.contains_key(&id) { + let slice = layout + .request_slice(id, range.clone()) + .map_err(Arc::new)? + .ok_or_else(invalid_request)?; + stats.cached_bytes += slice.len() as u64; + stats.cached_fragments += 1; + } + } + } + stats.cache_hit = pages.len(); + stats.cache_miss = missing.len(); + (self.observer.events)(PageEvent::Lookup { + hits: stats.cache_hit, + misses: stats.cache_miss, + }); + if !missing.is_empty() { + let (requests, receiver) = mpsc::unbounded(); + let batch_size = missing.len().max(1); + let mut loads = missing + .into_iter() + .map(|id| { + let requests = requests.clone(); + let key = (file.clone(), id); + Box::pin(future::maybe_done(async move { + let range = layout.page_range(id).ok_or_else(invalid_request)?; + let page = self + .cache + .try_get_with(&key, async { + let (send, receive) = + oneshot::channel::>>(); + requests + .unbounded_send((range.clone(), send)) + .map_err(|_| io::Error::other("Page read batch closed"))?; + let bytes = receive + .await + .map_err(|_| io::Error::other("Page read batch cancelled"))? + .map_err(|error: Arc| { + io::Error::new(error.kind(), error) + })?; + if bytes.len() as u64 != range.end - range.start { + return Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + "Invalid page response length", + )); + } + // Detach readers' shared backing buffers from cache accounting. + Ok(self + .create_cached_page(file_size, Bytes::copy_from_slice(&bytes))) + }) + .await?; + check_file_size(&page, file_size)?; + Ok::<_, Arc>((id, page.bytes.clone())) + })) + }) + .collect::>(); + // Poll every page before driving I/O, so executor fairness limits do not + // split one fetch's immediately available misses into separate batches. + let loads = future::poll_fn(move |cx| { + let mut complete = true; + for load in &mut loads { + complete &= load.as_mut().poll(cx).is_ready(); + if let Some(Err(error)) = load.as_mut().output_mut() { + return Poll::Ready(Err(error.clone())); + } + } + if !complete { + return Poll::Pending; + } + Poll::Ready( + loads + .iter_mut() + .map(|load| load.as_mut().take_output().ok_or_else(invalid_request)?) + .collect::, _>>(), + ) + }); + // The channel contains only loaders elected by Moka for this fetch. Waiters + // enqueue nothing; a waiter promoted after cancellation can enqueue later. + drop(requests); + let read_batches = async { + let mut batches = receiver.ready_chunks(batch_size); + while let Some(batch) = batches.next().await { + let (ranges, replies): (Vec<_>, Vec<_>) = batch.into_iter().unzip(); + let result = reader.read_vec(&ranges).await.and_then(|bytes| { + if bytes.len() == ranges.len() { + Ok(bytes) + } else { + Err(io::Error::new( + io::ErrorKind::InvalidData, + "Invalid page response count", + )) + } + }); + match result { + Ok(bytes) => { + for (reply, bytes) in replies.into_iter().zip(bytes) { + let _ = reply.send(Ok(bytes)); + } + } + Err(error) => { + let error = Arc::new(error); + for reply in replies { + let _ = reply.send(Err(error.clone())); + } + } + } + } + Ok::<(), Arc>(()) + }; + let (loaded, ()) = future::try_join(loads, read_batches).await?; + pages.extend(loaded); + } + let mut output = Vec::with_capacity(ranges.len()); + for range in ranges { + if range.is_empty() { + output.push(Bytes::new()); + continue; + } + let ids = layout.page_ids(range.clone()).map_err(Arc::new)?; + if ids.end - ids.start == 1 { + let slice = layout + .request_slice(ids.start, range.clone()) + .map_err(Arc::new)? + .ok_or_else(invalid_request)?; + let page = pages.get(&ids.start).ok_or_else(invalid_request)?; + if slice.end > page.len() { + return Err(invalid_request()); + } + output.push(page.slice(slice)); + continue; + } + let len = usize::try_from(range.end - range.start).map_err(|_| invalid_request())?; + let mut assembled = BytesMut::with_capacity(len); + for id in ids { + let slice = layout + .request_slice(id, range.clone()) + .map_err(Arc::new)? + .ok_or_else(invalid_request)?; + let page = pages.get(&id).ok_or_else(invalid_request)?; + if slice.end > page.len() { + return Err(invalid_request()); + } + assembled.extend_from_slice(&page[slice]); + } + if assembled.len() != len { + return Err(invalid_request()); + } + output.push(assembled.freeze()); + } + self.observe_residency(); + Ok((output, stats)) + } + + fn create_cached_page(&self, file_size: u64, bytes: Bytes) -> Arc { + (self.observer.events)(PageEvent::Loaded { + bytes: bytes.len() as u64, + }); + Arc::new(CachedPage { file_size, bytes }) + } + + fn observe_residency(&self) { + (self.observer.events)(PageEvent::Residency { + usage: self.cache.weighted_size(), + capacity: self.cache.capacity(), + }); + } +} + +struct CachedPage { + file_size: u64, + bytes: Bytes, +} + +fn invalid_request() -> Arc { + Arc::new(io::Error::new( + io::ErrorKind::InvalidInput, + "Invalid cached page range", + )) +} + +fn check_file_size(page: &CachedPage, expected: u64) -> Result<(), Arc> { + if page.file_size != expected { + return Err(Arc::new(io::Error::new( + io::ErrorKind::InvalidInput, + "Cached content identity used with inconsistent file sizes", + ))); + } + Ok(()) +} + +/// Counts the union of request ranges already validated by fetch. +fn unique_requested_bytes(ranges: &[Range]) -> u64 { + let mut ranges = ranges.to_vec(); + ranges.sort_unstable_by_key(|range| range.start); + let mut end = 0; + let mut bytes = 0; + for range in ranges { + bytes += range.end.saturating_sub(range.start.max(end)); + end = end.max(range.end); + } + bytes +} diff --git a/src/common/cache/src/paged/tests.rs b/src/common/cache/src/paged/tests.rs new file mode 100644 index 00000000000..556455ebfb1 --- /dev/null +++ b/src/common/cache/src/paged/tests.rs @@ -0,0 +1,492 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::future::Future; +use std::pin::Pin; +use std::sync::Mutex; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::time::Duration; + +use async_trait::async_trait; +use common_base::range_read::Metadata; +use tokio::sync::Semaphore; + +use super::*; + +#[derive(Default)] +struct Reader { + reads: Mutex>>, + batches: Mutex>>>, + gate: Option, + error: Option, +} + +impl Reader { + fn gated() -> Self { + Self { + gate: Some(Semaphore::new(0)), + ..Default::default() + } + } + + fn count(&self) -> usize { + self.reads.lock().unwrap().len() + } + + fn release(&self) { + self.gate.as_ref().unwrap().add_permits(64); + } +} + +#[async_trait] +impl RangeReader for Reader { + async fn metadata(&self) -> io::Result { + Ok(Metadata { content_length: 16 }) + } + + async fn read_vec(&self, ranges: &[Range]) -> io::Result> { + self.batches.lock().unwrap().push(ranges.to_vec()); + future::try_join_all(ranges.iter().map(|range| self.read(range.clone()))).await + } + + async fn read(&self, range: Range) -> io::Result { + self.reads.lock().unwrap().push(range.clone()); + if let Some(gate) = &self.gate { + gate.acquire().await.unwrap().forget(); + } + if let Some(kind) = self.error { + return Err(io::Error::new(kind, "backend failed")); + } + Ok(Bytes::from(range.map(|n| n as u8).collect::>())) + } +} + +struct Observation { + loaded: AtomicUsize, + bytes: crate::metrics::CacheSizeMetrics, +} + +impl Observation { + fn bytes(&self) -> i64 { + use prometheus::core::Collector; + self.bytes.collect()[0].get_metric()[0] + .get_gauge() + .as_ref() + .unwrap() + .value() as i64 + } +} + +fn cache() -> (PagedCache, Arc) { + let observation = Arc::new(Observation { + loaded: AtomicUsize::new(0), + bytes: crate::metrics::CacheSizeMetrics::new( + prometheus::Opts::new("test_page_bytes", "test"), + &[], + ) + .unwrap(), + }); + let events = observation.clone(); + let observer = PageCacheObserver { + bytes: observation.bytes.bind(&[]), + loads: None, + events: Arc::new(move |event| { + if let PageEvent::Loaded { .. } = event { + events.loaded.fetch_add(1, Ordering::Relaxed); + } + }), + }; + ( + PagedCache::new("pages", 1024, 4, |_, bytes| bytes.len() as u32, observer), + observation, + ) +} + +async fn start_read(mut future: Pin<&mut F>, reader: &Reader, count: usize) { + tokio::time::timeout(Duration::from_secs(5), async { + loop { + assert!(futures::poll!(future.as_mut()).is_pending()); + if reader.count() >= count { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); +} + +#[tokio::test] +async fn test_partial_hits_overlaps_and_eof_preserve_order() { + let (cache, observation) = cache(); + let reader = Reader::default(); + let (first, stats) = cache.fetch(1, 10, &[2..6, 4..5], &reader).await.unwrap(); + assert_eq!( + first, + vec![Bytes::from_static(&[2, 3, 4, 5]), Bytes::from_static(&[4])] + ); + assert_eq!((stats.cache_hit, stats.cache_miss), (0, 2)); + assert_eq!( + ( + stats.requested_bytes, + stats.unique_requested_bytes, + stats.aligned_bytes + ), + (5, 4, 8) + ); + let (second, stats) = cache + .fetch(1, 10, &[6..10, 1..3, 10..10], &reader) + .await + .unwrap(); + assert_eq!( + second, + vec![ + Bytes::from_static(&[6, 7, 8, 9]), + Bytes::from_static(&[1, 2]), + Bytes::new() + ] + ); + assert_eq!((stats.cache_hit, stats.cache_miss), (2, 1)); + assert_eq!((stats.cached_bytes, stats.cached_fragments), (4, 2)); + assert_eq!(observation.loaded.load(Ordering::Relaxed), 3); + assert_eq!( + cache + .known_file_size(&1, std::slice::from_ref(&(8..9))) + .await, + Some(10) + ); + assert_eq!(*reader.reads.lock().unwrap(), vec![0..4, 4..8, 8..10]); + cache.run_pending_tasks().await; + assert_eq!(observation.bytes(), 10); + cache.invalidate_if(|file| *file == 1).unwrap(); + cache.run_pending_tasks().await; + assert!(!cache.contains(&1, 0).await); + assert_eq!(observation.bytes(), 0); + // Returned values remain valid after invalidation, while cache accounting is released. + assert_eq!(&second[0][..], &[6, 7, 8, 9]); +} + +#[tokio::test] +async fn test_overlapping_fetches_share_pages_and_only_read_missing_bytes() { + let (cache, observation) = cache(); + let reader = Reader::gated(); + let first_range = 0..8; + let first_ranges = std::slice::from_ref(&first_range); + let second_range = 4..12; + let second_ranges = std::slice::from_ref(&second_range); + let mut first = Box::pin(cache.fetch(1, 16, first_ranges, &reader)); + start_read(first.as_mut(), &reader, 2).await; + let mut second = Box::pin(cache.fetch(1, 16, second_ranges, &reader)); + start_read(second.as_mut(), &reader, 3).await; + reader.release(); + let (first, second) = tokio::time::timeout(Duration::from_secs(5), async { + tokio::join!(first, second) + }) + .await + .unwrap(); + assert_eq!(&first.unwrap().0[0][..], &(0..8).collect::>()); + assert_eq!(&second.unwrap().0[0][..], &(4..12).collect::>()); + assert_eq!(*reader.reads.lock().unwrap(), vec![0..4, 4..8, 8..12]); + assert_eq!(observation.loaded.load(Ordering::Relaxed), 3); +} + +#[tokio::test] +async fn test_cancelling_owner_allows_waiter_to_finish() { + let (cache, observation) = cache(); + let reader = Reader::gated(); + let range = 0..4; + let ranges = std::slice::from_ref(&range); + let mut owner = Box::pin(cache.fetch(1, 16, ranges, &reader)); + start_read(owner.as_mut(), &reader, 1).await; + let mut waiter = Box::pin(cache.fetch(1, 16, ranges, &reader)); + assert!(futures::poll!(&mut waiter).is_pending()); + assert_eq!(reader.count(), 1); + drop(owner); + reader.release(); + let result = tokio::time::timeout(Duration::from_secs(5), waiter) + .await + .unwrap() + .unwrap(); + assert_eq!(&result.0[0][..], &[0, 1, 2, 3]); + assert_eq!(reader.count(), 2); + assert_eq!(observation.loaded.load(Ordering::Relaxed), 1); +} + +#[tokio::test] +async fn test_cancelling_waiter_does_not_cancel_owner() { + let (cache, _) = cache(); + let reader = Reader::gated(); + let range = 0..4; + let ranges = std::slice::from_ref(&range); + let mut owner = Box::pin(cache.fetch(1, 16, ranges, &reader)); + start_read(owner.as_mut(), &reader, 1).await; + let mut waiter = Box::pin(cache.fetch(1, 16, ranges, &reader)); + assert!(futures::poll!(&mut waiter).is_pending()); + drop(waiter); + reader.release(); + let result = tokio::time::timeout(Duration::from_secs(5), owner) + .await + .unwrap() + .unwrap(); + assert_eq!(&result.0[0][..], &[0, 1, 2, 3]); + assert_eq!(reader.count(), 1); +} + +#[tokio::test] +async fn test_backend_error_is_shared_and_not_cached() { + let (cache, observation) = cache(); + let reader = Reader { + error: Some(io::ErrorKind::BrokenPipe), + ..Reader::gated() + }; + let range = 0..4; + let ranges = std::slice::from_ref(&range); + let mut owner = Box::pin(cache.fetch(1, 16, ranges, &reader)); + start_read(owner.as_mut(), &reader, 1).await; + let mut waiter = Box::pin(cache.fetch(1, 16, ranges, &reader)); + assert!(futures::poll!(&mut waiter).is_pending()); + reader.release(); + let (owner, waiter) = tokio::time::timeout(Duration::from_secs(5), async { + tokio::join!(owner, waiter) + }) + .await + .unwrap(); + for result in [owner, waiter] { + let error = result.unwrap_err(); + assert_eq!(error.kind(), io::ErrorKind::BrokenPipe); + assert!(error.to_string().contains("backend failed")); + } + assert_eq!(reader.count(), 1); + assert_eq!(observation.loaded.load(Ordering::Relaxed), 0); + // The failed load was not cached. + let healthy = Reader::default(); + assert!(cache.fetch(1, 16, ranges, &healthy).await.is_ok()); + assert_eq!(healthy.count(), 1); +} + +#[tokio::test] +async fn test_completed_page_rejects_conflicting_file_size() { + let (cache, _) = cache(); + let reader = Reader::default(); + let range = 0..4; + let ranges = std::slice::from_ref(&range); + cache.fetch(1, 16, ranges, &reader).await.unwrap(); + assert!(cache.fetch(1, 12, ranges, &reader).await.is_err()); + assert_eq!(reader.count(), 1); + // A different identity remains independent. + assert!(cache.fetch(2, 12, ranges, &reader).await.is_ok()); +} + +#[tokio::test] +async fn test_shared_load_rejects_conflicting_file_size() { + let (cache, _) = cache(); + let reader = Reader::gated(); + let range = 0..4; + let ranges = std::slice::from_ref(&range); + let mut owner = Box::pin(cache.fetch(1, 16, ranges, &reader)); + start_read(owner.as_mut(), &reader, 1).await; + let mut conflicting = Box::pin(cache.fetch(1, 12, ranges, &reader)); + assert!(futures::poll!(&mut conflicting).is_pending()); + reader.release(); + let (owner, conflicting) = tokio::time::timeout(Duration::from_secs(5), async { + tokio::join!(owner, conflicting) + }) + .await + .unwrap(); + assert!(owner.is_ok()); + assert!(conflicting.is_err()); + assert_eq!(reader.count(), 1); +} + +#[tokio::test] +async fn test_empty_and_invalid_requests_do_not_issue_reads() { + let (cache, _) = cache(); + let reader = Reader::default(); + assert!(cache.fetch(1, 0, &[], &reader).await.unwrap().0.is_empty()); + assert_eq!( + cache.fetch(1, 0, &[0..0, 0..0], &reader).await.unwrap().0, + vec![Bytes::new(), Bytes::new()] + ); + for ranges in [vec![Range { start: 2, end: 1 }], vec![0..17, 0..0]] { + assert!(cache.fetch(1, 16, &ranges, &reader).await.is_err()); + } + assert_eq!(reader.count(), 0); +} + +// Keep the Send contract visible before a caller erases this future in async-trait. +fn send_fetch<'a, K, R>( + cache: &'a PagedCache, + file: K, + ranges: &'a [Range], + reader: &'a R, +) -> impl Future, PageFetchStats), Arc>> + Send + 'a +where + K: Hash + Eq + Clone + Send + Sync + 'static, + R: RangeReader + ?Sized, +{ + cache.fetch(file, 16, ranges, reader) +} + +#[tokio::test] +async fn test_generic_page_fetch_has_a_send_future() { + let (cache, _) = cache(); + let reader = Reader::default(); + let range = 1..7; + let ranges = std::slice::from_ref(&range); + let future: futures::future::BoxFuture<'_, _> = + Box::pin(send_fetch(&cache, 1, ranges, &reader)); + let (bytes, _) = future.await.unwrap(); + assert_eq!(bytes, vec![Bytes::from_static(&[1, 2, 3, 4, 5, 6])]); +} + +#[tokio::test] +async fn test_direct_page_rejects_wrong_lengths_without_caching() { + struct FixedReader(usize); + #[async_trait] + impl RangeReader for FixedReader { + async fn metadata(&self) -> io::Result { + unreachable!("fetch already has the immutable file size") + } + async fn read(&self, _: Range) -> io::Result { + Ok(Bytes::from(vec![7; self.0])) + } + } + for response_size in [3, 5] { + let (cache, observation) = cache(); + let range = 0..4; + let ranges = std::slice::from_ref(&range); + let error = cache + .fetch(1, 8, ranges, &FixedReader(response_size)) + .await + .unwrap_err(); + assert_eq!(error.kind(), io::ErrorKind::UnexpectedEof); + assert_eq!(observation.loaded.load(Ordering::Relaxed), 0); + assert!(!cache.contains(&1, 0).await); + let reader = Reader::default(); + assert!(cache.fetch(1, 8, ranges, &reader).await.is_ok()); + assert_eq!(reader.count(), 1); + } +} + +#[tokio::test] +async fn test_direct_page_detaches_shared_backing_buffer() { + struct SliceReader(Bytes); + #[async_trait] + impl RangeReader for SliceReader { + async fn metadata(&self) -> io::Result { + unreachable!("fetch already has the immutable file size") + } + async fn read(&self, range: Range) -> io::Result { + Ok(self.0.slice(range.start as usize..range.end as usize)) + } + } + let (cache, observation) = cache(); + let reader = SliceReader(Bytes::from(vec![7; 1024])); + let range = 0..4; + let (output, _) = cache + .fetch(1, 1024, std::slice::from_ref(&range), &reader) + .await + .unwrap(); + assert_eq!(output[0].as_ref(), [7; 4]); + assert_ne!(output[0].as_ptr(), reader.0.as_ptr()); + cache.run_pending_tasks().await; + assert_eq!(observation.bytes(), 4); +} + +#[tokio::test] +async fn test_direct_page_larger_than_eight_mib_is_not_split() { + let (base, _) = cache(); + let page_size = 9 * 1024 * 1024; + let cache = PagedCache::new( + "large_pages", + page_size * 2, + page_size, + |_: &(u64, u64), bytes: &Bytes| bytes.len() as u32, + base.observer.clone(), + ); + let reader = Reader::default(); + let range = 0..1; + let (output, _) = cache + .fetch(1, page_size, std::slice::from_ref(&range), &reader) + .await + .unwrap(); + assert_eq!(output[0].as_ref(), [0]); + assert_eq!( + reader.reads.lock().unwrap().as_slice(), + std::slice::from_ref(&(0..page_size)) + ); +} + +#[test] +fn test_unique_requested_bytes() { + for (ranges, expected) in [ + (vec![], 0), + (vec![4..4, 0..0], 0), + (vec![8..12, 2..6, 4..10, 2..6], 10), + (vec![0..4, 4..8, 12..16], 12), + (vec![0..u64::MAX, 1..2], u64::MAX), + ] { + assert_eq!(unique_requested_bytes(&ranges), expected); + } +} + +#[tokio::test] +async fn test_missing_pages_use_batch_reads() { + let (cache, _) = cache(); + let reader = Reader::default(); + let range = 0..512; + let ranges = std::slice::from_ref(&range); + let (bytes, _) = cache.fetch(1, 512, ranges, &reader).await.unwrap(); + assert_eq!( + bytes[0], + Bytes::from((0..512).map(|n| n as u8).collect::>()) + ); + assert_eq!( + *reader.batches.lock().unwrap(), + vec![(0..128).map(|n| n * 4..n * 4 + 4).collect::>()] + ); + cache.fetch(1, 512, ranges, &reader).await.unwrap(); + assert_eq!(reader.batches.lock().unwrap().len(), 1); +} + +#[tokio::test] +async fn test_shared_error_cancels_other_pending_batch() { + let (cache, _) = cache(); + let failed = Reader { + error: Some(io::ErrorKind::BrokenPipe), + ..Reader::gated() + }; + let blocked = Reader::gated(); + let owner_range = 0..4; + let waiter_range = 0..8; + let mut owner = Box::pin(cache.fetch(1, 16, std::slice::from_ref(&owner_range), &failed)); + start_read(owner.as_mut(), &failed, 1).await; + let mut waiter = Box::pin(cache.fetch(1, 16, std::slice::from_ref(&waiter_range), &blocked)); + start_read(waiter.as_mut(), &blocked, 1).await; + failed.release(); + assert_eq!(owner.await.unwrap_err().kind(), io::ErrorKind::BrokenPipe); + // The unrelated page stays blocked; the shared error must cancel its I/O. + let error = tokio::time::timeout(Duration::from_secs(5), waiter) + .await + .unwrap() + .unwrap_err(); + assert_eq!(error.kind(), io::ErrorKind::BrokenPipe); + let healthy = Reader::default(); + assert!( + cache + .fetch(1, 16, std::slice::from_ref(&waiter_range), &healthy) + .await + .is_ok() + ); +} From 045cf74c3f41e29d63b52c7e708b5f023fe2e039 Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Wed, 28 Oct 2026 17:09:26 +0000 Subject: [PATCH 131/239] perf(mito2): pipeline compaction and parallelize decoding and encoding (#9945) * perf(mito2): merge compaction inputs on a separate task Compaction read, merged and encoded each output on a single task, so parquet encoding waited for the merge and the merge waited for encoding. The merge now runs on its own compaction runtime task and hands merged batches to the SST writer through a bounded channel. The channel also closes when that task panics, so the reading side awaits the task before reporting the end of the input. A panic fails the compaction instead of producing a truncated SST. One output of 4.8M rows took 1.28s instead of 1.65s. Signed-off-by: Lei, HUANG * perf(mito2): parallelize decoding and encoding of compaction outputs Parquet encoding took about two thirds of a compaction output's time and ran on a single thread. When `experimental_compaction_output_parallelism` is above 1, the SST writer encodes the columns of each row group on up to that many blocking threads of the compaction runtime while the next row group is buffered, and the merge decodes up to that many input files ahead of it on compaction runtime tasks. The parallel writer feeds every column writer the same input as `ArrowWriter`, so the SSTs are byte-identical. The bytes of a row group reach the object store writer one row group later, so `max_file_size` may split a file one row group later. The option defaults to 1, which keeps today's CPU usage. One output of 4.8M rows took 0.44s with 4 and 0.38s with 8, against 1.28s with 1. Signed-off-by: Lei, HUANG * fix(mito2): cancel compaction producers before consumer polling Signed-off-by: Lei, HUANG * feat(mito2): support per-table compaction output parallelism Signed-off-by: Lei, HUANG * fix(mito2): drain encoding tasks before returning failures Abort remaining queued encoding tasks and await running tasks before propagating the first error, keeping row-group work inside the caller's compaction memory budget lifetime. Add deterministic error and panic tests for queued cancellation, running-task cleanup, budget retention, and original error preservation. Signed-off-by: Lei, HUANG * fix(mito2): drain encoders on SST write errors Add asynchronous abort cleanup to both SST write entry points so source errors cannot detach in-flight encoding work after its compaction memory budget is released. Share encoder task cleanup with the encoding-failure path and preserve the original write error. Cover source-read failures in flat and primary-key writers with deterministic running and queued tasks. Signed-off-by: Lei, HUANG * fix(mito2): drain compaction outputs on cancellation Stop scheduling new output batches on cancellation instead of aborting active writers. Await started outputs and retain successful SSTs, draining sibling tasks before propagating join failures. Add deterministic cancellation regressions covering blocking encoder lifetime, budget retention, output ownership, join failures, and cancellation before work starts. Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- config/config.md | 2 + config/datanode.example.toml | 7 + config/standalone.example.toml | 7 + .../src/bin/query_perf_fixture/direct_sst.rs | 1 + src/mito2/src/compaction/compactor.rs | 307 ++++++--- src/mito2/src/compaction/reader.rs | 16 + src/mito2/src/compaction/window.rs | 1 + src/mito2/src/config.rs | 5 + src/mito2/src/engine/alter_test.rs | 65 ++ src/mito2/src/engine/compaction_test.rs | 31 +- src/mito2/src/read.rs | 8 + src/mito2/src/read/scan_region.rs | 11 + src/mito2/src/read/seq_scan.rs | 24 +- src/mito2/src/read/stream.rs | 275 +++++++- src/mito2/src/region/options.rs | 65 ++ src/mito2/src/sst/parquet.rs | 60 +- src/mito2/src/sst/parquet/parallel_writer.rs | 648 ++++++++++++++++++ src/mito2/src/sst/parquet/writer.rs | 230 ++++++- src/mito2/src/worker/handle_alter.rs | 6 + src/store-api/src/mito_engine_options.rs | 7 + src/store-api/src/region_request.rs | 73 +- src/table/src/metadata.rs | 68 +- tests-integration/tests/http.rs | 1 + .../compaction_output_parallelism.result | 140 ++++ .../alter/compaction_output_parallelism.sql | 35 + .../compaction_output_parallelism/case.toml | 8 + .../compaction_output_parallelism/setup.sql | 7 + .../verify.result | 90 +++ .../compaction_output_parallelism/verify.sql | 15 + 29 files changed, 2072 insertions(+), 141 deletions(-) create mode 100644 src/mito2/src/sst/parquet/parallel_writer.rs create mode 100644 tests/cases/standalone/common/alter/compaction_output_parallelism.result create mode 100644 tests/cases/standalone/common/alter/compaction_output_parallelism.sql create mode 100644 tests/compatibility/cases/compaction_output_parallelism/case.toml create mode 100644 tests/compatibility/cases/compaction_output_parallelism/setup.sql create mode 100644 tests/compatibility/cases/compaction_output_parallelism/verify.result create mode 100644 tests/compatibility/cases/compaction_output_parallelism/verify.sql diff --git a/config/config.md b/config/config.md index 15cd4915f9d..711fe019a19 100644 --- a/config/config.md +++ b/config/config.md @@ -193,6 +193,7 @@ | `region_engine.mito.max_background_purges` | Integer | Auto | Max number of running background purge jobs (default: number of cpu cores). | | `region_engine.mito.experimental_compaction_memory_limit` | String | 0 | Memory budget for compaction tasks.
Supports absolute size (e.g., "2GiB", "512MB") or percentage of system memory (e.g., "50%").
Setting it to 0 or "unlimited" disables the limit. | | `region_engine.mito.experimental_compaction_on_exhausted` | String | wait | Behavior when compaction cannot acquire memory from the budget.
Options: "wait" (default, 10s), "wait()", "fail" | +| `region_engine.mito.experimental_compaction_output_parallelism` | Integer | 1 | Default per-output parallelism for input decoding and SST column encoding.
Each stage is bounded independently; this is not a global thread limit.
Values below 2 disable parallel decoding/encoding, but not the merge pipeline.
Tables can override it with 'compaction.experimental_output_parallelism' (positive integer). | | `region_engine.mito.auto_flush_interval` | String | `10m` | Interval to auto flush a region if it has not flushed yet. | | `region_engine.mito.global_write_buffer_size` | String | Auto | Global write buffer size for all regions. If not set, it's default to 1/8 of OS memory with a max limitation of 1GB. | | `region_engine.mito.global_write_buffer_reject_size` | String | Auto | Global write buffer size threshold to reject write requests. If not set, it's default to 2 times of `global_write_buffer_size`. | @@ -672,6 +673,7 @@ | `region_engine.mito.max_background_purges` | Integer | Auto | Max number of running background purge jobs (default: number of cpu cores). | | `region_engine.mito.experimental_compaction_memory_limit` | String | 0 | Memory budget for compaction tasks.
Supports absolute size (e.g., "2GiB", "512MB") or percentage of system memory (e.g., "50%").
Setting it to 0 or "unlimited" disables the limit. | | `region_engine.mito.experimental_compaction_on_exhausted` | String | wait | Behavior when compaction cannot acquire memory from the budget.
Options: "wait" (default, 10s), "wait()", "fail" | +| `region_engine.mito.experimental_compaction_output_parallelism` | Integer | 1 | Default per-output parallelism for input decoding and SST column encoding.
Each stage is bounded independently; this is not a global thread limit.
Values below 2 disable parallel decoding/encoding, but not the merge pipeline.
Tables can override it with 'compaction.experimental_output_parallelism' (positive integer). | | `region_engine.mito.auto_flush_interval` | String | `10m` | Interval to auto flush a region if it has not flushed yet. | | `region_engine.mito.global_write_buffer_size` | String | Auto | Global write buffer size for all regions. If not set, it's default to 1/8 of OS memory with a max limitation of 1GB. | | `region_engine.mito.global_write_buffer_reject_size` | String | Auto | Global write buffer size threshold to reject write requests. If not set, it's default to 2 times of `global_write_buffer_size` | diff --git a/config/datanode.example.toml b/config/datanode.example.toml index acc65582e81..a2457302a10 100644 --- a/config/datanode.example.toml +++ b/config/datanode.example.toml @@ -590,6 +590,13 @@ compress_manifest = false ## @toml2docs:none-default="wait" #+ experimental_compaction_on_exhausted = "wait" +## Default per-output parallelism for input decoding and SST column encoding. +## Each stage is bounded independently; this is not a global thread limit. +## Values below 2 disable parallel decoding/encoding, but not the merge pipeline. +## Tables can override it with 'compaction.experimental_output_parallelism' (positive integer). +## @toml2docs:none-default="1" +#+ experimental_compaction_output_parallelism = 1 + ## Interval to auto flush a region if it has not flushed yet. auto_flush_interval = "10m" diff --git a/config/standalone.example.toml b/config/standalone.example.toml index 5cfbb95729e..d047952538c 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -805,6 +805,13 @@ compress_manifest = false ## @toml2docs:none-default="wait" #+ experimental_compaction_on_exhausted = "wait" +## Default per-output parallelism for input decoding and SST column encoding. +## Each stage is bounded independently; this is not a global thread limit. +## Values below 2 disable parallel decoding/encoding, but not the merge pipeline. +## Tables can override it with 'compaction.experimental_output_parallelism' (positive integer). +## @toml2docs:none-default="1" +#+ experimental_compaction_output_parallelism = 1 + ## Interval to auto flush a region if it has not flushed yet. auto_flush_interval = "10m" diff --git a/src/cmd/src/bin/query_perf_fixture/direct_sst.rs b/src/cmd/src/bin/query_perf_fixture/direct_sst.rs index 27a6b4109cd..0b09d1634ca 100644 --- a/src/cmd/src/bin/query_perf_fixture/direct_sst.rs +++ b/src/cmd/src/bin/query_perf_fixture/direct_sst.rs @@ -470,6 +470,7 @@ pub(super) async fn run_direct_sst(args: DirectArgs) { row_group_size: scenario.layout.row_group_size, max_file_size: None, float_field_encoding: Default::default(), + ..Default::default() }; let infos = match format { FormatType::Flat => writer.write_all_flat(source, Some(sequence), &opts).await, diff --git a/src/mito2/src/compaction/compactor.rs b/src/mito2/src/compaction/compactor.rs index 527750bf33e..1673721ec9b 100644 --- a/src/mito2/src/compaction/compactor.rs +++ b/src/mito2/src/compaction/compactor.rs @@ -17,9 +17,9 @@ use std::sync::Arc; use std::time::Duration; use common_base::Plugins; -use common_base::cancellation::{CancellableFuture, CancellationHandle}; +use common_base::cancellation::CancellationHandle; use common_meta::key::SchemaMetadataManagerRef; -use common_telemetry::{debug, info, warn}; +use common_telemetry::{info, warn}; use common_time::TimeToLive; use either::Either; use itertools::Itertools; @@ -37,7 +37,7 @@ use crate::access_layer::{ }; use crate::cache::{CacheManager, CacheManagerRef}; use crate::compaction::picker::PickerOutput; -use crate::compaction::reader::CompactionSstReaderBuilder; +use crate::compaction::reader::{CompactionSstReaderBuilder, spawn_merge}; use crate::compaction::unit::CompactionUnit; use crate::compaction::{CompactionOutput, find_dynamic_options}; use crate::config::MitoConfig; @@ -59,6 +59,7 @@ use crate::sst::index::intermediate::IntermediateManager; use crate::sst::index::puffin_manager::PuffinManagerFactory; use crate::sst::location::region_dir_from_table_dir; use crate::sst::parquet::metadata::extract_primary_key_range; +use crate::sst::parquet::parallel_writer::EncodeRuntime; use crate::sst::parquet::{SstInfo, WriteOptions}; use crate::sst::version::{SstVersion, SstVersionRef}; @@ -590,8 +591,9 @@ impl SstMerger for DefaultSstMerger { filter_deleted: output.filter_deleted, time_range: output.output_time_range, merge_mode, + read_parallelism: write_opts.encode_parallelism, }; - let source = builder.build_flat_sst_reader().await?; + let source = spawn_merge(builder.build_flat_sst_reader().await?); let mut metrics = Metrics::new(WriteType::Compaction); let region_metadata = context.region_metadata.clone(); @@ -705,6 +707,10 @@ impl DefaultCompactor { /// /// This is the public entry point for external crates that want a /// cancellable compactor without opting into uncommitted-SST tracking. + /// + /// Cancellation stops new output batches. Already spawned outputs finish, + /// and their successful SSTs remain in the returned [`MergeOutput`]. This + /// can delay cancellation until the current batch completes. pub fn new_with_cancel_handle(cancel_handle: Arc) -> Self { Self { merger: DefaultSstMerger, @@ -753,6 +759,14 @@ impl DefaultCompactor { max_file_size: unit.max_file_size, row_group_size: context.region_options.row_group_size(), float_field_encoding: context.region_options.float_field_encoding, + encode_parallelism: context + .region_options + .compaction_output_parallelism_or( + context + .engine_config + .experimental_compaction_output_parallelism, + ), + encode_runtime: EncodeRuntime::Compact, }, ) .await?; @@ -796,6 +810,14 @@ where max_file_size: picker_output.max_file_size, row_group_size: compaction_region.region_options.row_group_size(), float_field_encoding: compaction_region.region_options.float_field_encoding, + encode_parallelism: compaction_region + .region_options + .compaction_output_parallelism_or( + compaction_region + .engine_config + .experimental_compaction_output_parallelism, + ), + encode_runtime: EncodeRuntime::Compact, }; let merger = self.merger.clone(); let context = context.clone(); @@ -820,7 +842,7 @@ where + picker_output.expired_ssts.len(), ); - while !tasks.is_empty() { + while !tasks.is_empty() && !self.cancel_handle.is_cancelled() { let mut chunk: Vec<(Vec, _)> = Vec::with_capacity(internal_parallelism); for _ in 0..internal_parallelism { if let Some(task) = tasks.pop() { @@ -835,61 +857,41 @@ where }) .collect(); - while let Some((inputs, mut handle)) = spawned.pop() { - let abort_handle = handle.abort_handle(); - match CancellableFuture::new(&mut handle, self.cancel_handle.clone()).await { - Ok(Ok(Ok((files, infos)))) => { + // Like merge_unit, cancellation must not drop a writer mid-output: + // its blocking encoders would detach and finalized SSTs could lose + // their owner. Join the whole batch and retain every successful output. + let mut join_error = None; + while let Some((inputs, handle)) = spawned.pop() { + match handle.await { + Ok(Ok((files, infos))) => { output_files.extend(files); if hook.is_some() { all_sst_infos.extend(infos); } compacted_inputs.extend(inputs); } - Ok(Ok(Err(e))) => { + Ok(Err(e)) => { warn!( e; "Failed to merge compaction output for region: {}, inputs: [{}]", region_id, inputs.iter().map(|f| f.file_id.to_string()).join(",") ); } - Ok(Err(e)) => { + Err(e) => { warn!( - "Region {} compaction task join error for inputs: [{}], skipping: {}", + "Region {} compaction task join error for inputs: [{}]: {}", region_id, inputs.iter().map(|f| f.file_id.to_string()).join(","), e ); - // If the cancel handle is cancelled, - // cancel the remaining tasks before returns the error. - if self.cancel_handle.is_cancelled() { - abort_handle.abort(); - for (_, handle) in &spawned { - handle.abort(); - } - for (_, handle) in spawned { - let _ = handle.await; - } - } - return Err(e).context(error::JoinSnafu); - } - Err(_) => { - debug!( - "Compaction merge cancelled for region: {}, aborting remaining {} spawned tasks", - region_id, - spawned.len(), - ); - abort_handle.abort(); - for (_, handle) in &spawned { - handle.abort(); - } - let _ = handle.await; - for (_, handle) in spawned { - let _ = handle.await; - } - break; + // Drain siblings before propagating the first join error. + join_error.get_or_insert(e); } } } + if let Some(e) = join_error { + return Err(e).context(error::JoinSnafu); + } if self.cancel_handle.is_cancelled() { info!("Compaction merge cancelled for region: {}", region_id); @@ -927,17 +929,22 @@ where #[cfg(test)] mod tests { use std::sync::atomic::{AtomicUsize, Ordering}; - use std::sync::{Arc, Mutex}; + use std::sync::{Arc, Mutex, Weak}; use std::time::Duration; + use datatypes::arrow::record_batch::RecordBatch; + use parquet::arrow::ArrowWriter; use store_api::storage::{FileId, RegionId}; - use tokio::time::sleep; + use tokio::sync::{mpsc, oneshot}; + use tokio::time::timeout; use super::{DefaultCompactor, *}; + use crate::compaction::memory_manager::new_compaction_memory_manager; use crate::compaction::picker::PickerOutput; use crate::error::Result; use crate::sst::file::FileHandle; use crate::sst::file_purger::NoopFilePurger; + use crate::test_util::sst_util::new_record_batch_by_range; fn dummy_file_meta() -> FileMeta { FileMeta { @@ -1244,6 +1251,7 @@ mod tests { filter_deleted: true, time_range: None, merge_mode: MergeMode::LastRow, + read_parallelism: 1, } .build_flat_sst_reader() .await @@ -1503,85 +1511,182 @@ mod tests { ); } + #[derive(Debug)] + struct PendingEncoding { + input_id: FileId, + output_id: FileId, + batch: Weak, + release: oneshot::Sender<()>, + } + + /// A merger with real blocking Parquet encoding and deterministic lifetime gates. #[derive(Clone)] - struct BlockingMerger { - call_idx: Arc, + struct GatedEncodingMerger { + started: mpsc::UnboundedSender, + panic_input: Option, } #[async_trait::async_trait] - impl SstMerger for BlockingMerger { + impl SstMerger for GatedEncodingMerger { async fn merge_single_output( &self, _context: CompactionContext, - _output: CompactionOutput, + output: CompactionOutput, _write_opts: WriteOptions, ) -> Result<(Vec, Vec)> { - self.call_idx.fetch_add(1, Ordering::SeqCst); - std::future::pending().await + let input_id = output.inputs[0].meta_ref().file_id; + let mut file = dummy_file_meta(); + let batch = Arc::new(new_record_batch_by_range(&["a", "b"], 0, 4)); + let (release, wait) = oneshot::channel(); + let pending = PendingEncoding { + input_id, + output_id: file.file_id, + batch: Arc::downgrade(&batch), + release, + }; + let started = self.started.clone(); + file.file_size = common_runtime::spawn_blocking_compact(move || { + started.send(pending).unwrap(); + wait.blocking_recv().unwrap(); + let mut bytes = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut bytes, batch.schema(), None)?; + writer.write(&batch)?; + writer.close()?; + Ok::<_, parquet::errors::ParquetError>(bytes.len() as u64) + }) + .await + .context(error::JoinSnafu)? + .context(error::WriteParquetSnafu)?; + if self.panic_input == Some(input_id) { + panic!("merge task panicked"); + } + Ok((vec![file], Vec::new())) } } - #[tokio::test(flavor = "multi_thread")] - async fn test_merge_ssts_cancels_spawned_tasks() { - common_telemetry::init_default_ut_logging(); - - let mut compaction_region = new_test_compaction_region().await; - compaction_region.max_parallelism = 2; - + #[rstest::rstest] + #[case::successful_outputs(false)] + #[case::join_error(true)] + #[tokio::test] + async fn test_merge_ssts_cancellation_waits_for_encoders(#[case] panic_merge: bool) { + let mut region = new_test_compaction_region().await; + region.max_parallelism = 2; + region.region_options.compaction_output_parallelism = Some(2); + let inputs: Vec<_> = (0..3).map(|_| dummy_file_meta()).collect(); + let picker_output = PickerOutput { + outputs: inputs + .iter() + .map(|input| CompactionOutput { + output_level: 1, + inputs: vec![new_file_handle(input.clone())], + filter_deleted: false, + output_time_range: None, + }) + .collect(), + ..Default::default() + }; let cancel_handle = Arc::new(CancellationHandle::default()); - let call_idx = Arc::new(AtomicUsize::new(0)); + let (started_tx, mut started_rx) = mpsc::unbounded_channel(); let compactor = DefaultCompactor { - merger: BlockingMerger { - call_idx: call_idx.clone(), + merger: GatedEncodingMerger { + started: started_tx, + // Outputs are popped from the back, then joined in reverse order. + panic_input: panic_merge.then_some(inputs[1].file_id), }, cancel_handle: cancel_handle.clone(), uncommitted: None, }; - - let picker_output = PickerOutput { - outputs: vec![ - CompactionOutput { - output_level: 1, - inputs: vec![new_file_handle(dummy_file_meta())], - filter_deleted: false, - output_time_range: None, - }, - CompactionOutput { - output_level: 1, - inputs: vec![new_file_handle(dummy_file_meta())], - filter_deleted: false, - output_time_range: None, - }, - CompactionOutput { - output_level: 1, - inputs: vec![new_file_handle(dummy_file_meta())], - filter_deleted: false, - output_time_range: None, - }, - ], - expired_ssts: vec![], - time_window_size: 3600, - max_file_size: None, - }; - - let task = tokio::spawn(async move { - compactor - .merge_ssts(&compaction_region, picker_output) - .await + let memory = new_compaction_memory_manager(1024 * 1024); + let budget = memory.try_acquire(memory.limit_bytes()).unwrap(); + let mut merge = Box::pin(async { + let _budget = budget; + compactor.merge_ssts(®ion, picker_output).await }); - - sleep(Duration::from_millis(100)).await; + assert!(futures::poll!(merge.as_mut()).is_pending()); + let mut encoders = Vec::new(); + for _ in 0..2 { + encoders.push( + timeout(Duration::from_secs(10), started_rx.recv()) + .await + .unwrap() + .unwrap(), + ); + } cancel_handle.cancel(); + assert!(futures::poll!(merge.as_mut()).is_pending()); + assert!(memory.try_acquire(memory.limit_bytes()).is_none()); + assert!( + encoders + .iter() + .all(|encoder| encoder.batch.upgrade().is_some()) + ); + let retained_batches: Vec<_> = encoders + .iter() + .map(|encoder| encoder.batch.clone()) + .collect(); + let mut expected_inputs: Vec<_> = encoders.iter().map(|encoder| encoder.input_id).collect(); + let mut expected_outputs: Vec<_> = + encoders.iter().map(|encoder| encoder.output_id).collect(); + for encoder in encoders { + encoder.release.send(()).unwrap(); + } + let result = timeout(Duration::from_secs(10), merge).await.unwrap(); + assert!( + retained_batches + .iter() + .all(|batch| batch.upgrade().is_none()) + ); + assert!(memory.try_acquire(memory.limit_bytes()).is_some()); + assert!( + started_rx.try_recv().is_err(), + "cancelled work must not start another output" + ); + if panic_merge { + assert!(matches!(result, Err(error::Error::Join { error, .. }) if error.is_panic())); + } else { + let merged = result.unwrap(); + let mut removed: Vec<_> = merged + .files_to_remove + .iter() + .map(|file| file.file_id) + .collect(); + let mut added: Vec<_> = merged + .files_to_add + .iter() + .map(|file| file.file_id) + .collect(); + for ids in [ + &mut removed, + &mut added, + &mut expected_inputs, + &mut expected_outputs, + ] { + ids.sort_by(|a, b| a.as_bytes().cmp(b.as_bytes())); + } + assert_eq!(expected_inputs, removed); + assert_eq!(expected_outputs, added); + assert!(merged.files_to_add.iter().all(|file| file.file_size > 0)); + } + } - let merge_output = task - .await - .expect("merge_ssts should stop after cancellation") - .unwrap(); - - let started = call_idx.load(Ordering::SeqCst); - - assert!(merge_output.files_to_add.is_empty()); - assert!(merge_output.files_to_remove.is_empty()); - assert_eq!(started, 2); + #[tokio::test] + async fn test_merge_ssts_cancelled_before_start_skips_outputs() { + let region = new_test_compaction_region().await; + let merger = MockMerger::new(vec![Ok(vec![dummy_file_meta()])]); + let calls = merger.call_idx.clone(); + let compactor = DefaultCompactor::with_merger(merger); + compactor.cancel_handle.cancel(); + let picker_output = PickerOutput { + outputs: vec![CompactionOutput { + output_level: 1, + inputs: vec![new_file_handle(dummy_file_meta())], + filter_deleted: false, + output_time_range: None, + }], + ..Default::default() + }; + let merged = compactor.merge_ssts(®ion, picker_output).await.unwrap(); + assert!(merged.is_empty()); + assert_eq!(0, calls.load(Ordering::SeqCst)); } } diff --git a/src/mito2/src/compaction/reader.rs b/src/mito2/src/compaction/reader.rs index 159ebe2f22a..0968f0e7a60 100644 --- a/src/mito2/src/compaction/reader.rs +++ b/src/mito2/src/compaction/reader.rs @@ -36,6 +36,7 @@ use crate::read::flat_projection::FlatProjectionMapper; use crate::read::read_columns::ReadColumns; use crate::read::scan_region::{PredicateGroup, ScanInput}; use crate::read::seq_scan::SeqScan; +use crate::read::stream::spawn_compact_stream; use crate::region::options::MergeMode; use crate::sst::file::FileHandle; use crate::sst::parquet::reader::MetadataCacheMetrics; @@ -51,6 +52,8 @@ pub(crate) struct CompactionSstReaderBuilder<'a> { pub(crate) filter_deleted: bool, pub(crate) time_range: Option, pub(crate) merge_mode: MergeMode, + /// Max number of inputs decoded concurrently ahead of the merge. + pub(crate) read_parallelism: usize, } impl CompactionSstReaderBuilder<'_> { @@ -131,6 +134,7 @@ impl CompactionSstReaderBuilder<'_> { .with_json2_rewrite_targets(targets) .with_files(self.inputs.to_vec()) .with_compaction(true) + .with_compaction_read_parallelism(self.read_parallelism) .with_batch_size(batch_size) .with_append_mode(self.append_mode) // We use special cache strategy for compaction. @@ -180,6 +184,17 @@ impl CompactionSstReaderBuilder<'_> { } } +/// Number of merged batches buffered between the merge task and the SST writer. +const MERGED_BATCH_CHANNEL_SIZE: usize = 2; + +/// Reads and merges `source` on a separate task, so the SST writer can encode +/// merged batches while the next ones are being read and merged. +pub(crate) fn spawn_merge(source: FlatSource) -> FlatSource { + let schema = source.schema().clone(); + let stream = spawn_compact_stream(source.into_stream(), MERGED_BATCH_CHANNEL_SIZE, None); + FlatSource::new_stream(schema, stream) +} + /// Converts time range to predicates so that rows outside the range will be filtered. fn time_range_to_predicate( range: TimestampRange, @@ -304,6 +319,7 @@ mod tests { filter_deleted: true, time_range: None, merge_mode: MergeMode::LastRow, + read_parallelism: 1, } .build_flat_sst_reader() .await diff --git a/src/mito2/src/compaction/window.rs b/src/mito2/src/compaction/window.rs index ab7ac0a5f9a..b2c4ceaa12a 100644 --- a/src/mito2/src/compaction/window.rs +++ b/src/mito2/src/compaction/window.rs @@ -356,6 +356,7 @@ mod tests { auto_flush_interval: None, compaction: Default::default(), compaction_override: false, + compaction_output_parallelism: None, storage: None, append_mode: false, skip_wal: false, diff --git a/src/mito2/src/config.rs b/src/mito2/src/config.rs index 98b4ba4be42..67df74d7160 100644 --- a/src/mito2/src/config.rs +++ b/src/mito2/src/config.rs @@ -129,6 +129,10 @@ pub struct MitoConfig { pub experimental_compaction_memory_limit: MemoryLimit, /// Behavior when compaction cannot acquire memory from the budget. pub experimental_compaction_on_exhausted: OnExhaustedPolicy, + /// Default per-output parallelism for input decoding and SST column encoding, + /// bounded independently per stage. Values below 2 disable both. + /// Tables can override it with `compaction.experimental_output_parallelism`. + pub experimental_compaction_output_parallelism: usize, // Flush configs: /// Interval to auto flush a region if it has not flushed yet (default 10 min). @@ -247,6 +251,7 @@ impl Default for MitoConfig { max_background_purges: get_total_cpu_cores(), experimental_compaction_memory_limit: MemoryLimit::Unlimited, experimental_compaction_on_exhausted: OnExhaustedPolicy::default(), + experimental_compaction_output_parallelism: 1, auto_flush_interval: Duration::from_secs(10 * 60), global_write_buffer_size: ReadableSize::gb(1), global_write_buffer_reject_size: ReadableSize::gb(2), diff --git a/src/mito2/src/engine/alter_test.rs b/src/mito2/src/engine/alter_test.rs index aee05880125..757da327d36 100644 --- a/src/mito2/src/engine/alter_test.rs +++ b/src/mito2/src/engine/alter_test.rs @@ -2073,6 +2073,68 @@ async fn test_alter_twcs_option_enables_compaction_override() { assert!(compaction_override()); } +#[tokio::test] +async fn test_alter_compaction_output_parallelism_preserves_existing_snapshots() { + let mut env = TestEnv::new().await; + let engine = env.create_engine(MitoConfig::default()).await; + let region_id = RegionId::new(1, 1); + engine + .handle_request( + region_id, + RegionRequest::Create( + CreateRequestBuilder::new() + .insert_option("compaction.experimental_output_parallelism", "4") + .build(), + ), + ) + .await + .unwrap(); + let region = engine.get_region(region_id).unwrap(); + let initial_version = region.version(); + assert_eq!( + Some(4), + initial_version.options.compaction_output_parallelism + ); + for (kind, expected) in [ + ( + AlterKind::SetRegionOptions { + options: vec![SetRegionOption::CompactionOutputParallelism(Some(1))], + }, + Some(1), + ), + ( + AlterKind::UnsetRegionOptions { + keys: vec![UnsetRegionOption::CompactionOutputParallelism], + }, + None, + ), + ] { + engine + .handle_request(region_id, RegionRequest::Alter(RegionAlterRequest { kind })) + .await + .unwrap(); + let version = region.version(); + assert_eq!(expected, version.options.compaction_output_parallelism); + assert!(!version.options.compaction_override); + assert_eq!( + initial_version.options.compaction, + version.options.compaction + ); + assert_eq!( + Some(4), + initial_version.options.compaction_output_parallelism + ); + } + // Direct engine callers must get the same validation as SQL callers. + let invalid = RegionRequest::Alter(RegionAlterRequest { + kind: AlterKind::SetRegionOptions { + options: vec![SetRegionOption::CompactionOutputParallelism(Some(0))], + }, + }); + assert!(engine.handle_request(region_id, invalid).await.is_err()); + assert_eq!(None, region.version().options.compaction_output_parallelism); +} + #[tokio::test] async fn test_mixed_region_options_are_published_after_flush() { let mut env = TestEnv::new().await; @@ -2122,6 +2184,7 @@ async fn test_mixed_region_options_are_published_after_flush() { "2h".to_string(), ), SetRegionOption::MaxRowGroupRowCount(Some(1024)), + SetRegionOption::CompactionOutputParallelism(Some(4)), ], }, }), @@ -2135,6 +2198,7 @@ async fn test_mixed_region_options_are_published_after_flush() { assert_eq!(None, version.options.ttl); assert!(!version.options.compaction_override); assert_eq!(None, version.options.max_row_group_row_count); + assert_eq!(None, version.options.compaction_output_parallelism); listener.wake_flush(); alter_job.await.unwrap(); @@ -2143,6 +2207,7 @@ async fn test_mixed_region_options_are_published_after_flush() { assert_eq!(Some(Duration::from_secs(500).into()), version.options.ttl); assert!(version.options.compaction_override); assert_eq!(Some(1024), version.options.max_row_group_row_count); + assert_eq!(Some(4), version.options.compaction_output_parallelism); } #[tokio::test(flavor = "multi_thread")] diff --git a/src/mito2/src/engine/compaction_test.rs b/src/mito2/src/engine/compaction_test.rs index b6594ddc136..ec9ddc5cf12 100644 --- a/src/mito2/src/engine/compaction_test.rs +++ b/src/mito2/src/engine/compaction_test.rs @@ -2890,16 +2890,29 @@ async fn test_worker_shutdown_fails_picking_waiter() { #[tokio::test] async fn test_compaction_region() { - test_compaction_region_with_format(false).await; - test_compaction_region_with_format(true).await; + for (engine_parallelism, table_parallelism) in [ + (1, None), + (4, None), + (1, Some(4)), + (4, Some(1)), + (1, Some(usize::MAX)), + ] { + test_compaction_region_with_format(false, engine_parallelism, table_parallelism).await; + test_compaction_region_with_format(true, engine_parallelism, table_parallelism).await; + } } -async fn test_compaction_region_with_format(flat_format: bool) { +async fn test_compaction_region_with_format( + flat_format: bool, + output_parallelism: usize, + table_parallelism: Option, +) { common_telemetry::init_default_ut_logging(); let mut env = TestEnv::new().await; let engine = env .create_engine(MitoConfig { default_flat_format: flat_format, + experimental_compaction_output_parallelism: output_parallelism, ..Default::default() }) .await; @@ -2916,10 +2929,16 @@ async fn test_compaction_region_with_format(flat_format: bool) { ) .await; - let request = CreateRequestBuilder::new() + let mut builder = CreateRequestBuilder::new() .insert_option("compaction.type", "twcs") - .insert_option("experimental_sst_float_field_encoding", "byte_stream_split") - .build(); + .insert_option("experimental_sst_float_field_encoding", "byte_stream_split"); + if let Some(parallelism) = table_parallelism { + builder = builder.insert_option( + "compaction.experimental_output_parallelism", + ¶llelism.to_string(), + ); + } + let request = builder.build(); let column_schemas = request .column_metadatas diff --git a/src/mito2/src/read.rs b/src/mito2/src/read.rs index 65d356d79f2..4d742121a4f 100644 --- a/src/mito2/src/read.rs +++ b/src/mito2/src/read.rs @@ -1013,6 +1013,14 @@ impl FlatSource { self.inner.next_batch().await } + /// Converts the source into a stream of its batches. + pub(crate) fn into_stream(self) -> BoxedRecordBatchStream { + match self.inner { + FlatSourceInner::Iter(iter) => Box::pin(futures::stream::iter(iter)), + FlatSourceInner::Stream(stream) => stream, + } + } + #[cfg(test)] pub(crate) fn take_iter(self) -> BoxedRecordBatchIterator { match self.inner { diff --git a/src/mito2/src/read/scan_region.rs b/src/mito2/src/read/scan_region.rs index da4e1c82bf9..31db8581b27 100644 --- a/src/mito2/src/read/scan_region.rs +++ b/src/mito2/src/read/scan_region.rs @@ -971,6 +971,9 @@ pub struct ScanInput { pub(crate) sequence_range: Option, /// Whether this scan is for compaction. pub(crate) compaction: bool, + /// Max number of compaction inputs decoded concurrently ahead of the merge. + /// Values below 2 decode inputs on the merging task. + pub(crate) compaction_read_parallelism: usize, /// Compaction-only JSON2 physical rewrite targets. json2_rewrite_targets: Json2RewriteTargets, /// Counters that should receive query-load metrics. @@ -1036,6 +1039,7 @@ impl ScanInput { snapshot_sequence: None, sequence_range: None, compaction: false, + compaction_read_parallelism: 1, json2_rewrite_targets: Arc::default(), query_stat_counters: None, #[cfg(feature = "enterprise")] @@ -1393,6 +1397,13 @@ impl ScanInputBuilder { self } + /// Sets the max number of compaction inputs decoded concurrently. + #[must_use] + pub(crate) fn with_compaction_read_parallelism(mut self, parallelism: usize) -> Self { + self.input.compaction_read_parallelism = parallelism; + self + } + /// Sets compaction-only JSON2 physical rewrite targets. #[must_use] pub(crate) fn with_json2_rewrite_targets(mut self, targets: Json2RewriteTargets) -> Self { diff --git a/src/mito2/src/read/seq_scan.rs b/src/mito2/src/read/seq_scan.rs index b57622e07e7..d8759ffaace 100644 --- a/src/mito2/src/read/seq_scan.rs +++ b/src/mito2/src/read/seq_scan.rs @@ -49,7 +49,7 @@ use crate::read::scan_util::{ SplitRecordBatchStream, compute_parallel_channel_size, prepare_flat_mem_stream, scan_flat_file_ranges, scan_flat_mem_ranges, should_split_flat_batches_for_merge, }; -use crate::read::stream::{ConvertBatchStream, ScanBatch, ScanBatchStream}; +use crate::read::stream::{ConvertBatchStream, ScanBatch, ScanBatchStream, spawn_compact_stream}; use crate::read::{BoxedRecordBatchStream, ScannerMetrics, scan_util}; use crate::region::options::MergeMode; use crate::sst::parquet::DEFAULT_READ_BATCH_SIZE; @@ -198,15 +198,19 @@ impl SeqScan { partition_ranges.len(), sources.len() ); - Self::build_flat_reader_from_sources( - stream_ctx, - sources, - None, - None, - false, - compute_parallel_channel_size(DEFAULT_READ_BATCH_SIZE), - ) - .await + let channel_size = compute_parallel_channel_size(DEFAULT_READ_BATCH_SIZE); + let parallelism = stream_ctx.input.compaction_read_parallelism; + if parallelism > 1 && sources.len() > 1 { + // More permits than inputs cannot increase concurrency. Clamping also + // keeps oversized user settings below the semaphore's permit limit. + let semaphore = Arc::new(Semaphore::new(parallelism.min(sources.len()))); + sources = sources + .into_iter() + .map(|source| spawn_compact_stream(source, channel_size, Some(semaphore.clone()))) + .collect(); + } + Self::build_flat_reader_from_sources(stream_ctx, sources, None, None, false, channel_size) + .await } /// Builds a flat reader to read sources that returns RecordBatch. diff --git a/src/mito2/src/read/stream.rs b/src/mito2/src/read/stream.rs index 0c882d374d2..19ed37cf65f 100644 --- a/src/mito2/src/read/stream.rs +++ b/src/mito2/src/read/stream.rs @@ -18,15 +18,19 @@ use std::sync::Arc; use std::task::{Context, Poll}; use std::time::Instant; +use async_stream::try_stream; use common_error::ext::BoxedError; use common_recordbatch::error::ExternalSnafu; use common_recordbatch::{DfRecordBatch, RecordBatch}; use futures::stream::BoxStream; use futures::{Stream, StreamExt}; use snafu::ResultExt; +use tokio::sync::{Semaphore, mpsc}; +use tokio::task::AbortHandle; use crate::cache::CacheStrategy; -use crate::error::Result; +use crate::error::{JoinSnafu, Result, UnexpectedSnafu}; +use crate::read::BoxedRecordBatchStream; use crate::read::flat_projection::FlatProjectionMapper; use crate::read::scan_util::PartitionMetrics; use crate::read::series_scan::SeriesBatch; @@ -122,3 +126,272 @@ impl Stream for ConvertBatchStream { Poll::Ready(Some(record_batch)) } } + +/// Polls `stream` on its own compaction runtime task, buffering up to +/// `capacity` batches ahead of the consumer. While polling, the task holds a +/// permit of `semaphore` if one is given. +/// +/// The channel also closes if the task panics, so the returned stream checks +/// how the task ended before reporting the end of the input. Dropping the +/// returned stream aborts the task. +pub(crate) fn spawn_compact_stream( + mut stream: BoxedRecordBatchStream, + capacity: usize, + semaphore: Option>, +) -> BoxedRecordBatchStream { + let (sender, mut receiver) = mpsc::channel(capacity); + let task = common_runtime::spawn_compact(async move { + loop { + // Releases the permit before sending, so a full channel never holds it. + let permit = match &semaphore { + Some(semaphore) => match semaphore.acquire().await { + Ok(permit) => Some(permit), + Err(e) => { + let reason = format!("Failed to acquire a stream permit: {e}"); + let _ = sender.send(UnexpectedSnafu { reason }.fail()).await; + break; + } + }, + None => None, + }; + let next = stream.next().await; + drop(permit); + let Some(result) = next else { + break; + }; + // Stops after forwarding the first error or once the consumer is gone. + let failed = result.is_err(); + if sender.send(result).await.is_err() || failed { + break; + } + } + }); + // Own the guard before the lazy stream can be dropped without being polled. + let abort = AbortOnDrop(task.abort_handle()); + Box::pin(try_stream! { + let _abort = abort; + while let Some(batch) = receiver.recv().await { + yield batch?; + } + task.await.context(JoinSnafu)?; + }) +} + +/// Aborts a task when its consumer is dropped before the end. +struct AbortOnDrop(AbortHandle); + +impl Drop for AbortOnDrop { + fn drop(&mut self) { + self.0.abort(); + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::task::Poll; + use std::time::Duration; + + use datatypes::arrow::array::{ArrayRef, AsArray, Int64Array}; + use datatypes::arrow::datatypes::{DataType, Field, Int64Type, Schema}; + use datatypes::arrow::record_batch::RecordBatch; + use tokio::sync::oneshot; + use tokio::time::timeout; + + use super::*; + use crate::error::{Error, InvalidRecordBatchSnafu}; + use crate::read::flat_merge::FlatMergeReader; + + /// Keeps the producer pending and exposes when its source is polled and dropped. + struct PendingStream { + polled: Option>, + dropped: Option>, + } + + impl PendingStream { + fn new() -> (Self, oneshot::Receiver<()>, oneshot::Receiver<()>) { + let (polled_tx, polled_rx) = oneshot::channel(); + let (dropped_tx, dropped_rx) = oneshot::channel(); + ( + Self { + polled: Some(polled_tx), + dropped: Some(dropped_tx), + }, + polled_rx, + dropped_rx, + ) + } + } + + impl Stream for PendingStream { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + if let Some(polled) = self.polled.take() { + let _ = polled.send(()); + } + Poll::Pending + } + } + + impl Drop for PendingStream { + fn drop(&mut self) { + if let Some(dropped) = self.dropped.take() { + let _ = dropped.send(()); + } + } + } + + async fn wait_for_signal(signal: oneshot::Receiver<()>) { + // Notifications establish ordering; the timeout only bounds a broken test. + timeout(Duration::from_secs(10), signal) + .await + .expect("timed out waiting for the source") + .expect("source dropped without sending the signal"); + } + + fn batch(value: i64) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)])); + let column = Arc::new(Int64Array::from(vec![value])) as ArrayRef; + RecordBatch::try_new(schema, vec![column]).unwrap() + } + + fn values_stream(values: std::ops::Range) -> BoxedRecordBatchStream { + Box::pin(futures::stream::iter(values.map(|value| Ok(batch(value))))) + } + + /// Drains `stream`, returning the values read before it ended or failed. + async fn drain(mut stream: BoxedRecordBatchStream) -> (Vec, Option) { + let mut values = Vec::new(); + while let Some(result) = stream.next().await { + match result { + Ok(batch) => values.push(batch.column(0).as_primitive::().value(0)), + Err(e) => return (values, Some(e)), + } + } + (values, None) + } + + #[tokio::test] + async fn test_spawn_compact_stream_forwards_batches_in_order() { + let (values, error) = drain(spawn_compact_stream(values_stream(0..5), 2, None)).await; + + assert_eq!(vec![0, 1, 2, 3, 4], values); + assert!(error.is_none()); + } + + #[tokio::test] + async fn test_spawn_compact_stream_stops_at_first_error() { + let stream = futures::stream::iter(vec![ + Ok(batch(0)), + InvalidRecordBatchSnafu { reason: "boom" }.fail(), + Ok(batch(2)), + ]); + + let (values, error) = drain(spawn_compact_stream(Box::pin(stream), 2, None)).await; + + assert_eq!(vec![0], values); + assert!( + matches!(error, Some(Error::InvalidRecordBatch { .. })), + "{error:?}" + ); + } + + #[tokio::test] + async fn test_spawn_compact_stream_reports_panic_instead_of_end() { + let mut polled = false; + let stream = futures::stream::poll_fn(move |_| { + if std::mem::replace(&mut polled, true) { + panic!("source panicked"); + } + Poll::Ready(Some(Ok(batch(0)))) + }); + + let (values, error) = drain(spawn_compact_stream(Box::pin(stream), 2, None)).await; + + // A truncated input must never look like a complete one. + assert_eq!(vec![0], values); + assert!(matches!(error, Some(Error::Join { .. })), "{error:?}"); + } + + #[rstest::rstest] + #[case::before_first_poll(false)] + #[case::after_first_poll(true)] + #[tokio::test] + async fn test_dropped_consumer_cancels_pending_source(#[case] poll_consumer: bool) { + let (source, polled, dropped) = PendingStream::new(); + let semaphore = Arc::new(Semaphore::new(1)); + let mut consumer = spawn_compact_stream(Box::pin(source), 1, Some(semaphore.clone())); + wait_for_signal(polled).await; + assert_eq!(0, semaphore.available_permits()); + if poll_consumer { + assert!(futures::poll!(consumer.next()).is_pending()); + } + + // A guard created inside the lazy stream misses the before-first-poll case. + drop(consumer); + wait_for_signal(dropped).await; + let _permit = timeout(Duration::from_secs(10), semaphore.acquire()) + .await + .expect("producer retained its permit") + .unwrap(); + } + + #[tokio::test] + async fn test_merge_initialization_error_cancels_unpolled_source() { + let (source, polled, dropped) = PendingStream::new(); + let consumer = spawn_compact_stream(Box::pin(source), 1, None); + wait_for_signal(polled).await; + let error_stream = Box::pin(futures::stream::iter(vec![ + InvalidRecordBatchSnafu { reason: "boom" }.fail(), + ])); + let sources = vec![spawn_compact_stream(error_stream, 1, None), consumer]; + + let result = FlatMergeReader::new(batch(0).schema(), sources, 1024, None).await; + + assert!(matches!(result, Err(Error::InvalidRecordBatch { .. }))); + wait_for_signal(dropped).await; + } + + #[tokio::test] + async fn test_cancelled_merge_initialization_drops_all_sources() { + let (first, first_polled, first_dropped) = PendingStream::new(); + let (second, second_polled, second_dropped) = PendingStream::new(); + let sources = vec![ + spawn_compact_stream(Box::pin(first), 1, None), + spawn_compact_stream(Box::pin(second), 1, None), + ]; + wait_for_signal(first_polled).await; + wait_for_signal(second_polled).await; + let mut initialization = + Box::pin(FlatMergeReader::new(batch(0).schema(), sources, 1024, None)); + assert!(futures::poll!(&mut initialization).is_pending()); + + // Only the first consumer has been polled by the sequential initialization. + drop(initialization); + + wait_for_signal(first_dropped).await; + wait_for_signal(second_dropped).await; + } + + #[tokio::test] + async fn test_spawn_compact_streams_share_semaphore() { + let semaphore = Arc::new(Semaphore::new(1)); + let streams: Vec<_> = (0..3) + .map(|i| { + spawn_compact_stream( + values_stream(i * 10..i * 10 + 3), + 1, + Some(semaphore.clone()), + ) + }) + .collect(); + + for (i, stream) in streams.into_iter().enumerate() { + let start = i as i64 * 10; + let (values, error) = drain(stream).await; + assert_eq!((start..start + 3).collect::>(), values); + assert!(error.is_none()); + } + } +} diff --git a/src/mito2/src/region/options.rs b/src/mito2/src/region/options.rs index 42f4adcd04a..58fdec15ad5 100644 --- a/src/mito2/src/region/options.rs +++ b/src/mito2/src/region/options.rs @@ -95,6 +95,11 @@ pub struct RegionOptions { /// Compaction options. pub compaction: CompactionOptions, pub compaction_override: bool, + /// Per-output decode and encode parallelism. Falls back to the engine config + /// when unset; 1 disables parallelism without disabling the merge pipeline. + /// Independent of TWCS options and their database-level inheritance. + #[serde(skip_serializing_if = "Option::is_none")] + pub compaction_output_parallelism: Option, /// Custom storage. Uses default storage if it is `None`. pub storage: Option, /// If append mode is enabled, the region keeps duplicate rows. @@ -162,6 +167,14 @@ impl RegionOptions { } ); } + if let Some(parallelism) = self.compaction_output_parallelism { + ensure!( + parallelism > 0, + InvalidRegionOptionsSnafu { + reason: "compaction.experimental_output_parallelism must be greater than 0", + } + ); + } if let Some(row_count) = self.max_row_group_row_count { ensure!( row_count > 0 && row_count <= MAX_ROW_GROUP_ROW_COUNT_LIMIT, @@ -243,6 +256,11 @@ impl RegionOptions { self.auto_flush_interval.unwrap_or(default) } + /// Returns the per-output compaction parallelism, falling back to the engine default. + pub fn compaction_output_parallelism_or(&self, default: usize) -> usize { + self.compaction_output_parallelism.unwrap_or(default) + } + /// Returns the `primary_key_encoding` if it is set, otherwise returns the default [`PrimaryKeyEncoding`]. pub fn primary_key_encoding(&self) -> PrimaryKeyEncoding { self.primary_key_encoding.unwrap_or_default() @@ -346,6 +364,7 @@ impl RegionOptions { auto_flush_interval: options.auto_flush_interval, compaction, compaction_override, + compaction_output_parallelism: options.compaction_output_parallelism, storage: options.storage, append_mode: options.append_mode, skip_wal: options.skip_wal, @@ -495,6 +514,9 @@ struct RegionOptionsWithoutEnum { ttl: Option, #[serde(with = "humantime_serde")] auto_flush_interval: Option, + #[serde(rename = "compaction.experimental_output_parallelism")] + #[serde_as(as = "NoneAsEmptyString")] + compaction_output_parallelism: Option, storage: Option, #[serde_as(as = "DisplayFromStr")] append_mode: bool, @@ -520,6 +542,7 @@ impl Default for RegionOptionsWithoutEnum { write_buffer_size: options.write_buffer_size, ttl: options.ttl, auto_flush_interval: options.auto_flush_interval, + compaction_output_parallelism: options.compaction_output_parallelism, storage: options.storage, append_mode: options.append_mode, skip_wal: options.skip_wal, @@ -1146,6 +1169,7 @@ mod tests { fallback_to_local: true, }), compaction_override: true, + compaction_output_parallelism: None, storage: Some("S3".to_string()), append_mode: false, skip_wal: false, @@ -1226,6 +1250,7 @@ mod tests { max_row_group_row_count: None, primary_key_encoding: None, write_buffer_size: Some(ReadableSize::mb(128)), + compaction_output_parallelism: Some(4), preserve_row_sequence: true, float_field_encoding: FloatFieldEncoding::default(), }; @@ -1238,6 +1263,8 @@ mod tests { let old_region_options_json_str = r#"{"ttl":null}"#; let got: RegionOptions = serde_json::from_str(old_region_options_json_str).unwrap(); assert_eq!(None, got.write_buffer_size); + assert_eq!(None, got.compaction_output_parallelism); + assert_eq!(4, got.compaction_output_parallelism_or(4)); assert!(!got.preserve_row_sequence); assert_eq!(FloatFieldEncoding::Default, got.float_field_encoding); let CompactionOptions::Twcs(twcs) = got.compaction; @@ -1305,12 +1332,50 @@ mod tests { max_row_group_row_count: None, primary_key_encoding: None, write_buffer_size: None, + compaction_output_parallelism: None, preserve_row_sequence: false, float_field_encoding: FloatFieldEncoding::default(), }; assert_eq!(options, got); } + #[test] + fn test_compaction_output_parallelism_preserves_strategy_inheritance() { + for parallelism in [None, Some(1), Some(4)] { + let mut map = make_map(&[]); + if let Some(parallelism) = parallelism { + map.insert( + "compaction.experimental_output_parallelism".to_string(), + parallelism.to_string(), + ); + } + let options = RegionOptions::try_from_options(RegionId::new(0, 0), &map).unwrap(); + assert_eq!(parallelism, options.compaction_output_parallelism); + assert_eq!( + parallelism.unwrap_or(8), + options.compaction_output_parallelism_or(8) + ); + assert!(!options.compaction_override); + + map.insert("compaction.type".to_string(), "twcs".to_string()); + map.insert("compaction.twcs.time_window".to_string(), "2h".to_string()); + let options = RegionOptions::try_from_options(RegionId::new(0, 0), &map).unwrap(); + assert_eq!(parallelism, options.compaction_output_parallelism); + assert_eq!( + Some(Duration::from_secs(7200)), + options.compaction.time_window() + ); + } + } + + #[test] + fn test_invalid_compaction_output_parallelism() { + for parallelism in ["0", "-1", "1.5", "invalid", "18446744073709551616"] { + let map = make_map(&[("compaction.experimental_output_parallelism", parallelism)]); + assert!(RegionOptions::try_from_options(RegionId::new(0, 0), &map).is_err()); + } + } + #[test] fn test_max_row_group_row_count() { // Default falls back to DEFAULT_ROW_GROUP_SIZE. diff --git a/src/mito2/src/sst/parquet.rs b/src/mito2/src/sst/parquet.rs index b1b459a595a..fec4ede14df 100644 --- a/src/mito2/src/sst/parquet.rs +++ b/src/mito2/src/sst/parquet.rs @@ -35,6 +35,7 @@ use crate::error::{CompressionOptionsSnafu, Result, WriteParquetSnafu}; use crate::sst::DEFAULT_WRITE_BUFFER_SIZE; use crate::sst::file::FileTimeRange; use crate::sst::index::IndexOutput; +use crate::sst::parquet::parallel_writer::EncodeRuntime; pub mod file_range; pub mod flat_format; @@ -44,6 +45,7 @@ pub(crate) mod index_reader; pub(crate) mod index_writer; pub(crate) mod json_align; pub mod metadata; +pub mod parallel_writer; pub mod prefilter; pub mod push_decoder; pub mod read_columns; @@ -190,6 +192,11 @@ pub struct WriteOptions { pub max_file_size: Option, /// Encoding policy for direct floating-point field columns. pub float_field_encoding: FloatFieldEncoding, + /// Maximum number of blocking tasks that encode the columns of a row group. + /// Values below 2 encode on the writing task. + pub encode_parallelism: usize, + /// Runtime whose blocking pool runs the parallel column encoding. + pub encode_runtime: EncodeRuntime, } impl Default for WriteOptions { @@ -199,6 +206,8 @@ impl Default for WriteOptions { row_group_size: DEFAULT_ROW_GROUP_SIZE, max_file_size: None, float_field_encoding: FloatFieldEncoding::default(), + encode_parallelism: 1, + encode_runtime: EncodeRuntime::default(), } } } @@ -492,7 +501,8 @@ mod tests { } } - struct NoopIndexBuilder; + /// Builds an indexer with no enabled indexes for SST writer tests. + pub(crate) struct NoopIndexBuilder; #[async_trait::async_trait] impl IndexerBuilder for NoopIndexBuilder { @@ -506,6 +516,54 @@ mod tests { } } + #[tokio::test] + async fn test_parallel_encoding_writes_identical_file() { + let mut env = TestEnv::new().await; + let object_store = env.init_object_store_manager(); + let metadata = Arc::new(sst_region_metadata()); + let mut contents = Vec::new(); + for encode_parallelism in [1, 3] { + let handle = sst_file_handle(0, 1000); + let file_path = FixedPathProvider { + region_file_id: handle.file_id(), + }; + let path = file_path.build_sst_file_path(handle.file_id()); + let source = new_flat_source_from_record_batches(vec![ + new_record_batch_by_range(&["a", "d"], 0, 60), + new_record_batch_by_range(&["b", "f"], 0, 40), + new_record_batch_by_range(&["b", "h"], 100, 200), + ]); + let write_opts = WriteOptions { + row_group_size: 50, + encode_parallelism, + ..Default::default() + }; + + let mut metrics = Metrics::new(WriteType::Flush); + let mut writer = ParquetWriter::new_with_object_store( + object_store.clone(), + metadata.clone(), + IndexConfig::default(), + NoopIndexBuilder, + file_path, + &mut metrics, + ) + .await; + let info = writer + .write_all_flat(source, None, &write_opts) + .await + .unwrap() + .remove(0); + assert_eq!(200, info.num_rows); + assert_eq!(4, info.num_row_groups); + contents.push(object_store.read(&path).await.unwrap().to_vec()); + } + assert!( + contents[0] == contents[1], + "parallel encoding changed the SST" + ); + } + #[tokio::test] async fn test_write_read() { let mut env = TestEnv::new().await; diff --git a/src/mito2/src/sst/parquet/parallel_writer.rs b/src/mito2/src/sst/parquet/parallel_writer.rs new file mode 100644 index 00000000000..d257574cb79 --- /dev/null +++ b/src/mito2/src/sst/parquet/parallel_writer.rs @@ -0,0 +1,648 @@ +// 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. + +//! Arrow to Parquet writer that encodes the columns of a row group in parallel. + +use std::mem; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use bytes::Bytes; +use common_runtime::JoinHandle; +use datatypes::arrow::datatypes::SchemaRef; +use datatypes::arrow::record_batch::RecordBatch; +use parquet::arrow::ArrowWriter; +use parquet::arrow::arrow_writer::{ + ArrowColumnChunk, ArrowColumnWriter, ArrowRowGroupWriterFactory, compute_leaves, +}; +use parquet::arrow::async_writer::AsyncFileWriter; +use parquet::errors::{ParquetError, Result}; +use parquet::file::metadata::ParquetMetaData; +use parquet::file::properties::WriterProperties; +use parquet::file::writer::SerializedFileWriter; + +/// Runtime whose blocking pool encodes the columns of a writer. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub enum EncodeRuntime { + /// The global runtime. + #[default] + Global, + /// The compaction runtime. + Compact, +} + +impl EncodeRuntime { + fn spawn(self, f: F) -> JoinHandle + where + F: FnOnce() -> R + Send + 'static, + R: Send + 'static, + { + match self { + EncodeRuntime::Global => common_runtime::spawn_blocking_global(f), + EncodeRuntime::Compact => common_runtime::spawn_blocking_compact(f), + } + } +} + +/// Column writers of one top-level field: its index and its leaf writers. +type FieldWriters = (usize, Vec<(usize, ArrowColumnWriter)>); + +/// Output of one encode task: column chunks by leaf index, and the time spent +/// on each field. +type EncodedFields = (Vec<(usize, ArrowColumnChunk)>, Vec<(usize, Duration)>); + +/// Writes Parquet like [`parquet::arrow::AsyncArrowWriter`], but encodes the +/// columns of each row group on up to `parallelism` blocking tasks. +/// +/// A row group encodes in the background while the next one is buffered, so at +/// most two row groups of input are held in memory. Every column writer sees the +/// same input as with [`ArrowWriter`], so the output is byte-for-byte identical. +/// The bytes of a row group reach the async writer when the next row group +/// starts or on [`Self::flush`], one row group later than with `AsyncArrowWriter`. +pub struct ParallelArrowWriter { + /// Assembles row groups and the footer into an in-memory buffer. + file_writer: SerializedFileWriter>, + /// Creates the column writers of each row group. + factory: Arc, + arrow_schema: SchemaRef, + /// Top-level field of each leaf column. + leaf_fields: Vec, + /// Encode time of each field in the last row group, used to balance tasks. + field_costs: Vec, + async_writer: W, + max_row_group_rows: usize, + parallelism: usize, + runtime: EncodeRuntime, + /// Slices of input buffered for the next row group. + buffered: Vec, + buffered_rows: usize, + /// Tasks encoding the previous row group. + encoding: Vec>>, + num_row_groups: usize, +} + +impl ParallelArrowWriter { + /// Returns whether the writer can reproduce [`ArrowWriter`]'s row groups for + /// `props`, which holds only when row groups are limited by row count alone. + pub fn supports(props: &WriterProperties) -> bool { + props.max_row_group_row_count().is_some() && props.max_row_group_bytes().is_none() + } + + /// Creates a writer. `props` must be [supported](Self::supports). + pub fn try_new( + async_writer: W, + arrow_schema: SchemaRef, + props: WriterProperties, + parallelism: usize, + runtime: EncodeRuntime, + ) -> Result { + let max_row_group_rows = match props.max_row_group_row_count() { + Some(rows) if Self::supports(&props) => rows.max(1), + _ => { + return Err(ParquetError::General( + "parallel writer requires a row-count-only row group limit".to_string(), + )); + } + }; + let (file_writer, factory) = + ArrowWriter::try_new(Vec::new(), arrow_schema.clone(), Some(props))? + .into_serialized_writer()?; + let schema = file_writer.schema_descr(); + let leaf_fields = (0..schema.num_columns()) + .map(|leaf| schema.get_column_root_idx(leaf)) + .collect(); + let num_fields = arrow_schema.fields().len(); + Ok(Self { + file_writer, + factory: Arc::new(factory), + arrow_schema, + leaf_fields, + field_costs: vec![Duration::ZERO; num_fields], + async_writer, + max_row_group_rows, + parallelism: parallelism.max(1), + runtime, + buffered: Vec::new(), + buffered_rows: 0, + encoding: Vec::new(), + num_row_groups: 0, + }) + } + + /// Buffers `batch`, starting to encode each row group once it is full. + pub async fn write(&mut self, batch: &RecordBatch) -> Result<()> { + let mut offset = 0; + while offset < batch.num_rows() { + let len = (self.max_row_group_rows - self.buffered_rows).min(batch.num_rows() - offset); + self.buffered.push(batch.slice(offset, len)); + self.buffered_rows += len; + offset += len; + if self.buffered_rows == self.max_row_group_rows { + self.encode_buffered().await?; + } + } + Ok(()) + } + + /// Writes all buffered rows as a row group, like [`ArrowWriter::flush`]. + pub async fn flush(&mut self) -> Result<()> { + if self.buffered_rows > 0 { + self.encode_buffered().await?; + } + self.finish_encoding().await + } + + /// Writes the remaining rows and the footer, then completes the async writer. + pub async fn close(mut self) -> Result { + self.flush().await?; + let metadata = self.file_writer.finish()?; + self.write_buffered_bytes().await?; + self.async_writer.complete().await?; + Ok(metadata) + } + + /// Discards buffered input and waits for encoders without writing a failed SST. + pub(crate) async fn abort(mut self) { + self.buffered.clear(); + abort_encoding_tasks(&mut self.encoding).await; + } + + /// Installs controlled tasks to test SST write-error cleanup independently of codec timing. + #[cfg(test)] + pub(crate) fn set_encoding_tasks_for_test( + &mut self, + tasks: Vec>>, + ) { + assert!(self.encoding.is_empty()); + self.encoding = tasks; + } + + /// Starts encoding the buffered rows once the previous row group is written. + async fn encode_buffered(&mut self) -> Result<()> { + self.finish_encoding().await?; + + let batches: Arc<[RecordBatch]> = mem::take(&mut self.buffered).into(); + self.buffered_rows = 0; + let writers = self.factory.create_column_writers(self.num_row_groups)?; + self.num_row_groups += 1; + + // One array produces all leaves of a field, so a field is never split. + let mut fields: Vec = (0..self.arrow_schema.fields().len()) + .map(|field| (field, Vec::new())) + .collect(); + for (leaf, writer) in writers.into_iter().enumerate() { + fields[self.leaf_fields[leaf]].1.push((leaf, writer)); + } + self.encoding = assign_fields(fields, &self.field_costs, self.parallelism) + .into_iter() + .map(|fields| { + let batches = batches.clone(); + let schema = self.arrow_schema.clone(); + self.runtime + .spawn(move || encode_fields(&schema, &batches, fields)) + }) + .collect(); + Ok(()) + } + + /// Waits for the row group being encoded and appends it to the file. + async fn finish_encoding(&mut self) -> Result<()> { + if self.encoding.is_empty() { + return Ok(()); + } + let mut chunks = Vec::with_capacity(self.leaf_fields.len()); + let mut tasks = mem::take(&mut self.encoding).into_iter(); + while let Some(task) = tasks.next() { + let encoded = task + .await + .map_err(|e| ParquetError::External(Box::new(e))) + .and_then(|encoded| encoded); + let (task_chunks, costs) = match encoded { + Ok(encoded) => encoded, + Err(error) => { + abort_encoding_tasks(tasks.as_mut_slice()).await; + return Err(error); + } + }; + chunks.extend(task_chunks); + for (field, cost) in costs { + self.field_costs[field] = cost; + } + } + chunks.sort_unstable_by_key(|(leaf, _)| *leaf); + let mut row_group = self.file_writer.next_row_group()?; + for (_, chunk) in chunks { + chunk.append_to_row_group(&mut row_group)?; + } + row_group.close()?; + self.write_buffered_bytes().await + } + + /// Moves the bytes assembled so far to the async writer. + async fn write_buffered_bytes(&mut self) -> Result<()> { + let bytes = mem::take(self.file_writer.inner_mut()); + if bytes.is_empty() { + return Ok(()); + } + self.async_writer.write(Bytes::from(bytes)).await + } +} + +/// Aborts queued blocking work and joins running work before its budget is released. +async fn abort_encoding_tasks(tasks: &mut [JoinHandle>]) { + for task in tasks.iter() { + task.abort(); + } + for task in tasks { + // Cleanup failures must not replace the original encoding or source error. + let _ = task.await; + } +} + +/// Splits `fields` into at most `parallelism` tasks, giving the most expensive +/// field to the least loaded task first. Fields cost nothing before any row +/// group is measured, which spreads them evenly. +fn assign_fields( + mut fields: Vec, + costs: &[Duration], + parallelism: usize, +) -> Vec> { + let num_tasks = parallelism.min(fields.len()).max(1); + fields.sort_by_key(|(field, _)| std::cmp::Reverse(costs[*field])); + let mut tasks: Vec<(Duration, Vec)> = (0..num_tasks) + .map(|_| (Duration::ZERO, Vec::new())) + .collect(); + for field in fields { + let key = |(load, fields): &(Duration, Vec)| (*load, fields.len()); + let mut lightest = 0; + for task in 1..tasks.len() { + if key(&tasks[task]) < key(&tasks[lightest]) { + lightest = task; + } + } + tasks[lightest].0 += costs[field.0]; + tasks[lightest].1.push(field); + } + tasks + .into_iter() + .map(|(_, fields)| fields) + .filter(|fields| !fields.is_empty()) + .collect() +} + +/// Encodes `fields` of all `batches` and closes their column writers. +fn encode_fields( + schema: &SchemaRef, + batches: &[RecordBatch], + fields: Vec, +) -> Result { + let mut chunks = Vec::new(); + let mut costs = Vec::with_capacity(fields.len()); + for (field, mut writers) in fields { + let start = Instant::now(); + for batch in batches { + let leaves = compute_leaves(schema.field(field), batch.column(field))?; + if leaves.len() != writers.len() { + return Err(ParquetError::General(format!( + "field {} has {} leaf arrays but {} column writers", + schema.field(field).name(), + leaves.len(), + writers.len() + ))); + } + for (leaf, (_, writer)) in leaves.iter().zip(writers.iter_mut()) { + writer.write(leaf)?; + } + } + for (leaf, writer) in writers { + chunks.push((leaf, writer.close()?)); + } + costs.push((field, start.elapsed())); + } + Ok((chunks, costs)) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::sync::atomic::{AtomicBool, Ordering}; + + use common_runtime::runtime::{BuilderBuild, RuntimeTrait}; + use datatypes::arrow::array::{ + ArrayRef, DictionaryArray, Float64Array, Int32Array, Int64Array, StringArray, StructArray, + }; + use datatypes::arrow::datatypes::{DataType, Field, Fields, Schema, UInt32Type}; + use parquet::basic::{Compression, ZstdLevel}; + use tokio::sync::oneshot; + use tokio::time::timeout; + + use super::*; + use crate::compaction::memory_manager::new_compaction_memory_manager; + + fn schema() -> SchemaRef { + let nested = Fields::from(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Utf8, true), + ]); + Arc::new(Schema::new(vec![ + Field::new("int", DataType::Int64, false), + Field::new("float", DataType::Float64, true), + Field::new( + "tag", + DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)), + false, + ), + Field::new("nested", DataType::Struct(nested), true), + ])) + } + + fn batch(start: usize, len: usize) -> RecordBatch { + let rows = start..start + len; + let ints = Int64Array::from_iter_values(rows.clone().map(|i| i as i64 * 7)); + let floats = Float64Array::from_iter( + rows.clone() + .map(|i| (i % 5 != 0).then_some((i % 13) as f64 / 3.0)), + ); + let tags: DictionaryArray = rows + .clone() + .map(|i| format!("tag-{}", i % 3)) + .collect::>() + .iter() + .map(String::as_str) + .collect(); + let nested = StructArray::from(vec![ + ( + Arc::new(Field::new("a", DataType::Int32, true)), + Arc::new(Int32Array::from_iter( + rows.clone().map(|i| (i % 4 != 0).then_some(i as i32)), + )) as ArrayRef, + ), + ( + Arc::new(Field::new("b", DataType::Utf8, true)), + Arc::new(StringArray::from_iter( + rows.map(|i| (i % 6 != 0).then(|| format!("v{i}"))), + )) as ArrayRef, + ), + ]); + RecordBatch::try_new( + schema(), + vec![ + Arc::new(ints), + Arc::new(floats), + Arc::new(tags), + Arc::new(nested), + ], + ) + .unwrap() + } + + fn props(row_group_rows: usize) -> WriterProperties { + WriterProperties::builder() + .set_compression(Compression::ZSTD(ZstdLevel::default())) + .set_max_row_group_row_count(Some(row_group_rows)) + .set_data_page_row_count_limit(16) + .build() + } + + /// One step of a write sequence: a batch of `len` rows, or a flush. + #[derive(Clone, Copy)] + enum Step { + Write(usize), + Flush, + } + + fn write_serial(steps: &[Step], props: WriterProperties) -> Vec { + let mut buf = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut buf, schema(), Some(props)).unwrap(); + let mut start = 0; + for step in steps { + match *step { + Step::Write(len) => { + writer.write(&batch(start, len)).unwrap(); + start += len; + } + Step::Flush => writer.flush().unwrap(), + } + } + writer.close().unwrap(); + buf + } + + async fn write_parallel( + steps: &[Step], + props: WriterProperties, + parallelism: usize, + ) -> Vec { + let mut buf = Vec::new(); + let mut writer = ParallelArrowWriter::try_new( + &mut buf, + schema(), + props, + parallelism, + EncodeRuntime::Global, + ) + .unwrap(); + let mut start = 0; + for step in steps { + match *step { + Step::Write(len) => { + writer.write(&batch(start, len)).await.unwrap(); + start += len; + } + Step::Flush => writer.flush().await.unwrap(), + } + } + writer.close().await.unwrap(); + buf + } + + #[tokio::test] + async fn test_parallel_writer_matches_arrow_writer() { + use Step::*; + let cases: Vec<(&str, Vec)> = vec![ + ("empty", vec![]), + ("partial row group", vec![Write(10)]), + ("exact row groups", vec![Write(64), Write(64)]), + ( + "batches across row groups", + vec![Write(7), Write(100), Write(3), Write(250), Write(1)], + ), + ("batch spanning many row groups", vec![Write(1000)]), + ( + "explicit flushes", + vec![Write(10), Flush, Flush, Write(70), Flush, Write(5)], + ), + ]; + for (name, steps) in cases { + let expected = write_serial(&steps, props(64)); + for parallelism in [1, 2, 3, 16] { + let actual = write_parallel(&steps, props(64), parallelism).await; + assert!( + expected == actual, + "case {name}, parallelism {parallelism}: {} bytes != {} bytes", + expected.len(), + actual.len() + ); + } + } + } + + #[rstest::rstest] + #[case::encode_error(false)] + #[case::encode_panic(true)] + #[tokio::test] + async fn test_encode_failure_drains_tasks_before_releasing_budget(#[case] panics: bool) { + let runtime: common_runtime::Runtime = common_runtime::Builder::default() + .worker_threads(1) + .max_blocking_threads(1) + .build() + .unwrap(); + let mut writer = ParallelArrowWriter::try_new( + Vec::new(), + schema(), + props(64), + 3, + EncodeRuntime::Compact, + ) + .unwrap(); + let memory = new_compaction_memory_manager(1024 * 1024); + let budget = memory.try_acquire(memory.limit_bytes()).unwrap(); + let batches: Arc<[RecordBatch]> = vec![batch(0, 4)].into(); + let retained_batches = Arc::downgrade(&batches); + + let failed = runtime.spawn_blocking(move || -> Result { + if panics { + panic!("encode task failed"); + } + Err(ParquetError::General("encode task failed".to_string())) + }); + let (started_tx, started_rx) = oneshot::channel(); + let (release_tx, release_rx) = oneshot::channel(); + let column = writer.factory.create_column_writers(0).unwrap().remove(0); + let running_batches = batches.clone(); + let running = runtime.spawn_blocking(move || { + started_tx.send(()).unwrap(); + release_rx.blocking_recv().unwrap(); + let _encoded = + encode_fields(&schema(), &running_batches, vec![(0, vec![(0, column)])])?; + Err(ParquetError::General("later encoding failure".to_string())) + }); + // With one blocking thread, this also proves the first task has ended. + timeout(Duration::from_secs(10), started_rx) + .await + .unwrap() + .unwrap(); + let queued_started = Arc::new(AtomicBool::new(false)); + let queued_started_task = queued_started.clone(); + let queued = runtime.spawn_blocking(move || { + queued_started_task.store(true, Ordering::SeqCst); + drop(batches); + Ok((Vec::new(), Vec::new())) + }); + writer.encoding = vec![failed, running, queued]; + + // Model the compaction caller's guard: it is released when flush returns. + let mut finish = Box::pin(async { + let _budget = budget; + writer.flush().await + }); + assert!(futures::poll!(finish.as_mut()).is_pending()); + assert!(memory.try_acquire(memory.limit_bytes()).is_none()); + assert!(retained_batches.upgrade().is_some()); + + release_tx.send(()).unwrap(); + let error = timeout(Duration::from_secs(10), finish) + .await + .unwrap() + .unwrap_err(); + if panics { + let ParquetError::External(source) = error else { + panic!("expected the original task panic, got {error}"); + }; + assert!( + source + .downcast_ref::() + .unwrap() + .is_panic() + ); + } else { + assert!( + matches!(error, ParquetError::General(message) if message == "encode task failed") + ); + } + assert!(!queued_started.load(Ordering::SeqCst)); + assert!(retained_batches.upgrade().is_none()); + assert!(memory.try_acquire(memory.limit_bytes()).is_some()); + assert!(writer.file_writer.flushed_row_groups().is_empty()); + assert!(writer.async_writer.is_empty()); + } + + #[test] + fn test_parallel_writer_requires_row_count_limit() { + assert!(ParallelArrowWriter::>::supports(&props(64))); + let byte_limited = WriterProperties::builder() + .set_max_row_group_bytes(Some(1024)) + .build(); + assert!(!ParallelArrowWriter::>::supports(&byte_limited)); + let unlimited = WriterProperties::builder() + .set_max_row_group_row_count(None) + .build(); + assert!(!ParallelArrowWriter::>::supports(&unlimited)); + assert!( + ParallelArrowWriter::try_new(Vec::new(), schema(), unlimited, 2, EncodeRuntime::Global) + .is_err() + ); + } + + #[test] + fn test_assign_fields_balances_measured_costs() { + let (file_writer, factory) = ArrowWriter::try_new(Vec::new(), schema(), Some(props(64))) + .unwrap() + .into_serialized_writer() + .unwrap(); + let schema_descr = file_writer.schema_descr(); + let mut fields: Vec = (0..4).map(|field| (field, Vec::new())).collect(); + for (leaf, writer) in factory + .create_column_writers(0) + .unwrap() + .into_iter() + .enumerate() + { + fields[schema_descr.get_column_root_idx(leaf)] + .1 + .push((leaf, writer)); + } + // The nested field owns both struct leaves. + assert_eq!(2, fields[3].1.len()); + + let costs = [10, 9, 1, 1].map(Duration::from_millis); + let tasks = assign_fields(fields, &costs, 2); + let mut loads: Vec<_> = tasks + .iter() + .map(|task| { + task.iter() + .map(|(field, _)| costs[*field]) + .sum::() + }) + .collect(); + loads.sort(); + assert_eq!( + vec![Duration::from_millis(10), Duration::from_millis(11)], + loads + ); + let mut assigned: Vec<_> = tasks.iter().flatten().map(|(field, _)| *field).collect(); + assigned.sort(); + assert_eq!(vec![0, 1, 2, 3], assigned); + } +} diff --git a/src/mito2/src/sst/parquet/writer.rs b/src/mito2/src/sst/parquet/writer.rs index 69a89e5802b..d7656d8c354 100644 --- a/src/mito2/src/sst/parquet/writer.rs +++ b/src/mito2/src/sst/parquet/writer.rs @@ -36,8 +36,9 @@ use datatypes::arrow::record_batch::RecordBatch; use datatypes::extension::json::is_json2_extension_type; use object_store::{FuturesAsyncWriter, ObjectStore}; use parquet::arrow::AsyncArrowWriter; +use parquet::arrow::async_writer::AsyncFileWriter; use parquet::basic::{Compression, Encoding, ZstdLevel}; -use parquet::file::metadata::KeyValue; +use parquet::file::metadata::{KeyValue, ParquetMetaData}; use parquet::file::properties::WriterProperties; use smallvec::smallvec; use snafu::{OptionExt, ResultExt}; @@ -59,6 +60,7 @@ use crate::sst::parquet::flat_format::{ FlatWriteFormat, primary_key_column_index, time_index_column_index, }; use crate::sst::parquet::format::{PrimaryKeyArray, PrimaryKeyWriteFormat}; +use crate::sst::parquet::parallel_writer::ParallelArrowWriter; use crate::sst::parquet::{ COLUMN_INDEX_TRUNCATE_LENGTH, PARQUET_METADATA_KEY, SstInfo, WriteOptions, apply_column_compression, apply_float_field_encoding, apply_internal_column_encoding, @@ -154,11 +156,48 @@ fn split_at_next_series( }) } +/// Arrow writer of the current SST file. +enum SstArrowWriter { + /// Encodes on the writing task. + Serial(AsyncArrowWriter), + /// Encodes the columns of each row group on blocking tasks. + Parallel(ParallelArrowWriter), +} + +impl SstArrowWriter { + async fn write(&mut self, batch: &RecordBatch) -> parquet::errors::Result<()> { + match self { + SstArrowWriter::Serial(writer) => writer.write(batch).await, + SstArrowWriter::Parallel(writer) => writer.write(batch).await, + } + } + + async fn flush(&mut self) -> parquet::errors::Result<()> { + match self { + SstArrowWriter::Serial(writer) => writer.flush().await, + SstArrowWriter::Parallel(writer) => writer.flush().await, + } + } + + async fn close(self) -> parquet::errors::Result { + match self { + SstArrowWriter::Serial(writer) => writer.close().await, + SstArrowWriter::Parallel(writer) => writer.close().await, + } + } + + async fn abort(self) { + if let SstArrowWriter::Parallel(writer) = self { + writer.abort().await; + } + } +} + /// Parquet SST writer. pub struct ParquetWriter<'a, F: WriterFactory, I: IndexerBuilder, P: FilePathProvider> { /// Path provider that creates SST and index file paths according to file id. path_provider: P, - writer: Option>>, + writer: Option>>, /// Current active file id. current_file: FileId, writer_factory: F, @@ -374,10 +413,7 @@ where ); let res = self.write_all_flat_inner(source, &converter, opts).await; if res.is_err() { - let file_id = self.current_file; - if let Some(cleaner) = &self.file_cleaner { - cleaner.clean_by_file_id(file_id).await; - } + self.abort_current_file().await; } res } @@ -402,10 +438,7 @@ where }; let res = self.write_all_flat_inner(source, &converter, opts).await; if res.is_err() { - let file_id = self.current_file; - if let Some(cleaner) = &self.file_cleaner { - cleaner.clean_by_file_id(file_id).await; - } + self.abort_current_file().await; } res } @@ -520,6 +553,16 @@ where .is_some_and(|max_size| self.bytes_written.load(Ordering::Relaxed) >= max_size) } + /// Joins encoders before dropping the writer and deleting its failed output. + async fn abort_current_file(&mut self) { + if let Some(writer) = self.writer.take() { + writer.abort().await; + } + if let Some(cleaner) = &self.file_cleaner { + cleaner.clean_by_file_id(self.current_file).await; + } + } + async fn abort_current_indexer(&mut self) { if let Some(indexer) = &mut self.current_indexer { indexer.abort().await; @@ -530,7 +573,7 @@ where &mut self, schema: &SchemaRef, opts: &WriteOptions, - ) -> Result<&mut AsyncArrowWriter>> { + ) -> Result<&mut SstArrowWriter>> { if let Some(ref mut w) = self.writer { Ok(w) } else { @@ -564,9 +607,26 @@ where .await?, self.bytes_written.clone(), ); - let arrow_writer = - AsyncArrowWriter::try_new(writer, maybe_wrap_schema(schema)?, Some(writer_props)) - .context(WriteParquetSnafu)?; + let schema = maybe_wrap_schema(schema)?; + let arrow_writer = if opts.encode_parallelism > 1 + && ParallelArrowWriter::>::supports(&writer_props) + { + SstArrowWriter::Parallel( + ParallelArrowWriter::try_new( + writer, + schema, + writer_props, + opts.encode_parallelism, + opts.encode_runtime, + ) + .context(WriteParquetSnafu)?, + ) + } else { + SstArrowWriter::Serial( + AsyncArrowWriter::try_new(writer, schema, Some(writer_props)) + .context(WriteParquetSnafu)?, + ) + }; self.writer = Some(arrow_writer); let indexer = self @@ -732,3 +792,145 @@ where Pin::new(&mut self.inner).poll_shutdown(cx) } } + +#[cfg(test)] +mod tests { + use std::sync::atomic::AtomicBool; + use std::time::Duration; + + use common_runtime::runtime::{BuilderBuild, RuntimeTrait}; + use parquet::arrow::ArrowWriter; + use parquet::errors::ParquetError; + use store_api::region_request::PathType; + use tokio::sync::oneshot; + use tokio::time::timeout; + + use super::*; + use crate::access_layer::{RegionFilePathFactory, WriteType}; + use crate::compaction::memory_manager::new_compaction_memory_manager; + use crate::error::Error; + use crate::sst::FormatType; + use crate::sst::parquet::tests::NoopIndexBuilder; + use crate::test_util::TestEnv; + use crate::test_util::sst_util::{new_record_batch_by_range, sst_region_metadata}; + + #[rstest::rstest] + #[case::flat(FormatType::Flat)] + #[case::primary_key(FormatType::PrimaryKey)] + #[tokio::test] + async fn test_source_error_waits_for_encoding_cleanup(#[case] format: FormatType) { + let mut env = TestEnv::new().await; + let metadata = Arc::new(sst_region_metadata()); + let mut metrics = Metrics::new(WriteType::Compaction); + let mut writer = ParquetWriter::new_with_object_store( + env.init_object_store_manager(), + metadata.clone(), + IndexConfig::default(), + NoopIndexBuilder, + RegionFilePathFactory::new("".to_string(), PathType::Bare), + &mut metrics, + ) + .await; + let opts = WriteOptions { + row_group_size: 8, + encode_parallelism: 2, + ..Default::default() + }; + let batch = new_record_batch_by_range(&["a", "b"], 0, 4); + let retained_field = Arc::downgrade(batch.column(2)); + let converter = match format { + FormatType::Flat => FlatBatchConverter::Flat(FlatWriteFormat::new( + metadata.clone(), + &FlatSchemaOptions::default(), + )), + FormatType::PrimaryKey => FlatBatchConverter::PrimaryKey { + format: PrimaryKeyWriteFormat::new(metadata.clone()), + num_fields: metadata.field_columns().count(), + }, + }; + let encoding_batch = converter.convert_batch(&batch).unwrap(); + let encoding_schema = encoding_batch.schema(); + let (_, factory) = ArrowWriter::try_new(Vec::new(), encoding_schema.clone(), None) + .unwrap() + .into_serialized_writer() + .unwrap(); + let columns = factory.create_column_writers(0).unwrap(); + let runtime: common_runtime::Runtime = common_runtime::Builder::default() + .worker_threads(1) + .max_blocking_threads(1) + .build() + .unwrap(); + let (started_tx, started_rx) = oneshot::channel(); + let (release_tx, release_rx) = oneshot::channel(); + let queued_batch = encoding_batch.clone(); + let running = runtime.spawn_blocking(move || { + started_tx.send(()).unwrap(); + release_rx.blocking_recv().unwrap(); + drop((encoding_batch, columns)); + Err(ParquetError::General( + "cleanup encoding failure".to_string(), + )) + }); + timeout(Duration::from_secs(10), started_rx) + .await + .unwrap() + .unwrap(); + let queued_started = Arc::new(AtomicBool::new(false)); + let queued_started_task = queued_started.clone(); + let queued = runtime.spawn_blocking(move || { + queued_started_task.store(true, Ordering::SeqCst); + drop(queued_batch); + Ok((Vec::new(), Vec::new())) + }); + let SstArrowWriter::Parallel(parallel) = writer + .maybe_init_writer(&encoding_schema, &opts) + .await + .unwrap() + else { + panic!("expected parallel encoding"); + }; + // Controlled real tasks represent the previous row group in flight. + parallel.set_encoding_tasks_for_test(vec![running, queued]); + let source = FlatSource::new_stream( + batch.schema(), + Box::pin(futures::stream::iter(vec![ + Ok(batch), + InvalidRecordBatchSnafu { + reason: "source read failed", + } + .fail(), + ])), + ); + let memory = new_compaction_memory_manager(1024 * 1024); + let budget = memory.try_acquire(memory.limit_bytes()).unwrap(); + let mut write = Box::pin(async { + let _budget = budget; + match format { + FormatType::Flat => writer.write_all_flat(source, None, &opts).await, + FormatType::PrimaryKey => { + writer + .write_all_flat_as_primary_key(source, None, &opts) + .await + } + } + }); + // The next row group is buffered, then the source fails. No codec timing + // or I/O wait can make the old indexer-only cleanup pass this assertion. + assert!(futures::poll!(write.as_mut()).is_pending()); + assert!(memory.try_acquire(memory.limit_bytes()).is_none()); + assert!(retained_field.upgrade().is_some()); + release_tx.send(()).unwrap(); + let error = timeout(Duration::from_secs(10), write) + .await + .unwrap() + .unwrap_err(); + assert!( + matches!(error, Error::InvalidRecordBatch { reason, .. } if reason == "source read failed") + ); + assert!(!queued_started.load(Ordering::SeqCst)); + assert!(retained_field.upgrade().is_none()); + assert!(memory.try_acquire(memory.limit_bytes()).is_some()); + assert!(writer.writer.is_none()); + assert_eq!(0, writer.bytes_written.load(Ordering::Relaxed)); + } +} diff --git a/src/mito2/src/worker/handle_alter.rs b/src/mito2/src/worker/handle_alter.rs index 1798f9a79bf..580e55f1bed 100644 --- a/src/mito2/src/worker/handle_alter.rs +++ b/src/mito2/src/worker/handle_alter.rs @@ -311,6 +311,11 @@ impl RegionWorkerLoop { current_options.auto_flush_interval = new_interval; } } + SetRegionOption::CompactionOutputParallelism(parallelism) => { + // New compactions use the updated options; running ones retain + // their snapshot. No memtable flush is needed. + current_options.compaction_output_parallelism = parallelism; + } SetRegionOption::MaxRowGroupRowCount(new_row_count) => { if let Some(row_count) = new_row_count { ensure!( @@ -431,6 +436,7 @@ fn new_region_options_on_empty_memtable( | SetRegionOption::Ttl(_) | SetRegionOption::Twsc(_, _) | SetRegionOption::AutoFlushInterval(_) + | SetRegionOption::CompactionOutputParallelism(_) | SetRegionOption::SkipWal(_) => (), SetRegionOption::Format(format_str) => { // Safety: handle_alter_region_options_fast() has validated this. diff --git a/src/store-api/src/mito_engine_options.rs b/src/store-api/src/mito_engine_options.rs index 457b33665c3..1f680aff995 100644 --- a/src/store-api/src/mito_engine_options.rs +++ b/src/store-api/src/mito_engine_options.rs @@ -29,6 +29,9 @@ pub const MERGE_MODE_KEY: &str = "merge_mode"; pub const TTL_KEY: &str = "ttl"; /// Option key for the per-table auto flush interval. pub const AUTO_FLUSH_INTERVAL_KEY: &str = "auto_flush_interval"; +/// Per-table override for each compaction output's decode and encode parallelism. +pub const COMPACTION_EXPERIMENTAL_OUTPUT_PARALLELISM: &str = + "compaction.experimental_output_parallelism"; /// Option key for compaction type. pub const COMPACTION_TYPE: &str = "compaction.type"; /// Option key for forcing compaction options override. @@ -166,6 +169,7 @@ pub fn is_mito_engine_option_key(key: &str) -> bool { "ttl", WRITE_BUFFER_SIZE_KEY, AUTO_FLUSH_INTERVAL_KEY, + COMPACTION_EXPERIMENTAL_OUTPUT_PARALLELISM, "storage", "index.inverted_index.ignore_column_ids", "index.inverted_index.segment_row_count", @@ -200,6 +204,9 @@ mod tests { assert!(is_mito_engine_option_key("ttl")); assert!(is_mito_engine_option_key("write_buffer_size")); assert!(is_mito_engine_option_key("auto_flush_interval")); + assert!(is_mito_engine_option_key( + "compaction.experimental_output_parallelism" + )); assert!(is_mito_engine_option_key("compaction.type")); assert!(is_mito_engine_option_key("compaction.override")); assert!(is_mito_engine_option_key( diff --git a/src/store-api/src/region_request.rs b/src/store-api/src/region_request.rs index 937fb9d950f..787d6c5afe1 100644 --- a/src/store-api/src/region_request.rs +++ b/src/store-api/src/region_request.rs @@ -60,11 +60,12 @@ use crate::metadata::{ use crate::metric_engine_consts::PHYSICAL_TABLE_METADATA_KEY; use crate::metrics; use crate::mito_engine_options::{ - APPEND_MODE_KEY, AUTO_FLUSH_INTERVAL_KEY, MAX_ROW_GROUP_ROW_COUNT, - MAX_ROW_GROUP_ROW_COUNT_LIMIT, PRESERVE_ROW_SEQUENCE, SKIP_WAL_KEY, SST_FORMAT_KEY, TTL_KEY, - TWCS_ACTIVE_WINDOW_L1_MERGE_TRIGGER, TWCS_ACTIVE_WINDOW_TRIGGER_FILE_NUM, - TWCS_INACTIVE_WINDOW_L1_MERGE_TRIGGER, TWCS_INACTIVE_WINDOW_TRIGGER_FILE_NUM, - TWCS_MAX_OUTPUT_FILE_SIZE, TWCS_TIME_WINDOW, TWCS_TRIGGER_FILE_NUM, WRITE_BUFFER_SIZE_KEY, + APPEND_MODE_KEY, AUTO_FLUSH_INTERVAL_KEY, COMPACTION_EXPERIMENTAL_OUTPUT_PARALLELISM, + MAX_ROW_GROUP_ROW_COUNT, MAX_ROW_GROUP_ROW_COUNT_LIMIT, PRESERVE_ROW_SEQUENCE, SKIP_WAL_KEY, + SST_FORMAT_KEY, TTL_KEY, TWCS_ACTIVE_WINDOW_L1_MERGE_TRIGGER, + TWCS_ACTIVE_WINDOW_TRIGGER_FILE_NUM, TWCS_INACTIVE_WINDOW_L1_MERGE_TRIGGER, + TWCS_INACTIVE_WINDOW_TRIGGER_FILE_NUM, TWCS_MAX_OUTPUT_FILE_SIZE, TWCS_TIME_WINDOW, + TWCS_TRIGGER_FILE_NUM, WRITE_BUFFER_SIZE_KEY, }; use crate::path_utils::table_dir; use crate::storage::{ColumnId, RegionId}; @@ -1603,6 +1604,8 @@ pub enum SetRegionOption { PreserveRowSequence(bool), // Whether to skip writing new WAL entries. SkipWal(bool), + /// Overrides per-output compaction parallelism; None restores the engine default. + CompactionOutputParallelism(Option), } #[derive(Serialize, Deserialize)] @@ -1616,6 +1619,7 @@ enum SetRegionOptionSerde { MaxRowGroupRowCount(Option), PreserveRowSequence(bool), SkipWal(bool), + CompactionOutputParallelism(Option), } impl Serialize for SetRegionOption { @@ -1638,6 +1642,9 @@ impl Serialize for SetRegionOption { Self::MaxRowGroupRowCount(value) => SetRegionOptionSerde::MaxRowGroupRowCount(*value), Self::PreserveRowSequence(value) => SetRegionOptionSerde::PreserveRowSequence(*value), Self::SkipWal(value) => SetRegionOptionSerde::SkipWal(*value), + Self::CompactionOutputParallelism(value) => { + SetRegionOptionSerde::CompactionOutputParallelism(*value) + } }; option.serialize(serializer) } @@ -1667,6 +1674,9 @@ impl From for SetRegionOption { SetRegionOptionSerde::MaxRowGroupRowCount(value) => Self::MaxRowGroupRowCount(value), SetRegionOptionSerde::PreserveRowSequence(value) => Self::PreserveRowSequence(value), SetRegionOptionSerde::SkipWal(value) => Self::SkipWal(value), + SetRegionOptionSerde::CompactionOutputParallelism(value) => { + Self::CompactionOutputParallelism(value) + } } } } @@ -1733,6 +1743,17 @@ impl TryFrom<&PbOption> for SetRegionOption { } Ok(Self::AutoFlushInterval(Some(interval))) } + COMPACTION_EXPERIMENTAL_OUTPUT_PARALLELISM => { + if value.is_empty() { + return Ok(Self::CompactionOutputParallelism(None)); + } + let parallelism = value + .parse::() + .ok() + .filter(|parallelism| *parallelism > 0) + .ok_or_else(|| InvalidSetRegionOptionRequestSnafu { key, value }.build())?; + Ok(Self::CompactionOutputParallelism(Some(parallelism))) + } MAX_ROW_GROUP_ROW_COUNT => { if value.is_empty() { return Ok(Self::MaxRowGroupRowCount(None)); @@ -1791,6 +1812,9 @@ impl From<&UnsetRegionOption> for SetRegionOption { UnsetRegionOption::MaxRowGroupRowCount => SetRegionOption::MaxRowGroupRowCount(None), UnsetRegionOption::WriteBufferSize => SetRegionOption::WriteBufferSize(None), UnsetRegionOption::PreserveRowSequence => SetRegionOption::PreserveRowSequence(false), + UnsetRegionOption::CompactionOutputParallelism => { + SetRegionOption::CompactionOutputParallelism(None) + } } } } @@ -1811,6 +1835,7 @@ impl TryFrom<&str> for UnsetRegionOption { TWCS_TIME_WINDOW => Ok(Self::TwcsTimeWindow), MAX_ROW_GROUP_ROW_COUNT => Ok(Self::MaxRowGroupRowCount), PRESERVE_ROW_SEQUENCE => Ok(Self::PreserveRowSequence), + COMPACTION_EXPERIMENTAL_OUTPUT_PARALLELISM => Ok(Self::CompactionOutputParallelism), _ => InvalidUnsetRegionOptionRequestSnafu { key }.fail(), } } @@ -1829,6 +1854,7 @@ pub enum UnsetRegionOption { WriteBufferSize, PreserveRowSequence, TwcsActiveWindowL1MergeTrigger, + CompactionOutputParallelism, } impl UnsetRegionOption { @@ -1845,6 +1871,7 @@ impl UnsetRegionOption { Self::TwcsTimeWindow => TWCS_TIME_WINDOW, Self::MaxRowGroupRowCount => MAX_ROW_GROUP_ROW_COUNT, Self::PreserveRowSequence => PRESERVE_ROW_SEQUENCE, + Self::CompactionOutputParallelism => COMPACTION_EXPERIMENTAL_OUTPUT_PARALLELISM, } } } @@ -2218,6 +2245,42 @@ mod tests { .unwrap_err(); } + #[test] + fn test_compaction_output_parallelism_set_unset_and_serde() { + for (raw, expected) in [("", None), ("1", Some(1)), ("4", Some(4))] { + let option = SetRegionOption::try_from(&PbOption { + key: COMPACTION_EXPERIMENTAL_OUTPUT_PARALLELISM.to_string(), + value: raw.to_string(), + }) + .unwrap(); + assert_eq!( + SetRegionOption::CompactionOutputParallelism(expected), + option + ); + let json = serde_json::to_string(&option).unwrap(); + assert_eq!( + option, + serde_json::from_str::(&json).unwrap() + ); + } + for raw in ["0", "-1", "1.5", "invalid", "18446744073709551616"] { + assert!( + SetRegionOption::try_from(&PbOption { + key: COMPACTION_EXPERIMENTAL_OUTPUT_PARALLELISM.to_string(), + value: raw.to_string(), + }) + .is_err() + ); + } + let unset = + UnsetRegionOption::try_from(COMPACTION_EXPERIMENTAL_OUTPUT_PARALLELISM).unwrap(); + assert_eq!(COMPACTION_EXPERIMENTAL_OUTPUT_PARALLELISM, unset.as_str()); + assert_eq!( + SetRegionOption::CompactionOutputParallelism(None), + (&unset).into() + ); + } + #[test] fn test_set_region_option_auto_flush_interval_try_from() { use std::time::Duration; diff --git a/src/table/src/metadata.rs b/src/table/src/metadata.rs index 57a4548160a..c94275c9e8a 100644 --- a/src/table/src/metadata.rs +++ b/src/table/src/metadata.rs @@ -31,9 +31,10 @@ use serde::{Deserialize, Deserializer, Serialize}; use snafu::{OptionExt, ResultExt, ensure}; use store_api::metric_engine_consts::PHYSICAL_TABLE_METADATA_KEY; use store_api::mito_engine_options::{ - APPEND_MODE_KEY, AUTO_FLUSH_INTERVAL_KEY, COMPACTION_TYPE, COMPACTION_TYPE_TWCS, - MAX_ROW_GROUP_ROW_COUNT, MERGE_MODE_KEY, PRESERVE_ROW_SEQUENCE, SKIP_WAL_KEY, SST_FORMAT_KEY, - TWCS_ACTIVE_WINDOW_TRIGGER_FILE_NUM, TWCS_TRIGGER_FILE_NUM, + APPEND_MODE_KEY, AUTO_FLUSH_INTERVAL_KEY, COMPACTION_EXPERIMENTAL_OUTPUT_PARALLELISM, + COMPACTION_TYPE, COMPACTION_TYPE_TWCS, MAX_ROW_GROUP_ROW_COUNT, MERGE_MODE_KEY, + PRESERVE_ROW_SEQUENCE, SKIP_WAL_KEY, SST_FORMAT_KEY, TWCS_ACTIVE_WINDOW_TRIGGER_FILE_NUM, + TWCS_TRIGGER_FILE_NUM, }; use store_api::region_request::{SetRegionOption, UnsetRegionOption}; use store_api::storage::{ColumnDescriptor, ColumnDescriptorBuilder, ColumnId}; @@ -415,6 +416,18 @@ impl TableMeta { new_options.extra_options.remove(AUTO_FLUSH_INTERVAL_KEY); } } + SetRegionOption::CompactionOutputParallelism(parallelism) => { + if let Some(parallelism) = parallelism { + new_options.extra_options.insert( + COMPACTION_EXPERIMENTAL_OUTPUT_PARALLELISM.to_string(), + parallelism.to_string(), + ); + } else { + new_options + .extra_options + .remove(COMPACTION_EXPERIMENTAL_OUTPUT_PARALLELISM); + } + } SetRegionOption::MaxRowGroupRowCount(row_count) => { if let Some(row_count) = row_count { new_options @@ -2571,6 +2584,55 @@ mod tests { ); } + #[test] + fn test_compaction_output_parallelism_persistence_and_unset() { + let mut meta = TableMetaBuilder::empty() + .schema(Arc::new(new_test_schema())) + .primary_key_indices(vec![0]) + .engine("mito") + .next_column_id(3) + .options(TableOptions::default()) + .build() + .unwrap(); + for (kind, expected) in [ + ( + AlterKind::SetTableOptions { + options: vec![SetRegionOption::CompactionOutputParallelism(Some(4))], + }, + Some("4"), + ), + ( + AlterKind::SetTableOptions { + options: vec![SetRegionOption::CompactionOutputParallelism(Some(1))], + }, + Some("1"), + ), + ( + AlterKind::UnsetTableOptions { + keys: vec![UnsetRegionOption::CompactionOutputParallelism], + }, + None, + ), + ] { + meta = meta + .builder_with_alter_kind("my_table", &kind) + .unwrap() + .build() + .unwrap(); + // Restart uses the persisted table options to reopen regions. + let json = serde_json::to_string(&meta.options).unwrap(); + let restored: TableOptions = serde_json::from_str(&json).unwrap(); + assert_eq!( + expected, + restored + .extra_options + .get(COMPACTION_EXPERIMENTAL_OUTPUT_PARALLELISM) + .map(String::as_str) + ); + assert!(!restored.extra_options.contains_key(COMPACTION_TYPE)); + } + } + #[test] fn test_set_auto_flush_interval() { let schema = Arc::new(new_test_schema()); diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 1827e8b6581..c13ada351ea 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -3573,6 +3573,7 @@ experimental_series_index_maintenance_interval = "5m" experimental_series_index_bucket_width = "5days" experimental_compaction_memory_limit = "unlimited" experimental_compaction_on_exhausted = "wait" +experimental_compaction_output_parallelism = 1 auto_flush_interval = "10m" default_region_write_buffer_size = "0KiB" enable_write_cache = false diff --git a/tests/cases/standalone/common/alter/compaction_output_parallelism.result b/tests/cases/standalone/common/alter/compaction_output_parallelism.result new file mode 100644 index 00000000000..d8e04bc7837 --- /dev/null +++ b/tests/cases/standalone/common/alter/compaction_output_parallelism.result @@ -0,0 +1,140 @@ +-- Execution parallelism does not require compaction.type. +CREATE TABLE compaction_output_parallelism ( + ts TIMESTAMP TIME INDEX, + val INT +) WITH ('compaction.experimental_output_parallelism' = '4'); + +Affected Rows: 0 + +INSERT INTO compaction_output_parallelism VALUES (0, 10); + +Affected Rows: 1 + +ADMIN FLUSH_TABLE('compaction_output_parallelism'); + ++----------------------------------------------------+ +| ADMIN FLUSH_TABLE('compaction_output_parallelism') | ++----------------------------------------------------+ +| 0 | ++----------------------------------------------------+ + +INSERT INTO compaction_output_parallelism VALUES (1000, 20); + +Affected Rows: 1 + +ADMIN FLUSH_TABLE('compaction_output_parallelism'); + ++----------------------------------------------------+ +| ADMIN FLUSH_TABLE('compaction_output_parallelism') | ++----------------------------------------------------+ +| 0 | ++----------------------------------------------------+ + +ADMIN COMPACT_TABLE('compaction_output_parallelism'); + ++------------------------------------------------------+ +| ADMIN COMPACT_TABLE('compaction_output_parallelism') | ++------------------------------------------------------+ +| 0 | ++------------------------------------------------------+ + +-- SQLNESS ARG restart=true +SHOW CREATE TABLE compaction_output_parallelism; + ++-------------------------------+--------------------------------------------------------------+ +| Table | Create Table | ++-------------------------------+--------------------------------------------------------------+ +| compaction_output_parallelism | CREATE TABLE IF NOT EXISTS "compaction_output_parallelism" ( | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | "val" INT NULL, | +| | TIME INDEX ("ts") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | 'compaction.experimental_output_parallelism' = '4' | +| | ) | ++-------------------------------+--------------------------------------------------------------+ + +SELECT * FROM compaction_output_parallelism ORDER BY ts; + ++---------------------+-----+ +| ts | val | ++---------------------+-----+ +| 1970-01-01T00:00:00 | 10 | +| 1970-01-01T00:00:01 | 20 | ++---------------------+-----+ + +-- A table can disable parallel encoding even when the engine default enables it. +ALTER TABLE compaction_output_parallelism SET 'compaction.experimental_output_parallelism' = '1'; + +Affected Rows: 0 + +SHOW CREATE TABLE compaction_output_parallelism; + ++-------------------------------+--------------------------------------------------------------+ +| Table | Create Table | ++-------------------------------+--------------------------------------------------------------+ +| compaction_output_parallelism | CREATE TABLE IF NOT EXISTS "compaction_output_parallelism" ( | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | "val" INT NULL, | +| | TIME INDEX ("ts") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | 'compaction.experimental_output_parallelism' = '1' | +| | ) | ++-------------------------------+--------------------------------------------------------------+ + +ALTER TABLE compaction_output_parallelism SET 'compaction.experimental_output_parallelism' = '0'; + +Error: 1004(InvalidArguments), Invalid set table option request: Invalid set region option request, key: compaction.experimental_output_parallelism, value: 0 + +ALTER TABLE compaction_output_parallelism SET 'compaction.experimental_output_parallelism' = '-1'; + +Error: 1004(InvalidArguments), Invalid set table option request: Invalid set region option request, key: compaction.experimental_output_parallelism, value: -1 + +ALTER TABLE compaction_output_parallelism SET 'compaction.experimental_output_parallelism' = 'invalid'; + +Error: 1004(InvalidArguments), Invalid set table option request: Invalid set region option request, key: compaction.experimental_output_parallelism, value: invalid + +-- UNSET removes the override without introducing compaction.type or compaction.override. +ALTER TABLE compaction_output_parallelism UNSET 'compaction.experimental_output_parallelism'; + +Affected Rows: 0 + +-- SQLNESS ARG restart=true +SHOW CREATE TABLE compaction_output_parallelism; + ++-------------------------------+--------------------------------------------------------------+ +| Table | Create Table | ++-------------------------------+--------------------------------------------------------------+ +| compaction_output_parallelism | CREATE TABLE IF NOT EXISTS "compaction_output_parallelism" ( | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | "val" INT NULL, | +| | TIME INDEX ("ts") | +| | ) | +| | | +| | ENGINE=mito | +| | | ++-------------------------------+--------------------------------------------------------------+ + +SELECT * FROM compaction_output_parallelism ORDER BY ts; + ++---------------------+-----+ +| ts | val | ++---------------------+-----+ +| 1970-01-01T00:00:00 | 10 | +| 1970-01-01T00:00:01 | 20 | ++---------------------+-----+ + +DROP TABLE compaction_output_parallelism; + +Affected Rows: 0 + +CREATE TABLE invalid_parallelism (ts TIMESTAMP TIME INDEX) +WITH ('compaction.experimental_output_parallelism' = '0'); + +Error: 1004(InvalidArguments), Invalid region options, compaction.experimental_output_parallelism must be greater than 0 + diff --git a/tests/cases/standalone/common/alter/compaction_output_parallelism.sql b/tests/cases/standalone/common/alter/compaction_output_parallelism.sql new file mode 100644 index 00000000000..7772c197379 --- /dev/null +++ b/tests/cases/standalone/common/alter/compaction_output_parallelism.sql @@ -0,0 +1,35 @@ +-- Execution parallelism does not require compaction.type. +CREATE TABLE compaction_output_parallelism ( + ts TIMESTAMP TIME INDEX, + val INT +) WITH ('compaction.experimental_output_parallelism' = '4'); + +INSERT INTO compaction_output_parallelism VALUES (0, 10); +ADMIN FLUSH_TABLE('compaction_output_parallelism'); +INSERT INTO compaction_output_parallelism VALUES (1000, 20); +ADMIN FLUSH_TABLE('compaction_output_parallelism'); +ADMIN COMPACT_TABLE('compaction_output_parallelism'); + +-- SQLNESS ARG restart=true +SHOW CREATE TABLE compaction_output_parallelism; + +SELECT * FROM compaction_output_parallelism ORDER BY ts; + +-- A table can disable parallel encoding even when the engine default enables it. +ALTER TABLE compaction_output_parallelism SET 'compaction.experimental_output_parallelism' = '1'; +SHOW CREATE TABLE compaction_output_parallelism; + +ALTER TABLE compaction_output_parallelism SET 'compaction.experimental_output_parallelism' = '0'; +ALTER TABLE compaction_output_parallelism SET 'compaction.experimental_output_parallelism' = '-1'; +ALTER TABLE compaction_output_parallelism SET 'compaction.experimental_output_parallelism' = 'invalid'; + +-- UNSET removes the override without introducing compaction.type or compaction.override. +ALTER TABLE compaction_output_parallelism UNSET 'compaction.experimental_output_parallelism'; +-- SQLNESS ARG restart=true +SHOW CREATE TABLE compaction_output_parallelism; + +SELECT * FROM compaction_output_parallelism ORDER BY ts; +DROP TABLE compaction_output_parallelism; + +CREATE TABLE invalid_parallelism (ts TIMESTAMP TIME INDEX) +WITH ('compaction.experimental_output_parallelism' = '0'); diff --git a/tests/compatibility/cases/compaction_output_parallelism/case.toml b/tests/compatibility/cases/compaction_output_parallelism/case.toml new file mode 100644 index 00000000000..9e3bc1d8e91 --- /dev/null +++ b/tests/compatibility/cases/compaction_output_parallelism/case.toml @@ -0,0 +1,8 @@ +name = "compaction_output_parallelism" +reason = "Verify old table options reopen with inherited output parallelism and support setting and clearing the new per-table override." +introduced_by = "PR #9945" +topologies = ["distributed", "standalone"] +from_range = ["*"] +to_range = [">=v1.4.0-alpha.0"] +features = ["table", "metadata", "alter", "compaction"] +owner = "storage" diff --git a/tests/compatibility/cases/compaction_output_parallelism/setup.sql b/tests/compatibility/cases/compaction_output_parallelism/setup.sql new file mode 100644 index 00000000000..a3dd0a23f92 --- /dev/null +++ b/tests/compatibility/cases/compaction_output_parallelism/setup.sql @@ -0,0 +1,7 @@ +CREATE TABLE compaction_parallelism ( + ts TIMESTAMP TIME INDEX, + val INT +) ENGINE=mito; + +INSERT INTO compaction_parallelism VALUES (0, 10), (1000, 20); +ADMIN FLUSH_TABLE('compaction_parallelism'); diff --git a/tests/compatibility/cases/compaction_output_parallelism/verify.result b/tests/compatibility/cases/compaction_output_parallelism/verify.result new file mode 100644 index 00000000000..646761ec763 --- /dev/null +++ b/tests/compatibility/cases/compaction_output_parallelism/verify.result @@ -0,0 +1,90 @@ +SELECT * FROM compaction_parallelism ORDER BY ts; + ++---------------------+-----+ +| ts | val | ++---------------------+-----+ +| 1970-01-01T00:00:00 | 10 | +| 1970-01-01T00:00:01 | 20 | ++---------------------+-----+ + +ALTER TABLE compaction_parallelism SET 'compaction.experimental_output_parallelism' = '4'; + +Affected Rows: 0 + +INSERT INTO compaction_parallelism VALUES (2000, 30); + +Affected Rows: 1 + +ADMIN FLUSH_TABLE('compaction_parallelism'); + ++---------------------------------------------+ +| ADMIN FLUSH_TABLE('compaction_parallelism') | ++---------------------------------------------+ +| 0 | ++---------------------------------------------+ + +ADMIN COMPACT_TABLE('compaction_parallelism'); + ++-----------------------------------------------+ +| ADMIN COMPACT_TABLE('compaction_parallelism') | ++-----------------------------------------------+ +| 0 | ++-----------------------------------------------+ + +-- SQLNESS ARG restart=true +SHOW CREATE TABLE compaction_parallelism; + ++------------------------+-------------------------------------------------------+ +| Table | Create Table | ++------------------------+-------------------------------------------------------+ +| compaction_parallelism | CREATE TABLE IF NOT EXISTS "compaction_parallelism" ( | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | "val" INT NULL, | +| | TIME INDEX ("ts") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | 'compaction.experimental_output_parallelism' = '4' | +| | ) | ++------------------------+-------------------------------------------------------+ + +SELECT * FROM compaction_parallelism ORDER BY ts; + ++---------------------+-----+ +| ts | val | ++---------------------+-----+ +| 1970-01-01T00:00:00 | 10 | +| 1970-01-01T00:00:01 | 20 | +| 1970-01-01T00:00:02 | 30 | ++---------------------+-----+ + +ALTER TABLE compaction_parallelism UNSET 'compaction.experimental_output_parallelism'; + +Affected Rows: 0 + +-- SQLNESS ARG restart=true +SHOW CREATE TABLE compaction_parallelism; + ++------------------------+-------------------------------------------------------+ +| Table | Create Table | ++------------------------+-------------------------------------------------------+ +| compaction_parallelism | CREATE TABLE IF NOT EXISTS "compaction_parallelism" ( | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | "val" INT NULL, | +| | TIME INDEX ("ts") | +| | ) | +| | | +| | ENGINE=mito | +| | | ++------------------------+-------------------------------------------------------+ + +SELECT * FROM compaction_parallelism ORDER BY ts; + ++---------------------+-----+ +| ts | val | ++---------------------+-----+ +| 1970-01-01T00:00:00 | 10 | +| 1970-01-01T00:00:01 | 20 | +| 1970-01-01T00:00:02 | 30 | ++---------------------+-----+ diff --git a/tests/compatibility/cases/compaction_output_parallelism/verify.sql b/tests/compatibility/cases/compaction_output_parallelism/verify.sql new file mode 100644 index 00000000000..814a4c45ceb --- /dev/null +++ b/tests/compatibility/cases/compaction_output_parallelism/verify.sql @@ -0,0 +1,15 @@ +SELECT * FROM compaction_parallelism ORDER BY ts; + +ALTER TABLE compaction_parallelism SET 'compaction.experimental_output_parallelism' = '4'; +INSERT INTO compaction_parallelism VALUES (2000, 30); +ADMIN FLUSH_TABLE('compaction_parallelism'); +ADMIN COMPACT_TABLE('compaction_parallelism'); + +-- SQLNESS ARG restart=true +SHOW CREATE TABLE compaction_parallelism; +SELECT * FROM compaction_parallelism ORDER BY ts; + +ALTER TABLE compaction_parallelism UNSET 'compaction.experimental_output_parallelism'; +-- SQLNESS ARG restart=true +SHOW CREATE TABLE compaction_parallelism; +SELECT * FROM compaction_parallelism ORDER BY ts; From b21cd16f1d62a213bb7b90bb8198e3840ce7593f Mon Sep 17 00:00:00 2001 From: Dragon Roar <2181719471@qq.com> Date: Wed, 28 Oct 2026 17:09:29 +0000 Subject: [PATCH 132/239] feat(datanode): probe object storage at startup (#9954) * feat(datanode): self-test object storage at startup Signed-off-by: lczllx <2181719471@qq.com> * fix(datanode): harden the startup object storage self-test Signed-off-by: lczllx <2181719471@qq.com> * test(datanode): cover the File store skip in the object store self-test Signed-off-by: lczllx <2181719471@qq.com> * refactor(datanode): rename the object storage self-test to probe Signed-off-by: lczllx <2181719471@qq.com> * test(datanode): disable object store self-test in sqlness and integration tests Signed-off-by: lczllx <2181719471@qq.com> * docs(datanode): correct what the object storage probe claims to cover Signed-off-by: lczllx <2181719471@qq.com> --------- Signed-off-by: lczllx <2181719471@qq.com> --- Cargo.lock | 1 + config/config.md | 2 + config/datanode.example.toml | 8 + config/standalone.example.toml | 8 + src/datanode/Cargo.toml | 1 + src/datanode/src/config.rs | 22 +++ src/datanode/src/datanode.rs | 164 ++++++++++++++++- src/datanode/src/error.rs | 28 ++- src/datanode/src/store.rs | 221 ++++++++++++++++++++++- tests-integration/src/test_util.rs | 1 + tests-integration/tests/http.rs | 2 + tests/conf/datanode-test.toml.template | 1 + tests/conf/standalone-test.toml.template | 1 + 13 files changed, 447 insertions(+), 13 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index f68dfbe2334..d40ca8b5baa 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4566,6 +4566,7 @@ dependencies = [ "prometheus 0.14.0", "prost 0.14.1", "query", + "rand 0.9.4", "reqwest 0.13.4", "serde", "serde_json", diff --git a/config/config.md b/config/config.md index 711fe019a19..cd2b518f12c 100644 --- a/config/config.md +++ b/config/config.md @@ -153,6 +153,7 @@ | `storage` | -- | -- | The data storage options. | | `storage.data_home` | String | `./greptimedb_data` | The working home directory. | | `storage.copy_root` | String | `./greptimedb_data/copy` | Root directory for standalone SQL access to local files.
Relative SQL paths are resolved below this directory. Absolute paths are accepted only when
they are inside this directory. Defaults to `/copy`.
Distributed deployments always reject SQL access to local files.
Upgrade note: COPY commands and existing external tables that reference paths outside this
directory will fail. Move those files below the copy root, set this option to a dedicated
directory containing them, or migrate the files to object storage before upgrading. | +| `storage.probe_on_startup` | Bool | `true` | Whether to probe every remote object storage when the datanode starts.
The probe writes, reads, overwrites, lists and deletes an object under
`.greptime_probe/`, and a failure aborts startup. The probe lives at the storage
root, so the credentials must allow writes and deletes below the root. Disable it
when the credentials are scoped to narrower prefixes or when the storage may be
unavailable at startup. | | `storage.type` | String | `File` | The storage type used to store the data.
- `File`: the data is stored in the local file system.
- `S3`: the data is stored in the S3 object storage.
- `Gcs`: the data is stored in the Google Cloud Storage.
- `Azblob`: the data is stored in the Azure Blob Storage.
- `Oss`: the data is stored in the Aliyun OSS.
- `Hdfs`: the data is stored in the Hadoop Distributed File System. | | `storage.bucket` | String | Unset | The S3 bucket name.
**It's only used when the storage type is `S3`, `Oss` and `Gcs`**. | | `storage.root` | String | Unset | The directory or object prefix under which data is stored.
**It's only used when the storage type is `S3`, `Oss`, `Gcs`, `Azblob` and `Hdfs`**. | @@ -631,6 +632,7 @@ | `query.experimental_spill_compression` | String | `uncompressed` | Compression for spilled data files: "uncompressed" (default), "lz4_frame", "zstd".
Ignored unless mode is "custom". | | `storage` | -- | -- | The data storage options. | | `storage.data_home` | String | `./greptimedb_data` | The working home directory. | +| `storage.probe_on_startup` | Bool | `true` | Whether to probe every remote object storage when the datanode starts.
The probe writes, reads, overwrites, lists and deletes an object under
`.greptime_probe/`, and a failure aborts startup. The probe lives at the storage
root, so the credentials must allow writes and deletes below the root. Disable it
when the credentials are scoped to narrower prefixes or when the storage may be
unavailable at startup. | | `storage.type` | String | `File` | The storage type used to store the data.
- `File`: the data is stored in the local file system.
- `S3`: the data is stored in the S3 object storage.
- `Gcs`: the data is stored in the Google Cloud Storage.
- `Azblob`: the data is stored in the Azure Blob Storage.
- `Oss`: the data is stored in the Aliyun OSS.
- `Hdfs`: the data is stored in the Hadoop Distributed File System. | | `storage.bucket` | String | Unset | The S3 bucket name.
**It's only used when the storage type is `S3`, `Oss` and `Gcs`**. | | `storage.root` | String | Unset | The directory or object prefix under which data is stored.
**It's only used when the storage type is `S3`, `Oss`, `Gcs`, `Azblob` and `Hdfs`**. | diff --git a/config/datanode.example.toml b/config/datanode.example.toml index a2457302a10..a7e9417717f 100644 --- a/config/datanode.example.toml +++ b/config/datanode.example.toml @@ -375,6 +375,14 @@ memory_pool_size = "50%" ## The working home directory. data_home = "./greptimedb_data" +## Whether to probe every remote object storage when the datanode starts. +## The probe writes, reads, overwrites, lists and deletes an object under +## `.greptime_probe/`, and a failure aborts startup. The probe lives at the storage +## root, so the credentials must allow writes and deletes below the root. Disable it +## when the credentials are scoped to narrower prefixes or when the storage may be +## unavailable at startup. +probe_on_startup = true + ## The storage type used to store the data. ## - `File`: the data is stored in the local file system. ## - `S3`: the data is stored in the S3 object storage. diff --git a/config/standalone.example.toml b/config/standalone.example.toml index d047952538c..0fa5907f95b 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -603,6 +603,14 @@ data_home = "./greptimedb_data" ## directory containing them, or migrate the files to object storage before upgrading. #+ copy_root = "./greptimedb_data/copy" +## Whether to probe every remote object storage when the datanode starts. +## The probe writes, reads, overwrites, lists and deletes an object under +## `.greptime_probe/`, and a failure aborts startup. The probe lives at the storage +## root, so the credentials must allow writes and deletes below the root. Disable it +## when the credentials are scoped to narrower prefixes or when the storage may be +## unavailable at startup. +probe_on_startup = true + ## The storage type used to store the data. ## - `File`: the data is stored in the local file system. ## - `S3`: the data is stored in the S3 object storage. diff --git a/src/datanode/Cargo.toml b/src/datanode/Cargo.toml index 0720a576b9c..45cdc783197 100644 --- a/src/datanode/Cargo.toml +++ b/src/datanode/Cargo.toml @@ -61,6 +61,7 @@ partition.workspace = true prometheus.workspace = true prost.workspace = true query.workspace = true +rand.workspace = true reqwest.workspace = true serde.workspace = true serde_json.workspace = true diff --git a/src/datanode/src/config.rs b/src/datanode/src/config.rs index 34a41ce2801..cc7ebf8a0ce 100644 --- a/src/datanode/src/config.rs +++ b/src/datanode/src/config.rs @@ -44,6 +44,14 @@ pub struct StorageConfig { /// Defaults to `/copy` when `data_home` is a local path. /// Distributed deployments always disable SQL access to local files. pub copy_root: Option, + /// Whether to probe every remote object storage when the datanode starts. + /// + /// The probe writes, reads, overwrites, lists and deletes an object under + /// `.greptime_probe/`, and a failure aborts startup. The probe lives at the + /// storage root, so the credentials must allow writes and deletes below the root. + /// Disable it when the credentials are scoped to narrower prefixes or when the + /// storage may be unavailable at startup. + pub probe_on_startup: bool, #[serde(flatten)] pub store: ObjectStoreConfig, /// Object storage providers @@ -62,6 +70,7 @@ impl Default for StorageConfig { Self { data_home: DEFAULT_DATA_HOME.to_string(), copy_root: None, + probe_on_startup: true, store: ObjectStoreConfig::default(), providers: vec![], } @@ -203,6 +212,19 @@ mod tests { let _parsed: DatanodeOptions = toml::from_str(&toml_string).unwrap(); } + #[test] + fn test_probe_on_startup() { + assert!(DatanodeOptions::default().storage.probe_on_startup); + + let toml_str = r#" + [storage] + type = "File" + probe_on_startup = false + "#; + let opts: DatanodeOptions = toml::from_str(toml_str).unwrap(); + assert!(!opts.storage.probe_on_startup); + } + #[test] fn test_secstr() { let toml_str = r#" diff --git a/src/datanode/src/datanode.rs b/src/datanode/src/datanode.rs index c304aaa6e00..6ebe7ffe05d 100644 --- a/src/datanode/src/datanode.rs +++ b/src/datanode/src/datanode.rs @@ -44,6 +44,7 @@ use common_wal::config::kafka::DatanodeKafkaConfig; use common_wal::config::object_store::{ObjectStoreWalConfig, STANDALONE_GENERATION}; use common_wal::config::raft_engine::RaftEngineConfig; use file_engine::engine::FileRegionEngine; +use futures::future::join_all; use log_store::ObjectStoreLogStore; use log_store::kafka::log_store::KafkaLogStore; use log_store::noop::log_store::NoopLogStore; @@ -466,20 +467,83 @@ impl DatanodeBuilder { .await } - /// Builds [ObjectStoreManager] from [StorageConfig]. - pub async fn build_object_store_manager(cfg: &StorageConfig) -> Result { + /// Builds [ObjectStoreManager] from [StorageConfig] and probes every remote + /// object store it contains, unless `probe_on_startup` is disabled. + /// + /// The probe fails the build when a store rejects any of the operations the + /// engines rely on, so a misconfigured storage surfaces here instead of failing + /// on the first write. + pub async fn build_object_store_manager( + cfg: &StorageConfig, + node_id: u64, + ) -> Result { let object_store = store::new_object_store(cfg.store.clone(), &cfg.data_home).await?; let default_name = cfg.store.config_name(); - let mut object_store_manager = ObjectStoreManager::new(default_name, object_store); - for store in &cfg.providers { - object_store_manager.add( - store.config_name(), - store::new_object_store(store.clone(), &cfg.data_home).await?, - ); + let mut checks = Vec::with_capacity(1 + cfg.providers.len()); + if cfg.store.is_object_storage() { + checks.push((default_name.to_string(), object_store.clone())); } + let mut object_store_manager = ObjectStoreManager::new(default_name, object_store); + for provider in &cfg.providers { + let name = provider.config_name(); + let object_store = store::new_object_store(provider.clone(), &cfg.data_home).await?; + if provider.is_object_storage() { + checks.push((name.to_string(), object_store.clone())); + } + object_store_manager.add(name, object_store); + } + + if cfg.probe_on_startup && !checks.is_empty() { + let prefix = store::probe_key_prefix(node_id); + Self::run_probes(&checks, &prefix).await?; + } + Ok(Arc::new(object_store_manager)) } + /// Probes every store in `checks` under its own key prefix. + /// + /// Providers often share one object storage, so when there is more than one + /// store each probe starts after a random delay rather than all at once. + /// + /// Each check is bounded by [`store::PROBE_TIMEOUT`], so a store that accepts + /// connections but never answers fails startup instead of blocking it, after a + /// best-effort removal of its probe. + async fn run_probes(checks: &[(String, ObjectStore)], prefix: &str) -> Result<()> { + let stagger = checks.len() > 1; + let checks = checks + .iter() + .enumerate() + .map(|(index, (name, object_store))| { + let prefix = format!("{prefix}/{index}"); + async move { + if stagger { + tokio::time::sleep(store::probe_start_delay()).await; + } + match tokio::time::timeout( + store::PROBE_TIMEOUT, + store::probe_object_store(object_store, name, &prefix), + ) + .await + { + Ok(result) => result, + Err(_) => { + store::remove_probe_object(object_store, &prefix).await; + error::ObjectStoreProbeTimeoutSnafu { + name, + timeout: store::PROBE_TIMEOUT, + } + .fail() + } + } + } + }); + for result in join_all(checks).await { + result?; + } + Ok(()) + } + #[cfg(test)] /// Open all regions belong to this datanode. async fn initialize_region_server( @@ -573,7 +637,11 @@ impl DatanodeBuilder { ); region_server.install_remote_dyn_filter_receiver_injector(&self.plugins); - let object_store_manager = Self::build_object_store_manager(&opts.storage).await?; + let object_store_manager = Self::build_object_store_manager( + &opts.storage, + opts.node_id.context(MissingNodeIdSnafu)?, + ) + .await?; let engines = self .build_store_engines( object_store_manager, @@ -1074,6 +1142,7 @@ mod tests { use cache::build_datanode_layered_cache_registry; use common_base::Plugins; use common_base::readable_size::ReadableSize; + use common_base::secrets::SecretString; use common_config::Configurable; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; @@ -1092,6 +1161,7 @@ mod tests { use mito2::engine::MITO_ENGINE_NAME; use mito2::test_util::{CreateRequestBuilder, build_rows, rows_schema}; use object_store::ObjectStore; + use object_store::config::{S3Config, S3Connection}; use object_store::manager::ObjectStoreManager; use partition::manager::PartitionRuleManagerRef; use query::region_query::{RegionQueryHandlerFactory, RegionQueryHandlerFactoryRef}; @@ -1100,7 +1170,7 @@ mod tests { use store_api::region_request::{PathType, RegionOpenRequest, RegionPutRequest, RegionRequest}; use store_api::storage::RegionId; - use crate::config::{DatanodeOptions, RegionEngineConfig, StorageConfig}; + use crate::config::{DatanodeOptions, ObjectStoreConfig, RegionEngineConfig, StorageConfig}; use crate::datanode::{ Datanode, DatanodeBuilder, validate_object_store_wal_config, validate_region_engine_config, wal_object_store, @@ -1460,6 +1530,80 @@ mod tests { ); } + #[tokio::test] + async fn test_object_store_probe_can_be_disabled() { + // The endpoint does not accept connections, so this only succeeds when the + // probe does not run. + let storage = StorageConfig { + store: ObjectStoreConfig::S3(S3Config { + connection: S3Connection { + bucket: "test".to_string(), + endpoint: Some("http://127.0.0.1:1".to_string()), + region: Some("us-east-1".to_string()), + access_key_id: SecretString::from("test".to_string()), + secret_access_key: SecretString::from("test".to_string()), + ..Default::default() + }, + ..Default::default() + }), + probe_on_startup: false, + ..Default::default() + }; + + DatanodeBuilder::build_object_store_manager(&storage, 1) + .await + .unwrap(); + } + + #[tokio::test] + async fn test_file_store_is_not_probed() { + // The probe is enabled by default, but a File store must skip it. + let data_home = create_temp_dir("object-store-probe-file"); + let storage = StorageConfig { + data_home: data_home.path().to_str().unwrap().to_string(), + ..Default::default() + }; + + DatanodeBuilder::build_object_store_manager(&storage, 1) + .await + .unwrap(); + + // The probe would leave this directory behind if it ran. + assert!(!data_home.path().join(".greptime_probe").exists()); + } + + #[tokio::test] + async fn test_object_store_probe_failure_aborts_build() { + // The endpoint does not accept connections, so the probe write fails and the + // build must fail with it. + let storage = StorageConfig { + store: ObjectStoreConfig::S3(S3Config { + connection: S3Connection { + bucket: "test".to_string(), + endpoint: Some("http://127.0.0.1:1".to_string()), + region: Some("us-east-1".to_string()), + access_key_id: SecretString::from("test".to_string()), + secret_access_key: SecretString::from("test".to_string()), + ..Default::default() + }, + ..Default::default() + }), + ..Default::default() + }; + + let err = DatanodeBuilder::build_object_store_manager(&storage, 1) + .await + .unwrap_err(); + assert_matches!( + err, + Error::ObjectStoreProbe { + name, + step: "write", + .. + } if name == storage.store.config_name() + ); + } + struct CountingRegionQueryHandlerFactory(Arc); impl RegionQueryHandlerFactory for CountingRegionQueryHandlerFactory { diff --git a/src/datanode/src/error.rs b/src/datanode/src/error.rs index aaf1ff4d8df..d5037c2874e 100644 --- a/src/datanode/src/error.rs +++ b/src/datanode/src/error.rs @@ -14,6 +14,7 @@ use std::any::Any; use std::sync::Arc; +use std::time::Duration; use common_error::define_into_tonic_status; use common_error::ext::{BoxedError, ErrorExt, RetryHint}; @@ -472,6 +473,24 @@ pub enum Error { location: Location, }, + #[snafu(display("Failed to run object storage probe for provider '{name}' at step '{step}'"))] + ObjectStoreProbe { + name: String, + step: &'static str, + #[snafu(source)] + error: object_store::Error, + #[snafu(implicit)] + location: Location, + }, + + #[snafu(display("Object storage probe for provider '{name}' timed out after {timeout:?}"))] + ObjectStoreProbeTimeout { + name: String, + timeout: Duration, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Not yet implemented: {what}"))] NotYetImplemented { what: String }, } @@ -554,6 +573,9 @@ impl ErrorExt for Error { SerializeJson { .. } => StatusCode::Internal, ObjectStore { source, .. } => source.status_code(), + ObjectStoreProbe { .. } | ObjectStoreProbeTimeout { .. } => { + StatusCode::StorageUnavailable + } } } @@ -569,7 +591,8 @@ impl ErrorExt for Error { RegionNotReady { .. } | RegionNotFound { .. } => RetryHint::RetryAfterRouteRefresh, RegionBusy { .. } | ConcurrentQueryLimiterClosed { .. } - | ConcurrentQueryLimiterTimeout { .. } => RetryHint::Retryable, + | ConcurrentQueryLimiterTimeout { .. } + | ObjectStoreProbeTimeout { .. } => RetryHint::Retryable, NewPlanDecoder { source, .. } | ExecuteLogicalPlan { source, .. } => { source.retry_hint() } @@ -593,6 +616,9 @@ impl ErrorExt for Error { BuildMetricEngine { source, .. } => source.retry_hint(), ListStorageSsts { source, .. } => source.retry_hint(), ObjectStore { source, .. } => source.retry_hint(), + ObjectStoreProbe { error, .. } => { + object_store::error::retry_hint_from_opendal_error(error) + } _ => RetryHint::NonRetryable, } } diff --git a/src/datanode/src/store.rs b/src/datanode/src/store.rs index 3369de5aedc..484a5ec26c2 100644 --- a/src/datanode/src/store.rs +++ b/src/datanode/src/store.rs @@ -14,11 +14,14 @@ //! object storage utilities +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + use common_telemetry::{info, warn}; +use futures::TryStreamExt; use object_store::factory::new_raw_object_store; use object_store::util::{clean_temp_dir, join_dir, with_instrument_layers, with_retry_layers}; -use object_store::{ATOMIC_WRITE_DIR, ObjectStore}; -use snafu::prelude::*; +use object_store::{ATOMIC_WRITE_DIR, Error as ObjectStoreError, ErrorKind, ObjectStore}; +use snafu::{IntoError, prelude::*}; use crate::config::ObjectStoreConfig; use crate::error::{self, Result}; @@ -84,3 +87,217 @@ pub async fn new_object_store(store: ObjectStoreConfig, data_home: &str) -> Resu let object_store = with_instrument_layers(object_store, true); Ok(object_store) } + +/// The content written by [`probe_object_store`]. +const PROBE_CONTENT: &[u8] = b"greptimedb object store probe"; + +/// The content appended by [`probe_object_store`] to exercise read-modify-write. +const PROBE_SUFFIX: &[u8] = b" - read-modify-write"; + +/// The directory that holds the objects created by [`probe_object_store`]. +const PROBE_DIR: &str = ".greptime_probe"; + +/// The budget of one [`probe_object_store`] call, including the retries an object +/// store makes on top of its 30s request timeout. +pub(crate) const PROBE_TIMEOUT: Duration = Duration::from_secs(180); + +/// The maximum random delay before a probe starts. It is short enough that startup +/// waits at most this delay, and it spreads out the probes of providers that share +/// an object storage. +const PROBE_STAGGER: Duration = Duration::from_millis(500); + +/// The budget of the best-effort cleanup after a timed-out probe, kept short so an +/// unresponsive store cannot block startup a second time. +const PROBE_CLEANUP_TIMEOUT: Duration = Duration::from_secs(10); + +/// Returns a random delay in `[0, PROBE_STAGGER)`. +pub(crate) fn probe_start_delay() -> Duration { + Duration::from_nanos(rand::random_range(0..PROBE_STAGGER.as_nanos() as u64)) +} + +/// Returns a unique key prefix for [`probe_object_store`], so datanodes sharing +/// the same bucket and root do not collide. +pub(crate) fn probe_key_prefix(node_id: u64) -> String { + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_nanos(); + format!("{PROBE_DIR}/{node_id}-{}-{nanos}", std::process::id()) +} + +/// Returns the object key [`probe_object_store`] probes under `prefix`. +fn probe_object_key(prefix: &str) -> String { + format!("{prefix}/probe") +} + +/// Checks that `store` supports the operations the storage engines rely on: write, +/// read, overwrite, list and delete. +/// +/// The probe object is removed before returning, after a failed step too. A caller +/// that abandons the check early must call [`remove_probe_object`]. +pub(crate) async fn probe_object_store( + store: &ObjectStore, + name: &str, + prefix: &str, +) -> Result<()> { + let key = probe_object_key(prefix); + + let check = async { + store + .write(&key, PROBE_CONTENT) + .await + .context(error::ObjectStoreProbeSnafu { + name, + step: "write", + })?; + verify_content(store, name, &key, PROBE_CONTENT).await?; + + // Object stores cannot append in place, so this reads the object, appends to + // it and overwrites it, exercising an overwrite of an existing object. + let mut content = store + .read(&key) + .await + .context(error::ObjectStoreProbeSnafu { name, step: "read" })? + .to_vec(); + content.extend_from_slice(PROBE_SUFFIX); + store + .write(&key, content.clone()) + .await + .context(error::ObjectStoreProbeSnafu { + name, + step: "overwrite", + })?; + verify_content(store, name, &key, &content).await?; + + let mut lister = store + .lister_with(&format!("{prefix}/")) + .await + .context(error::ObjectStoreProbeSnafu { name, step: "list" })?; + while let Some(entry) = lister + .try_next() + .await + .context(error::ObjectStoreProbeSnafu { name, step: "list" })? + { + if entry.path() == key { + return Ok(()); + } + } + + Err( + error::ObjectStoreProbeSnafu { name, step: "list" }.into_error(ObjectStoreError::new( + ErrorKind::Unexpected, + format!("'{key}' is not visible in the listing"), + )), + ) + } + .await; + + match store.delete(&key).await { + Ok(()) => check, + Err(err) if check.is_err() => { + warn!( + err; + "Failed to remove the object storage probe object '{}' of provider '{}'", + key, + name + ); + check + } + Err(err) => Err(err).context(error::ObjectStoreProbeSnafu { + name, + step: "delete", + }), + } +} + +async fn verify_content(store: &ObjectStore, name: &str, key: &str, expected: &[u8]) -> Result<()> { + let got = store + .read(key) + .await + .context(error::ObjectStoreProbeSnafu { name, step: "read" })?; + if got.to_vec() != expected { + return Err( + error::ObjectStoreProbeSnafu { name, step: "read" }.into_error(ObjectStoreError::new( + ErrorKind::Unexpected, + format!("content of '{key}' differs from what was written"), + )), + ); + } + Ok(()) +} + +/// Removes the probe object of a probe that did not finish. +/// +/// It has its own short budget, and failures only log: the caller reports the timeout. +pub(crate) async fn remove_probe_object(store: &ObjectStore, prefix: &str) { + let key = probe_object_key(prefix); + match tokio::time::timeout(PROBE_CLEANUP_TIMEOUT, store.delete(&key)).await { + Ok(Ok(())) => {} + Ok(Err(err)) => warn!( + err; + "Failed to remove the object storage probe object '{}'", key + ), + Err(_) => warn!( + "Timed out removing the object storage probe object '{}'", + key + ), + } +} + +#[cfg(test)] +mod tests { + use common_test_util::temp_dir::create_temp_dir; + use object_store::services::Fs; + + use super::*; + + fn new_fs_store(root: &str) -> ObjectStore { + ObjectStore::new(Fs::default().root(root)).unwrap() + } + + #[tokio::test] + async fn test_probe_object_store() { + let dir = create_temp_dir("object-store-probe"); + let store = new_fs_store(dir.path().to_str().unwrap()); + let prefix = probe_key_prefix(1); + + probe_object_store(&store, "fs", &prefix).await.unwrap(); + + // The lister also returns the listed directory itself, so only check that the + // probe object is gone. + let key = format!("{prefix}/probe"); + let mut lister = store.lister_with(&format!("{prefix}/")).await.unwrap(); + while let Some(entry) = lister.try_next().await.unwrap() { + assert_ne!(entry.path(), key); + } + } + + #[tokio::test] + async fn test_probe_object_store_write_failure() { + let dir = create_temp_dir("object-store-probe-write-failure"); + let root = dir.path().join("not-a-directory"); + std::fs::write(&root, b"").unwrap(); + let store = new_fs_store(root.to_str().unwrap()); + + let err = probe_object_store(&store, "broken", &probe_key_prefix(1)) + .await + .unwrap_err(); + assert!( + err.to_string() + .contains("probe for provider 'broken' at step 'write'") + ); + } + + #[tokio::test] + async fn test_remove_probe_object() { + let dir = create_temp_dir("object-store-probe-cleanup"); + let store = new_fs_store(dir.path().to_str().unwrap()); + let prefix = probe_key_prefix(1); + let key = probe_object_key(&prefix); + store.write(&key, PROBE_CONTENT).await.unwrap(); + + remove_probe_object(&store, &prefix).await; + + assert!(store.read(&key).await.is_err()); + } +} diff --git a/tests-integration/src/test_util.rs b/tests-integration/src/test_util.rs index 3dd451b4716..da8d8add687 100644 --- a/tests-integration/src/test_util.rs +++ b/tests-integration/src/test_util.rs @@ -578,6 +578,7 @@ pub(crate) fn create_datanode_opts( storage: StorageConfig { data_home: home_dir, copy_root: None, + probe_on_startup: false, providers, store: default_store, }, diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index c13ada351ea..c7ba62a473f 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -3390,6 +3390,7 @@ pub async fn test_config_api(store_type: StorageType) { let storage = if store_type != StorageType::File { format!( r#"[storage] +probe_on_startup = false type = "{}" providers = [] @@ -3404,6 +3405,7 @@ skip_ssl_validation = false"#, } else { format!( r#"[storage] +probe_on_startup = false type = "{}" providers = []"#, store_type diff --git a/tests/conf/datanode-test.toml.template b/tests/conf/datanode-test.toml.template index 907fcee4ddf..7b0dcebf081 100644 --- a/tests/conf/datanode-test.toml.template +++ b/tests/conf/datanode-test.toml.template @@ -33,6 +33,7 @@ overwrite_entry_start_id = true [storage] type = 'File' data_home = '{data_home}' +probe_on_startup = false [meta_client_options] metasrv_addrs = ['{addrs.metasrv_addr}'] diff --git a/tests/conf/standalone-test.toml.template b/tests/conf/standalone-test.toml.template index 8d7e307d511..149d53ffde6 100644 --- a/tests/conf/standalone-test.toml.template +++ b/tests/conf/standalone-test.toml.template @@ -28,6 +28,7 @@ broker_endpoints = {kafka_wal_broker_endpoints | unescaped} type = 'File' data_home = '{data_home}' copy_root = '{copy_root}' +probe_on_startup = false [grpc] bind_addr = '{addrs.grpc_addr}' From efed847ca6dadd37a76fbbaca2a454621c505a92 Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Wed, 28 Oct 2026 17:09:33 +0000 Subject: [PATCH 133/239] fix!: enforce sparse primary key schema (#10020) * fix: preserve required sorting for ordinary sparse scans Signed-off-by: shuiyisong * fix: enforce sparse primary key schema Signed-off-by: shuiyisong * test: fix Signed-off-by: shuiyisong --------- Signed-off-by: shuiyisong --- .../function/src/scalars/primary_key.rs | 16 ++- src/mito2/benches/bench_index_update.rs | 3 +- src/mito2/benches/bench_pk_tag_column.rs | 6 +- src/mito2/src/engine/create_test.rs | 40 ++++++ src/mito2/src/memtable/bulk/part.rs | 102 ++++--------- src/mito2/src/read/compat.rs | 64 +++++++-- src/mito2/src/sst/parquet/flat_format.rs | 9 +- src/store-api/src/metadata.rs | 134 +++++++++++++++++- src/table/src/table/scan.rs | 14 ++ .../create/create_with_compression.result | 2 +- 10 files changed, 294 insertions(+), 96 deletions(-) diff --git a/src/common/function/src/scalars/primary_key.rs b/src/common/function/src/scalars/primary_key.rs index f090575f166..37dc5bacde5 100644 --- a/src/common/function/src/scalars/primary_key.rs +++ b/src/common/function/src/scalars/primary_key.rs @@ -354,7 +354,11 @@ mod tests { builder.primary_key_encoding(encoding); for (id, name, ty) in columns { builder.push_column_metadata(ColumnMetadata { - column_schema: ColumnSchema::new((*name).to_string(), ty.clone(), true), + column_schema: ColumnSchema::new( + (*name).to_string(), + ty.clone(), + !ReservedColumnId::is_reserved(*id), + ), semantic_type: SemanticType::Tag, column_id: *id, }); @@ -440,6 +444,16 @@ mod tests { #[test] fn test_decode_sparse_primary_key() { let columns = vec![ + ( + ReservedColumnId::table_id(), + "__table_id", + ConcreteDataType::uint32_datatype(), + ), + ( + ReservedColumnId::tsid(), + "__tsid", + ConcreteDataType::uint64_datatype(), + ), (10, "k0", ConcreteDataType::string_datatype()), (11, "k1", ConcreteDataType::string_datatype()), ]; diff --git a/src/mito2/benches/bench_index_update.rs b/src/mito2/benches/bench_index_update.rs index a76241fa53f..0d0b1ef6c23 100644 --- a/src/mito2/benches/bench_index_update.rs +++ b/src/mito2/benches/bench_index_update.rs @@ -60,7 +60,8 @@ fn input(shape: &Shape) -> (RegionMetadataRef, HashSet, RecordBatch) { let mut pk = vec![ReservedColumnId::table_id(), ReservedColumnId::tsid()]; let mut inverted_columns = HashSet::new(); let mut add_tag = |id, name: String, data_type, indexed| { - let mut schema = ColumnSchema::new(name, data_type, true).with_inverted_index(false); + let mut schema = ColumnSchema::new(name, data_type, !ReservedColumnId::is_reserved(id)) + .with_inverted_index(false); if indexed && shape.inverted { inverted_columns.insert(id); } diff --git a/src/mito2/benches/bench_pk_tag_column.rs b/src/mito2/benches/bench_pk_tag_column.rs index 9e191a69511..2279067dfdc 100644 --- a/src/mito2/benches/bench_pk_tag_column.rs +++ b/src/mito2/benches/bench_pk_tag_column.rs @@ -77,7 +77,11 @@ fn metadata(tags: u32) -> RegionMetadataRef { primary_key.push(column_id); builder.push_column_metadata(ColumnMetadata { column_id, - column_schema: ColumnSchema::new(name, data_type, true), + column_schema: ColumnSchema::new( + name, + data_type, + !ReservedColumnId::is_reserved(column_id), + ), semantic_type: SemanticType::Tag, }); } diff --git a/src/mito2/src/engine/create_test.rs b/src/mito2/src/engine/create_test.rs index 05696ad6109..57c0c3fa7c2 100644 --- a/src/mito2/src/engine/create_test.rs +++ b/src/mito2/src/engine/create_test.rs @@ -117,6 +117,46 @@ async fn test_engine_column_compression_validation() { assert_eq!(StatusCode::InvalidArguments, err.status_code()); } +#[tokio::test] +async fn test_engine_create_sparse_requires_reserved_columns() { + let mut env = TestEnv::new().await; + let engine = env.create_engine(MitoConfig::default()).await; + let region_id = RegionId::new(1, 1); + let mut request = CreateRequestBuilder::new() + .insert_option("primary_key_encoding", "sparse") + .insert_option("sst_format", "flat") + .build(); + let error = engine + .handle_request(region_id, RegionRequest::Create(request.clone())) + .await + .unwrap_err(); + assert_eq!(StatusCode::InvalidArguments, error.status_code()); + assert!( + error + .output_msg() + .contains("sparse primary keys must start with (__table_id, __tsid)") + ); + assert!(!engine.is_region_exists(region_id)); + + let metadata = sst_region_metadata_with_encoding(PrimaryKeyEncoding::Sparse); + request.column_metadatas = metadata.column_metadatas; + request.primary_key = metadata.primary_key; + engine + .handle_request(region_id, RegionRequest::Create(request.clone())) + .await + .unwrap(); + let expected = engine.get_region(region_id).unwrap().metadata(); + reopen_region( + &engine, + region_id, + request.table_dir, + false, + request.options, + ) + .await; + assert_eq!(expected, engine.get_region(region_id).unwrap().metadata()); +} + #[tokio::test] async fn test_engine_create_new_region() { test_engine_create_new_region_with_format(false).await; diff --git a/src/mito2/src/memtable/bulk/part.rs b/src/mito2/src/memtable/bulk/part.rs index 9bcc4238b45..5852d072197 100644 --- a/src/mito2/src/memtable/bulk/part.rs +++ b/src/mito2/src/memtable/bulk/part.rs @@ -1858,19 +1858,17 @@ impl MultiBulkPart { #[cfg(test)] mod tests { - use api::v1::{Row, SemanticType, WriteHint}; + use api::v1::{Row, WriteHint}; use datafusion_common::ScalarValue; use datatypes::arrow::array::{ BinaryArray, DictionaryArray, Float64Array, TimestampMillisecondArray, }; use datatypes::arrow::datatypes::UInt32Type; use datatypes::prelude::{ConcreteDataType, Value}; - use datatypes::schema::ColumnSchema; use mito_codec::row_converter::build_primary_key_codec; use rand::rngs::StdRng; use rand::{Rng, SeedableRng}; - use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; - use store_api::storage::RegionId; + use store_api::metadata::RegionMetadataBuilder; use store_api::storage::consts::ReservedColumnId; use table::predicate::Predicate; @@ -2594,47 +2592,30 @@ mod tests { KeyValues::new(metadata.as_ref(), mutation).unwrap() } + fn sparse_metadata_for_test() -> RegionMetadataRef { + let mut metadata = metadata_for_test().as_ref().clone(); + metadata.column_metadatas[1].column_schema.data_type = ConcreteDataType::string_datatype(); + let mut builder = RegionMetadataBuilder::from_existing(metadata); + let reserved = crate::test_util::sst_util::sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + ); + for column in reserved.column_metadatas.into_iter().take(2) { + builder.push_column_metadata(column); + } + builder + .primary_key(vec![ + ReservedColumnId::table_id(), + ReservedColumnId::tsid(), + 0, + 1, + ]) + .primary_key_encoding(PrimaryKeyEncoding::Sparse); + Arc::new(builder.build().unwrap()) + } + #[test] fn test_bulk_part_converter_sparse_primary_key_encoding() { - use api::v1::SemanticType; - use datatypes::schema::ColumnSchema; - use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; - use store_api::storage::RegionId; - - let mut builder = RegionMetadataBuilder::new(RegionId::new(123, 456)); - builder - .push_column_metadata(ColumnMetadata { - column_schema: ColumnSchema::new("k0", ConcreteDataType::string_datatype(), false), - semantic_type: SemanticType::Tag, - column_id: 0, - }) - .push_column_metadata(ColumnMetadata { - column_schema: ColumnSchema::new("k1", ConcreteDataType::string_datatype(), false), - semantic_type: SemanticType::Tag, - column_id: 1, - }) - .push_column_metadata(ColumnMetadata { - column_schema: ColumnSchema::new( - "ts", - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ), - semantic_type: SemanticType::Timestamp, - column_id: 2, - }) - .push_column_metadata(ColumnMetadata { - column_schema: ColumnSchema::new("v0", ConcreteDataType::int64_datatype(), true), - semantic_type: SemanticType::Field, - column_id: 3, - }) - .push_column_metadata(ColumnMetadata { - column_schema: ColumnSchema::new("v1", ConcreteDataType::float64_datatype(), true), - semantic_type: SemanticType::Field, - column_id: 4, - }) - .primary_key(vec![0, 1]) - .primary_key_encoding(PrimaryKeyEncoding::Sparse); - let metadata = Arc::new(builder.build().unwrap()); + let metadata = sparse_metadata_for_test(); let primary_key_codec = build_primary_key_codec(&metadata); let schema = to_flat_sst_arrow_schema( @@ -2923,40 +2904,7 @@ mod tests { #[test] fn test_convert_bulk_part_sparse_encoding() { - let mut builder = RegionMetadataBuilder::new(RegionId::new(123, 456)); - builder - .push_column_metadata(ColumnMetadata { - column_schema: ColumnSchema::new("k0", ConcreteDataType::string_datatype(), false), - semantic_type: SemanticType::Tag, - column_id: 0, - }) - .push_column_metadata(ColumnMetadata { - column_schema: ColumnSchema::new("k1", ConcreteDataType::string_datatype(), false), - semantic_type: SemanticType::Tag, - column_id: 1, - }) - .push_column_metadata(ColumnMetadata { - column_schema: ColumnSchema::new( - "ts", - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ), - semantic_type: SemanticType::Timestamp, - column_id: 2, - }) - .push_column_metadata(ColumnMetadata { - column_schema: ColumnSchema::new("v0", ConcreteDataType::int64_datatype(), true), - semantic_type: SemanticType::Field, - column_id: 3, - }) - .push_column_metadata(ColumnMetadata { - column_schema: ColumnSchema::new("v1", ConcreteDataType::float64_datatype(), true), - semantic_type: SemanticType::Field, - column_id: 4, - }) - .primary_key(vec![0, 1]) - .primary_key_encoding(PrimaryKeyEncoding::Sparse); - let metadata = Arc::new(builder.build().unwrap()); + let metadata = sparse_metadata_for_test(); let primary_key_codec = build_primary_key_codec(&metadata); diff --git a/src/mito2/src/read/compat.rs b/src/mito2/src/read/compat.rs index 6f9a4a36e50..1d4a7f5457b 100644 --- a/src/mito2/src/read/compat.rs +++ b/src/mito2/src/read/compat.rs @@ -1344,8 +1344,19 @@ mod tests { ], &[], ); + let reserved = crate::test_util::sst_util::sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + ); + actual_metadata.primary_key = reserved.primary_key.into_iter().take(2).collect(); + actual_metadata + .column_metadatas + .extend(reserved.column_metadatas.into_iter().take(2)); actual_metadata.primary_key_encoding = PrimaryKeyEncoding::Sparse; - let actual_metadata = Arc::new(actual_metadata); + let actual_metadata = Arc::new( + RegionMetadataBuilder::from_existing(actual_metadata) + .build() + .unwrap(), + ); let mut expected_metadata = (*actual_metadata).clone(); for column in expected_metadata.column_metadatas.iter_mut() { @@ -1373,7 +1384,10 @@ mod tests { .unwrap() .unwrap(); - let sparse_k1 = encode_sparse_key(&[]); + let mut sparse_k1 = Vec::new(); + SparsePrimaryKeyCodec::schemaless() + .encode_internal(42, 7, &mut sparse_k1) + .unwrap(); let input_columns: Vec = vec![ Arc::new(Int64Array::from(vec![100, 200])), Arc::new(TimestampMillisecondArray::from_iter_values([1000, 2000])), @@ -1384,8 +1398,10 @@ mod tests { OpType::Put as u8, ])), ]; - let input_schema = - to_flat_sst_arrow_schema(&actual_metadata, &FlatSchemaOptions::default()); + let input_schema = to_flat_sst_arrow_schema( + &actual_metadata, + &FlatSchemaOptions::from_encoding(PrimaryKeyEncoding::Sparse), + ); let input_batch = RecordBatch::try_new(input_schema, input_columns).unwrap(); let result = compat_batch.compat(input_batch).unwrap(); @@ -1402,8 +1418,10 @@ mod tests { OpType::Put as u8, ])), ]; - let output_schema = - to_flat_sst_arrow_schema(&expected_metadata, &FlatSchemaOptions::default()); + let output_schema = to_flat_sst_arrow_schema( + &expected_metadata, + &FlatSchemaOptions::from_encoding(PrimaryKeyEncoding::Sparse), + ); let expected_batch = RecordBatch::try_new(output_schema, expected_columns).unwrap(); assert_eq!(expected_batch, result); @@ -1520,6 +1538,16 @@ mod tests { columns[0] = json; let batch = RecordBatch::try_new(Arc::new(Schema::new(fields)), columns)?; for encoding in [PrimaryKeyEncoding::Dense, PrimaryKeyEncoding::Sparse] { + let batch = if encoding == PrimaryKeyEncoding::Sparse { + let mut key = Vec::new(); + SparsePrimaryKeyCodec::schemaless().encode_internal(42, 7, &mut key)?; + let mut columns = batch.columns().to_vec(); + columns[primary_key_column_index(batch.num_columns())] = + build_flat_test_pk_array(&[&key, &key]); + RecordBatch::try_new(batch.schema(), columns)? + } else { + batch.clone() + }; for missing_json in [false, true] { let source = if missing_json { new_metadata( @@ -1533,12 +1561,12 @@ mod tests { } else { metadata.as_ref().clone() }; - let mut builder = RegionMetadataBuilder::from_existing(source); - builder.primary_key_encoding(encoding); - let actual = Arc::new(builder.build()?); - let mut builder = RegionMetadataBuilder::from_existing(metadata.as_ref().clone()); - builder.primary_key_encoding(encoding); - builder.push_column_metadata(ColumnMetadata { + let mut actual_builder = RegionMetadataBuilder::from_existing(source); + actual_builder.primary_key_encoding(encoding); + let mut expected_builder = + RegionMetadataBuilder::from_existing(metadata.as_ref().clone()); + expected_builder.primary_key_encoding(encoding); + expected_builder.push_column_metadata(ColumnMetadata { column_schema: ColumnSchema::new( "added", ConcreteDataType::int64_datatype(), @@ -1550,7 +1578,17 @@ mod tests { semantic_type: SemanticType::Field, column_id: 2, }); - let expected = Arc::new(builder.build()?); + if encoding == PrimaryKeyEncoding::Sparse { + let reserved = test_util::sst_util::sst_region_metadata_with_encoding(encoding); + for builder in [&mut actual_builder, &mut expected_builder] { + builder.primary_key(reserved.primary_key[..2].to_vec()); + for column in reserved.column_metadatas.iter().take(2) { + builder.push_column_metadata(column.clone()); + } + } + } + let actual = Arc::new(actual_builder.build()?); + let expected = Arc::new(expected_builder.build()?); let read_columns = ReadColumns::new([0, 1, 2]); let mapper = FlatProjectionMapper::new_with_json2_rewrite_targets( &expected, diff --git a/src/mito2/src/sst/parquet/flat_format.rs b/src/mito2/src/sst/parquet/flat_format.rs index 3abba4d2db6..3f40f7ab3a1 100644 --- a/src/mito2/src/sst/parquet/flat_format.rs +++ b/src/mito2/src/sst/parquet/flat_format.rs @@ -1355,7 +1355,14 @@ mod tests { column_id: col_id, }); - let primary_key: Vec = (0..num_tags as u32).collect(); + let mut primary_key: Vec = (0..num_tags as u32).collect(); + if encoding == PrimaryKeyEncoding::Sparse { + let reserved = crate::test_util::sst_util::sst_region_metadata_with_encoding(encoding); + primary_key.splice(0..0, reserved.primary_key.into_iter().take(2)); + for column in reserved.column_metadatas.into_iter().take(2) { + builder.push_column_metadata(column); + } + } builder.primary_key(primary_key); builder.primary_key_encoding(encoding); builder.build().unwrap() diff --git a/src/store-api/src/metadata.rs b/src/store-api/src/metadata.rs index 55ff240ee4a..f2bac9785f3 100644 --- a/src/store-api/src/metadata.rs +++ b/src/store-api/src/metadata.rs @@ -31,6 +31,7 @@ use common_macro::stack_trace_debug; use datatypes::arrow; use datatypes::arrow::datatypes::FieldRef; use datatypes::extension::json::json2_metadata_with_updated_settings; +use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, FulltextOptions, Schema, SchemaRef}; use datatypes::types::TimestampType; use itertools::Itertools; @@ -39,10 +40,11 @@ use serde::{Deserialize, Deserializer, Serialize}; use snafu::{Location, OptionExt, ResultExt, Snafu, ensure}; use crate::codec::PrimaryKeyEncoding; +use crate::metric_engine_consts::{DATA_SCHEMA_TABLE_ID_COLUMN_NAME, DATA_SCHEMA_TSID_COLUMN_NAME}; use crate::region_request::{ AddColumn, AddColumnLocation, AlterKind, ModifyColumnType, SetIndexOption, UnsetIndexOption, }; -use crate::storage::consts::is_internal_column; +use crate::storage::consts::{ReservedColumnId, is_internal_column}; use crate::storage::{ColumnId, RegionId}; pub type Result = std::result::Result; @@ -476,6 +478,16 @@ impl RegionMetadata { } ); + if self.primary_key_encoding == PrimaryKeyEncoding::Sparse { + ensure!( + self.primary_key + .starts_with(&[ReservedColumnId::table_id(), ReservedColumnId::tsid()]), + InvalidMetaSnafu { + reason: "sparse primary keys must start with (__table_id, __tsid)", + } + ); + } + if !self.primary_key.is_empty() { let mut pk_ids = HashSet::with_capacity(self.primary_key.len()); // Checks column ids in the primary key is valid. @@ -523,6 +535,27 @@ impl RegionMetadata { } ); + if self.primary_key_encoding == PrimaryKeyEncoding::Sparse { + let schema = &column.column_schema; + let valid = if *column_id == ReservedColumnId::table_id() { + schema.name == DATA_SCHEMA_TABLE_ID_COLUMN_NAME + && schema.data_type == ConcreteDataType::uint32_datatype() + && !schema.is_nullable() + } else if *column_id == ReservedColumnId::tsid() { + schema.name == DATA_SCHEMA_TSID_COLUMN_NAME + && schema.data_type == ConcreteDataType::uint64_datatype() + && !schema.is_nullable() + } else { + !ReservedColumnId::is_reserved(*column_id) && schema.data_type.is_string() + }; + ensure!( + valid, + InvalidMetaSnafu { + reason: format!("invalid sparse primary key column {}", schema.name), + } + ); + } + pk_ids.insert(column_id); } } @@ -1388,6 +1421,105 @@ mod test { builder.build().unwrap() } + #[test] + fn test_sparse_primary_key_requires_reserved_columns() { + for primary_key in [vec![], vec![1]] { + let mut metadata = build_test_region_metadata(); + if primary_key.is_empty() { + metadata.column_metadatas[0].semantic_type = SemanticType::Field; + } + metadata.primary_key = primary_key; + let mut builder = RegionMetadataBuilder::from_existing(metadata); + builder.primary_key_encoding(PrimaryKeyEncoding::Sparse); + let error = builder.build().unwrap_err(); + assert_eq!(error.status_code(), StatusCode::InvalidArguments); + assert!( + error + .to_string() + .contains("sparse primary keys must start with (__table_id, __tsid)") + ); + } + } + + #[test] + fn test_sparse_reserved_primary_key_contract() { + let mut builder = RegionMetadataBuilder::from_existing(build_test_region_metadata()); + builder.primary_key_encoding(PrimaryKeyEncoding::Sparse); + builder.column_metadatas[0].column_schema.data_type = ConcreteDataType::string_datatype(); + for (name, column_id, data_type) in [ + ( + DATA_SCHEMA_TABLE_ID_COLUMN_NAME, + ReservedColumnId::table_id(), + ConcreteDataType::uint32_datatype(), + ), + ( + DATA_SCHEMA_TSID_COLUMN_NAME, + ReservedColumnId::tsid(), + ConcreteDataType::uint64_datatype(), + ), + ] { + builder.push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new(name, data_type, false), + semantic_type: SemanticType::Tag, + column_id, + }); + } + builder.primary_key(vec![ + ReservedColumnId::table_id(), + ReservedColumnId::tsid(), + 1, + ]); + let metadata = builder.build().unwrap(); + let json = serde_json::to_string(&metadata).unwrap(); + assert_eq!( + metadata, + serde_json::from_str::(&json).unwrap() + ); + + for case in 0..12 { + let mut builder = RegionMetadataBuilder::from_existing(metadata.clone()); + match case { + 0 => builder.primary_key.swap(0, 1), + 1 => { + builder.primary_key.remove(0); + builder.column_metadatas.remove(3); + } + 2 => builder.column_metadatas[3].column_schema.set_nullable(), + 3 => { + builder.column_metadatas[3].column_schema.data_type = + ConcreteDataType::uint64_datatype() + } + 4 => builder.column_metadatas[4].column_schema.name = "wrong_name".into(), + 5 => { + builder.column_metadatas[0].column_schema.data_type = + ConcreteDataType::int64_datatype() + } + 6 => { + builder.primary_key.remove(1); + builder.column_metadatas.remove(4); + } + 7 => builder.column_metadatas[4].column_schema.set_nullable(), + 8 => { + builder.column_metadatas[4].column_schema.data_type = + ConcreteDataType::uint32_datatype() + } + 9 => builder.column_metadatas[3].semantic_type = SemanticType::Field, + 10 => builder.column_metadatas[4].semantic_type = SemanticType::Field, + 11 => { + builder.primary_key[0] = 10; + builder.primary_key[1] = 11; + builder.column_metadatas[3].column_id = 10; + builder.column_metadatas[4].column_id = 11; + } + _ => unreachable!(), + } + assert_eq!( + builder.build().unwrap_err().status_code(), + StatusCode::InvalidArguments + ); + } + } + #[test] fn test_region_metadata() { let region_metadata = build_test_region_metadata(); diff --git a/src/table/src/table/scan.rs b/src/table/src/table/scan.rs index cae9f5f042c..2dd501e678f 100644 --- a/src/table/src/table/scan.rs +++ b/src/table/src/table/scan.rs @@ -717,6 +717,7 @@ mod test { use datatypes::schema::{ColumnSchema, Schema, SchemaRef}; use datatypes::vectors::{Int32Vector, TimestampMillisecondVector, VectorRef}; use futures::TryStreamExt; + use store_api::codec::PrimaryKeyEncoding; use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; use store_api::metric_engine_consts::DATA_SCHEMA_TABLE_ID_COLUMN_NAME; use store_api::region_engine::{ @@ -1108,7 +1109,14 @@ mod test { /// order or hash partition columns. #[test] fn test_metric_projection_claims_only_projected_row_key_prefix() { + for encoding in [PrimaryKeyEncoding::Dense, PrimaryKeyEncoding::Sparse] { + check_metric_projection_ordering(encoding); + } + } + + fn check_metric_projection_ordering(encoding: PrimaryKeyEncoding) { let mut builder = RegionMetadataBuilder::new(RegionId::new(1234, 5691)); + builder.primary_key_encoding(encoding); for (name, column_id) in [ ( DATA_SCHEMA_TABLE_ID_COLUMN_NAME, @@ -1139,6 +1147,12 @@ mod test { let metadata = Arc::new(builder.build().unwrap()); for (projection, partitioned_by_tsid, claimed, unclaimed) in [ + ( + vec![DATA_SCHEMA_TABLE_ID_COLUMN_NAME, "A", "ts"], + false, + vec![vec![DATA_SCHEMA_TABLE_ID_COLUMN_NAME]], + vec![vec!["A"], vec![DATA_SCHEMA_TABLE_ID_COLUMN_NAME, "ts"]], + ), // The logical projection drops the row key prefix, so `A` is not a prefix of the // row key and neither `A` nor `ts` is ordered. ( diff --git a/tests/cases/standalone/common/create/create_with_compression.result b/tests/cases/standalone/common/create/create_with_compression.result index 77a21288c94..a1c91a6e7ea 100644 --- a/tests/cases/standalone/common/create/create_with_compression.result +++ b/tests/cases/standalone/common/create/create_with_compression.result @@ -109,7 +109,7 @@ CREATE TABLE bad ( host STRING PRIMARY KEY COMPRESSION WITH (level = 3) ) WITH ('primary_key_encoding' = 'sparse'); -Error: 1004(InvalidArguments), Invalid region options, COMPRESSION on tag column host requires flat SST format and dense primary key encoding +Error: 1004(InvalidArguments), Invalid metadata, sparse primary keys must start with (__table_id, __tsid) CREATE TABLE compressed_tags ( ts TIMESTAMP TIME INDEX, From e87766d7ef64a33ea02fce31d2cb2700903f0cef Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Thu, 29 Oct 2026 00:00:07 +0000 Subject: [PATCH 134/239] ci: reduce repeated Windows and fuzz tests on pull requests (#10043) * ci: reduce repeated Windows and fuzz tests on pull requests Signed-off-by: WenyXu * ci: use the nightly merge label as the enqueue opt-in Signed-off-by: WenyXu * ci: log nightly merge decisions for dry-run verification Signed-off-by: WenyXu * ci: gate expensive jobs on GitHub script tests Signed-off-by: WenyXu * ci: grant pull request write permission for nightly label sync Signed-off-by: WenyXu * ci: run platform checks on schedule or manual dispatch Signed-off-by: WenyXu * test(ci): update PR gate count for scheduled platform checks Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- .../actions/github-script-tests/action.yml | 52 +++++++++ .github/ci/README.md | 12 +- .github/heavy-ci.md | 28 +++++ .github/nightly-merge.md | 6 +- .github/scripts/heavy-ci-policy.test.cjs | 105 ++++++++++++++++++ .github/scripts/tests/nightly-merge.test.cjs | 23 +++- .github/scripts/tests/test_ci_slash.py | 2 +- .github/workflows/checks.yml | 41 +------ .github/workflows/integration.yml | 62 +++++++---- .github/workflows/nightly-merge.yml | 14 ++- .github/workflows/rust.yml | 71 ++++-------- 11 files changed, 290 insertions(+), 126 deletions(-) create mode 100644 .github/actions/github-script-tests/action.yml create mode 100644 .github/heavy-ci.md create mode 100644 .github/scripts/heavy-ci-policy.test.cjs diff --git a/.github/actions/github-script-tests/action.yml b/.github/actions/github-script-tests/action.yml new file mode 100644 index 00000000000..4e3d797a95c --- /dev/null +++ b/.github/actions/github-script-tests/action.yml @@ -0,0 +1,52 @@ +name: GitHub Script Tests +description: Validate CI scripts before starting expensive jobs +runs: + using: composite + steps: + - uses: actions/setup-node@v4 + with: + node-version: 24 + - name: Run Node GitHub script tests + shell: bash + run: node --test .github/scripts/query-regression-comment.test.cjs .github/scripts/heavy-ci-policy.test.cjs .github/scripts/tests/nightly-merge.test.cjs + - name: Setup uv + uses: astral-sh/setup-uv@v3 + with: + version: "latest" + - name: Run check-version script tests + shell: bash + run: uv run --no-project python .github/scripts/check-version-test.py + - name: Run release Linux architecture tests + shell: bash + run: python3 .github/scripts/release-linux-architectures-test.py + - name: Run previous-release-tag script tests + shell: bash + run: python3 .github/scripts/previous-release-tag-test.py + - name: Run enterprise license check tests + shell: bash + run: python3 scripts/check-enterprise-license-test.py + - name: Run fuzz orchestration script tests + shell: bash + run: .github/scripts/run-fuzz-targets-test.sh + - name: Run CI comment-command tests + shell: bash + run: python3 .github/scripts/tests/test_ci_slash.py + - name: Run CI scope gate tests + shell: bash + run: python3 .github/ci/tests/test_ci_scope.py + - name: Run metric snapshot benchmark tests + shell: bash + run: python3 -m unittest discover -s .github/scripts/tests -p 'test_metric_snapshot*.py' + - name: Run query-regression Python tooling tests + shell: bash + run: | + python3 tests/perf/test_query_regression_runner_compaction_toctou.py + python3 tests/perf/test_query_regression_runner_otlp_trace_load.py + python3 tests/perf/test_query_regression_summary_otlp.py + python3 tests/perf/test_query_regression_case_selection.py + python3 tests/perf/test_query_regression_nightly_refs.py + python3 tests/perf/test_query_regression_slash.py + python3 tests/perf/test_aliyun_ecs_runner_scripts.py + python3 tests/perf/test_ci_runner_scripts.py + python3 tests/perf/test_agent_observability_summary.py + python3 tests/perf/test_tracesbench_summary.py diff --git a/.github/ci/README.md b/.github/ci/README.md index 07de1fcae7b..1ead30842ec 100644 --- a/.github/ci/README.md +++ b/.github/ci/README.md @@ -1,9 +1,10 @@ # Integration CI admission gate -`integration.yml` runs its expensive jobs — fuzz, export/import E2E, the sqlness -KvBackend legs — on every non-draft pull request. Most pull requests cannot -affect any of them. This directory decides which of those jobs start, using only -the diff. +`integration.yml` filters its expensive jobs — fuzz, export/import E2E, and the +sqlness KvBackend legs — by the diff. Fuzz additionally follows the event policy +in [Heavy CI scheduling](../heavy-ci.md): PR opening or ready-for-review events, +scheduled runs, and explicit manual profiles. A scope verdict cannot override +that event policy. The gate can only *skip* work, never add it, and it fails open: when the verdict is missing or uncertain, every gated job runs. @@ -53,7 +54,8 @@ three minutes, less than the cost of getting its verdict wrong. First match wins: 1. **Forced** — `merge_group`, `schedule`, `workflow_dispatch`, the `ci:full` - label or `CI_SCOPE_FORCE_ALL`. Everything runs. + label or `CI_SCOPE_FORCE_ALL`. All scope outputs allow execution; job event + conditions and manual profile selection still apply. 2. **Too large to judge** — more than 300 changed files. Everything runs. 3. **Always-run file changed** — `Cargo.toml`, `Cargo.lock`, `tests-fuzz/**`, this directory, or anything else that changes what the suites build or how diff --git a/.github/heavy-ci.md b/.github/heavy-ci.md new file mode 100644 index 00000000000..4b0d17bc3a3 --- /dev/null +++ b/.github/heavy-ci.md @@ -0,0 +1,28 @@ +# Heavy CI scheduling + +Checks, Rust CI, and Integration CI each run the shared GitHub script test +action before expensive builds and tests. A failed, skipped, or cancelled +script-test job blocks those jobs, including manually dispatched runs. Each +workflow runs its own preflight so it does not wait on another workflow run. + +Rust CI keeps the Windows compilation check on PRs and merge groups. Windows +nextest and Windows SQLness run in the existing Nightly CI workflow, which also +supports manual dispatch. Its existing schedule is Tuesday through Friday at +07:00 Asia/Shanghai, not every day. No new Windows test job is added. + +Integration CI admits fuzz when a non-draft PR is opened or marked ready for +review, selecting suites with the existing diff-scope classifier. Its weekday +schedule (23:00 Asia/Shanghai, Monday through Friday) runs all fuzz suites. Ordinary pushes and reopening do not run fuzz; core +integration checks still run. Merge groups continue to skip fuzz. + +Re-running an ordinary push-triggered Integration CI run does not enable fuzz. +A new push can cancel an earlier ready-for-review run. Use an explicit fuzz +profile through the CI entry points below when later changes need fuzz coverage. + +For a same-repository branch, use **Run workflow** and select `fuzz_profile` +(`standalone`, `distributed`, `chaos`, or `all`). The default `none` runs core +integration without fuzz. The existing `/ci fuzz ...` commands remain available +for same-repository draft PRs with their existing authorization rules. + +Validate job admission with `node --test .github/scripts/heavy-ci-policy.test.cjs`. +These tests verify triggering decisions, not runner availability or performance. diff --git a/.github/nightly-merge.md b/.github/nightly-merge.md index a8699a9b159..08b523dbab2 100644 --- a/.github/nightly-merge.md +++ b/.github/nightly-merge.md @@ -2,14 +2,16 @@ The PR template opts new PRs into the `merge nightly` label. The Nightly Merge workflow reconciles the label on PR creation, edits, and reopening. Existing PRs -without the checkbox are not opted in. Unchecking removes the label; it does not +without the checkbox can opt in by adding the label manually; label sync leaves +their labels unchanged. Unchecking removes the label; it does not request an immediate merge or cancel an existing queue entry/auto-merge request. Every day at 02:10, 03:10, 04:10, 05:10, 06:10, 07:10, and 08:10 Asia/Shanghai (UTC+8), the workflow enqueues up to five eligible PRs targeting `main` per run, oldest first. Already queued PRs do not count toward this limit. GitHub may delay scheduled runs. -PRs must still have the checkbox and label, be approved and clean, and have all +The label alone controls nightly opt-in; the enqueue job does not check the PR +body. PRs must be approved and clean, and have all review threads resolved. Drafts are skipped. Already queued PRs and manually enabled auto-merge requests are left alone; Maintainers can use **Merge when ready** to proceed earlier. Unready PRs wait for a later diff --git a/.github/scripts/heavy-ci-policy.test.cjs b/.github/scripts/heavy-ci-policy.test.cjs new file mode 100644 index 00000000000..093ff8c636b --- /dev/null +++ b/.github/scripts/heavy-ci-policy.test.cjs @@ -0,0 +1,105 @@ +const { readFileSync } = require('node:fs'); +const { resolve } = require('node:path'); +const assert = require('node:assert/strict'); +const { test } = require('node:test'); +const workflows = resolve(__dirname, '../workflows'); +const integration = readFileSync(resolve(workflows, 'integration.yml'), 'utf8'); +const suites = ['fuzztest', 'unstable-fuzztest', 'distributed-fuzztest', 'distributed-fuzztest-with-chaos']; +const prerequisites = ['prepare-fuzz-targets', 'build-greptime-ci']; +function enabled(job, { event = 'pull_request', action = 'synchronize', draft = false, attempt = 1, profile = 'none', repository = 'GreptimeTeam/greptimedb', outputs = {}, results = {}, cancelled = false, workflow = integration } = {}) { + const block = workflow.split(` ${job}:\n`)[1].split(/\n [a-z][\w-]*:/)[0]; + if (block.includes("needs['github-script-tests']")) { + assert.match(block, /^ needs: \[.*github-script-tests.*\]/m); + } + const expression = block.match(/^ if: \$\{\{ (.+) \}\}/m)[1]; + // These job guards use only operators whose semantics match JS for this fixture. + return new Function('github', 'inputs', 'needs', 'cancelled', `return (${expression});`)( + { repository, event_name: event, event: { action, pull_request: { draft } }, run_attempt: attempt }, + { fuzz_profile: profile, ci_command: false }, + { scope: { outputs }, ...Object.fromEntries(['github-script-tests', 'build', ...prerequisites].map(name => [name, { result: results[name] || 'success' }])) }, + () => cancelled + ); +} +test('fuzz event policy on ready PR admission and schedule, not every update', () => { + for (const job of [...suites, ...prerequisites]) { + for (const action of ['opened', 'ready_for_review']) assert.equal(enabled(job, { action }), true, job); + for (const action of ['synchronize', 'reopened']) assert.equal(enabled(job, { action }), false, job); + assert.equal(enabled(job, { event: 'schedule' }), true, job); + assert.equal(enabled(job, { event: 'merge_group' }), false, job); + assert.equal(enabled(job, { attempt: 2 }), false, job); + assert.equal(enabled(job, { action: 'opened', draft: true }), false, job); + assert.equal(enabled(job, { attempt: 2, draft: true }), false, job); + assert.equal(enabled(job, { event: 'schedule', repository: 'fork/greptimedb' }), false, job); + } +}); +test('manual profiles select only their suite and all required build jobs', () => { + const expected = { none: [], standalone: suites.slice(0, 2), distributed: [suites[2]], chaos: [suites[3]], all: suites }; + for (const [profile, jobs] of Object.entries(expected)) { + assert.deepEqual(suites.filter(job => enabled(job, { event: 'workflow_dispatch', profile })), jobs); + for (const job of prerequisites) assert.equal(enabled(job, { event: 'workflow_dispatch', profile }), profile !== 'none'); + } +}); +test('core integration tests still run on PR updates and merge groups', () => { + for (const job of ['build', 'sqlness', 'export-import-v2-e2e', 'run-multi-lang-tests']) { + assert.equal(enabled(job), true, job); + assert.equal(enabled(job, { event: 'merge_group' }), true, job); + } +}); +test('Windows compilation stays in Rust CI; test coverage remains in Nightly CI', () => { + const rust = readFileSync(resolve(workflows, 'rust.yml'), 'utf8'); + const nightly = readFileSync(resolve(workflows, 'nightly-ci.yml'), 'utf8'); + assert.match(rust, /^ check-windows:/m); + assert.doesNotMatch(rust, /^ test-windows:/m); + assert.match(nightly, /^ test-on-windows:/m); + assert.match(nightly, /^ sqlness-windows:/m); + assert.match(nightly, /cargo nextest run -F dashboard/); + assert.match(nightly, /^ workflow_dispatch:/m); +}); + +test('scope exclusions and dependency failures compose with the event policy', () => { + const gates = { + 'prepare-fuzz-targets': 'any_fuzz', 'build-greptime-ci': 'build_ci_needed', + fuzztest: 'fuzz_standalone', 'unstable-fuzztest': 'fuzz_unstable', + 'distributed-fuzztest': 'fuzz_distributed', 'distributed-fuzztest-with-chaos': 'fuzz_chaos' + }; + for (const [job, gate] of Object.entries(gates)) { + assert.equal(enabled(job, { action: 'opened', outputs: { [gate]: 'false' } }), false, job); + assert.equal(enabled(job, { action: 'opened', outputs: {} }), true, job); + assert.equal(enabled(job, { outputs: { [gate]: 'true' } }), false, job); + assert.equal(enabled(job, { action: 'opened', cancelled: true }), false, job); + } + for (const job of suites) { + for (const result of ['failure', 'skipped', 'cancelled']) { + assert.equal(enabled(job, { action: 'opened', results: { 'prepare-fuzz-targets': result } }), false, job); + const build = job === 'fuzztest' ? 'build' : 'build-greptime-ci'; + assert.equal(enabled(job, { action: 'opened', results: { [build]: result } }), false, job); + } + } +}); + +test('script tests must pass before any expensive integration job can start', () => { + const jobs = ['build', 'sqlness', 'sqlness-kvbackend', 'export-import-v2-e2e', + 'run-multi-lang-tests', 'compat', ...prerequisites, ...suites]; + for (const job of jobs) { + assert.equal(enabled(job, { action: 'opened' }), true, job); + for (const result of ['failure', 'skipped', 'cancelled']) { + assert.equal(enabled(job, { action: 'opened', results: { 'github-script-tests': result } }), false, job); + } + } +}); + +test('script tests gate expensive Rust and Checks jobs', () => { + for (const [file, jobs] of [ + ['rust.yml', ['clippy', 'check-riscv64', 'check-windows', 'test', 'coverage']], + ['checks.yml', ['check']] + ]) { + const workflow = readFileSync(resolve(workflows, file), 'utf8'); + for (const job of jobs) { + const event = job === 'coverage' ? 'merge_group' : 'workflow_dispatch'; + assert.equal(enabled(job, { workflow, event }), true, job); + for (const result of ['failure', 'skipped', 'cancelled']) { + assert.equal(enabled(job, { workflow, event, results: { 'github-script-tests': result } }), false, job); + } + } + } +}); diff --git a/.github/scripts/tests/nightly-merge.test.cjs b/.github/scripts/tests/nightly-merge.test.cjs index 885876f5610..753b458ebce 100644 --- a/.github/scripts/tests/nightly-merge.test.cjs +++ b/.github/scripts/tests/nightly-merge.test.cjs @@ -61,16 +61,22 @@ async function run({ sync = false, body = checked, labels = ['merge nightly'], const summary = { addHeading() { return this; }, addTable(rows) { calls.push(['summary', rows]); return this; }, async write() {} }; await scripts[sync ? 0 : 1](github, { repo: { owner: 'GreptimeTeam', repo: 'greptimedb' }, payload: { pull_request: { number: 1 } } }, - { summary, warning() {}, setFailed: message => calls.push(['failed', message]) }, + { summary, info() {}, warning() {}, setFailed: message => calls.push(['failed', message]) }, { env: { DRY_RUN: String(dryRun) } }); return calls; } +test('label sync grants PR label write permission without checking out PR code', () => { + const syncJob = workflow.split(' sync-label:\n')[1].split(' enqueue:\n')[0]; + assert.match(syncJob, /^ pull-requests: write$/m); + assert.doesNotMatch(syncJob, /uses: actions\/checkout/); +}); + test('template opts in and checkbox edits reconcile only the nightly label', async () => { assert.ok(readFileSync(resolve(__dirname, '../../pull_request_template.md'), 'utf8').includes(checked)); assert.deepEqual(await run({ sync: true, labels: ['bug'] }), [['add', ['merge nightly']]]); assert.deepEqual(await run({ sync: true }), []); - for (const body of [checked.replace('[x]', '[ ]'), '', null, `> ${checked}`]) { + for (const body of [checked.replace('[x]', '[ ]')]) { assert.deepEqual(await run({ sync: true, body }), [['remove', 'merge nightly']]); } assert.deepEqual(await run({ sync: true, body: checked.replace('[x]', '[X]'), labels: [] }), [['add', ['merge nightly']]]); @@ -83,7 +89,7 @@ test('eligible PR is enqueued with its checked head; dry run performs no mutatio test('each eligibility gate independently prevents enqueueing', async () => { for (const options of [ - { body: checked.replace('[x]', '[ ]') }, { labels: [] }, + { labels: [] }, { current: { draft: true } }, { current: { state: 'closed' } }, { current: { base: { ref: 'release/v1.3' } } }, { detail: { reviewDecision: 'REVIEW_REQUIRED' } }, { detail: { reviewDecision: 'CHANGES_REQUESTED' } }, { detail: { mergeStateStatus: 'BLOCKED' } }, { detail: { mergeStateStatus: 'UNKNOWN' } }, @@ -106,3 +112,14 @@ test('per-run cap and individual failures do not prevent later candidates', asyn assert.equal(calls.filter(c => c[0] === 'enqueue').length, 2); assert.equal(calls.filter(c => c[0] === 'failed').length, 1); }); + +test('label-only opt-in survives PRs without a template and does not require a checkbox', async () => { + for (const body of ['', null, 'Existing PR description', `> ${checked}`]) { + assert.deepEqual(await run({ sync: true, body }), []); + assert.deepEqual(await run({ sync: true, body, labels: [] }), []); + assert.equal((await run({ body })).filter(call => call[0] === 'enqueue').length, 1); + assert.equal((await run({ body, labels: [] })).some(call => call[0] === 'enqueue'), false); + } + // The label is authoritative for admission, even before checkbox synchronization. + assert.equal((await run({ body: checked.replace('[x]', '[ ]') })).some(call => call[0] === 'enqueue'), true); +}); diff --git a/.github/scripts/tests/test_ci_slash.py b/.github/scripts/tests/test_ci_slash.py index e551a4f5297..2d9c7bccc84 100644 --- a/.github/scripts/tests/test_ci_slash.py +++ b/.github/scripts/tests/test_ci_slash.py @@ -136,7 +136,7 @@ class CiSlashTest(unittest.TestCase): def test_fork_rerun_gates_cover_standard_jobs(self): import re - for name, count in [("rust", 7), ("integration", 7), ("checks", 6), ("docs", 3)]: + for name, count in [("rust", 6), ("integration", 8), ("checks", 6), ("docs", 3)]: workflow = (SCRIPT.parents[1] / "workflows" / (name + ".yml")).read_text() gates = re.findall(r"^ if:.*github.run_attempt > 1.*$", workflow, re.MULTILINE) self.assertEqual(len(gates), count, name) diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml index 941f36d9726..1f14d8851cc 100644 --- a/.github/workflows/checks.yml +++ b/.github/workflows/checks.yml @@ -71,43 +71,7 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: actions/setup-node@v4 - with: - node-version: 24 - - name: Run Node GitHub script tests - run: node --test .github/scripts/query-regression-comment.test.cjs .github/scripts/tests/nightly-merge.test.cjs - - name: Setup uv - uses: astral-sh/setup-uv@v3 - with: - version: "latest" - - name: Run check-version script tests - run: uv run --no-project python .github/scripts/check-version-test.py - - name: Run release Linux architecture tests - run: python3 .github/scripts/release-linux-architectures-test.py - - name: Run previous-release-tag script tests - run: python3 .github/scripts/previous-release-tag-test.py - - name: Run enterprise license check tests - run: python3 scripts/check-enterprise-license-test.py - - name: Run fuzz orchestration script tests - run: .github/scripts/run-fuzz-targets-test.sh - - name: Run CI comment-command tests - run: python3 .github/scripts/tests/test_ci_slash.py - - name: Run CI scope gate tests - run: python3 .github/ci/tests/test_ci_scope.py - - name: Run metric snapshot benchmark tests - run: python3 -m unittest discover -s .github/scripts/tests -p 'test_metric_snapshot*.py' - - name: Run query-regression Python tooling tests - run: | - python3 tests/perf/test_query_regression_runner_compaction_toctou.py - python3 tests/perf/test_query_regression_runner_otlp_trace_load.py - python3 tests/perf/test_query_regression_summary_otlp.py - python3 tests/perf/test_query_regression_case_selection.py - python3 tests/perf/test_query_regression_nightly_refs.py - python3 tests/perf/test_query_regression_slash.py - python3 tests/perf/test_aliyun_ecs_runner_scripts.py - python3 tests/perf/test_ci_runner_scripts.py - python3 tests/perf/test_agent_observability_summary.py - python3 tests/perf/test_tracesbench_summary.py + - uses: ./.github/actions/github-script-tests scope-rules: if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} @@ -127,7 +91,8 @@ jobs: run: python3 .github/ci/ci_scope_rules.py --check-paths check: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} + if: ${{ needs['github-script-tests'].result == 'success' && github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} + needs: [github-script-tests] name: Check runs-on: ${{ matrix.os }} strategy: diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml index 6106e5a2860..00c3f52a034 100644 --- a/.github/workflows/integration.yml +++ b/.github/workflows/integration.yml @@ -21,7 +21,7 @@ on: default: false type: boolean fuzz_profile: - description: Fuzz profile selected by /ci fuzz + description: Fuzz profile for a manual run or /ci fuzz required: false default: none type: choice @@ -42,6 +42,17 @@ concurrency: cancel-in-progress: true jobs: + github-script-tests: + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} + name: GitHub Script Tests (Integration CI) + runs-on: ubuntu-latest + timeout-minutes: 5 + steps: + - uses: actions/checkout@v4 + with: + persist-credentials: false + - uses: ./.github/actions/github-script-tests + # Which jobs this gate skips, and why: .github/ci/README.md. scope: name: CI Scope @@ -94,7 +105,8 @@ jobs: python3 .github/ci/ci_scope.py "${args[@]}" build: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} + if: ${{ needs['github-script-tests'].result == 'success' && github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} + needs: [github-script-tests] name: Build GreptimeDB binaries runs-on: ${{ matrix.os }} strategy: @@ -155,9 +167,9 @@ jobs: version: current sqlness: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} + if: ${{ needs['github-script-tests'].result == 'success' && github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} name: Sqlness Test (${{ matrix.mode.name }}) - needs: build + needs: [github-script-tests, build] runs-on: ${{ matrix.os }} strategy: fail-fast: false @@ -213,9 +225,9 @@ jobs: # required checks, so they are the only ones this gate may skip. They run # the same cases as the `Basic` leg, so the diff only matters when it can # reach the backend itself; the merge queue runs them regardless. - if: ${{ !cancelled() && github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) && needs.build.result == 'success' && needs.scope.outputs.sqlness_kvbackend != 'false' }} + if: ${{ !cancelled() && needs['github-script-tests'].result == 'success' && github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) && needs.build.result == 'success' && needs.scope.outputs.sqlness_kvbackend != 'false' }} name: Sqlness Test (${{ matrix.mode.name }}) - needs: [scope, build] + needs: [github-script-tests, scope, build] runs-on: ${{ matrix.os }} strategy: fail-fast: false @@ -267,9 +279,9 @@ jobs: retention-days: 3 export-import-v2-e2e: - if: ${{ !cancelled() && github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) && needs.build.result == 'success' && needs.scope.outputs.export_import_v2 != 'false' }} + if: ${{ !cancelled() && needs['github-script-tests'].result == 'success' && github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) && needs.build.result == 'success' && needs.scope.outputs.export_import_v2 != 'false' }} name: Export/Import V2 E2E Test - needs: [scope, build] + needs: [github-script-tests, scope, build] runs-on: ubuntu-latest timeout-minutes: 60 steps: @@ -337,9 +349,9 @@ jobs: retention-days: 3 run-multi-lang-tests: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} + if: ${{ needs['github-script-tests'].result == 'success' && github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} name: Run Multi-language SDK Tests - needs: build + needs: [github-script-tests, build] uses: ./.github/workflows/run-multi-lang-tests.yml with: artifact-name: bins @@ -362,9 +374,9 @@ jobs: run: python3 .github/scripts/update-compat-versions.py --check --published-only --check-anchors compat: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name == 'merge_group' || github.event_name == 'workflow_dispatch' || github.event_name == 'schedule' || (github.event_name == 'pull_request' && (github.event.action == 'ready_for_review' || (github.event.action == 'opened' && github.event.pull_request.draft == false)))) }} + if: ${{ needs['github-script-tests'].result == 'success' && github.repository == 'GreptimeTeam/greptimedb' && (github.event_name == 'merge_group' || github.event_name == 'workflow_dispatch' || github.event_name == 'schedule' || (github.event_name == 'pull_request' && (github.event.action == 'ready_for_review' || (github.event.action == 'opened' && github.event.pull_request.draft == false)))) }} name: Compatibility Test (recent releases) - needs: build + needs: [github-script-tests, build] runs-on: ubuntu-latest timeout-minutes: 30 steps: @@ -401,10 +413,12 @@ jobs: path: /tmp/sqlness-compat* retention-days: 3 + # Fuzz runs on opening a ready PR, marking it ready, manual dispatch, and + # the scheduled main run. Synchronize events keep the core integration tests. prepare-fuzz-targets: - if: ${{ !cancelled() && github.repository == 'GreptimeTeam/greptimedb' && ((github.event_name == 'pull_request' && github.event.pull_request.draft == false) || (github.event_name == 'workflow_dispatch' && inputs.ci_command) || github.event_name == 'merge_group') && (github.event_name == 'pull_request' || github.event_name == 'merge_group' || inputs.fuzz_profile != 'none') && needs.scope.outputs.any_fuzz != 'false' }} + if: ${{ !cancelled() && needs['github-script-tests'].result == 'success' && github.repository == 'GreptimeTeam/greptimedb' && (github.event_name == 'schedule' || (github.event_name == 'pull_request' && github.event.pull_request.draft == false && (github.event.action == 'opened' || github.event.action == 'ready_for_review')) || (github.event_name == 'workflow_dispatch' && (inputs.fuzz_profile != 'none'))) && needs.scope.outputs.any_fuzz != 'false' }} name: Prepare Fuzz Targets - needs: [scope] + needs: [github-script-tests, scope] runs-on: ubuntu-latest timeout-minutes: 60 steps: @@ -517,9 +531,9 @@ jobs: if-no-files-found: error fuzztest: - if: ${{ !cancelled() && github.repository == 'GreptimeTeam/greptimedb' && ((github.event_name == 'pull_request' && github.event.pull_request.draft == false) || (github.event_name == 'workflow_dispatch' && inputs.ci_command) || github.event_name == 'merge_group') && (github.event_name == 'pull_request' || github.event_name == 'merge_group' || inputs.fuzz_profile == 'standalone' || inputs.fuzz_profile == 'all') && needs.build.result == 'success' && needs['prepare-fuzz-targets'].result == 'success' && needs.scope.outputs.fuzz_standalone != 'false' }} + if: ${{ !cancelled() && needs['github-script-tests'].result == 'success' && github.repository == 'GreptimeTeam/greptimedb' && (github.event_name == 'schedule' || (github.event_name == 'pull_request' && github.event.pull_request.draft == false && (github.event.action == 'opened' || github.event.action == 'ready_for_review')) || (github.event_name == 'workflow_dispatch' && (inputs.fuzz_profile == 'standalone' || inputs.fuzz_profile == 'all'))) && needs.build.result == 'success' && needs['prepare-fuzz-targets'].result == 'success' && needs.scope.outputs.fuzz_standalone != 'false' }} name: Fuzz Test (Standalone, ${{ matrix.group.name }}) - needs: [build, prepare-fuzz-targets, scope] + needs: [github-script-tests, build, prepare-fuzz-targets, scope] runs-on: ubuntu-latest timeout-minutes: 60 strategy: @@ -598,9 +612,9 @@ jobs: artifact-name: fuzz-standalone-${{ matrix.group.id }} unstable-fuzztest: - if: ${{ !cancelled() && github.repository == 'GreptimeTeam/greptimedb' && ((github.event_name == 'pull_request' && github.event.pull_request.draft == false) || (github.event_name == 'workflow_dispatch' && inputs.ci_command) || github.event_name == 'merge_group') && (github.event_name == 'pull_request' || github.event_name == 'merge_group' || inputs.fuzz_profile == 'standalone' || inputs.fuzz_profile == 'all') && needs['build-greptime-ci'].result == 'success' && needs['prepare-fuzz-targets'].result == 'success' && needs.scope.outputs.fuzz_unstable != 'false' }} + if: ${{ !cancelled() && needs['github-script-tests'].result == 'success' && github.repository == 'GreptimeTeam/greptimedb' && (github.event_name == 'schedule' || (github.event_name == 'pull_request' && github.event.pull_request.draft == false && (github.event.action == 'opened' || github.event.action == 'ready_for_review')) || (github.event_name == 'workflow_dispatch' && (inputs.fuzz_profile == 'standalone' || inputs.fuzz_profile == 'all'))) && needs['build-greptime-ci'].result == 'success' && needs['prepare-fuzz-targets'].result == 'success' && needs.scope.outputs.fuzz_unstable != 'false' }} name: Unstable Fuzz Test - needs: [build-greptime-ci, prepare-fuzz-targets, scope] + needs: [github-script-tests, build-greptime-ci, prepare-fuzz-targets, scope] runs-on: ubuntu-latest timeout-minutes: 60 steps: @@ -658,9 +672,9 @@ jobs: artifact-name: fuzz-unstable-standalone build-greptime-ci: - if: ${{ !cancelled() && github.repository == 'GreptimeTeam/greptimedb' && ((github.event_name == 'pull_request' && github.event.pull_request.draft == false) || (github.event_name == 'workflow_dispatch' && inputs.ci_command && inputs.fuzz_profile != 'none') || github.event_name == 'merge_group') && needs.scope.outputs.build_ci_needed != 'false' }} + if: ${{ !cancelled() && needs['github-script-tests'].result == 'success' && github.repository == 'GreptimeTeam/greptimedb' && (github.event_name == 'schedule' || (github.event_name == 'pull_request' && github.event.pull_request.draft == false && (github.event.action == 'opened' || github.event.action == 'ready_for_review')) || (github.event_name == 'workflow_dispatch' && (inputs.fuzz_profile != 'none'))) && needs.scope.outputs.build_ci_needed != 'false' }} name: Build GreptimeDB binary (profile-CI) - needs: [scope] + needs: [github-script-tests, scope] runs-on: ${{ matrix.os }} strategy: matrix: @@ -719,10 +733,10 @@ jobs: version: current distributed-fuzztest: - if: ${{ !cancelled() && github.repository == 'GreptimeTeam/greptimedb' && ((github.event_name == 'pull_request' && github.event.pull_request.draft == false) || (github.event_name == 'workflow_dispatch' && inputs.ci_command) || github.event_name == 'merge_group') && (github.event_name == 'pull_request' || github.event_name == 'merge_group' || inputs.fuzz_profile == 'distributed' || inputs.fuzz_profile == 'all') && needs['build-greptime-ci'].result == 'success' && needs['prepare-fuzz-targets'].result == 'success' && needs.scope.outputs.fuzz_distributed != 'false' }} + if: ${{ !cancelled() && needs['github-script-tests'].result == 'success' && github.repository == 'GreptimeTeam/greptimedb' && (github.event_name == 'schedule' || (github.event_name == 'pull_request' && github.event.pull_request.draft == false && (github.event.action == 'opened' || github.event.action == 'ready_for_review')) || (github.event_name == 'workflow_dispatch' && (inputs.fuzz_profile == 'distributed' || inputs.fuzz_profile == 'all'))) && needs['build-greptime-ci'].result == 'success' && needs['prepare-fuzz-targets'].result == 'success' && needs.scope.outputs.fuzz_distributed != 'false' }} name: Fuzz Test (Distributed, ${{ matrix.mode.name }}, ${{ matrix.group.name }}) runs-on: ubuntu-latest - needs: [build-greptime-ci, prepare-fuzz-targets, scope] + needs: [github-script-tests, build-greptime-ci, prepare-fuzz-targets, scope] timeout-minutes: 60 strategy: fail-fast: false @@ -958,10 +972,10 @@ jobs: docker system prune -f distributed-fuzztest-with-chaos: - if: ${{ !cancelled() && github.repository == 'GreptimeTeam/greptimedb' && ((github.event_name == 'pull_request' && github.event.pull_request.draft == false) || (github.event_name == 'workflow_dispatch' && inputs.ci_command) || github.event_name == 'merge_group') && (github.event_name == 'pull_request' || github.event_name == 'merge_group' || inputs.fuzz_profile == 'chaos' || inputs.fuzz_profile == 'all') && needs['build-greptime-ci'].result == 'success' && needs['prepare-fuzz-targets'].result == 'success' && needs.scope.outputs.fuzz_chaos != 'false' }} + if: ${{ !cancelled() && needs['github-script-tests'].result == 'success' && github.repository == 'GreptimeTeam/greptimedb' && (github.event_name == 'schedule' || (github.event_name == 'pull_request' && github.event.pull_request.draft == false && (github.event.action == 'opened' || github.event.action == 'ready_for_review')) || (github.event_name == 'workflow_dispatch' && (inputs.fuzz_profile == 'chaos' || inputs.fuzz_profile == 'all'))) && needs['build-greptime-ci'].result == 'success' && needs['prepare-fuzz-targets'].result == 'success' && needs.scope.outputs.fuzz_chaos != 'false' }} name: Fuzz Test with Chaos (Distributed, ${{ matrix.mode.name }}, ${{ matrix.group.name }}) runs-on: ubuntu-latest - needs: [build-greptime-ci, prepare-fuzz-targets, scope] + needs: [github-script-tests, build-greptime-ci, prepare-fuzz-targets, scope] timeout-minutes: 60 strategy: fail-fast: false diff --git a/.github/workflows/nightly-merge.yml b/.github/workflows/nightly-merge.yml index ff6d5874ff1..6188d12cd3f 100644 --- a/.github/workflows/nightly-merge.yml +++ b/.github/workflows/nightly-merge.yml @@ -23,7 +23,8 @@ jobs: runs-on: ubuntu-latest timeout-minutes: 5 permissions: - pull-requests: read + # PR label mutations require pull-request write permission. + pull-requests: write issues: write concurrency: group: nightly-merge-label-${{ github.event.pull_request.number }} @@ -36,7 +37,10 @@ jobs: const { owner, repo } = context.repo; const issue_number = context.payload.pull_request.number; const { data: pr } = await github.rest.pulls.get({ owner, repo, pull_number: issue_number }); - const selected = /^- \[[xX]\] Merge during the nightly window \s*$/m.test(pr.body || ''); + const checkbox = /^- \[([ xX])\] Merge during the nightly window \s*$/m.exec(pr.body || ''); + // PRs without the template manage their opt-in through the label. + if (!checkbox) return; + const selected = checkbox[1].toLowerCase() === 'x'; const label = 'merge nightly'; const labeled = pr.labels.some(item => item.name === label); if (selected && !labeled) { @@ -73,11 +77,10 @@ jobs: if (!candidate.pull_request || enqueued >= 5) continue; const number = candidate.number; try { - // Fetch current body/labels, not the earlier list snapshot. + // Fetch current labels and PR state, not the earlier list snapshot. const { data: current } = await github.rest.pulls.get({ owner, repo, pull_number: number }); if (current.state !== 'open' || current.draft || current.base.ref !== 'main' || - !current.labels.some(label => label.name === 'merge nightly') || - !/^- \[[xX]\] Merge during the nightly window \s*$/m.test(current.body || '')) { + !current.labels.some(label => label.name === 'merge nightly')) { rows.push([String(number), 'Skipped: not opted in, draft, closed, or not targeting main']); continue; } @@ -134,5 +137,6 @@ jobs: core.setFailed('Some PRs could not be processed; see workflow logs.'); } } + for (const [number, result] of rows) core.info(`PR #${number}: ${result}`); await core.summary.addHeading(dryRun ? 'Nightly merge (dry run)' : 'Nightly merge') .addTable([['PR', 'Result'], ...rows]).write(); diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index d1b93a60e43..842fa306c4e 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -31,6 +31,17 @@ concurrency: cancel-in-progress: true jobs: + github-script-tests: + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} + name: GitHub Script Tests (Rust CI) + runs-on: ubuntu-latest + timeout-minutes: 5 + steps: + - uses: actions/checkout@v4 + with: + persist-credentials: false + - uses: ./.github/actions/github-script-tests + fmt: if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} name: Rustfmt @@ -50,7 +61,8 @@ jobs: run: make fmt-check clippy: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} + if: ${{ needs['github-script-tests'].result == 'success' && github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} + needs: [github-script-tests] name: Clippy runs-on: ubuntu-latest timeout-minutes: 60 @@ -94,7 +106,9 @@ jobs: run: make check-udeps check-riscv64: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} + # Keep platform coverage on the weekday schedule and allow manual checks. + if: ${{ needs['github-script-tests'].result == 'success' && github.repository == 'GreptimeTeam/greptimedb' && (github.event_name == 'schedule' || github.event_name == 'workflow_dispatch') }} + needs: [github-script-tests] name: Check riscv64 build runs-on: ubuntu-latest timeout-minutes: 90 @@ -120,7 +134,9 @@ jobs: run: cargo check --locked --workspace --all-targets --features servers/dashboard --target riscv64gc-unknown-linux-gnu check-windows: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} + # Keep platform coverage on the weekday schedule and allow manual checks. + if: ${{ needs['github-script-tests'].result == 'success' && github.repository == 'GreptimeTeam/greptimedb' && (github.event_name == 'schedule' || github.event_name == 'workflow_dispatch') }} + needs: [github-script-tests] name: Check Windows build runs-on: windows-2022 timeout-minutes: 60 @@ -148,48 +164,6 @@ jobs: env: RUSTFLAGS: "-D warnings" - test-windows: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false) }} - name: Test Windows - runs-on: windows-2022 - timeout-minutes: 90 - env: - CARGO_INCREMENTAL: 0 - CARGO_PROFILE_DEV_DEBUG: 0 - CARGO_PROFILE_TEST_DEBUG: 0 - RUSTFLAGS: "-D warnings -C linker=lld-link" - RUST_BACKTRACE: 1 - RUST_MIN_STACK: 8388608 # 8MB - steps: - - run: git config --global core.autocrlf false - - uses: actions/checkout@v4 - with: - persist-credentials: false - - uses: ./.github/actions/setup-cyborg - - uses: taiki-e/install-action@v2 - with: - tool: protoc@3.23.4 - - uses: KyleMayes/install-llvm-action@v1 - with: - version: "14.0" - - uses: actions-rust-lang/setup-rust-toolchain@v1 - with: - cache: false - - name: Rust Cache - uses: Swatinem/rust-cache@v2 - with: - shared-key: "test-windows" - cache-all-crates: "true" - save-if: ${{ github.ref == 'refs/heads/main' }} - - name: Install Cargo Nextest - uses: taiki-e/install-action@nextest - - name: Install WSL distribution - uses: Vampire/setup-wsl@v2 - with: - distribution: Ubuntu-22.04 - - name: Run Windows tests - run: cargo nextest run --locked --workspace -F dashboard --no-fail-fast - conflict-check: if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} name: Check for conflict @@ -202,10 +176,10 @@ jobs: uses: olivernybroe/action-conflict-finder@v4.0 test: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && github.event_name != 'merge_group' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} + if: ${{ needs['github-script-tests'].result == 'success' && github.repository == 'GreptimeTeam/greptimedb' && github.event_name != 'merge_group' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} runs-on: ubuntu-22.04-arm timeout-minutes: 60 - needs: [conflict-check, clippy, fmt, check-udeps] + needs: [github-script-tests, conflict-check, clippy, fmt, check-udeps] steps: - uses: actions/checkout@v4 with: @@ -258,7 +232,8 @@ jobs: UNITTEST_LOG_DIR: "__unittest_logs" coverage: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && github.event_name == 'merge_group' }} + if: ${{ needs['github-script-tests'].result == 'success' && github.repository == 'GreptimeTeam/greptimedb' && github.event_name == 'merge_group' }} + needs: [github-script-tests] runs-on: ubuntu-22.04-8-cores timeout-minutes: 60 steps: From ec240ad7e64f92f0a1c308d9c3f6834ba605edb2 Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Thu, 29 Oct 2026 05:45:07 +0000 Subject: [PATCH 135/239] feat(mito): preserve ordinary sparse keys across writes and schema changes (#9920) * feat(mito): support ordinary sparse append tables Ordinary append tables need sparse keys for optional series tags, including JSON2 values written through row and bulk paths. This change validates sparse tag schemas and defaults, keeps key identity stable when nullable tags are added, selects flat SST storage, and preserves append-mode semantics when reopening regions. Includes storage, bulk-write, schema-evolution and repartition tests, plus SQL and mixed-version fixtures. Existing metric-engine sparse key handling remains supported. This is the storage prerequisite for exemplar ingestion; it introduces no exemplar protocol endpoint. Signed-off-by: shuiyisong * fix(mito): preserve reserved sparse key semantics Signed-off-by: shuiyisong * fix(mito): clarify sparse errors and avoid panics Signed-off-by: shuiyisong * perf(mito): reduce sparse key validation overhead Signed-off-by: shuiyisong * fix: sqlness Signed-off-by: shuiyisong * perf(mito-codec): reduce sparse validation cost Signed-off-by: shuiyisong * refactor(mito): remove ordinary sparse support Signed-off-by: shuiyisong --------- Signed-off-by: shuiyisong --- src/datatypes/src/extension/json.rs | 3 +- src/metric-engine/src/engine/read.rs | 121 ++++- src/mito-codec/src/row_converter.rs | 27 +- src/mito-codec/src/row_converter/sparse.rs | 326 ++++++++++-- .../src/row_converter/sparse/checked.rs | 13 +- src/mito2/Cargo.toml | 4 + src/mito2/benches/bench_sparse_write.rs | 391 +++++++++++++++ src/mito2/src/compaction/scheduler_test.rs | 1 + .../src/engine/bulk_insert_validation_test.rs | 309 +++++++++++- src/mito2/src/engine/create_test.rs | 23 +- src/mito2/src/engine/scan_test.rs | 265 ++++++++++ src/mito2/src/manifest/manager.rs | 7 +- src/mito2/src/memtable/bulk/part.rs | 468 ++++++++++++++++-- src/mito2/src/read/compat.rs | 274 +++++++--- src/mito2/src/region.rs | 3 + src/mito2/src/region/opener.rs | 36 +- src/mito2/src/region/options.rs | 44 +- src/mito2/src/request.rs | 18 +- src/mito2/src/sst/primary_key.rs | 23 +- src/mito2/src/test_util.rs | 1 + src/mito2/src/test_util/scheduler_util.rs | 1 + src/store-api/src/mito_engine_options.rs | 2 +- src/store-api/src/region_engine.rs | 11 +- src/table/src/table/scan.rs | 42 +- tests-integration/src/tests/instance_test.rs | 37 ++ .../common/types/json/exemplar_sparse.result | 47 ++ .../common/types/json/exemplar_sparse.sql | 15 + .../cases/exemplar_legacy_keys/case.toml | 8 + .../cases/exemplar_legacy_keys/setup.sql | 17 + .../cases/exemplar_legacy_keys/verify.result | 85 ++++ .../cases/exemplar_legacy_keys/verify.sql | 13 + 31 files changed, 2370 insertions(+), 265 deletions(-) create mode 100644 src/mito2/benches/bench_sparse_write.rs create mode 100644 tests/cases/standalone/common/types/json/exemplar_sparse.result create mode 100644 tests/cases/standalone/common/types/json/exemplar_sparse.sql create mode 100644 tests/compatibility/cases/exemplar_legacy_keys/case.toml create mode 100644 tests/compatibility/cases/exemplar_legacy_keys/setup.sql create mode 100644 tests/compatibility/cases/exemplar_legacy_keys/verify.result create mode 100644 tests/compatibility/cases/exemplar_legacy_keys/verify.sql diff --git a/src/datatypes/src/extension/json.rs b/src/datatypes/src/extension/json.rs index 2e86d7791c5..ef72f40e997 100644 --- a/src/datatypes/src/extension/json.rs +++ b/src/datatypes/src/extension/json.rs @@ -117,7 +117,8 @@ fn parse_version(metadata: &str) -> crate::error::Result> { } /// Returns the remainder field of a JSON2 v2 root. -pub(crate) fn json2_remainder_field(field: &Field) -> crate::error::Result> { +/// Validates the layout version and remainder type, returning `None` for legacy layouts. +pub fn json2_remainder_field(field: &Field) -> crate::error::Result> { let layout = Json2PhysicalLayout::try_from_root(field)?; if !layout.is_version_2() { return Ok(None); diff --git a/src/metric-engine/src/engine/read.rs b/src/metric-engine/src/engine/read.rs index 8d8d32fcc6c..b10fcfb12ad 100644 --- a/src/metric-engine/src/engine/read.rs +++ b/src/metric-engine/src/engine/read.rs @@ -16,7 +16,8 @@ use std::sync::Arc; use api::v1::SemanticType; use common_telemetry::{debug, error, tracing}; -use datafusion::logical_expr::{self, Expr}; +use datafusion::common::ScalarValue; +use datafusion::logical_expr::{self, Expr, Operator}; use snafu::{OptionExt, ResultExt}; use store_api::metadata::{RegionMetadataBuilder, RegionMetadataRef}; use store_api::metric_engine_consts::DATA_SCHEMA_TABLE_ID_COLUMN_NAME; @@ -60,10 +61,31 @@ impl MetricEngineInner { .with_label_values(&["read_physical"]) .start_timer(); - self.mito + // PromQL scans physical regions with a table-ID filter to access `__tsid`. + let logical_region = request + .filters + .iter() + .flat_map(logical_expr::utils::split_conjunction) + .any(|expr| { + let Expr::BinaryExpr(expr) = expr else { + return false; + }; + expr.op == Operator::Eq + && matches!( + (expr.left.as_ref(), expr.right.as_ref()), + (Expr::Column(column), Expr::Literal(ScalarValue::UInt32(Some(_)), _)) + | (Expr::Literal(ScalarValue::UInt32(Some(_)), _), Expr::Column(column)) + if column.name == DATA_SCHEMA_TABLE_ID_COLUMN_NAME + ) + }); + let mut scanner = self + .mito .handle_query(region_id, request) .await - .context(MitoReadOperationSnafu) + .context(MitoReadOperationSnafu)?; + scanner.set_logical_region(logical_region); + + Ok(scanner) } async fn read_logical_region( @@ -399,6 +421,99 @@ mod test { ); } + #[tokio::test] + async fn test_single_logical_region_scan_properties() { + let env = TestEnv::new().await; + env.init_metric_region().await; + + let physical_region_id = env.default_physical_region_id(); + let logical_region_id = env.default_logical_region_id(); + let table_id = logical_expr::col(DATA_SCHEMA_TABLE_ID_COLUMN_NAME); + let table_filter = table_id + .clone() + .eq(logical_expr::lit(logical_region_id.table_id())); + let job_filter = logical_expr::col("job").eq(logical_expr::lit("api")); + for (filters, logical_region) in [ + (vec![], false), + (vec![table_filter.clone()], true), + ( + vec![logical_expr::lit(logical_region_id.table_id()).eq(table_id.clone())], + true, + ), + (vec![job_filter.clone(), table_filter.clone()], true), + (vec![job_filter.clone().and(table_filter.clone())], true), + ( + vec![ + table_filter.clone().and( + job_filter + .clone() + .or(logical_expr::col("job").eq(logical_expr::lit("db"))), + ), + ], + true, + ), + ( + vec![ + table_filter + .clone() + .or(table_id.clone().eq(logical_expr::lit(4u32))), + ], + false, + ), + (vec![table_filter.clone().or(job_filter.clone())], false), + ( + vec![table_id.clone().not_eq(logical_expr::lit(4u32))], + false, + ), + (vec![table_id.clone().gt(logical_expr::lit(4u32))], false), + ( + vec![table_id.eq(logical_expr::lit(ScalarValue::UInt32(None)))], + false, + ), + (vec![job_filter], false), + ] { + let scanner = env + .metric() + .handle_query( + physical_region_id, + ScanRequest { + filters: filters.clone(), + distribution: Some(TimeSeriesDistribution::PerSeries), + ..Default::default() + }, + ) + .await + .unwrap(); + assert_eq!( + scanner.properties().is_logical_region(), + logical_region, + "{filters:?}" + ); + } + + let scanner = env + .metric() + .handle_query(logical_region_id, ScanRequest::default()) + .await + .unwrap(); + assert!(scanner.properties().is_logical_region()); + + // A table-ID filter alone cannot grant metric guarantees to a Mito scan. + let scanner = env + .mito() + .handle_query( + utils::to_data_region_id(physical_region_id), + ScanRequest { + filters: vec![table_filter], + distribution: Some(TimeSeriesDistribution::PerSeries), + ..Default::default() + }, + ) + .await + .unwrap(); + assert!(!scanner.properties().is_logical_region()); + } + #[tokio::test] async fn test_transform_scan_req() { let env = TestEnv::new().await; diff --git a/src/mito-codec/src/row_converter.rs b/src/mito-codec/src/row_converter.rs index 60bbf1935e0..8a42c1f1190 100644 --- a/src/mito-codec/src/row_converter.rs +++ b/src/mito-codec/src/row_converter.rs @@ -175,24 +175,35 @@ pub fn build_primary_key_codec_with_fields( /// Finds the checked boundary of an Option, shared by Dense and Sparse. /// This validates framing, not UTF-8; consumers decoding strings validate UTF-8. pub(crate) fn encoded_string_len(bytes: &[u8]) -> memcomparable::Result { + encoded_string_info::(bytes).map(|(len, _)| len) +} + +/// Checks string framing and optionally detects ASCII in the same pass. +/// Framing-only readers compile out payload inspection and ignore the ASCII result. +pub(crate) fn encoded_string_info( + bytes: &[u8], +) -> memcomparable::Result<(usize, bool)> { match bytes.first().copied().ok_or(memcomparable::Error::Eof)? { - 0 => return Ok(1), + 0 => return Ok((1, true)), 1 => {} marker => return Err(memcomparable::Error::InvalidTagEncoding(marker as usize)), } match bytes.get(1).copied().ok_or(memcomparable::Error::Eof)? { - 0 => return Ok(2), + 0 => return Ok((2, true)), 1 => {} marker => return Err(memcomparable::Error::InvalidBytesEncoding(marker)), } - let mut end = 2; - loop { - let chunk = bytes.get(end..end + 9).ok_or(memcomparable::Error::Eof)?; - end += 9; - match chunk[8] { - 1..=8 => return Ok(end), + let mut ascii_bits = 0_u64; + for (index, chunk) in bytes[2..].as_chunks::<9>().0.iter().enumerate() { + let [word @ .., marker] = *chunk; + if CHECK_ASCII { + ascii_bits |= u64::from_ne_bytes(word); + } + match marker { + 1..=8 => return Ok((2 + (index + 1) * 9, ascii_bits & 0x8080_8080_8080_8080 == 0)), 9 => {} marker => return Err(memcomparable::Error::InvalidBytesEncoding(marker)), } } + Err(memcomparable::Error::Eof) } diff --git a/src/mito-codec/src/row_converter/sparse.rs b/src/mito-codec/src/row_converter/sparse.rs index 5497a7dc034..9cb16f406cd 100644 --- a/src/mito-codec/src/row_converter/sparse.rs +++ b/src/mito-codec/src/row_converter/sparse.rs @@ -23,7 +23,7 @@ use datatypes::prelude::ConcreteDataType; use datatypes::value::{Value, ValueRef}; use memcomparable::{Deserializer, Serializer}; use serde::{Deserialize, Serialize}; -use snafu::{ResultExt, ensure}; +use snafu::{OptionExt, ResultExt, ensure}; use store_api::codec::PrimaryKeyEncoding; use store_api::metadata::RegionMetadataRef; use store_api::storage::ColumnId; @@ -39,30 +39,28 @@ use crate::row_converter::dense::SortField; pub use crate::row_converter::sparse::checked::SparsePrimaryKeyView; use crate::row_converter::{CompositeValues, PrimaryKeyCodec, PrimaryKeyFilter}; -/// A codec for sparse key of metrics. +/// A codec for sparse keys with reserved table ID and TSID columns. /// /// ## Encoding format /// Each primary key is encoded as a sequence of `(column_id, value)` pairs: -/// - The first two fields are always the reserved `table_id` (uint32) and `tsid` (uint64). -/// - User-defined labels follow, sorted by **column name** in lexicographical order. -/// - Null values are omitted (not encoded). +/// - Keys start with reserved `table_id` (uint32) and `tsid` (uint64). +/// - String tags follow the reserved prefix; null tags are omitted. /// /// The `column_id` is encoded as a 4-byte big-endian integer, and the value is encoded /// using memcomparable serialization. /// -/// `decode_leftmost` always decodes the first value from the encoded bytes (i.e., the -/// `table_id` field). +/// `decode_leftmost` returns the reserved `table_id` value. /// /// ## Requirements -/// It requires the input primary key columns are sorted by the column name in lexicographical order. -/// It encodes the column id of the physical region. +/// Encoders supply columns in the layout's order using stable physical-region IDs. #[derive(Clone, Debug)] pub struct SparsePrimaryKeyCodec { inner: Arc, } -#[derive(Debug)] +#[derive(Debug, Clone)] struct SparsePrimaryKeyCodecInner { + fields: Option>, // Internal fields table_id_field: SortField, // Internal fields @@ -146,6 +144,66 @@ const TSID_VALUE_OFFSET: usize = COLUMN_ID_ENCODE_SIZE + 5 + COLUMN_ID_ENCODE_SI /// Byte offset where tag columns start (after 9-byte table_id + 13-byte tsid entries). const TAGS_START_OFFSET: usize = COLUMN_ID_ENCODE_SIZE + 5 + COLUMN_ID_ENCODE_SIZE + 9; +/// Validates complete sparse keys without materializing their values. +/// +/// Scratch buffers retain capacity between keys; validation does not depend on region metadata. +#[derive(Default)] +pub struct SparsePrimaryKeyValidator { + column_ids: Vec, + label_buf: Vec, +} + +impl SparsePrimaryKeyValidator { + /// Checks the prefix, all label framing and UTF-8, and rejects duplicate or reserved labels. + pub fn validate(&mut self, bytes: &[u8]) -> Result<()> { + self.column_ids.clear(); + validate_prefix(bytes)?; + let mut remaining = &bytes[TAGS_START_OFFSET..]; + while !remaining.is_empty() { + let id = remaining.first_chunk::().context( + InvalidSparsePrimaryKeySnafu { + reason: "truncated label column id", + }, + )?; + let column_id = u32::from_be_bytes(*id); + ensure!( + !ReservedColumnId::is_reserved(column_id), + InvalidSparsePrimaryKeySnafu { + reason: "reserved column id in sparse labels", + } + ); + self.column_ids.push(column_id); + let encoded = &remaining[COLUMN_ID_ENCODE_SIZE..]; + let (length, ascii) = crate::row_converter::encoded_string_info::(encoded) + .context(DeserializeFieldSnafu)?; + let encoded = &encoded[..length]; + // Framing is already checked; ASCII needs no unchunking to establish valid UTF-8. + if !ascii { + self.label_buf.clear(); + self.label_buf.reserve(length); + Deserializer::new(&encoded[1..]) + .read_bytes_into(&mut self.label_buf) + .context(DeserializeFieldSnafu)?; + std::str::from_utf8(&self.label_buf).map_err(|_| { + InvalidSparsePrimaryKeySnafu { + reason: "label is not valid UTF-8", + } + .build() + })?; + } + remaining = &remaining[COLUMN_ID_ENCODE_SIZE + length..]; + } + self.column_ids.sort_unstable(); + ensure!( + !self.column_ids.windows(2).any(|pair| pair[0] == pair[1]), + InvalidSparsePrimaryKeySnafu { + reason: "duplicate column id in sparse key", + } + ); + Ok(()) + } +} + /// Inline capacity for the small-vec fast path of [`SparseOffsetsCache`]. /// /// Most sparse primary keys carry only a handful of tags, so a linear scan @@ -225,6 +283,7 @@ impl SparsePrimaryKeyCodec { let columns = columns_ids.collect(); Self { inner: Arc::new(SparsePrimaryKeyCodecInner { + fields: None, table_id_field: SortField::new(ConcreteDataType::uint32_datatype()), tsid_field: SortField::new(ConcreteDataType::uint64_datatype()), label_field: SortField::new(ConcreteDataType::string_datatype()), @@ -235,7 +294,16 @@ impl SparsePrimaryKeyCodec { /// Creates a new [`SparsePrimaryKeyCodec`] instance. pub fn new(region_metadata: &RegionMetadataRef) -> Self { - Self::from_columns(region_metadata.primary_key_columns().map(|c| c.column_id)) + let fields = region_metadata + .primary_key_columns() + .map(|c| { + ( + c.column_id, + SortField::new(c.column_schema.data_type.clone()), + ) + }) + .collect(); + Self::with_fields(fields) } /// Returns a new [`SparsePrimaryKeyCodec`] instance. @@ -244,6 +312,7 @@ impl SparsePrimaryKeyCodec { pub fn schemaless() -> Self { Self { inner: Arc::new(SparsePrimaryKeyCodecInner { + fields: None, table_id_field: SortField::new(ConcreteDataType::uint32_datatype()), tsid_field: SortField::new(ConcreteDataType::uint64_datatype()), label_field: SortField::new(ConcreteDataType::string_datatype()), @@ -252,27 +321,20 @@ impl SparsePrimaryKeyCodec { } } - /// Creates a new [`SparsePrimaryKeyCodec`] instance with additional label `fields`. + /// Creates a schema-aware sparse codec. pub fn with_fields(fields: Vec<(ColumnId, SortField)>) -> Self { - Self { - inner: Arc::new(SparsePrimaryKeyCodecInner { - columns: Some(fields.iter().map(|f| f.0).collect()), - table_id_field: SortField::new(ConcreteDataType::uint32_datatype()), - tsid_field: SortField::new(ConcreteDataType::uint64_datatype()), - label_field: SortField::new(ConcreteDataType::string_datatype()), - }), - } + let mut codec = Self::from_columns(fields.iter().map(|f| f.0)); + Arc::make_mut(&mut codec.inner).fields = Some(fields); + codec } /// Returns the field of the given column id. fn get_field(&self, column_id: ColumnId) -> Option<&SortField> { - // if the `columns` is not specified, all unknown columns is primary key(label field). if let Some(columns) = &self.inner.columns && !columns.contains(&column_id) { return None; } - match column_id { RESERVED_COLUMN_ID_TABLE_ID => Some(&self.inner.table_id_field), RESERVED_COLUMN_ID_TSID => Some(&self.inner.tsid_field), @@ -405,32 +467,50 @@ impl SparsePrimaryKeyCodec { /// Decodes the given bytes into a [`SparseValues`]. fn decode_sparse(&self, bytes: &[u8]) -> Result { - let mut deserializer = Deserializer::new(bytes); let mut values = SparseValues::with_capacity(16); - - let column_id = u32::deserialize(&mut deserializer).context(DeserializeFieldSnafu)?; - let value = self.inner.table_id_field.deserialize(&mut deserializer)?; - values.insert(column_id, value); - - let column_id = u32::deserialize(&mut deserializer).context(DeserializeFieldSnafu)?; - let value = self.inner.tsid_field.deserialize(&mut deserializer)?; - values.insert(column_id, value); - while deserializer.has_remaining() { - let column_id = u32::deserialize(&mut deserializer).context(DeserializeFieldSnafu)?; - let value = self.inner.label_field.deserialize(&mut deserializer)?; + let (table_id, tsid) = self.decode_ids(bytes)?; + values.insert(RESERVED_COLUMN_ID_TABLE_ID, Value::UInt32(table_id)); + values.insert(RESERVED_COLUMN_ID_TSID, Value::UInt64(tsid)); + let mut remaining = &bytes[TAGS_START_OFFSET..]; + while !remaining.is_empty() { + let id = remaining.first_chunk::().context( + InvalidSparsePrimaryKeySnafu { + reason: "truncated label column id", + }, + )?; + let column_id = u32::from_be_bytes(*id); + ensure!( + !ReservedColumnId::is_reserved(column_id), + InvalidSparsePrimaryKeySnafu { + reason: "reserved column id in sparse labels" + } + ); + let encoded = &remaining[COLUMN_ID_ENCODE_SIZE..]; + let length = + crate::row_converter::encoded_string_len(encoded).context(DeserializeFieldSnafu)?; + let value = self + .inner + .label_field + .deserialize(&mut Deserializer::new(&encoded[..length]))?; + ensure!( + values.get(&column_id).is_none(), + InvalidSparsePrimaryKeySnafu { + reason: "duplicate column id in sparse key", + } + ); values.insert(column_id, value); + remaining = &encoded[length..]; } - Ok(values) } - /// Decodes the given bytes into a [`Value`]. + /// Decodes the reserved table ID from the key prefix. fn decode_leftmost(&self, bytes: &[u8]) -> Result> { - let mut deserializer = Deserializer::new(bytes); - // Skip the column id. - deserializer.advance(COLUMN_ID_ENCODE_SIZE); - let value = self.inner.table_id_field.deserialize(&mut deserializer)?; - Ok(Some(value)) + if bytes.is_empty() { + return Ok(None); + } + self.decode_ids(bytes) + .map(|(table_id, _)| Some(Value::UInt32(table_id))) } /// Returns the offset of the given column id in the given primary key. @@ -529,11 +609,18 @@ impl SparsePrimaryKeyCodec { } impl PrimaryKeyCodec for SparsePrimaryKeyCodec { - fn encode_key_value(&self, _key_value: &KeyValue, _buffer: &mut Vec) -> Result<()> { - UnsupportedOperationSnafu { - err_msg: "The encode_key_value method is not supported in SparsePrimaryKeyCodec.", - } - .fail() + fn encode_key_value(&self, key_value: &KeyValue, buffer: &mut Vec) -> Result<()> { + let fields = self + .inner + .fields + .as_ref() + .context(UnsupportedOperationSnafu { + err_msg: "raw row encoding requires primary key metadata", + })?; + self.encode_to_vec( + fields.iter().map(|f| f.0).zip(key_value.primary_keys()), + buffer, + ) } fn encode_values(&self, values: &[(ColumnId, Value)], buffer: &mut Vec) -> Result<()> { @@ -581,6 +668,21 @@ impl PrimaryKeyCodec for SparsePrimaryKeyCodec { } } +fn validate_prefix(pk: &[u8]) -> Result<()> { + ensure!( + pk.len() >= TAGS_START_OFFSET + && pk[..COLUMN_ID_ENCODE_SIZE] == RESERVED_COLUMN_ID_TABLE_ID.to_be_bytes() + && pk[TSID_VALUE_OFFSET - COLUMN_ID_ENCODE_SIZE..TSID_VALUE_OFFSET] + == RESERVED_COLUMN_ID_TSID.to_be_bytes() + && pk[TABLE_ID_VALUE_OFFSET] == 1 + && pk[TSID_VALUE_OFFSET] == 1, + InvalidSparsePrimaryKeySnafu { + reason: "invalid table_id/tsid prefix" + } + ); + Ok(()) +} + #[cfg(test)] mod tests { use std::sync::Arc; @@ -846,6 +948,12 @@ mod tests { assert!(!buffer.is_empty()); let result = codec.decode_leftmost(&buffer).unwrap().unwrap(); assert_eq!(result, Value::UInt32(42)); + assert_eq!(codec.decode_leftmost(&[]).unwrap(), None); + assert!( + codec + .decode_leftmost(&buffer[..TAGS_START_OFFSET - 1]) + .is_err() + ); } #[test] @@ -863,6 +971,132 @@ mod tests { assert!(codec.decode_ids(&buffer[..buffer.len() - 1]).is_err()); } + #[test] + fn test_sparse_validator_matches_decoder() { + let codec = SparsePrimaryKeyCodec::schemaless(); + let mut validator = SparsePrimaryKeyValidator::default(); + let mut prefix = Vec::new(); + codec.encode_internal(0, 0, &mut prefix).unwrap(); + let mut cases = vec![prefix.clone()]; + let mut max_prefix = Vec::new(); + codec + .encode_internal(u32::MAX, u64::MAX, &mut max_prefix) + .unwrap(); + cases.push(max_prefix); + + for label in ["", "a", "12345678", "123456789", "1234567🦀末尾", "日本語"] + .into_iter() + .map(String::from) + .chain([ + "a".repeat(128), + format!("{}🦀", "a".repeat(127)), + format!("🦀{}", "a".repeat(127)), + ]) + { + let mut key = prefix.clone(); + codec + .encode_raw_tag_value([(37, label.as_bytes())].into_iter(), &mut key) + .unwrap(); + assert!(validator.validate(&key).is_ok()); + cases.push(key); + } + + // Null labels, unknown IDs and arbitrary label order are accepted by the decoder. + let mut labels = prefix.clone(); + labels.extend_from_slice(&12_u32.to_be_bytes()); + labels.push(0); + codec + .encode_raw_tag_value( + [(999, b"unknown".as_slice()), (0, b"first".as_slice())].into_iter(), + &mut labels, + ) + .unwrap(); + assert!(validator.validate(&labels).is_ok()); + cases.push(labels.clone()); + + // Padding outside the significant bytes has never been required to be zero. + let mut padded = cases[3].clone(); + let padding_offset = padded.len() - 2; + padded[padding_offset] = 0xff; + assert!(codec.decode(&padded).is_ok()); + assert!(validator.validate(&padded).is_ok()); + cases.push(padded); + + for id in [ + 12, + RESERVED_COLUMN_ID_TABLE_ID, + RESERVED_COLUMN_ID_TSID, + u32::MAX, + ] { + let mut invalid = labels.clone(); + invalid.extend_from_slice(&id.to_be_bytes()); + invalid.push(0); + assert!(codec.decode(&invalid).is_err()); + assert!(validator.validate(&invalid).is_err()); + cases.push(invalid); + } + for label in [ + b"\xff".as_slice(), + b"\xc0\xaf", + b"\xed\xa0\x80", + b"1234567\xf4\x90\x80\x80", + ] { + let mut invalid = prefix.clone(); + codec + .encode_raw_tag_value([(37, label)].into_iter(), &mut invalid) + .unwrap(); + assert!(codec.decode(&invalid).is_err()); + assert!(validator.validate(&invalid).is_err()); + cases.push(invalid); + } + + // Reuse the same scratch state across successes and failures, including every + // truncation and byte substitution in this corpus of framing and UTF-8 boundaries. + let mut compare = |bytes: &[u8]| { + assert_eq!( + validator.validate(bytes).is_ok(), + codec.decode(bytes).is_ok(), + "{bytes:02x?}" + ); + }; + for key in cases { + compare(&key); + for end in 0..key.len() { + compare(&key[..end]); + } + for offset in 0..key.len() { + let mut mutated = key.clone(); + for byte in 0..=u8::MAX { + mutated[offset] = byte; + compare(&mutated); + } + } + } + } + + #[test] + fn test_sparse_validator_reuses_state_for_unordered_columns() { + let codec = SparsePrimaryKeyCodec::schemaless(); + let mut validator = SparsePrimaryKeyValidator::default(); + for num_tags in [512, 32, 1] { + let mut key = Vec::new(); + codec.encode_internal(0, 0, &mut key).unwrap(); + codec + .encode_raw_tag_value( + (0..num_tags).map(|index| (index * 17 % num_tags, "1234567🦀".as_bytes())), + &mut key, + ) + .unwrap(); + assert!(validator.validate(&key).is_ok()); + + let mut duplicate = key.clone(); + duplicate.extend_from_slice(&0_u32.to_be_bytes()); + duplicate.push(0); + assert!(validator.validate(&duplicate).is_err()); + assert!(validator.validate(&key).is_ok()); + } + } + #[test] fn test_has_column() { let region_metadata = test_region_metadata(); diff --git a/src/mito-codec/src/row_converter/sparse/checked.rs b/src/mito-codec/src/row_converter/sparse/checked.rs index fec89b74e0d..3a3aac097c6 100644 --- a/src/mito-codec/src/row_converter/sparse/checked.rs +++ b/src/mito-codec/src/row_converter/sparse/checked.rs @@ -21,6 +21,7 @@ use crate::row_converter::encoded_string_len; use crate::row_converter::sparse::{ COLUMN_ID_ENCODE_SIZE, RESERVED_COLUMN_ID_TABLE_ID, RESERVED_COLUMN_ID_TSID, SparseOffsetsCache, TABLE_ID_VALUE_OFFSET, TAGS_START_OFFSET, TSID_VALUE_OFFSET, + validate_prefix, }; /// A borrowing view of one sparse primary key, with checked, lazy field lookup. @@ -35,17 +36,7 @@ pub struct SparsePrimaryKeyView<'a, 'b> { impl<'a, 'b> SparsePrimaryKeyView<'a, 'b> { /// Validates the fixed prefix and binds fresh offset discovery to this key. pub fn new(pk: &'a [u8], cache: &'b mut SparseOffsetsCache) -> Result { - ensure!( - pk.len() >= TAGS_START_OFFSET - && pk[..COLUMN_ID_ENCODE_SIZE] == RESERVED_COLUMN_ID_TABLE_ID.to_be_bytes() - && pk[TSID_VALUE_OFFSET - COLUMN_ID_ENCODE_SIZE..TSID_VALUE_OFFSET] - == RESERVED_COLUMN_ID_TSID.to_be_bytes() - && pk[TABLE_ID_VALUE_OFFSET] == 1 - && pk[TSID_VALUE_OFFSET] == 1, - InvalidSparsePrimaryKeySnafu { - reason: "invalid table_id/tsid prefix" - } - ); + validate_prefix(pk)?; cache.clear(); Ok(Self { pk, cache }) } diff --git a/src/mito2/Cargo.toml b/src/mito2/Cargo.toml index ba6cbe5d8cb..a5e26fcd60f 100644 --- a/src/mito2/Cargo.toml +++ b/src/mito2/Cargo.toml @@ -159,3 +159,7 @@ required-features = ["testing"] name = "bench_compaction_merge" harness = false required-features = ["test"] + +[[bench]] +name = "bench_sparse_write" +harness = false diff --git a/src/mito2/benches/bench_sparse_write.rs b/src/mito2/benches/bench_sparse_write.rs new file mode 100644 index 00000000000..4a641b46a1e --- /dev/null +++ b/src/mito2/benches/bench_sparse_write.rs @@ -0,0 +1,391 @@ +// 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. + +//! Measures the sparse row and bulk conversion paths used by ingestion and WAL replay. +//! Input preparation, TSID calculation, protocol decoding, and I/O are excluded. + +use std::hint::black_box; +use std::sync::Arc; + +use api::helper::{ColumnDataTypeWrapper, to_grpc_value}; +use api::v1::value::ValueData; +use api::v1::{ColumnDataType, Mutation, OpType, Row, Rows, SemanticType, WriteHint}; +use criterion::{BatchSize, BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; +use datatypes::arrow::array::{ + ArrayRef, BinaryArray, BinaryDictionaryBuilder, Float64Array, TimestampMillisecondArray, +}; +use datatypes::arrow::compute; +use datatypes::arrow::datatypes::{DataType, Field, Schema, UInt32Type}; +use datatypes::arrow::record_batch::RecordBatch; +use datatypes::data_type::ConcreteDataType; +use datatypes::schema::ColumnSchema; +use mito_codec::key_values::KeyValues; +use mito_codec::row_converter::{self, CompositeValues, PrimaryKeyCodec, SparsePrimaryKeyCodec}; +use mito2::memtable::bulk::part::{self, BulkPart, BulkPartConverter}; +use mito2::sst::{self, FlatSchemaOptions}; +use store_api::codec::PrimaryKeyEncoding; +use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder, RegionMetadataRef}; +use store_api::metric_engine_consts::{ + DATA_SCHEMA_TABLE_ID_COLUMN_NAME, DATA_SCHEMA_TSID_COLUMN_NAME, +}; +use store_api::storage::RegionId; +use store_api::storage::consts::{PRIMARY_KEY_COLUMN_NAME, ReservedColumnId}; + +const ROWS_PER_BATCH: usize = 4096; +const TABLE_ID: u32 = 1024; + +fn build_metadata(num_tags: u32) -> RegionMetadataRef { + let mut builder = RegionMetadataBuilder::new(RegionId::new(TABLE_ID, 0)); + let mut primary_key = vec![ReservedColumnId::table_id(), ReservedColumnId::tsid()]; + for (name, data_type, column_id) in [ + ( + DATA_SCHEMA_TABLE_ID_COLUMN_NAME, + ConcreteDataType::uint32_datatype(), + ReservedColumnId::table_id(), + ), + ( + DATA_SCHEMA_TSID_COLUMN_NAME, + ConcreteDataType::uint64_datatype(), + ReservedColumnId::tsid(), + ), + ] { + builder.push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new(name, data_type, false), + semantic_type: SemanticType::Tag, + column_id, + }); + } + for index in 0..num_tags { + let column_id = index * 17 % num_tags; + builder.push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + format!("tag_{column_id}"), + ConcreteDataType::string_datatype(), + true, + ), + semantic_type: SemanticType::Tag, + column_id, + }); + primary_key.push(column_id); + } + builder + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: num_tags, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new("value", ConcreteDataType::float64_datatype(), true), + semantic_type: SemanticType::Field, + column_id: num_tags + 1, + }) + .primary_key(primary_key) + .primary_key_encoding(PrimaryKeyEncoding::Sparse); + Arc::new(builder.build().unwrap()) +} + +fn build_keys(num_series: usize, num_tags: u32, value_bytes: usize, unicode: bool) -> Vec> { + let codec = SparsePrimaryKeyCodec::schemaless(); + (0..num_series) + .map(|series| { + let tags: Vec<_> = (0..num_tags) + .map(|index| { + // Label name order need not match physical column ID order. + let column_id = index * 17 % num_tags; + let mut value = format!("{series:08x}{column_id:08x}").into_bytes(); + assert!(value.len() <= value_bytes); + value.resize(value_bytes, b'x'); + if unicode { + // The code point crosses an eight-byte encoding chunk boundary. + let start = value_bytes - 9; + value[start..start + 4].copy_from_slice("🦀".as_bytes()); + } + (column_id, value) + }) + .collect(); + let mut key = Vec::new(); + // Permute distinct TSIDs so even one-row-per-series input needs sorting. + let tsid = (series * 17 % num_series) as u64; + codec.encode_internal(TABLE_ID, tsid, &mut key).unwrap(); + codec + .encode_raw_tag_value( + tags.iter().map(|(id, value)| (*id, value.as_slice())), + &mut key, + ) + .unwrap(); + assert_eq!(codec.decode_ids(&key).unwrap(), (TABLE_ID, tsid)); + key + }) + .collect() +} + +fn build_rows(metadata: &RegionMetadataRef, keys: &[Vec]) -> KeyValues { + let schema = [ + ( + PRIMARY_KEY_COLUMN_NAME, + ColumnDataType::Binary, + SemanticType::Tag, + ), + ( + "ts", + ColumnDataType::TimestampMillisecond, + SemanticType::Timestamp, + ), + ("value", ColumnDataType::Float64, SemanticType::Field), + ] + .into_iter() + .map(|(name, data_type, semantic_type)| api::v1::ColumnSchema { + column_name: name.to_string(), + datatype: data_type as i32, + semantic_type: semantic_type as i32, + ..Default::default() + }) + .collect(); + let rows = (0..ROWS_PER_BATCH) + .map(|row| Row { + values: vec![ + ValueData::BinaryValue(keys[row % keys.len()].clone()).into(), + ValueData::TimestampMillisecondValue((row / keys.len()) as i64 * 1000).into(), + ValueData::F64Value(row as f64).into(), + ], + }) + .collect(); + KeyValues::new( + metadata, + Mutation { + op_type: OpType::Put as i32, + sequence: 1, + rows: Some(Rows { schema, rows }), + write_hint: Some(WriteHint { + primary_key_encoding: api::v1::PrimaryKeyEncoding::Sparse.into(), + }), + }, + ) + .unwrap() +} + +fn build_bulk(keys: &[Vec], dictionary: bool) -> BulkPart { + let key_column: ArrayRef = if dictionary { + let mut builder = BinaryDictionaryBuilder::::new(); + for row in 0..ROWS_PER_BATCH { + builder.append(&keys[row % keys.len()]).unwrap(); + } + let array = builder.finish(); + assert_eq!(array.values().len(), keys.len()); + Arc::new(array) + } else { + Arc::new(BinaryArray::from_iter_values( + (0..ROWS_PER_BATCH).map(|row| &keys[row % keys.len()]), + )) + }; + let columns: Vec = vec![ + key_column, + Arc::new(TimestampMillisecondArray::from_iter_values( + (0..ROWS_PER_BATCH).map(|row| (row / keys.len()) as i64 * 1000), + )), + Arc::new(Float64Array::from_iter_values( + (0..ROWS_PER_BATCH).map(|row| row as f64), + )), + ]; + let schema = Arc::new(Schema::new( + [PRIMARY_KEY_COLUMN_NAME, "ts", "value"] + .into_iter() + .zip(&columns) + .map(|(name, column)| Field::new(name, column.data_type().clone(), false)) + .collect::>(), + )); + BulkPart { + batch: RecordBatch::try_new(schema, columns).unwrap(), + min_timestamp: 0, + max_timestamp: (ROWS_PER_BATCH / keys.len() - 1) as i64 * 1000, + sequence: 1, + min_sequence: 1, + timestamp_index: 1, + raw_data: None, + } +} + +fn build_raw_rows(metadata: &RegionMetadataRef, keys: &[Vec]) -> KeyValues { + let codec = SparsePrimaryKeyCodec::schemaless(); + let tags: Vec> = keys + .iter() + .map(|key| { + let CompositeValues::Sparse(values) = codec.decode(key).unwrap() else { + unreachable!() + }; + metadata + .primary_key + .iter() + .map(|id| to_grpc_value(values.get(id).unwrap().clone())) + .collect() + }) + .collect(); + let schema = metadata + .column_metadatas + .iter() + .map(|column| api::v1::ColumnSchema { + column_name: column.column_schema.name.clone(), + datatype: ColumnDataTypeWrapper::try_from(column.column_schema.data_type.clone()) + .unwrap() + .datatype() as i32, + semantic_type: column.semantic_type as i32, + ..Default::default() + }) + .collect(); + let rows = (0..ROWS_PER_BATCH) + .map(|row| { + let mut values = tags[row % tags.len()].clone(); + values.push( + ValueData::TimestampMillisecondValue((row / tags.len()) as i64 * 1000).into(), + ); + values.push(ValueData::F64Value(row as f64).into()); + Row { values } + }) + .collect(); + KeyValues::new( + metadata, + Mutation { + op_type: OpType::Put as i32, + sequence: 1, + rows: Some(Rows { schema, rows }), + write_hint: None, + }, + ) + .unwrap() +} + +// Check every output row outside timing, including repeated series and the sort order. +fn check_output(part: &BulkPart, keys: &[Vec]) { + assert_eq!(part.batch.num_rows(), ROWS_PER_BATCH); + let encoded = compute::cast( + part.batch.column_by_name(PRIMARY_KEY_COLUMN_NAME).unwrap(), + &DataType::Binary, + ) + .unwrap(); + let encoded = encoded.as_any().downcast_ref::().unwrap(); + let timestamps = part + .batch + .column_by_name("ts") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let values = part + .batch + .column_by_name("value") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let rows_per_series = ROWS_PER_BATCH / keys.len(); + let mut series_order: Vec<_> = (0..keys.len()).collect(); + series_order.sort_unstable_by_key(|&series| &keys[series]); + for row in 0..ROWS_PER_BATCH { + let series = series_order[row / rows_per_series]; + let sample = row % rows_per_series; + assert_eq!(encoded.value(row), keys[series]); + assert_eq!(timestamps.value(row), sample as i64 * 1000); + assert_eq!(values.value(row), (sample * keys.len() + series) as f64); + } +} + +fn bench_sparse_write(c: &mut Criterion) { + let mut group = c.benchmark_group("sparse_write"); + group.throughput(Throughput::Elements(ROWS_PER_BATCH as u64)); + + for (num_tags, value_bytes, unicode) in [ + (8, 16, false), + (32, 16, false), + (32, 128, false), + (32, 128, true), + (512, 16, true), + ] { + let metadata = build_metadata(num_tags); + let codec = row_converter::build_primary_key_codec(&metadata); + let schema = sst::to_flat_sst_arrow_schema( + &metadata, + &FlatSchemaOptions::from_encoding(PrimaryKeyEncoding::Sparse), + ); + for num_series in [64, ROWS_PER_BATCH] { + let suffix = if unicode { "_unicode" } else { "" }; + let name = format!("{num_tags}tags_{value_bytes}bytes_{num_series}series{suffix}"); + let keys = build_keys(num_series, num_tags, value_bytes, unicode); + let rows = build_rows(&metadata, &keys); + let new_converter = || { + BulkPartConverter::new( + &metadata, + schema.clone(), + ROWS_PER_BATCH, + codec.clone(), + true, + ) + }; + + let mut converter = new_converter(); + converter.append_key_values(&rows).unwrap(); + check_output(&converter.convert().unwrap(), &keys); + group.bench_function(BenchmarkId::new("rows", &name), |b| { + b.iter_batched( + new_converter, + |mut converter| { + converter.append_key_values(black_box(&rows)).unwrap(); + black_box(converter.convert().unwrap()) + }, + // A wide-label batch can own tens of MiB; do not retain many outputs. + BatchSize::PerIteration, + ); + }); + + for (input, dictionary) in [("bulk_binary", false), ("bulk_dictionary", true)] { + let part = build_bulk(&keys, dictionary); + let convert = |part| { + part::convert_bulk_part(part, &metadata, codec.clone(), schema.clone(), true) + .unwrap() + .unwrap() + }; + check_output(&convert(part.clone()), &keys); + group.bench_function(BenchmarkId::new(input, &name), |b| { + b.iter_batched( + || part.clone(), + |part| black_box(convert(part)), + BatchSize::PerIteration, + ); + }); + } + + let raw_rows = build_raw_rows(&metadata, &keys); + let mut converter = new_converter(); + converter.append_key_values(&raw_rows).unwrap(); + check_output(&converter.convert().unwrap(), &keys); + group.bench_function(BenchmarkId::new("raw_rows", &name), |b| { + b.iter_batched( + new_converter, + |mut converter| { + converter.append_key_values(black_box(&raw_rows)).unwrap(); + black_box(converter.convert().unwrap()) + }, + BatchSize::PerIteration, + ); + }); + } + } + group.finish(); +} + +criterion_group!(benches, bench_sparse_write); +criterion_main!(benches); diff --git a/src/mito2/src/compaction/scheduler_test.rs b/src/mito2/src/compaction/scheduler_test.rs index 2cd556ad7c3..8f00fd80817 100644 --- a/src/mito2/src/compaction/scheduler_test.rs +++ b/src/mito2/src/compaction/scheduler_test.rs @@ -1846,6 +1846,7 @@ async fn test_compaction_bypass_in_staging_mode() { manifest_cache: None, }, FormatType::PrimaryKey, + false, &Default::default(), ) .await diff --git a/src/mito2/src/engine/bulk_insert_validation_test.rs b/src/mito2/src/engine/bulk_insert_validation_test.rs index 9652152031f..31b5473b73e 100644 --- a/src/mito2/src/engine/bulk_insert_validation_test.rs +++ b/src/mito2/src/engine/bulk_insert_validation_test.rs @@ -22,8 +22,8 @@ use common_error::status_code::StatusCode; use common_recordbatch::{DfRecordBatch, RecordBatches}; use common_test_util::flight::encode_to_flight_data; use datatypes::arrow::array::{ - ArrayRef, Float64Array, Int64Array, LargeStringArray, ListArray, StringArray, StringViewArray, - TimestampMillisecondArray, TimestampSecondArray, + ArrayRef, BinaryArray, Float64Array, Int64Array, LargeStringArray, ListArray, StringArray, + StringViewArray, TimestampMillisecondArray, TimestampSecondArray, }; use datatypes::arrow::datatypes::{Field, Int64Type, Schema}; use datatypes::prelude::ConcreteDataType; @@ -91,6 +91,101 @@ async fn scan_num_rows(engine: &MitoEngine, region_id: RegionId) -> usize { batches.iter().map(|batch| batch.num_rows()).sum() } +#[tokio::test] +async fn test_encoded_sparse_keys_reject_whole_batch_and_replay() { + use api::v1::value::ValueData; + use api::v1::{ColumnDataType, WriteHint}; + use mito_codec::row_converter::SparsePrimaryKeyCodec; + + for bulk in [false, true] { + let mut env = TestEnv::new().await; + let engine = env.create_engine(MitoConfig::default()).await; + let region_id = RegionId::new(1, 1); + let mut create = CreateRequestBuilder::new() + .insert_option("primary_key_encoding", "sparse") + .build(); + let reserved = crate::test_util::sst_util::sst_region_metadata_with_encoding( + store_api::codec::PrimaryKeyEncoding::Sparse, + ); + create + .column_metadatas + .extend(reserved.column_metadatas.into_iter().take(2)); + create + .primary_key + .splice(0..0, reserved.primary_key.into_iter().take(2)); + let table_dir = create.table_dir.clone(); + let options = create.options.clone(); + let mut schema = rows_schema(&create); + schema.truncate(3); + schema[0].column_name = "__primary_key".to_string(); + schema[0].datatype = ColumnDataType::Binary as i32; + schema.swap(1, 2); + engine + .handle_request(region_id, RegionRequest::Create(create)) + .await + .unwrap(); + + let codec = SparsePrimaryKeyCodec::schemaless(); + let mut key = Vec::new(); + codec.encode_internal(42, 7, &mut key).unwrap(); + codec + .encode_raw_tag_value([(0, "1234567🦀".as_bytes())].into_iter(), &mut key) + .unwrap(); + let mut malformed = key.clone(); + malformed.push(0); // A truncated label ID after a valid Unicode label. + let request = |second: &[u8]| { + if bulk { + bulk_request( + region_id, + vec![ + column( + "__primary_key", + Arc::new(BinaryArray::from_vec(vec![key.as_slice(), second])), + ), + field(), + ts(), + ], + ) + } else { + let mut rows = build_rows(0, 2); + for (row, key) in rows.iter_mut().zip([key.as_slice(), second]) { + row.values[0].value_data = Some(ValueData::BinaryValue(key.to_vec())); + row.values.swap(1, 2); + } + RegionRequest::Put(RegionPutRequest { + rows: Rows { + schema: schema.clone(), + rows, + }, + hint: Some(WriteHint { + primary_key_encoding: api::v1::PrimaryKeyEncoding::Sparse as i32, + }), + skip_wal: false, + partition_expr_version: None, + }) + } + }; + let error = engine + .handle_request(region_id, request(&malformed)) + .await + .unwrap_err(); + assert_eq!(StatusCode::InvalidArguments, error.status_code()); + assert_eq!(0, scan_num_rows(&engine, region_id).await); + // Byte validation follows WAL append; replay must also discard the malformed batch. + reopen_region(&engine, region_id, table_dir.clone(), true, options.clone()).await; + assert_eq!(0, scan_num_rows(&engine, region_id).await); + engine + .handle_request(region_id, request(&key)) + .await + .unwrap(); + assert_eq!(2, scan_num_rows(&engine, region_id).await); + + reopen_region(&engine, region_id, table_dir, true, options).await; + assert_eq!(2, scan_num_rows(&engine, region_id).await); + engine.stop().await.unwrap(); + } +} + /// Batches that do not match the region schema are rejected before they are /// appended to the WAL, and the region keeps accepting writes. #[tokio::test] @@ -128,6 +223,28 @@ async fn test_bulk_insert_rejects_batches_not_matching_schema() { column("ts", Arc::new(TimestampSecondArray::from(vec![0, 1]))), ], vec![tag(), field()], + vec![field(), ts()], + vec![ + column("tag_0", Arc::new(Int64Array::from(vec![1, 2]))), + field(), + ts(), + ], + vec![ + column( + "__primary_key", + Arc::new(Float64Array::from(vec![1.0, 2.0])), + ), + field(), + ts(), + ], + vec![ + column( + "__primary_key", + Arc::new(BinaryArray::from(vec![None::<&[u8]>; 2])), + ), + field(), + ts(), + ], vec![ tag(), column("field_0", Arc::new(StringArray::from(vec!["1", "2"]))), @@ -301,6 +418,7 @@ async fn test_bulk_insert_fills_columns_missing_from_aligned_batch() { let engine = env.create_engine(MitoConfig::default()).await; let region_id = RegionId::new(1, i as u32 + 1); let request = CreateRequestBuilder::new() + .tag_num(2) .insert_option("sst_format", sst_format) .build(); let table_dir = request.table_dir.clone(); @@ -325,5 +443,192 @@ async fn test_bulk_insert_fills_columns_missing_from_aligned_batch() { assert_eq!(2, response.affected_rows, "{sst_format}"); reopen_region(&engine, region_id, table_dir, true, options).await; assert_eq!(2, scan_num_rows(&engine, region_id).await, "{sst_format}"); + let stream = engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + for batch in batches.iter() { + let tag = batch.column_by_name("tag_1").unwrap(); + assert_eq!(tag.null_count(), batch.num_rows()); + } + } +} + +#[tokio::test] +async fn test_bulk_insert_json2_physical_layout() { + use datatypes::arrow::array::{NullArray, StructArray}; + use datatypes::arrow::datatypes::DataType; + use datatypes::extension::json::{Json2ExtensionType, JsonMetadata}; + use datatypes::json::JsonSettings; + use datatypes::types::json_type::JsonNativeType; + + for encoding in ["dense", "sparse"] { + for layout in ["v2", "legacy", "missing"] { + let mut env = TestEnv::new().await; + let engine = env.create_engine(MitoConfig::default()).await; + let region_id = RegionId::new(1, 1); + let settings = JsonSettings::default(); + let mut request = CreateRequestBuilder::new() + .field_datatype(ConcreteDataType::json2(JsonNativeType::Null)) + .insert_option("sst_format", "flat") + .insert_option("primary_key_encoding", encoding) + .build(); + if encoding == "sparse" { + let reserved = crate::test_util::sst_util::sst_region_metadata_with_encoding( + store_api::codec::PrimaryKeyEncoding::Sparse, + ); + request + .column_metadatas + .extend(reserved.column_metadatas.into_iter().take(2)); + request + .primary_key + .splice(0..0, reserved.primary_key.into_iter().take(2)); + } + let key_columns = if encoding == "sparse" { + let codec = mito_codec::row_converter::SparsePrimaryKeyCodec::schemaless(); + let keys: Vec<_> = (1..=2) + .map(|tsid| { + let mut key = Vec::new(); + codec.encode_internal(42, tsid, &mut key).unwrap(); + key + }) + .collect(); + vec![ + column( + "__primary_key", + Arc::new(BinaryArray::from_iter_values(&keys)), + ), + ts(), + ] + } else { + vec![tag(), ts()] + }; + let json_column = &mut request.column_metadatas[1].column_schema; + json_column.with_extension_type(&Json2ExtensionType::new(Arc::new(JsonMetadata::new( + settings.clone(), + )))); + let mut columns = key_columns.clone(); + if layout != "missing" { + let (array, metadata) = if layout == "legacy" { + let array = Arc::new(StructArray::from(vec![( + Arc::new(Field::new("trace_id", DataType::Utf8, true)), + Arc::new(StringArray::from(vec!["abc", "def"])) as ArrayRef, + )])) as ArrayRef; + (array, JsonMetadata::new_v1(settings.clone())) + } else { + let mut builder = json_column.create_mutable_vector(2); + for id in ["abc", "def"] { + let value = settings + .encode(serde_json::json!({"trace_id": id})) + .unwrap(); + builder.try_push_value_ref(&value.as_value_ref()).unwrap(); + } + ( + builder.to_vector().to_arrow_array(), + JsonMetadata::new(settings.clone()), + ) + }; + let field = Field::new("field_0", array.data_type().clone(), true) + .with_extension_type(Json2ExtensionType::new(Arc::new(metadata))); + columns.push((field, array)); + } + let table_dir = request.table_dir.clone(); + let options = request.options.clone(); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + if layout == "v2" { + let malformed = Arc::new(StructArray::from(vec![( + Arc::new(Field::new( + datatypes::json::JSON2_REMAINDER_FIELD_NAME, + DataType::Int64, + true, + )), + Arc::new(Int64Array::from(vec![1, 2])) as ArrayRef, + )])) as ArrayRef; + for malformed in [ + malformed, + Arc::new(StringArray::from(vec!["abc", "def"])) as ArrayRef, + ] { + let malformed_field = + Field::new("field_0", malformed.data_type().clone(), true) + .with_extension_type(Json2ExtensionType::new(Arc::new( + JsonMetadata::new(settings.clone()), + ))); + let error = engine + .handle_request( + region_id, + bulk_request( + region_id, + [key_columns.clone(), vec![(malformed_field, malformed)]].concat(), + ), + ) + .await + .unwrap_err(); + assert_eq!(error.status_code(), StatusCode::InvalidArguments); + } + let mut invalid_layout = + serde_json::to_value(JsonMetadata::new(settings.clone())).unwrap(); + invalid_layout["layout_version"] = serde_json::json!(3); + let mut null_field = Field::new("field_0", DataType::Null, true) + .with_extension_type(Json2ExtensionType::new(Arc::new(JsonMetadata::new( + settings.clone(), + )))); + null_field.metadata_mut().insert( + arrow_schema::extension::EXTENSION_TYPE_METADATA_KEY.to_string(), + invalid_layout.to_string(), + ); + let error = engine + .handle_request( + region_id, + bulk_request( + region_id, + [ + key_columns.clone(), + vec![(null_field, Arc::new(NullArray::new(2)) as ArrayRef)], + ] + .concat(), + ), + ) + .await + .unwrap_err(); + assert_eq!(error.status_code(), StatusCode::InvalidArguments); + } + let response = engine + .handle_request(region_id, bulk_request(region_id, columns)) + .await + .unwrap(); + assert_eq!(response.affected_rows, 2); + for phase in 0..3 { + if phase == 1 { + reopen_region(&engine, region_id, table_dir.clone(), true, options.clone()) + .await; + } else if phase == 2 { + crate::test_util::flush_region(&engine, region_id, None).await; + } + let stream = engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 2); + let printed = batches.pretty_print().unwrap(); + if layout != "missing" { + assert!( + printed.contains("abc") && printed.contains("def"), + "{encoding}/{layout}/{phase}: {printed}" + ); + } else { + for batch in batches.iter() { + assert_eq!( + batch.column_by_name("field_0").unwrap().null_count(), + batch.num_rows() + ); + } + } + } + } } } diff --git a/src/mito2/src/engine/create_test.rs b/src/mito2/src/engine/create_test.rs index 57c0c3fa7c2..895d15b4195 100644 --- a/src/mito2/src/engine/create_test.rs +++ b/src/mito2/src/engine/create_test.rs @@ -510,7 +510,7 @@ async fn test_engine_create_preserves_manifest_on_open_error() -> WhateverResult } #[tokio::test] -async fn test_engine_create_json2_sparse_requires_flat() -> WhateverResult<()> { +async fn test_engine_create_json2_sparse_forces_flat() -> WhateverResult<()> { let mut env = TestEnv::with_prefix("create-json2-sparse").await; let engine = env.create_engine(MitoConfig::default()).await; let region_id = RegionId::new(1, 1); @@ -529,21 +529,18 @@ async fn test_engine_create_json2_sparse_requires_flat() -> WhateverResult<()> { request.column_metadatas = metadata.column_metadatas.clone(); request.primary_key = metadata.primary_key.clone(); request.column_metadatas[4].column_schema = json_column; - let err = engine - .handle_request(region_id, RegionRequest::Create(request.clone())) - .await - .unwrap_err(); - assert_eq!(StatusCode::InvalidArguments, err.status_code()); - assert!( - err.output_msg() - .contains("JSON2 columns require sst_format=flat") - ); - assert!(!engine.is_region_exists(region_id)); - - request.options.insert("sst_format".into(), "flat".into()); engine .handle_request(region_id, RegionRequest::Create(request.clone())) .await?; + assert_eq!( + Some(FormatType::Flat), + engine + .get_region(region_id) + .unwrap() + .version() + .options + .sst_format + ); let schema = rows_schema(&request); let mut pk_schema = schema[0].clone(); pk_schema.column_name = PRIMARY_KEY_COLUMN_NAME.into(); diff --git a/src/mito2/src/engine/scan_test.rs b/src/mito2/src/engine/scan_test.rs index 81b9b016772..0651fd2dd4c 100644 --- a/src/mito2/src/engine/scan_test.rs +++ b/src/mito2/src/engine/scan_test.rs @@ -73,6 +73,271 @@ use crate::test_util; use crate::test_util::sst_util::{new_sparse_primary_key, sst_region_metadata_with_encoding}; use crate::test_util::{CreateRequestBuilder, TestEnv, reopen_region}; +#[tokio::test] +async fn test_sparse_upsert_after_adding_nullable_tag() -> WhateverResult<()> { + use datatypes::schema::ColumnSchema; + use store_api::metadata::ColumnMetadata; + use store_api::region_request::AddColumn; + + let mut env = TestEnv::new().await; + let engine = env.create_engine(MitoConfig::default()).await; + let region_id = RegionId::new(1024, 0); + let mut request = CreateRequestBuilder::new() + .tag_num(1) + .field_num(1) + .insert_option("primary_key_encoding", "sparse") + .build(); + let reserved = sst_region_metadata_with_encoding(PrimaryKeyEncoding::Sparse); + request + .column_metadatas + .extend(reserved.column_metadatas.into_iter().take(2)); + request + .primary_key + .splice(0..0, reserved.primary_key.into_iter().take(2)); + let schema = test_util::rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await?; + let mut rows = Rows { + schema, + rows: vec![row(vec![ + ValueData::StringValue("api".into()), + ValueData::F64Value(1.0), + ValueData::TimestampMillisecondValue(1000), + ValueData::U32Value(42), + ValueData::U64Value(7), + ])], + }; + test_util::put_rows(&engine, region_id, rows.clone()).await; + test_util::flush_region(&engine, region_id, None).await; + engine + .handle_request( + region_id, + RegionRequest::Alter(RegionAlterRequest { + kind: AlterKind::AddColumns { + columns: vec![AddColumn { + column_metadata: ColumnMetadata { + column_schema: ColumnSchema::new( + "zone", + ConcreteDataType::string_datatype(), + true, + ), + semantic_type: SemanticType::Tag, + column_id: 3, + }, + location: None, + }], + }, + }), + ) + .await?; + + // The old SST and the new write must still identify the same series before compaction. + rows.rows[0].values[1].value_data = Some(ValueData::F64Value(2.0)); + test_util::put_rows(&engine, region_id, rows).await; + for flush in [false, true] { + if flush { + test_util::flush_region(&engine, region_id, None).await; + } + let stream = engine + .scan_to_stream( + region_id, + ScanRequest { + projection: Some(vec![0, 1, 2, 5]), + ..Default::default() + }, + ) + .await?; + let batches = RecordBatches::try_collect(stream).await?; + let expected = "\ ++-------+---------+---------------------+------+ +| tag_0 | field_0 | ts | zone | ++-------+---------+---------------------+------+ +| api | 2.0 | 1970-01-01T00:00:01 | | ++-------+---------+---------------------+------+"; + assert_eq!(batches.pretty_print()?, expected); + } + engine.stop().await?; + Ok(()) +} + +#[tokio::test] +async fn test_exemplar_sparse_storage_lifecycle() -> WhateverResult<()> { + use datatypes::schema::ColumnSchema; + use store_api::metadata::ColumnMetadata; + use store_api::region_request::AddColumn; + use store_api::storage::consts::ReservedColumnId; + + for default_flat_format in [false, true] { + let mut env = TestEnv::new().await; + let engine = env + .create_engine(MitoConfig { + default_flat_format, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1024, 0); + let mut request = CreateRequestBuilder::new() + .tag_num(2) + .field_num(2) + .insert_option("primary_key_encoding", "sparse") + .insert_option("append_mode", "true") + .insert_option("sst_format", "primary_key") + .build(); + request.column_metadatas[3].column_schema.data_type = + ConcreteDataType::json2(JsonNativeType::Object(JsonObjectType::new())); + request.column_metadatas[3] + .column_schema + .with_extension_type(&Json2ExtensionType::default()); + for (name, column_id, data_type) in [ + ( + "__table_id", + ReservedColumnId::table_id(), + ConcreteDataType::uint32_datatype(), + ), + ( + "__tsid", + ReservedColumnId::tsid(), + ConcreteDataType::uint64_datatype(), + ), + ] { + request.column_metadatas.push(ColumnMetadata { + column_schema: ColumnSchema::new(name, data_type, false), + semantic_type: SemanticType::Tag, + column_id, + }); + } + request.primary_key.splice( + 0..0, + [ReservedColumnId::table_id(), ReservedColumnId::tsid()], + ); + let table_dir = request.table_dir.clone(); + let schema = test_util::rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await?; + let mut rows = Rows { + schema: schema.clone(), + rows: vec![row(vec![ + ValueData::StringValue("metric".into()), + ValueData::StringValue("api".into()), + ValueData::F64Value(1.5), + ValueData::JsonValue(encode_json_value(JsonValue::from( + json!({"trace_id":"abc"}), + ))), + ValueData::TimestampMillisecondValue(1), + ])], + }; + rows.rows[0].values.extend([ + ValueData::U32Value(42).into(), + ValueData::U64Value(7).into(), + ]); + // Equal series/timestamp observations are preserved in append mode. + rows.rows.push(rows.rows[0].clone()); + let mut null_row = rows.rows[0].clone(); + null_row.values[0].value_data = None; + null_row.values[1].value_data = None; + // Different source metrics can have the same TSID. + null_row.values[5] = ValueData::U32Value(43).into(); + rows.rows.push(null_row); + test_util::put_rows(&engine, region_id, rows.clone()).await; + + // Reopen without encoding options, replaying raw rows from WAL. + reopen_region(&engine, region_id, table_dir.clone(), true, HashMap::new()).await; + assert_eq!( + engine.get_primary_key_encoding(region_id).unwrap(), + PrimaryKeyEncoding::Sparse + ); + test_util::flush_region(&engine, region_id, None).await; + engine + .handle_request( + region_id, + RegionRequest::Alter(RegionAlterRequest { + kind: AlterKind::AddColumns { + columns: vec![AddColumn { + column_metadata: ColumnMetadata { + column_schema: ColumnSchema::new( + "new_tag", + ConcreteDataType::string_datatype(), + true, + ), + semantic_type: SemanticType::Tag, + column_id: 5, + }, + location: None, + }], + }, + }), + ) + .await?; + // A stale writer omits the newly added nullable tag. + test_util::put_rows(&engine, region_id, rows).await; + test_util::flush_region(&engine, region_id, None).await; + engine + .handle_request( + region_id, + RegionRequest::Compact(RegionCompactRequest { + options: compact_request::Options::StrictWindow(StrictWindow { + window_seconds: 86400, + }), + ..Default::default() + }), + ) + .await?; + reopen_region( + &engine, + region_id, + table_dir.clone(), + true, + HashMap::from([ + ("primary_key_encoding".into(), "dense".into()), + ("sst_format".into(), "primary_key".into()), + ]), + ) + .await; + reopen_region(&engine, region_id, table_dir, true, HashMap::new()).await; + assert_eq!( + engine.get_primary_key_encoding(region_id).unwrap(), + PrimaryKeyEncoding::Sparse + ); + let mut cases = vec![ + (vec![], 6), + (vec![col("tag_0").eq(lit("metric"))], 4), + (vec![col("tag_0").is_null()], 2), + ]; + cases.extend([ + (vec![col("__table_id").eq(lit(42u32))], 4), + (vec![col("__table_id").eq(lit(43u32))], 2), + (vec![col("__tsid").eq(lit(7u64))], 6), + ]); + for (filters, expected) in cases { + let scanner = engine + .scanner( + region_id, + ScanRequest { + filters, + json_type_hint: HashMap::from([( + "field_1".into(), + JsonNativeType::Variant, + )]), + ..Default::default() + }, + ) + .await?; + let batches = RecordBatches::try_collect(scanner.scan().await?).await?; + assert_eq!( + batches.iter().map(|b| b.num_rows()).sum::(), + expected + ); + for batch in batches.iter() { + let label_index = batch.schema.column_index_by_name("new_tag").unwrap(); + assert_eq!(batch.column(label_index).null_count(), batch.num_rows()); + } + } + } + Ok(()) +} + #[tokio::test] async fn test_nullable_label_filters_in_memtables_and_ssts() { use datafusion_expr::{BinaryExpr as LogicalBinaryExpr, Cast, Expr}; diff --git a/src/mito2/src/manifest/manager.rs b/src/mito2/src/manifest/manager.rs index 436488f3c3c..5471b1883a9 100644 --- a/src/mito2/src/manifest/manager.rs +++ b/src/mito2/src/manifest/manager.rs @@ -170,6 +170,7 @@ impl RegionManifestManager { flushed_entry_id: u64, options: RegionManifestOptions, sst_format: FormatType, + append_mode: bool, stats: &ManifestStats, ) -> Result { // construct storage @@ -195,7 +196,7 @@ impl RegionManifestManager { RegionChange { metadata: metadata.clone(), sst_format, - append_mode: None, + append_mode: Some(append_mode), }, ); let manifest = manifest_builder.try_build()?; @@ -209,7 +210,7 @@ impl RegionManifestManager { let mut actions = vec![RegionMetaAction::Change(RegionChange { metadata, sst_format, - append_mode: None, + append_mode: Some(append_mode), })]; if flushed_entry_id > 0 { actions.push(RegionMetaAction::Edit(RegionEdit { @@ -1068,6 +1069,6 @@ mod test { // get manifest size again let manifest_size = manager.manifest_usage(); - assert_eq!(manifest_size, 1397); + assert_eq!(manifest_size, 1398); } } diff --git a/src/mito2/src/memtable/bulk/part.rs b/src/mito2/src/memtable/bulk/part.rs index 5852d072197..50954f1c6db 100644 --- a/src/mito2/src/memtable/bulk/part.rs +++ b/src/mito2/src/memtable/bulk/part.rs @@ -45,13 +45,14 @@ use datatypes::prelude::{MutableVector, Vector}; use datatypes::value::ValueRef; use datatypes::vectors::Helper; use mito_codec::key_values::{KeyValue, KeyValues}; +use mito_codec::row_converter::sparse::SparsePrimaryKeyValidator; use mito_codec::row_converter::{PrimaryKeyCodec, SortField, build_primary_key_codec_with_fields}; use parquet::arrow::ArrowWriter; use parquet::basic::{Compression, ZstdLevel}; use parquet::file::metadata::ParquetMetaData; use parquet::file::properties::WriterProperties; use smallvec::SmallVec; -use snafu::{OptionExt, ResultExt}; +use snafu::{OptionExt, ResultExt, ensure}; use store_api::codec::PrimaryKeyEncoding; use store_api::metadata::{RegionMetadata, RegionMetadataRef}; use store_api::mito_engine_options::FloatFieldEncoding; @@ -223,7 +224,7 @@ impl BulkPart { let mut columns_to_fill = Vec::new(); for column_meta in ®ion_metadata.column_metadatas { // Sparse bulk batches already carry tags in the encoded primary key. - if region_metadata.primary_key_encoding == PrimaryKeyEncoding::Sparse + if batch_columns.contains(PRIMARY_KEY_COLUMN_NAME) && column_meta.semantic_type == SemanticType::Tag { continue; @@ -266,13 +267,14 @@ impl BulkPart { ), })?; let arrow_array = default_vector.to_arrow_array(); - column_meta.column_schema.data_type.as_arrow_type(); - - new_fields.push(Arc::new(Field::new( - column_meta.column_schema.name.clone(), - column_meta.column_schema.data_type.as_arrow_type(), - column_meta.column_schema.is_nullable(), - ))); + new_fields.push(Arc::new( + Field::new( + column_meta.column_schema.name.clone(), + arrow_array.data_type().clone(), + column_meta.column_schema.is_nullable(), + ) + .with_metadata(column_meta.column_schema.metadata().clone()), + )); new_columns.push(arrow_array); } @@ -598,6 +600,8 @@ pub struct BulkPartConverter { key_buf: Vec, /// Primary key array builder. key_array_builder: PrimaryKeyArrayBuilder, + /// Whether any input key was supplied pre-encoded and still needs validation. + has_pre_encoded_keys: bool, /// Builders for non-primary key columns. value_builder: ValueBuilder, /// Builders for individual primary key columns. @@ -644,6 +648,7 @@ impl BulkPartConverter { primary_key_codec, key_buf: Vec::new(), key_array_builder: PrimaryKeyArrayBuilder::new(), + has_pre_encoded_keys: false, value_builder: ValueBuilder::new(region_metadata, capacity), primary_key_column_builders, min_ts: i64::MAX, @@ -664,10 +669,16 @@ impl BulkPartConverter { /// Appends a [KeyValue] to builders. /// - /// If the primary key uses sparse encoding, callers must encoded the primary key in the [KeyValue]. + /// Metric-sparse hints carry encoded keys; raw tags use the destination codec. fn append_key_value(&mut self, kv: &KeyValue) -> Result<()> { // Handles primary key based on encoding type - if self.primary_key_codec.encoding() == PrimaryKeyEncoding::Sparse { + if kv.primary_key_encoding() == PrimaryKeyEncoding::Sparse { + ensure!( + self.primary_key_codec.encoding() == PrimaryKeyEncoding::Sparse, + error::InvalidRecordBatchSnafu { + reason: "pre-encoded primary keys require a sparse destination layout", + } + ); // For sparse encoding, the primary key is already encoded in the KeyValue // Gets the first (and only) primary key value which contains the encoded key let mut primary_keys = kv.primary_keys(); @@ -682,13 +693,15 @@ impl BulkPartConverter { self.key_array_builder .append(encoded) .context(ComputeArrowSnafu)?; + self.has_pre_encoded_keys = true; } else { - self.key_array_builder - .append("") - .context(ComputeArrowSnafu)?; + return error::InvalidRecordBatchSnafu { + reason: "encoded metric primary key must not be null", + } + .fail(); } } else { - // For dense encoding, we need to encode the primary key columns + // Encode raw primary key columns with the destination codec. self.key_buf.clear(); self.primary_key_codec .encode_key_value(kv, &mut self.key_buf) @@ -752,6 +765,9 @@ impl BulkPartConverter { columns.push(values.timestamp.to_arrow_array()); // Primary key. let pk_array = self.key_array_builder.finish(); + if self.has_pre_encoded_keys { + validate_sparse_dictionary(&pk_array)?; + } columns.push(Arc::new(pk_array)); // Sequence and op type. columns.push(values.sequence.to_arrow_array()); @@ -776,6 +792,35 @@ impl BulkPartConverter { } } +/// Validates each referenced key once; sliced dictionaries may retain unused invalid values. +fn validate_sparse_dictionary(keys: &PrimaryKeyArray) -> Result<()> { + let values = keys + .values() + .as_any() + .downcast_ref::() + .context(error::InvalidRecordBatchSnafu { + reason: "encoded primary keys must be binary", + })?; + ensure!( + keys.logical_null_count() == 0, + error::InvalidRecordBatchSnafu { + reason: "encoded primary key must not be null", + } + ); + let mut validated = vec![false; values.len()]; + let mut validator = SparsePrimaryKeyValidator::default(); + for &index in keys.keys().values() { + let index = index as usize; + if !validated[index] { + validator + .validate(values.value(index)) + .context(EncodeSnafu)?; + validated[index] = true; + } + } + Ok(()) +} + fn new_primary_key_column_builders( metadata: &RegionMetadata, capacity: usize, @@ -972,12 +1017,12 @@ fn lexsort_primary_key_indices(batch: &RecordBatch) -> Result { datatypes::arrow::compute::lexsort_to_indices(&sort_columns, None).context(ComputeArrowSnafu) } -/// Converts a `BulkPart` that is unordered and without encoded primary keys into a `BulkPart` +/// Converts an unordered `BulkPart` with raw tags or encoded primary keys into a `BulkPart` /// with the same format as produced by [BulkPartConverter]. /// /// This function takes a `BulkPart` where: -/// - For dense encoding: Primary key columns may be stored as individual columns -/// - For sparse encoding: The `__primary_key` column should already be present with encoded keys +/// - Raw primary key columns are stored individually and encoded with the destination codec +/// - Sparse destinations also accept encoded keys in a `__primary_key` column /// - The batch may not be sorted /// /// And produces a `BulkPart` where: @@ -1023,13 +1068,36 @@ pub fn convert_bulk_part( // Determines the structure of the input batch by looking up columns by name let mut output_columns = Vec::new(); - // Extracts primary key columns if we need to encode them (dense encoding) - let pk_array = if is_sparse { - // For sparse encoding, the input should already have the __primary_key column - // We need to find it in the input batch - None + // Encoded input is explicit; otherwise encode raw tags with the destination codec. + let pk_dictionary: ArrayRef = if column_indices.contains_key(PRIMARY_KEY_COLUMN_NAME) { + ensure!( + is_sparse, + error::InvalidRecordBatchSnafu { + reason: "pre-encoded bulk keys require a sparse destination layout" + } + ); + let column = part.batch.column(column_indices[PRIMARY_KEY_COLUMN_NAME]); + ensure!( + column.data_type() == &ArrowDataType::Binary + || matches!(column.data_type(), ArrowDataType::Dictionary(_, value) if value.as_ref() == &ArrowDataType::Binary), + error::InvalidRecordBatchSnafu { + reason: "encoded primary keys must be binary" + } + ); + let target_type = ArrowDataType::Dictionary( + Box::new(ArrowDataType::UInt32), + Box::new(ArrowDataType::Binary), + ); + let column = arrow::compute::cast(column, &target_type).context(ComputeArrowSnafu)?; + let keys = column.as_any().downcast_ref::().context( + error::InvalidRecordBatchSnafu { + reason: "encoded primary keys must be a UInt32 binary dictionary", + }, + )?; + validate_sparse_dictionary(keys)?; + column } else { - // For dense encoding, extract and encode primary key columns by name + // Extract and encode primary key columns by name. let pk_vectors: Result> = region_metadata .primary_key_columns() .map(|col_meta| { @@ -1068,7 +1136,7 @@ pub fn convert_bulk_part( .context(ComputeArrowSnafu)?; } - Some(key_array_builder.finish()) + Arc::new(key_array_builder.finish()) }; // Adds primary key columns if storing them (only for dense encoding) @@ -1121,24 +1189,6 @@ pub fn convert_bulk_part( output_columns.push(part.batch.column(*ts_col_idx).clone()); // Adds encoded primary key dictionary column - let pk_dictionary = if let Some(pk_dict_array) = pk_array { - Arc::new(pk_dict_array) as ArrayRef - } else { - let pk_col_idx = - column_indices - .get(PRIMARY_KEY_COLUMN_NAME) - .context(ColumnNotFoundSnafu { - column: PRIMARY_KEY_COLUMN_NAME, - })?; - let col = part.batch.column(*pk_col_idx); - - // Casts to dictionary type if needed - let target_type = ArrowDataType::Dictionary( - Box::new(ArrowDataType::UInt32), - Box::new(ArrowDataType::Binary), - ); - arrow::compute::cast(col, &target_type).context(ComputeArrowSnafu)? - }; output_columns.push(pk_dictionary); let sequence_array = UInt64Array::from(vec![part.sequence; num_rows]); @@ -1147,6 +1197,46 @@ pub fn convert_bulk_part( let op_type_array = UInt8Array::from(vec![OpType::Put as u8; num_rows]); output_columns.push(Arc::new(op_type_array) as ArrayRef); + // JSON2 arrays carry a data-dependent shape and their own physical layout version. + // Keep that layout metadata while retaining the destination's column IDs. + let schema = if region_metadata + .field_columns() + .any(|c| c.column_schema.data_type.is_json2()) + { + let fields = schema + .fields() + .iter() + .map(|field| { + if !datatypes::extension::json::is_json2_extension_type(field) { + return field.clone(); + } + let source = input_schema.field(column_indices[field.name().as_str()]); + let mut metadata = field.metadata().clone(); + if datatypes::extension::json::is_json2_extension_type(source) { + for key in [ + arrow_schema::extension::EXTENSION_TYPE_NAME_KEY, + arrow_schema::extension::EXTENSION_TYPE_METADATA_KEY, + ] { + if let Some(value) = source.metadata().get(key) { + metadata.insert(key.to_string(), value.clone()); + } else { + metadata.remove(key); + } + } + } + Arc::new( + field + .as_ref() + .clone() + .with_data_type(source.data_type().clone()) + .with_metadata(metadata), + ) + }) + .collect::>(); + Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone())) + } else { + schema + }; let batch = RecordBatch::try_new(schema, output_columns).context(NewRecordBatchSnafu)?; // Sorts the batch by (primary_key, timestamp, sequence desc) @@ -2490,7 +2580,6 @@ mod tests { #[allow(clippy::too_many_arguments)] fn build_key_values_with_sparse_encoding( metadata: &RegionMetadataRef, - primary_key_codec: &Arc, table_id: u32, tsid: u64, k0: String, @@ -2507,7 +2596,7 @@ mod tests { (1, Value::String(k1.clone().into())), ]; let mut encoded_key = Vec::new(); - primary_key_codec + mito_codec::row_converter::SparsePrimaryKeyCodec::schemaless() .encode_values(&pk_values, &mut encoded_key) .unwrap(); assert!(!encoded_key.is_empty()); @@ -2627,12 +2716,16 @@ mod tests { assert_eq!(primary_key_codec.encoding(), PrimaryKeyEncoding::Sparse); let capacity = 100; - let mut converter = - BulkPartConverter::new(&metadata, schema, capacity, primary_key_codec.clone(), true); + let mut converter = BulkPartConverter::new( + &metadata, + schema.clone(), + capacity, + primary_key_codec.clone(), + true, + ); let key_values1 = build_key_values_with_sparse_encoding( &metadata, - &primary_key_codec, 2048u32, // table_id 100u64, // tsid "key11".to_string(), @@ -2642,9 +2735,8 @@ mod tests { 1, ); - let key_values2 = build_key_values_with_sparse_encoding( + let mut key_values2 = build_key_values_with_sparse_encoding( &metadata, - &primary_key_codec, 4096u32, // table_id 200u64, // tsid "key12".to_string(), @@ -2653,6 +2745,101 @@ mod tests { vec![Some(3.0)].into_iter(), 2, ); + key_values2.mutation.op_type = OpType::Delete as i32; + + let mut raw_mutation = key_values1.mutation.clone(); + raw_mutation.write_hint = None; + let rows = raw_mutation.rows.as_mut().unwrap(); + rows.schema.splice( + 0..1, + ["k0", "k1"].map(|name| api::v1::ColumnSchema { + column_name: name.to_string(), + datatype: api::v1::ColumnDataType::String as i32, + semantic_type: SemanticType::Tag as i32, + ..Default::default() + }), + ); + for row in &mut rows.rows { + row.values.splice( + 0..1, + ["raw🦀", ""] + .map(|value| api::v1::value::ValueData::StringValue(value.to_string()).into()), + ); + } + rows.schema.extend( + [ + ("__table_id", api::v1::ColumnDataType::Uint32), + ("__tsid", api::v1::ColumnDataType::Uint64), + ] + .map(|(name, datatype)| api::v1::ColumnSchema { + column_name: name.into(), + datatype: datatype as i32, + semantic_type: SemanticType::Tag as i32, + ..Default::default() + }), + ); + for row in &mut rows.rows { + row.values.extend([ + api::v1::value::ValueData::U32Value(2048).into(), + api::v1::value::ValueData::U64Value(100).into(), + ]); + } + let raw = KeyValues::new(&metadata, raw_mutation).unwrap(); + let mut raw_converter = BulkPartConverter::new( + &metadata, + schema.clone(), + capacity, + primary_key_codec.clone(), + true, + ); + raw_converter.append_key_values(&raw).unwrap(); + let raw_part = raw_converter.convert().unwrap(); + assert_eq!(2, raw_part.num_rows()); + let keys = arrow::compute::cast( + raw_part + .batch + .column_by_name(PRIMARY_KEY_COLUMN_NAME) + .unwrap(), + &ArrowDataType::Binary, + ) + .unwrap(); + let keys = keys.as_any().downcast_ref::().unwrap(); + for key in keys.iter().flatten() { + let decoded = primary_key_codec.decode(key).unwrap().into_sparse(); + assert_eq!( + Some(&Value::UInt32(2048)), + decoded.get(&ReservedColumnId::table_id()) + ); + assert_eq!( + Some(&Value::UInt64(100)), + decoded.get(&ReservedColumnId::tsid()) + ); + assert_eq!(Some(&Value::from("raw🦀")), decoded.get(&0)); + assert_eq!(Some(&Value::from("")), decoded.get(&1)); + } + + for op_type in [OpType::Put, OpType::Delete] { + let mut mutation = key_values1.mutation.clone(); + mutation.op_type = op_type as i32; + mutation.rows.as_mut().unwrap().rows[1].values[0].value_data = Some( + api::v1::value::ValueData::BinaryValue(b"truncated".to_vec()), + ); + let invalid = KeyValues::new(&metadata, mutation).unwrap(); + // Raw rows before or after encoded input must not disable its validation. + for batches in [vec![&invalid], vec![&raw, &invalid], vec![&invalid, &raw]] { + let mut invalid_converter = BulkPartConverter::new( + &metadata, + schema.clone(), + capacity, + primary_key_codec.clone(), + true, + ); + for batch in batches { + invalid_converter.append_key_values(batch).unwrap(); + } + assert!(invalid_converter.convert().is_err()); + } + } converter.append_key_values(&key_values1).unwrap(); converter.append_key_values(&key_values2).unwrap(); @@ -2664,6 +2851,26 @@ mod tests { assert_eq!(bulk_part.max_timestamp, 2000); assert_eq!(bulk_part.sequence, 2); assert_eq!(bulk_part.timestamp_index, bulk_part.batch.num_columns() - 4); + assert_eq!( + bulk_part + .batch + .column_by_name("__sequence") + .unwrap() + .as_ref(), + &UInt64Array::from(vec![1, 2, 2]) as &dyn Array, + ); + assert_eq!( + bulk_part + .batch + .column_by_name("__op_type") + .unwrap() + .as_ref(), + &UInt8Array::from(vec![ + OpType::Put as u8, + OpType::Put as u8, + OpType::Delete as u8 + ]) as &dyn Array, + ); // For sparse encoding, primary key columns should NOT be stored individually // even when store_primary_key_columns is true, because sparse encoding @@ -2909,10 +3116,16 @@ mod tests { let primary_key_codec = build_primary_key_codec(&metadata); // Create input batch with __primary_key column (sparse encoding) - let pk_array = Arc::new(arrow::array::BinaryArray::from(vec![ - b"encoded_key_1".as_slice(), - b"encoded_key_2".as_slice(), - ])); + let keys = (1..=2) + .map(|tsid| { + let mut key = Vec::new(); + mito_codec::row_converter::SparsePrimaryKeyCodec::schemaless() + .encode_internal(123, tsid, &mut key) + .unwrap(); + key + }) + .collect::>(); + let pk_array = Arc::new(BinaryArray::from_iter_values(&keys)); let v0_array = Arc::new(arrow::array::Int64Array::from(vec![100, 200])); let v1_array = Arc::new(arrow::array::Float64Array::from(vec![1.0, 2.0])); let ts_array = Arc::new(TimestampMillisecondArray::from(vec![1000, 2000])); @@ -2947,6 +3160,153 @@ mod tests { &FlatSchemaOptions::from_encoding(metadata.primary_key_encoding), ); + let dictionary_values = Arc::new(BinaryArray::from(vec![ + Some(keys[0].as_slice()), + Some(keys[1].as_slice()), + Some(b"truncated".as_slice()), + None, + Some(keys[0].as_slice()), + ])); + let with_keys = |keys: ArrayRef| { + let mut input = part.clone(); + let mut fields = input.batch.schema().fields().to_vec(); + fields[0] = Arc::new(Field::new( + PRIMARY_KEY_COLUMN_NAME, + keys.data_type().clone(), + true, + )); + let mut columns = input.batch.columns().to_vec(); + columns[0] = keys; + input.batch = RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap(); + input + }; + // Slicing retains malformed/null dictionary values outside the selected rows. + let sliced = DictionaryArray::try_new( + UInt32Array::from(vec![2, 0, 1, 3]), + dictionary_values.clone(), + ) + .unwrap() + .slice(1, 2); + for key_type in [ + ArrowDataType::Int8, + ArrowDataType::Int16, + ArrowDataType::Int32, + ArrowDataType::Int64, + ArrowDataType::UInt8, + ArrowDataType::UInt16, + ArrowDataType::UInt32, + ArrowDataType::UInt64, + ] { + let keys = arrow::compute::cast( + &sliced, + &ArrowDataType::Dictionary(Box::new(key_type), Box::new(ArrowDataType::Binary)), + ) + .unwrap(); + let converted = convert_bulk_part( + with_keys(keys), + &metadata, + primary_key_codec.clone(), + output_schema.clone(), + true, + ) + .unwrap() + .unwrap(); + let keys = arrow::compute::cast( + converted + .batch + .column_by_name(PRIMARY_KEY_COLUMN_NAME) + .unwrap(), + &ArrowDataType::Binary, + ) + .unwrap(); + assert_eq!(keys.as_ref(), part.batch.column(0).as_ref()); + for (name, index) in [("v0", 1), ("v1", 2), ("ts", 3)] { + assert_eq!( + converted.batch.column_by_name(name).unwrap().as_ref(), + part.batch.column(index).as_ref() + ); + } + assert_eq!( + converted + .batch + .column_by_name("__sequence") + .unwrap() + .as_ref(), + &UInt64Array::from(vec![7, 7]) as &dyn Array + ); + assert_eq!( + converted + .batch + .column_by_name("__op_type") + .unwrap() + .as_ref(), + &UInt8Array::from(vec![OpType::Put as u8; 2]) as &dyn Array + ); + } + // Equal bytes under distinct dictionary slots remain separate observations. + let repeated = + DictionaryArray::try_new(UInt32Array::from(vec![0, 4]), dictionary_values.clone()) + .unwrap(); + assert_eq!( + convert_bulk_part( + with_keys(Arc::new(repeated)), + &metadata, + primary_key_codec.clone(), + output_schema.clone(), + true + ) + .unwrap() + .unwrap() + .num_rows(), + 2 + ); + for indices in [ + vec![Some(0), Some(2)], + vec![Some(0), Some(3)], + vec![Some(0), None], + ] { + let dictionary = + DictionaryArray::try_new(UInt32Array::from(indices), dictionary_values.clone()) + .unwrap(); + assert!( + convert_bulk_part( + with_keys(Arc::new(dictionary)), + &metadata, + primary_key_codec.clone(), + output_schema.clone(), + true + ) + .is_err() + ); + } + + for invalid_keys in [ + Arc::new(arrow::array::StringArray::from(vec!["invalid"; 2])) as ArrayRef, + Arc::new(BinaryArray::from(vec![None::<&[u8]>; 2])), + Arc::new(BinaryArray::from_iter_values([b"truncated"; 2])), + ] { + let mut invalid = part.clone(); + let mut fields = invalid.batch.schema().fields().to_vec(); + fields[0] = Arc::new(Field::new( + PRIMARY_KEY_COLUMN_NAME, + invalid_keys.data_type().clone(), + true, + )); + let mut columns = invalid.batch.columns().to_vec(); + columns[0] = invalid_keys; + invalid.batch = RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap(); + assert!( + convert_bulk_part( + invalid, + &metadata, + primary_key_codec.clone(), + output_schema.clone(), + true, + ) + .is_err() + ); + } + let result = convert_bulk_part( part, &metadata, diff --git a/src/mito2/src/read/compat.rs b/src/mito2/src/read/compat.rs index 1d4a7f5457b..377b096db26 100644 --- a/src/mito2/src/read/compat.rs +++ b/src/mito2/src/read/compat.rs @@ -113,7 +113,7 @@ impl FlatCompatBatch { if actual.primary_key_encoding == PrimaryKeyEncoding::Sparse && compaction { // Special handling for sparse encoding in compaction. - return FlatCompatBatch::try_new_compact_sparse(mapper, actual); + return FlatCompatBatch::try_new_compact_sparse(mapper, actual, &actual_schema); } let (index_or_defaults, fields) = Self::compute_index_and_fields( @@ -225,33 +225,49 @@ impl FlatCompatBatch { fn try_new_compact_sparse( mapper: &FlatProjectionMapper, actual: &RegionMetadataRef, + actual_schema: &[(ColumnId, ConcreteDataType)], ) -> Result> { // Currently, we don't support converting sparse encoding back to dense encoding in // flat format. ensure!( - mapper.metadata().primary_key_encoding == PrimaryKeyEncoding::Sparse, + mapper.metadata().primary_key_encoding == actual.primary_key_encoding, UnsupportedOperationSnafu { err_msg: "Flat format doesn't support converting sparse encoding back to dense encoding" } ); - // For sparse encoding, we don't need to check the primary keys. - // Since this is for compaction, we always read all columns. - let actual_schema: Vec<_> = actual - .field_columns() - .chain([actual.time_index_column()]) - .map(|col| (col.column_id, col.column_schema.data_type.clone())) - .collect(); - let expect_schema: Vec<_> = mapper - .metadata() - .field_columns() - .chain([mapper.metadata().time_index_column()]) - .map(|col| (col.column_id, col.column_schema.data_type.clone())) - .collect(); + // Encoded sparse tags are not materialized during compaction. Preserve + // the projected JSON2 types: metadata only describes their logical shape. + let mut actual_fields = Vec::with_capacity(actual_schema.len()); + for (id, data_type) in actual_schema { + let column = actual + .column_by_id(*id) + .with_context(|| CompatReaderSnafu { + region_id: actual.region_id, + reason: format!("column {id} is absent from region metadata"), + })?; + if column.semantic_type != SemanticType::Tag { + actual_fields.push((*id, data_type.clone())); + } + } + let mut expect_fields = Vec::with_capacity(mapper.batch_schema().len()); + for (id, data_type) in mapper.batch_schema() { + let column = + mapper + .metadata() + .column_by_id(*id) + .with_context(|| CompatReaderSnafu { + region_id: mapper.metadata().region_id, + reason: format!("column {id} is absent from region metadata"), + })?; + if column.semantic_type != SemanticType::Tag { + expect_fields.push((*id, data_type.clone())); + } + } let (index_or_defaults, fields) = Self::compute_index_and_fields( - &actual_schema, - &expect_schema, + &actual_fields, + &expect_fields, mapper.metadata(), mapper.input_arrow_schema(true).fields(), )?; @@ -552,14 +568,12 @@ impl FlatCompatPrimaryKey { debug_assert!(!fields.is_empty()); // Create converter to append values. - let converter = Some(build_primary_key_codec_with_fields( - expect.primary_key_encoding, - fields.into_iter(), - )); + let converter = + build_primary_key_codec_with_fields(expect.primary_key_encoding, fields.into_iter()); Ok(Self { rewriter, - converter, + converter: Some(converter), values, }) } @@ -675,6 +689,7 @@ mod tests { use store_api::codec::PrimaryKeyEncoding; use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; use store_api::storage::RegionId; + use store_api::storage::consts::ReservedColumnId; use super::*; use crate::read::flat_projection::FlatProjectionMapper; @@ -683,6 +698,7 @@ mod tests { use crate::sst::parquet::{Json2RewriteTargets, Json2TargetLayout}; use crate::sst::{FlatSchemaOptions, to_flat_sst_arrow_schema}; use crate::test_util; + use crate::test_util::sst_util::sst_region_metadata_with_encoding; /// Creates a new [RegionMetadata]. fn new_metadata( @@ -691,7 +707,7 @@ mod tests { ) -> RegionMetadata { let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 1)); for (id, semantic_type, data_type) in semantic_types { - let column_schema = match semantic_type { + let mut column_schema = match semantic_type { SemanticType::Tag => { ColumnSchema::new(format!("tag_{id}"), data_type.clone(), true) } @@ -700,6 +716,11 @@ mod tests { } SemanticType::Timestamp => ColumnSchema::new("ts", data_type.clone(), false), }; + if data_type.is_json2() { + column_schema.with_extension_type( + &datatypes::extension::json::Json2ExtensionType::default(), + ); + } builder.push_column_metadata(ColumnMetadata { column_schema, @@ -711,6 +732,24 @@ mod tests { builder.build().unwrap() } + fn with_sparse_encoding(metadata: RegionMetadata) -> RegionMetadata { + let reserved = sst_region_metadata_with_encoding(PrimaryKeyEncoding::Sparse); + let primary_key = reserved + .primary_key + .into_iter() + .take(2) + .chain(metadata.primary_key.iter().copied()) + .collect(); + let mut builder = RegionMetadataBuilder::from_existing(metadata); + for column in reserved.column_metadatas.into_iter().take(2) { + builder.push_column_metadata(column); + } + builder + .primary_key(primary_key) + .primary_key_encoding(PrimaryKeyEncoding::Sparse); + builder.build().unwrap() + } + /// Encode primary key. fn encode_key(keys: &[Option<&str>]) -> Vec { let fields = (0..keys.len()) @@ -727,10 +766,7 @@ mod tests { /// Encode sparse primary key. fn encode_sparse_key(keys: &[(ColumnId, Option<&str>)]) -> Vec { - let fields = (0..keys.len()) - .map(|_| (1, SortField::new(ConcreteDataType::string_datatype()))) - .collect(); - let converter = SparsePrimaryKeyCodec::with_fields(fields); + let converter = SparsePrimaryKeyCodec::schemaless(); let row = keys .iter() .map(|(id, str_opt)| match str_opt { @@ -739,6 +775,7 @@ mod tests { }) .collect::>(); let mut buffer = vec![]; + converter.encode_internal(42, 7, &mut buffer).unwrap(); converter.encode_value_refs(&row, &mut buffer).unwrap(); buffer } @@ -839,10 +876,15 @@ mod tests { assert_eq!(expected_batch, result); } - #[test] - fn test_flat_compat_batch_uses_projected_json2_type() -> Result<()> { + #[rstest::rstest] + #[case(PrimaryKeyEncoding::Dense, false)] + #[case(PrimaryKeyEncoding::Sparse, true)] + fn test_flat_compat_batch_uses_projected_json2_type( + #[case] encoding: PrimaryKeyEncoding, + #[case] compaction: bool, + ) -> Result<()> { let json2 = ConcreteDataType::json2(JsonNativeType::object()); - let actual_metadata = Arc::new(new_metadata( + let mut actual_metadata = Arc::new(new_metadata( &[ ( 0, @@ -853,7 +895,7 @@ mod tests { ], &[], )); - let expected_metadata = Arc::new(new_metadata( + let mut expected_metadata = Arc::new(new_metadata( &[ ( 0, @@ -865,6 +907,10 @@ mod tests { ], &[], )); + if encoding == PrimaryKeyEncoding::Sparse { + actual_metadata = Arc::new(with_sparse_encoding((*actual_metadata).clone())); + expected_metadata = Arc::new(with_sparse_encoding((*expected_metadata).clone())); + } let read_columns = ReadColumns::new([0, 1, 2]) .with_json_target_types(BTreeMap::from([(1, JsonNativeType::Variant)])); let mapper = FlatProjectionMapper::new_with_read_columns( @@ -874,7 +920,7 @@ mod tests { )?; let read_format = FlatReadFormat::new(actual_metadata, read_columns, None, "test", false)?; - let compat = FlatCompatBatch::try_new(&mapper, &read_format, false)?.unwrap(); + let compat = FlatCompatBatch::try_new(&mapper, &read_format, compaction)?.unwrap(); let json_index = mapper .batch_schema() .iter() @@ -887,6 +933,10 @@ mod tests { .. } )); + assert_eq!( + compat.arrow_schema.field(json_index).data_type(), + &datatypes::arrow::datatypes::DataType::Binary + ); Ok(()) } @@ -995,6 +1045,7 @@ mod tests { ], &[1], ); + actual_metadata = with_sparse_encoding(actual_metadata); actual_metadata.primary_key_encoding = PrimaryKeyEncoding::Dense; let actual_metadata = Arc::new(actual_metadata); @@ -1011,13 +1062,20 @@ mod tests { ], &[1, 3], ); - expected_metadata.primary_key_encoding = PrimaryKeyEncoding::Sparse; + expected_metadata = with_sparse_encoding(expected_metadata); let expected_metadata = Arc::new(expected_metadata); let mapper = FlatProjectionMapper::all(&expected_metadata).unwrap(); let read_format = FlatReadFormat::new( actual_metadata.clone(), - ReadColumns::new([0, 1, 2, 3]), + ReadColumns::new([ + 0, + 1, + 2, + 3, + ReservedColumnId::table_id(), + ReservedColumnId::tsid(), + ]), None, "test", false, @@ -1034,8 +1092,20 @@ mod tests { tag1_builder.append_value("tag1"); let tag1_dict_array = Arc::new(tag1_builder.finish()); - let k1 = encode_key(&[Some("tag1")]); + let mut k1 = Vec::new(); + build_primary_key_codec(&actual_metadata) + .encode_values( + &[ + (ReservedColumnId::table_id(), Value::UInt32(42)), + (ReservedColumnId::tsid(), Value::UInt64(7)), + (1, Value::from("tag1")), + ], + &mut k1, + ) + .unwrap(); let input_columns: Vec = vec![ + Arc::new(datatypes::arrow::array::UInt32Array::from(vec![42, 42])), + Arc::new(UInt64Array::from(vec![7, 7])), tag1_dict_array.clone(), Arc::new(Int64Array::from(vec![100, 200])), Arc::new(TimestampMillisecondArray::from_iter_values([1000, 2000])), @@ -1057,6 +1127,8 @@ mod tests { null_tag_builder.append_nulls(2); let null_tag_dict_array = Arc::new(null_tag_builder.finish()); let expected_columns: Vec = vec![ + Arc::new(datatypes::arrow::array::UInt32Array::from(vec![42, 42])), + Arc::new(UInt64Array::from(vec![7, 7])), tag1_dict_array.clone(), null_tag_dict_array, Arc::new(Int64Array::from(vec![100, 200])), @@ -1075,6 +1147,52 @@ mod tests { assert_eq!(expected_batch, result); } + #[test] + fn test_compat_sparse_primary_key_with_new_nullable_tag() { + let actual = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let mut builder = RegionMetadataBuilder::from_existing((*actual).clone()); + let mut primary_key = actual.primary_key.clone(); + primary_key.push(4); + builder + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "new_tag", + ConcreteDataType::string_datatype(), + true, + ), + semantic_type: SemanticType::Tag, + column_id: 4, + }) + .primary_key(primary_key); + let expected = Arc::new(builder.build().unwrap()); + let compat = FlatCompatPrimaryKey::new(&expected, &actual).unwrap(); + + let values = vec![ + (ReservedColumnId::table_id(), Value::UInt32(23)), + (ReservedColumnId::tsid(), Value::UInt64(42)), + (0, Value::from("api")), + (1, Value::Null), + ]; + let mut key = Vec::new(); + build_primary_key_codec(&actual) + .encode_values(&values, &mut key) + .unwrap(); + + // Adding a NULL sparse tag must preserve the existing key, including metric IDs. + let dictionary_key = build_flat_test_pk_array(&[&key, &key]); + let binary_key: ArrayRef = Arc::new(BinaryArray::from(vec![ + Some(key.as_slice()), + None, + Some(key.as_slice()), + ])); + for input in [dictionary_key, binary_key] { + let result = compat.compat(&input).unwrap(); + assert_eq!(result.as_ref(), input.as_ref()); + } + } + #[test] fn test_compat_primary_key_with_different_encoding_only() { let mut actual_metadata = new_metadata( @@ -1089,6 +1207,7 @@ mod tests { ], &[1], ); + actual_metadata = with_sparse_encoding(actual_metadata); actual_metadata.primary_key_encoding = PrimaryKeyEncoding::Dense; let actual_metadata = Arc::new(actual_metadata); @@ -1098,8 +1217,14 @@ mod tests { let mapper = FlatProjectionMapper::all(&expected_metadata).unwrap(); let read_format = FlatReadFormat::new( - actual_metadata, - ReadColumns::new([0, 1, 2]), + actual_metadata.clone(), + ReadColumns::new([ + 0, + 1, + 2, + ReservedColumnId::table_id(), + ReservedColumnId::tsid(), + ]), None, "test", false, @@ -1109,7 +1234,17 @@ mod tests { .unwrap() .unwrap(); - let dense_key = encode_key(&[Some("tag1")]); + let mut dense_key = Vec::new(); + build_primary_key_codec(&actual_metadata) + .encode_values( + &[ + (ReservedColumnId::table_id(), Value::UInt32(42)), + (ReservedColumnId::tsid(), Value::UInt64(7)), + (1, Value::from("tag1")), + ], + &mut dense_key, + ) + .unwrap(); let sparse_key = encode_sparse_key(&[(1, Some("tag1"))]); let dictionary_key = build_flat_test_pk_array(&[&dense_key, &dense_key]); @@ -1344,19 +1479,8 @@ mod tests { ], &[], ); - let reserved = crate::test_util::sst_util::sst_region_metadata_with_encoding( - PrimaryKeyEncoding::Sparse, - ); - actual_metadata.primary_key = reserved.primary_key.into_iter().take(2).collect(); - actual_metadata - .column_metadatas - .extend(reserved.column_metadatas.into_iter().take(2)); - actual_metadata.primary_key_encoding = PrimaryKeyEncoding::Sparse; - let actual_metadata = Arc::new( - RegionMetadataBuilder::from_existing(actual_metadata) - .build() - .unwrap(), - ); + actual_metadata = with_sparse_encoding(actual_metadata); + let actual_metadata = Arc::new(actual_metadata); let mut expected_metadata = (*actual_metadata).clone(); for column in expected_metadata.column_metadatas.iter_mut() { @@ -1440,7 +1564,7 @@ mod tests { ], &[], ); - actual_metadata.primary_key_encoding = PrimaryKeyEncoding::Sparse; + actual_metadata = with_sparse_encoding(actual_metadata); let actual_metadata = Arc::new(actual_metadata); let mut expected_metadata = new_metadata( @@ -1455,7 +1579,7 @@ mod tests { ], &[], ); - expected_metadata.primary_key_encoding = PrimaryKeyEncoding::Sparse; + expected_metadata = with_sparse_encoding(expected_metadata); let expected_metadata = Arc::new(expected_metadata); let mapper = FlatProjectionMapper::all(&expected_metadata).unwrap(); @@ -1483,8 +1607,10 @@ mod tests { OpType::Put as u8, ])), ]; - let input_schema = - to_flat_sst_arrow_schema(&actual_metadata, &FlatSchemaOptions::default()); + let input_schema = to_flat_sst_arrow_schema( + &actual_metadata, + &FlatSchemaOptions::from_encoding(PrimaryKeyEncoding::Sparse), + ); let input_batch = RecordBatch::try_new(input_schema, input_columns).unwrap(); let result = compat_batch.compat(input_batch).unwrap(); @@ -1500,8 +1626,10 @@ mod tests { OpType::Put as u8, ])), ]; - let output_schema = - to_flat_sst_arrow_schema(&expected_metadata, &FlatSchemaOptions::default()); + let output_schema = to_flat_sst_arrow_schema( + &expected_metadata, + &FlatSchemaOptions::from_encoding(PrimaryKeyEncoding::Sparse), + ); let expected_batch = RecordBatch::try_new(output_schema, expected_columns).unwrap(); assert_eq!(expected_batch, result); @@ -1561,12 +1689,18 @@ mod tests { } else { metadata.as_ref().clone() }; - let mut actual_builder = RegionMetadataBuilder::from_existing(source); - actual_builder.primary_key_encoding(encoding); - let mut expected_builder = - RegionMetadataBuilder::from_existing(metadata.as_ref().clone()); - expected_builder.primary_key_encoding(encoding); - expected_builder.push_column_metadata(ColumnMetadata { + let actual = Arc::new(if encoding == PrimaryKeyEncoding::Sparse { + with_sparse_encoding(source) + } else { + source + }); + let expected = if encoding == PrimaryKeyEncoding::Sparse { + with_sparse_encoding(metadata.as_ref().clone()) + } else { + metadata.as_ref().clone() + }; + let mut builder = RegionMetadataBuilder::from_existing(expected); + builder.push_column_metadata(ColumnMetadata { column_schema: ColumnSchema::new( "added", ConcreteDataType::int64_datatype(), @@ -1578,21 +1712,11 @@ mod tests { semantic_type: SemanticType::Field, column_id: 2, }); - if encoding == PrimaryKeyEncoding::Sparse { - let reserved = test_util::sst_util::sst_region_metadata_with_encoding(encoding); - for builder in [&mut actual_builder, &mut expected_builder] { - builder.primary_key(reserved.primary_key[..2].to_vec()); - for column in reserved.column_metadatas.iter().take(2) { - builder.push_column_metadata(column.clone()); - } - } - } - let actual = Arc::new(actual_builder.build()?); - let expected = Arc::new(expected_builder.build()?); + let expected = Arc::new(builder.build()?); let read_columns = ReadColumns::new([0, 1, 2]); let mapper = FlatProjectionMapper::new_with_json2_rewrite_targets( &expected, - vec![0, 1, 2], + vec![0, 1, expected.column_index_by_id(2).unwrap()], read_columns.clone(), &targets, )?; diff --git a/src/mito2/src/region.rs b/src/mito2/src/region.rs index 923f223c8c7..819f32cf7e8 100644 --- a/src/mito2/src/region.rs +++ b/src/mito2/src/region.rs @@ -2282,6 +2282,7 @@ mod tests { manifest_cache: None, }, FormatType::PrimaryKey, + false, &Default::default(), ) .await @@ -2835,6 +2836,7 @@ mod tests { manifest_cache: None, }, FormatType::PrimaryKey, + false, &Default::default(), ) .await @@ -2903,6 +2905,7 @@ mod tests { manifest_cache: None, }, FormatType::PrimaryKey, + false, &Default::default(), ) .await diff --git a/src/mito2/src/region/opener.rs b/src/mito2/src/region/opener.rs index d10545198f7..ba4fe615951 100644 --- a/src/mito2/src/region/opener.rs +++ b/src/mito2/src/region/opener.rs @@ -403,6 +403,7 @@ impl RegionOpener { flushed_entry_id, region_manifest_options, sst_format, + options.append_mode, &self.stats, ) .await?; @@ -795,6 +796,20 @@ pub(crate) fn version_builder_from_manifest( /// Updates region options by persistent options. pub(crate) fn sanitize_region_options(manifest: &RegionManifest, options: &mut RegionOptions) { + if let Some(encoding) = options.primary_key_encoding + && encoding != manifest.metadata.primary_key_encoding + { + common_telemetry::warn!( + "Overriding primary key encoding from {:?} to {:?} for region {}", + encoding, + manifest.metadata.primary_key_encoding, + manifest.metadata.region_id, + ); + } + options.primary_key_encoding = Some(manifest.metadata.primary_key_encoding); + if manifest.metadata.primary_key_encoding == PrimaryKeyEncoding::Sparse { + options.sst_format = Some(FormatType::Flat); + } // Caller options take precedence over the manifest, except JSON2 always requires flat output. let format = options.sst_format.unwrap_or(manifest.sst_format); // Older versions allowed JSON2 with primary-key output; recover without rewriting its history. @@ -1009,9 +1024,7 @@ where // columns like the write path, otherwise the memtable rejects the // batch. Skips sparse batches as they don't carry tag columns by // design. - if region_metadata.primary_key_encoding != PrimaryKeyEncoding::Sparse { - part.fill_missing_columns(®ion_metadata)?; - } + part.fill_missing_columns(®ion_metadata)?; rows_replayed += part.num_rows(); // During replay, we should adopt the sequence from WAL. let bulk_sequence_from_wal = part.sequence; @@ -1447,6 +1460,7 @@ fn can_load_cache(state: RegionRoleState) -> bool { mod tests { use std::collections::HashMap; use std::sync::Arc; + use store_api::codec::PrimaryKeyEncoding; use arrow_schema::extension::ExtensionType; use common_base::readable_size::ReadableSize; @@ -1627,6 +1641,22 @@ mod tests { )); } + #[test] + fn test_sparse_reopen_uses_persisted_encoding_and_flat_sst() { + for supplied in [None, Some(PrimaryKeyEncoding::Dense)] { + let mut manifest = build_test_manifest(FormatType::PrimaryKey); + Arc::make_mut(&mut manifest.metadata).primary_key_encoding = PrimaryKeyEncoding::Sparse; + let mut options = RegionOptions { + primary_key_encoding: supplied, + sst_format: Some(FormatType::PrimaryKey), + ..Default::default() + }; + sanitize_region_options(&manifest, &mut options); + assert_eq!(options.primary_key_encoding(), PrimaryKeyEncoding::Sparse); + assert_eq!(options.sst_format, Some(FormatType::Flat)); + } + } + #[test] fn test_sanitize_region_options_options_format_wins() { // Manifest persisted PrimaryKey, but the re-parsed options now request Flat diff --git a/src/mito2/src/region/options.rs b/src/mito2/src/region/options.rs index 58fdec15ad5..85cb79928cd 100644 --- a/src/mito2/src/region/options.rs +++ b/src/mito2/src/region/options.rs @@ -120,7 +120,7 @@ pub struct RegionOptions { /// Max number of rows in a parquet row group. Uses [DEFAULT_ROW_GROUP_SIZE] if `None`. #[serde(skip_serializing_if = "Option::is_none")] pub max_row_group_row_count: Option, - /// Internal primary key encoding override used by metric-engine. + /// Primary key encoding. Sparse regions require reserved table ID and TSID columns. #[serde(skip_serializing_if = "Option::is_none")] pub primary_key_encoding: Option, /// Per-region write buffer size. A positive size flushes/stalls this region @@ -313,6 +313,19 @@ impl RegionOptions { None }; + let primary_key_encoding = options_map + .get(PRIMARY_KEY_ENCODING) + .or_else(|| options_map.get(MEMTABLE_PARTITION_TREE_PRIMARY_KEY_ENCODING)) + .map(|v| match v.to_lowercase().as_str() { + "dense" => Ok(PrimaryKeyEncoding::Dense), + "sparse" => Ok(PrimaryKeyEncoding::Sparse), + _ => Err(InvalidRegionOptionsSnafu { + reason: format!("Invalid primary key encoding: {v}"), + } + .build()), + }) + .transpose()?; + // The partition tree memtable has been removed. Besides falling back to // the default memtable, also override the SST format to flat. let mut sst_format = options.sst_format; @@ -327,7 +340,9 @@ impl RegionOptions { // Bulk memtable produces flat-encoded ranges and flushes them through // `put_sst()`, so the SST format must be flat to match. - if matches!(memtable, Some(MemtableOptions::Bulk(_))) { + if matches!(memtable, Some(MemtableOptions::Bulk(_))) + || primary_key_encoding == Some(PrimaryKeyEncoding::Sparse) + { if let Some(format) = sst_format && format != FormatType::Flat { @@ -346,18 +361,6 @@ impl RegionOptions { .map(|v| matches!(v.to_lowercase().as_str(), "true" | "1")) .unwrap_or(false); let compaction_override = has_compaction_type || compaction_override_flag; - let primary_key_encoding = options_map - .get(PRIMARY_KEY_ENCODING) - .or_else(|| options_map.get(MEMTABLE_PARTITION_TREE_PRIMARY_KEY_ENCODING)) - .map(|v| match v.to_lowercase().as_str() { - "dense" => Ok(PrimaryKeyEncoding::Dense), - "sparse" => Ok(PrimaryKeyEncoding::Sparse), - _ => Err(InvalidRegionOptionsSnafu { - reason: format!("Invalid primary key encoding: {v}"), - } - .build()), - }) - .transpose()?; let opts = RegionOptions { ttl: options.ttl, @@ -1029,6 +1032,19 @@ mod tests { assert_eq!(expect, options); } + #[test] + fn test_sparse_forces_flat_sst() { + for format in [None, Some("primary_key"), Some("flat")] { + let mut map = make_map(&[("primary_key_encoding", "sparse")]); + if let Some(format) = format { + map.insert("sst_format".into(), format.into()); + } + let options = RegionOptions::try_from_options(RegionId::new(0, 0), &map).unwrap(); + assert_eq!(options.primary_key_encoding(), PrimaryKeyEncoding::Sparse); + assert_eq!(options.sst_format, Some(FormatType::Flat)); + } + } + #[test] fn test_primary_key_encoding() { // New top-level key. diff --git a/src/mito2/src/request.rs b/src/mito2/src/request.rs index d7e12a5da60..9c0178710f1 100644 --- a/src/mito2/src/request.rs +++ b/src/mito2/src/request.rs @@ -526,6 +526,7 @@ impl RequestColumn for RowColumn<'_> { /// A column of the batch of a bulk insert. struct BulkColumn<'a> { array: &'a ArrayRef, + field: &'a datatypes::arrow::datatypes::Field, } impl RequestColumn for BulkColumn<'_> { @@ -539,6 +540,11 @@ impl RequestColumn for BulkColumn<'_> { actual => actual.is_string(), }; } + if data_type.is_json2() && datatypes::extension::json::is_json2_extension_type(self.field) { + return matches!(actual, ArrowDataType::Struct(_) | ArrowDataType::Null) + && datatypes::extension::json::json2_remainder_field(self.field).is_ok() + && datatypes::types::json_type::JsonNativeType::try_from(actual).is_ok(); + } *actual == data_type.as_arrow_type() } @@ -559,7 +565,7 @@ impl RequestColumn for BulkColumn<'_> { /// [WriteRequest::check_schema] applies to rows: column names are unique, every /// column is a column of the region with its Arrow type and contains nulls only /// if the region column is nullable, and a missing column is nullable or has a -/// default value. Tags may use any string encoding. Sparse batches carry the +/// default value. Tags may use any string encoding. Sparse batches may carry the /// tags in one encoded primary key column instead. Returns whether the batch /// misses columns that the writer fills. pub fn check_bulk_insert_batch(metadata: &RegionMetadata, batch: &RecordBatch) -> Result { @@ -572,7 +578,7 @@ pub fn check_bulk_insert_batch(metadata: &RegionMetadata, batch: &RecordBatch) - primary_key.replace(array).is_some() } else { columns - .insert(field.name().as_str(), BulkColumn { array }) + .insert(field.name().as_str(), BulkColumn { array, field }) .is_some() }; ensure!( @@ -583,11 +589,7 @@ pub fn check_bulk_insert_batch(metadata: &RegionMetadata, batch: &RecordBatch) - } ); } - if sparse { - let primary_key = primary_key.with_context(|| InvalidRequestSnafu { - region_id, - reason: format!("missing column {PRIMARY_KEY_COLUMN_NAME}"), - })?; + if let Some(primary_key) = primary_key { // The writer casts the encoded primary key to a dictionary of binary. let type_matches = match primary_key.data_type() { ArrowDataType::Binary => true, @@ -609,7 +611,7 @@ pub fn check_bulk_insert_batch(metadata: &RegionMetadata, batch: &RecordBatch) - } let mut needs_fill = false; check_request_columns(region_id, metadata, columns, |column| { - if sparse && column.semantic_type == SemanticType::Tag { + if primary_key.is_some() && column.semantic_type == SemanticType::Tag { Ok(()) } else { needs_fill = true; diff --git a/src/mito2/src/sst/primary_key.rs b/src/mito2/src/sst/primary_key.rs index 7752a652f80..2d3f8c90e3b 100644 --- a/src/mito2/src/sst/primary_key.rs +++ b/src/mito2/src/sst/primary_key.rs @@ -654,12 +654,27 @@ mod tests { #[test] fn test_sparse_ranges_keep_their_original_encoding() { use mito_codec::row_converter::SparsePrimaryKeyCodec; - let mut metadata = (*metadata(&[Value::from(""), Value::Null])).clone(); - metadata.primary_key_encoding = PrimaryKeyEncoding::Sparse; - let metadata = Arc::new(metadata); + let metadata = Arc::new( + crate::test_util::sst_util::sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + ), + ); let mut bytes = Vec::new(); SparsePrimaryKeyCodec::new(&metadata) - .encode_values(&[(0, Value::from("a"))], &mut bytes) + .encode_values( + &[ + ( + store_api::storage::consts::ReservedColumnId::table_id(), + Value::UInt32(42), + ), + ( + store_api::storage::consts::ReservedColumnId::tsid(), + Value::UInt64(7), + ), + (0, Value::from("a")), + ], + &mut bytes, + ) .unwrap(); let key = Bytes::from(bytes); let mapper = PrimaryKeyRangeMapper::new(metadata); diff --git a/src/mito2/src/test_util.rs b/src/mito2/src/test_util.rs index f265494058a..1d742505df4 100644 --- a/src/mito2/src/test_util.rs +++ b/src/mito2/src/test_util.rs @@ -841,6 +841,7 @@ impl TestEnv { 0, manifest_opts, FormatType::PrimaryKey, + false, &Default::default(), ) .await diff --git a/src/mito2/src/test_util/scheduler_util.rs b/src/mito2/src/test_util/scheduler_util.rs index f97886b4ea2..3185a9895a3 100644 --- a/src/mito2/src/test_util/scheduler_util.rs +++ b/src/mito2/src/test_util/scheduler_util.rs @@ -147,6 +147,7 @@ impl SchedulerEnv { manifest_cache: None, }, FormatType::PrimaryKey, + false, &Default::default(), ) .await diff --git a/src/store-api/src/mito_engine_options.rs b/src/store-api/src/mito_engine_options.rs index 1f680aff995..6a0d2b13d02 100644 --- a/src/store-api/src/mito_engine_options.rs +++ b/src/store-api/src/mito_engine_options.rs @@ -182,7 +182,7 @@ pub fn is_mito_engine_option_key(key: &str) -> bool { MEMTABLE_PARTITION_TREE_INDEX_MAX_KEYS_PER_SHARD, MEMTABLE_PARTITION_TREE_DATA_FREEZE_THRESHOLD, MEMTABLE_PARTITION_TREE_FORK_DICTIONARY_BYTES, - // We don't allow to create a mito table with sparse primary key encoding directly. + crate::metric_engine_consts::PRIMARY_KEY_ENCODING, APPEND_MODE_KEY, MERGE_MODE_KEY, SST_FORMAT_KEY, diff --git a/src/store-api/src/region_engine.rs b/src/store-api/src/region_engine.rs index 2a2bc7b7f1f..fb3c7f01b90 100644 --- a/src/store-api/src/region_engine.rs +++ b/src/store-api/src/region_engine.rs @@ -318,7 +318,8 @@ pub struct ScannerProperties { /// The target partitions of the scanner. 0 indicates using the number of partitions as target partitions. target_partitions: usize, - /// Whether the scanner is scanning a logical region. + /// Whether the scanner reads one logical metric region, including a physical + /// metric region constrained by a table-ID filter. logical_region: bool, /// Region id that should receive query-load metrics for this scanner. @@ -392,7 +393,8 @@ impl ScannerProperties { self.total_rows_is_exact } - /// Returns whether the scanner is scanning a logical region. + /// Returns whether the scanner reads one logical metric region, including a + /// physical metric region constrained by a table-ID filter. pub fn is_logical_region(&self) -> bool { self.logical_region } @@ -406,7 +408,7 @@ impl ScannerProperties { } } - /// Sets whether the scanner is reading a logical region. + /// Sets whether the scanner reads one logical metric region. pub fn set_logical_region(&mut self, logical_region: bool) { self.logical_region = logical_region; } @@ -535,7 +537,8 @@ pub trait RegionScanner: Debug + DisplayAs + Send { /// The default is a no-op and does not make inherently one-shot sources replayable. fn reset_state(&mut self) {} - /// Sets whether the scanner is reading a logical region. + /// Sets whether the scanner reads one logical metric region, including a + /// physical metric region constrained by a table-ID filter. fn set_logical_region(&mut self, logical_region: bool); /// Sets the region id that should receive query-load metrics. diff --git a/src/table/src/table/scan.rs b/src/table/src/table/scan.rs index 2dd501e678f..80878fae5bc 100644 --- a/src/table/src/table/scan.rs +++ b/src/table/src/table/scan.rs @@ -185,16 +185,14 @@ impl RegionScanExec { } orderings.push(pk_time_ordering); - // Metric engine physical tables keep tag columns alongside `__tsid`. When `__tsid` - // is present, the scan output is also ordered by `(__tsid, time index)` which can - // help eliminate redundant sorts for promql plans (e.g. SeriesDivide). - // - // `pk_time_ordering` and the potential `(__tsid, time index)` ordering are actually - // the same thing, thus present in one `OrderingEquivalenceClass`. - if let (Ok(tsid_col), Some(ts)) = ( - Column::new_with_schema(DATA_SCHEMA_TSID_COLUMN_NAME, &arrow_schema), - ts_col.clone(), - ) { + // Only logical metric scans constrain the physical key to one source table. + // A shared table can have the same TSID under different table IDs. + if scanner_props.is_logical_region() + && let (Ok(tsid_col), Some(ts)) = ( + Column::new_with_schema(DATA_SCHEMA_TSID_COLUMN_NAME, &arrow_schema), + ts_col.clone(), + ) + { orderings.push(vec![ PhysicalSortExpr::new( Arc::new(tsid_col) as _, @@ -225,13 +223,13 @@ impl RegionScanExec { let partitioning = match request.distribution { Some(TimeSeriesDistribution::PerSeries) => { match Column::new_with_schema(DATA_SCHEMA_TSID_COLUMN_NAME, &arrow_schema) { - Ok(tsid_col) => { + Ok(tsid_col) if scanner_props.is_logical_region() => { Partitioning::Hash(vec![Arc::new(tsid_col) as _], num_output_partition) } - Err(_) if pk_columns.len() == pk_names.len() => { + _ if pk_columns.len() == pk_names.len() => { Partitioning::Hash(pk_columns.clone(), num_output_partition) } - Err(_) => Partitioning::UnknownPartitioning(num_output_partition), + _ => Partitioning::UnknownPartitioning(num_output_partition), } } Some(TimeSeriesDistribution::TimeWindowed) | None => { @@ -1146,10 +1144,11 @@ mod test { ]); let metadata = Arc::new(builder.build().unwrap()); - for (projection, partitioned_by_tsid, claimed, unclaimed) in [ + for (projection, logical_region, partitioned_by_tsid, claimed, unclaimed) in [ ( vec![DATA_SCHEMA_TABLE_ID_COLUMN_NAME, "A", "ts"], false, + false, vec![vec![DATA_SCHEMA_TABLE_ID_COLUMN_NAME]], vec![vec!["A"], vec![DATA_SCHEMA_TABLE_ID_COLUMN_NAME, "ts"]], ), @@ -1157,6 +1156,7 @@ mod test { // row key and neither `A` nor `ts` is ordered. ( vec!["A", "marker", "ts"], + true, false, vec![], vec![vec!["A"], vec!["marker"], vec!["ts"], vec!["A", "ts"]], @@ -1165,12 +1165,24 @@ mod test { ( vec![DATA_SCHEMA_TSID_COLUMN_NAME, "A", "marker", "ts"], true, + true, vec![ vec![DATA_SCHEMA_TSID_COLUMN_NAME], vec![DATA_SCHEMA_TSID_COLUMN_NAME, "ts"], ], vec![vec!["A"], vec!["ts"], vec!["A", "ts"]], ), + // Shared exemplar tables project TSID too, but span multiple source tables. + ( + vec![DATA_SCHEMA_TSID_COLUMN_NAME, "A", "marker", "ts"], + false, + false, + vec![], + vec![ + vec![DATA_SCHEMA_TSID_COLUMN_NAME], + vec![DATA_SCHEMA_TSID_COLUMN_NAME, "ts"], + ], + ), ] { let schema = Arc::new(Schema::new( projection @@ -1179,7 +1191,7 @@ mod test { .collect(), )); let mut properties = ScannerProperties::new(vec![vec![], vec![]], false, 0); - properties.set_logical_region(true); + properties.set_logical_region(logical_region); let scanner = Box::new(RepeatableScanner { batches: RecordBatches::try_new(schema, vec![]).unwrap(), properties, diff --git a/tests-integration/src/tests/instance_test.rs b/tests-integration/src/tests/instance_test.rs index 3b3dca0bbd0..4f6cd5a3fab 100644 --- a/tests-integration/src/tests/instance_test.rs +++ b/tests-integration/src/tests/instance_test.rs @@ -43,6 +43,43 @@ use crate::tests::test_util::{ standalone_instance_case, standalone_with_multiple_object_stores, }; +#[apply(both_instances_cases)] +async fn test_sparse_primary_key_schema_validation(instance: Arc) { + let instance = instance.frontend(); + for columns in [ + "host STRING PRIMARY KEY, ts TIMESTAMP TIME INDEX, val DOUBLE", + "host BIGINT PRIMARY KEY, ts TIMESTAMP TIME INDEX, val DOUBLE", + "ts TIMESTAMP TIME INDEX, val DOUBLE", + "__table_id INT UNSIGNED NOT NULL, __tsid BIGINT UNSIGNED NOT NULL, ts TIMESTAMP TIME INDEX, PRIMARY KEY (__table_id, __tsid)", + ] { + let error = try_execute_sql( + &instance, + &format!("CREATE TABLE sparse_tags ({columns}) WITH ('primary_key_encoding'='sparse')"), + ) + .await + .unwrap_err(); + assert!( + error + .output_msg() + .contains("sparse primary keys must start with (__table_id, __tsid)"), + "{error:?}" + ); + } + + // Reusing the name verifies that rejected creation did not persist the invalid schema. + execute_sql( + &instance, + "CREATE TABLE sparse_tags (host STRING PRIMARY KEY, ts TIMESTAMP TIME INDEX, val DOUBLE)", + ) + .await; + execute_sql(&instance, "INSERT INTO sparse_tags VALUES ('a', 1000, 2.0)").await; + let output = execute_sql(&instance, "SELECT COUNT(*) FROM sparse_tags").await; + assert_eq!( + output.data.pretty_print().await, + "+----------+\n| count(*) |\n+----------+\n| 1 |\n+----------+" + ); +} + #[apply(both_instances_cases)] async fn test_create_database_and_insert_query(instance: Arc) { let instance = instance.frontend(); diff --git a/tests/cases/standalone/common/types/json/exemplar_sparse.result b/tests/cases/standalone/common/types/json/exemplar_sparse.result new file mode 100644 index 00000000000..6836e12813f --- /dev/null +++ b/tests/cases/standalone/common/types/json/exemplar_sparse.result @@ -0,0 +1,47 @@ +-- A shared table's TSID does not imply ordering across source table IDs. +CREATE TABLE exemplar_tsid_order ( + __table_id INT UNSIGNED, + __tsid BIGINT UNSIGNED, + greptime_metric STRING, + ts TIMESTAMP TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY (__table_id, __tsid, greptime_metric) +) WITH ('append_mode'='true'); + +Affected Rows: 0 + +INSERT INTO exemplar_tsid_order VALUES + (1, 7, 'cpu', 20, 1), (2, 7, 'memory', 10, 2); + +Affected Rows: 2 + +SELECT __tsid, ts, greptime_value FROM exemplar_tsid_order ORDER BY __tsid, ts; + ++--------+-------------------------+----------------+ +| __tsid | ts | greptime_value | ++--------+-------------------------+----------------+ +| 7 | 1970-01-01T00:00:00.010 | 2.0 | +| 7 | 1970-01-01T00:00:00.020 | 1.0 | ++--------+-------------------------+----------------+ + +ADMIN FLUSH_TABLE('exemplar_tsid_order'); + ++------------------------------------------+ +| ADMIN FLUSH_TABLE('exemplar_tsid_order') | ++------------------------------------------+ +| 0 | ++------------------------------------------+ + +SELECT __tsid, ts, greptime_value FROM exemplar_tsid_order ORDER BY __tsid, ts; + ++--------+-------------------------+----------------+ +| __tsid | ts | greptime_value | ++--------+-------------------------+----------------+ +| 7 | 1970-01-01T00:00:00.010 | 2.0 | +| 7 | 1970-01-01T00:00:00.020 | 1.0 | ++--------+-------------------------+----------------+ + +DROP TABLE exemplar_tsid_order; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/types/json/exemplar_sparse.sql b/tests/cases/standalone/common/types/json/exemplar_sparse.sql new file mode 100644 index 00000000000..20958fc20e3 --- /dev/null +++ b/tests/cases/standalone/common/types/json/exemplar_sparse.sql @@ -0,0 +1,15 @@ +-- A shared table's TSID does not imply ordering across source table IDs. +CREATE TABLE exemplar_tsid_order ( + __table_id INT UNSIGNED, + __tsid BIGINT UNSIGNED, + greptime_metric STRING, + ts TIMESTAMP TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY (__table_id, __tsid, greptime_metric) +) WITH ('append_mode'='true'); +INSERT INTO exemplar_tsid_order VALUES + (1, 7, 'cpu', 20, 1), (2, 7, 'memory', 10, 2); +SELECT __tsid, ts, greptime_value FROM exemplar_tsid_order ORDER BY __tsid, ts; +ADMIN FLUSH_TABLE('exemplar_tsid_order'); +SELECT __tsid, ts, greptime_value FROM exemplar_tsid_order ORDER BY __tsid, ts; +DROP TABLE exemplar_tsid_order; diff --git a/tests/compatibility/cases/exemplar_legacy_keys/case.toml b/tests/compatibility/cases/exemplar_legacy_keys/case.toml new file mode 100644 index 00000000000..89b629d1766 --- /dev/null +++ b/tests/compatibility/cases/exemplar_legacy_keys/case.toml @@ -0,0 +1,8 @@ +name = "exemplar_legacy_keys" +reason = "Reserved sparse key validation must preserve existing dense and metric-engine sparse data." +introduced_by = "persistent_metric_exemplars" +topologies = ["standalone", "distributed"] +from_range = [">=v1.2.1"] +to_range = [">=v1.4.0-alpha.0"] +features = ["metric", "sparse", "table", "read"] +owner = "mito2" diff --git a/tests/compatibility/cases/exemplar_legacy_keys/setup.sql b/tests/compatibility/cases/exemplar_legacy_keys/setup.sql new file mode 100644 index 00000000000..c063ecf04a5 --- /dev/null +++ b/tests/compatibility/cases/exemplar_legacy_keys/setup.sql @@ -0,0 +1,17 @@ +CREATE TABLE dense_keys (ts TIMESTAMP TIME INDEX, host STRING PRIMARY KEY, val DOUBLE); +INSERT INTO dense_keys (ts, host, val) VALUES (1, 'api', 1.5), (2, NULL, 2.5); +ADMIN FLUSH_TABLE('dense_keys'); +CREATE TABLE metric_keys (ts TIMESTAMP TIME INDEX, val DOUBLE) ENGINE=metric +WITH ('physical_metric_table'='', 'memtable.type'='bulk', 'primary_key_encoding'='sparse'); +CREATE TABLE logical_keys (ts TIMESTAMP TIME INDEX, host STRING PRIMARY KEY, val DOUBLE) ENGINE=metric +WITH ('on_physical_table'='metric_keys'); +INSERT INTO logical_keys (ts, host, val) VALUES (1, 'api', 3.5), (2, NULL, 4.5); +ADMIN FLUSH_TABLE('metric_keys'); +-- Historical OTLP infinity labels must retain their identity after upgrade. +CREATE TABLE otlp_upgrade_bucket ( + greptime_timestamp TIMESTAMP(3) TIME INDEX, + le STRING PRIMARY KEY, + greptime_value DOUBLE +); +INSERT INTO otlp_upgrade_bucket VALUES (10000, 'inf', 1), (20000, 'inf', 2); +ADMIN FLUSH_TABLE('otlp_upgrade_bucket'); diff --git a/tests/compatibility/cases/exemplar_legacy_keys/verify.result b/tests/compatibility/cases/exemplar_legacy_keys/verify.result new file mode 100644 index 00000000000..59d7b04ccac --- /dev/null +++ b/tests/compatibility/cases/exemplar_legacy_keys/verify.result @@ -0,0 +1,85 @@ +SELECT ts, host, val FROM dense_keys ORDER BY ts; + ++-------------------------+------+-----+ +| ts | host | val | ++-------------------------+------+-----+ +| 1970-01-01T00:00:00.001 | api | 1.5 | +| 1970-01-01T00:00:00.002 | | 2.5 | ++-------------------------+------+-----+ + +SELECT ts, host, val FROM logical_keys ORDER BY ts; + ++-------------------------+------+-----+ +| ts | host | val | ++-------------------------+------+-----+ +| 1970-01-01T00:00:00.001 | api | 3.5 | +| 1970-01-01T00:00:00.002 | | 4.5 | ++-------------------------+------+-----+ + +INSERT INTO dense_keys (ts, host, val) VALUES (3, 'new', 5.5); + +Affected Rows: 1 + +INSERT INTO logical_keys (ts, host, val) VALUES (3, 'new', 6.5); + +Affected Rows: 1 + +ADMIN FLUSH_TABLE('dense_keys'); + ++---------------------------------+ +| ADMIN FLUSH_TABLE('dense_keys') | ++---------------------------------+ +| 0 | ++---------------------------------+ + +ADMIN FLUSH_TABLE('metric_keys'); + ++----------------------------------+ +| ADMIN FLUSH_TABLE('metric_keys') | ++----------------------------------+ +| 0 | ++----------------------------------+ + +SELECT COUNT(*) FROM dense_keys; + ++----------+ +| count(*) | ++----------+ +| 3 | ++----------+ + +SELECT COUNT(*) FROM logical_keys; + ++----------+ +| count(*) | ++----------+ +| 3 | ++----------+ + +-- New OTLP writes retain the historical infinity label across upgrade. +INSERT INTO otlp_upgrade_bucket VALUES (30000, 'inf', 3), (40000, 'inf', 4); + +Affected Rows: 2 + +SELECT greptime_timestamp, le, greptime_value FROM otlp_upgrade_bucket ORDER BY greptime_timestamp; + ++---------------------+-----+----------------+ +| greptime_timestamp | le | greptime_value | ++---------------------+-----+----------------+ +| 1970-01-01T00:00:10 | inf | 1.0 | +| 1970-01-01T00:00:20 | inf | 2.0 | +| 1970-01-01T00:00:30 | inf | 3.0 | +| 1970-01-01T00:00:40 | inf | 4.0 | ++---------------------+-----+----------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (10, 40, '10s') otlp_upgrade_bucket{le="inf"}; + ++---------------------+-----+----------------+ +| greptime_timestamp | le | greptime_value | ++---------------------+-----+----------------+ +| 1970-01-01T00:00:10 | inf | 1.0 | +| 1970-01-01T00:00:20 | inf | 2.0 | +| 1970-01-01T00:00:30 | inf | 3.0 | +| 1970-01-01T00:00:40 | inf | 4.0 | ++---------------------+-----+----------------+ diff --git a/tests/compatibility/cases/exemplar_legacy_keys/verify.sql b/tests/compatibility/cases/exemplar_legacy_keys/verify.sql new file mode 100644 index 00000000000..115119d9686 --- /dev/null +++ b/tests/compatibility/cases/exemplar_legacy_keys/verify.sql @@ -0,0 +1,13 @@ +SELECT ts, host, val FROM dense_keys ORDER BY ts; +SELECT ts, host, val FROM logical_keys ORDER BY ts; +INSERT INTO dense_keys (ts, host, val) VALUES (3, 'new', 5.5); +INSERT INTO logical_keys (ts, host, val) VALUES (3, 'new', 6.5); +ADMIN FLUSH_TABLE('dense_keys'); +ADMIN FLUSH_TABLE('metric_keys'); +SELECT COUNT(*) FROM dense_keys; +SELECT COUNT(*) FROM logical_keys; +-- New OTLP writes retain the historical infinity label across upgrade. +INSERT INTO otlp_upgrade_bucket VALUES (30000, 'inf', 3), (40000, 'inf', 4); +SELECT greptime_timestamp, le, greptime_value FROM otlp_upgrade_bucket ORDER BY greptime_timestamp; +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (10, 40, '10s') otlp_upgrade_bucket{le="inf"}; From 7d42feb0236955eb4a1608858033cb40a16dfd8d Mon Sep 17 00:00:00 2001 From: LFC <990479+MichaelScofield@users.noreply.github.com> Date: Thu, 29 Oct 2026 16:24:50 +0800 Subject: [PATCH 136/239] chore(ci): update compatibility versions (#10051) * chore(ci): update compatibility versions Signed-off-by: greptimedb-ci * test(compat): bound the fulltext bloom case below v1.3.0 Signed-off-by: discord9 (cherry picked from commit ea570c510ed068d2c68ff8bb926165f1b1b91c5d) Signed-off-by: WenyXu --------- Signed-off-by: greptimedb-ci Signed-off-by: discord9 Signed-off-by: WenyXu Co-authored-by: greptimedb-ci Co-authored-by: discord9 --- tests/compatibility/cases/fulltext_bloom_tokens/case.toml | 4 ++-- tests/compatibility/ci.toml | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/compatibility/cases/fulltext_bloom_tokens/case.toml b/tests/compatibility/cases/fulltext_bloom_tokens/case.toml index d796adb51ef..04f720795c0 100644 --- a/tests/compatibility/cases/fulltext_bloom_tokens/case.toml +++ b/tests/compatibility/cases/fulltext_bloom_tokens/case.toml @@ -1,8 +1,8 @@ name = "fulltext_bloom_tokens" -reason = "Verify matches_term and LIKE return every matching row from fulltext bloom SSTs written by older binaries, whose tokens don't follow matches_term boundaries, alone and mixed with SSTs written by the current binary." +reason = "Verify matches_term and LIKE return every matching row from fulltext bloom SSTs written by older binaries, whose tokens don't follow matches_term boundaries, alone and mixed with SSTs written by the current binary. The upper bound is required because the current binary and v1.3.0 both write bloom version 2 (#9470, backported to v1.3 by #9529), so from v1.3.0 there is no version-1 bloom left to mix in and the version-2 SST count below no longer matches." introduced_by = "fix: fulltext bloom index drops rows matched by matches_term" topologies = ["distributed", "standalone"] -from_range = [">=v1.0.0"] +from_range = [" Date: Thu, 29 Oct 2026 16:34:41 +0800 Subject: [PATCH 137/239] feat(meta): add drain migration batch operations (#10035) * feat(meta): add drain migration batch operations Build migrations through the shared manager and collect batch results with route reconciliation and failure strategies. Keep the existing fence-and-plan flow while preparing tested operations for drain orchestration. Signed-off-by: WenyXu * fix(meta): gate procedure receiver helpers to unit tests Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- src/meta-srv/src/error.rs | 13 + src/meta-srv/src/metasrv/builder.rs | 2 + src/meta-srv/src/procedure/drain_datanode.rs | 355 ++----------- .../drain_datanode/drain_datanode_end.rs | 37 ++ .../drain_datanode/drain_datanode_start.rs | 39 ++ .../src/procedure/drain_datanode/manager.rs | 87 ++++ .../src/procedure/drain_datanode/migration.rs | 491 ++++++++++++++++++ .../src/procedure/drain_datanode/test_util.rs | 184 +++++++ .../src/procedure/drain_datanode/utils.rs | 174 ++++++- .../src/procedure/region_migration/manager.rs | 31 +- src/meta-srv/src/procedure/repartition.rs | 4 +- .../src/procedure/repartition/test_util.rs | 61 +-- src/meta-srv/src/procedure/test_util.rs | 71 +++ 13 files changed, 1175 insertions(+), 374 deletions(-) create mode 100644 src/meta-srv/src/procedure/drain_datanode/migration.rs create mode 100644 src/meta-srv/src/procedure/drain_datanode/test_util.rs diff --git a/src/meta-srv/src/error.rs b/src/meta-srv/src/error.rs index 8ceb39ed56e..65d86f9b8ce 100644 --- a/src/meta-srv/src/error.rs +++ b/src/meta-srv/src/error.rs @@ -117,6 +117,18 @@ pub enum Error { reason: String, }, + #[snafu(display( + "Drain datanode {} aborted, failed regions: {:?}", + node_id, + failed_regions + ))] + DrainDatanodeAbort { + node_id: DatanodeId, + failed_regions: Vec, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display( "Another procedure is operating the region: {} on peer: {}", region_id, @@ -1262,6 +1274,7 @@ impl ErrorExt for Error { | Error::MigrationAbort { .. } | Error::MigrationRunning { .. } | Error::DrainDatanodeRunning { .. } + | Error::DrainDatanodeAbort { .. } | Error::RegionMigrated { .. } => StatusCode::Unexpected, Error::TableNotFound { .. } => StatusCode::TableNotFound, Error::SaveClusterInfo { source, .. } diff --git a/src/meta-srv/src/metasrv/builder.rs b/src/meta-srv/src/metasrv/builder.rs index bc10a27af85..44de98e7502 100644 --- a/src/meta-srv/src/metasrv/builder.rs +++ b/src/meta-srv/src/metasrv/builder.rs @@ -416,6 +416,8 @@ impl MetasrvBuilder { procedure_manager.clone(), table_metadata_manager.clone(), node_scheduling_manager.clone(), + region_migration_manager.clone(), + selector_ctx.clone(), )); drain_datanode_manager.try_start()?; let affinity_resolver = plugins diff --git a/src/meta-srv/src/procedure/drain_datanode.rs b/src/meta-srv/src/procedure/drain_datanode.rs index 41bdb85c217..aa8d333478e 100644 --- a/src/meta-srv/src/procedure/drain_datanode.rs +++ b/src/meta-srv/src/procedure/drain_datanode.rs @@ -15,8 +15,12 @@ pub(crate) mod drain_datanode_end; pub(crate) mod drain_datanode_start; pub(crate) mod manager; +mod migration; pub(crate) mod utils; +#[cfg(test)] +mod test_util; + use std::any::Any; use std::fmt::Debug; @@ -24,7 +28,7 @@ use common_meta::DatanodeId; use common_meta::key::TableMetadataManagerRef; use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, Error as ProcedureError, LockKey, Procedure, + Context as ProcedureContext, Error as ProcedureError, LockKey, Procedure, ProcedureId, Result as ProcedureResult, Status, }; use common_telemetry::{error, info}; @@ -33,11 +37,24 @@ use snafu::ResultExt; use store_api::storage::RegionId; use crate::error::Result; +use crate::metasrv::SelectorContext; use crate::node_scheduling::NodeSchedulingManagerRef; use crate::procedure::drain_datanode::drain_datanode_start::DrainDatanodeStart; use crate::procedure::drain_datanode::manager::{ DrainDatanodeProcedureGuard, DrainDatanodeProcedureTracker, }; +use crate::procedure::region_migration::manager::RegionMigrationManagerRef; + +/// The strategy of a drain run when a region cannot be migrated. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +pub enum DrainDatanodeStrategy { + /// Aborts the procedure once every migration of the batch is settled and at least one + /// of them failed. + FailFast, + /// Drains every remaining region and records the failed ones. + #[default] + TryAll, +} /// The regions placed on a datanode, split by role. #[derive(Debug, Clone, Default, PartialEq, Eq)] @@ -60,6 +77,15 @@ pub struct PersistentContext { pub target_node_id: Option, } +/// A dispatched migration sub-procedure that is not settled yet. +#[derive(Debug, Clone, Copy)] +pub struct SubprocedureMeta { + /// The region migrated by the sub-procedure. + pub region_id: RegionId, + /// The procedure id of the sub-procedure. + pub procedure_id: ProcedureId, +} + /// The in-memory data of [`DrainDatanodeProcedure`]. /// /// It is not persisted: a restarted procedure rebuilds the plan by rescanning the datanode. @@ -67,6 +93,12 @@ pub struct PersistentContext { pub struct VolatileContext { /// The drain plan, `None` until the first state builds it. pub plan: Option, + /// The migration sub-procedures of the current batch that are not settled yet. + pub inflight_subprocedures: Vec, + /// The number of regions confirmed to have left the source while collecting batch results. + pub migrated: usize, + /// The regions given up by [`DrainDatanodeStrategy::TryAll`] in this run. + pub failed_regions: Vec, } /// The task of [`DrainDatanodeProcedure`]. @@ -89,6 +121,10 @@ pub struct Context { pub table_metadata_manager: TableMetadataManagerRef, /// The node scheduling manager, used to fence the drained datanode. pub node_scheduling_manager: NodeSchedulingManagerRef, + /// Builds region migration procedures. + pub region_migration_manager: RegionMigrationManagerRef, + /// Resolves destinations for region migrations. + pub selector_ctx: SelectorContext, } impl Context { @@ -99,6 +135,8 @@ impl Context { pub(crate) fn new( table_metadata_manager: TableMetadataManagerRef, node_scheduling_manager: NodeSchedulingManagerRef, + region_migration_manager: RegionMigrationManagerRef, + selector_ctx: SelectorContext, persistent_ctx: PersistentContext, ) -> Self { Self { @@ -106,6 +144,8 @@ impl Context { volatile_ctx: VolatileContext::default(), table_metadata_manager, node_scheduling_manager, + region_migration_manager, + selector_ctx, } } } @@ -128,6 +168,15 @@ impl Context { } } +impl Context { + /// Removes `region_id` from the current plan. A later scan may add it back. + pub(crate) fn forget_leader_region(&mut self, region_id: RegionId) { + if let Some(plan) = self.volatile_ctx.plan.as_mut() { + plan.leader_regions.retain(|region| *region != region_id); + } + } +} + impl VolatileContext { /// Returns the leader regions still to migrate. pub(crate) fn leader_regions(&self) -> &[RegionId] { @@ -289,225 +338,12 @@ impl Procedure for DrainDatanodeProcedure { #[cfg(test)] mod tests { - use std::collections::HashMap; - - use common_meta::key::TableMetadataManagerRef; - use common_meta::key::table_route::TableRouteValue; - use common_meta::key::test_utils::new_test_table_info; - use common_meta::peer::Peer; - use common_meta::rpc::router::{Region, RegionRoute}; - use super::*; - use crate::error; use crate::key::NodeSchedulingRole; - use crate::procedure::drain_datanode::drain_datanode_end::DrainDatanodeEnd; - use crate::procedure::drain_datanode::manager::{ - DrainDatanodeManager, DrainDatanodeProcedureTracker, + use crate::procedure::drain_datanode::test_util::{ + TestEnv, follower_region, leader_region, new_persistent_context, run_until_done, }; use crate::procedure::region_migration::test_util::TestingEnv; - use crate::test_util::{create_meta_peer_client, new_node_scheduling_discovery, put_datanodes}; - - fn new_persistent_context(target_node_id: Option) -> PersistentContext { - PersistentContext { - node_id: 1, - target_node_id, - } - } - - fn leader_region(region_id: RegionId, leader_id: DatanodeId) -> RegionRoute { - RegionRoute { - region: Region::new_test(region_id), - leader_peer: Some(Peer::empty(leader_id)), - follower_peers: vec![], - ..Default::default() - } - } - - fn follower_region( - region_id: RegionId, - leader_id: DatanodeId, - follower_id: DatanodeId, - ) -> RegionRoute { - RegionRoute { - region: Region::new_test(region_id), - leader_peer: Some(Peer::empty(leader_id)), - follower_peers: vec![Peer::empty(follower_id)], - ..Default::default() - } - } - - struct TestEnv { - env: TestingEnv, - table_metadata_manager: TableMetadataManagerRef, - node_scheduling_manager: NodeSchedulingManagerRef, - } - - impl TestEnv { - async fn new() -> Self { - let env = TestingEnv::new(); - // Shares one backend between the scheduling manager and the peer discovery, so - // the selector observes the scheduling marks written by the procedure. - let meta_peer_client = create_meta_peer_client(); - put_datanodes( - &meta_peer_client, - vec![ - Peer { - id: 1, - addr: "dn1:3001".to_string(), - }, - Peer { - id: 2, - addr: "dn2:3001".to_string(), - }, - ], - ) - .await; - let (node_scheduling_manager, _discovery) = - new_node_scheduling_discovery(&meta_peer_client); - let table_metadata_manager = env.table_metadata_manager().clone(); - - Self { - env, - table_metadata_manager, - node_scheduling_manager, - } - } - - fn context(&self, persistent_ctx: PersistentContext) -> Context { - Context::new( - self.table_metadata_manager.clone(), - self.node_scheduling_manager.clone(), - persistent_ctx, - ) - } - - fn table_metadata_manager(&self) -> &TableMetadataManagerRef { - &self.table_metadata_manager - } - - async fn create_table(&self, region_routes: Vec) { - create_table(&self.table_metadata_manager, region_routes).await; - } - } - - async fn create_table( - table_metadata_manager: &TableMetadataManagerRef, - region_routes: Vec, - ) { - table_metadata_manager - .create_table_metadata( - new_test_table_info(1024), - TableRouteValue::physical(region_routes), - HashMap::default(), - ) - .await - .unwrap(); - } - - async fn run_until_done( - state: &mut Box, - ctx: &mut Context, - procedure_ctx: &ProcedureContext, - ) -> Result { - for _ in 0..10 { - let (next, status) = state.next(ctx, procedure_ctx).await?; - *state = next; - if !matches!(status, Status::Executing { .. }) { - return Ok(status); - } - } - - panic!("the drain procedure did not finish in 10 steps"); - } - - #[tokio::test] - async fn test_build_datanode_plan() { - let env = TestEnv::new().await; - let first = RegionId::new(1024, 1); - let second = RegionId::new(1024, 2); - env.create_table(vec![leader_region(first, 1), follower_region(second, 2, 1)]) - .await; - - let plan = utils::build_datanode_plan(env.table_metadata_manager(), 1) - .await - .unwrap(); - - assert_eq!(vec![first], plan.leader_regions); - assert_eq!(vec![second], plan.follower_regions); - } - - #[tokio::test] - async fn test_build_datanode_plan_of_empty_datanode() { - let env = TestEnv::new().await; - - let plan = utils::build_datanode_plan(env.table_metadata_manager(), 1) - .await - .unwrap(); - - assert!(plan.leader_regions.is_empty()); - assert!(plan.follower_regions.is_empty()); - } - - #[tokio::test] - async fn test_build_datanode_plan_reads_the_table_route() { - let env = TestEnv::new().await; - let region_id = RegionId::new(1024, 1); - // The datanode is a follower of the region, so it is not a leader even though the - // datanode table index of the datanode contains the table. - env.create_table(vec![follower_region(region_id, 2, 1)]) - .await; - - let plan = utils::build_datanode_plan(env.table_metadata_manager(), 1) - .await - .unwrap(); - - assert!(plan.leader_regions.is_empty()); - assert_eq!(vec![region_id], plan.follower_regions); - } - - #[tokio::test] - async fn test_drain_datanode_start_transitions_to_end() { - let env = TestEnv::new().await; - env.create_table(vec![leader_region(RegionId::new(1024, 1), 1)]) - .await; - - let mut ctx = env.context(new_persistent_context(Some(2))); - let procedure_ctx = TestingEnv::procedure_context(); - - let (next, status) = DrainDatanodeStart - .next(&mut ctx, &procedure_ctx) - .await - .unwrap(); - - assert!(next.as_any().downcast_ref::().is_some()); - // Registering the procedure is the only persisted step of the run. - assert!(matches!(status, Status::Executing { persist: true, .. })); - assert_eq!( - vec![RegionId::new(1024, 1)], - ctx.volatile_ctx.leader_regions() - ); - assert!( - !ctx.node_scheduling_manager - .is_schedulable(NodeSchedulingRole::Datanode, 1) - .await - .unwrap() - ); - } - - #[tokio::test] - async fn test_drain_datanode_end_transitions_to_done() { - let env = TestEnv::new().await; - let mut ctx = env.context(new_persistent_context(Some(2))); - let procedure_ctx = TestingEnv::procedure_context(); - - let (next, status) = DrainDatanodeEnd - .next(&mut ctx, &procedure_ctx) - .await - .unwrap(); - - assert!(next.as_any().downcast_ref::().is_some()); - assert!(matches!(status, Status::Done { output: None })); - } #[tokio::test] async fn test_drain_datanode_fences_and_collects_the_plan() { @@ -562,91 +398,4 @@ mod tests { // The plan is volatile, so it is not dumped. assert!(!json.contains("leader_regions")); } - - #[tokio::test] - async fn test_drain_datanode_end_finishes() { - let env = TestEnv::new().await; - let mut ctx = env.context(new_persistent_context(Some(2))); - let procedure_ctx = TestingEnv::procedure_context(); - - let (next, status) = DrainDatanodeEnd - .next(&mut ctx, &procedure_ctx) - .await - .unwrap(); - - assert!(next.as_any().downcast_ref::().is_some()); - assert!(matches!(status, Status::Done { .. })); - } - - #[tokio::test] - async fn test_submit_procedure_rejects_same_target() { - let env = TestEnv::new().await; - let manager = DrainDatanodeManager::new( - env.env.procedure_manager().clone(), - env.table_metadata_manager.clone(), - env.node_scheduling_manager.clone(), - ); - - let error = manager - .submit_procedure( - common_procedure::ProcedureContext::default(), - DrainDatanodeProcedureTask { - node_id: 1, - target_node_id: Some(1), - }, - ) - .await - .unwrap_err(); - - assert!(matches!(error, error::Error::InvalidArguments { .. })); - } - - #[tokio::test] - async fn test_submit_procedure_rejects_a_running_drain() { - let env = TestEnv::new().await; - let manager = DrainDatanodeManager::new( - env.env.procedure_manager().clone(), - env.table_metadata_manager.clone(), - env.node_scheduling_manager.clone(), - ); - let task = DrainDatanodeProcedureTask { - node_id: 1, - target_node_id: Some(2), - }; - - // A drain of datanode 1 is running, for example one recovered after a restart. - let _guard = manager.tracker().insert_running_procedure(&task).unwrap(); - - let error = manager - .submit_procedure(common_procedure::ProcedureContext::default(), task) - .await - .unwrap_err(); - - assert!(matches!(error, error::Error::DrainDatanodeRunning { .. })); - assert!(!error.is_retryable()); - } - - #[tokio::test] - async fn test_drain_datanode_tracker_takes_the_place_back_on_reload() { - let env = TestEnv::new().await; - let procedure = - DrainDatanodeProcedure::new(env.context(new_persistent_context(Some(2))), vec![]); - let json = procedure.dump().unwrap(); - let tracker = DrainDatanodeProcedureTracker::default(); - - // Reloading takes the place back, so a second drain of the datanode is rejected. - let reloaded = DrainDatanodeProcedure::from_json(&json, &tracker, |persistent_ctx| { - env.context(persistent_ctx) - }) - .unwrap(); - let task = DrainDatanodeProcedureTask { - node_id: 1, - target_node_id: Some(2), - }; - assert!(tracker.insert_running_procedure(&task).is_none()); - - // Dropping the reloaded procedure releases the place. - drop(reloaded); - assert!(tracker.insert_running_procedure(&task).is_some()); - } } diff --git a/src/meta-srv/src/procedure/drain_datanode/drain_datanode_end.rs b/src/meta-srv/src/procedure/drain_datanode/drain_datanode_end.rs index 1ff5e3041e1..b2880a1c948 100644 --- a/src/meta-srv/src/procedure/drain_datanode/drain_datanode_end.rs +++ b/src/meta-srv/src/procedure/drain_datanode/drain_datanode_end.rs @@ -39,3 +39,40 @@ impl State for DrainDatanodeEnd { self } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::procedure::drain_datanode::test_util::{TestEnv, new_persistent_context}; + use crate::procedure::region_migration::test_util::TestingEnv; + + #[tokio::test] + async fn test_drain_datanode_end_transitions_to_done() { + let env = TestEnv::new().await; + let mut ctx = env.context(new_persistent_context(Some(2))); + let procedure_ctx = TestingEnv::procedure_context(); + + let (next, status) = DrainDatanodeEnd + .next(&mut ctx, &procedure_ctx) + .await + .unwrap(); + + assert!(next.as_any().downcast_ref::().is_some()); + assert!(matches!(status, Status::Done { output: None })); + } + + #[tokio::test] + async fn test_drain_datanode_end_finishes() { + let env = TestEnv::new().await; + let mut ctx = env.context(new_persistent_context(Some(2))); + let procedure_ctx = TestingEnv::procedure_context(); + + let (next, status) = DrainDatanodeEnd + .next(&mut ctx, &procedure_ctx) + .await + .unwrap(); + + assert!(next.as_any().downcast_ref::().is_some()); + assert!(matches!(status, Status::Done { .. })); + } +} diff --git a/src/meta-srv/src/procedure/drain_datanode/drain_datanode_start.rs b/src/meta-srv/src/procedure/drain_datanode/drain_datanode_start.rs index 0e657ebbd00..4cd28c7853f 100644 --- a/src/meta-srv/src/procedure/drain_datanode/drain_datanode_start.rs +++ b/src/meta-srv/src/procedure/drain_datanode/drain_datanode_start.rs @@ -59,3 +59,42 @@ impl State for DrainDatanodeStart { self } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::procedure::drain_datanode::test_util::{ + TestEnv, leader_region, new_persistent_context, + }; + use crate::procedure::region_migration::test_util::TestingEnv; + use store_api::storage::RegionId; + + #[tokio::test] + async fn test_drain_datanode_start_transitions_to_end() { + let env = TestEnv::new().await; + env.create_table(vec![leader_region(RegionId::new(1024, 1), 1)]) + .await; + + let mut ctx = env.context(new_persistent_context(Some(2))); + let procedure_ctx = TestingEnv::procedure_context(); + + let (next, status) = DrainDatanodeStart + .next(&mut ctx, &procedure_ctx) + .await + .unwrap(); + + assert!(next.as_any().downcast_ref::().is_some()); + // Registering the procedure is the only persisted step of the run. + assert!(matches!(status, Status::Executing { persist: true, .. })); + assert_eq!( + vec![RegionId::new(1024, 1)], + ctx.volatile_ctx.leader_regions() + ); + assert!( + !ctx.node_scheduling_manager + .is_schedulable(NodeSchedulingRole::Datanode, 1) + .await + .unwrap() + ); + } +} diff --git a/src/meta-srv/src/procedure/drain_datanode/manager.rs b/src/meta-srv/src/procedure/drain_datanode/manager.rs index 1211356bea0..f262cae2ed1 100644 --- a/src/meta-srv/src/procedure/drain_datanode/manager.rs +++ b/src/meta-srv/src/procedure/drain_datanode/manager.rs @@ -23,11 +23,13 @@ use common_telemetry::info; use snafu::ResultExt; use crate::error::{self, Result}; +use crate::metasrv::SelectorContext; use crate::node_scheduling::NodeSchedulingManagerRef; use crate::procedure::drain_datanode::{ Context as DrainDatanodeContext, DrainDatanodeProcedure, DrainDatanodeProcedureTask, PersistentContext, }; +use crate::procedure::region_migration::manager::RegionMigrationManagerRef; pub type DrainDatanodeManagerRef = Arc; @@ -85,6 +87,8 @@ pub struct DrainDatanodeManager { procedure_manager: ProcedureManagerRef, table_metadata_manager: TableMetadataManagerRef, node_scheduling_manager: NodeSchedulingManagerRef, + region_migration_manager: RegionMigrationManagerRef, + selector_ctx: SelectorContext, tracker: DrainDatanodeProcedureTracker, } @@ -94,11 +98,15 @@ impl DrainDatanodeManager { procedure_manager: ProcedureManagerRef, table_metadata_manager: TableMetadataManagerRef, node_scheduling_manager: NodeSchedulingManagerRef, + region_migration_manager: RegionMigrationManagerRef, + selector_ctx: SelectorContext, ) -> Self { Self { procedure_manager, table_metadata_manager, node_scheduling_manager, + region_migration_manager, + selector_ctx, tracker: DrainDatanodeProcedureTracker::default(), } } @@ -114,6 +122,8 @@ impl DrainDatanodeManager { DrainDatanodeContext::new( self.table_metadata_manager.clone(), self.node_scheduling_manager.clone(), + self.region_migration_manager.clone(), + self.selector_ctx.clone(), persistent_ctx, ) } @@ -162,6 +172,8 @@ impl DrainDatanodeManager { pub(crate) fn try_start(&self) -> Result<()> { let table_metadata_manager = self.table_metadata_manager.clone(); let node_scheduling_manager = self.node_scheduling_manager.clone(); + let region_migration_manager = self.region_migration_manager.clone(); + let selector_ctx = self.selector_ctx.clone(); let tracker = self.tracker.clone(); self.procedure_manager .register_loader( @@ -169,10 +181,14 @@ impl DrainDatanodeManager { Box::new(move |json| { let table_metadata_manager = table_metadata_manager.clone(); let node_scheduling_manager = node_scheduling_manager.clone(); + let region_migration_manager = region_migration_manager.clone(); + let selector_ctx = selector_ctx.clone(); let factory = move |persistent_ctx| { DrainDatanodeContext::new( table_metadata_manager, node_scheduling_manager, + region_migration_manager, + selector_ctx, persistent_ctx, ) }; @@ -185,3 +201,74 @@ impl DrainDatanodeManager { }) } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::procedure::drain_datanode::test_util::{TestEnv, new_persistent_context}; + use common_procedure::Procedure; + + #[tokio::test] + async fn test_submit_procedure_rejects_same_target() { + let env = TestEnv::new().await; + let manager = env.drain_manager(); + + let error = manager + .submit_procedure( + common_procedure::ProcedureContext::default(), + DrainDatanodeProcedureTask { + node_id: 1, + target_node_id: Some(1), + }, + ) + .await + .unwrap_err(); + + assert!(matches!(error, error::Error::InvalidArguments { .. })); + } + + #[tokio::test] + async fn test_submit_procedure_rejects_a_running_drain() { + let env = TestEnv::new().await; + let manager = env.drain_manager(); + let task = DrainDatanodeProcedureTask { + node_id: 1, + target_node_id: Some(2), + }; + + // A drain of datanode 1 is running, for example one recovered after a restart. + let _guard = manager.tracker().insert_running_procedure(&task).unwrap(); + + let error = manager + .submit_procedure(common_procedure::ProcedureContext::default(), task) + .await + .unwrap_err(); + + assert!(matches!(error, error::Error::DrainDatanodeRunning { .. })); + assert!(!error.is_retryable()); + } + + #[tokio::test] + async fn test_drain_datanode_tracker_takes_the_place_back_on_reload() { + let env = TestEnv::new().await; + let procedure = + DrainDatanodeProcedure::new(env.context(new_persistent_context(Some(2))), vec![]); + let json = procedure.dump().unwrap(); + let tracker = DrainDatanodeProcedureTracker::default(); + + // Reloading takes the place back, so a second drain of the datanode is rejected. + let reloaded = DrainDatanodeProcedure::from_json(&json, &tracker, |persistent_ctx| { + env.context(persistent_ctx) + }) + .unwrap(); + let task = DrainDatanodeProcedureTask { + node_id: 1, + target_node_id: Some(2), + }; + assert!(tracker.insert_running_procedure(&task).is_none()); + + // Dropping the reloaded procedure releases the place. + drop(reloaded); + assert!(tracker.insert_running_procedure(&task).is_some()); + } +} diff --git a/src/meta-srv/src/procedure/drain_datanode/migration.rs b/src/meta-srv/src/procedure/drain_datanode/migration.rs new file mode 100644 index 00000000000..1767f7494ad --- /dev/null +++ b/src/meta-srv/src/procedure/drain_datanode/migration.rs @@ -0,0 +1,491 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use common_meta::DatanodeId; +use common_meta::key::table_route::TableRouteValue; +use common_meta::peer::Peer; +use common_procedure::{ + Context as ProcedureContext, ProcedureContext as SubmissionContext, ProcedureWithId, watcher, +}; +use common_telemetry::{error, info, warn}; +use snafu::ResultExt; +use store_api::storage::RegionId; + +use crate::error::{self, Result}; +use crate::metrics::METRIC_META_REGION_MIGRATION_FAIL; +use crate::procedure::drain_datanode::{Context, DrainDatanodeStrategy, utils}; +use crate::procedure::region_migration::DEFAULT_REGION_MIGRATION_TIMEOUT; +use crate::procedure::region_migration::manager::{ + RegionMigrationProcedureTask, RegionMigrationTriggerReason, +}; + +impl Context { + /// Waits for every migration in the batch before applying the drain strategy. + /// + /// Keep the batch and its results unchanged until all fallible route checks succeed, + /// so a metadata retry cannot lose failures or apply progress twice. + pub async fn wait_for_inflight_subprocedures( + &mut self, + procedure_ctx: &ProcedureContext, + strategy: DrainDatanodeStrategy, + ) -> Result<()> { + let node_id = self.persistent_ctx.node_id; + let mut migrated = Vec::new(); + let mut failed = Vec::new(); + for meta in &self.volatile_ctx.inflight_subprocedures { + let receiver = procedure_ctx + .provider + .procedure_state_receiver(meta.procedure_id) + .await + .inspect_err(|e| { + error!(e; "Failed to get state receiver for migration sub-procedure {}", meta.procedure_id); + }) + .context(error::QueryProcedureSnafu)?; + let Some(mut receiver) = receiver else { + warn!( + "Migration sub-procedure {} of region {} is unknown, keep the region pending", + meta.procedure_id, meta.region_id + ); + continue; + }; + match watcher::wait(&mut receiver).await { + Ok(_) => migrated.push(meta.region_id), + Err(e) => { + error!(e; "Migration sub-procedure {} of region {} failed", meta.procedure_id, meta.region_id); + failed.push(meta.region_id); + } + } + } + + let failed_count = failed.len(); + let mut failed_regions = Vec::new(); + for region_id in failed { + // A concurrent failover may have achieved the goal despite the child failure. + if utils::leads_region(&self.table_metadata_manager, node_id, region_id).await? { + failed_regions.push(region_id); + } else { + migrated.push(region_id); + } + } + + METRIC_META_REGION_MIGRATION_FAIL.inc_by(failed_count as u64); + for region_id in migrated { + self.forget_leader_region(region_id); + self.volatile_ctx.migrated += 1; + info!("Region {} left datanode {}", region_id, node_id); + } + self.volatile_ctx.inflight_subprocedures.clear(); + match strategy { + DrainDatanodeStrategy::FailFast if !failed_regions.is_empty() => { + error::DrainDatanodeAbortSnafu { + node_id, + failed_regions, + } + .fail() + } + DrainDatanodeStrategy::TryAll => { + for ®ion_id in &failed_regions { + self.forget_leader_region(region_id); + } + self.volatile_ctx.failed_regions.extend(failed_regions); + Ok(()) + } + DrainDatanodeStrategy::FailFast => Ok(()), + } + } +} + +impl Context { + /// Builds the migration sub-procedure of `region_id`. + /// + /// Returns `Ok(None)` when the region is no longer led by the drained datanode. + /// Otherwise returns the sub-procedure and its target. + pub async fn build_migration_procedure( + &self, + procedure_ctx: &ProcedureContext, + region_id: RegionId, + ) -> Result> { + let node_id = self.persistent_ctx.node_id; + let Some(table_route) = self + .table_metadata_manager + .table_route_manager() + .table_route_storage() + .get(region_id.table_id()) + .await + .context(error::TableMetadataManagerSnafu)? + else { + return Ok(None); + }; + let TableRouteValue::Physical(table_route) = table_route else { + return Ok(None); + }; + let Some(region_route) = table_route + .region_routes + .iter() + .find(|region_route| region_route.region.id == region_id) + else { + return Ok(None); + }; + if region_route.leader_peer.as_ref().map(|peer| peer.id) != Some(node_id) { + return Ok(None); + } + + let to_peer = + utils::select_migration_target(self, region_id, ®ion_route.follower_peers).await?; + + let submission_context = match procedure_ctx.event_context.clone() { + Some(event_context) => SubmissionContext::from_event_context(event_context), + None => SubmissionContext::default(), + }; + let target_node_id = to_peer.id; + let child = self + .region_migration_manager + .build_subprocedure( + submission_context, + RegionMigrationProcedureTask::new( + region_id, + Peer::empty(node_id), + to_peer, + DEFAULT_REGION_MIGRATION_TIMEOUT, + RegionMigrationTriggerReason::Manual, + ), + ) + .await?; + + Ok(Some((child, target_node_id))) + } +} +#[cfg(test)] +mod tests { + use crate::error; + use crate::procedure::drain_datanode::test_util::{ + TestEnv, leader_region, leader_regions, new_persistent_context, seed_plan, + }; + use crate::procedure::drain_datanode::{DrainDatanodeStrategy, SubprocedureMeta}; + use crate::procedure::test_util::{procedure_context_with_receivers, procedure_state_receiver}; + use common_error::mock::MockError; + use common_error::status_code::StatusCode; + use common_procedure::{Error as ProcedureError, ProcedureId, ProcedureState}; + use std::collections::HashMap; + use std::sync::Arc; + use store_api::storage::RegionId; + use tokio::sync::watch; + #[tokio::test] + async fn test_build_migration_procedure_checks_route_and_guard() { + for leader_id in [1, 2] { + let env = TestEnv::new().await; + let region_id = RegionId::new(1024, 1); + env.create_table(vec![leader_region(region_id, leader_id)]) + .await; + let ctx = env.context(new_persistent_context(Some(2))); + let procedure_ctx = + crate::procedure::region_migration::test_util::TestingEnv::procedure_context(); + let child = ctx + .build_migration_procedure(&procedure_ctx, region_id) + .await + .unwrap(); + if leader_id == 2 { + assert!(child.is_none()); + assert!(!env.region_migration_manager.tracker().contains(region_id)); + continue; + } + let (child, target) = child.unwrap(); + assert_eq!(2, target); + assert!(env.region_migration_manager.tracker().contains(region_id)); + let error = ctx + .build_migration_procedure(&procedure_ctx, region_id) + .await + .unwrap_err(); + assert!(matches!(error, error::Error::MigrationRunning { .. })); + drop(child); + assert!(!env.region_migration_manager.tracker().contains(region_id)); + } + } + + #[tokio::test] + async fn test_drain_datanode_waits_for_migrated_batch() { + let env = TestEnv::new().await; + let region_id = RegionId::new(1024, 1); + env.create_table(vec![leader_region(region_id, 1)]).await; + + let mut ctx = env.context(new_persistent_context(Some(2))); + seed_plan(&mut ctx, vec![region_id]); + ctx.volatile_ctx.inflight_subprocedures = vec![SubprocedureMeta { + region_id, + procedure_id: ProcedureId::random(), + }]; + let procedure_ctx = procedure_context_with_receivers(HashMap::from([( + ctx.volatile_ctx.inflight_subprocedures[0].procedure_id, + procedure_state_receiver(ProcedureState::Done { output: None }), + )])); + + ctx.wait_for_inflight_subprocedures(&procedure_ctx, DrainDatanodeStrategy::FailFast) + .await + .unwrap(); + + assert_eq!(1, ctx.volatile_ctx.migrated); + assert!(leader_regions(&ctx).is_empty()); + assert!(ctx.volatile_ctx.inflight_subprocedures.is_empty()); + } + + #[tokio::test] + async fn test_drain_datanode_wait_failed_child_fail_fast() { + let env = TestEnv::new().await; + let region_id = RegionId::new(1024, 1); + env.create_table(vec![leader_region(region_id, 1)]).await; + + let mut ctx = env.context(new_persistent_context(Some(2))); + seed_plan(&mut ctx, vec![region_id]); + ctx.volatile_ctx.inflight_subprocedures = vec![SubprocedureMeta { + region_id, + procedure_id: ProcedureId::random(), + }]; + let procedure_ctx = procedure_context_with_receivers(HashMap::from([( + ctx.volatile_ctx.inflight_subprocedures[0].procedure_id, + procedure_state_receiver(ProcedureState::failed(Arc::new(ProcedureError::external( + MockError::new(StatusCode::Internal), + )))), + )])); + + let error = ctx + .wait_for_inflight_subprocedures(&procedure_ctx, DrainDatanodeStrategy::FailFast) + .await + .unwrap_err(); + + // The whole batch is settled before the procedure aborts. + assert!(matches!(error, error::Error::DrainDatanodeAbort { .. })); + assert_eq!(0, ctx.volatile_ctx.migrated); + assert_eq!(vec![region_id], leader_regions(&ctx)); + assert!(ctx.volatile_ctx.inflight_subprocedures.is_empty()); + } + + #[tokio::test] + async fn test_drain_datanode_wait_failed_child_try_all() { + let env = TestEnv::new().await; + let region_id = RegionId::new(1024, 1); + env.create_table(vec![leader_region(region_id, 1)]).await; + + let mut ctx = env.context(new_persistent_context(Some(2))); + seed_plan(&mut ctx, vec![region_id]); + ctx.volatile_ctx.inflight_subprocedures = vec![SubprocedureMeta { + region_id, + procedure_id: ProcedureId::random(), + }]; + let procedure_ctx = procedure_context_with_receivers(HashMap::from([( + ctx.volatile_ctx.inflight_subprocedures[0].procedure_id, + procedure_state_receiver(ProcedureState::failed(Arc::new(ProcedureError::external( + MockError::new(StatusCode::Internal), + )))), + )])); + + ctx.wait_for_inflight_subprocedures(&procedure_ctx, DrainDatanodeStrategy::TryAll) + .await + .unwrap(); + + assert_eq!(vec![region_id], ctx.volatile_ctx.failed_regions); + assert_eq!(0, ctx.volatile_ctx.migrated); + assert!(leader_regions(&ctx).is_empty()); + } + + #[tokio::test] + async fn test_drain_datanode_wait_failed_child_counts_migrated_when_moved() { + let env = TestEnv::new().await; + let region_id = RegionId::new(1024, 1); + // The route reports another datanode as the leader, which is what a failover that + // raced with the migration leaves behind. + env.create_table(vec![leader_region(region_id, 2)]).await; + + let mut ctx = env.context(new_persistent_context(Some(2))); + seed_plan(&mut ctx, vec![region_id]); + ctx.volatile_ctx.inflight_subprocedures = vec![SubprocedureMeta { + region_id, + procedure_id: ProcedureId::random(), + }]; + let procedure_ctx = procedure_context_with_receivers(HashMap::from([( + ctx.volatile_ctx.inflight_subprocedures[0].procedure_id, + procedure_state_receiver(ProcedureState::failed(Arc::new(ProcedureError::external( + MockError::new(StatusCode::Internal), + )))), + )])); + + ctx.wait_for_inflight_subprocedures(&procedure_ctx, DrainDatanodeStrategy::FailFast) + .await + .unwrap(); + + assert_eq!(1, ctx.volatile_ctx.migrated); + assert!(ctx.volatile_ctx.failed_regions.is_empty()); + assert!(leader_regions(&ctx).is_empty()); + } + + #[tokio::test] + async fn test_drain_datanode_wait_unknown_child_keeps_pending() { + let env = TestEnv::new().await; + let region_id = RegionId::new(1024, 1); + env.create_table(vec![leader_region(region_id, 1)]).await; + + let mut ctx = env.context(new_persistent_context(Some(2))); + seed_plan(&mut ctx, vec![region_id]); + ctx.volatile_ctx.inflight_subprocedures = vec![SubprocedureMeta { + region_id, + procedure_id: ProcedureId::random(), + }]; + let procedure_ctx = procedure_context_with_receivers(HashMap::new()); + + ctx.wait_for_inflight_subprocedures(&procedure_ctx, DrainDatanodeStrategy::FailFast) + .await + .unwrap(); + + // The region stays pending and is dispatched again instead of being reported as + // failed. + assert_eq!(vec![region_id], leader_regions(&ctx)); + assert_eq!(0, ctx.volatile_ctx.migrated); + assert!(ctx.volatile_ctx.failed_regions.is_empty()); + } + + #[tokio::test] + async fn test_drain_datanode_waits_for_entire_batch() { + for strategy in [ + DrainDatanodeStrategy::FailFast, + DrainDatanodeStrategy::TryAll, + ] { + let env = TestEnv::new().await; + let failed_region = RegionId::new(1024, 1); + let succeeded_region = RegionId::new(1024, 2); + env.create_table(vec![ + leader_region(failed_region, 1), + leader_region(succeeded_region, 1), + ]) + .await; + let mut ctx = env.context(new_persistent_context(Some(2))); + seed_plan(&mut ctx, vec![failed_region, succeeded_region]); + let failed_id = ProcedureId::random(); + let pending_id = ProcedureId::random(); + ctx.volatile_ctx.inflight_subprocedures = vec![ + SubprocedureMeta { + region_id: failed_region, + procedure_id: failed_id, + }, + SubprocedureMeta { + region_id: succeeded_region, + procedure_id: pending_id, + }, + ]; + let (sender, receiver) = watch::channel(ProcedureState::Running); + let procedure_ctx = procedure_context_with_receivers(HashMap::from([ + ( + failed_id, + procedure_state_receiver(ProcedureState::failed(Arc::new( + ProcedureError::external(MockError::new(StatusCode::Internal)), + ))), + ), + (pending_id, receiver), + ])); + let mut waiting = + Box::pin(ctx.wait_for_inflight_subprocedures(&procedure_ctx, strategy)); + assert!(futures::poll!(&mut waiting).is_pending()); + sender.send(ProcedureState::Done { output: None }).unwrap(); + let result = waiting.await; + match strategy { + DrainDatanodeStrategy::FailFast => assert!( + matches!(result, Err(error::Error::DrainDatanodeAbort { failed_regions, .. }) if failed_regions == vec![failed_region]) + ), + DrainDatanodeStrategy::TryAll => { + result.unwrap(); + assert_eq!(vec![failed_region], ctx.volatile_ctx.failed_regions); + } + } + assert_eq!(1, ctx.volatile_ctx.migrated); + assert!(ctx.volatile_ctx.inflight_subprocedures.is_empty()); + } + } + + #[tokio::test] + async fn test_drain_datanode_wait_preserves_batch_on_metadata_retry() { + use common_meta::key::TableMetadataManager; + use common_meta::kv_backend::test_util::MockKvBackendBuilder; + + for strategy in [ + DrainDatanodeStrategy::FailFast, + DrainDatanodeStrategy::TryAll, + ] { + let env = TestEnv::new().await; + let failed_region = RegionId::new(1024, 1); + let succeeded_region = RegionId::new(1024, 2); + env.create_table(vec![ + leader_region(failed_region, 1), + leader_region(succeeded_region, 1), + ]) + .await; + let mut ctx = env.context(new_persistent_context(Some(2))); + seed_plan(&mut ctx, vec![failed_region, succeeded_region]); + let failed_id = ProcedureId::random(); + let succeeded_id = ProcedureId::random(); + ctx.volatile_ctx.inflight_subprocedures = vec![ + SubprocedureMeta { + region_id: succeeded_region, + procedure_id: succeeded_id, + }, + SubprocedureMeta { + region_id: failed_region, + procedure_id: failed_id, + }, + ]; + let procedure_ctx = procedure_context_with_receivers(HashMap::from([ + ( + succeeded_id, + procedure_state_receiver(ProcedureState::Done { output: None }), + ), + ( + failed_id, + procedure_state_receiver(ProcedureState::failed(Arc::new( + ProcedureError::external(MockError::new(StatusCode::Internal)), + ))), + ), + ])); + ctx.table_metadata_manager = Arc::new(TableMetadataManager::new(Arc::new( + MockKvBackendBuilder::default() + .range_fn(Arc::new(|_| { + common_meta::error::ElectionNoLeaderSnafu.fail() + })) + .build() + .unwrap(), + ))); + let error = ctx + .wait_for_inflight_subprocedures(&procedure_ctx, strategy) + .await + .unwrap_err(); + assert!(error.is_retryable()); + assert_eq!(2, ctx.volatile_ctx.inflight_subprocedures.len()); + assert_eq!(0, ctx.volatile_ctx.migrated); + assert!(ctx.volatile_ctx.failed_regions.is_empty()); + assert_eq!(vec![failed_region, succeeded_region], leader_regions(&ctx)); + + ctx.table_metadata_manager = env.table_metadata_manager().clone(); + let result = ctx + .wait_for_inflight_subprocedures(&procedure_ctx, strategy) + .await; + match strategy { + DrainDatanodeStrategy::FailFast => assert!( + matches!(result, Err(error::Error::DrainDatanodeAbort { failed_regions, .. }) if failed_regions == vec![failed_region]) + ), + DrainDatanodeStrategy::TryAll => { + result.unwrap(); + assert_eq!(vec![failed_region], ctx.volatile_ctx.failed_regions); + } + } + assert_eq!(1, ctx.volatile_ctx.migrated); + assert!(ctx.volatile_ctx.inflight_subprocedures.is_empty()); + } + } +} diff --git a/src/meta-srv/src/procedure/drain_datanode/test_util.rs b/src/meta-srv/src/procedure/drain_datanode/test_util.rs new file mode 100644 index 00000000000..aebdb16e0f6 --- /dev/null +++ b/src/meta-srv/src/procedure/drain_datanode/test_util.rs @@ -0,0 +1,184 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use crate::metasrv::SelectorContext; +use crate::procedure::drain_datanode::DrainDatanodePlan; +use crate::procedure::drain_datanode::manager::DrainDatanodeManager; +use crate::procedure::region_migration::manager::{ + RegionMigrationManager, RegionMigrationManagerRef, +}; +use common_meta::affinity::NoopAffinity; +use common_procedure::ProcedureManagerRef; +use std::sync::Arc; + +use common_meta::DatanodeId; +use common_meta::key::TableMetadataManagerRef; +use common_meta::key::test_utils::new_test_table_info; +use common_meta::peer::Peer; +use common_meta::rpc::router::{Region, RegionRoute}; +use common_procedure::{Context as ProcedureContext, Status}; +use store_api::storage::RegionId; + +use crate::error::Result; +use crate::node_scheduling::NodeSchedulingManagerRef; +use crate::procedure::drain_datanode::{Context, PersistentContext, State}; +use crate::procedure::region_migration::test_util::TestingEnv; +use crate::test_util::{create_meta_peer_client, new_node_scheduling_discovery, put_datanodes}; + +pub(super) fn new_persistent_context(target_node_id: Option) -> PersistentContext { + PersistentContext { + node_id: 1, + target_node_id, + } +} + +pub(super) fn leader_region(region_id: RegionId, leader_id: DatanodeId) -> RegionRoute { + RegionRoute { + region: Region::new_test(region_id), + leader_peer: Some(Peer::empty(leader_id)), + follower_peers: vec![], + ..Default::default() + } +} + +pub(super) fn follower_region( + region_id: RegionId, + leader_id: DatanodeId, + follower_id: DatanodeId, +) -> RegionRoute { + RegionRoute { + region: Region::new_test(region_id), + leader_peer: Some(Peer::empty(leader_id)), + follower_peers: vec![Peer::empty(follower_id)], + ..Default::default() + } +} + +pub(super) struct TestEnv { + pub(super) env: TestingEnv, + pub(super) table_metadata_manager: TableMetadataManagerRef, + pub(super) node_scheduling_manager: NodeSchedulingManagerRef, + pub(super) region_migration_manager: RegionMigrationManagerRef, + pub(super) selector_ctx: SelectorContext, +} + +impl TestEnv { + pub(super) async fn new() -> Self { + let env = TestingEnv::new(); + // Shares one backend between the scheduling manager and the peer discovery, so + // the selector observes the scheduling marks written by the procedure. + let meta_peer_client = create_meta_peer_client(); + put_datanodes( + &meta_peer_client, + vec![ + Peer { + id: 1, + addr: "dn1:3001".to_string(), + }, + Peer { + id: 2, + addr: "dn2:3001".to_string(), + }, + Peer { + id: 3, + addr: "dn3:3001".to_string(), + }, + ], + ) + .await; + let (node_scheduling_manager, discovery) = new_node_scheduling_discovery(&meta_peer_client); + let selector_ctx = SelectorContext { + affinity_fallback: Default::default(), + affinity: Arc::new(NoopAffinity), + peer_discovery: Arc::new(discovery), + }; + let region_migration_manager = Arc::new(RegionMigrationManager::new( + env.procedure_manager().clone(), + env.context_factory(), + )); + + let table_metadata_manager = env.table_metadata_manager().clone(); + + Self { + env, + table_metadata_manager, + node_scheduling_manager, + region_migration_manager, + selector_ctx, + } + } + + pub(super) fn context(&self, persistent_ctx: PersistentContext) -> Context { + Context::new( + self.table_metadata_manager.clone(), + self.node_scheduling_manager.clone(), + self.region_migration_manager.clone(), + self.selector_ctx.clone(), + persistent_ctx, + ) + } + + pub(super) fn table_metadata_manager(&self) -> &TableMetadataManagerRef { + &self.table_metadata_manager + } + + pub(super) async fn create_table(&self, region_routes: Vec) { + self.env + .create_physical_table_metadata(new_test_table_info(1024), region_routes) + .await; + } + + pub(super) fn procedure_manager(&self) -> &ProcedureManagerRef { + self.env.procedure_manager() + } + + pub(super) fn drain_manager(&self) -> DrainDatanodeManager { + DrainDatanodeManager::new( + self.procedure_manager().clone(), + self.table_metadata_manager.clone(), + self.node_scheduling_manager.clone(), + self.region_migration_manager.clone(), + self.selector_ctx.clone(), + ) + } +} + +pub(super) async fn run_until_done( + state: &mut Box, + ctx: &mut Context, + procedure_ctx: &ProcedureContext, +) -> Result { + for _ in 0..10 { + let (next, status) = state.next(ctx, procedure_ctx).await?; + *state = next; + if !matches!(status, Status::Executing { .. }) { + return Ok(status); + } + } + + panic!("the drain procedure did not finish in 10 steps"); +} + +/// Seeds the volatile plan for migration batch tests. +pub(super) fn seed_plan(ctx: &mut Context, leader_regions: Vec) { + ctx.volatile_ctx.plan = Some(DrainDatanodePlan { + leader_regions, + follower_regions: vec![], + }); +} + +/// Returns the leader regions the procedure still has to migrate. +pub(super) fn leader_regions(ctx: &Context) -> Vec { + ctx.volatile_ctx.leader_regions().to_vec() +} diff --git a/src/meta-srv/src/procedure/drain_datanode/utils.rs b/src/meta-srv/src/procedure/drain_datanode/utils.rs index a43ba815e18..8f5dfa70af2 100644 --- a/src/meta-srv/src/procedure/drain_datanode/utils.rs +++ b/src/meta-srv/src/procedure/drain_datanode/utils.rs @@ -17,13 +17,15 @@ use std::collections::BTreeSet; use common_meta::DatanodeId; use common_meta::key::TableMetadataManagerRef; use common_meta::key::table_route::TableRouteValue; +use common_meta::peer::Peer; use common_meta::rpc::router::region_distribution; use futures::TryStreamExt; use snafu::ResultExt; use store_api::storage::RegionId; use crate::error::{self, Result}; -use crate::procedure::drain_datanode::DrainDatanodePlan; +use crate::procedure::drain_datanode::{Context, DrainDatanodePlan}; +use crate::selector::common; /// Builds the drain plan of `node_id`. /// @@ -81,3 +83,173 @@ pub(crate) async fn build_datanode_plan( follower_regions: follower_regions.into_iter().collect(), }) } + +/// Returns true if `node_id` still leads `region_id`. +/// +/// The table route is authoritative, so a region that is not listed anymore is reported as +/// not led by the datanode. +pub(crate) async fn leads_region( + table_metadata_manager: &TableMetadataManagerRef, + node_id: DatanodeId, + region_id: RegionId, +) -> Result { + let Some(table_route) = table_metadata_manager + .table_route_manager() + .table_route_storage() + .get(region_id.table_id()) + .await + .context(error::TableMetadataManagerSnafu)? + else { + return Ok(false); + }; + let TableRouteValue::Physical(table_route) = table_route else { + return Ok(false); + }; + let Some(region_route) = table_route + .region_routes + .iter() + .find(|region_route| region_route.region.id == region_id) + else { + return Ok(false); + }; + + Ok(region_route.leader_peer.as_ref().map(|peer| peer.id) == Some(node_id)) +} + +/// Selects the migration destination of `region_id` off the drained datanode. +/// +/// Returns the pinned target when the procedure has one, otherwise the schedulable +/// datanode with the smallest peer id. Candidates come from the +/// peer discovery, so the returned peer carries its address and was not marked +/// unschedulable; the drained datanode and the existing followers are excluded. +pub(crate) async fn select_migration_target( + ctx: &Context, + region_id: RegionId, + follower_peers: &[Peer], +) -> Result { + let node_id = ctx.persistent_ctx.node_id; + let target_node_id = ctx.persistent_ctx.target_node_id; + let target = common::active_schedulable_datanodes(&ctx.selector_ctx, None) + .await? + .into_iter() + .map(|datanode| datanode.peer) + .filter(|peer| peer.id != node_id) + .filter(|peer| target_node_id.is_none_or(|target| peer.id == target)) + .filter(|peer| !follower_peers.iter().any(|follower| follower.id == peer.id)) + .min_by_key(|peer| peer.id); + + match target { + Some(peer) => Ok(peer), + None => error::DrainDatanodeAbortSnafu { + node_id, + failed_regions: vec![region_id], + } + .fail(), + } +} + +#[cfg(test)] +mod tests { + use crate::error; + use crate::key::NodeSchedulingRole; + use crate::procedure::drain_datanode::test_util::new_persistent_context; + use crate::procedure::drain_datanode::test_util::{TestEnv, follower_region, leader_region}; + use crate::procedure::drain_datanode::utils; + use common_meta::peer::Peer; + use store_api::storage::RegionId; + + #[tokio::test] + async fn test_build_datanode_plan() { + let env = TestEnv::new().await; + let first = RegionId::new(1024, 1); + let second = RegionId::new(1024, 2); + env.create_table(vec![leader_region(first, 1), follower_region(second, 2, 1)]) + .await; + + let plan = utils::build_datanode_plan(env.table_metadata_manager(), 1) + .await + .unwrap(); + + assert_eq!(vec![first], plan.leader_regions); + assert_eq!(vec![second], plan.follower_regions); + } + + #[tokio::test] + async fn test_build_datanode_plan_of_empty_datanode() { + let env = TestEnv::new().await; + + let plan = utils::build_datanode_plan(env.table_metadata_manager(), 1) + .await + .unwrap(); + + assert!(plan.leader_regions.is_empty()); + assert!(plan.follower_regions.is_empty()); + } + + #[tokio::test] + async fn test_build_datanode_plan_reads_the_table_route() { + let env = TestEnv::new().await; + let region_id = RegionId::new(1024, 1); + // The datanode is a follower of the region, so it is not a leader even though the + // datanode table index of the datanode contains the table. + env.create_table(vec![follower_region(region_id, 2, 1)]) + .await; + + let plan = utils::build_datanode_plan(env.table_metadata_manager(), 1) + .await + .unwrap(); + + assert!(plan.leader_regions.is_empty()); + assert_eq!(vec![region_id], plan.follower_regions); + } + #[tokio::test] + async fn test_select_migration_target() { + let env = TestEnv::new().await; + let region_id = RegionId::new(1024, 1); + + let mut ctx = env.context(new_persistent_context(None)); + let peer = utils::select_migration_target(&ctx, region_id, &[]) + .await + .unwrap(); + assert_eq!(2, peer.id); + + // The existing followers and the drained datanode are excluded. + let peer = utils::select_migration_target(&ctx, region_id, &[Peer::empty(2)]) + .await + .unwrap(); + assert_eq!(3, peer.id); + assert_eq!("dn3:3001", peer.addr); + + // A pinned target is resolved to the active peer, so it carries its address. + ctx.persistent_ctx.target_node_id = Some(2); + let peer = utils::select_migration_target(&ctx, region_id, &[]) + .await + .unwrap(); + assert_eq!(2, peer.id); + assert_eq!("dn2:3001", peer.addr); + + // A target that is already a follower, the drained datanode itself or an + // unschedulable datanode is rejected. + ctx.persistent_ctx.target_node_id = Some(2); + let error = utils::select_migration_target(&ctx, region_id, &[Peer::empty(2)]) + .await + .unwrap_err(); + assert!(matches!(error, error::Error::DrainDatanodeAbort { .. })); + + ctx.persistent_ctx.target_node_id = Some(1); + let error = utils::select_migration_target(&ctx, region_id, &[]) + .await + .unwrap_err(); + assert!(matches!(error, error::Error::DrainDatanodeAbort { .. })); + + env.node_scheduling_manager + .mark_unschedulable(NodeSchedulingRole::Datanode, 3) + .await + .unwrap(); + ctx.persistent_ctx.target_node_id = Some(3); + let error = utils::select_migration_target(&ctx, region_id, &[]) + .await + .unwrap_err(); + assert!(matches!(error, error::Error::DrainDatanodeAbort { .. })); + } +} diff --git a/src/meta-srv/src/procedure/region_migration/manager.rs b/src/meta-srv/src/procedure/region_migration/manager.rs index 8a115cd9f8e..998e13fad20 100644 --- a/src/meta-srv/src/procedure/region_migration/manager.rs +++ b/src/meta-srv/src/procedure/region_migration/manager.rs @@ -537,17 +537,15 @@ impl RegionMigrationManager { Ok(procedure_id) } - /// Submits a new region migration procedure. - pub async fn submit_procedure( + /// Builds a new region migration procedure for `task` without submitting it. + /// + /// Returns the built [`ProcedureWithId`] so that a parent procedure can run it as a + /// child procedure. + pub async fn build_subprocedure( &self, procedure_context: ProcedureContext, mut task: RegionMigrationProcedureTask, - ) -> Result> { - if let Some(event_context) = procedure_context.event_context.as_ref() { - task.trigger_reason = - RegionMigrationTriggerReason::from_trigger_reason(event_context.reason); - } - + ) -> Result { let Some(guard) = self.insert_running_procedure(&task) else { return error::MigrationRunningSnafu { region_id: task.region_id, @@ -608,6 +606,23 @@ impl RegionMigrationManager { ); let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); + Ok(procedure_with_id) + } + + /// Submits a new region migration procedure. + pub async fn submit_procedure( + &self, + procedure_context: ProcedureContext, + mut task: RegionMigrationProcedureTask, + ) -> Result> { + if let Some(event_context) = procedure_context.event_context.as_ref() { + task.trigger_reason = + RegionMigrationTriggerReason::from_trigger_reason(event_context.reason); + } + + let procedure_with_id = self + .build_subprocedure(procedure_context, task.clone()) + .await?; let procedure_id = procedure_with_id.id; info!("Starting region migration procedure {procedure_id} for {task}"); let mut watcher = self diff --git a/src/meta-srv/src/procedure/repartition.rs b/src/meta-srv/src/procedure/repartition.rs index b137eac2166..9011a122b90 100644 --- a/src/meta-srv/src/procedure/repartition.rs +++ b/src/meta-srv/src/procedure/repartition.rs @@ -1021,6 +1021,7 @@ impl RepartitionProcedureFactory for DefaultRepartitionProcedureFactory { #[cfg(test)] mod tests { + use crate::procedure::test_util::{procedure_context_with_receivers, procedure_state_receiver}; use std::collections::HashMap; use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; @@ -1057,8 +1058,7 @@ mod tests { use crate::procedure::repartition::repartition_end::RepartitionEnd; use crate::procedure::repartition::test_util::{ TestingEnv, assert_parent_state, current_parent_region_routes, extract_subprocedure_ids, - new_parent_context, procedure_context_with_receivers, procedure_state_receiver, range_expr, - test_region_route, test_region_wal_options, + new_parent_context, range_expr, test_region_route, test_region_wal_options, }; use crate::procedure::repartition::update_partition_metadata::{ PartitionMetadataUpdate, UpdatePartitionMetadata, diff --git a/src/meta-srv/src/procedure/repartition/test_util.rs b/src/meta-srv/src/procedure/repartition/test_util.rs index 2ce494ee7c9..1fe3553286e 100644 --- a/src/meta-srv/src/procedure/repartition/test_util.rs +++ b/src/meta-srv/src/procedure/repartition/test_util.rs @@ -28,16 +28,13 @@ use common_meta::rpc::router::{Region, RegionRoute}; use common_meta::sequence::SequenceBuilder; use common_meta::test_util::new_ddl_context_with_kv_backend; use common_meta::wal_provider::RegionWalOptions; -use common_procedure::{ - Context as ProcedureContext, ContextProvider, ProcedureId, ProcedureState, Status, -}; +use common_procedure::{Context as ProcedureContext, ProcedureId, Status}; use common_procedure_test::MockContextProvider; use common_wal::options::{KafkaWalOptions, WalOptions}; use datatypes::value::Value; use partition::expr::{PartitionExpr, col}; use store_api::storage::{RegionId, RegionNumber, TableId}; use table::table_name::TableName; -use tokio::sync::watch; use uuid::Uuid; use crate::cache_invalidator::MetasrvCacheInvalidator; @@ -249,59 +246,3 @@ pub fn extract_subprocedure_ids(status: Status) -> Vec { .map(|procedure| procedure.id) .collect() } - -pub fn procedure_state_receiver(state: ProcedureState) -> watch::Receiver { - let (tx, rx) = watch::channel(ProcedureState::Running); - tx.send(state).unwrap(); - rx -} - -pub fn procedure_context_with_receivers( - receivers: HashMap>, -) -> ProcedureContext { - ProcedureContext { - procedure_id: ProcedureId::random(), - provider: Arc::new(ProcedureStateReceiverProvider { - receivers, - inner: MockContextProvider::default(), - }), - event_context: None, - } -} - -struct ProcedureStateReceiverProvider { - receivers: HashMap>, - inner: MockContextProvider, -} - -#[async_trait::async_trait] -impl ContextProvider for ProcedureStateReceiverProvider { - async fn procedure_state( - &self, - procedure_id: ProcedureId, - ) -> common_procedure::Result> { - self.inner.procedure_state(procedure_id).await - } - - async fn procedure_state_receiver( - &self, - procedure_id: ProcedureId, - ) -> common_procedure::Result>> { - Ok(self.receivers.get(&procedure_id).cloned()) - } - - async fn try_put_poison( - &self, - key: &common_procedure::PoisonKey, - procedure_id: ProcedureId, - ) -> common_procedure::Result<()> { - self.inner.try_put_poison(key, procedure_id).await - } - - async fn acquire_lock( - &self, - key: &common_procedure::StringKey, - ) -> common_procedure::local::DynamicKeyLockGuard { - self.inner.acquire_lock(key).await - } -} diff --git a/src/meta-srv/src/procedure/test_util.rs b/src/meta-srv/src/procedure/test_util.rs index 4208c6e52e7..385768d033e 100644 --- a/src/meta-srv/src/procedure/test_util.rs +++ b/src/meta-srv/src/procedure/test_util.rs @@ -13,6 +13,13 @@ // limitations under the License. use std::collections::HashMap; +#[cfg(test)] +use std::sync::Arc; + +#[cfg(test)] +use common_procedure::{Context, ContextProvider, ProcedureId, ProcedureState}; +#[cfg(test)] +use common_procedure_test::MockContextProvider; use api::v1::meta::mailbox_message::Payload; use api::v1::meta::{HeartbeatResponse, MailboxMessage}; @@ -412,3 +419,67 @@ pub async fn update_in_memory_region_flushed_entry_id( Ok(()) } + +/// Returns a receiver whose initial result is visible to [`common_procedure::watcher::wait`]. +#[cfg(test)] +pub fn procedure_state_receiver( + state: ProcedureState, +) -> tokio::sync::watch::Receiver { + let (tx, rx) = tokio::sync::watch::channel(ProcedureState::Running); + tx.send(state).unwrap(); + rx +} + +/// Returns an execution context that resolves child state receivers by procedure id. +#[cfg(test)] +pub fn procedure_context_with_receivers( + receivers: HashMap>, +) -> Context { + Context { + procedure_id: ProcedureId::random(), + provider: Arc::new(ProcedureStateReceiverProvider { + receivers, + inner: MockContextProvider::default(), + }), + event_context: None, + } +} + +#[cfg(test)] +struct ProcedureStateReceiverProvider { + receivers: HashMap>, + inner: MockContextProvider, +} + +#[cfg(test)] +#[async_trait::async_trait] +impl ContextProvider for ProcedureStateReceiverProvider { + async fn procedure_state( + &self, + procedure_id: ProcedureId, + ) -> common_procedure::Result> { + self.inner.procedure_state(procedure_id).await + } + + async fn procedure_state_receiver( + &self, + procedure_id: ProcedureId, + ) -> common_procedure::Result>> { + Ok(self.receivers.get(&procedure_id).cloned()) + } + + async fn try_put_poison( + &self, + key: &common_procedure::PoisonKey, + procedure_id: ProcedureId, + ) -> common_procedure::Result<()> { + self.inner.try_put_poison(key, procedure_id).await + } + + async fn acquire_lock( + &self, + key: &common_procedure::StringKey, + ) -> common_procedure::local::DynamicKeyLockGuard { + self.inner.acquire_lock(key).await + } +} From c7ea80558f9cf53a8feaa4fb6fdcdd07060f4ddf Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Thu, 29 Oct 2026 02:58:37 -0700 Subject: [PATCH 138/239] fix(mito2): preserve live index files across remote compactions (#10046) fix(mito2): reuse executor-owned compaction index resources Signed-off-by: Lei, HUANG --- src/mito2/src/compaction/compactor.rs | 230 +++++++++++++++----------- src/mito2/src/engine/open_test.rs | 141 +++++++++++++++- src/mito2/src/engine/region_hook.rs | 4 +- 3 files changed, 266 insertions(+), 109 deletions(-) diff --git a/src/mito2/src/compaction/compactor.rs b/src/mito2/src/compaction/compactor.rs index 1673721ec9b..4c89343b61c 100644 --- a/src/mito2/src/compaction/compactor.rs +++ b/src/mito2/src/compaction/compactor.rs @@ -201,28 +201,26 @@ pub struct OpenCompactionRegionRequest { pub plugins: Plugins, } -/// Open a compaction region from a compaction request. -/// It's simple version of RegionOpener::open(). -pub async fn open_compaction_region( - req: &OpenCompactionRegionRequest, - mito_config: &MitoConfig, +/// Opens manifest-backed compaction regions using executor-owned index resources. +/// +/// Create one opener per auxiliary directory before accepting jobs and reuse it +/// for the executor's lifetime. Initialization recovers staging files and removes +/// stale intermediate files, so repeating it can delete another job's live files. +#[derive(Clone)] +pub struct CompactionRegionOpener { + mito_config: Arc, object_store_manager: ObjectStoreManagerRef, - ttl_provider: Either, -) -> Result { - let object_store = { - let name = &req.region_options.storage; - if let Some(name) = name { - object_store_manager - .find(name) - .with_context(|| ObjectStoreNotFoundSnafu { - object_store: name.clone(), - })? - } else { - object_store_manager.default_object_store() - } - }; + puffin_manager_factory: PuffinManagerFactory, + intermediate_manager: IntermediateManager, +} - let access_layer = { +impl CompactionRegionOpener { + /// Initializes the executor's index resources. The config must be sanitized + /// and its auxiliary directory must not be owned by another engine or opener. + pub async fn new( + mito_config: MitoConfig, + object_store_manager: ObjectStoreManagerRef, + ) -> Result { let puffin_manager_factory = PuffinManagerFactory::new( &mito_config.index.aux_path, mito_config.index.staging_size.as_bytes(), @@ -231,95 +229,131 @@ pub async fn open_compaction_region( ) .await?; let intermediate_manager = - IntermediateManager::init_fs(mito_config.index.aux_path.clone()).await?; + IntermediateManager::init_fs(&mito_config.index.aux_path).await?; - Arc::new(AccessLayer::new( + Ok(Self { + mito_config: Arc::new(mito_config), + object_store_manager, + puffin_manager_factory, + intermediate_manager, + }) + } + + /// Opens a region without reinitializing or cleaning shared index resources. + pub async fn open( + &self, + req: &OpenCompactionRegionRequest, + ttl_provider: Either, + ) -> Result { + let mito_config = &self.mito_config; + let object_store = { + let name = &req.region_options.storage; + if let Some(name) = name { + self.object_store_manager + .find(name) + .with_context(|| ObjectStoreNotFoundSnafu { + object_store: name.clone(), + })? + } else { + self.object_store_manager.default_object_store() + } + }; + + let access_layer = Arc::new(AccessLayer::new( &req.table_dir, req.path_type, object_store.clone(), - puffin_manager_factory, - intermediate_manager, - )) - }; + self.puffin_manager_factory.clone(), + self.intermediate_manager.clone(), + )); - let manifest_manager = { - let region_dir = region_dir_from_table_dir(&req.table_dir, req.region_id, req.path_type); - let region_manifest_options = - RegionManifestOptions::new(mito_config, ®ion_dir, object_store); + let manifest_manager = { + let region_dir = + region_dir_from_table_dir(&req.table_dir, req.region_id, req.path_type); + let region_manifest_options = + RegionManifestOptions::new(mito_config, ®ion_dir, object_store); - RegionManifestManager::open(region_manifest_options, &Default::default()) - .await? - .with_context(|| EmptyRegionDirSnafu { - region_id: req.region_id, - region_dir: region_dir_from_table_dir(&req.table_dir, req.region_id, req.path_type), - })? - }; + RegionManifestManager::open(region_manifest_options, &Default::default()) + .await? + .with_context(|| EmptyRegionDirSnafu { + region_id: req.region_id, + region_dir: region_dir_from_table_dir( + &req.table_dir, + req.region_id, + req.path_type, + ), + })? + }; - let manifest = manifest_manager.manifest(); - let region_metadata = manifest.metadata.clone(); - let hook: Option = req.plugins.get(); - let manifest_ctx = Arc::new(ManifestContext::new( - manifest_manager, - RegionRoleState::Leader(RegionLeaderState::Writable), - hook, - )); + let manifest = manifest_manager.manifest(); + let region_metadata = manifest.metadata.clone(); + let hook: Option = req.plugins.get(); + let manifest_ctx = Arc::new(ManifestContext::new( + manifest_manager, + RegionRoleState::Leader(RegionLeaderState::Writable), + hook, + )); - let file_purger = { - let purge_scheduler = Arc::new(LocalScheduler::new(mito_config.max_background_purges)); - Arc::new(LocalFilePurger::new( - purge_scheduler.clone(), - access_layer.clone(), - None, - )) - }; + let file_purger = { + let purge_scheduler = Arc::new(LocalScheduler::new(mito_config.max_background_purges)); + Arc::new(LocalFilePurger::new( + purge_scheduler.clone(), + access_layer.clone(), + None, + )) + }; - let current_version = { - let mut ssts = SstVersion::new(region_metadata.clone()); - ssts.add_files(file_purger.clone(), manifest.files.values().cloned()); - CompactionVersion { - metadata: region_metadata.clone(), - options: req.region_options.clone(), - ssts: Arc::new(ssts), - // Remote execution uses an already picked output. A manifest-only - // opener cannot prove that a new LastNonNull pick is memtable-safe. - memtable_min_sequence: Some(0), - compaction_time_window: manifest.compaction_time_window, - } - }; + let current_version = { + let mut ssts = SstVersion::new(region_metadata.clone()); + ssts.add_files(file_purger.clone(), manifest.files.values().cloned()); + CompactionVersion { + metadata: region_metadata.clone(), + options: req.region_options.clone(), + ssts: Arc::new(ssts), + // Remote execution uses an already picked output. A manifest-only + // opener cannot prove that a new LastNonNull pick is memtable-safe. + memtable_min_sequence: Some(0), + compaction_time_window: manifest.compaction_time_window, + } + }; - let ttl = match ttl_provider { - // Use the specified ttl. - Either::Left(ttl) => ttl, - // Get the ttl from the schema metadata manager. - Either::Right(schema_metadata_manager) => { - let (_, ttl) = - find_dynamic_options(req.region_id, &req.region_options, &schema_metadata_manager) - .await - .unwrap_or_else(|e| { - warn!(e; "Failed to get ttl for region: {}", region_metadata.region_id); - ( - crate::region::options::CompactionOptions::default(), - TimeToLive::default(), - ) - }); - ttl - } - }; + let ttl = match ttl_provider { + // Use the specified ttl. + Either::Left(ttl) => ttl, + // Get the ttl from the schema metadata manager. + Either::Right(schema_metadata_manager) => { + let (_, ttl) = find_dynamic_options( + req.region_id, + &req.region_options, + &schema_metadata_manager, + ) + .await + .unwrap_or_else(|e| { + warn!(e; "Failed to get ttl for region: {}", region_metadata.region_id); + ( + crate::region::options::CompactionOptions::default(), + TimeToLive::default(), + ) + }); + ttl + } + }; - Ok(CompactionRegion { - region_id: req.region_id, - region_options: req.region_options.clone(), - engine_config: Arc::new(mito_config.clone()), - region_metadata: region_metadata.clone(), - cache_manager: Arc::new(CacheManager::default()), - access_layer, - manifest_ctx, - current_version, - file_purger: Some(file_purger), - ttl: Some(ttl), - max_parallelism: req.max_parallelism, - plugins: req.plugins.clone(), - }) + Ok(CompactionRegion { + region_id: req.region_id, + region_options: req.region_options.clone(), + engine_config: self.mito_config.clone(), + region_metadata: region_metadata.clone(), + cache_manager: Arc::new(CacheManager::default()), + access_layer, + manifest_ctx, + current_version, + file_purger: Some(file_purger), + ttl: Some(ttl), + max_parallelism: req.max_parallelism, + plugins: req.plugins.clone(), + }) + } } impl CompactionRegion { diff --git a/src/mito2/src/engine/open_test.rs b/src/mito2/src/engine/open_test.rs index 126dde56697..13915dc1ef7 100644 --- a/src/mito2/src/engine/open_test.rs +++ b/src/mito2/src/engine/open_test.rs @@ -33,7 +33,7 @@ use store_api::region_request::{ use store_api::storage::{RegionId, ScanRequest}; use tokio::sync::{Notify, oneshot}; -use crate::compaction::compactor::{OpenCompactionRegionRequest, open_compaction_region}; +use crate::compaction::compactor::{CompactionRegionOpener, OpenCompactionRegionRequest}; use crate::config::MitoConfig; use crate::engine::flush_test::MockRegionHook; use crate::engine::listener::EventListener; @@ -1037,16 +1037,139 @@ async fn test_open_compaction_region_with_format(flat_format: bool) { plugins: Plugins::new(), }; - let compaction_region = open_compaction_region( - &req, - &mito_config, - object_store_manager.clone(), - Either::Right(schema_metadata_manager), - ) - .await - .unwrap(); + mito_config.index.aux_path = env + .data_home() + .join("compactor-index") + .display() + .to_string(); + let opener = CompactionRegionOpener::new(mito_config, object_store_manager) + .await + .unwrap(); + let compaction_region = opener + .open(&req, Either::Right(schema_metadata_manager)) + .await + .unwrap(); assert_eq!(region_id, compaction_region.region_id); + compaction_region.stop_purger_scheduler().await.unwrap(); +} + +// Regression: per-job initialization used to delete another compaction's live +// atomic-write, external-sort, and staging files under the shared aux_path. +#[tokio::test] +async fn test_compaction_region_opens_preserve_live_index_files() { + use common_time::TimeToLive; + use futures::{AsyncReadExt, AsyncWriteExt}; + use index::external_provider::ExternalTempFileProvider; + use store_api::storage::FileId; + + use crate::sst::index::intermediate::{IntermediateLocation, TempFileProvider}; + + let mut env = TestEnv::new().await; + let engine = env.create_engine(MitoConfig::default()).await; + let mut requests = Vec::new(); + for (index, path_type) in [PathType::Bare, PathType::Data].into_iter().enumerate() { + let region_id = RegionId::new(100 + index as u32, 0); + let mut create = CreateRequestBuilder::new() + .table_dir(&format!("data/{}/", region_id.table_id())) + .build(); + create.path_type = path_type; + engine + .handle_request(region_id, RegionRequest::Create(create.clone())) + .await + .unwrap(); + put_rows( + &engine, + region_id, + Rows { + schema: rows_schema(&create), + rows: build_rows(0, 3), + }, + ) + .await; + flush_region(&engine, region_id, None).await; + requests.push(OpenCompactionRegionRequest { + region_id, + table_dir: create.table_dir, + path_type, + region_options: RegionOptions::default(), + max_parallelism: 1, + plugins: Plugins::new(), + }); + } + + let mut config = MitoConfig::default(); + config.sanitize(env.data_home().to_str().unwrap()).unwrap(); + let aux_path = env.data_home().join("compactor-index"); + config.index.aux_path = aux_path.display().to_string(); + let opener = CompactionRegionOpener::new(config, env.get_object_store_manager().unwrap()) + .await + .unwrap(); + let first = opener + .open(&requests[0], Either::Left(TimeToLive::default())) + .await + .unwrap(); + let provider = TempFileProvider::new( + IntermediateLocation::new(&first.region_id, &FileId::random()), + first.access_layer.intermediate_manager().clone(), + ); + let mut committed = provider.create("sort", "committed").await.unwrap(); + committed.write_all(b"committed sort input").await.unwrap(); + committed.close().await.unwrap(); + let mut pending = provider.create("sort", "pending").await.unwrap(); + pending.write_all(b"pending sort input").await.unwrap(); + pending.flush().await.unwrap(); + let staging_file = aux_path.join("staging/inflight.tmp"); + tokio::fs::write(&staging_file, b"staging input") + .await + .unwrap(); + + let cloned_opener = opener.clone(); + let (second, reopened) = tokio::try_join!( + opener.open(&requests[1], Either::Left(TimeToLive::default())), + cloned_opener.open(&requests[0], Either::Left(TimeToLive::default())), + ) + .unwrap(); + assert_eq!(second.access_layer.path_type(), PathType::Data); + assert_eq!(reopened.access_layer.path_type(), PathType::Bare); + let missing = OpenCompactionRegionRequest { + region_id: RegionId::new(999, 0), + table_dir: "missing/".into(), + ..requests[0].clone() + }; + assert!(matches!( + opener + .open(&missing, Either::Left(TimeToLive::default())) + .await, + Err(error::Error::EmptyRegionDir { .. }) + )); + drop(cloned_opener); + drop(opener); + + // Open regions retain the resources even after the opener itself is dropped. + pending.close().await.unwrap(); + let mut contents = Vec::new(); + for (_, mut reader) in provider.read_all("sort").await.unwrap() { + let mut bytes = Vec::new(); + reader.read_to_end(&mut bytes).await.unwrap(); + contents.push(bytes); + } + contents.sort(); + assert_eq!( + contents, + vec![ + b"committed sort input".to_vec(), + b"pending sort input".to_vec() + ] + ); + assert_eq!( + tokio::fs::read(&staging_file).await.unwrap(), + b"staging input" + ); + provider.cleanup().await.unwrap(); + for region in [first, second, reopened] { + region.stop_purger_scheduler().await.unwrap(); + } } /// Returns a dummy fetcher that returns the expr_json only once. diff --git a/src/mito2/src/engine/region_hook.rs b/src/mito2/src/engine/region_hook.rs index 90a15c6aa60..c8b736d2894 100644 --- a/src/mito2/src/engine/region_hook.rs +++ b/src/mito2/src/engine/region_hook.rs @@ -62,7 +62,7 @@ //! | Operations during staging | N/A | ❌ (staging dir) | //! | Async index build | ❌ (index files only) | ✅ Yes | //! -//! ¹ Remote compaction runs on a dedicated compactor node via `open_compaction_region()`; +//! ¹ Remote compaction runs on a dedicated compactor node via `CompactionRegionOpener::open()`; //! pass plugins via `OpenCompactionRegionRequest` to enable hooks there. `sst_infos` is //! `#[serde(skip)]` over the wire, so each [`SstInfo`] is rebuilt from [`FileMeta`] with //! empty footer/index — see [`SstFileInfo`] for field-level detail. @@ -328,7 +328,7 @@ pub trait RegionHook: Send + Sync + Debug { /// Fires once when a region becomes active via a create or open request — /// the natural counterpart to [`on_region_closed`] / [`on_region_dropped`]. /// It does **not** fire for the compactor's transient compaction regions - /// (`open_compaction_region`), nor for the internal reopen performed during + /// (`CompactionRegionOpener::open`), nor for the internal reopen performed during /// follower catch-up / leadership promotion. /// /// On the **create** path it runs inline in the region worker loop; on the From ee55c0c2faccb899e81d3feb2cfdd9a690c55149 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Thu, 29 Oct 2026 10:03:35 +0000 Subject: [PATCH 139/239] feat: share concurrent SST metadata loads (#9729) Signed-off-by: WenyXu Co-authored-by: Lei, HUANG <6406592+v0y4g3r@users.noreply.github.com> --- Cargo.lock | 1 + Cargo.toml | 1 + src/mito2/Cargo.toml | 1 + src/mito2/src/access_layer.rs | 29 +- src/mito2/src/cache.rs | 879 +++++++++++++++--- src/mito2/src/cache/test_util.rs | 33 +- src/mito2/src/compaction/reader.rs | 3 + src/mito2/src/error.rs | 36 +- src/mito2/src/metrics.rs | 50 +- src/mito2/src/read/pruner.rs | 18 +- src/mito2/src/read/range.rs | 35 +- src/mito2/src/read/scan_region.rs | 13 +- src/mito2/src/read/unordered_scan.rs | 9 +- src/mito2/src/region/opener.rs | 8 +- src/mito2/src/sst/file.rs | 4 +- .../src/sst/index/inverted_index/creator.rs | 3 +- src/mito2/src/sst/parquet/reader.rs | 334 ++++++- 17 files changed, 1265 insertions(+), 192 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index d40ca8b5baa..59c8421e955 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8798,6 +8798,7 @@ dependencies = [ "bytes", "chrono", "common-base", + "common-cache", "common-datasource", "common-error", "common-function", diff --git a/Cargo.toml b/Cargo.toml index a5e3db37960..1f84eb7dbcb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -301,6 +301,7 @@ client = { path = "src/client" } cmd = { path = "src/cmd", default-features = false } common-base = { path = "src/common/base" } common-batcher = { path = "src/common/batcher" } +common-cache = { path = "src/common/cache" } common-catalog = { path = "src/common/catalog" } common-config = { path = "src/common/config" } common-datasource = { path = "src/common/datasource" } diff --git a/src/mito2/Cargo.toml b/src/mito2/Cargo.toml index a5e26fcd60f..411cf643145 100644 --- a/src/mito2/Cargo.toml +++ b/src/mito2/Cargo.toml @@ -28,6 +28,7 @@ bytemuck.workspace = true bytes.workspace = true chrono.workspace = true common-base.workspace = true +common-cache.workspace = true common-datasource.workspace = true common-error.workspace = true common-function.workspace = true diff --git a/src/mito2/src/access_layer.rs b/src/mito2/src/access_layer.rs index 3067c6098a1..7a80687958b 100644 --- a/src/mito2/src/access_layer.rs +++ b/src/mito2/src/access_layer.rs @@ -446,26 +446,29 @@ impl AccessLayer { let parquet_metadata = parquet_metadata.clone(); let region_metadata = region_metadata.clone(); let cache_manager = cache_manager.clone(); - // Compact cache preparation is best-effort. Run the entire operation in one - // detached blocking task so it neither blocks an async worker nor delays the - // SST write. - runtime.spawn_blocking(move || { - match prepare_sst_meta_sync( - &file_path, - Arc::unwrap_or_clone(parquet_metadata), - Some(region_metadata), - page_index_policy, - ) { + // Prepare and cache metadata in the background without delaying the SST write. + let task_runtime = runtime.clone(); + runtime.spawn(async move { + let preparation = task_runtime.spawn_blocking(move || { + prepare_sst_meta_sync( + &file_path, + Arc::unwrap_or_clone(parquet_metadata), + Some(region_metadata), + page_index_policy, + ) + }).await.context(crate::error::JoinSnafu).and_then(|result| result); + match preparation { Ok(SstMetaPreparation::Prepared(metadata)) => { - cache_manager.put_prepared_sst_meta(file_id, metadata, true); + cache_manager.put_compact_sst_meta(file_id, metadata.compact); + cache_manager.put_sst_meta_data(file_id, metadata.decoded).await; } Ok(SstMetaPreparation::DecodedOnly { encoding_error, .. }) => warn!( encoding_error; "Failed to encode parquet metadata for cache, file: {}", - file_path + file_id.file_id() ), Err(err) => { - warn!(err; "Failed to cache parquet metadata for {}", file_path); + warn!(err; "Failed to cache parquet metadata for {}", file_id.file_id()); } } }); diff --git a/src/mito2/src/cache.rs b/src/mito2/src/cache.rs index c4b8ffa0d73..4cff42043ce 100644 --- a/src/mito2/src/cache.rs +++ b/src/mito2/src/cache.rs @@ -241,8 +241,8 @@ pub(crate) struct CompactSstMeta { /// Both forms produced by decoding metadata after a cache miss. #[derive(Debug)] pub(crate) struct PreparedSstMeta { - compact: Arc, - decoded: Arc, + pub(crate) compact: Arc, + pub(crate) decoded: Arc, } impl PreparedSstMeta { @@ -731,17 +731,96 @@ impl CacheStrategy { } } - /// Gets fused SST metadata with cache metrics tracking. - pub(crate) async fn get_sst_meta_data( + #[cfg(test)] + pub(crate) async fn get_or_load_sst_meta( &self, file_id: RegionFileId, + page_index_policy: PageIndexPolicy, + load: F, + ) -> Result<(Arc, bool)> + where + F: std::future::Future, bool)>>, + { + self.get_or_load_observed_sst_meta(file_id, page_index_policy, async { + (load.await, SstMetaResolution::default()) + }) + .await + .0 + } + + /// Shares the resolved source on success and failure, without sharing I/O accounting. + pub(crate) async fn get_or_load_observed_sst_meta( + &self, + file_id: RegionFileId, + page_index_policy: PageIndexPolicy, + load: F, + ) -> (Result<(Arc, bool)>, SstMetaResolution) + where + F: std::future::Future, bool)>, SstMetaResolution)>, + { + let cache = match self { + Self::EnableAll(manager) | Self::Compaction(manager) => { + manager.sst_decoded_meta_cache.as_ref() + } + Self::Disabled => None, + }; + let Some(cache) = cache else { + return load.await; + }; + let key = SstDecodedMetaKey::new(file_id, page_index_policy); + let result = cache + .try_get_with(&key, async { + let (result, resolution) = load.await; + let (metadata, cache_miss) = result.map_err(|source| SstMetaLoadError { + source: Arc::new(source), + resolution, + })?; + let value = SstMetaLoadResult::new(metadata, cache_miss, resolution); + Ok::<_, SstMetaLoadError>(value) + }) + .await; + match result { + Ok(value) => ( + Ok((value.metadata.clone(), value.cache_miss)), + value.resolution, + ), + Err(error) => ( + Err(error.source.clone()).context(crate::error::LoadSstMetaSnafu), + error.resolution, + ), + } + } + + /// Observes memory residency for this caller before sharing decode or I/O work. + pub(crate) async fn observe_sst_meta( + &self, + file_id: RegionFileId, + page_index_policy: PageIndexPolicy, + metrics: &mut MetadataCacheMetrics, + ) -> SstMetaObservation { + match self { + Self::EnableAll(manager) | Self::Compaction(manager) => { + manager + .observe_sst_meta(file_id, page_index_policy, metrics) + .await + } + Self::Disabled => SstMetaObservation::Miss, + } + } + + /// Resolves an observed metadata entry and updates this caller's cache metrics. + pub(crate) async fn resolve_sst_meta( + &self, + observation: SstMetaObservation, + file_id: RegionFileId, metrics: &mut MetadataCacheMetrics, page_index_policy: PageIndexPolicy, ) -> Option> { match self { - CacheStrategy::EnableAll(cache_manager) | CacheStrategy::Compaction(cache_manager) => { - cache_manager - .get_sst_meta_data( + Self::EnableAll(manager) | Self::Compaction(manager) => { + manager + .resolve_sst_meta( + observation, file_id, metrics, page_index_policy, @@ -749,7 +828,7 @@ impl CacheStrategy { ) .await } - CacheStrategy::Disabled => { + Self::Disabled => { metrics.cache_miss += 1; None } @@ -757,45 +836,47 @@ impl CacheStrategy { } /// Calls [CacheManager::get_sst_meta_data_from_mem_cache()]. - pub(crate) fn get_sst_meta_data_from_mem_cache( + pub(crate) async fn get_sst_meta_data_from_mem_cache( &self, file_id: RegionFileId, page_index_policy: PageIndexPolicy, ) -> Option> { match self { CacheStrategy::EnableAll(cache_manager) | CacheStrategy::Compaction(cache_manager) => { - cache_manager.get_sst_meta_data_from_mem_cache(file_id, page_index_policy) + cache_manager + .get_sst_meta_data_from_mem_cache(file_id, page_index_policy) + .await } CacheStrategy::Disabled => None, } } /// Calls [CacheManager::get_parquet_meta_data_from_mem_cache()]. - pub fn get_parquet_meta_data_from_mem_cache( + pub async fn get_parquet_meta_data_from_mem_cache( &self, file_id: RegionFileId, ) -> Option> { self.get_sst_meta_data_from_mem_cache(file_id, PageIndexPolicy::Skip) + .await .map(|metadata| metadata.parquet_metadata()) } - /// Puts compact and decoded forms of SST metadata into the shared cache capacity. - pub(crate) fn put_prepared_sst_meta( + /// Calls [CacheManager::put_compact_sst_meta()]. + pub(crate) fn put_compact_sst_meta( &self, file_id: RegionFileId, - metadata: PreparedSstMeta, - retain_decoded: bool, + metadata: Arc, ) { match self { CacheStrategy::EnableAll(cache_manager) | CacheStrategy::Compaction(cache_manager) => { - cache_manager.put_prepared_sst_meta(file_id, metadata, retain_decoded); + cache_manager.put_compact_sst_meta(file_id, metadata); } CacheStrategy::Disabled => {} } } /// Calls [CacheManager::put_parquet_meta_data()]. - pub fn put_parquet_meta_data( + pub async fn put_parquet_meta_data( &self, file_id: RegionFileId, metadata: Arc, @@ -803,7 +884,9 @@ impl CacheStrategy { ) { match self { CacheStrategy::EnableAll(cache_manager) | CacheStrategy::Compaction(cache_manager) => { - cache_manager.put_parquet_meta_data(file_id, metadata, region_metadata); + cache_manager + .put_parquet_meta_data(file_id, metadata, region_metadata) + .await; } CacheStrategy::Disabled => {} } @@ -827,13 +910,13 @@ impl CacheStrategy { } /// Calls [CacheManager::remove_parquet_meta_data()]. - pub fn remove_parquet_meta_data(&self, file_id: RegionFileId) { + pub async fn remove_parquet_meta_data(&self, file_id: RegionFileId) { match self { CacheStrategy::EnableAll(cache_manager) => { - cache_manager.remove_parquet_meta_data(file_id); + cache_manager.remove_parquet_meta_data(file_id).await; } CacheStrategy::Compaction(cache_manager) => { - cache_manager.remove_parquet_meta_data(file_id); + cache_manager.remove_parquet_meta_data(file_id).await; } CacheStrategy::Disabled => {} } @@ -1100,8 +1183,8 @@ impl CacheManager { CacheManagerBuilder::default() } - /// Gets fused SST metadata with metrics tracking. - /// Tries in-memory cache first, then file cache, updating metrics accordingly. + /// Gets fused SST metadata with metrics tracking for direct cache tests. + #[cfg(test)] pub(crate) async fn get_sst_meta_data( &self, file_id: RegionFileId, @@ -1109,42 +1192,54 @@ impl CacheManager { page_index_policy: PageIndexPolicy, runtime: &Runtime, ) -> Option> { + let observation = self + .observe_sst_meta(file_id, page_index_policy, metrics) + .await; + self.resolve_sst_meta(observation, file_id, metrics, page_index_policy, runtime) + .await + } + + async fn observe_sst_meta( + &self, + file_id: RegionFileId, + page_index_policy: PageIndexPolicy, + metrics: &mut MetadataCacheMetrics, + ) -> SstMetaObservation { let cache_key = SstMetaKey(file_id.region_id(), file_id.file_id()); let compact = self .get_compact_sst_meta(&cache_key) .filter(|metadata| metadata.satisfies_page_index_policy(page_index_policy)); let decoded = self - .get_decoded_sst_meta(&cache_key) + .get_decoded_sst_meta(file_id, page_index_policy) + .await .filter(|metadata| metadata.satisfies_page_index_policy(page_index_policy)); + let observation = match (compact, decoded) { + (_, Some(decoded)) => SstMetaObservation::Decoded(decoded), + (Some(compact), None) => SstMetaObservation::Compact(compact), + (None, None) => SstMetaObservation::Miss, + }; + observation.record(metrics); + observation + } - if let Some(compact) = compact { - CACHE_HIT.with_label_values(&[SST_META_TYPE]).inc(); - metrics.mem_cache_hit += 1; - if let Some(decoded) = decoded { - CACHE_HIT.with_label_values(&[SST_META_DECODED_TYPE]).inc(); - return Some(decoded); - } - - CACHE_MISS.with_label_values(&[SST_META_DECODED_TYPE]).inc(); - match compact.decode(runtime).await { - Ok(decoded) => { - self.put_sst_meta_data(file_id, decoded.clone()); - return Some(decoded); - } + async fn resolve_sst_meta( + &self, + observation: SstMetaObservation, + file_id: RegionFileId, + metrics: &mut MetadataCacheMetrics, + page_index_policy: PageIndexPolicy, + runtime: &Runtime, + ) -> Option> { + match observation { + SstMetaObservation::Decoded(decoded) => return Some(decoded), + SstMetaObservation::Compact(compact) => match compact.decode(runtime).await { + Ok(decoded) => return Some(decoded), Err(err) => { warn!(err; "Failed to decode compact SST metadata, region_id: {}, file_id: {}", file_id.region_id(), file_id.file_id()); - self.remove_parquet_meta_data(file_id); + self.remove_parquet_meta_data(file_id).await; } - } - } else if let Some(decoded) = decoded { - // Metadata produced by a new SST writer can enter the decoded tier before its compact - // representation is prepared. - CACHE_HIT.with_label_values(&[SST_META_TYPE]).inc(); - CACHE_HIT.with_label_values(&[SST_META_DECODED_TYPE]).inc(); - metrics.mem_cache_hit += 1; - return Some(decoded); - } else { - CACHE_MISS.with_label_values(&[SST_META_TYPE]).inc(); + }, + SstMetaObservation::Miss => {} } let key = IndexKey::new(file_id.region_id(), file_id.file_id(), FileType::Parquet); @@ -1159,7 +1254,7 @@ impl CacheManager { let decoded = metadata.decoded(); match metadata { SstMetaPreparation::Prepared(metadata) => { - self.put_prepared_sst_meta(file_id, metadata, true); + self.put_compact_sst_meta(file_id, metadata.compact); } SstMetaPreparation::DecodedOnly { encoding_error, .. } => { warn!( @@ -1187,55 +1282,61 @@ impl CacheManager { /// Gets cached fused SST metadata from in-memory cache. /// This method does not perform I/O. - pub(crate) fn get_sst_meta_data_from_mem_cache( + pub(crate) async fn get_sst_meta_data_from_mem_cache( &self, file_id: RegionFileId, page_index_policy: PageIndexPolicy, ) -> Option> { - let key = SstMetaKey(file_id.region_id(), file_id.file_id()); let value = self - .get_decoded_sst_meta(&key) - .filter(|metadata| metadata.satisfies_page_index_policy(page_index_policy)); + .get_sst_meta_data_from_mem_cache_inner(file_id, page_index_policy) + .await; update_hit_miss(value, SST_META_DECODED_TYPE) } + async fn get_sst_meta_data_from_mem_cache_inner( + &self, + file_id: RegionFileId, + page_index_policy: PageIndexPolicy, + ) -> Option> { + self.get_decoded_sst_meta(file_id, page_index_policy).await + } + /// Gets cached [ParquetMetaData] from in-memory cache. /// This method does not perform I/O. - pub fn get_parquet_meta_data_from_mem_cache( + pub async fn get_parquet_meta_data_from_mem_cache( &self, file_id: RegionFileId, ) -> Option> { self.get_sst_meta_data_from_mem_cache(file_id, PageIndexPolicy::Skip) + .await .map(|metadata| metadata.parquet_metadata()) } /// Puts fused SST metadata into the cache. - pub(crate) fn put_sst_meta_data(&self, file_id: RegionFileId, metadata: Arc) { + pub(crate) async fn put_sst_meta_data( + &self, + file_id: RegionFileId, + metadata: Arc, + ) { if let Some(cache) = &self.sst_decoded_meta_cache { - let key = SstMetaKey(file_id.region_id(), file_id.file_id()); - CACHE_BYTES - .with_label_values(&[SST_META_DECODED_TYPE]) - .add(decoded_meta_cache_weight(&key, &metadata).into()); - cache.insert(key, metadata); + let key = SstDecodedMetaKey::new(file_id, metadata.page_index_policy); + let value = SstMetaLoadResult::new(metadata, false, SstMetaResolution::default()); + cache.insert(key, value).await; } } - /// Puts a compact metadata entry and optionally retains its decoded acceleration entry. - pub(crate) fn put_prepared_sst_meta( + /// Puts a compact SST metadata entry into the cache. + pub(crate) fn put_compact_sst_meta( &self, file_id: RegionFileId, - metadata: PreparedSstMeta, - retain_decoded: bool, + metadata: Arc, ) { let key = SstMetaKey(file_id.region_id(), file_id.file_id()); if let Some(cache) = &self.sst_meta_cache { CACHE_BYTES .with_label_values(&[SST_META_TYPE]) - .add(meta_cache_weight(&key, &metadata.compact).into()); - cache.insert(key.clone(), metadata.compact); - } - if retain_decoded { - self.put_sst_meta_data(file_id, metadata.decoded); + .add(meta_cache_weight(&key, &metadata).into()); + cache.insert(key, metadata); } } @@ -1243,8 +1344,26 @@ impl CacheManager { self.sst_meta_cache.as_ref()?.get(key) } - fn get_decoded_sst_meta(&self, key: &SstMetaKey) -> Option> { - self.sst_decoded_meta_cache.as_ref()?.get(key) + async fn get_decoded_sst_meta( + &self, + file_id: RegionFileId, + requested: PageIndexPolicy, + ) -> Option> { + let cache = self.sst_decoded_meta_cache.as_ref()?; + for policy in [ + PageIndexPolicy::Skip, + PageIndexPolicy::Optional, + PageIndexPolicy::Required, + ] { + // The key describes the loading request, while compact metadata may + // already contain stronger page-index guarantees. + if let Some(value) = cache.get(&SstDecodedMetaKey::new(file_id, policy)).await + && value.metadata.satisfies_page_index_policy(requested) + { + return Some(value.metadata.clone()); + } + } + None } /// Gets and decodes only the compact tier without promoting into the decoded cache. @@ -1264,14 +1383,14 @@ impl CacheManager { Ok(metadata) => Some(metadata), Err(err) => { warn!(err; "Failed to decode compact SST metadata, region_id: {}, file_id: {}", file_id.region_id(), file_id.file_id()); - self.remove_parquet_meta_data(file_id); + self.remove_parquet_meta_data(file_id).await; None } } } /// Puts [ParquetMetaData] into the cache. - pub fn put_parquet_meta_data( + pub async fn put_parquet_meta_data( &self, file_id: RegionFileId, metadata: Arc, @@ -1288,7 +1407,7 @@ impl CacheManager { Arc::unwrap_or_clone(metadata), region_metadata, ) { - Ok(metadata) => self.put_sst_meta_data(file_id, Arc::new(metadata)), + Ok(metadata) => self.put_sst_meta_data(file_id, Arc::new(metadata)).await, Err(err) => warn!( err; "Failed to decode region metadata while caching parquet metadata, region_id: {}, file_id: {}", file_id.region_id(), @@ -1299,13 +1418,21 @@ impl CacheManager { } /// Removes [ParquetMetaData] from the cache. - pub fn remove_parquet_meta_data(&self, file_id: RegionFileId) { + pub async fn remove_parquet_meta_data(&self, file_id: RegionFileId) { let key = SstMetaKey(file_id.region_id(), file_id.file_id()); if let Some(cache) = &self.sst_meta_cache { cache.remove(&key); } if let Some(cache) = &self.sst_decoded_meta_cache { - cache.remove(&key); + for policy in [ + PageIndexPolicy::Skip, + PageIndexPolicy::Optional, + PageIndexPolicy::Required, + ] { + cache + .invalidate(&SstDecodedMetaKey::new(file_id, policy)) + .await; + } } } @@ -1636,19 +1763,18 @@ impl CacheManagerBuilder { .build() }); let sst_decoded_meta_cache = (decoded_meta_capacity != 0).then(|| { - Cache::builder() - .max_capacity(decoded_meta_capacity) - .weigher(decoded_meta_cache_weight) - .eviction_listener(|k, v, cause| { - let size = decoded_meta_cache_weight(&k, &v); - CACHE_BYTES - .with_label_values(&[SST_META_DECODED_TYPE]) - .sub(size.into()); + common_cache::Cache::new( + "sst_decoded_metadata", + decoded_meta_capacity, + decoded_meta_cache_weight, + |_k, _v, cause| { CACHE_EVICTION - .with_label_values(&[SST_META_DECODED_TYPE, removal_cause_str(cause)]) + .with_label_values(&[SST_META_DECODED_TYPE, cause.as_str()]) .inc(); - }) - .build() + }, + ) + .with_size_metric(CACHE_BYTES.bind(&[SST_META_DECODED_TYPE])) + .with_metrics(crate::metrics::cache_loading_metrics(SST_META_DECODED_TYPE)) }); let vector_cache = (self.vector_cache_size != 0).then(|| { Cache::builder() @@ -1683,7 +1809,7 @@ impl CacheManagerBuilder { let prefilter_result_cache = (self.prefilter_result_cache_size != 0) .then(|| new_prefilter_result_cache(self.prefilter_result_cache_size)); let puffin_metadata_cache = - PuffinMetadataCache::new(self.puffin_metadata_size, &CACHE_BYTES); + PuffinMetadataCache::new(self.puffin_metadata_size, CACHE_BYTES.gauges()); let selector_result_cache = (self.selector_result_cache_size != 0).then(|| { Cache::builder() .max_capacity(self.selector_result_cache_size) @@ -1730,9 +1856,8 @@ fn meta_cache_weight(k: &SstMetaKey, v: &Arc) -> u32 { u32::try_from(size).unwrap_or(u32::MAX) } -fn decoded_meta_cache_weight(k: &SstMetaKey, v: &Arc) -> u32 { - let size = k.estimated_size() + v.parquet_metadata_size + v.region_metadata_weight; - u32::try_from(size).unwrap_or(u32::MAX) +fn decoded_meta_cache_weight(k: &SstDecodedMetaKey, v: &Arc) -> u32 { + SstMetaLoadResult::weight(k, &v.metadata) } fn vector_cache_weight(_k: &(ConcreteDataType, Value), v: &VectorRef) -> u32 { @@ -2119,10 +2244,96 @@ impl SelectorResultValue { } } -/// Maps (region id, file id) to fused SST metadata. +/// A caller's memory observation. Holding the entry avoids a second lookup after coordination. +pub(crate) enum SstMetaObservation { + Decoded(Arc), + Compact(Arc), + Miss, +} + +impl SstMetaObservation { + fn record(&self, metrics: &mut MetadataCacheMetrics) { + match self { + Self::Decoded(_) | Self::Compact(_) => { + metrics.mem_cache_hit += 1; + CACHE_HIT.with_label_values(&[SST_META_TYPE]).inc(); + match self { + Self::Decoded(_) => CACHE_HIT.with_label_values(&[SST_META_DECODED_TYPE]).inc(), + _ => CACHE_MISS.with_label_values(&[SST_META_DECODED_TYPE]).inc(), + } + } + Self::Miss => CACHE_MISS.with_label_values(&[SST_META_TYPE]).inc(), + } + } +} + +/// Resolved cache source, not physical I/O. Also returned when the subsequent load fails. +/// A caller that observed a memory miss can find a completed value at the native recheck; +/// it reuses that value's source without executing its initializer. This does not prove +/// that the caller waited for an in-flight load or performed a file-cache lookup. +/// Ordinary memory hits return before this resolution is applied. +/// A default value means no file-cache resolution was needed (or known). +#[derive(Clone, Copy, Default)] +pub(crate) struct SstMetaResolution { + pub(crate) file_cache_hit: usize, + pub(crate) cache_miss: usize, +} + +struct SstMetaLoadError { + source: Arc, + resolution: SstMetaResolution, +} + +/// Metadata returned by a shared load. Cache hits ignore the original load's miss flag. +struct SstMetaLoadResult { + metadata: Arc, + /// Preserves background-download eligibility for all requests sharing a storage load. + cache_miss: bool, + resolution: SstMetaResolution, +} + +impl SstMetaLoadResult { + fn weight(key: &SstDecodedMetaKey, metadata: &CachedSstMeta) -> u32 { + let size = size_of_val(key) + + size_of::() + + metadata.parquet_metadata_size + + metadata.region_metadata_weight; + u32::try_from(size).unwrap_or(u32::MAX) + } + + fn new( + metadata: Arc, + cache_miss: bool, + resolution: SstMetaResolution, + ) -> Arc { + Arc::new(Self { + metadata, + cache_miss, + resolution, + }) + } +} + +/// Separates loads that require different page-index guarantees. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +struct SstDecodedMetaKey { + file: SstMetaKey, + policy: mem::Discriminant, +} + +impl SstDecodedMetaKey { + fn new(file: RegionFileId, policy: PageIndexPolicy) -> Self { + Self { + file: SstMetaKey(file.region_id(), file.file_id()), + policy: mem::discriminant(&policy), + } + } +} + +/// Maps (region id, file id) to compact SST metadata. type SstMetaCache = Cache>; /// Maps (region id, file id) to decoded SST metadata retained as an acceleration tier. -type SstDecodedMetaCache = Cache>; +type SstDecodedMetaCache = common_cache::Cache>; /// Maps [Value] to a vector that holds this value repeatedly. /// /// e.g. `"hello" => ["hello", "hello", "hello"]` @@ -2169,7 +2380,7 @@ mod tests { let file_id = RegionFileId::new(region_id, FileId::random()); let metadata = parquet_meta(); let mut metrics = MetadataCacheMetrics::default(); - cache.put_parquet_meta_data(file_id, metadata, None); + cache.put_parquet_meta_data(file_id, metadata, None).await; assert!( cache .get_sst_meta_data( @@ -2191,21 +2402,365 @@ mod tests { .is_none() ); - cache.put_page_ranges( - file_id.file_id(), - 1, - &[Range { start: 0, end: 5 }], - &[Bytes::from_static(b"abcde")], - ); - assert!( - cache - .get_page_ranges(file_id.file_id(), 1, &[Range { start: 0, end: 5 }]) - .is_none() - ); + assert!(cache.page_cache.is_none()); assert!(cache.write_cache().is_none()); } + #[tokio::test] + async fn test_sst_meta_load_coalescing() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + let manager = Arc::new(CacheManager::builder().sst_meta_cache_size(2).build()); + let strategy = CacheStrategy::EnableAll(manager.clone()); + let (metadata, region) = sst_parquet_meta(); + let metadata = Arc::new( + CachedSstMeta::try_new("test.parquet", Arc::unwrap_or_clone(metadata)).unwrap(), + ); + let file_id = RegionFileId::new(region.region_id, FileId::random()); + let calls = AtomicUsize::new(0); + let gate = tokio::sync::Notify::new(); + let load = || async { + calls.fetch_add(1, Ordering::Relaxed); + gate.notified().await; + Ok((metadata.clone(), true)) + }; + let mut first = + Box::pin(strategy.get_or_load_sst_meta(file_id, PageIndexPolicy::Skip, load())); + let mut second = + Box::pin(strategy.get_or_load_sst_meta(file_id, PageIndexPolicy::Skip, load())); + assert!(futures::poll!(&mut first).is_pending()); + assert!(futures::poll!(&mut second).is_pending()); + assert_eq!(1, calls.load(Ordering::Relaxed)); + gate.notify_waiters(); + let (first, second) = tokio::join!(first, second); + let first = first.unwrap(); + let second = second.unwrap(); + assert!(Arc::ptr_eq(&first.0, &second.0)); + assert!( + first.1 && second.1, + "all waiters preserve the remote miss flag" + ); + + // Moka applies capacity limits during maintenance, not synchronously on insertion. + manager + .sst_decoded_meta_cache + .as_ref() + .unwrap() + .run_pending_tasks() + .await; + // An entry larger than the decoded reservation must not be retained. + let mut next = + Box::pin(strategy.get_or_load_sst_meta(file_id, PageIndexPolicy::Skip, load())); + assert!(futures::poll!(&mut next).is_pending()); + assert_eq!(2, calls.load(Ordering::Relaxed)); + drop(next); + manager + .sst_decoded_meta_cache + .as_ref() + .unwrap() + .run_pending_tasks() + .await; + assert_eq!( + 0, + manager + .sst_decoded_meta_cache + .as_ref() + .unwrap() + .entry_count() + ); + } + + #[tokio::test] + async fn test_compact_sst_meta_stronger_policy_reused_from_skip_key() { + for policy in [PageIndexPolicy::Optional, PageIndexPolicy::Required] { + let manager = Arc::new( + CacheManager::builder() + .sst_meta_cache_size(1024 * 1024) + .build(), + ); + let strategy = CacheStrategy::EnableAll(manager.clone()); + let (metadata, region) = sst_parquet_meta(); + let file_id = RegionFileId::new(region.region_id, FileId::random()); + let prepared = prepare_sst_meta( + "test.parquet", + with_offset_indexes(Arc::unwrap_or_clone(metadata)), + None, + policy, + &common_runtime::global_runtime(), + ) + .await + .unwrap(); + let SstMetaPreparation::Prepared(prepared) = prepared else { + panic!("valid metadata should produce a compact entry"); + }; + manager.put_compact_sst_meta(file_id, prepared.compact); + let mut metrics = MetadataCacheMetrics::default(); + let observation = strategy + .observe_sst_meta(file_id, PageIndexPolicy::Skip, &mut metrics) + .await; + assert!(matches!(&observation, SstMetaObservation::Compact(_))); + let (decoded, _) = strategy + .get_or_load_sst_meta(file_id, PageIndexPolicy::Skip, async { + let value = strategy + .resolve_sst_meta( + observation, + file_id, + &mut MetadataCacheMetrics::default(), + PageIndexPolicy::Skip, + ) + .await + .expect("compact entry must decode without storage I/O"); + Ok((value, false)) + }) + .await + .unwrap(); + assert!(decoded.satisfies_page_index_policy(policy)); + let cache = manager.sst_decoded_meta_cache.as_ref().unwrap(); + cache.run_pending_tasks().await; + let weight = cache.weighted_size(); + assert_eq!(cache.entry_count(), 1); + + // Losing the compact copy must not force another decode or storage load. + manager + .sst_meta_cache + .as_ref() + .unwrap() + .invalidate(&SstMetaKey(file_id.region_id(), file_id.file_id())); + for requested in [PageIndexPolicy::Optional, policy] { + let mut metrics = MetadataCacheMetrics::default(); + let observation = strategy + .observe_sst_meta(file_id, requested, &mut metrics) + .await; + let SstMetaObservation::Decoded(cached) = observation else { + panic!("the completed Skip-key entry already satisfies this request"); + }; + assert!(Arc::ptr_eq(&decoded, &cached)); + assert_eq!(metrics.mem_cache_hit, 1); + assert_eq!(metrics.cache_miss, 0); + } + cache.run_pending_tasks().await; + assert_eq!(cache.entry_count(), 1); + assert_eq!(cache.weighted_size(), weight); + assert!( + cache + .get(&SstDecodedMetaKey::new(file_id, PageIndexPolicy::Optional)) + .await + .is_none() + ); + assert!( + cache + .get(&SstDecodedMetaKey::new(file_id, PageIndexPolicy::Required)) + .await + .is_none() + ); + } + } + + #[tokio::test] + async fn test_async_sst_meta_policy_reuse_and_invalidation() { + let manager = Arc::new( + CacheManager::builder() + .sst_meta_cache_size(1024 * 1024) + .build(), + ); + let strategy = CacheStrategy::EnableAll(manager.clone()); + let (parquet_metadata, region) = sst_parquet_meta(); + let file_id = RegionFileId::new(region.region_id, FileId::random()); + let policies = [ + PageIndexPolicy::Skip, + PageIndexPolicy::Optional, + PageIndexPolicy::Required, + ]; + for loaded in policies { + let metadata = Arc::new( + CachedSstMeta::try_new_with_page_index_policy( + "test.parquet", + (*parquet_metadata).clone(), + None, + loaded, + ) + .unwrap(), + ); + let (result, miss) = strategy + .get_or_load_sst_meta(file_id, loaded, async { Ok((metadata.clone(), true)) }) + .await + .unwrap(); + assert!(Arc::ptr_eq(&metadata, &result)); + assert!(miss); + for requested in policies { + let mut metrics = MetadataCacheMetrics::default(); + let cached = strategy + .observe_sst_meta(file_id, requested, &mut metrics) + .await; + let cached = match cached { + SstMetaObservation::Decoded(value) => Some(value), + _ => None, + }; + assert_eq!( + satisfies_page_index_policy(loaded, requested), + cached.is_some() + ); + if let Some(cached) = cached { + assert!(Arc::ptr_eq(&metadata, &cached)); + assert_eq!(1, metrics.mem_cache_hit); + } + } + manager.remove_parquet_meta_data(file_id).await; + } + + // Populate all policy keys before removing the file, not just the last inserted key. + for loaded in policies { + let metadata = Arc::new( + CachedSstMeta::try_new_with_page_index_policy( + "test.parquet", + (*parquet_metadata).clone(), + None, + loaded, + ) + .unwrap(), + ); + manager.put_sst_meta_data(file_id, metadata).await; + } + manager.remove_parquet_meta_data(file_id).await; + for requested in policies { + assert!( + manager + .get_sst_meta_data_from_mem_cache(file_id, requested) + .await + .is_none() + ); + } + } + + #[tokio::test] + async fn test_sst_meta_load_cancellation_and_policy_isolation() { + let strategy = CacheStrategy::EnableAll(Arc::new( + CacheManager::builder().sst_meta_cache_size(2000).build(), + )); + let (metadata, region) = sst_parquet_meta(); + let metadata = Arc::new( + CachedSstMeta::try_new("test.parquet", Arc::unwrap_or_clone(metadata)).unwrap(), + ); + let file_id = RegionFileId::new(region.region_id, FileId::random()); + let mut first = Box::pin(strategy.get_or_load_sst_meta( + file_id, + PageIndexPolicy::Skip, + std::future::pending(), + )); + assert!(futures::poll!(&mut first).is_pending()); + // A different policy or file does not wait for this load. + for (id, policy) in [ + (file_id, PageIndexPolicy::Optional), + (file_id, PageIndexPolicy::Required), + ( + RegionFileId::new(region.region_id, FileId::random()), + PageIndexPolicy::Skip, + ), + ] { + let mut independent = Box::pin( + strategy.get_or_load_sst_meta(id, policy, async { Ok((metadata.clone(), false)) }), + ); + assert!(futures::poll!(&mut independent).is_ready()); + } + let mut waiter = Box::pin(strategy.get_or_load_sst_meta( + file_id, + PageIndexPolicy::Skip, + async { Ok((metadata.clone(), false)) }, + )); + assert!(futures::poll!(&mut waiter).is_pending()); + drop(first); + let result = tokio::time::timeout(std::time::Duration::from_secs(5), waiter) + .await + .unwrap() + .unwrap(); + assert!(Arc::ptr_eq(&metadata, &result.0)); + } + + #[tokio::test] + async fn test_sst_meta_load_failure_retry_and_disabled_cache() { + use common_error::ext::ErrorExt; + + let manager = Arc::new(CacheManager::builder().sst_meta_cache_size(2000).build()); + let strategy = CacheStrategy::EnableAll(manager); + let (metadata, region) = sst_parquet_meta(); + let metadata = Arc::new( + CachedSstMeta::try_new("test.parquet", Arc::unwrap_or_clone(metadata)).unwrap(), + ); + let file_id = RegionFileId::new(region.region_id, FileId::random()); + let gate = tokio::sync::Notify::new(); + let mut first = Box::pin(strategy.get_or_load_observed_sst_meta( + file_id, + PageIndexPolicy::Skip, + async { + gate.notified().await; + ( + UnexpectedSnafu { + reason: "load failed", + } + .fail(), + SstMetaResolution { + file_cache_hit: 0, + cache_miss: 1, + }, + ) + }, + )); + let mut waiter = Box::pin(strategy.get_or_load_observed_sst_meta( + file_id, + PageIndexPolicy::Skip, + async { panic!("waiter must share the failure") }, + )); + assert!(futures::poll!(&mut first).is_pending()); + assert!(futures::poll!(&mut waiter).is_pending()); + gate.notify_waiters(); + let (first, first_resolution) = first.await; + let (second, second_resolution) = waiter.await; + for resolution in [first_resolution, second_resolution] { + assert_eq!(1, resolution.cache_miss); + assert_eq!(0, resolution.file_cache_hit); + } + let first = first.unwrap_err(); + let second = second.unwrap_err(); + assert_eq!(first.status_code(), second.status_code()); + match (first, second) { + ( + crate::error::Error::LoadSstMeta { source: a, .. }, + crate::error::Error::LoadSstMeta { source: b, .. }, + ) => assert!(Arc::ptr_eq(&a, &b)), + errors => panic!("unexpected errors: {errors:?}"), + } + assert!( + strategy + .get_or_load_sst_meta(file_id, PageIndexPolicy::Skip, async { + Ok((metadata.clone(), true)) + }) + .await + .is_ok() + ); + + for strategy in [ + CacheStrategy::Disabled, + CacheStrategy::EnableAll(Arc::new(CacheManager::default())), + CacheStrategy::EnableAll(Arc::new( + CacheManager::builder().sst_meta_cache_size(1).build(), + )), + ] { + let mut blocked = Box::pin(strategy.get_or_load_sst_meta( + file_id, + PageIndexPolicy::Skip, + std::future::pending(), + )); + assert!(futures::poll!(&mut blocked).is_pending()); + let mut independent = Box::pin(strategy.get_or_load_sst_meta( + file_id, + PageIndexPolicy::Skip, + async { Ok((metadata.clone(), true)) }, + )); + assert!(futures::poll!(&mut independent).is_ready()); + } + } + #[test] fn test_sst_meta_cache_splits_capacity() { let cache = CacheManager::builder().sst_meta_cache_size(101).build(); @@ -2220,13 +2775,8 @@ mod tests { .max_capacity() ); assert_eq!( - Some(50), - cache - .sst_decoded_meta_cache - .as_ref() - .unwrap() - .policy() - .max_capacity() + 50, + cache.sst_decoded_meta_cache.as_ref().unwrap().capacity() ); } @@ -2248,7 +2798,7 @@ mod tests { .is_none() ); let (metadata, region_metadata) = sst_parquet_meta(); - cache.put_parquet_meta_data(file_id, metadata, None); + cache.put_parquet_meta_data(file_id, metadata, None).await; let cached = cache .get_sst_meta_data( file_id, @@ -2270,7 +2820,7 @@ mod tests { .all(|key_value| key_value.key != PARQUET_METADATA_KEY) }) ); - cache.remove_parquet_meta_data(file_id); + cache.remove_parquet_meta_data(file_id).await; assert!( cache .get_sst_meta_data( @@ -2292,7 +2842,9 @@ mod tests { let file_id = RegionFileId::new(region_id, FileId::random()); let (metadata, region_metadata) = sst_parquet_meta(); - cache.put_parquet_meta_data(file_id, metadata, Some(region_metadata.clone())); + cache + .put_parquet_meta_data(file_id, metadata, Some(region_metadata.clone())) + .await; let cached = cache .get_sst_meta_data( @@ -2306,15 +2858,7 @@ mod tests { assert!(Arc::ptr_eq(®ion_metadata, &cached.region_metadata())); } - #[tokio::test] - async fn test_compact_sst_meta_round_trip() { - let cache = CacheManager::builder() - .sst_meta_cache_size(1024 * 1024) - .build(); - let region_id = RegionId::new(1, 1); - let file_id = RegionFileId::new(region_id, FileId::random()); - let (metadata, expected_region_metadata) = sst_parquet_meta(); - let metadata = Arc::unwrap_or_clone(metadata); + fn with_offset_indexes(metadata: ParquetMetaData) -> ParquetMetaData { let offset_indexes = metadata .row_groups() .iter() @@ -2333,10 +2877,22 @@ mod tests { .collect() }) .collect(); - let metadata = metadata + metadata .into_builder() .set_offset_index(Some(offset_indexes)) + .build() + } + + #[tokio::test] + async fn test_compact_sst_meta_round_trip() { + let cache = CacheManager::builder() + .sst_meta_cache_size(1024 * 1024) .build(); + let region_id = RegionId::new(1, 1); + let file_id = RegionFileId::new(region_id, FileId::random()); + let (metadata, expected_region_metadata) = sst_parquet_meta(); + let metadata = Arc::unwrap_or_clone(metadata); + let metadata = with_offset_indexes(metadata); let expected_rows = metadata.file_metadata().num_rows(); let prepared = prepare_sst_meta( "test.parquet", @@ -2352,7 +2908,7 @@ mod tests { }; // Retain only the compact form so the lookup must exercise decompression and decoding. - cache.put_prepared_sst_meta(file_id, prepared, false); + cache.put_compact_sst_meta(file_id, prepared.compact); let mut metrics = MetadataCacheMetrics::default(); let cached = cache .get_sst_meta_data( @@ -2371,6 +2927,16 @@ mod tests { cached.parquet_metadata.file_metadata().num_rows() ); assert!(cached.parquet_metadata.offset_index().is_some()); + assert!( + cache + .sst_decoded_meta_cache + .as_ref() + .unwrap() + .get(&SstDecodedMetaKey::new(file_id, PageIndexPolicy::Required)) + .await + .is_none(), + "compact inspection must not fill the decoded reservation" + ); assert_eq!(expected_region_metadata, cached.region_metadata()); assert!( cached @@ -2383,8 +2949,8 @@ mod tests { ); } - #[test] - fn test_cache_encoding_failure_preserves_decoded_metadata() { + #[tokio::test] + async fn test_cache_encoding_failure_preserves_decoded_metadata() { let (metadata, expected_region_metadata) = sst_parquet_meta(); let expected_rows = metadata.file_metadata().num_rows(); let cache_encoding: Result<(Bytes, usize)> = Err(UnexpectedSnafu { @@ -2418,6 +2984,36 @@ mod tests { .to_string() .contains("injected compact metadata encoding failure") ); + let file_id = RegionFileId::new(expected_region_metadata.region_id, FileId::random()); + let manager = Arc::new( + CacheManager::builder() + .sst_meta_cache_size(1024 * 1024) + .build(), + ); + let strategy = CacheStrategy::EnableAll(manager.clone()); + let (shared, miss) = strategy + .get_or_load_sst_meta(file_id, PageIndexPolicy::Skip, async { + Ok((decoded.clone(), true)) + }) + .await + .unwrap(); + assert!(miss); + assert!(Arc::ptr_eq(&decoded, &shared)); + assert!( + manager + .get_compact_sst_meta(&SstMetaKey(file_id.region_id(), file_id.file_id())) + .is_none() + ); + let mut metrics = MetadataCacheMetrics::default(); + let SstMetaObservation::Decoded(cached) = strategy + .observe_sst_meta(file_id, PageIndexPolicy::Skip, &mut metrics) + .await + else { + panic!("decoded fallback must be retained without a compact entry"); + }; + assert!(Arc::ptr_eq(&decoded, &cached)); + assert_eq!(1, metrics.mem_cache_hit); + assert_eq!(0, metrics.cache_miss); } #[tokio::test] @@ -2436,7 +3032,7 @@ mod tests { ) .unwrap(), ); - cache.put_sst_meta_data(file_id, skip_metadata); + cache.put_sst_meta_data(file_id, skip_metadata).await; let mut metrics = MetadataCacheMetrics::default(); assert!( @@ -2461,7 +3057,7 @@ mod tests { ) .unwrap(), ); - cache.put_sst_meta_data(file_id, optional_metadata); + cache.put_sst_meta_data(file_id, optional_metadata).await; let mut metrics = MetadataCacheMetrics::default(); assert!( @@ -2500,12 +3096,19 @@ mod tests { let cached = Arc::new( CachedSstMeta::try_new("test.parquet", Arc::unwrap_or_clone(metadata)).unwrap(), ); - let key = SstMetaKey(region_metadata.region_id, FileId::random()); + let key = SstDecodedMetaKey::new( + RegionFileId::new(region_metadata.region_id, FileId::random()), + PageIndexPolicy::Skip, + ); + let value = SstMetaLoadResult::new(cached.clone(), false, SstMetaResolution::default()); assert!(cached.region_metadata_weight > json_len); assert_eq!( - decoded_meta_cache_weight(&key, &cached) as usize, - key.estimated_size() + cached.parquet_metadata_size + cached.region_metadata_weight + decoded_meta_cache_weight(&key, &value) as usize, + size_of_val(&key) + + size_of::() + + cached.parquet_metadata_size + + cached.region_metadata_weight ); assert_eq!( cached.parquet_metadata_size, @@ -2521,9 +3124,13 @@ mod tests { CachedSstMeta::try_new("test.parquet", Arc::unwrap_or_clone(metadata)).unwrap(); cached.region_metadata_weight = u32::MAX as usize + 1; let cached = Arc::new(cached); - let key = SstMetaKey(region_metadata.region_id, FileId::random()); + let key = SstDecodedMetaKey::new( + RegionFileId::new(region_metadata.region_id, FileId::random()), + PageIndexPolicy::Skip, + ); + let value = SstMetaLoadResult::new(cached.clone(), false, SstMetaResolution::default()); - assert_eq!(u32::MAX, decoded_meta_cache_weight(&key, &cached)); + assert_eq!(u32::MAX, decoded_meta_cache_weight(&key, &value)); } #[test] diff --git a/src/mito2/src/cache/test_util.rs b/src/mito2/src/cache/test_util.rs index 82cabe24297..13c711b35a6 100644 --- a/src/mito2/src/cache/test_util.rs +++ b/src/mito2/src/cache/test_util.rs @@ -15,11 +15,13 @@ //! Utilities for testing cache. use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; use bytes::Bytes; use datatypes::arrow::array::{ArrayRef, Int64Array}; use datatypes::arrow::record_batch::RecordBatch; use object_store::ObjectStore; +use object_store::layers::mock; use object_store::services::Fs; use parquet::arrow::ArrowWriter; use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; @@ -31,6 +33,33 @@ use store_api::metadata::RegionMetadataRef; use crate::sst::parquet::PARQUET_METADATA_KEY; use crate::test_util::sst_util::sst_region_metadata; +/// Holds storage reads until every concurrent caller has been polled. +pub(crate) struct GatedStorageReader { + pub(crate) inner: mock::oio::Reader, + pub(crate) gate: Arc, + pub(crate) reads: Arc, +} + +impl mock::oio::Read for GatedStorageReader { + async fn open( + &self, + range: mock::BytesRange, + ) -> object_store::Result<(mock::RpRead, Box)> { + self.reads.fetch_add(1, Ordering::Relaxed); + self.gate.acquire().await.unwrap().forget(); + self.inner.open(range).await + } + + async fn read( + &self, + range: mock::BytesRange, + ) -> object_store::Result<(mock::RpRead, object_store::Buffer)> { + self.reads.fetch_add(1, Ordering::Relaxed); + self.gate.acquire().await.unwrap().forget(); + self.inner.read(range).await + } +} + /// Returns a parquet meta data. pub(crate) fn parquet_meta() -> Arc { let file_data = parquet_file_data(); @@ -60,7 +89,9 @@ fn parquet_file_data() -> Vec { parquet_file_data_inner(None) } -fn parquet_file_data_with_region_metadata(region_metadata: &RegionMetadataRef) -> Vec { +pub(crate) fn parquet_file_data_with_region_metadata( + region_metadata: &RegionMetadataRef, +) -> Vec { let json = region_metadata.to_json().unwrap(); let key_value = KeyValue::new(PARQUET_METADATA_KEY.to_string(), json); parquet_file_data_inner(Some(vec![key_value])) diff --git a/src/mito2/src/compaction/reader.rs b/src/mito2/src/compaction/reader.rs index 0968f0e7a60..0a7f30ee34f 100644 --- a/src/mito2/src/compaction/reader.rs +++ b/src/mito2/src/compaction/reader.rs @@ -21,6 +21,7 @@ use common_time::range::TimestampRange; use common_time::timestamp::TimeUnit; use datafusion_common::ScalarValue; use datafusion_expr::Expr; +use futures::FutureExt; use parquet::file::metadata::{PageIndexPolicy, ParquetMetaData}; use snafu::OptionExt; use store_api::metadata::RegionMetadataRef; @@ -67,8 +68,10 @@ impl CompactionSstReaderBuilder<'_> { let schema = scan_input.mapper.output_schema(); let schema = rewrite_json2_schema(schema.arrow_schema(), &plans); + // Erase the nested scan future before embedding it in the compaction task. let stream = SeqScan::new(scan_input) .build_flat_reader_for_compaction() + .boxed() .await?; Ok(FlatSource::new_stream(schema, stream)) } diff --git a/src/mito2/src/error.rs b/src/mito2/src/error.rs index a336fb89673..91402711136 100644 --- a/src/mito2/src/error.rs +++ b/src/mito2/src/error.rs @@ -198,6 +198,13 @@ pub enum Error { location: Location, }, + #[snafu(display("Failed to load SST metadata"))] + LoadSstMeta { + source: Arc, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Failed to write parquet file"))] WriteParquet { #[snafu(source)] @@ -1455,6 +1462,7 @@ impl Error { match self { Error::OpenDal { error, .. } => error.kind() == ErrorKind::NotFound, Error::ManifestDeltaNotFound { .. } => true, + Error::LoadSstMeta { source, .. } => source.is_object_not_found(), _ => false, } } @@ -1577,7 +1585,7 @@ impl ErrorExt for Error { OpenRegion { source, .. } => source.status_code(), WriteParquet { .. } => StatusCode::StorageUnavailable, - WriteGroup { source, .. } => source.status_code(), + WriteGroup { source, .. } | LoadSstMeta { source, .. } => source.status_code(), InvalidBatch { .. } => StatusCode::InvalidArguments, InvalidRecordBatch { .. } => StatusCode::InvalidArguments, ConvertVector { source, .. } => source.status_code(), @@ -1732,6 +1740,7 @@ impl ErrorExt for Error { | External { source, .. } => source.retry_hint(), OpenRegion { source, .. } + | LoadSstMeta { source, .. } | WriteGroup { source, .. } | FlushRegion { source, .. } | BuildIndexAsync { source, .. } @@ -1861,6 +1870,31 @@ mod tests { } } + #[test] + fn test_shared_sst_metadata_preserves_error_semantics() { + for source in [ + OpenDalSnafu {} + .into_error(object_store::Error::new(ErrorKind::NotFound, "missing SST")), + OpenDalSnafu {}.into_error(object_store::Error::new( + ErrorKind::Unexpected, + "read failed", + )), + RejectWriteSnafu { + region_id: RegionId::new(1, 1), + } + .build(), + ] { + let expected_not_found = source.is_object_not_found(); + let expected_status = source.status_code(); + let expected_retry = source.retry_hint(); + let source = Arc::new(source); + let wrapped = LoadSstMetaSnafu.into_error(source); + assert_eq!(expected_not_found, wrapped.is_object_not_found()); + assert_eq!(expected_status, wrapped.status_code()); + assert_eq!(expected_retry, wrapped.retry_hint()); + } + } + #[test] fn test_manifest_update_persistence() { let rejected_kinds = [ diff --git a/src/mito2/src/metrics.rs b/src/mito2/src/metrics.rs index f7af9d54ffc..acb6a9da4c4 100644 --- a/src/mito2/src/metrics.rs +++ b/src/mito2/src/metrics.rs @@ -272,8 +272,44 @@ lazy_static! { ).unwrap(); } +pub(crate) fn cache_loading_metrics( + cache_type: &'static str, +) -> std::sync::Arc { + use common_cache::metrics::{CacheMetrics, OperationMetrics}; + + let operation = |stage| OperationMetrics { + started: CACHE_NATIVE_LOAD_EVENTS.with_label_values(&[cache_type, stage, "started"]), + succeeded: CACHE_NATIVE_LOAD_EVENTS.with_label_values(&[cache_type, stage, "success"]), + failed: CACHE_NATIVE_LOAD_EVENTS.with_label_values(&[cache_type, stage, "error"]), + cancelled: CACHE_NATIVE_LOAD_EVENTS.with_label_values(&[cache_type, stage, "cancelled"]), + active: CACHE_NATIVE_LOAD_ACTIVE.with_label_values(&[cache_type, stage]), + elapsed: CACHE_NATIVE_LOAD_ELAPSED.with_label_values(&[cache_type, stage]), + }; + std::sync::Arc::new(CacheMetrics { + requests: operation("request"), + loads: operation("loader"), + }) +} + // Cache metrics. lazy_static! { + /// Native cache request and actual loader lifecycle; request is not a waiter count. + pub static ref CACHE_NATIVE_LOAD_EVENTS: IntCounterVec = register_int_counter_vec!( + "greptime_mito_cache_native_load_events_total", + "Native cache request and executed loader outcomes; requests include hits and shared loads", + &[TYPE_LABEL, STAGE_LABEL, "event"] + ).unwrap(); + pub static ref CACHE_NATIVE_LOAD_ACTIVE: IntGaugeVec = register_int_gauge_vec!( + "greptime_mito_cache_native_load_active", + "Active native cache requests or executed loaders, not waiting requests", + &[TYPE_LABEL, STAGE_LABEL] + ).unwrap(); + pub static ref CACHE_NATIVE_LOAD_ELAPSED: HistogramVec = register_histogram_vec!( + "greptime_mito_cache_native_load_seconds", + "Native cache request or executed loader wall time including suspension and cancellation, not CPU time", + &[TYPE_LABEL, STAGE_LABEL], exponential_buckets(0.0001, 10.0, 6).unwrap() + ).unwrap(); + /// Cache hit counter. pub static ref CACHE_HIT: IntCounterVec = register_int_counter_vec!( "greptime_mito_cache_hit", @@ -289,12 +325,14 @@ lazy_static! { ) .unwrap(); /// Cache size in bytes. - pub static ref CACHE_BYTES: IntGaugeVec = register_int_gauge_vec!( - "greptime_mito_cache_bytes", - "mito cache bytes", - &[TYPE_LABEL] - ) - .unwrap(); + pub static ref CACHE_BYTES: common_cache::metrics::CacheSizeMetrics = { + let metrics = common_cache::metrics::CacheSizeMetrics::new( + prometheus::Opts::new("greptime_mito_cache_bytes", "mito cache bytes; backend cache weights are approximate"), + &[TYPE_LABEL], + ).unwrap(); + prometheus::register(Box::new(metrics.clone())).unwrap(); + metrics + }; /// Download bytes counter in the write cache. pub static ref WRITE_CACHE_DOWNLOAD_BYTES_TOTAL: IntCounter = register_int_counter!( "mito_write_cache_download_bytes_total", diff --git a/src/mito2/src/read/pruner.rs b/src/mito2/src/read/pruner.rs index ef09f5f48ab..47689db161c 100644 --- a/src/mito2/src/read/pruner.rs +++ b/src/mito2/src/read/pruner.rs @@ -810,7 +810,7 @@ mod tests { use super::*; use crate::read::flat_projection::FlatProjectionMapper; - use crate::read::range::RowGroupIndex; + use crate::read::range::{RangeMeta, RowGroupIndex}; use crate::read::scan_region::{PredicateGroup, ScanInput}; use crate::read::scan_util::PartitionMetrics; use crate::sst::file::{FileHandle, FileMeta}; @@ -853,7 +853,13 @@ mod tests { .with_files(files) .with_append_mode(true) .build(); - let stream_ctx = Arc::new(StreamContext::unordered_scan_ctx(input)); + let query_start = input.query_start.unwrap_or_else(Instant::now); + let ranges = RangeMeta::unordered_scan_ranges(&input).await; + let stream_ctx = Arc::new(StreamContext::unordered_scan_ctx( + input, + ranges, + query_start, + )); let pruner = Arc::new(Pruner::new_with_options( stream_ctx, 1, @@ -912,7 +918,13 @@ mod tests { .with_predicate(predicate) .with_append_mode(true) .build(); - let stream_ctx = Arc::new(StreamContext::unordered_scan_ctx(input)); + let query_start = input.query_start.unwrap_or_else(Instant::now); + let ranges = RangeMeta::unordered_scan_ranges(&input).await; + let stream_ctx = Arc::new(StreamContext::unordered_scan_ctx( + input, + ranges, + query_start, + )); let pruner = Arc::new(Pruner::new(stream_ctx, 1)); (env, pruner) } diff --git a/src/mito2/src/read/range.rs b/src/mito2/src/read/range.rs index bf47e045cad..5e6da0a5986 100644 --- a/src/mito2/src/read/range.rs +++ b/src/mito2/src/read/range.rs @@ -15,6 +15,7 @@ //! Structs for partition ranges. use common_time::Timestamp; +use futures::future::BoxFuture; use smallvec::{SmallVec, smallvec}; use store_api::region_engine::PartitionRange; use store_api::storage::TimeSeriesDistribution; @@ -105,20 +106,24 @@ impl RangeMeta { } /// Creates a list of ranges from the `input` for unordered scan. - pub(crate) fn unordered_scan_ranges(input: &ScanInput) -> Vec { - let mut ranges = Vec::with_capacity(input.memtables.len() + input.files.len()); - Self::push_unordered_mem_ranges(&input.memtables, &mut ranges); - Self::push_unordered_file_ranges( - input.memtables.len(), - &input.files, - &input.cache_strategy, - &mut ranges, - ); + pub(crate) fn unordered_scan_ranges(input: &ScanInput) -> BoxFuture<'_, Vec> { + // Keep Moka's async lookup type out of callers' nested scanner futures. + Box::pin(async move { + let mut ranges = Vec::with_capacity(input.memtables.len() + input.files.len()); + Self::push_unordered_mem_ranges(&input.memtables, &mut ranges); + Self::push_unordered_file_ranges( + input.memtables.len(), + &input.files, + &input.cache_strategy, + &mut ranges, + ) + .await; - #[cfg(feature = "enterprise")] - Self::push_extension_ranges(input, &mut ranges); + #[cfg(feature = "enterprise")] + Self::push_extension_ranges(input, &mut ranges); - ranges + ranges + }) } /// Returns true if the time range of given `meta` overlaps with the time range of this meta. @@ -203,7 +208,7 @@ impl RangeMeta { } } - fn push_unordered_file_ranges( + async fn push_unordered_file_ranges( num_memtables: usize, files: &[FileHandle], cache: &CacheStrategy, @@ -213,7 +218,9 @@ impl RangeMeta { for (i, file) in files.iter().enumerate() { let file_index = num_memtables + i; // Get parquet meta from the cache. - let parquet_meta = cache.get_parquet_meta_data_from_mem_cache(file.file_id()); + let parquet_meta = cache + .get_parquet_meta_data_from_mem_cache(file.file_id()) + .await; if let Some(parquet_meta) = parquet_meta { // Scans each row group. for row_group_index in 0..file.meta_ref().num_row_groups { diff --git a/src/mito2/src/read/scan_region.rs b/src/mito2/src/read/scan_region.rs index 31db8581b27..6dd6fe81c63 100644 --- a/src/mito2/src/read/scan_region.rs +++ b/src/mito2/src/read/scan_region.rs @@ -435,7 +435,10 @@ impl ScanRegion { #[tracing::instrument(skip_all, fields(region_id = %self.region_id()))] pub(crate) async fn unordered_scan(self) -> Result { let input = self.scan_input().await?; - Ok(UnorderedScan::new(input)) + let query_start = input.query_start.unwrap_or_else(Instant::now); + let ranges = RangeMeta::unordered_scan_ranges(&input).await; + let stream_ctx = StreamContext::unordered_scan_ctx(input, ranges, query_start); + Ok(UnorderedScan::new(stream_ctx)) } /// Scans by series. @@ -2041,9 +2044,11 @@ impl StreamContext { } /// Creates a new [StreamContext] for [UnorderedScan]. - pub(crate) fn unordered_scan_ctx(input: ScanInput) -> Self { - let query_start = input.query_start.unwrap_or_else(Instant::now); - let ranges = RangeMeta::unordered_scan_ranges(&input); + pub(crate) fn unordered_scan_ctx( + input: ScanInput, + ranges: Vec, + query_start: Instant, + ) -> Self { READ_SST_COUNT.observe(input.num_files() as f64); Self { input, diff --git a/src/mito2/src/read/unordered_scan.rs b/src/mito2/src/read/unordered_scan.rs index 05387be0ce1..bda554e4c95 100644 --- a/src/mito2/src/read/unordered_scan.rs +++ b/src/mito2/src/read/unordered_scan.rs @@ -35,7 +35,9 @@ use store_api::region_engine::{ use crate::error::{PartitionOutOfRangeSnafu, Result}; use crate::read::pruner::{PartitionPruner, Pruner}; -use crate::read::scan_region::{ScanInput, StreamContext}; +#[cfg(test)] +use crate::read::scan_region::ScanInput; +use crate::read::scan_region::StreamContext; use crate::read::scan_util::{ PartitionMetrics, PartitionMetricsList, prepare_flat_mem_stream, scan_flat_file_ranges, scan_flat_mem_ranges, @@ -59,7 +61,8 @@ pub struct UnorderedScan { impl UnorderedScan { /// Creates a new [UnorderedScan]. - pub(crate) fn new(input: ScanInput) -> Self { + pub(crate) fn new(stream_ctx: StreamContext) -> Self { + let input = &stream_ctx.input; let mut properties = ScannerProperties::default() .with_append_mode(input.append_mode) .with_total_rows(input.total_rows()) @@ -67,7 +70,7 @@ impl UnorderedScan { if let Some(counters) = input.query_stat_counters.clone() { properties.set_query_stat_counters(counters); } - let stream_ctx = Arc::new(StreamContext::unordered_scan_ctx(input)); + let stream_ctx = Arc::new(stream_ctx); properties.partitions = vec![stream_ctx.partition_ranges()]; // Create the shared pruner with number of workers equal to CPU cores. diff --git a/src/mito2/src/region/opener.rs b/src/mito2/src/region/opener.rs index ba4fe615951..30eb6b99cbc 100644 --- a/src/mito2/src/region/opener.rs +++ b/src/mito2/src/region/opener.rs @@ -1299,7 +1299,7 @@ async fn preload_parquet_meta_cache_for_files( } match metadata { SstMetaPreparation::Prepared(metadata) => { - cache_manager.put_prepared_sst_meta(file_id, metadata, false); + cache_manager.put_compact_sst_meta(file_id, metadata.compact); loaded += 1; } SstMetaPreparation::DecodedOnly { encoding_error, .. } => warn!( @@ -1342,7 +1342,7 @@ async fn preload_parquet_meta_cache_for_files( Ok(SstMetaPreparation::Prepared(metadata)) => { // Preload the compact canonical entry only. Query traffic promotes decoded // entries into the separate acceleration tier according to actual demand. - cache_manager.put_prepared_sst_meta(file_id, metadata, false); + cache_manager.put_compact_sst_meta(file_id, metadata.compact); loaded += 1; } Ok(SstMetaPreparation::DecodedOnly { encoding_error, .. }) => warn!( @@ -1801,6 +1801,7 @@ mod tests { assert!( cache_manager .get_parquet_meta_data_from_mem_cache(region_file_id) + .await .is_none() ); @@ -1877,6 +1878,7 @@ mod tests { assert!( cache_manager .get_parquet_meta_data_from_mem_cache(region_file_id) + .await .is_none() ); @@ -1896,6 +1898,7 @@ mod tests { assert!( cache_manager .get_parquet_meta_data_from_mem_cache(region_file_id) + .await .is_none() ); } @@ -1950,6 +1953,7 @@ mod tests { assert!( cache_manager .get_parquet_meta_data_from_mem_cache(region_file_id) + .await .is_none() ); diff --git a/src/mito2/src/sst/file.rs b/src/mito2/src/sst/file.rs index a627513d310..6422819607b 100644 --- a/src/mito2/src/sst/file.rs +++ b/src/mito2/src/sst/file.rs @@ -873,7 +873,9 @@ pub async fn delete_files( if let Some(cache) = &cache_manager { for (file_id, _) in file_ids { cache.invalidate_range_results(*file_id); - cache.remove_parquet_meta_data(RegionFileId::new(region_id, *file_id)); + cache + .remove_parquet_meta_data(RegionFileId::new(region_id, *file_id)) + .await; } } let mut attempted_files = Vec::with_capacity(file_ids.len()); diff --git a/src/mito2/src/sst/index/inverted_index/creator.rs b/src/mito2/src/sst/index/inverted_index/creator.rs index 19c55253b58..28242e79105 100644 --- a/src/mito2/src/sst/index/inverted_index/creator.rs +++ b/src/mito2/src/sst/index/inverted_index/creator.rs @@ -636,7 +636,8 @@ mod tests { move |expr| { let _d = &d; let cache = Arc::new(InvertedIndexCache::new(10, 10, 100)); - let puffin_metadata_cache = Arc::new(PuffinMetadataCache::new(10, &CACHE_BYTES)); + let puffin_metadata_cache = + Arc::new(PuffinMetadataCache::new(10, CACHE_BYTES.gauges())); let applier = InvertedIndexApplierBuilder::new( table_dir.clone(), PathType::Bare, diff --git a/src/mito2/src/sst/parquet/reader.rs b/src/mito2/src/sst/parquet/reader.rs index 1c98444946c..2ddb66a53f8 100644 --- a/src/mito2/src/sst/parquet/reader.rs +++ b/src/mito2/src/sst/parquet/reader.rs @@ -53,7 +53,10 @@ use store_api::storage::{ColumnId, FileId}; use table::predicate::Predicate; use crate::cache::index::result_cache::PredicateKey; -use crate::cache::{CacheStrategy, CachedSstMeta, SstMetaPreparation, prepare_sst_meta}; +use crate::cache::{ + CacheStrategy, CachedSstMeta, SstMetaObservation, SstMetaPreparation, SstMetaResolution, + prepare_sst_meta, +}; use crate::error::{ ParquetToArrowSchemaSnafu, ReadDataPartSnafu, Result, SerializePartitionExprSnafu, UnexpectedSnafu, @@ -822,14 +825,66 @@ impl ParquetReaderBuilder { .with_label_values(&["read_parquet_metadata"]) .start_timer(); + let observation = self + .cache_strategy + .observe_sst_meta(self.file_handle.file_id(), page_index_policy, cache_metrics) + .await; + if let SstMetaObservation::Decoded(metadata) = observation { + cache_metrics.metadata_load_cost += start.elapsed(); + return Ok((metadata, false)); + } + + let mut loader_metrics = MetadataCacheMetrics::default(); + // Keep the shared loader future off the stack of nested series-index builders. + // Decoded metadata hits above do not need this allocation. + let (result, resolution) = Box::pin(self.cache_strategy.get_or_load_observed_sst_meta( + self.file_handle.file_id(), + page_index_policy, + async { + let result = self + .read_parquet_metadata_inner( + file_path, + file_size, + &mut loader_metrics, + page_index_policy, + observation, + ) + .await; + ( + result, + SstMetaResolution { + file_cache_hit: loader_metrics.file_cache_hit, + cache_miss: loader_metrics.cache_miss, + }, + ) + }, + )) + .await; + // These are resolution counts, including a known miss before a failed remote read. + // Only the selected loader contributes actual I/O counts and bytes. + cache_metrics.file_cache_hit += resolution.file_cache_hit; + cache_metrics.cache_miss += resolution.cache_miss; + cache_metrics.num_reads += loader_metrics.num_reads; + cache_metrics.bytes_read += loader_metrics.bytes_read; + cache_metrics.metadata_load_cost += start.elapsed(); + result + } + + async fn read_parquet_metadata_inner( + &self, + file_path: &str, + file_size: u64, + cache_metrics: &mut MetadataCacheMetrics, + page_index_policy: PageIndexPolicy, + observation: SstMetaObservation, + ) -> Result<(Arc, bool)> { let file_id = self.file_handle.file_id(); // Tries to get from cache with metrics tracking. if let Some(metadata) = self .cache_strategy - .get_sst_meta_data(file_id, cache_metrics, page_index_policy) + .resolve_sst_meta(observation, file_id, cache_metrics, page_index_policy) .await { - cache_metrics.metadata_load_cost += start.elapsed(); return Ok((metadata, false)); } @@ -852,7 +907,7 @@ impl ParquetReaderBuilder { match metadata { SstMetaPreparation::Prepared(metadata) => { self.cache_strategy - .put_prepared_sst_meta(file_id, metadata, true); + .put_compact_sst_meta(file_id, metadata.compact); } SstMetaPreparation::DecodedOnly { encoding_error, .. } => { warn!( @@ -872,7 +927,6 @@ impl ParquetReaderBuilder { )?) }; - cache_metrics.metadata_load_cost += start.elapsed(); Ok((decoded, true)) } @@ -1635,13 +1689,13 @@ impl ReaderFilterMetrics { pub struct MetadataCacheMetrics { /// Number of memory cache hits for parquet metadata. pub mem_cache_hit: usize, - /// Number of file cache hits for parquet metadata. + /// Successful file-cache resolutions, including results shared with concurrent callers. pub file_cache_hit: usize, /// Number of cache misses for parquet metadata. pub cache_miss: usize, /// Duration to load parquet metadata. pub metadata_load_cost: Duration, - /// Number of read operations performed. + /// Physical reads executed by this caller; shared waiters do not repeat them. pub num_reads: usize, /// Total bytes read from storage. pub bytes_read: u64, @@ -2470,10 +2524,12 @@ mod tests { use datatypes::extension::json::Json2ExtensionType; use datatypes::prelude::ConcreteDataType; use datatypes::schema::ColumnSchema; + use object_store::layers::mock::MockLayerBuilder; use object_store::services::Memory; use parquet::arrow::ArrowWriter; use parquet::arrow::arrow_reader::RowSelector; use parquet::file::properties::WriterProperties; + use std::sync::atomic::{AtomicUsize, Ordering}; use store_api::codec::PrimaryKeyEncoding; use store_api::metadata::{ColumnMetadata, RegionMetadata, RegionMetadataBuilder}; use store_api::region_request::PathType; @@ -2482,12 +2538,276 @@ mod tests { use super::*; use crate::cache::CacheManager; + use crate::cache::test_util::{GatedStorageReader, parquet_file_data_with_region_metadata}; use crate::sst::parquet::metadata::MetadataLoader; use crate::sst::parquet::prefilter::{build_reader_filter_plan, execute_prefilter}; use crate::sst::parquet::read_columns::{ParquetReadColumn, ParquetReadColumns}; use crate::sst::parquet::row_group::ParquetFetchMetrics; use crate::test_util::sst_util::{sst_file_handle, sst_region_metadata}; + #[tokio::test] + async fn test_shared_file_metadata_resolution_counts_without_duplicate_io() { + use crate::cache::file_cache::{FileType, IndexKey, IndexValue}; + use crate::cache::write_cache::WriteCache; + use crate::test_util::TestEnv; + use common_base::readable_size::ReadableSize; + + let env = TestEnv::new().await; + let metadata = Arc::new(sst_region_metadata()); + let bytes = parquet_file_data_with_region_metadata(&metadata); + let handle = sst_file_handle(0, 3); + let reads = Arc::new(AtomicUsize::new(0)); + let gate = Arc::new(tokio::sync::Semaphore::new(0)); + let layer = MockLayerBuilder::default() + .reader_factory(Arc::new({ + let reads = reads.clone(); + let gate = gate.clone(); + move |_, _, inner| { + Box::new(GatedStorageReader { + inner, + reads: reads.clone(), + gate: gate.clone(), + }) + } + })) + .build() + .unwrap(); + let local = ObjectStore::new(Memory::default()).unwrap().layer(layer); + let write_cache = Arc::new( + WriteCache::new( + local.clone(), + ReadableSize::mb(10), + None, + None, + false, + env.get_puffin_manager(), + env.get_intermediate_manager(), + None, + ) + .await + .unwrap(), + ); + let file_cache = write_cache.file_cache(); + let key = IndexKey::new( + handle.region_id(), + handle.file_id().file_id(), + FileType::Parquet, + ); + local + .write(&file_cache.cache_file_path(key), bytes.clone()) + .await + .unwrap(); + file_cache + .put( + key, + IndexValue { + file_size: bytes.len() as _, + }, + ) + .await; + let cache = Arc::new( + CacheManager::builder() + .sst_meta_cache_size(1024 * 1024) + .write_cache(Some(write_cache)) + .build(), + ); + // The remote store is empty: success must come from the local file cache. + let builder = ParquetReaderBuilder::new( + "file_resolution".to_string(), + PathType::Bare, + handle.clone(), + ObjectStore::new(Memory::default()).unwrap(), + ) + .cache(CacheStrategy::EnableAll(cache)); + let path = handle.file_path("file_resolution", PathType::Bare); + let mut requests = (0..8) + .map(|_| { + Box::pin(async { + let mut metrics = MetadataCacheMetrics::default(); + let result = builder + .read_parquet_metadata( + &path, + bytes.len() as u64, + &mut metrics, + PageIndexPolicy::Skip, + ) + .await; + (result, metrics) + }) + }) + .collect::>(); + for request in &mut requests { + assert!(futures::poll!(request).is_pending()); + } + assert_eq!(1, reads.load(Ordering::Relaxed)); + gate.add_permits(32); + let results = + tokio::time::timeout(Duration::from_secs(10), futures::future::join_all(requests)) + .await + .unwrap(); + assert_eq!(1, reads.load(Ordering::Relaxed)); + assert_eq!( + 1, + results + .iter() + .map(|(_, metrics)| metrics.num_reads) + .sum::() + ); + for (result, metrics) in results { + assert!(!result.unwrap().1); + assert_eq!(0, metrics.mem_cache_hit); + assert_eq!(1, metrics.file_cache_hit); + assert_eq!(0, metrics.cache_miss); + } + let mut warm = MetadataCacheMetrics::default(); + builder + .read_parquet_metadata(&path, bytes.len() as u64, &mut warm, PageIndexPolicy::Skip) + .await + .unwrap(); + assert_eq!(1, warm.mem_cache_hit); + assert_eq!( + 0, warm.file_cache_hit, + "hot hits must not replay the original file source" + ); + assert_eq!(0, warm.num_reads); + } + + #[tokio::test] + #[allow( + clippy::print_stdout, + reason = "records the controlled storage-read comparison" + )] + async fn test_metadata_load_coalescing_storage_reads() { + const CONCURRENCY: usize = 8; + let metadata = Arc::new(sst_region_metadata()); + let bytes = parquet_file_data_with_region_metadata(&metadata); + let file_handle = sst_file_handle(0, 3); + let file_path = file_handle.file_path("metadata_coalescing", PathType::Bare); + // Same bytes, file, cache budget and concurrency; only bypass the coordinator + // in the control arm. Repeat with fresh caches to verify cold behavior. + for _ in 0..3 { + for coalesce in [false, true] { + let reads = Arc::new(AtomicUsize::new(0)); + let gate = Arc::new(tokio::sync::Semaphore::new(0)); + let layer = MockLayerBuilder::default() + .reader_factory(Arc::new({ + let reads = reads.clone(); + let gate = gate.clone(); + move |_, _, inner| { + Box::new(GatedStorageReader { + inner, + reads: reads.clone(), + gate: gate.clone(), + }) + } + })) + .build() + .unwrap(); + let store = ObjectStore::new(Memory::default()).unwrap().layer(layer); + store.write(&file_path, bytes.clone()).await.unwrap(); + let cache = Arc::new( + CacheManager::builder() + .sst_meta_cache_size(1024 * 1024) + .build(), + ); + let builder = ParquetReaderBuilder::new( + "metadata_coalescing".to_string(), + PathType::Bare, + file_handle.clone(), + store, + ) + .cache(CacheStrategy::EnableAll(cache.clone())); + let mut requests = (0..CONCURRENCY) + .map(|_| { + Box::pin(async { + let mut metrics = MetadataCacheMetrics::default(); + let result = if coalesce { + builder + .read_parquet_metadata( + &file_path, + bytes.len() as u64, + &mut metrics, + PageIndexPolicy::Skip, + ) + .await + } else { + let observation = builder + .cache_strategy + .observe_sst_meta( + file_handle.file_id(), + PageIndexPolicy::Skip, + &mut metrics, + ) + .await; + builder + .read_parquet_metadata_inner( + &file_path, + bytes.len() as u64, + &mut metrics, + PageIndexPolicy::Skip, + observation, + ) + .await + }; + (result.unwrap(), metrics) + }) + }) + .collect::>(); + for request in &mut requests { + assert!(futures::poll!(request).is_pending()); + } + let expected_reads = if coalesce { 1 } else { CONCURRENCY }; + assert_eq!(expected_reads, reads.load(Ordering::Relaxed)); + gate.add_permits(CONCURRENCY * 4); + let results = tokio::time::timeout( + Duration::from_secs(10), + futures::future::join_all(requests), + ) + .await + .unwrap(); + assert_eq!(expected_reads, reads.load(Ordering::Relaxed)); + assert_eq!( + expected_reads, + results.iter().map(|(_, m)| m.num_reads).sum::() + ); + for ((value, miss), request_metrics) in &results { + assert_eq!(1, request_metrics.cache_miss); + assert_eq!(0, request_metrics.mem_cache_hit); + assert_eq!(0, request_metrics.file_cache_hit); + assert!( + *miss, + "each request must preserve background-download eligibility" + ); + assert_eq!(3, value.parquet_metadata().file_metadata().num_rows()); + if coalesce { + assert!(Arc::ptr_eq(value, &results[0].0.0)); + } + } + let mut warm_metrics = MetadataCacheMetrics::default(); + let (_, miss) = tokio::time::timeout( + Duration::from_secs(5), + builder.read_parquet_metadata( + &file_path, + bytes.len() as u64, + &mut warm_metrics, + PageIndexPolicy::Skip, + ), + ) + .await + .unwrap() + .unwrap(); + assert!(!miss); + assert_eq!(1, warm_metrics.mem_cache_hit); + assert_eq!(0, warm_metrics.cache_miss); + assert_eq!(0, warm_metrics.num_reads); + assert_eq!(expected_reads, reads.load(Ordering::Relaxed)); + println!( + "coalesce={coalesce} concurrency={CONCURRENCY} cold_storage_reads={expected_reads} warm_storage_reads=0" + ); + } + } + } + async fn prefilter_test_builder( object_store: ObjectStore, predicate: Predicate, From d33c66418379463c5e7c928f3a132517d2a1fa9b Mon Sep 17 00:00:00 2001 From: fys <40801205+fengys1996@users.noreply.github.com> Date: Thu, 29 Oct 2026 14:29:00 +0000 Subject: [PATCH 140/239] feat(json2): support skipping indexes (#9967) * feat: support skipping indexes on JSON2 type hints Signed-off-by: fys * chore: remove unused code files Signed-off-by: fys * refactor(mito): pass build type to bloom filter indexer Signed-off-by: fys * refactor(mito): share bloom filter field and JSON update loop Signed-off-by: fys * refactor: better name * refactor(mito2): simplify JSON path bloom index handling Signed-off-by: fys * refactor(mito2): expose index targets from creators Signed-off-by: fys * docs(mito2): note JSON path index rebuild risk Signed-off-by: fys * chore: update greptime-proto * chore: update greptime-proto --------- Signed-off-by: fys --- Cargo.lock | 2 +- Cargo.toml | 2 +- src/api/src/v1/column_def.rs | 17 + src/common/grpc-expr/src/alter.rs | 42 ++ src/datatypes/src/json.rs | 47 ++ src/datatypes/src/vectors/json/array.rs | 1 + src/datatypes/src/vectors/json/builder.rs | 5 + src/index/src/target.rs | 13 + src/mito2/benches/bench_dense_index_update.rs | 12 +- src/mito2/benches/bench_index_update.rs | 11 +- src/mito2/src/cache.rs | 12 +- src/mito2/src/cache/index.rs | 15 +- .../src/cache/index/bloom_filter_index.rs | 93 ++- src/mito2/src/cache/index/result_cache.rs | 34 +- src/mito2/src/compaction/json2.rs | 9 + src/mito2/src/engine/puffin_index.rs | 2 +- src/mito2/src/memtable/bulk.rs | 1 + src/mito2/src/sst/index.rs | 10 +- .../src/sst/index/bloom_filter/applier.rs | 272 +++++++-- .../sst/index/bloom_filter/applier/builder.rs | 92 +-- .../src/sst/index/bloom_filter/creator.rs | 207 ++++--- src/mito2/src/sst/index/column.rs | 110 ++++ .../src/sst/index/fulltext_index/applier.rs | 2 +- src/mito2/src/sst/index/indexer/finish.rs | 17 +- .../index/inverted_index/applier/builder.rs | 85 +-- .../src/sst/index/inverted_index/creator.rs | 15 +- src/mito2/src/sst/parquet/json_align.rs | 1 + src/operator/src/expr_helper.rs | 21 +- src/pipeline/src/etl/transform.rs | 1 + .../src/etl/transform/transformer/greptime.rs | 2 + .../transform/transformer/greptime/coerce.rs | 1 + src/query/src/datafusion/json_expr_planner.rs | 3 + src/sql/src/parsers/create_parser/json.rs | 90 ++- src/sql/src/statements.rs | 1 + src/sql/src/statements/create.rs | 21 + src/store-api/src/region_request.rs | 24 + .../types/json/json2_inverted_index.result | 9 +- .../types/json/json2_inverted_index.sql | 7 +- .../types/json/json2_skipping_index.result | 550 ++++++++++++++++++ .../types/json/json2_skipping_index.sql | 190 ++++++ 40 files changed, 1725 insertions(+), 324 deletions(-) create mode 100644 tests/cases/standalone/common/types/json/json2_skipping_index.result create mode 100644 tests/cases/standalone/common/types/json/json2_skipping_index.sql diff --git a/Cargo.lock b/Cargo.lock index 59c8421e955..3b7ab8897ec 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6156,7 +6156,7 @@ dependencies = [ [[package]] name = "greptime-proto" version = "0.1.0" -source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=3b5ba5e7422e1b96500cf052bb94e978f4660c05#3b5ba5e7422e1b96500cf052bb94e978f4660c05" +source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=9eb24ddd94ddb9295bff9e596834832a7d8ba008#9eb24ddd94ddb9295bff9e596834832a7d8ba008" dependencies = [ "prost 0.14.1", "prost-types 0.14.1", diff --git a/Cargo.toml b/Cargo.toml index 1f84eb7dbcb..e1271bdecd1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -163,7 +163,7 @@ fs2 = "0.4" fst = "0.4.7" futures = "0.3" futures-util = "0.3" -greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "3b5ba5e7422e1b96500cf052bb94e978f4660c05" } +greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "9eb24ddd94ddb9295bff9e596834832a7d8ba008" } hdfs-native = "0.14.6" hex = "0.4" hostname = "0.4.0" diff --git a/src/api/src/v1/column_def.rs b/src/api/src/v1/column_def.rs index 73c56339474..b1d2f0841d5 100644 --- a/src/api/src/v1/column_def.rs +++ b/src/api/src/v1/column_def.rs @@ -280,6 +280,23 @@ pub fn as_skipping_index_type(skipping_index_type: PbSkippingIndexType) -> Skipp } } +/// Validates skipping options received for a JSON2 type hint. +pub fn json_skipping_options_from_proto( + options: crate::v1::JsonSkippingIndexOptions, +) -> datatypes::error::Result { + let index_type = PbSkippingIndexType::try_from(options.index_type).map_err(|_| { + datatypes::error::InvalidSkippingIndexOptionSnafu { + msg: format!("Unknown skipping index type: {}", options.index_type), + } + .build() + })?; + SkippingIndexOptions::new( + options.granularity, + options.false_positive_rate, + as_skipping_index_type(index_type), + ) +} + #[cfg(test)] mod tests { diff --git a/src/common/grpc-expr/src/alter.rs b/src/common/grpc-expr/src/alter.rs index ad3a9937887..dda1d84015a 100644 --- a/src/common/grpc-expr/src/alter.rs +++ b/src/common/grpc-expr/src/alter.rs @@ -75,6 +75,16 @@ fn json_settings_from_proto(settings: api::v1::JsonSettings) -> Result>>()?; @@ -426,6 +436,7 @@ mod tests { path: vec!["service".into()], data_type: ColumnDataType::String as i32, inverted_index: enabled, + skipping_index: None, ..Default::default() }], max_auto_expanded_paths: Some(0), @@ -435,6 +446,36 @@ mod tests { } } + #[test] + fn test_json_settings_skipping_options_from_proto() { + for (granularity, rate, index_type, valid) in [ + (8, 0.05, 0, true), + (0, 0.05, 0, false), + (8, 0.0, 0, false), + (8, 0.05, 9, false), + ] { + let settings = json_settings_from_proto(api::v1::JsonSettings { + type_hints: vec![api::v1::JsonTypeHint { + path: vec!["service".into()], + data_type: ColumnDataType::String as i32, + skipping_index: Some(api::v1::JsonSkippingIndexOptions { + granularity, + index_type, + false_positive_rate: rate, + }), + ..Default::default() + }], + max_auto_expanded_paths: Some(0), + }); + assert_eq!(settings.is_ok(), valid); + if let Ok(settings) = settings { + let options = settings.type_hints()[0].skipping_index.as_ref().unwrap(); + assert_eq!(options.granularity, granularity); + assert_eq!(options.false_positive_rate(), rate); + } + } + } + #[test] fn test_json_settings_path_validation_from_proto() { for (path, valid_path) in [ @@ -454,6 +495,7 @@ mod tests { path: path.iter().map(|part| (*part).to_string()).collect(), data_type: ColumnDataType::String as i32, inverted_index, + skipping_index: None, ..Default::default() }], max_auto_expanded_paths: Some(0), diff --git a/src/datatypes/src/json.rs b/src/datatypes/src/json.rs index 78af6b240bb..2c86170fbce 100644 --- a/src/datatypes/src/json.rs +++ b/src/datatypes/src/json.rs @@ -93,6 +93,9 @@ pub struct JsonTypeHint { #[serde(rename = "type")] pub data_type: ConcreteDataType, pub inverted_index: bool, + /// Optional Bloom Filter configuration for this path. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub skipping_index: Option, } /// Specifies how encoding handles a value that does not match its JSON2 type hint. @@ -215,6 +218,13 @@ fn validate_type_hints(type_hints: &[JsonTypeHint]) -> Result<()> { let mut object = JsonObjectType::new(); for hint in type_hints { + if let Some(options) = &hint.skipping_index { + crate::schema::SkippingIndexOptions::new( + options.granularity, + options.false_positive_rate(), + options.index_type.clone(), + )?; + } if hint.path.len() > JSON2_MAX_STRUCTURED_DEPTH { return InvalidJson2SettingsSnafu { reason: format!( @@ -738,6 +748,24 @@ mod tests { &struct_value.items()[index] } + #[test] + fn test_skipping_index_settings_roundtrip() { + let old = r#"{"type_hints":[{"path":["service","name"],"type":{"String":{}},"inverted_index":true}]}"#; + let legacy: JsonSettings = serde_json::from_str(old).unwrap(); + assert!(legacy.type_hints()[0].skipping_index.is_none()); + let mut hints = legacy.type_hints().to_vec(); + hints[0].skipping_index = Some(crate::schema::SkippingIndexOptions::default()); + let settings = JsonSettings::try_new(hints.clone(), Some(0)).unwrap(); + let json = serde_json::to_string(&settings).unwrap(); + assert_eq!( + settings, + serde_json::from_str::(&json).unwrap() + ); + assert!(!settings.equivalent(&legacy)); + hints[0].skipping_index.as_mut().unwrap().granularity = 0; + assert!(JsonSettings::try_new(hints, Some(0)).is_err()); + } + #[test] fn test_json_settings_equivalent() -> Result<()> { let hints = vec![ @@ -745,11 +773,13 @@ mod tests { path: vec!["user".to_string(), "name".to_string()], data_type: ConcreteDataType::string_datatype(), inverted_index: false, + skipping_index: None, }, JsonTypeHint { path: vec!["count".to_string()], data_type: ConcreteDataType::int64_datatype(), inverted_index: false, + skipping_index: None, }, ]; let settings = JsonSettings::try_new(hints.clone(), Some(10))?; @@ -821,11 +851,13 @@ mod tests { path: vec!["user".to_string(), "age".to_string()], data_type: ConcreteDataType::int64_datatype(), inverted_index: true, + skipping_index: None, }, JsonTypeHint { path: vec!["user".to_string(), "name".to_string()], data_type: ConcreteDataType::string_datatype(), inverted_index: false, + skipping_index: None, }, ], None, @@ -841,11 +873,13 @@ mod tests { path: vec!["user".to_string(), "age".to_string()], data_type: ConcreteDataType::int64_datatype(), inverted_index: true, + skipping_index: None, }, JsonTypeHint { path: vec!["user".to_string(), "name".to_string()], data_type: ConcreteDataType::string_datatype(), inverted_index: false, + skipping_index: None, }, ]; @@ -885,6 +919,7 @@ mod tests { path: path.iter().map(|part| (*part).to_string()).collect(), data_type: ConcreteDataType::string_datatype(), inverted_index, + skipping_index: None, }; let serialized = json!({"type_hints": [hint.clone()]}); assert_eq!( @@ -919,6 +954,7 @@ mod tests { path, data_type, inverted_index: false, + skipping_index: None, }; assert!( JsonSettings::try_new( @@ -982,6 +1018,7 @@ mod tests { path: vec!["kind".to_string()], data_type: ConcreteDataType::int64_datatype(), inverted_index: false, + skipping_index: None, }], None, )?; @@ -1012,11 +1049,13 @@ mod tests { path: vec!["to_string".to_string()], data_type: ConcreteDataType::string_datatype(), inverted_index: false, + skipping_index: None, }, JsonTypeHint { path: vec!["to_int".to_string()], data_type: ConcreteDataType::int64_datatype(), inverted_index: false, + skipping_index: None, }, ], None, @@ -1051,6 +1090,7 @@ mod tests { path: vec!["user".to_string(), "age".to_string()], data_type: ConcreteDataType::int64_datatype(), inverted_index: false, + skipping_index: None, }], None, )?; @@ -1212,6 +1252,7 @@ mod tests { path: vec!["age".to_string()], data_type: ConcreteDataType::int64_datatype(), inverted_index: false, + skipping_index: None, }]; let settings = JsonSettings::try_new(type_hints, None)?; @@ -1244,6 +1285,7 @@ mod tests { path: vec!["count".to_string()], data_type: ConcreteDataType::uint64_datatype(), inverted_index: false, + skipping_index: None, }]; let settings = JsonSettings::try_new(type_hints, None)?; @@ -1278,6 +1320,7 @@ mod tests { path: vec!["user".to_string(), "age".to_string()], data_type: ConcreteDataType::int64_datatype(), inverted_index: false, + skipping_index: None, }]; let settings = JsonSettings::try_new(type_hints, None)?; @@ -1303,6 +1346,7 @@ mod tests { path: vec!["user".to_string(), "name".to_string()], data_type: ConcreteDataType::string_datatype(), inverted_index: false, + skipping_index: None, }]; let settings = JsonSettings::try_new(type_hints, None)?; @@ -1328,6 +1372,7 @@ mod tests { path: vec!["user".to_string(), "age".to_string()], data_type: ConcreteDataType::int64_datatype(), inverted_index: false, + skipping_index: None, }]; let settings = JsonSettings::try_new(type_hints, None)?; @@ -1342,11 +1387,13 @@ mod tests { path: vec!["user".to_string(), "age".to_string()], data_type: ConcreteDataType::int64_datatype(), inverted_index: false, + skipping_index: None, }, JsonTypeHint { path: vec!["user".to_string(), "name".to_string()], data_type: ConcreteDataType::string_datatype(), inverted_index: false, + skipping_index: None, }, ]; let settings = JsonSettings::try_new(type_hints, None)?; diff --git a/src/datatypes/src/vectors/json/array.rs b/src/datatypes/src/vectors/json/array.rs index 5b6265c7b37..48644ad5c58 100644 --- a/src/datatypes/src/vectors/json/array.rs +++ b/src/datatypes/src/vectors/json/array.rs @@ -643,6 +643,7 @@ mod test { path: vec!["kind".to_string()], data_type: ConcreteDataType::string_datatype(), inverted_index: false, + skipping_index: None, }], Some(0), )?; diff --git a/src/datatypes/src/vectors/json/builder.rs b/src/datatypes/src/vectors/json/builder.rs index 2b9c44a326a..515bd8ec3e5 100644 --- a/src/datatypes/src/vectors/json/builder.rs +++ b/src/datatypes/src/vectors/json/builder.rs @@ -1061,16 +1061,19 @@ mod tests { path: vec!["kind".to_string()], data_type: ConcreteDataType::string_datatype(), inverted_index: false, + skipping_index: None, }, JsonTypeHint { path: vec!["commit".to_string(), "operation".to_string()], data_type: ConcreteDataType::string_datatype(), inverted_index: false, + skipping_index: None, }, JsonTypeHint { path: vec!["time_us".to_string()], data_type: ConcreteDataType::int64_datatype(), inverted_index: false, + skipping_index: None, }, ], Some(0), @@ -1166,6 +1169,7 @@ mod tests { path: vec!["hint".to_string()], data_type: ConcreteDataType::string_datatype(), inverted_index: false, + skipping_index: None, }], Some(2), )?; @@ -1331,6 +1335,7 @@ mod tests { path: vec!["nested".to_string(), "hinted".to_string()], data_type: ConcreteDataType::string_datatype(), inverted_index: false, + skipping_index: None, }], Some(8), )?; diff --git a/src/index/src/target.rs b/src/index/src/target.rs index 7a5b3811b44..679510e2aef 100644 --- a/src/index/src/target.rs +++ b/src/index/src/target.rs @@ -64,6 +64,19 @@ impl JsonPathTarget { } impl IndexTarget { + /// Estimates the memory used by this target, including JSON path strings. + pub fn approx_mem_size(&self) -> usize { + std::mem::size_of::() + + match self { + Self::ColumnId(_) => 0, + Self::JsonPath(target) => target + .path + .iter() + .map(|part| std::mem::size_of::() + part.len()) + .sum::(), + } + } + /// Creates an index target for a JSON object path. pub fn new_json_path(column_id: ColumnId, path: Vec) -> Result { validate_json_path(&path)?; diff --git a/src/mito2/benches/bench_dense_index_update.rs b/src/mito2/benches/bench_dense_index_update.rs index 9f827c5ab80..17d9b5f8cef 100644 --- a/src/mito2/benches/bench_dense_index_update.rs +++ b/src/mito2/benches/bench_dense_index_update.rs @@ -41,8 +41,8 @@ use datatypes::value::Value; use index::target::IndexTarget; use mito_codec::row_converter::{DensePrimaryKeyCodec, PrimaryKeyCodecExt}; use mito2::read::{Batch, BatchBuilder}; -use mito2::sst::index::Indexer; use mito2::sst::index::intermediate::IntermediateManager; +use mito2::sst::index::{IndexBuildType, Indexer}; use mito2::sst::{FlatSchemaOptions, to_flat_sst_arrow_schema}; use mito2::test_util::bench_util::{BloomFilterIndexer, InvertedIndexer}; use store_api::codec::PrimaryKeyEncoding; @@ -367,8 +367,14 @@ fn bench_dense_index_update(c: &mut Criterion) { ) }); let bloom = if shape.bloom { - BloomFilterIndexer::new(file_id, &metadata, intermediate.clone(), None) - .unwrap() + BloomFilterIndexer::new( + file_id, + &metadata, + intermediate.clone(), + None, + IndexBuildType::Flush, + ) + .unwrap() } else { None }; diff --git a/src/mito2/benches/bench_index_update.rs b/src/mito2/benches/bench_index_update.rs index 0d0b1ef6c23..e8a54dbea3d 100644 --- a/src/mito2/benches/bench_index_update.rs +++ b/src/mito2/benches/bench_index_update.rs @@ -35,6 +35,7 @@ use datatypes::data_type::ConcreteDataType; use datatypes::schema::{ColumnSchema, SkippingIndexOptions}; use index::target::IndexTarget; use mito_codec::row_converter::SparsePrimaryKeyCodec; +use mito2::sst::index::IndexBuildType; use mito2::sst::index::intermediate::IntermediateManager; use mito2::sst::{FlatSchemaOptions, to_flat_sst_arrow_schema}; use mito2::test_util::bench_util::{BloomFilterIndexer, InvertedIndexer}; @@ -250,8 +251,14 @@ fn bench_index_update(c: &mut Criterion) { ) }); let mut bloom = if shape.bloom { - BloomFilterIndexer::new(file_id, &metadata, intermediate.clone(), None) - .unwrap() + BloomFilterIndexer::new( + file_id, + &metadata, + intermediate.clone(), + None, + IndexBuildType::Flush, + ) + .unwrap() } else { None }; diff --git a/src/mito2/src/cache.rs b/src/mito2/src/cache.rs index 4cff42043ce..2d63a09329f 100644 --- a/src/mito2/src/cache.rs +++ b/src/mito2/src/cache.rs @@ -3601,10 +3601,10 @@ mod tests { let bloom_key = ( index_id.file_id(), index_id.version, - column_id, + Arc::new(::index::target::IndexTarget::ColumnId(column_id)), Tag::Skipping, ); - bloom_cache.put_metadata(bloom_key, Arc::new(BloomFilterMeta::default())); + bloom_cache.put_metadata(bloom_key.clone(), Arc::new(BloomFilterMeta::default())); inverted_cache.put_metadata( (index_id.file_id(), index_id.version), Arc::new(InvertedIndexMetas::default()), @@ -3619,7 +3619,7 @@ mod tests { }); puffin_metadata_cache.put_metadata(file_id_str.clone(), metadata); - assert!(bloom_cache.get_metadata(bloom_key).is_some()); + assert!(bloom_cache.get_metadata(bloom_key.clone()).is_some()); assert!( inverted_cache .get_metadata((index_id.file_id(), index_id.version)) @@ -3630,7 +3630,7 @@ mod tests { cache.evict_puffin_cache(index_id).await; - assert!(bloom_cache.get_metadata(bloom_key).is_none()); + assert!(bloom_cache.get_metadata(bloom_key.clone()).is_none()); assert!( inverted_cache .get_metadata((index_id.file_id(), index_id.version)) @@ -3640,7 +3640,7 @@ mod tests { assert!(puffin_metadata_cache.get_metadata(&file_id_str).is_none()); // Refill caches and evict via CacheStrategy to ensure delegation works. - bloom_cache.put_metadata(bloom_key, Arc::new(BloomFilterMeta::default())); + bloom_cache.put_metadata(bloom_key.clone(), Arc::new(BloomFilterMeta::default())); inverted_cache.put_metadata( (index_id.file_id(), index_id.version), Arc::new(InvertedIndexMetas::default()), @@ -3661,7 +3661,7 @@ mod tests { let strategy = CacheStrategy::EnableAll(cache.clone()); strategy.evict_puffin_cache(index_id).await; - assert!(bloom_cache.get_metadata(bloom_key).is_none()); + assert!(bloom_cache.get_metadata(bloom_key.clone()).is_none()); assert!( inverted_cache .get_metadata((index_id.file_id(), index_id.version)) diff --git a/src/mito2/src/cache/index.rs b/src/mito2/src/cache/index.rs index 26f9a3d5f96..ab813c0353d 100644 --- a/src/mito2/src/cache/index.rs +++ b/src/mito2/src/cache/index.rs @@ -151,7 +151,7 @@ where impl IndexCache where - K: Hash + Eq + Clone + Copy + Send + Sync + 'static, + K: Hash + Eq + Clone + Send + Sync + 'static, M: Send + Sync + 'static, { pub fn get_metadata(&self, key: K) -> Option> { @@ -195,7 +195,7 @@ where let last_index = page_keys.len() - 1; // TODO: Avoid copy as much as possible. for (i, page_key) in page_keys.iter().enumerate() { - match self.get_page(key, *page_key) { + match self.get_page(key.clone(), *page_key) { Some(page) => { CACHE_HIT.with_label_values(&[INDEX_CONTENT_TYPE]).inc(); metrics.cache_hit += 1; @@ -222,7 +222,7 @@ where let page_key = page_keys[i]; metrics.page_bytes += page.len() as u64; data[i] = page.clone(); - self.put_page(key, page_key, page.clone()); + self.put_page(key.clone(), page_key, page.clone()); } } let buffer = Buffer::from_iter(data); @@ -260,7 +260,7 @@ where continue; } metrics.num_pages += 1; - match self.get_page(key, page_key) { + match self.get_page(key.clone(), page_key) { Some(page) => { CACHE_HIT.with_label_values(&[INDEX_CONTENT_TYPE]).inc(); metrics.cache_hit += 1; @@ -289,7 +289,7 @@ where debug_assert_eq!(loaded.len(), missing.len()); for (page_id, page) in missing.into_iter().zip(loaded) { metrics.page_bytes += page.len() as u64; - self.put_page(key, PageKey { page_id }, page.clone()); + self.put_page(key.clone(), PageKey { page_id }, page.clone()); pages.insert(page_id, page); } } @@ -320,10 +320,11 @@ where fn put_page(&self, key: K, page_key: PageKey, value: Bytes) { // Clones the value to ensure it doesn't reference a larger buffer. let value = Bytes::from(value.to_vec()); + let key = (key, page_key); CACHE_BYTES .with_label_values(&[INDEX_CONTENT_TYPE]) - .add((self.weight_of_content)(&(key, page_key), &value).into()); - self.index.insert((key, page_key), value); + .add((self.weight_of_content)(&key, &value).into()); + self.index.insert(key, value); } /// Invalidates all cache entries whose keys satisfy `predicate`. diff --git a/src/mito2/src/cache/index/bloom_filter_index.rs b/src/mito2/src/cache/index/bloom_filter_index.rs index 57f0103c1bb..0e81817a754 100644 --- a/src/mito2/src/cache/index/bloom_filter_index.rs +++ b/src/mito2/src/cache/index/bloom_filter_index.rs @@ -21,7 +21,8 @@ use async_trait::async_trait; use bytes::Bytes; use index::bloom_filter::error::Result; use index::bloom_filter::reader::{BloomFilterReadMetrics, BloomFilterReader}; -use store_api::storage::{ColumnId, FileId, IndexVersion}; +use index::target::IndexTarget; +use store_api::storage::{FileId, IndexVersion}; use crate::cache::index::{INDEX_METADATA_TYPE, IndexCache, PageKey}; use crate::metrics::{CACHE_HIT, CACHE_MISS}; @@ -35,7 +36,7 @@ pub enum Tag { Fulltext, } -pub type BloomFilterIndexKey = (FileId, IndexVersion, ColumnId, Tag); +pub type BloomFilterIndexKey = (FileId, IndexVersion, Arc, Tag); /// Cache for bloom filter index. pub type BloomFilterIndexCache = IndexCache; @@ -64,7 +65,7 @@ impl BloomFilterIndexCache { fn bloom_filter_index_metadata_weight(k: &BloomFilterIndexKey, meta: &Arc) -> u32 { let base = k.0.as_bytes().len() + std::mem::size_of::() - + std::mem::size_of::() + + k.2.approx_mem_size() + std::mem::size_of::() + std::mem::size_of::(); @@ -76,14 +77,14 @@ fn bloom_filter_index_metadata_weight(k: &BloomFilterIndexKey, meta: &Arc u32 { - (k.0.as_bytes().len() + std::mem::size_of::() + v.len()) as u32 + (k.0.as_bytes().len() + k.2.approx_mem_size() + v.len()) as u32 } /// Bloom filter index blob reader with cache. pub struct CachedBloomFilterIndexBlobReader { file_id: FileId, index_version: IndexVersion, - column_id: ColumnId, + target: Arc, tag: Tag, blob_size: u64, inner: R, @@ -95,7 +96,7 @@ impl CachedBloomFilterIndexBlobReader { pub fn new( file_id: FileId, index_version: IndexVersion, - column_id: ColumnId, + target: IndexTarget, tag: Tag, blob_size: u64, inner: R, @@ -104,7 +105,7 @@ impl CachedBloomFilterIndexBlobReader { Self { file_id, index_version, - column_id, + target: Arc::new(target), tag, blob_size, inner, @@ -126,7 +127,12 @@ impl BloomFilterReader for CachedBloomFilterIndexBl let (result, cache_metrics) = self .cache .get_or_load( - (self.file_id, self.index_version, self.column_id, self.tag), + ( + self.file_id, + self.index_version, + self.target.clone(), + self.tag, + ), self.blob_size, offset, size, @@ -158,7 +164,12 @@ impl BloomFilterReader for CachedBloomFilterIndexBl let (pages, total_cache_metrics) = self .cache .get_or_load_vec( - (self.file_id, self.index_version, self.column_id, self.tag), + ( + self.file_id, + self.index_version, + self.target.clone(), + self.tag, + ), self.blob_size, ranges, move |ranges| async move { inner.read_vec(&ranges, None).await }, @@ -182,10 +193,12 @@ impl BloomFilterReader for CachedBloomFilterIndexBl &self, metrics: Option<&mut BloomFilterReadMetrics>, ) -> Result> { - if let Some(cached) = - self.cache - .get_metadata((self.file_id, self.index_version, self.column_id, self.tag)) - { + if let Some(cached) = self.cache.get_metadata(( + self.file_id, + self.index_version, + self.target.clone(), + self.tag, + )) { CACHE_HIT.with_label_values(&[INDEX_METADATA_TYPE]).inc(); if let Some(m) = metrics { m.cache_hit += 1; @@ -194,7 +207,12 @@ impl BloomFilterReader for CachedBloomFilterIndexBl } else { let meta = self.inner.metadata(metrics).await?; self.cache.put_metadata( - (self.file_id, self.index_version, self.column_id, self.tag), + ( + self.file_id, + self.index_version, + self.target.clone(), + self.tag, + ), meta.clone(), ); CACHE_MISS.with_label_values(&[INDEX_METADATA_TYPE]).inc(); @@ -206,11 +224,49 @@ impl BloomFilterReader for CachedBloomFilterIndexBl #[cfg(test)] mod test { use rand::{Rng, RngCore}; + use store_api::storage::ColumnId; use super::*; const FUZZ_REPEAT_TIMES: usize = 100; + #[test] + fn json_paths_have_independent_cache_entries() { + let cache = BloomFilterIndexCache::new(1024 * 1024, 1024 * 1024, 1024); + let file_id = FileId::random(); + let first = ( + file_id, + 0, + Arc::new(IndexTarget::new_json_path(1, vec!["a".into(), "b".into()]).unwrap()), + Tag::Skipping, + ); + let second = ( + file_id, + 0, + Arc::new(IndexTarget::new_json_path(1, vec!["a.b".into()]).unwrap()), + Tag::Skipping, + ); + cache.put_metadata( + first.clone(), + Arc::new(BloomFilterMeta { + row_count: 1, + ..Default::default() + }), + ); + cache.put_metadata( + second.clone(), + Arc::new(BloomFilterMeta { + row_count: 2, + ..Default::default() + }), + ); + assert_eq!(cache.get_metadata(first.clone()).unwrap().row_count, 1); + assert_eq!(cache.get_metadata(second.clone()).unwrap().row_count, 2); + cache.invalidate_file(file_id); + assert!(cache.get_metadata(first).is_none()); + assert!(cache.get_metadata(second).is_none()); + } + #[test] fn bloom_filter_metadata_weight_counts_vec_contents() { let file_id = FileId::parse_str("00000000-0000-0000-0000-000000000001").unwrap(); @@ -239,13 +295,18 @@ mod test { }; let weight = bloom_filter_index_metadata_weight( - &(file_id, version, column_id, tag), + &( + file_id, + version, + Arc::new(IndexTarget::ColumnId(column_id)), + tag, + ), &Arc::new(meta.clone()), ); let base = file_id.as_bytes().len() + std::mem::size_of::() - + std::mem::size_of::() + + std::mem::size_of::() + std::mem::size_of::() + std::mem::size_of::(); let expected_dynamic = meta.segment_loc_indices.len() * std::mem::size_of::() diff --git a/src/mito2/src/cache/index/result_cache.rs b/src/mito2/src/cache/index/result_cache.rs index 4448d9e0b08..75dd9794ef2 100644 --- a/src/mito2/src/cache/index/result_cache.rs +++ b/src/mito2/src/cache/index/result_cache.rs @@ -128,7 +128,7 @@ impl PredicateKey { } /// Creates a new bloom filter key. - pub fn new_bloom(predicates: Arc>>) -> Self { + pub fn new_bloom(predicates: Arc>>) -> Self { Self::Bloom(BloomFilterKey::new(predicates)) } @@ -195,21 +195,24 @@ impl FulltextIndexKey { /// Key for bloom filter queries. #[derive(Debug, Clone, PartialEq, Eq, Hash, Default)] pub struct BloomFilterKey { - predicates: Arc>>, + predicates: Arc>>, mem_usage: usize, } impl BloomFilterKey { /// Creates a new bloom filter key with the given predicates. /// Calculates memory usage based on the size of predicate lists. - pub fn new(predicates: Arc>>) -> Self { + pub fn new(predicates: Arc>>) -> Self { let mem_usage = predicates - .values() - .map(|predicates| { - predicates - .iter() - .map(|predicate| predicate.list.iter().map(|list| list.len()).sum::()) - .sum::() + .iter() + .map(|(target, predicates)| { + target.approx_mem_size() + + predicates + .iter() + .map(|predicate| { + predicate.list.iter().map(|list| list.len()).sum::() + }) + .sum::() }) .sum(); Self { @@ -233,16 +236,7 @@ impl InvertedIndexKey { let mem_usage = predicates .iter() .map(|(target, predicates)| { - let target_size = size_of::() - + match target { - IndexTarget::ColumnId(_) => 0, - IndexTarget::JsonPath(target) => target - .path() - .iter() - .map(|part| size_of::() + part.len()) - .sum(), - }; - target_size + target.approx_mem_size() + predicates .iter() .map(|predicate| match predicate { @@ -485,7 +479,7 @@ mod tests { let predicate2 = BloomInListPredicate { list: BTreeSet::from([b"test1".to_vec(), b"test2".to_vec()]), }; - predicates2.insert(1, vec![predicate2]); + predicates2.insert(IndexTarget::ColumnId(1), vec![predicate2]); let key2 = PredicateKey::new_bloom(Arc::new(predicates2)); let selection2 = Arc::new(RowGroupSelection::from_row_ids( [1, 2, 3].into_iter().collect(), diff --git a/src/mito2/src/compaction/json2.rs b/src/mito2/src/compaction/json2.rs index cb2a9d34ec8..faa7eb6c434 100644 --- a/src/mito2/src/compaction/json2.rs +++ b/src/mito2/src/compaction/json2.rs @@ -283,6 +283,7 @@ fn select_dynamic_hints( path: path.iter().map(|x| (*x).to_owned()).collect(), data_type: ConcreteDataType::from_arrow_type(&stat.data_type.as_arrow_type()), inverted_index: false, + skipping_index: None, }) .collect() } @@ -378,6 +379,7 @@ mod tests { path: vec!["hint".to_string()], data_type: ConcreteDataType::string_datatype(), inverted_index: false, + skipping_index: None, }], Some(2), )?; @@ -434,6 +436,7 @@ mod tests { path: vec!["kind".to_string()], data_type: ConcreteDataType::string_datatype(), inverted_index: false, + skipping_index: None, }], Some(0), )?; @@ -488,6 +491,7 @@ mod tests { path: vec!["kind".to_string()], data_type: ConcreteDataType::string_datatype(), inverted_index: false, + skipping_index: None, }], Some(0), )?; @@ -508,11 +512,13 @@ mod tests { path: vec!["kind".to_string()], data_type: ConcreteDataType::string_datatype(), inverted_index: false, + skipping_index: None, }, JsonTypeHint { path: vec!["source_only".to_string()], data_type: ConcreteDataType::int64_datatype(), inverted_index: false, + skipping_index: None, }, ], None, @@ -579,6 +585,7 @@ mod tests { path: vec!["kind".to_string()], data_type: ConcreteDataType::string_datatype(), inverted_index: false, + skipping_index: None, }], Some(0), )?; @@ -588,11 +595,13 @@ mod tests { path: vec!["kind".to_string()], data_type: ConcreteDataType::string_datatype(), inverted_index: false, + skipping_index: None, }, JsonTypeHint { path: vec!["promoted".to_string()], data_type: ConcreteDataType::int64_datatype(), inverted_index: false, + skipping_index: None, }, ], Some(0), diff --git a/src/mito2/src/engine/puffin_index.rs b/src/mito2/src/engine/puffin_index.rs index cef4f3a449b..975569aa213 100644 --- a/src/mito2/src/engine/puffin_index.rs +++ b/src/mito2/src/engine/puffin_index.rs @@ -316,7 +316,7 @@ async fn try_read_bloom_meta( CachedBloomFilterIndexBlobReader::new( region_index_id.file_id(), region_index_id.version, - column_id, + IndexTarget::ColumnId(column_id), tag, blob_size, bloom_reader, diff --git a/src/mito2/src/memtable/bulk.rs b/src/mito2/src/memtable/bulk.rs index e323fb6b340..e1de64c1122 100644 --- a/src/mito2/src/memtable/bulk.rs +++ b/src/mito2/src/memtable/bulk.rs @@ -2053,6 +2053,7 @@ mod tests { path: vec!["id".to_string()], data_type: ConcreteDataType::int64_datatype(), inverted_index: false, + skipping_index: None, }], Some(0), )?; diff --git a/src/mito2/src/sst/index.rs b/src/mito2/src/sst/index.rs index fe39af344d2..881a72ead83 100644 --- a/src/mito2/src/sst/index.rs +++ b/src/mito2/src/sst/index.rs @@ -308,12 +308,16 @@ impl Indexer { let requested: HashSet<_> = self .inverted_indexer .iter() - .flat_map(|indexer| indexer.indexed_column_ids()) + .flat_map(|indexer| indexer.indexed_targets()) .chain( self.bloom_filter_indexer .iter() - .flat_map(|indexer| indexer.column_ids()), + .flat_map(|indexer| indexer.indexed_targets()), ) + .filter_map(|target| match target { + IndexTarget::ColumnId(id) => Some(*id), + IndexTarget::JsonPath(_) => None, + }) .collect(); if metadata.primary_key.iter().all(|id| requested.contains(id)) { self.dense_pk_decoder = Some(DensePrimaryKeyCodec::new(metadata)); @@ -590,6 +594,7 @@ impl IndexerBuilderImpl { /// Only flush and compaction builds collect these targets; manual and schema-change /// rebuilds return no JSON targets. fn collect_json_inverted_index_targets(&self) -> Vec { + // FIXME: Rebuilds replace the whole Puffin file and can drop existing JSON path indexes. if !matches!( self.build_type, IndexBuildType::Flush | IndexBuildType::Compact @@ -702,6 +707,7 @@ impl IndexerBuilderImpl { &self.metadata, self.intermediate_manager.clone(), mem_limit, + self.build_type.clone(), ); let err = match indexer { diff --git a/src/mito2/src/sst/index/bloom_filter/applier.rs b/src/mito2/src/sst/index/bloom_filter/applier.rs index decf50470b7..0f97e89a0fa 100644 --- a/src/mito2/src/sst/index/bloom_filter/applier.rs +++ b/src/mito2/src/sst/index/bloom_filter/applier.rs @@ -33,7 +33,6 @@ use puffin::puffin_manager::{PuffinManager, PuffinReader}; use snafu::ResultExt; use store_api::metadata::RegionMetadataRef; use store_api::region_request::PathType; -use store_api::storage::ColumnId; use crate::access_layer::{RegionFilePathFactory, WriteCachePathProvider}; use crate::cache::file_cache::{FileCacheRef, FileType, IndexKey}; @@ -134,10 +133,10 @@ pub struct BloomFilterIndexApplier { /// Bloom filter predicates. /// For each column, the value will be retained only if it contains __all__ predicates. - default_predicates: Arc>>, + default_predicates: Arc>>, - /// Expected predicate column types from the latest region metadata. - expected_predicate_col_types: BTreeMap, + /// Expected predicate target types from the latest region metadata. + expected_predicate_target_types: BTreeMap, } impl BloomFilterIndexApplier { @@ -149,8 +148,8 @@ impl BloomFilterIndexApplier { path_type: PathType, object_store: ObjectStore, puffin_manager_factory: PuffinManagerFactory, - predicates: BTreeMap>, - expected_predicate_col_types: BTreeMap, + predicates: BTreeMap>, + expected_predicate_target_types: BTreeMap, ) -> Self { let default_predicates = Arc::new(predicates); Self { @@ -162,7 +161,7 @@ impl BloomFilterIndexApplier { puffin_metadata_cache: None, bloom_filter_index_cache: None, default_predicates, - expected_predicate_col_types, + expected_predicate_target_types, } } @@ -209,7 +208,7 @@ impl BloomFilterIndexApplier { &self, file_id: RegionIndexId, file_size_hint: Option, - predicates: &BTreeMap>, + predicates: &BTreeMap>, row_groups: impl Iterator, mut metrics: Option<&mut BloomFilterIndexApplyMetrics>, ) -> Result>)>> { @@ -235,7 +234,7 @@ impl BloomFilterIndexApplier { for (column_id, predicates) in predicates { let blob = match self - .blob_reader(file_id, *column_id, file_size_hint, metrics.as_deref_mut()) + .blob_reader(file_id, column_id, file_size_hint, metrics.as_deref_mut()) .await? { Some(blob) => blob, @@ -251,7 +250,7 @@ impl BloomFilterIndexApplier { let reader = CachedBloomFilterIndexBlobReader::new( file_id.file_id(), file_id.version, - *column_id, + column_id.clone(), Tag::Skipping, blob_size, BloomFilterReaderImpl::new(blob), @@ -296,7 +295,7 @@ impl BloomFilterIndexApplier { async fn blob_reader( &self, file_id: RegionIndexId, - column_id: ColumnId, + column_id: &IndexTarget, file_size_hint: Option, metrics: Option<&mut BloomFilterIndexApplyMetrics>, ) -> Result> { @@ -338,7 +337,7 @@ impl BloomFilterIndexApplier { async fn cached_blob_reader( &self, file_id: RegionIndexId, - column_id: ColumnId, + column_id: &IndexTarget, file_size_hint: Option, ) -> Result> { let Some(file_cache) = &self.file_cache else { @@ -358,7 +357,7 @@ impl BloomFilterIndexApplier { file_cache.local_store(), WriteCachePathProvider::new(file_cache.clone()), ); - let blob_name = Self::column_blob_name(column_id); + let blob_name = Self::target_blob_name(column_id); let reader = puffin_manager .reader(&file_id) @@ -375,8 +374,8 @@ impl BloomFilterIndexApplier { } // TODO(ruihang): use the same util with the code in creator - fn column_blob_name(column_id: ColumnId) -> String { - let target_key = IndexTarget::ColumnId(column_id).encode(); + fn target_blob_name(column_id: &IndexTarget) -> String { + let target_key = column_id.encode(); format!("{INDEX_BLOB_TYPE}-{target_key}") } @@ -384,7 +383,7 @@ impl BloomFilterIndexApplier { async fn remote_blob_reader( &self, file_id: RegionIndexId, - column_id: ColumnId, + column_id: &IndexTarget, file_size_hint: Option, ) -> Result { let path_factory = RegionFilePathFactory::new(self.table_dir.clone(), self.path_type); @@ -403,7 +402,7 @@ impl BloomFilterIndexApplier { .build(self.object_store.clone(), path_factory) .with_puffin_metadata_cache(self.puffin_metadata_cache.clone()); - let blob_name = Self::column_blob_name(column_id); + let blob_name = Self::target_blob_name(column_id); puffin_manager .reader(&file_id) @@ -439,22 +438,18 @@ impl BloomFilterIndexApplier { pub fn compatible_predicate_for_sst( &self, sst_metadata: &RegionMetadataRef, - ) -> Option>>> { + ) -> Option>>> { let mut has_type_mismatch = false; let mut compatible_col_ids = Vec::new(); - for (col_id, expected) in &self.expected_predicate_col_types { - let Some(sst_col) = sst_metadata.column_by_id(*col_id) else { - has_type_mismatch = true; - continue; - }; - - if sst_col.column_schema.data_type != *expected { + for (col_id, expected) in &self.expected_predicate_target_types { + if crate::sst::index::column::skipping_target_data_type(sst_metadata, col_id).as_ref() + != Some(expected) + { has_type_mismatch = true; continue; } - - compatible_col_ids.push(*col_id); + compatible_col_ids.push(col_id.clone()); } if compatible_col_ids.is_empty() { @@ -500,11 +495,196 @@ mod tests { use super::*; use crate::sst::file::RegionFileId; + use crate::sst::index::IndexBuildType; use crate::sst::index::bloom_filter::creator::BloomFilterIndexer; use crate::sst::index::bloom_filter::creator::tests::{ mock_object_store, mock_region_metadata, new_batch, new_intm_mgr, }; + #[tokio::test] + async fn test_json_skipping_build_and_prune_both_formats() { + use crate::read::{Batch, BatchColumn}; + use common_function::scalars::json::json_get::JsonGetWithType; + use common_function::scalars::udf::create_udf; + use datafusion_common::ScalarValue; + use datafusion_expr::expr::ScalarFunction; + use datatypes::arrow::array::{ArrayRef, Int64Array, StructArray}; + use datatypes::arrow::buffer::NullBuffer; + use datatypes::arrow::datatypes::{Field, Schema}; + use datatypes::arrow::record_batch::RecordBatch; + use datatypes::schema::SkippingIndexOptions; + use datatypes::vectors::{Helper, UInt8Vector, UInt64Vector}; + + for flat in [false, true] { + let (dir, factory) = + PuffinManagerFactory::new_for_test_async("json_skipping_index").await; + let mut metadata = (*mock_region_metadata()).clone(); + let column = &mut metadata.column_metadatas[2].column_schema; + column.data_type = + ConcreteDataType::json2(datatypes::types::json_type::JsonNativeType::Object( + datatypes::types::json_type::JsonObjectType::new(), + )); + column.with_extension_type(&datatypes::extension::json::Json2ExtensionType::new( + Arc::new(datatypes::extension::json::JsonMetadata::new( + datatypes::json::JsonSettings::try_new( + vec![datatypes::json::JsonTypeHint { + path: vec!["value".into()], + data_type: ConcreteDataType::int64_datatype(), + inverted_index: false, + skipping_index: Some( + SkippingIndexOptions::new(2, 0.01, Default::default()).unwrap(), + ), + }], + Some(0), + ) + .unwrap(), + )), + )); + // Other ordinary indexed columns are absent in this synthetic batch. + for col in &mut metadata.column_metadatas { + col.column_schema.unset_skipping_options().unwrap(); + } + let metadata = Arc::new(metadata); + let file_id = FileId::random(); + let target = IndexTarget::new_json_path(3, vec!["value".into()]).unwrap(); + let mut creator = BloomFilterIndexer::new( + file_id, + &metadata, + new_intm_mgr(dir.path().to_string_lossy()).await, + None, + IndexBuildType::Flush, + ) + .unwrap() + .unwrap(); + let values: ArrayRef = Arc::new(Int64Array::from(vec![Some(1), None, Some(1), None])); + let root: ArrayRef = Arc::new(StructArray::new( + vec![Field::new("value", values.data_type().clone(), true)].into(), + vec![values], + Some(NullBuffer::from(vec![true, false, true, true])), + )); + for start in [0, 2] { + let array = root.slice(start, 2); + if flat { + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "field_u64", + array.data_type().clone(), + true, + )])), + vec![array], + ) + .unwrap(); + creator.update_flat(&batch).await.unwrap(); + } else { + let mut batch = Batch::new( + vec![], + Arc::new(UInt64Vector::from_vec(vec![0, 1])), + Arc::new(UInt64Vector::from_vec(vec![0, 0])), + Arc::new(UInt8Vector::from_vec(vec![1, 1])), + vec![BatchColumn { + column_id: 3, + data: Helper::try_into_vector(array).unwrap(), + }], + ) + .unwrap(); + creator.update(&mut batch).await.unwrap(); + } + } + + assert_eq!(creator.indexed_targets().collect::>(), vec![&target]); + let object_store = mock_object_store(); + let manager = factory.build( + object_store.clone(), + RegionFilePathFactory::new("json_test".into(), PathType::Bare), + ); + let index_id = RegionIndexId::new(RegionFileId::new(metadata.region_id, file_id), 0); + let mut writer = manager.writer(&index_id).await.unwrap(); + let (rows, bytes) = creator.finish(&mut writer).await.unwrap(); + assert_eq!(rows, 4); + assert!(bytes > 0); + writer.finish().await.unwrap(); + let json_expr = |path: &str, ty| { + Expr::ScalarFunction(ScalarFunction::new_udf( + Arc::new(create_udf(Arc::new(JsonGetWithType::default()))), + vec![col("field_u64"), lit(path), Expr::Literal(ty, None)], + )) + }; + let cache = Arc::new( + crate::cache::index::bloom_filter_index::BloomFilterIndexCache::new( + 1024 * 1024, + 1024 * 1024, + 1024, + ), + ); + for expr in [ + json_expr("$.value", ScalarValue::Int64(None)).eq(lit(9_i64)), + json_expr("$.value", ScalarValue::Int64(None)) + .in_list(vec![lit(8_i64), lit(9_i64)], false), + ] { + let applier = BloomFilterIndexApplierBuilder::new( + "json_test".into(), + PathType::Bare, + object_store.clone(), + &metadata, + factory.clone(), + ) + .with_bloom_filter_index_cache(Some(cache.clone())) + .build(&[expr]) + .unwrap() + .unwrap(); + let predicates = applier.compatible_predicate_for_sst(&metadata).unwrap(); + assert!(predicates.contains_key(&target)); + // Run twice to cover cold reads and cached index reads. + for _ in 0..2 { + let result = applier + .apply(index_id, None, &predicates, [(4, true)].into_iter(), None) + .await + .unwrap(); + assert!(result[0].1.is_empty()); + } + assert!( + applier + .compatible_predicate_for_sst(&mock_region_metadata()) + .is_none() + ); + } + let applier = BloomFilterIndexApplierBuilder::new( + "json_test".into(), + PathType::Bare, + object_store.clone(), + &metadata, + factory.clone(), + ) + .build(&[json_expr("$.value", ScalarValue::Int64(None)).eq(lit(1_i64))]) + .unwrap() + .unwrap(); + let predicates = applier.compatible_predicate_for_sst(&metadata).unwrap(); + let result = applier + .apply(index_id, None, &predicates, [(4, true)].into_iter(), None) + .await + .unwrap(); + assert_eq!(result[0].1.len(), 1); + assert_eq!(result[0].1[0], 0..4); + for expr in [ + json_expr("$.value", ScalarValue::Float64(None)), + json_expr("$.missing", ScalarValue::Int64(None)), + ] { + assert!( + BloomFilterIndexApplierBuilder::new( + "json_test".into(), + PathType::Bare, + object_store.clone(), + &metadata, + factory.clone() + ) + .build(&[expr.eq(lit(1_i64))]) + .unwrap() + .is_none() + ); + } + } + } + #[tokio::test] async fn test_compatible_predicate_for_sst() { let (_d, puffin_manager_factory) = @@ -513,13 +693,15 @@ mod tests { let table_dir = "table_dir".to_string(); let predicates = BTreeMap::from_iter([( - 1, + IndexTarget::ColumnId(1), vec![InListPredicate { list: BTreeSet::from_iter([Bytes::from("foo")]), }], )]); - let expected_predicate_col_types = - BTreeMap::from_iter([(1, ConcreteDataType::string_datatype())]); + let expected_predicate_target_types = BTreeMap::from_iter([( + IndexTarget::ColumnId(1), + ConcreteDataType::string_datatype(), + )]); let applier = BloomFilterIndexApplier::new( table_dir, @@ -527,7 +709,7 @@ mod tests { object_store, puffin_manager_factory, predicates, - expected_predicate_col_types, + expected_predicate_target_types, ); let predicates = applier.compatible_predicate_for_sst(&mock_region_metadata()); assert!(predicates.is_some()); @@ -541,13 +723,13 @@ mod tests { let table_dir = "table_dir".to_string(); let predicates = BTreeMap::from_iter([( - 1, + IndexTarget::ColumnId(1), vec![InListPredicate { list: BTreeSet::from_iter([Bytes::from("foo")]), }], )]); - let expected_predicate_col_types = - BTreeMap::from_iter([(1, ConcreteDataType::int64_datatype())]); + let expected_predicate_target_types = + BTreeMap::from_iter([(IndexTarget::ColumnId(1), ConcreteDataType::int64_datatype())]); let applier = BloomFilterIndexApplier::new( table_dir, @@ -555,7 +737,7 @@ mod tests { object_store, puffin_manager_factory, predicates, - expected_predicate_col_types, + expected_predicate_target_types, ); let predicates = applier.compatible_predicate_for_sst(&mock_region_metadata()); assert!(predicates.is_none()); @@ -572,21 +754,24 @@ mod tests { // Column 3 (field_u64): expected int64 — SST has uint64 (mismatched). let predicates = BTreeMap::from_iter([ ( - 1, + IndexTarget::ColumnId(1), vec![InListPredicate { list: BTreeSet::from_iter([Bytes::from("foo")]), }], ), ( - 3, + IndexTarget::ColumnId(3), vec![InListPredicate { list: BTreeSet::from_iter([Bytes::from("bar")]), }], ), ]); - let expected_predicate_col_types = BTreeMap::from_iter([ - (1, ConcreteDataType::string_datatype()), - (3, ConcreteDataType::int64_datatype()), // intentional mismatch + let expected_predicate_target_types = BTreeMap::from_iter([ + ( + IndexTarget::ColumnId(1), + ConcreteDataType::string_datatype(), + ), + (IndexTarget::ColumnId(3), ConcreteDataType::int64_datatype()), // intentional mismatch ]); let applier = BloomFilterIndexApplier::new( @@ -595,18 +780,18 @@ mod tests { object_store, puffin_manager_factory, predicates, - expected_predicate_col_types, + expected_predicate_target_types, ); let result = applier.compatible_predicate_for_sst(&mock_region_metadata()); // The subset containing only the compatible column must be returned. let result = result.expect("expected Some with compatible subset"); assert!( - result.contains_key(&1), + result.contains_key(&IndexTarget::ColumnId(1)), "compatible column 1 must be present" ); assert!( - !result.contains_key(&3), + !result.contains_key(&IndexTarget::ColumnId(3)), "mismatched column 3 must be absent" ); assert_eq!(result.len(), 1, "only the compatible predicate must remain"); @@ -686,6 +871,7 @@ mod tests { ®ion_metadata, intm_mgr, memory_usage_threshold, + IndexBuildType::Flush, ) .unwrap() .unwrap(); diff --git a/src/mito2/src/sst/index/bloom_filter/applier/builder.rs b/src/mito2/src/sst/index/bloom_filter/applier/builder.rs index 4b8a11b0975..29012c951f8 100644 --- a/src/mito2/src/sst/index/bloom_filter/applier/builder.rs +++ b/src/mito2/src/sst/index/bloom_filter/applier/builder.rs @@ -15,13 +15,14 @@ use std::collections::{BTreeMap, BTreeSet}; use common_telemetry::warn; -use datafusion_common::{Column, ScalarValue}; +use datafusion_common::ScalarValue; use datafusion_expr::expr::InList; use datafusion_expr::{BinaryExpr, Expr, Operator}; use datatypes::data_type::ConcreteDataType; use datatypes::value::Value; use index::Bytes; use index::bloom_filter::applier::InListPredicate; +use index::target::IndexTarget; use mito_codec::index::IndexValueCodec; use mito_codec::row_converter::SortField; use object_store::ObjectStore; @@ -29,7 +30,6 @@ use puffin::puffin_manager::cache::PuffinMetadataCacheRef; use snafu::{OptionExt, ResultExt}; use store_api::metadata::RegionMetadata; use store_api::region_request::PathType; -use store_api::storage::ColumnId; use crate::cache::file_cache::FileCacheRef; use crate::cache::index::bloom_filter_index::BloomFilterIndexCacheRef; @@ -46,7 +46,7 @@ pub struct BloomFilterIndexApplierBuilder<'a> { file_cache: Option, puffin_metadata_cache: Option, bloom_filter_index_cache: Option, - predicates: BTreeMap>, + predicates: BTreeMap>, } impl<'a> BloomFilterIndexApplierBuilder<'a> { @@ -101,14 +101,14 @@ impl<'a> BloomFilterIndexApplierBuilder<'a> { return Ok(None); } - let expected_predicate_column_types = self.expected_predicate_column_types(); + let expected_predicate_target_types = self.expected_predicate_target_types(); let applier = BloomFilterIndexApplier::new( self.table_dir, self.path_type, self.object_store, self.puffin_manager_factory, self.predicates, - expected_predicate_column_types, + expected_predicate_target_types, ) .with_file_cache(self.file_cache) .with_puffin_metadata_cache(self.puffin_metadata_cache) @@ -139,22 +139,23 @@ impl<'a> BloomFilterIndexApplierBuilder<'a> { } } - /// Returns `(column_id, data_type)` pairs for predicate columns. - fn expected_predicate_column_types(&self) -> BTreeMap { + /// Returns the encoding type of each predicate target. + fn expected_predicate_target_types(&self) -> BTreeMap { self.predicates .keys() .filter_map(|col_id| { - let col = self.metadata.column_by_id(*col_id)?; - Some((*col_id, col.column_schema.data_type.clone())) + let ty = + crate::sst::index::column::skipping_target_data_type(self.metadata, col_id)?; + Some((col_id.clone(), ty)) }) .collect() } /// Helper function to get the column id and type - fn column_id_and_type( + fn column_target_and_type( &self, column_name: &str, - ) -> Result> { + ) -> Result> { let column = self .metadata .column_by_name(column_name) @@ -163,11 +164,20 @@ impl<'a> BloomFilterIndexApplierBuilder<'a> { })?; Ok(Some(( - column.column_id, + IndexTarget::ColumnId(column.column_id), column.column_schema.data_type.clone(), ))) } + fn expr_to_index_target(&self, expr: &Expr) -> Result> { + if let Expr::Column(column) = expr { + return self.column_target_and_type(&column.name); + } + crate::sst::index::column::json_index_target(self.metadata, expr, |hint| { + hint.skipping_index.is_some() + }) + } + /// Collects an equality expression (column = value) fn collect_eq(&mut self, left: &Expr, right: &Expr) -> Result<()> { let Some((col, lit)) = Self::eq_expr_col_lit(left, right)? else { @@ -176,7 +186,7 @@ impl<'a> BloomFilterIndexApplierBuilder<'a> { if lit.is_null() { return Ok(()); } - let Some((column_id, data_type)) = self.column_id_and_type(&col.name)? else { + let Some((column_id, data_type)) = self.expr_to_index_target(col)? else { return Ok(()); }; let value = encode_lit(lit, data_type)?; @@ -192,15 +202,11 @@ impl<'a> BloomFilterIndexApplierBuilder<'a> { /// Collects an in list expression in the form of `column IN (lit, lit, ...)`. fn collect_in_list(&mut self, in_list: &InList) -> Result<()> { - // Only collect InList predicates if they reference a column - let Expr::Column(column) = &in_list.expr.as_ref() else { - return Ok(()); - }; if in_list.list.is_empty() || in_list.negated { return Ok(()); } - let Some((column_id, data_type)) = self.column_id_and_type(&column.name)? else { + let Some((column_id, data_type)) = self.expr_to_index_target(&in_list.expr)? else { return Ok(()); }; @@ -252,13 +258,13 @@ impl<'a> BloomFilterIndexApplierBuilder<'a> { if lit.is_null() { return Ok(()); } - let Some((column_id, data_type)) = self.column_id_and_type(&col.name)? else { + let Some((column_id, data_type)) = self.expr_to_index_target(col)? else { return Ok(()); }; let mut inlist = BTreeSet::new(); inlist.insert(encode_lit(lit, data_type.clone())?); - if Self::collect_or_eq_list_rec(&col.name, &data_type, or_list, &mut inlist)? { + if Self::collect_or_eq_list_rec(col, &data_type, or_list, &mut inlist)? { self.predicates .entry(column_id) .or_default() @@ -269,7 +275,7 @@ impl<'a> BloomFilterIndexApplierBuilder<'a> { } fn collect_or_eq_list_rec( - column_name: &str, + column_expr: &Expr, data_type: &ConcreteDataType, expr: &Expr, inlist: &mut BTreeSet, @@ -277,9 +283,9 @@ impl<'a> BloomFilterIndexApplierBuilder<'a> { if let Expr::BinaryExpr(BinaryExpr { left, op, right }) = expr { match op { Operator::Or => { - let r = Self::collect_or_eq_list_rec(column_name, data_type, left, inlist)? + let r = Self::collect_or_eq_list_rec(column_expr, data_type, left, inlist)? .then(|| { - Self::collect_or_eq_list_rec(column_name, data_type, right, inlist) + Self::collect_or_eq_list_rec(column_expr, data_type, right, inlist) }) .transpose()? .unwrap_or(false); @@ -289,7 +295,7 @@ impl<'a> BloomFilterIndexApplierBuilder<'a> { let Some((col, lit)) = Self::eq_expr_col_lit(left, right)? else { return Ok(false); }; - if lit.is_null() || column_name != col.name { + if lit.is_null() || column_expr != col { return Ok(false); } let bytes = encode_lit(lit, data_type.clone())?; @@ -307,10 +313,10 @@ impl<'a> BloomFilterIndexApplierBuilder<'a> { fn eq_expr_col_lit<'b>( left: &'b Expr, right: &'b Expr, - ) -> Result> { + ) -> Result> { let (col, lit) = match (left, right) { - (Expr::Column(col), Expr::Literal(lit, _)) => (col, lit), - (Expr::Literal(lit, _), Expr::Column(col)) => (col, lit), + (expr, Expr::Literal(lit, _)) => (expr, lit), + (Expr::Literal(lit, _), expr) => (expr, lit), _ => return Ok(None), }; Ok(Some((col, lit))) @@ -404,7 +410,7 @@ mod tests { let predicates = result.unwrap().default_predicates; assert_eq!(predicates.len(), 1); - let column_predicates = predicates.get(&1).unwrap(); + let column_predicates = predicates.get(&IndexTarget::ColumnId(1)).unwrap(); assert_eq!(column_predicates.len(), 1); let expected = encode_lit( @@ -419,7 +425,9 @@ mod tests { i.lit() } - fn build_bloom_predicates(exprs: &[Expr]) -> Option>> { + fn build_bloom_predicates( + exprs: &[Expr], + ) -> Option>> { let (_d, factory) = PuffinManagerFactory::new_for_test_block("bloom_builder_"); let metadata = test_region_metadata(); BloomFilterIndexApplierBuilder::new( @@ -460,7 +468,7 @@ mod tests { assert_eq!( build_bloom_predicates(&[expr]), Some(BTreeMap::from([( - 2, + IndexTarget::ColumnId(2), vec![int64_inlist_predicate([1, 2, 3])] )])) ); @@ -491,7 +499,10 @@ mod tests { assert_eq!( build_bloom_predicates(&[expr]), - Some(BTreeMap::from([(2, vec![int64_inlist_predicate([1, 3])])])) + Some(BTreeMap::from([( + IndexTarget::ColumnId(2), + vec![int64_inlist_predicate([1, 3])] + )])) ); } @@ -523,7 +534,10 @@ mod tests { assert_eq!( build_bloom_predicates(&[expr]), - Some(BTreeMap::from([(2, vec![int64_inlist_predicate([42])])])) + Some(BTreeMap::from([( + IndexTarget::ColumnId(2), + vec![int64_inlist_predicate([42])] + )])) ); } @@ -560,7 +574,7 @@ mod tests { assert!(result.is_some()); let predicates = result.unwrap().default_predicates; - let column_predicates = predicates.get(&2).unwrap(); + let column_predicates = predicates.get(&IndexTarget::ColumnId(2)).unwrap(); assert_eq!(column_predicates.len(), 1); assert_eq!(column_predicates[0].list.len(), 3); } @@ -590,7 +604,7 @@ mod tests { assert!(result.is_some()); let predicates = result.unwrap().default_predicates; - let column_predicates = predicates.get(&1).unwrap(); + let column_predicates = predicates.get(&IndexTarget::ColumnId(1)).unwrap(); assert_eq!(column_predicates.len(), 1); assert_eq!(column_predicates[0].list.len(), 4); let or_chain_predicates = &column_predicates[0].list; @@ -655,8 +669,8 @@ mod tests { let predicates = result.unwrap().default_predicates; assert_eq!(predicates.len(), 2); - assert!(predicates.contains_key(&1)); - assert!(predicates.contains_key(&2)); + assert!(predicates.contains_key(&IndexTarget::ColumnId(1))); + assert!(predicates.contains_key(&IndexTarget::ColumnId(2))); } #[test] @@ -692,8 +706,8 @@ mod tests { assert!(result.is_some()); let predicates = result.unwrap().default_predicates; - assert!(!predicates.contains_key(&1)); // Null equality should be ignored - let column2_predicates = predicates.get(&2).unwrap(); + assert!(!predicates.contains_key(&IndexTarget::ColumnId(1))); // Null equality should be ignored + let column2_predicates = predicates.get(&IndexTarget::ColumnId(2)).unwrap(); assert_eq!(column2_predicates[0].list.len(), 2); } @@ -761,7 +775,7 @@ mod tests { assert!(result.is_some()); let predicates = result.unwrap().default_predicates; - let column_predicates = predicates.get(&1).unwrap(); + let column_predicates = predicates.get(&IndexTarget::ColumnId(1)).unwrap(); assert_eq!(column_predicates.len(), 2); } } diff --git a/src/mito2/src/sst/index/bloom_filter/creator.rs b/src/mito2/src/sst/index/bloom_filter/creator.rs index ea727f1b7b4..cc96e596866 100644 --- a/src/mito2/src/sst/index/bloom_filter/creator.rs +++ b/src/mito2/src/sst/index/bloom_filter/creator.rs @@ -17,9 +17,10 @@ use std::sync::Arc; use std::sync::atomic::AtomicUsize; use api::v1::SemanticType; +use arrow_schema::extension::ExtensionType; use common_telemetry::{debug, warn}; use datatypes::arrow::record_batch::RecordBatch; -use datatypes::schema::SkippingIndexType; +use datatypes::extension::json::Json2ExtensionType; use datatypes::vectors::Helper; use index::bloom_filter::creator::BloomFilterCreator; use index::target::IndexTarget; @@ -28,7 +29,7 @@ use mito_codec::row_converter::sparse::SparsePrimaryKeyView; use mito_codec::row_converter::{SortField, SparseOffsetsCache}; use puffin::puffin_manager::{PuffinWriter, PutOptions}; use smallvec::SmallVec; -use snafu::{ResultExt, ensure}; +use snafu::{OptionExt, ResultExt, ensure}; use store_api::codec::PrimaryKeyEncoding; use store_api::metadata::RegionMetadataRef; use store_api::storage::{ColumnId, FileId}; @@ -39,15 +40,15 @@ use crate::error::{ OperateAbortedIndexSnafu, PuffinAddBlobSnafu, PushBloomFilterValueSnafu, Result, }; use crate::read::Batch; -use crate::sst::index::TYPE_BLOOM_FILTER_INDEX; use crate::sst::index::bloom_filter::INDEX_BLOB_TYPE; -use crate::sst::index::column::column_index_rows; +use crate::sst::index::column::{column_index_rows, extract_json_type_hint_values}; use crate::sst::index::intermediate::{ IntermediateLocation, IntermediateManager, TempFileProvider, }; use crate::sst::index::primary_key::PrimaryKeyRuns; use crate::sst::index::puffin_manager::SstPuffinWriter; use crate::sst::index::statistics::{ByteCount, RowCount, Statistics}; +use crate::sst::index::{IndexBuildType, TYPE_BLOOM_FILTER_INDEX}; /// The buffer size for the pipe used to send index data to the puffin blob. const PIPE_BUFFER_SIZE_FOR_SENDING_BLOB: usize = 8192; @@ -55,7 +56,7 @@ const PIPE_BUFFER_SIZE_FOR_SENDING_BLOB: usize = 8192; /// The indexer for the bloom filter index. pub struct BloomFilterIndexer { /// The bloom filter creators. - creators: HashMap, + creators: HashMap, /// The provider for intermediate files. temp_file_provider: Arc, @@ -82,11 +83,13 @@ pub struct BloomFilterIndexer { impl BloomFilterIndexer { /// Creates a new bloom filter indexer. + /// JSON paths are included only for flush and compaction builds. pub fn new( sst_file_id: FileId, metadata: &RegionMetadataRef, intermediate_manager: IntermediateManager, memory_usage_threshold: Option, + build_type: IndexBuildType, ) -> Result> { let mut creators = HashMap::new(); @@ -105,19 +108,39 @@ impl BloomFilterIndexer { column_name: &column.column_schema.name, })?; - let options = match options { - Some(options) if options.index_type == SkippingIndexType::BloomFilter => options, - _ => continue, - }; - - let creator = BloomFilterCreator::new( - options.granularity as _, - options.false_positive_rate(), - temp_file_provider.clone(), - global_memory_usage.clone(), - memory_usage_threshold, - ); - creators.insert(column.column_id, creator); + let mut targets = Vec::new(); + if let Some(options) = options { + targets.push((IndexTarget::ColumnId(column.column_id), options)); + } + // Only flush and compaction builds collect JSON targets; manual and schema-change + // rebuilds skip them because existing SSTs may not materialize the current hints. + // FIXME: Rebuilds replace the whole Puffin file and can drop existing JSON path indexes. + if matches!(build_type, IndexBuildType::Flush | IndexBuildType::Compact) + && column.column_schema.data_type.is_json2() + && let Ok(Some(extension)) = + column.column_schema.extension_type::() + { + for hint in extension.metadata().json_settings().type_hints() { + let Some(options) = &hint.skipping_index else { + continue; + }; + if let Ok(target) = + IndexTarget::new_json_path(column.column_id, hint.path.clone()) + { + targets.push((target, options.clone())); + } + } + } + for (target, options) in targets { + let creator = BloomFilterCreator::new( + options.granularity as _, + options.false_positive_rate(), + temp_file_provider.clone(), + global_memory_usage.clone(), + memory_usage_threshold, + ); + creators.insert(target, creator); + } } if creators.is_empty() { @@ -237,33 +260,33 @@ impl BloomFilterIndexer { let n = batch.num_rows(); guard.inc_row_count(n); - for (col_id, creator) in &mut self.creators { - match self.codec.pk_col_info(*col_id) { - // tags - Some(col_info) => { - let pk_idx = col_info.idx; - let field = &col_info.field; - let elems = batch - .pk_col_value(self.codec.decoder(), pk_idx, *col_id)? - .filter(|v| !v.is_null()) - .map(|v| { - let mut buf = vec![]; - IndexValueCodec::encode_nonnull_value( - v.as_value_ref(), - field, - &mut buf, - ) - .context(EncodeSnafu)?; - Ok(buf) - }) - .transpose()?; - creator - .push_n_row_elems(n, elems) - .await - .context(PushBloomFilterValueSnafu)?; - } - // fields - None => { + for (target, creator) in &mut self.creators { + let (vector, field) = match target { + IndexTarget::ColumnId(col_id) => { + if let Some(col_info) = self.codec.pk_col_info(*col_id) { + let pk_idx = col_info.idx; + let field = &col_info.field; + let elems = batch + .pk_col_value(self.codec.decoder(), pk_idx, *col_id)? + .filter(|v| !v.is_null()) + .map(|v| { + let mut buf = vec![]; + IndexValueCodec::encode_nonnull_value( + v.as_value_ref(), + field, + &mut buf, + ) + .context(EncodeSnafu)?; + Ok(buf) + }) + .transpose()?; + creator + .push_n_row_elems(n, elems) + .await + .context(PushBloomFilterValueSnafu)?; + continue; + } + let Some(values) = batch.field_col_value(*col_id) else { debug!( "Column {} not found in the batch during building bloom filter index", @@ -271,24 +294,33 @@ impl BloomFilterIndexer { ); continue; }; - let sort_field = SortField::new(values.data.data_type()); - for i in 0..n { - let value = values.data.get_ref(i); - let elems = (!value.is_null()) - .then(|| { - let mut buf = vec![]; - IndexValueCodec::encode_nonnull_value(value, &sort_field, &mut buf) - .context(EncodeSnafu)?; - Ok(buf) - }) - .transpose()?; - - creator - .push_row_elems(elems) - .await - .context(PushBloomFilterValueSnafu)?; - } + let field = SortField::new(values.data.data_type()); + (values.data.clone(), field) } + IndexTarget::JsonPath(path) => { + let root = batch + .field_col_value(path.column_id()) + .with_context(|| crate::error::InvalidRecordBatchSnafu { + reason: format!( + "Missing JSON index column {} in batch", + path.column_id() + ), + })? + .data + .to_arrow_array(); + let vector = extract_json_type_hint_values(&root, path.path())?; + let field = SortField::new(vector.data_type()); + (vector, field) + } + }; + for row in 0..n { + let value = + IndexValueCodec::encode_value(vector.get_ref(row), &field, &mut self.value_buf) + .context(EncodeSnafu)?; + creator + .push_n_row_elem(1, value) + .await + .context(PushBloomFilterValueSnafu)?; } } @@ -305,14 +337,23 @@ impl BloomFilterIndexer { let mut sparse_columns: SmallVec<[(ColumnId, &mut BloomFilterCreator); 8]> = SmallVec::new(); - for (col_id, creator) in &mut self.creators { + for (target, creator) in &mut self.creators { + let col_id = match target { + IndexTarget::ColumnId(id) => *id, + IndexTarget::JsonPath(path) => path.column_id(), + }; // Safety: `creators` are created from the metadata so it won't be None. - let column_meta = self.metadata.column_by_id(*col_id).unwrap(); + let column_meta = self.metadata.column_by_id(col_id).unwrap(); let column_name = &column_meta.column_schema.name; if let Some(column_array) = batch.column_by_name(column_name) { // Convert Arrow array to VectorRef - let vector = Helper::try_into_vector(column_array.clone()) - .context(crate::error::ConvertVectorSnafu)?; + let vector = match target { + IndexTarget::ColumnId(_) => Helper::try_into_vector(column_array.clone()) + .context(crate::error::ConvertVectorSnafu)?, + IndexTarget::JsonPath(path) => { + extract_json_type_hint_values(column_array, path.path())? + } + }; let sort_field = SortField::new(vector.data_type()); for (row, count) in column_index_rows(batch, column_meta.semantic_type) { @@ -328,9 +369,14 @@ impl BloomFilterIndexer { .await .context(PushBloomFilterValueSnafu)?; } + } else if matches!(target, IndexTarget::JsonPath(_)) { + return crate::error::InvalidRecordBatchSnafu { + reason: format!("Missing JSON index column {column_name} in batch"), + } + .fail(); } else if is_sparse && column_meta.semantic_type == SemanticType::Tag { - if self.codec.pk_col_info(*col_id).is_some() { - sparse_columns.push((*col_id, creator)); + if self.codec.pk_col_info(col_id).is_some() { + sparse_columns.push((col_id, creator)); } } else { debug!( @@ -402,13 +448,13 @@ impl BloomFilterIndexer { /// /// TODO(zhongzc): duplicate with `mito2::sst::index::inverted_index::creator::InvertedIndexCreator` async fn do_finish_single_creator( - col_id: &ColumnId, + target: &IndexTarget, creator: &mut BloomFilterCreator, puffin_writer: &mut SstPuffinWriter, ) -> Result { let (tx, rx) = tokio::io::duplex(PIPE_BUFFER_SIZE_FOR_SENDING_BLOB); - let target_key = IndexTarget::ColumnId(*col_id).encode(); + let target_key = target.encode(); let blob_name = format!("{INDEX_BLOB_TYPE}-{target_key}"); let (index_finish, puffin_add_blob) = futures::join!( creator.finish(tx.compat_write()), @@ -446,9 +492,9 @@ impl BloomFilterIndexer { .load(std::sync::atomic::Ordering::Relaxed) } - /// Returns the column ids to be indexed. - pub fn column_ids(&self) -> impl Iterator + use<'_> { - self.creators.keys().copied() + /// Returns the targets indexed by this creator. + pub fn indexed_targets(&self) -> impl Iterator { + self.creators.keys() } } @@ -457,7 +503,7 @@ pub(crate) mod tests { use api::v1::SemanticType; use datatypes::data_type::ConcreteDataType; - use datatypes::schema::{ColumnSchema, SkippingIndexOptions}; + use datatypes::schema::{ColumnSchema, SkippingIndexOptions, SkippingIndexType}; use datatypes::value::ValueRef; use datatypes::vectors::{UInt8Vector, UInt64Vector}; use index::bloom_filter::reader::{BloomFilterReader, BloomFilterReaderImpl}; @@ -599,10 +645,15 @@ pub(crate) mod tests { let memory_usage_threshold = Some(1024); let file_id = FileId::random(); - let mut indexer = - BloomFilterIndexer::new(file_id, ®ion_metadata, intm_mgr, memory_usage_threshold) - .unwrap() - .unwrap(); + let mut indexer = BloomFilterIndexer::new( + file_id, + ®ion_metadata, + intm_mgr, + memory_usage_threshold, + IndexBuildType::Flush, + ) + .unwrap() + .unwrap(); // push 20 rows let mut batch = new_batch("tag1", 0..10); diff --git a/src/mito2/src/sst/index/column.rs b/src/mito2/src/sst/index/column.rs index 2818efbeeda..ea1993fc33c 100644 --- a/src/mito2/src/sst/index/column.rs +++ b/src/mito2/src/sst/index/column.rs @@ -13,10 +13,17 @@ // limitations under the License. use api::v1::SemanticType; +use arrow_schema::extension::ExtensionType; +use common_function::scalars::json::json_get::{JsonGetWithType, json_object_path}; +use datafusion_expr::Expr; use datatypes::arrow::array::{Array, ArrayRef, StructArray}; use datatypes::arrow::record_batch::RecordBatch; +use datatypes::data_type::{ConcreteDataType, DataType}; +use datatypes::extension::json::Json2ExtensionType; use datatypes::vectors::{Helper, VectorRef}; +use index::target::IndexTarget; use snafu::{OptionExt, ResultExt}; +use store_api::metadata::RegionMetadata; use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME; use crate::error::{ConvertVectorSnafu, InvalidRecordBatchSnafu, Result}; @@ -83,6 +90,109 @@ pub(crate) fn column_index_rows( }) } +/// Resolves a typed JSON access when its hint enables the requested index. +pub(crate) fn json_index_target( + metadata: &RegionMetadata, + expr: &Expr, + enabled: impl Fn(&datatypes::json::JsonTypeHint) -> bool, +) -> Result> { + let Expr::ScalarFunction(func) = expr else { + return Ok(None); + }; + + if func.func.name() != JsonGetWithType::NAME || !(2..=3).contains(&func.args.len()) { + return Ok(None); + } + + let Expr::Column(column) = &func.args[0] else { + return Ok(None); + }; + + let Some(col) = metadata.column_by_name(&column.name) else { + return Ok(None); + }; + + if !col.column_schema.data_type.is_json2() { + return Ok(None); + } + + let Some(path_literal) = func.args[1] + .as_literal() + .and_then(|v| v.try_as_str()) + .flatten() + else { + return Ok(None); + }; + + let Some(path) = json_object_path(path_literal).ok().flatten() else { + return Ok(None); + }; + + let Ok(Some(extension)) = col.column_schema.extension_type::() else { + return Ok(None); + }; + + let Some(hint) = extension + .metadata() + .json_settings() + .type_hints() + .iter() + .find(|hint| hint.path == path && enabled(hint)) + else { + return Ok(None); + }; + + if let Some(type_arg) = func.args.get(2) { + let Expr::Literal(value, _) = type_arg else { + return Ok(None); + }; + if !value.is_null() { + return Ok(None); + } + // A type conversion may change comparison semantics and invalidate index pruning. + let result_type = value.data_type(); + if result_type != hint.data_type.as_arrow_type() + && !(result_type.is_string() && hint.data_type.is_string()) + { + return Ok(None); + } + } + + if !(hint.data_type.is_numeric() || hint.data_type.is_boolean() || hint.data_type.is_string()) { + return Ok(None); + } + Ok(IndexTarget::new_json_path(col.column_id, path) + .ok() + .map(|target| (target, hint.data_type.clone()))) +} + +/// Resolves a skipping target's encoding type from an SST's own hints. +pub(crate) fn skipping_target_data_type( + metadata: &RegionMetadata, + target: &IndexTarget, +) -> Option { + let path = match target { + IndexTarget::ColumnId(id) => { + return metadata + .column_by_id(*id) + .map(|col| col.column_schema.data_type.clone()); + } + IndexTarget::JsonPath(path) => path, + }; + let column = metadata.column_by_id(path.column_id())?; + let extension = column + .column_schema + .extension_type::() + .ok()??; + extension + .metadata() + .json_settings() + .type_hints() + .iter() + .find(|hint| hint.path == path.path() && hint.skipping_index.is_some()) + .map(|hint| hint.data_type.clone()) +} + #[cfg(test)] mod tests { use std::sync::Arc; diff --git a/src/mito2/src/sst/index/fulltext_index/applier.rs b/src/mito2/src/sst/index/fulltext_index/applier.rs index 4a18835dab4..792c6319f85 100644 --- a/src/mito2/src/sst/index/fulltext_index/applier.rs +++ b/src/mito2/src/sst/index/fulltext_index/applier.rs @@ -432,7 +432,7 @@ impl FulltextIndexApplier { let reader = CachedBloomFilterIndexBlobReader::new( file_id.file_id(), file_id.version, - column_id, + IndexTarget::ColumnId(column_id), Tag::Fulltext, blob_size, BloomFilterReaderImpl::new(range_reader), diff --git a/src/mito2/src/sst/index/indexer/finish.rs b/src/mito2/src/sst/index/indexer/finish.rs index 1a6d90ac4f5..801b3a8eeb2 100644 --- a/src/mito2/src/sst/index/indexer/finish.rs +++ b/src/mito2/src/sst/index/indexer/finish.rs @@ -13,6 +13,7 @@ // limitations under the License. use common_telemetry::{debug, warn}; +use index::target::IndexTarget; use puffin::puffin_manager::{PuffinManager, PuffinWriter}; use store_api::storage::ColumnId; @@ -126,7 +127,13 @@ impl Indexer { return true; }; - let column_ids = indexer.indexed_column_ids().collect(); + let column_ids = indexer + .indexed_targets() + .filter_map(|target| match target { + IndexTarget::ColumnId(id) => Some(*id), + IndexTarget::JsonPath(_) => None, + }) + .collect(); let err = match indexer.finish(puffin_writer).await { Ok((row_count, byte_count)) => { self.fill_inverted_index_output( @@ -203,7 +210,13 @@ impl Indexer { return true; }; - let column_ids = indexer.column_ids().collect(); + let column_ids = indexer + .indexed_targets() + .filter_map(|target| match target { + IndexTarget::ColumnId(id) => Some(*id), + IndexTarget::JsonPath(_) => None, + }) + .collect(); let err = match indexer.finish(puffin_writer).await { Ok((row_count, byte_count)) => { self.fill_bloom_filter_output( diff --git a/src/mito2/src/sst/index/inverted_index/applier/builder.rs b/src/mito2/src/sst/index/inverted_index/applier/builder.rs index 62d6f69e5b4..501f90a082e 100644 --- a/src/mito2/src/sst/index/inverted_index/applier/builder.rs +++ b/src/mito2/src/sst/index/inverted_index/applier/builder.rs @@ -20,12 +20,13 @@ mod regex_match; use std::collections::{BTreeMap, HashSet}; -use arrow_schema::extension::ExtensionType; -use common_function::scalars::json::json_get::{JsonGetWithType, json_object_path}; +#[cfg(test)] +use common_function::scalars::json::json_get::JsonGetWithType; use common_telemetry::warn; use datafusion_common::ScalarValue; use datafusion_expr::{BinaryExpr, Expr, Operator}; -use datatypes::data_type::{ConcreteDataType, DataType}; +use datatypes::data_type::ConcreteDataType; +#[cfg(test)] use datatypes::extension::json::Json2ExtensionType; use datatypes::value::Value; use index::inverted_index::search::predicate::Predicate; @@ -242,77 +243,9 @@ impl<'a> InvertedIndexApplierBuilder<'a> { .map(|(id, ty)| (IndexTarget::ColumnId(id), ty))); } - let Expr::ScalarFunction(func) = expr else { - return Ok(None); - }; - - if func.func.name() != JsonGetWithType::NAME || !(2..=3).contains(&func.args.len()) { - return Ok(None); - } - - let Some(name) = Self::column_name(&func.args[0]) else { - return Ok(None); - }; - - let Some(col) = self.metadata.column_by_name(name) else { - return Ok(None); - }; - - if !col.column_schema.data_type.is_json2() { - return Ok(None); - } - - let Some(path_literal) = func.args[1] - .as_literal() - .and_then(|v| v.try_as_str()) - .flatten() - else { - return Ok(None); - }; - - let Some(path) = json_object_path(path_literal).ok().flatten() else { - return Ok(None); - }; - - let Ok(Some(extension)) = col.column_schema.extension_type::() else { - return Ok(None); - }; - - let Some(hint) = extension - .metadata() - .json_settings() - .type_hints() - .iter() - .find(|hint| hint.path == path && hint.inverted_index) - else { - return Ok(None); - }; - - if let Some(type_arg) = func.args.get(2) { - let Expr::Literal(value, _) = type_arg else { - return Ok(None); - }; - if !value.is_null() { - return Ok(None); - } - // A type conversion may change comparison semantics and invalidate index pruning. - let result_type = value.data_type(); - if result_type != hint.data_type.as_arrow_type() - && !(result_type.is_string() && hint.data_type.is_string()) - { - return Ok(None); - } - } - - if !(hint.data_type.is_numeric() - || hint.data_type.is_boolean() - || hint.data_type.is_string()) - { - return Ok(None); - } - Ok(IndexTarget::new_json_path(col.column_id, path) - .ok() - .map(|target| (target, hint.data_type.clone()))) + crate::sst::index::column::json_index_target(self.metadata, expr, |hint| { + hint.inverted_index + }) } /// Helper function to get a non-null literal. @@ -456,12 +389,14 @@ mod tests { path: vec!["n".into()], data_type: ConcreteDataType::float64_datatype(), inverted_index: true, + skipping_index: None, }], vec![], vec![JsonTypeHint { path: vec!["n".into()], data_type: ConcreteDataType::int64_datatype(), inverted_index: false, + skipping_index: None, }], ] { let mut sst = metadata.clone(); @@ -720,11 +655,13 @@ mod tests { path: path.into_iter().map(String::from).collect(), data_type: ConcreteDataType::string_datatype(), inverted_index, + skipping_index: None, }) .chain(std::iter::once(JsonTypeHint { path: vec!["n".into()], data_type: ConcreteDataType::int64_datatype(), inverted_index: true, + skipping_index: None, })) .collect(); column.with_extension_type(&Json2ExtensionType::new(Arc::new(JsonMetadata::new( diff --git a/src/mito2/src/sst/index/inverted_index/creator.rs b/src/mito2/src/sst/index/inverted_index/creator.rs index 28242e79105..239fb608c1d 100644 --- a/src/mito2/src/sst/index/inverted_index/creator.rs +++ b/src/mito2/src/sst/index/inverted_index/creator.rs @@ -451,14 +451,9 @@ impl InvertedIndexer { self.temp_file_provider.cleanup().await } - /// Returns column-level index IDs, excluding JSON path targets. - pub fn indexed_column_ids(&self) -> impl Iterator + '_ { - self.index_targets - .iter() - .filter_map(|(target, _)| match target { - IndexTarget::ColumnId(col_id) => Some(*col_id), - IndexTarget::JsonPath(_) => None, - }) + /// Returns the targets indexed by this creator. + pub fn indexed_targets(&self) -> impl Iterator { + self.index_targets.iter().map(|(target, _)| target) } pub fn memory_usage(&self) -> usize { @@ -718,6 +713,7 @@ mod tests { path: vec!["value".into()], data_type: ConcreteDataType::int64_datatype(), inverted_index: true, + skipping_index: None, }], Some(0), ) @@ -769,8 +765,7 @@ mod tests { creator.update(&mut batch).await.unwrap(); } } - // JSON paths must not be advertised as ordinary column indexes. - assert_eq!(creator.indexed_column_ids().count(), 0); + assert_eq!(creator.indexed_targets().collect::>(), vec![&target]); let store = mock_object_store(); let manager = factory.build( store.clone(), diff --git a/src/mito2/src/sst/parquet/json_align.rs b/src/mito2/src/sst/parquet/json_align.rs index c03c7f1d62b..aa891dcef27 100644 --- a/src/mito2/src/sst/parquet/json_align.rs +++ b/src/mito2/src/sst/parquet/json_align.rs @@ -836,6 +836,7 @@ mod tests { path: vec!["kind".to_string()], data_type: ConcreteDataType::string_datatype(), inverted_index: false, + skipping_index: None, }], Some(0), ) diff --git a/src/operator/src/expr_helper.rs b/src/operator/src/expr_helper.rs index 563eb00789e..1fe7be2e51f 100644 --- a/src/operator/src/expr_helper.rs +++ b/src/operator/src/expr_helper.rs @@ -806,6 +806,17 @@ fn json_settings_to_proto(settings: JsonSettings) -> Result { data_type: data_type as i32, datatype_extension, inverted_index: hint.inverted_index, + skipping_index: hint.skipping_index.map(|options| { + api::v1::JsonSkippingIndexOptions { + granularity: options.granularity, + index_type: match options.index_type { + SkippingIndexType::BloomFilter => { + PbSkippingIndexType::BloomFilter.into() + } + }, + false_positive_rate: options.false_positive_rate(), + } + }), }) }) .collect::>>()?; @@ -1852,7 +1863,7 @@ SELECT max(c1), min(c2) FROM schema_2.table_2;"; #[test] fn test_to_alter_set_json_settings_expr() { - let sql = "ALTER TABLE monitor MODIFY COLUMN payload JSON2 (service STRING INVERTED INDEX, message STRING);"; + let sql = "ALTER TABLE monitor MODIFY COLUMN payload JSON2 (service STRING INVERTED INDEX SKIPPING INDEX WITH (granularity='8'), message STRING);"; let stmt = ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) .unwrap() @@ -1874,6 +1885,14 @@ SELECT max(c1), min(c2) FROM schema_2.table_2;"; assert_eq!(2, settings.type_hints.len()); assert_eq!(["service"], &settings.type_hints[0].path[..]); assert!(settings.type_hints[0].inverted_index); + assert_eq!( + settings.type_hints[0] + .skipping_index + .as_ref() + .unwrap() + .granularity, + 8 + ); assert!(!settings.type_hints[1].inverted_index); let sql = "ALTER TABLE monitor MODIFY COLUMN payload JSON2;"; diff --git a/src/pipeline/src/etl/transform.rs b/src/pipeline/src/etl/transform.rs index 0869a4be0fc..09cb0a3aa0b 100644 --- a/src/pipeline/src/etl/transform.rs +++ b/src/pipeline/src/etl/transform.rs @@ -538,6 +538,7 @@ fn parse_json2_type_hint(value: &yaml_rust::Yaml) -> Result { path, data_type, inverted_index, + skipping_index: None, }) } diff --git a/src/pipeline/src/etl/transform/transformer/greptime.rs b/src/pipeline/src/etl/transform/transformer/greptime.rs index a2f2fc2101e..87430cd1f98 100644 --- a/src/pipeline/src/etl/transform/transformer/greptime.rs +++ b/src/pipeline/src/etl/transform/transformer/greptime.rs @@ -1190,6 +1190,7 @@ transform: path: vec!["age".to_string()], data_type: ConcreteDataType::int64_datatype(), inverted_index: false, + skipping_index: None, }], None, ) @@ -1205,6 +1206,7 @@ transform: path: vec!["age".to_string()], data_type: ConcreteDataType::string_datatype(), inverted_index: false, + skipping_index: None, }], None, ) diff --git a/src/pipeline/src/etl/transform/transformer/greptime/coerce.rs b/src/pipeline/src/etl/transform/transformer/greptime/coerce.rs index a63b5ca076e..2e350b3b9a8 100644 --- a/src/pipeline/src/etl/transform/transformer/greptime/coerce.rs +++ b/src/pipeline/src/etl/transform/transformer/greptime/coerce.rs @@ -774,6 +774,7 @@ mod tests { path: vec!["age".to_string()], data_type: ConcreteDataType::int64_datatype(), inverted_index: false, + skipping_index: None, }], None, ) diff --git a/src/query/src/datafusion/json_expr_planner.rs b/src/query/src/datafusion/json_expr_planner.rs index a60e6155445..79af2405ff0 100644 --- a/src/query/src/datafusion/json_expr_planner.rs +++ b/src/query/src/datafusion/json_expr_planner.rs @@ -753,6 +753,7 @@ mod tests { path: vec!["payload".to_string(), "cpu".to_string()], data_type: ConcreteDataType::int64_datatype(), inverted_index: false, + skipping_index: None, }], None, ) @@ -787,11 +788,13 @@ mod tests { path: vec!["payload".to_string(), "cpu".to_string()], data_type: ConcreteDataType::int64_datatype(), inverted_index: false, + skipping_index: None, }, JsonTypeHint { path: vec!["payload.cpu".to_string()], data_type: ConcreteDataType::string_datatype(), inverted_index: false, + skipping_index: None, }, ], None, diff --git a/src/sql/src/parsers/create_parser/json.rs b/src/sql/src/parsers/create_parser/json.rs index 204f294a82f..d13ea0fa53f 100644 --- a/src/sql/src/parsers/create_parser/json.rs +++ b/src/sql/src/parsers/create_parser/json.rs @@ -157,6 +157,7 @@ fn parse_json2_type_hint(parser: &mut Parser<'_>) -> Result { let data_type = validate_json2_type_hint_type(data_type)?; let mut inverted_index = false; + let mut skipping_index = None; loop { if parser.parse_keywords(&[Keyword::NOT, Keyword::NULL]) @@ -208,10 +209,39 @@ fn parse_json2_type_hint(parser: &mut Parser<'_>) -> Result { } else if let Token::Word(word) = parser.peek_token().token && word.value.eq_ignore_ascii_case(SKIPPING) { - return InvalidSqlSnafu { - msg: "JSON2 type hint SKIPPING INDEX is not supported yet".to_string(), + parser.next_token(); + ensure!( + parser.parse_keyword(Keyword::INDEX), + InvalidSqlSnafu { + msg: "expect INDEX after SKIPPING keyword", + } + ); + ensure!( + skipping_index.is_none(), + InvalidSqlSnafu { + msg: "duplicated SKIPPING INDEX option for JSON2 type hint", + } + ); + let options = parser + .parse_options(Keyword::WITH) + .context(SyntaxSnafu)? + .into_iter() + .map(crate::util::parse_option_string) + .collect::>>()?; + for (key, _) in &options { + ensure!( + crate::parsers::utils::validate_column_skipping_index_create_option(key), + InvalidSqlSnafu { + msg: format!("invalid SKIPPING INDEX option: {key}"), + } + ); } - .fail(); + skipping_index = Some( + datatypes::schema::SkippingIndexOptions::try_from( + crate::statements::OptionMap::new(options).into_map(), + ) + .context(crate::error::SetSkippingIndexOptionSnafu)?, + ); } else if matches!(parser.peek_token().token, Token::Comma | Token::RParen) { break; } else { @@ -225,6 +255,7 @@ fn parse_json2_type_hint(parser: &mut Parser<'_>) -> Result { path, data_type, inverted_index, + skipping_index, }) } @@ -310,6 +341,59 @@ mod tests { create_table.columns.remove(0) } + #[test] + fn test_json2_skipping_index() { + let column = parse_json2_column( + "CREATE TABLE logs (attrs JSON2 (service.name STRING SKIPPING INDEX WITH (granularity='8', false_positive_rate='0.05') INVERTED INDEX, status BIGINT SKIPPING INDEX), ts TIMESTAMP TIME INDEX)", + ); + let options = column.extensions.json2_options.as_ref().unwrap(); + let settings = options.build_json_settings().unwrap(); + let hint = &settings.type_hints()[0]; + assert!(hint.inverted_index); + assert_eq!(hint.skipping_index.as_ref().unwrap().granularity, 8); + assert_eq!( + hint.skipping_index.as_ref().unwrap().false_positive_rate(), + 0.05 + ); + assert_eq!( + settings.type_hints()[1].skipping_index, + Some(datatypes::schema::SkippingIndexOptions::default()) + ); + let rebuilt = parse_json2_column(&format!( + "CREATE TABLE logs (attrs JSON2 {options}, ts TIMESTAMP TIME INDEX)" + )); + assert_eq!( + settings, + rebuilt + .extensions + .json2_options + .unwrap() + .build_json_settings() + .unwrap() + ); + for option in [ + "SKIPPING INDEX SKIPPING INDEX", + "SKIPPING", + "SKIPPING INDEX WITH (granularity='0')", + "SKIPPING INDEX WITH (false_positive_rate='0')", + "SKIPPING INDEX WITH (type='unknown')", + "SKIPPING INDEX WITH (unknown='1')", + ] { + let sql = format!( + "CREATE TABLE logs (attrs JSON2 (a STRING {option}), ts TIMESTAMP TIME INDEX)" + ); + assert!( + ParserContext::create_with_dialect( + &sql, + &GreptimeDbDialect {}, + ParseOptions::default() + ) + .is_err(), + "{sql}" + ); + } + } + #[test] fn test_parse_json2_type_hint_path() { assert_eq!( diff --git a/src/sql/src/statements.rs b/src/sql/src/statements.rs index 34a4c483cae..ffb00eba313 100644 --- a/src/sql/src/statements.rs +++ b/src/sql/src/statements.rs @@ -733,6 +733,7 @@ mod tests { path: vec!["kind".to_string()], data_type: ConcreteDataType::string_datatype(), inverted_index: false, + skipping_index: None, }], None, )?)?; diff --git a/src/sql/src/statements/create.rs b/src/sql/src/statements/create.rs index 54c2bdac19c..d01ac85a5de 100644 --- a/src/sql/src/statements/create.rs +++ b/src/sql/src/statements/create.rs @@ -201,6 +201,15 @@ impl Display for Json2Options { } f.write_char(' ')?; write_data_type(f, &hint.data_type)?; + if let Some(options) = &hint.skipping_index { + write!( + f, + " SKIPPING INDEX WITH (granularity = '{}', false_positive_rate = '{}', type = '{}')", + options.granularity, + options.false_positive_rate(), + options.index_type + )?; + } if hint.inverted_index { f.write_str(" INVERTED INDEX")?; } @@ -219,6 +228,7 @@ impl Json2Options { path: hint.path.clone(), data_type: sql_data_type_to_concrete_data_type(&hint.data_type)?, inverted_index: hint.inverted_index, + skipping_index: hint.skipping_index.clone(), }) }) .collect::>>()?; @@ -234,6 +244,8 @@ pub struct JsonTypeHint { pub path: Vec, pub data_type: DataType, pub inverted_index: bool, + /// Skipping index options for this hinted path. + pub skipping_index: Option, } impl Column { @@ -354,6 +366,7 @@ impl ColumnExtensions { path: hint.path, data_type, inverted_index: hint.inverted_index, + skipping_index: hint.skipping_index, }) }) .collect::>>()?; @@ -818,6 +831,7 @@ mod tests { path: vec!["x".repeat(8192)], data_type: DataType::Text, inverted_index: true, + skipping_index: None, }], }); assert!(!create.consume_sql_size_budget(&mut 4096)); @@ -838,6 +852,7 @@ mod tests { path: vec![path.into()], data_type: DataType::Text, inverted_index: false, + skipping_index: None, }], }; assert_eq!( @@ -1163,26 +1178,31 @@ ENGINE=mito path: vec!["i".to_string()], data_type: ConcreteDataType::int64_datatype(), inverted_index: false, + skipping_index: None, }, DatatypeJsonTypeHint { path: vec!["f".to_string()], data_type: ConcreteDataType::float64_datatype(), inverted_index: false, + skipping_index: None, }, DatatypeJsonTypeHint { path: vec!["u".to_string()], data_type: ConcreteDataType::uint64_datatype(), inverted_index: false, + skipping_index: None, }, DatatypeJsonTypeHint { path: vec!["s".to_string()], data_type: ConcreteDataType::string_datatype(), inverted_index: false, + skipping_index: None, }, DatatypeJsonTypeHint { path: vec!["b".to_string()], data_type: ConcreteDataType::boolean_datatype(), inverted_index: false, + skipping_index: None, }, ], None, @@ -1209,6 +1229,7 @@ ENGINE=mito path: vec!["u".to_string()], data_type: ConcreteDataType::date_datatype(), inverted_index: false, + skipping_index: None, }], None, ) diff --git a/src/store-api/src/region_request.rs b/src/store-api/src/region_request.rs index 787d6c5afe1..cf98ead1857 100644 --- a/src/store-api/src/region_request.rs +++ b/src/store-api/src/region_request.rs @@ -1571,6 +1571,16 @@ fn json_settings_from_proto(settings: v1::JsonSettings) -> Result path: hint.path, data_type, inverted_index: hint.inverted_index, + skipping_index: hint + .skipping_index + .map(api::v1::column_def::json_skipping_options_from_proto) + .transpose() + .map_err(|err| { + InvalidRawRegionRequestSnafu { + err: err.to_string(), + } + .build() + })?, }) }) .collect::>>()?; @@ -2521,6 +2531,7 @@ mod tests { path: path.iter().map(|part| (*part).to_string()).collect(), data_type: ColumnDataType::String as i32, inverted_index, + skipping_index: None, ..Default::default() }], max_auto_expanded_paths: Some(0), @@ -2606,6 +2617,11 @@ mod tests { data_type: ColumnDataType::String as i32, datatype_extension: None, inverted_index: true, + skipping_index: Some(v1::JsonSkippingIndexOptions { + granularity: 8, + index_type: 0, + false_positive_rate: 0.05, + }), }], max_auto_expanded_paths: Some(10), }), @@ -2624,6 +2640,14 @@ mod tests { assert_eq!(Some(10), settings.max_auto_expanded_paths()); assert_eq!(1, settings.type_hints().len()); assert!(settings.type_hints()[0].inverted_index); + assert_eq!( + settings.type_hints()[0] + .skipping_index + .as_ref() + .unwrap() + .granularity, + 8 + ); } #[test] diff --git a/tests/cases/standalone/common/types/json/json2_inverted_index.result b/tests/cases/standalone/common/types/json/json2_inverted_index.result index c03155e3961..8d9ee9c3ece 100644 --- a/tests/cases/standalone/common/types/json/json2_inverted_index.result +++ b/tests/cases/standalone/common/types/json/json2_inverted_index.result @@ -136,14 +136,7 @@ CREATE TABLE json2_inverted_invalid ( Error: 2000(InvalidSyntax), Invalid SQL, error: unsupported JSON2 type hint data type: TIMESTAMP; supported types: STRING, BIGINT, BIGINT UNSIGNED, DOUBLE, BOOLEAN; supported aliases: INT64, UINT64, FLOAT64 --- JSON hints support only inverted indexes. -CREATE TABLE json2_inverted_invalid ( - ts TIMESTAMP TIME INDEX, - attrs JSON2 (message STRING SKIPPING INDEX) -); - -Error: 2000(InvalidSyntax), Invalid SQL, error: JSON2 type hint SKIPPING INDEX is not supported yet - +-- JSON hints do not support fulltext indexes. CREATE TABLE json2_inverted_invalid ( ts TIMESTAMP TIME INDEX, attrs JSON2 (message STRING FULLTEXT INDEX) diff --git a/tests/cases/standalone/common/types/json/json2_inverted_index.sql b/tests/cases/standalone/common/types/json/json2_inverted_index.sql index 866a99a143e..aad6e414cdd 100644 --- a/tests/cases/standalone/common/types/json/json2_inverted_index.sql +++ b/tests/cases/standalone/common/types/json/json2_inverted_index.sql @@ -66,12 +66,7 @@ CREATE TABLE json2_inverted_invalid ( attrs JSON2 (event_time TIMESTAMP INVERTED INDEX) ); --- JSON hints support only inverted indexes. -CREATE TABLE json2_inverted_invalid ( - ts TIMESTAMP TIME INDEX, - attrs JSON2 (message STRING SKIPPING INDEX) -); - +-- JSON hints do not support fulltext indexes. CREATE TABLE json2_inverted_invalid ( ts TIMESTAMP TIME INDEX, attrs JSON2 (message STRING FULLTEXT INDEX) diff --git a/tests/cases/standalone/common/types/json/json2_skipping_index.result b/tests/cases/standalone/common/types/json/json2_skipping_index.result new file mode 100644 index 00000000000..1b861271457 --- /dev/null +++ b/tests/cases/standalone/common/types/json/json2_skipping_index.result @@ -0,0 +1,550 @@ +-- Query indexed JSON paths alongside ordinary columns in flat SSTs. +CREATE TABLE json2_skipping_query ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 ( + max_auto_expanded_paths = 0, + service.name STRING SKIPPING INDEX WITH (granularity='1'), + "service.name" STRING SKIPPING INDEX WITH (granularity='1'), + status BIGINT SKIPPING INDEX WITH (granularity='1'), + success BOOLEAN SKIPPING INDEX WITH (granularity='1'), + duration DOUBLE SKIPPING INDEX WITH (granularity='1'), + bytes BIGINT UNSIGNED SKIPPING INDEX WITH (granularity='1'), + unindexed STRING + ), + ordinary STRING SKIPPING INDEX WITH (granularity='1') +) WITH ('append_mode' = 'true', 'sst_format' = 'flat'); + +Affected Rows: 0 + +INSERT INTO json2_skipping_query VALUES + (1, '{"service":{"name":"api"},"service.name":"literal","status":200,"success":true,"duration":1.5,"bytes":18446744073709551615,"unindexed":"x"}', 'a'), + (2, '{"service":{"name":"worker"},"service.name":"api","status":500,"success":false,"duration":2.5,"bytes":42,"unindexed":"y"}', 'b'), + (3, '{}', 'c'), (4, NULL, 'd'); + +Affected Rows: 4 + +ADMIN FLUSH_TABLE('json2_skipping_query'); + ++-------------------------------------------+ +| ADMIN FLUSH_TABLE('json2_skipping_query') | ++-------------------------------------------+ +| 0 | ++-------------------------------------------+ + +-- Nested and literal dotted paths resolve to different targets. +SELECT ts FROM json2_skipping_query WHERE attrs.service.name = 'api' ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | ++-------------------------+ + +SELECT ts FROM json2_skipping_query WHERE json_get(attrs, '$."service.name"') = 'api' ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.002 | ++-------------------------+ + +-- Integer predicates support equality, reversed ranges, BETWEEN, and IN. +SELECT ts FROM json2_skipping_query WHERE attrs.status::BIGINT = 200 ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | ++-------------------------+ + +SELECT ts FROM json2_skipping_query WHERE 400 < attrs.status::BIGINT ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.002 | ++-------------------------+ + +SELECT ts FROM json2_skipping_query WHERE attrs.status::BIGINT BETWEEN 100 AND 300 ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | ++-------------------------+ + +SELECT ts FROM json2_skipping_query WHERE attrs.status::BIGINT IN (200, 500) ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | +| 1970-01-01T00:00:00.002 | ++-------------------------+ + +-- OR predicates work within one path and across different paths. +SELECT ts FROM json2_skipping_query WHERE attrs.service.name = 'api' OR attrs.service.name = 'worker' ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | +| 1970-01-01T00:00:00.002 | ++-------------------------+ + +SELECT ts FROM json2_skipping_query WHERE attrs.service.name = 'api' OR json_get(attrs, '$."service.name"') = 'api' ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | +| 1970-01-01T00:00:00.002 | ++-------------------------+ + +-- Regex, boolean, floating-point, and unsigned predicates return matching rows. +SELECT ts FROM json2_skipping_query WHERE attrs.service.name ~ '^a.*[ip]$' ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | ++-------------------------+ + +SELECT ts FROM json2_skipping_query WHERE attrs.success::BOOLEAN = true ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | ++-------------------------+ + +SELECT ts FROM json2_skipping_query WHERE attrs.duration::DOUBLE > 2.0 ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.002 | ++-------------------------+ + +SELECT ts FROM json2_skipping_query WHERE attrs.bytes::BIGINT UNSIGNED = 18446744073709551615::BIGINT UNSIGNED ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | ++-------------------------+ + +-- Ordinary and JSON index predicates can be combined. +SELECT ts FROM json2_skipping_query WHERE ordinary = 'b' AND attrs.status::BIGINT > 400 ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.002 | ++-------------------------+ + +-- Missing paths and null roots remain visible to IS NULL. +SELECT ts FROM json2_skipping_query WHERE attrs.service.name IS NULL ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.003 | +| 1970-01-01T00:00:00.004 | ++-------------------------+ + +-- Unindexed hints remain queryable. +SELECT ts FROM json2_skipping_query WHERE attrs.unindexed = 'x' ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | ++-------------------------+ + +-- Implicit comparison and explicit conversion preserve query results. +SELECT ts FROM json2_skipping_query WHERE attrs.status = '200' ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | ++-------------------------+ + +SELECT ts FROM json2_skipping_query WHERE attrs.duration::BIGINT = 1 ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | ++-------------------------+ + +-- Verify equality and IN actually prune, including warm-cache reads. +-- SQLNESS REPLACE "metrics_per_partition":.*("rows_bloom_filtered":\d+).*metrics= "metrics_per_partition":{$1} metrics= +-- SQLNESS REPLACE "metrics_per_partition":\s*\[.*metrics= "metrics_per_partition":REDACTED metrics= +-- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED +-- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" +-- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED +-- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED +-- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE ,\s"dyn_filters":\s\[.* REDACTED +-- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED +-- SQLNESS REPLACE "max_part_range":\s*\{[^}]*\} "max_part_range":REDACTED +EXPLAIN ANALYZE VERBOSE SELECT ts FROM json2_skipping_query WHERE attrs.service.name = 'absent'; + ++-+-+-+ +| stage | node | plan_| ++-+-+-+ +| 0_| 0_|_CooperativeExec metrics=REDACTED_| +|_|_|_MergeScanExec: REDACTED +|_|_|_| +| 1_| 0_|_FilterExec: json_get(attrs@1, $.service.name, NULL) = absent, projection=[ts@0] metrics=REDACTED_| +|_|_|_CooperativeExec metrics=REDACTED_| +|_|_|_UnorderedScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["ts", "attrs"], "filters": ["json_get(attrs, Utf8(\"$.service.name\"), Utf8View(NULL)) = Utf8View(\"absent\")"], "time_range": "TimestampRange{{[#,#)}}", "file_summary": {"files":1,"rows":4,"row_groups":1,"size":REDACTED,"index_size":REDACTED}, "max_part_range":REDACTED, "files": [{"file_id":"REDACTED","time_range_start":"REDACTED","time_range_end":"REDACTED","rows":4,"size":REDACTED,"index_size":REDACTED}], "flat_format":REDACTED, "metrics_per_partition":{"rows_bloom_filtered":4} metrics=REDACTED | +|_|_|_| +|_|_| Total rows: REDACTED_| ++-+-+-+ + +-- SQLNESS REPLACE "metrics_per_partition":.*("rows_bloom_filtered":\d+).*metrics= "metrics_per_partition":{$1} metrics= +-- SQLNESS REPLACE "metrics_per_partition":\s*\[.*metrics= "metrics_per_partition":REDACTED metrics= +-- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED +-- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" +-- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED +-- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED +-- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE ,\s"dyn_filters":\s\[.* REDACTED +-- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED +-- SQLNESS REPLACE "max_part_range":\s*\{[^}]*\} "max_part_range":REDACTED +EXPLAIN ANALYZE VERBOSE SELECT ts FROM json2_skipping_query WHERE attrs.status::BIGINT IN (900, 901); + ++-+-+-+ +| stage | node | plan_| ++-+-+-+ +| 0_| 0_|_CooperativeExec metrics=REDACTED_| +|_|_|_MergeScanExec: REDACTED +|_|_|_| +| 1_| 0_|_FilterExec: json_get(attrs@1, $.status, NULL) IN (SET) ([900, 901]), projection=[ts@0] metrics=REDACTED_| +|_|_|_CooperativeExec metrics=REDACTED_| +|_|_|_UnorderedScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["ts", "attrs"], "filters": ["json_get(attrs, Utf8(\"$.status\"), Int64(NULL)) IN ([Int64(900), Int64(901)])"], "time_range": "TimestampRange{{[#,#)}}", "file_summary": {"files":1,"rows":4,"row_groups":1,"size":REDACTED,"index_size":REDACTED}, "max_part_range":REDACTED, "files": [{"file_id":"REDACTED","time_range_start":"REDACTED","time_range_end":"REDACTED","rows":4,"size":REDACTED,"index_size":REDACTED}], "flat_format":REDACTED, "metrics_per_partition":{"rows_bloom_filtered":4} metrics=REDACTED | +|_|_|_| +|_|_| Total rows: REDACTED_| ++-+-+-+ + +-- SQLNESS REPLACE "metrics_per_partition":.*("rows_bloom_filtered":\d+).*metrics= "metrics_per_partition":{$1} metrics= +-- SQLNESS REPLACE "metrics_per_partition":\s*\[.*metrics= "metrics_per_partition":REDACTED metrics= +-- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED +-- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" +-- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED +-- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED +-- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE ,\s"dyn_filters":\s\[.* REDACTED +-- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED +-- SQLNESS REPLACE "max_part_range":\s*\{[^}]*\} "max_part_range":REDACTED +EXPLAIN ANALYZE VERBOSE SELECT ts FROM json2_skipping_query WHERE attrs.service.name = 'absent'; + ++-+-+-+ +| stage | node | plan_| ++-+-+-+ +| 0_| 0_|_CooperativeExec metrics=REDACTED_| +|_|_|_MergeScanExec: REDACTED +|_|_|_| +| 1_| 0_|_FilterExec: json_get(attrs@1, $.service.name, NULL) = absent, projection=[ts@0] metrics=REDACTED_| +|_|_|_CooperativeExec metrics=REDACTED_| +|_|_|_UnorderedScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["ts", "attrs"], "filters": ["json_get(attrs, Utf8(\"$.service.name\"), Utf8View(NULL)) = Utf8View(\"absent\")"], "time_range": "TimestampRange{{[#,#)}}", "file_summary": {"files":1,"rows":4,"row_groups":1,"size":REDACTED,"index_size":REDACTED}, "max_part_range":REDACTED, "files": [{"file_id":"REDACTED","time_range_start":"REDACTED","time_range_end":"REDACTED","rows":4,"size":REDACTED,"index_size":REDACTED}], "flat_format":REDACTED, "metrics_per_partition":{"rows_bloom_filtered":4} metrics=REDACTED | +|_|_|_| +|_|_| Total rows: REDACTED_| ++-+-+-+ + +INSERT INTO json2_skipping_query VALUES (5, '{"service":{"name":"api"},"status":201}', 'e'); + +Affected Rows: 1 + +ADMIN FLUSH_TABLE('json2_skipping_query'); + ++-------------------------------------------+ +| ADMIN FLUSH_TABLE('json2_skipping_query') | ++-------------------------------------------+ +| 0 | ++-------------------------------------------+ + +ADMIN COMPACT_TABLE('json2_skipping_query', 'swcs', '86400'); + ++--------------------------------------------------------------+ +| ADMIN COMPACT_TABLE('json2_skipping_query', 'swcs', '86400') | ++--------------------------------------------------------------+ +| 0 | ++--------------------------------------------------------------+ + +-- SQLNESS REPLACE "metrics_per_partition":.*("rows_bloom_filtered":\d+).*metrics= "metrics_per_partition":{$1} metrics= +-- SQLNESS REPLACE "metrics_per_partition":\s*\[.*metrics= "metrics_per_partition":REDACTED metrics= +-- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED +-- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" +-- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED +-- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED +-- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE ,\s"dyn_filters":\s\[.* REDACTED +-- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED +-- SQLNESS REPLACE "max_part_range":\s*\{[^}]*\} "max_part_range":REDACTED +EXPLAIN ANALYZE VERBOSE SELECT ts FROM json2_skipping_query WHERE attrs.service.name = 'absent'; + ++-+-+-+ +| stage | node | plan_| ++-+-+-+ +| 0_| 0_|_CooperativeExec metrics=REDACTED_| +|_|_|_MergeScanExec: REDACTED +|_|_|_| +| 1_| 0_|_FilterExec: json_get(attrs@1, $.service.name, NULL) = absent, projection=[ts@0] metrics=REDACTED_| +|_|_|_CooperativeExec metrics=REDACTED_| +|_|_|_UnorderedScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["ts", "attrs"], "filters": ["json_get(attrs, Utf8(\"$.service.name\"), Utf8View(NULL)) = Utf8View(\"absent\")"], "time_range": "TimestampRange{{[#,#)}}", "file_summary": {"files":1,"rows":5,"row_groups":1,"size":REDACTED,"index_size":REDACTED}, "max_part_range":REDACTED, "files": [{"file_id":"REDACTED","time_range_start":"REDACTED","time_range_end":"REDACTED","rows":5,"size":REDACTED,"index_size":REDACTED}], "flat_format":REDACTED, "metrics_per_partition":{"rows_bloom_filtered":5} metrics=REDACTED | +|_|_|_| +|_|_| Total rows: REDACTED_| ++-+-+-+ + +SELECT ts FROM json2_skipping_query WHERE attrs.service.name = 'api' ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | +| 1970-01-01T00:00:00.005 | ++-------------------------+ + +SELECT ts FROM json2_skipping_query WHERE attrs.status::BIGINT IN (200, 201, 500) ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | +| 1970-01-01T00:00:00.002 | +| 1970-01-01T00:00:00.005 | ++-------------------------+ + +-- Disable hints' indexes: results remain identical across existing SSTs. +ALTER TABLE json2_skipping_query MODIFY COLUMN attrs JSON2 (max_auto_expanded_paths = 0, service.name STRING, "service.name" STRING, status BIGINT, success BOOLEAN, duration DOUBLE, bytes BIGINT UNSIGNED, unindexed STRING); + +Affected Rows: 0 + +SELECT ts FROM json2_skipping_query WHERE attrs.service.name = 'api' ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | +| 1970-01-01T00:00:00.005 | ++-------------------------+ + +SELECT ts FROM json2_skipping_query WHERE attrs.status::BIGINT IN (200, 201, 500) ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | +| 1970-01-01T00:00:00.002 | +| 1970-01-01T00:00:00.005 | ++-------------------------+ + +DROP TABLE json2_skipping_query; + +Affected Rows: 0 + +-- Old SST without a path index, new SST with both index kinds, then a type change. +CREATE TABLE json2_skipping_evolution (ts TIMESTAMP TIME INDEX, j JSON2 (n BIGINT)) WITH ('append_mode'='true'); + +Affected Rows: 0 + +INSERT INTO json2_skipping_evolution VALUES (1, '{"n":42}'); + +Affected Rows: 1 + +ADMIN FLUSH_TABLE('json2_skipping_evolution'); + ++-----------------------------------------------+ +| ADMIN FLUSH_TABLE('json2_skipping_evolution') | ++-----------------------------------------------+ +| 0 | ++-----------------------------------------------+ + +ALTER TABLE json2_skipping_evolution MODIFY COLUMN j JSON2 (n BIGINT SKIPPING INDEX WITH (granularity='1') INVERTED INDEX); + +Affected Rows: 0 + +SHOW CREATE TABLE json2_skipping_evolution; + ++--------------------------+---------------------------------------------------------------------------------------------------------------------+ +| Table | Create Table | ++--------------------------+---------------------------------------------------------------------------------------------------------------------+ +| json2_skipping_evolution | CREATE TABLE IF NOT EXISTS "json2_skipping_evolution" ( | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | "j" JSON2( | +| | max_auto_expanded_paths = 100, | +| | "n" BIGINT SKIPPING INDEX WITH (granularity = '1', false_positive_rate = '0.01', type = 'BLOOM') INVERTED INDEX | +| | ) NULL, | +| | TIME INDEX ("ts") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | append_mode = 'true' | +| | ) | ++--------------------------+---------------------------------------------------------------------------------------------------------------------+ + +INSERT INTO json2_skipping_evolution VALUES (2, '{"n":42}'); + +Affected Rows: 1 + +ADMIN FLUSH_TABLE('json2_skipping_evolution'); + ++-----------------------------------------------+ +| ADMIN FLUSH_TABLE('json2_skipping_evolution') | ++-----------------------------------------------+ +| 0 | ++-----------------------------------------------+ + +SELECT ts FROM json2_skipping_evolution WHERE json_get(j, '$.n') = 42 ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | +| 1970-01-01T00:00:00.002 | ++-------------------------+ + +ALTER TABLE json2_skipping_evolution MODIFY COLUMN j JSON2 (n DOUBLE SKIPPING INDEX WITH (granularity='1')); + +Affected Rows: 0 + +INSERT INTO json2_skipping_evolution VALUES (3, '{"n":42.0}'); + +Affected Rows: 1 + +ADMIN FLUSH_TABLE('json2_skipping_evolution'); + ++-----------------------------------------------+ +| ADMIN FLUSH_TABLE('json2_skipping_evolution') | ++-----------------------------------------------+ +| 0 | ++-----------------------------------------------+ + +SELECT ts FROM json2_skipping_evolution WHERE json_get(j, '$.n') = 42.0 ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | +| 1970-01-01T00:00:00.002 | +| 1970-01-01T00:00:00.003 | ++-------------------------+ + +-- SQLNESS ARG restart=true +SHOW CREATE TABLE json2_skipping_evolution; + ++--------------------------+------------------------------------------------------------------------------------------------------+ +| Table | Create Table | ++--------------------------+------------------------------------------------------------------------------------------------------+ +| json2_skipping_evolution | CREATE TABLE IF NOT EXISTS "json2_skipping_evolution" ( | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | "j" JSON2( | +| | max_auto_expanded_paths = 100, | +| | "n" DOUBLE SKIPPING INDEX WITH (granularity = '1', false_positive_rate = '0.01', type = 'BLOOM') | +| | ) NULL, | +| | TIME INDEX ("ts") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | append_mode = 'true' | +| | ) | ++--------------------------+------------------------------------------------------------------------------------------------------+ + +SELECT ts FROM json2_skipping_evolution WHERE json_get(j, '$.n') = 42.0 ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | +| 1970-01-01T00:00:00.002 | +| 1970-01-01T00:00:00.003 | ++-------------------------+ + +ADMIN COMPACT_TABLE('json2_skipping_evolution', 'swcs', '86400'); + ++------------------------------------------------------------------+ +| ADMIN COMPACT_TABLE('json2_skipping_evolution', 'swcs', '86400') | ++------------------------------------------------------------------+ +| 0 | ++------------------------------------------------------------------+ + +-- SQLNESS REPLACE "metrics_per_partition":.*("rows_bloom_filtered":\d+).*metrics= "metrics_per_partition":{$1} metrics= +-- SQLNESS REPLACE "metrics_per_partition":\s*\[.*metrics= "metrics_per_partition":REDACTED metrics= +-- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED +-- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" +-- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED +-- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED +-- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE ,\s"dyn_filters":\s\[.* REDACTED +-- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED +-- SQLNESS REPLACE "max_part_range":\s*\{[^}]*\} "max_part_range":REDACTED +EXPLAIN ANALYZE VERBOSE SELECT ts FROM json2_skipping_evolution WHERE json_get(j, '$.n') = 99.0; + ++-+-+-+ +| stage | node | plan_| ++-+-+-+ +| 0_| 0_|_CooperativeExec metrics=REDACTED_| +|_|_|_MergeScanExec: REDACTED +|_|_|_| +| 1_| 0_|_FilterExec: json_get(j@1, $.n, NULL) = 99, projection=[ts@0] metrics=REDACTED_| +|_|_|_CooperativeExec metrics=REDACTED_| +|_|_|_UnorderedScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["ts", "j"], "filters": ["json_get(j, Utf8(\"$.n\"), Float64(NULL)) = Float64(99)"], "time_range": "TimestampRange{{[#,#)}}", "file_summary": {"files":1,"rows":3,"row_groups":1,"size":REDACTED,"index_size":REDACTED}, "max_part_range":REDACTED, "files": [{"file_id":"REDACTED","time_range_start":"REDACTED","time_range_end":"REDACTED","rows":3,"size":REDACTED,"index_size":REDACTED}], "flat_format":REDACTED, "metrics_per_partition":{"rows_bloom_filtered":3} metrics=REDACTED | +|_|_|_| +|_|_| Total rows: REDACTED_| ++-+-+-+ + +SELECT ts FROM json2_skipping_evolution WHERE json_get(j, '$.n') = 42.0 ORDER BY ts; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.001 | +| 1970-01-01T00:00:00.002 | +| 1970-01-01T00:00:00.003 | ++-------------------------+ + +DROP TABLE json2_skipping_evolution; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/types/json/json2_skipping_index.sql b/tests/cases/standalone/common/types/json/json2_skipping_index.sql new file mode 100644 index 00000000000..4f3a12838a5 --- /dev/null +++ b/tests/cases/standalone/common/types/json/json2_skipping_index.sql @@ -0,0 +1,190 @@ +-- Query indexed JSON paths alongside ordinary columns in flat SSTs. +CREATE TABLE json2_skipping_query ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 ( + max_auto_expanded_paths = 0, + service.name STRING SKIPPING INDEX WITH (granularity='1'), + "service.name" STRING SKIPPING INDEX WITH (granularity='1'), + status BIGINT SKIPPING INDEX WITH (granularity='1'), + success BOOLEAN SKIPPING INDEX WITH (granularity='1'), + duration DOUBLE SKIPPING INDEX WITH (granularity='1'), + bytes BIGINT UNSIGNED SKIPPING INDEX WITH (granularity='1'), + unindexed STRING + ), + ordinary STRING SKIPPING INDEX WITH (granularity='1') +) WITH ('append_mode' = 'true', 'sst_format' = 'flat'); + +INSERT INTO json2_skipping_query VALUES + (1, '{"service":{"name":"api"},"service.name":"literal","status":200,"success":true,"duration":1.5,"bytes":18446744073709551615,"unindexed":"x"}', 'a'), + (2, '{"service":{"name":"worker"},"service.name":"api","status":500,"success":false,"duration":2.5,"bytes":42,"unindexed":"y"}', 'b'), + (3, '{}', 'c'), (4, NULL, 'd'); + +ADMIN FLUSH_TABLE('json2_skipping_query'); + +-- Nested and literal dotted paths resolve to different targets. +SELECT ts FROM json2_skipping_query WHERE attrs.service.name = 'api' ORDER BY ts; + +SELECT ts FROM json2_skipping_query WHERE json_get(attrs, '$."service.name"') = 'api' ORDER BY ts; + +-- Integer predicates support equality, reversed ranges, BETWEEN, and IN. +SELECT ts FROM json2_skipping_query WHERE attrs.status::BIGINT = 200 ORDER BY ts; + +SELECT ts FROM json2_skipping_query WHERE 400 < attrs.status::BIGINT ORDER BY ts; + +SELECT ts FROM json2_skipping_query WHERE attrs.status::BIGINT BETWEEN 100 AND 300 ORDER BY ts; + +SELECT ts FROM json2_skipping_query WHERE attrs.status::BIGINT IN (200, 500) ORDER BY ts; + +-- OR predicates work within one path and across different paths. +SELECT ts FROM json2_skipping_query WHERE attrs.service.name = 'api' OR attrs.service.name = 'worker' ORDER BY ts; + +SELECT ts FROM json2_skipping_query WHERE attrs.service.name = 'api' OR json_get(attrs, '$."service.name"') = 'api' ORDER BY ts; + +-- Regex, boolean, floating-point, and unsigned predicates return matching rows. +SELECT ts FROM json2_skipping_query WHERE attrs.service.name ~ '^a.*[ip]$' ORDER BY ts; + +SELECT ts FROM json2_skipping_query WHERE attrs.success::BOOLEAN = true ORDER BY ts; + +SELECT ts FROM json2_skipping_query WHERE attrs.duration::DOUBLE > 2.0 ORDER BY ts; + +SELECT ts FROM json2_skipping_query WHERE attrs.bytes::BIGINT UNSIGNED = 18446744073709551615::BIGINT UNSIGNED ORDER BY ts; + +-- Ordinary and JSON index predicates can be combined. +SELECT ts FROM json2_skipping_query WHERE ordinary = 'b' AND attrs.status::BIGINT > 400 ORDER BY ts; + +-- Missing paths and null roots remain visible to IS NULL. +SELECT ts FROM json2_skipping_query WHERE attrs.service.name IS NULL ORDER BY ts; + +-- Unindexed hints remain queryable. +SELECT ts FROM json2_skipping_query WHERE attrs.unindexed = 'x' ORDER BY ts; + +-- Implicit comparison and explicit conversion preserve query results. +SELECT ts FROM json2_skipping_query WHERE attrs.status = '200' ORDER BY ts; + +SELECT ts FROM json2_skipping_query WHERE attrs.duration::BIGINT = 1 ORDER BY ts; + +-- Verify equality and IN actually prune, including warm-cache reads. +-- SQLNESS REPLACE "metrics_per_partition":.*("rows_bloom_filtered":\d+).*metrics= "metrics_per_partition":{$1} metrics= +-- SQLNESS REPLACE "metrics_per_partition":\s*\[.*metrics= "metrics_per_partition":REDACTED metrics= +-- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED +-- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" +-- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED +-- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED +-- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE ,\s"dyn_filters":\s\[.* REDACTED +-- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED +-- SQLNESS REPLACE "max_part_range":\s*\{[^}]*\} "max_part_range":REDACTED +EXPLAIN ANALYZE VERBOSE SELECT ts FROM json2_skipping_query WHERE attrs.service.name = 'absent'; + +-- SQLNESS REPLACE "metrics_per_partition":.*("rows_bloom_filtered":\d+).*metrics= "metrics_per_partition":{$1} metrics= +-- SQLNESS REPLACE "metrics_per_partition":\s*\[.*metrics= "metrics_per_partition":REDACTED metrics= +-- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED +-- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" +-- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED +-- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED +-- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE ,\s"dyn_filters":\s\[.* REDACTED +-- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED +-- SQLNESS REPLACE "max_part_range":\s*\{[^}]*\} "max_part_range":REDACTED +EXPLAIN ANALYZE VERBOSE SELECT ts FROM json2_skipping_query WHERE attrs.status::BIGINT IN (900, 901); + +-- SQLNESS REPLACE "metrics_per_partition":.*("rows_bloom_filtered":\d+).*metrics= "metrics_per_partition":{$1} metrics= +-- SQLNESS REPLACE "metrics_per_partition":\s*\[.*metrics= "metrics_per_partition":REDACTED metrics= +-- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED +-- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" +-- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED +-- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED +-- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE ,\s"dyn_filters":\s\[.* REDACTED +-- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED +-- SQLNESS REPLACE "max_part_range":\s*\{[^}]*\} "max_part_range":REDACTED +EXPLAIN ANALYZE VERBOSE SELECT ts FROM json2_skipping_query WHERE attrs.service.name = 'absent'; + +INSERT INTO json2_skipping_query VALUES (5, '{"service":{"name":"api"},"status":201}', 'e'); +ADMIN FLUSH_TABLE('json2_skipping_query'); +ADMIN COMPACT_TABLE('json2_skipping_query', 'swcs', '86400'); +-- SQLNESS REPLACE "metrics_per_partition":.*("rows_bloom_filtered":\d+).*metrics= "metrics_per_partition":{$1} metrics= +-- SQLNESS REPLACE "metrics_per_partition":\s*\[.*metrics= "metrics_per_partition":REDACTED metrics= +-- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED +-- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" +-- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED +-- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED +-- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE ,\s"dyn_filters":\s\[.* REDACTED +-- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED +-- SQLNESS REPLACE "max_part_range":\s*\{[^}]*\} "max_part_range":REDACTED +EXPLAIN ANALYZE VERBOSE SELECT ts FROM json2_skipping_query WHERE attrs.service.name = 'absent'; +SELECT ts FROM json2_skipping_query WHERE attrs.service.name = 'api' ORDER BY ts; +SELECT ts FROM json2_skipping_query WHERE attrs.status::BIGINT IN (200, 201, 500) ORDER BY ts; + +-- Disable hints' indexes: results remain identical across existing SSTs. +ALTER TABLE json2_skipping_query MODIFY COLUMN attrs JSON2 (max_auto_expanded_paths = 0, service.name STRING, "service.name" STRING, status BIGINT, success BOOLEAN, duration DOUBLE, bytes BIGINT UNSIGNED, unindexed STRING); +SELECT ts FROM json2_skipping_query WHERE attrs.service.name = 'api' ORDER BY ts; +SELECT ts FROM json2_skipping_query WHERE attrs.status::BIGINT IN (200, 201, 500) ORDER BY ts; +DROP TABLE json2_skipping_query; + +-- Old SST without a path index, new SST with both index kinds, then a type change. +CREATE TABLE json2_skipping_evolution (ts TIMESTAMP TIME INDEX, j JSON2 (n BIGINT)) WITH ('append_mode'='true'); +INSERT INTO json2_skipping_evolution VALUES (1, '{"n":42}'); +ADMIN FLUSH_TABLE('json2_skipping_evolution'); +ALTER TABLE json2_skipping_evolution MODIFY COLUMN j JSON2 (n BIGINT SKIPPING INDEX WITH (granularity='1') INVERTED INDEX); +SHOW CREATE TABLE json2_skipping_evolution; +INSERT INTO json2_skipping_evolution VALUES (2, '{"n":42}'); +ADMIN FLUSH_TABLE('json2_skipping_evolution'); +SELECT ts FROM json2_skipping_evolution WHERE json_get(j, '$.n') = 42 ORDER BY ts; +ALTER TABLE json2_skipping_evolution MODIFY COLUMN j JSON2 (n DOUBLE SKIPPING INDEX WITH (granularity='1')); +INSERT INTO json2_skipping_evolution VALUES (3, '{"n":42.0}'); +ADMIN FLUSH_TABLE('json2_skipping_evolution'); +SELECT ts FROM json2_skipping_evolution WHERE json_get(j, '$.n') = 42.0 ORDER BY ts; +-- SQLNESS ARG restart=true +SHOW CREATE TABLE json2_skipping_evolution; +SELECT ts FROM json2_skipping_evolution WHERE json_get(j, '$.n') = 42.0 ORDER BY ts; +ADMIN COMPACT_TABLE('json2_skipping_evolution', 'swcs', '86400'); +-- SQLNESS REPLACE "metrics_per_partition":.*("rows_bloom_filtered":\d+).*metrics= "metrics_per_partition":{$1} metrics= +-- SQLNESS REPLACE "metrics_per_partition":\s*\[.*metrics= "metrics_per_partition":REDACTED metrics= +-- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED +-- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" +-- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED +-- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED +-- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE ,\s"dyn_filters":\s\[.* REDACTED +-- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED +-- SQLNESS REPLACE "max_part_range":\s*\{[^}]*\} "max_part_range":REDACTED +EXPLAIN ANALYZE VERBOSE SELECT ts FROM json2_skipping_evolution WHERE json_get(j, '$.n') = 99.0; +SELECT ts FROM json2_skipping_evolution WHERE json_get(j, '$.n') = 42.0 ORDER BY ts; +DROP TABLE json2_skipping_evolution; From 8ecf49dca017adf873d31f7d4e6f01ebd4cef1cb Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Thu, 29 Oct 2026 19:16:19 +0000 Subject: [PATCH 141/239] ci: authenticate ECR Public pulls in benchmark workflows (#10056) * ci: authenticate agent benchmark pulls to ECR Public Signed-off-by: WenyXu * ci: authenticate remaining ECR Public benchmark pulls Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- .github/workflows/agent-observability.yml | 24 +++++++++++++++++++++++ .github/workflows/oswal-s3-bench.yml | 24 +++++++++++++++++++++++ .github/workflows/tracesbench.yml | 24 +++++++++++++++++++++++ .github/workflows/vmbench-long-range.yml | 24 +++++++++++++++++++++++ 4 files changed, 96 insertions(+) diff --git a/.github/workflows/agent-observability.yml b/.github/workflows/agent-observability.yml index d265801d4a0..90840536594 100644 --- a/.github/workflows/agent-observability.yml +++ b/.github/workflows/agent-observability.yml @@ -144,6 +144,18 @@ jobs: python3 greptimedb/tests/perf/test_aliyun_ecs_runner_scripts.py python3 greptimedb/tests/perf/test_ci_runner_scripts.py python3 greptimedb/tests/perf/test_agent_observability_summary.py + - name: Configure AWS credentials for ECR Public + if: startsWith(env.RUNTIME_IMAGE, 'public.ecr.aws/') + uses: aws-actions/configure-aws-credentials@v4 + with: + aws-access-key-id: ${{ secrets.AWS_ECR_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_ECR_SECRET_ACCESS_KEY }} + aws-region: us-east-1 + - name: Login to Amazon ECR Public + if: startsWith(env.RUNTIME_IMAGE, 'public.ecr.aws/') + uses: aws-actions/amazon-ecr-login@v2 + with: + registry-type: public - name: Check runtime supports shared version-comparison data if: inputs.compare_greptimedb run: | @@ -285,6 +297,18 @@ jobs: docker info # The workspace (and Docker state) must not use a small /tmp filesystem. df -h "$GITHUB_WORKSPACE" /var/lib/docker + - name: Configure AWS credentials for ECR Public + if: startsWith(env.RUNTIME_IMAGE, 'public.ecr.aws/') + uses: aws-actions/configure-aws-credentials@v4 + with: + aws-access-key-id: ${{ secrets.AWS_ECR_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_ECR_SECRET_ACCESS_KEY }} + aws-region: us-east-1 + - name: Login to Amazon ECR Public + if: startsWith(env.RUNTIME_IMAGE, 'public.ecr.aws/') + uses: aws-actions/amazon-ecr-login@v2 + with: + registry-type: public - name: Resolve image tags and verify runtime revision shell: bash run: | diff --git a/.github/workflows/oswal-s3-bench.yml b/.github/workflows/oswal-s3-bench.yml index e61754cb1f1..92243c89e58 100644 --- a/.github/workflows/oswal-s3-bench.yml +++ b/.github/workflows/oswal-s3-bench.yml @@ -117,6 +117,18 @@ jobs: --timeout-minutes "$BENCHMARK_TIMEOUT_MINUTES" --ttl-hours "$JANITOR_TTL_HOURS" python3 greptimedb/.github/scripts/oswal-s3-bench.py image --tag "$GREPTIMEDB_TAG" --image "$GREPTIMEDB_IMAGE_INPUT" \ --images "$OSWAL_IMAGES" + - name: Configure AWS credentials for ECR Public + if: startsWith(env.RUNTIME_IMAGE, 'public.ecr.aws/') + uses: aws-actions/configure-aws-credentials@v4 + with: + aws-access-key-id: ${{ secrets.AWS_ECR_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_ECR_SECRET_ACCESS_KEY }} + aws-region: us-east-1 + - name: Login to Amazon ECR Public + if: startsWith(env.RUNTIME_IMAGE, 'public.ecr.aws/') + uses: aws-actions/amazon-ecr-login@v2 + with: + registry-type: public - name: Pull every GreptimeDB image and record its version id: images env: @@ -200,6 +212,18 @@ jobs: run: | bash greptimedb/.github/scripts/pull-benchmark-image.sh greptime/minio:2024 docker compose -f greptimedb/tests-integration/fixtures/docker-compose.yml up -d --wait minio + - name: Configure AWS credentials for ECR Public + if: startsWith(env.RUNTIME_IMAGE, 'public.ecr.aws/') + uses: aws-actions/configure-aws-credentials@v4 + with: + aws-access-key-id: ${{ secrets.AWS_ECR_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_ECR_SECRET_ACCESS_KEY }} + aws-region: us-east-1 + - name: Login to Amazon ECR Public + if: startsWith(env.RUNTIME_IMAGE, 'public.ecr.aws/') + uses: aws-actions/amazon-ecr-login@v2 + with: + registry-type: public - name: Pull and check images, write the plan run: | bash greptimedb/.github/scripts/oswal-s3-bench-run.sh pull diff --git a/.github/workflows/tracesbench.yml b/.github/workflows/tracesbench.yml index 049faaef1a7..fd64ebfcf37 100644 --- a/.github/workflows/tracesbench.yml +++ b/.github/workflows/tracesbench.yml @@ -180,6 +180,18 @@ jobs: esac fi python3 greptimedb/.github/scripts/ci-runner-config.py --provider "$PROVIDER" --instance-type "$instance_type" + - name: Configure AWS credentials for ECR Public + if: startsWith(env.RUNTIME_IMAGE, 'public.ecr.aws/') + uses: aws-actions/configure-aws-credentials@v4 + with: + aws-access-key-id: ${{ secrets.AWS_ECR_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_ECR_SECRET_ACCESS_KEY }} + aws-region: us-east-1 + - name: Login to Amazon ECR Public + if: startsWith(env.RUNTIME_IMAGE, 'public.ecr.aws/') + uses: aws-actions/amazon-ecr-login@v2 + with: + registry-type: public - name: Validate budgets and shared runtime id: runtime run: | @@ -314,6 +326,18 @@ jobs: ref: ${{ env.RESOLVED_O11YBENCH_REF }} path: o11ybench persist-credentials: false + - name: Configure AWS credentials for ECR Public + if: startsWith(env.RUNTIME_IMAGE, 'public.ecr.aws/') + uses: aws-actions/configure-aws-credentials@v4 + with: + aws-access-key-id: ${{ secrets.AWS_ECR_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_ECR_SECRET_ACCESS_KEY }} + aws-region: us-east-1 + - name: Login to Amazon ECR Public + if: startsWith(env.RUNTIME_IMAGE, 'public.ecr.aws/') + uses: aws-actions/amazon-ecr-login@v2 + with: + registry-type: public - name: Check runner and runtime identity id: runtime run: | diff --git a/.github/workflows/vmbench-long-range.yml b/.github/workflows/vmbench-long-range.yml index aaf2dd7f437..d4e3ad124c1 100644 --- a/.github/workflows/vmbench-long-range.yml +++ b/.github/workflows/vmbench-long-range.yml @@ -132,6 +132,18 @@ jobs: run: | python3 greptimedb/tests/perf/test_aliyun_ecs_runner_scripts.py python3 greptimedb/tests/perf/test_ci_runner_scripts.py + - name: Configure AWS credentials for ECR Public + if: startsWith(env.RUNTIME_IMAGE, 'public.ecr.aws/') + uses: aws-actions/configure-aws-credentials@v4 + with: + aws-access-key-id: ${{ secrets.AWS_ECR_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_ECR_SECRET_ACCESS_KEY }} + aws-region: us-east-1 + - name: Login to Amazon ECR Public + if: startsWith(env.RUNTIME_IMAGE, 'public.ecr.aws/') + uses: aws-actions/amazon-ecr-login@v2 + with: + registry-type: public - name: Validate budgets and build warm plan id: plan env: @@ -279,6 +291,18 @@ jobs: with: name: ${{ needs.validate.outputs.plan_artifact }} path: ${{ env.METADATA_ROOT }} + - name: Configure AWS credentials for ECR Public + if: startsWith(env.RUNTIME_IMAGE, 'public.ecr.aws/') + uses: aws-actions/configure-aws-credentials@v4 + with: + aws-access-key-id: ${{ secrets.AWS_ECR_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_ECR_SECRET_ACCESS_KEY }} + aws-region: us-east-1 + - name: Login to Amazon ECR Public + if: startsWith(env.RUNTIME_IMAGE, 'public.ecr.aws/') + uses: aws-actions/amazon-ecr-login@v2 + with: + registry-type: public - name: Check runner and resolve database images run: | set -euo pipefail From 00cda3853aab98d9c95be133ba37ef8962ad9590 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Fri, 30 Oct 2026 00:11:18 +0000 Subject: [PATCH 142/239] ci: reduce duplicate and PR-scoped Rust cache writes (#10047) * ci: reduce duplicate and PR-scoped Rust cache writes Signed-off-by: WenyXu * ci: warm shared caches through scheduled integration jobs Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- .../actions/github-script-tests/action.yml | 2 +- .github/ci/cache-policy.md | 45 +++++ .../scripts/tests/ci-cache-policy.test.cjs | 155 ++++++++++++++++++ .github/workflows/checks.yml | 4 + .github/workflows/integration.yml | 24 ++- .github/workflows/nightly-ci.yml | 7 + .github/workflows/rust.yml | 12 +- 7 files changed, 237 insertions(+), 12 deletions(-) create mode 100644 .github/ci/cache-policy.md create mode 100644 .github/scripts/tests/ci-cache-policy.test.cjs diff --git a/.github/actions/github-script-tests/action.yml b/.github/actions/github-script-tests/action.yml index 4e3d797a95c..5721aead096 100644 --- a/.github/actions/github-script-tests/action.yml +++ b/.github/actions/github-script-tests/action.yml @@ -8,7 +8,7 @@ runs: node-version: 24 - name: Run Node GitHub script tests shell: bash - run: node --test .github/scripts/query-regression-comment.test.cjs .github/scripts/heavy-ci-policy.test.cjs .github/scripts/tests/nightly-merge.test.cjs + run: node --test .github/scripts/tests/ci-cache-policy.test.cjs .github/scripts/query-regression-comment.test.cjs .github/scripts/heavy-ci-policy.test.cjs .github/scripts/tests/nightly-merge.test.cjs - name: Setup uv uses: astral-sh/setup-uv@v3 with: diff --git a/.github/ci/cache-policy.md b/.github/ci/cache-policy.md new file mode 100644 index 00000000000..d1eb053ec1c --- /dev/null +++ b/.github/ci/cache-policy.md @@ -0,0 +1,45 @@ +# CI cache policy + +Checks, Rust CI, Integration CI, and Nightly CI use one explicit Rust cache per +build job. Disable the implicit cache in `setup-rust-toolchain` when adding an +explicit cache: otherwise its independent save policy can write PR-scoped +caches even when the explicit cache only saves on main. + +Integration CI's existing weekday schedule runs the fuzz builds and suites, +including the profile-CI binary and Kafka helper, to populate main caches. +To warm these caches manually after deployment, run Integration CI on main with +`fuzz_profile: all`; no internal `ci_command` marker is required. + +Rust CI's existing weekday schedule warms the Windows and RISC-V compilation +caches. Nightly CI's existing schedule warms its Windows runtime test caches. +A manual Rust CI run on main can warm the compilation caches sooner. These jobs +keep their separate keys and existing runner types; no extra warmup job is added. +Coverage caches are not warmed by these schedules. +PR, merge-group, and non-main manual runs restore caches without saving them. +Fuzz binaries still flow to downstream jobs through workflow artifacts on every +run; only the persistent binary cache is main-written. Re-running a PR cannot +reuse a newly saved PR-specific binary cache under this policy. + +Keep helper version variables outside the `CARGO*` and `RUST*` prefixes unless +they affect the compilation being cached. rust-cache includes those prefixes in +its environment hash. Fuzz tool versions remain explicitly included in the fuzz +cache keys. Compiler versions, Cargo manifests, flags, OS, and architecture +remain part of the existing cache isolation. + +This policy reduces duplicate and PR-only writes; it does not guarantee that all +build caches fit in the repository quota. Renamed environment variables cause +an initial cold restore until main produces the new keys. Existing caches are +not deleted by this change. Coverage warming and additional runner capacity are +not introduced. + +Validate the policy with: + +```sh +node --test .github/scripts/tests/ci-cache-policy.test.cjs +``` + +After deployment, compare the same commit and runner image across repeated runs. +Record exact versus fallback restores, archive size, restore/save duration, and +compilation duration. Verify main writes the keys consumed by PRs and that no +implicit post-action uploads remain. Local policy tests do not establish a +runtime speedup or long-term cache hit rate. diff --git a/.github/scripts/tests/ci-cache-policy.test.cjs b/.github/scripts/tests/ci-cache-policy.test.cjs new file mode 100644 index 00000000000..ba6817424c6 --- /dev/null +++ b/.github/scripts/tests/ci-cache-policy.test.cjs @@ -0,0 +1,155 @@ +const assert = require('node:assert/strict'); +const { readFileSync } = require('node:fs'); +const { resolve } = require('node:path'); +const { test } = require('node:test'); + +const directory = resolve(__dirname, '../../workflows'); +const workflows = ['checks', 'rust', 'integration', 'nightly-ci'].map(name => ({ + name, + text: readFileSync(resolve(directory, name + '.yml'), 'utf8'), +})); +function jobs(text) { + return [...text.matchAll(/^ ([\w-]+):\n([\s\S]*?)(?=^ [\w-]+:\n|(?![\s\S]))/gm)]; +} +function steps(body) { + return body.split(/^ - /m).slice(1); +} + +test('build caches have one controller and lightweight checks do not cache targets', () => { + let checked = 0; + for (const { name, text } of workflows) { + for (const [, job, body] of jobs(text)) { + if (!body.includes('Swatinem/rust-cache@') && !['fmt', 'check-udeps', 'toml'].includes(job)) continue; + for (const step of steps(body).filter(s => s.includes('uses: actions-rust-lang/setup-rust-toolchain@'))) { + assert.match(step, /^ cache: false$/m, name + '/' + job); + checked++; + } + } + } + assert.ok(checked >= 10, 'expected build and lightweight jobs to be covered'); +}); + +test('only main writes caches, and the nightly read-only cache stays read-only', () => { + let checked = 0; + for (const { name, text } of workflows) { + for (const [, job, body] of jobs(text)) { + for (const step of steps(body)) { + if (!step.includes('uses: Swatinem/rust-cache@') && !step.includes('uses: actions/cache/save@')) continue; + const condition = step.includes('uses: Swatinem/rust-cache@') + ? step.match(/^ save-if: (.+)$/m)?.[1] + : step.match(/^ if: (.+)$/m)?.[1]; + assert.ok(condition, name + '/' + job); + const expression = condition.replace(/^\$\{\{\s*|\s*\}\}$/g, '') + .replaceAll('steps.fuzz-binaries.outputs.cache-hit', 'hit'); + const enabled = new Function('github', 'hit', 'return (' + expression + ');'); + for (const [event, ref] of [ + ['pull_request', 'refs/pull/123/merge'], + ['merge_group', 'refs/heads/gh-readonly-queue/main/pr-123'], + ['workflow_dispatch', 'refs/heads/feature'], + ]) assert.equal(enabled({ event_name: event, ref }, 'false'), false, name + '/' + job); + assert.equal(enabled({ event_name: 'schedule', ref: 'refs/heads/main' }, 'false'), + condition !== 'false', name + '/' + job); + if (step.includes('uses: actions/cache/save@')) { + assert.equal(enabled({ ref: 'refs/heads/main' }, 'true'), false); + } + checked++; + } + } + } + assert.ok(checked >= 10); +}); + +test('helper tool versions do not fragment unrelated Rust caches', () => { + for (const { text } of workflows) { + assert.doesNotMatch(text, /CARGO_(SHEAR|FUZZ)_VERSION/); + } + const integration = workflows.find(w => w.name === 'integration').text; + assert.match(integration, /prefix-key: "v1-fuzz-\$\{\{ env.FUZZ_VERSION \}\}"/); + assert.match(integration, /cargo install cargo-fuzz --version "\$\{FUZZ_VERSION\}"/); +}); + +// Evaluate the actual job guard as well as the cache's save condition. A valid +// save-if alone cannot prove that a scheduled cache producer is reachable. +function warmupJob(workflow, job, { + event = 'schedule', ref = 'refs/heads/main', + profile = event === 'workflow_dispatch' ? 'all' : undefined, + outputs = {}, failed = '', cancelled = false, draft = false, +} = {}) { + const text = workflows.find(w => w.name === workflow).text; + assert.match(text, /^ schedule:$/m); + const body = jobs(text).find(([, name]) => name === job)?.[2]; + assert.ok(body, workflow + '/' + job); + const github = { repository: 'GreptimeTeam/greptimedb', event_name: event, ref, + event: { pull_request: { draft } }, run_attempt: 1 }; + const needs = { scope: { outputs }, ...Object.fromEntries( + ['github-script-tests', 'build', 'prepare-fuzz-targets', 'build-greptime-ci'].map(name => + [name, { result: name === failed ? 'failure' : 'success' }])) }; + const evaluate = expression => new Function('github', 'inputs', 'needs', 'cancelled', + 'return (' + expression.replace(/^\$\{\{\s*|\s*\}\}$/g, '') + ');')( + github, { fuzz_profile: profile, ci_command: false }, needs, () => cancelled); + const admitted = evaluate(body.match(/^ if: (.+)$/m)[1]); + const saves = [...body.matchAll(/^ save-if: (.+)$/gm)] + .map(match => admitted && evaluate(match[1])); + return { admitted, saves }; +} + +const fuzzJobs = ['prepare-fuzz-targets', 'build-greptime-ci', 'fuzztest', + 'unstable-fuzztest', 'distributed-fuzztest', 'distributed-fuzztest-with-chaos']; + +test('main schedules can execute cache producers and their fuzz consumers', () => { + for (const job of fuzzJobs) { + // Scheduled events do not supply workflow_dispatch inputs. + const result = warmupJob('integration', job, { profile: undefined }); + assert.equal(result.admitted, true, job); + assert.ok(result.saves.every(Boolean), job); + } + for (const [workflow, job] of [ + ['integration', 'build'], ['integration', 'export-import-v2-e2e'], + ['integration', 'prepare-fuzz-targets'], ['integration', 'build-greptime-ci'], + ['integration', 'distributed-fuzztest'], + ['rust', 'check-windows'], ['rust', 'check-riscv64'], + ['nightly-ci', 'sqlness-windows'], ['nightly-ci', 'test-on-windows'], + ]) { + for (const event of ['schedule', 'workflow_dispatch']) { + assert.deepEqual(warmupJob(workflow, job, { event }), + { admitted: true, saves: [true] }, workflow + '/' + job + '/' + event); + } + } +}); + +test('manual warming respects selected profiles without an internal command marker', () => { + const expected = { + none: [], + standalone: fuzzJobs.slice(0, 4), + distributed: [fuzzJobs[0], fuzzJobs[1], fuzzJobs[4]], + chaos: [fuzzJobs[0], fuzzJobs[1], fuzzJobs[5]], + all: fuzzJobs, + }; + for (const [profile, wanted] of Object.entries(expected)) { + assert.deepEqual(fuzzJobs.filter(job => warmupJob('integration', job, + { event: 'workflow_dispatch', profile }).admitted), wanted, profile); + } + for (const job of fuzzJobs) { + const result = warmupJob('integration', job, + { event: 'workflow_dispatch', ref: 'refs/heads/feature' }); + assert.ok(result.saves.every(value => value === false), job); + assert.equal(warmupJob('integration', job, + { event: 'pull_request', ref: 'refs/pull/1/merge', draft: true }).admitted, false, job); + } +}); + +test('warmup keeps scope exclusions, dependency failures and cancellation gates', () => { + const gates = ['any_fuzz', 'build_ci_needed', 'fuzz_standalone', + 'fuzz_unstable', 'fuzz_distributed', 'fuzz_chaos']; + for (const [index, job] of fuzzJobs.entries()) { + assert.equal(warmupJob('integration', job, { failed: 'github-script-tests' }).admitted, false, job); + assert.equal(warmupJob('integration', job, { cancelled: true }).admitted, false, job); + assert.equal(warmupJob('integration', job, + { outputs: { [gates[index]]: 'false' } }).admitted, false, job); + } + for (const job of fuzzJobs.slice(2)) { + for (const failed of ['github-script-tests', 'prepare-fuzz-targets', job === 'fuzztest' ? 'build' : 'build-greptime-ci']) { + assert.equal(warmupJob('integration', job, { failed }).admitted, false, job); + } + } +}); diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml index 1f14d8851cc..0d38a844539 100644 --- a/.github/workflows/checks.yml +++ b/.github/workflows/checks.yml @@ -107,6 +107,8 @@ jobs: with: tool: protoc@3.23.4 - uses: actions-rust-lang/setup-rust-toolchain@v1 + with: + cache: false - name: Rust Cache uses: Swatinem/rust-cache@v2 with: @@ -130,6 +132,8 @@ jobs: with: persist-credentials: false - uses: actions-rust-lang/setup-rust-toolchain@v1 + with: + cache: false - name: Install taplo uses: taiki-e/install-action@v2 with: diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml index 00c3f52a034..134fdfde6ba 100644 --- a/.github/workflows/integration.yml +++ b/.github/workflows/integration.yml @@ -30,7 +30,7 @@ on: name: Integration CI env: - CARGO_FUZZ_VERSION: "0.13.2" + FUZZ_VERSION: "0.13.2" # Fuzz targets require a nightly toolchain (cargo-fuzz needs # `-Zsanitizer=fuzzer`). This pin decouples the fuzz workflow from the # workspace default in rust-toolchain.toml, so it keeps working when the @@ -121,6 +121,8 @@ jobs: with: tool: protoc@3.23.4 - uses: actions-rust-lang/setup-rust-toolchain@v1 + with: + cache: false - name: Build Binaries Cache id: build-binaries-cache uses: Swatinem/rust-cache@v2 @@ -292,6 +294,8 @@ jobs: with: tool: protoc@3.23.4 - uses: actions-rust-lang/setup-rust-toolchain@v1 + with: + cache: false - name: Rust Cache uses: Swatinem/rust-cache@v2 with: @@ -436,18 +440,18 @@ jobs: id: fuzz-dependencies-cache uses: Swatinem/rust-cache@v2 with: - prefix-key: "v1-fuzz-${{ env.CARGO_FUZZ_VERSION }}" + prefix-key: "v1-fuzz-${{ env.FUZZ_VERSION }}" shared-key: "fuzz-dependencies" workspaces: ". -> target" cache-targets: "true" cache-all-crates: "true" - save-if: ${{ github.event_name == 'pull_request' }} + save-if: ${{ github.ref == 'refs/heads/main' }} - name: Look up prebuilt fuzz binaries id: fuzz-binaries uses: actions/cache/restore@v4 with: path: ${{ runner.temp }}/greptime-fuzz-binaries - key: v1-fuzz-binaries-${{ env.CARGO_FUZZ_VERSION }}-${{ env.FUZZ_RUST_TOOLCHAIN }}-all-${{ hashFiles('Cargo.toml', 'Cargo.lock', '.cargo/**', 'tests-fuzz/**', 'src/**') }} + key: v1-fuzz-binaries-${{ env.FUZZ_VERSION }}-${{ env.FUZZ_RUST_TOOLCHAIN }}-all-${{ hashFiles('Cargo.toml', 'Cargo.lock', '.cargo/**', 'tests-fuzz/**', 'src/**') }} - name: Report fuzz cache lookup shell: bash run: | @@ -462,7 +466,7 @@ jobs: env: RUSTUP_TOOLCHAIN: ${{ env.FUZZ_RUST_TOOLCHAIN }} run: | - cargo install cargo-fuzz --version "${CARGO_FUZZ_VERSION}" --locked + cargo install cargo-fuzz --version "${FUZZ_VERSION}" --locked - name: Build fuzz target binaries if: steps.fuzz-binaries.outputs.cache-hit != 'true' shell: bash @@ -511,11 +515,11 @@ jobs: build_targets false "${FUZZ_TARGETS}" build_targets true "${FUZZ_UNSTABLE_TARGETS}" - name: Save prebuilt fuzz binaries - if: steps.fuzz-binaries.outputs.cache-hit != 'true' + if: github.ref == 'refs/heads/main' && steps.fuzz-binaries.outputs.cache-hit != 'true' uses: actions/cache/save@v4 with: path: ${{ runner.temp }}/greptime-fuzz-binaries - key: v1-fuzz-binaries-${{ env.CARGO_FUZZ_VERSION }}-${{ env.FUZZ_RUST_TOOLCHAIN }}-all-${{ hashFiles('Cargo.toml', 'Cargo.lock', '.cargo/**', 'tests-fuzz/**', 'src/**') }} + key: v1-fuzz-binaries-${{ env.FUZZ_VERSION }}-${{ env.FUZZ_RUST_TOOLCHAIN }}-all-${{ hashFiles('Cargo.toml', 'Cargo.lock', '.cargo/**', 'tests-fuzz/**', 'src/**') }} - name: Pack prebuilt fuzz binaries shell: bash run: | @@ -688,6 +692,8 @@ jobs: with: tool: protoc@3.23.4 - uses: actions-rust-lang/setup-rust-toolchain@v1 + with: + cache: false - name: Build Profile-CI Cache id: build-greptime-ci-cache uses: Swatinem/rust-cache@v2 @@ -696,7 +702,7 @@ jobs: shared-key: "build-greptime-ci" cache-all-crates: "true" cache-workspace-crates: "true" - save-if: ${{ github.event_name == 'pull_request' }} + save-if: ${{ github.ref == 'refs/heads/main' }} - name: Report profile-CI cache lookup shell: bash run: | @@ -819,7 +825,7 @@ jobs: workspaces: ". -> target" cache-targets: "true" cache-all-crates: "true" - save-if: ${{ github.event_name == 'pull_request' }} + save-if: ${{ github.ref == 'refs/heads/main' }} - if: matrix.mode.kafkaWalHelper name: Report Kafka WAL helper cache lookup shell: bash diff --git a/.github/workflows/nightly-ci.yml b/.github/workflows/nightly-ci.yml index f6158a8d7ad..ffd1b01d477 100644 --- a/.github/workflows/nightly-ci.yml +++ b/.github/workflows/nightly-ci.yml @@ -57,8 +57,12 @@ jobs: with: tool: protoc@3.23.4 - uses: actions-rust-lang/setup-rust-toolchain@v1 + with: + cache: false - name: Rust Cache uses: Swatinem/rust-cache@v2 + with: + save-if: ${{ github.ref == 'refs/heads/main' }} - name: Run sqlness run: make sqlness-test env: @@ -143,9 +147,12 @@ jobs: - name: Install Rust toolchain uses: actions-rust-lang/setup-rust-toolchain@v1 with: + cache: false components: llvm-tools-preview - name: Rust Cache uses: Swatinem/rust-cache@v2 + with: + save-if: ${{ github.ref == 'refs/heads/main' }} - name: Install Cargo Nextest uses: taiki-e/install-action@nextest - name: Install WSL distribution diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 842fa306c4e..d66bc042b5b 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -24,7 +24,7 @@ on: name: Rust CI env: - CARGO_SHEAR_VERSION: "1.13.4" + SHEAR_VERSION: "1.13.4" concurrency: group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }} @@ -56,6 +56,7 @@ jobs: tool: protoc@3.23.4 - uses: actions-rust-lang/setup-rust-toolchain@v1 with: + cache: false components: rustfmt - name: Check format run: make fmt-check @@ -75,6 +76,7 @@ jobs: tool: protoc@3.23.4 - uses: actions-rust-lang/setup-rust-toolchain@v1 with: + cache: false components: clippy - name: Rust Cache uses: Swatinem/rust-cache@v2 @@ -97,10 +99,12 @@ jobs: with: persist-credentials: false - uses: actions-rust-lang/setup-rust-toolchain@v1 + with: + cache: false - name: Install cargo-shear uses: taiki-e/install-action@v2 with: - tool: cargo-shear@${{ env.CARGO_SHEAR_VERSION }} + tool: cargo-shear@${{ env.SHEAR_VERSION }} fallback: none - name: Check unused dependencies run: make check-udeps @@ -122,6 +126,8 @@ jobs: - name: Install riscv64 cross toolchain run: sudo apt-get update && sudo apt-get install -y gcc-riscv64-linux-gnu g++-riscv64-linux-gnu - uses: actions-rust-lang/setup-rust-toolchain@v1 + with: + cache: false - name: Add riscv64 target run: rustup target add riscv64gc-unknown-linux-gnu - name: Rust Cache @@ -153,6 +159,8 @@ jobs: with: version: "14.0" - uses: actions-rust-lang/setup-rust-toolchain@v1 + with: + cache: false - name: Rust Cache uses: Swatinem/rust-cache@v2 with: From db03b2341e7c635576e4efbe31b9e4ffabf1f152 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Fri, 30 Oct 2026 00:12:42 +0000 Subject: [PATCH 143/239] feat(meta): orchestrate datanode draining with detached migrations (#9961) * feat(meta): orchestrate datanode draining with detached migrations Signed-off-by: WenyXu * fix(meta): recheck drain source after migration route races Signed-off-by: WenyXu * fix(meta): skip dropped tables during datanode drain Signed-off-by: WenyXu * fix(meta): skip missing table info during datanode drain Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- src/meta-srv/src/error.rs | 15 +- src/meta-srv/src/procedure/drain_datanode.rs | 498 ++++++++++++-- .../procedure/drain_datanode/check_drained.rs | 233 +++++++ .../src/procedure/drain_datanode/dispatch.rs | 629 ++++++++++++++++++ .../drain_datanode/drain_datanode_end.rs | 69 +- .../drain_datanode/drain_datanode_start.rs | 43 +- .../src/procedure/drain_datanode/manager.rs | 101 ++- .../src/procedure/drain_datanode/migration.rs | 40 +- .../src/procedure/drain_datanode/test_util.rs | 162 ++++- .../src/procedure/drain_datanode/utils.rs | 43 +- .../procedure/region_migration/test_util.rs | 11 + 11 files changed, 1728 insertions(+), 116 deletions(-) create mode 100644 src/meta-srv/src/procedure/drain_datanode/check_drained.rs create mode 100644 src/meta-srv/src/procedure/drain_datanode/dispatch.rs diff --git a/src/meta-srv/src/error.rs b/src/meta-srv/src/error.rs index 65d86f9b8ce..8e6c9097dcc 100644 --- a/src/meta-srv/src/error.rs +++ b/src/meta-srv/src/error.rs @@ -129,6 +129,18 @@ pub enum Error { location: Location, }, + #[snafu(display( + "Drain datanode {} exceeded the maximum rescan rounds: {}", + node_id, + max_rounds + ))] + DrainDatanodeRescanExceeded { + node_id: DatanodeId, + max_rounds: u32, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display( "Another procedure is operating the region: {} on peer: {}", region_id, @@ -1273,9 +1285,10 @@ impl ErrorExt for Error { | Error::RegionRouteNotFound { .. } | Error::MigrationAbort { .. } | Error::MigrationRunning { .. } + | Error::RegionMigrated { .. } | Error::DrainDatanodeRunning { .. } | Error::DrainDatanodeAbort { .. } - | Error::RegionMigrated { .. } => StatusCode::Unexpected, + | Error::DrainDatanodeRescanExceeded { .. } => StatusCode::Unexpected, Error::TableNotFound { .. } => StatusCode::TableNotFound, Error::SaveClusterInfo { source, .. } | Error::InvalidClusterInfoFormat { source, .. } diff --git a/src/meta-srv/src/procedure/drain_datanode.rs b/src/meta-srv/src/procedure/drain_datanode.rs index aa8d333478e..87da8108428 100644 --- a/src/meta-srv/src/procedure/drain_datanode.rs +++ b/src/meta-srv/src/procedure/drain_datanode.rs @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +pub(crate) mod check_drained; +pub(crate) mod dispatch; pub(crate) mod drain_datanode_end; pub(crate) mod drain_datanode_start; pub(crate) mod manager; @@ -23,6 +25,7 @@ mod test_util; use std::any::Any; use std::fmt::Debug; +use std::time::Duration; use common_meta::DatanodeId; use common_meta::key::TableMetadataManagerRef; @@ -36,7 +39,7 @@ use serde::{Deserialize, Serialize}; use snafu::ResultExt; use store_api::storage::RegionId; -use crate::error::Result; +use crate::error::{self, Result}; use crate::metasrv::SelectorContext; use crate::node_scheduling::NodeSchedulingManagerRef; use crate::procedure::drain_datanode::drain_datanode_start::DrainDatanodeStart; @@ -45,6 +48,15 @@ use crate::procedure::drain_datanode::manager::{ }; use crate::procedure::region_migration::manager::RegionMigrationManagerRef; +/// The default number of leader regions migrated concurrently. +pub const DEFAULT_PARALLELISM: usize = 4; +/// The maximum number of leader regions migrated concurrently. +pub const MAX_PARALLELISM: usize = 32; +/// The default budget of the whole drain. +pub const DEFAULT_DRAIN_TIMEOUT: Duration = Duration::from_secs(20 * 60); +/// The maximum number of plan build attempts, including the first scan and failed reads. +pub const MAX_RESCAN_ROUNDS: u32 = 3; + /// The strategy of a drain run when a region cannot be migrated. #[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] pub enum DrainDatanodeStrategy { @@ -67,14 +79,20 @@ pub struct DrainDatanodePlan { /// The persistent data of [`DrainDatanodeProcedure`]. /// -/// Only what cannot be derived from the cluster metadata is persisted: the plan is rebuilt -/// from the authoritative table routes after a restart. +/// Only what cannot be derived from the cluster metadata is persisted: the plan and the +/// progress of a run are rebuilt from the authoritative table routes after a restart. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct PersistentContext { /// The datanode to drain. pub node_id: DatanodeId, - /// The pinned destination of leader regions. + /// The pinned destination of leader regions. Resolved by the selector when `None`. pub target_node_id: Option, + /// The number of leader regions migrated concurrently. + pub parallelism: usize, + /// The error handling strategy. + pub strategy: DrainDatanodeStrategy, + /// The absolute deadline of the whole drain, in milliseconds since the Unix epoch. + pub deadline_ms: i64, } /// A dispatched migration sub-procedure that is not settled yet. @@ -88,10 +106,17 @@ pub struct SubprocedureMeta { /// The in-memory data of [`DrainDatanodeProcedure`]. /// -/// It is not persisted: a restarted procedure rebuilds the plan by rescanning the datanode. +/// It is not persisted: a restart rebuilds the plan by rescanning the datanode, which is +/// bounded by [`MAX_RESCAN_ROUNDS`], and re-dispatches the regions that are still led by +/// the drained datanode and have no running migration. Recovered migrations own their +/// lifecycle and may finish after this drain procedure ends. #[derive(Debug, Clone, Default)] pub struct VolatileContext { - /// The drain plan, `None` until the first state builds it. + /// The drain plan, `None` until it is built by the first state that runs. + /// + /// [`DrainDatanodePlan::leader_regions`] is the remaining work: a region leaves it when + /// its migration succeeded, when it was given up by [`DrainDatanodeStrategy::TryAll`], + /// when the drained datanode stopped leading it, or when another migration owns it. pub plan: Option, /// The migration sub-procedures of the current batch that are not settled yet. pub inflight_subprocedures: Vec, @@ -99,6 +124,8 @@ pub struct VolatileContext { pub migrated: usize, /// The regions given up by [`DrainDatanodeStrategy::TryAll`] in this run. pub failed_regions: Vec, + /// The number of plan build attempts in this run, including failed reads. + pub rescan_round: u32, } /// The task of [`DrainDatanodeProcedure`]. @@ -106,8 +133,14 @@ pub struct VolatileContext { pub struct DrainDatanodeProcedureTask { /// The datanode to drain. pub node_id: DatanodeId, - /// The pinned destination of leader regions. + /// The pinned destination of leader regions. Resolved by the selector when `None`. pub target_node_id: Option, + /// The number of leader regions migrated concurrently. + pub parallelism: usize, + /// The error handling strategy. + pub strategy: DrainDatanodeStrategy, + /// The budget of the whole drain. `None` uses [`DEFAULT_DRAIN_TIMEOUT`]. + pub timeout: Option, } /// The context of [`DrainDatanodeProcedure`]. @@ -121,22 +154,23 @@ pub struct Context { pub table_metadata_manager: TableMetadataManagerRef, /// The node scheduling manager, used to fence the drained datanode. pub node_scheduling_manager: NodeSchedulingManagerRef, - /// Builds region migration procedures. + /// The region migration manager, used to build migration sub-procedures. pub region_migration_manager: RegionMigrationManagerRef, - /// Resolves destinations for region migrations. + /// The selector context, used to resolve the destination of a leader region. pub selector_ctx: SelectorContext, } impl Context { /// Builds the context of a run. /// - /// The per-run state starts empty: the plan is rebuilt from the table routes when the - /// procedure needs it. + /// The per-run state starts empty: the plan and the progress are rebuilt from the + /// table routes when the procedure needs them. pub(crate) fn new( table_metadata_manager: TableMetadataManagerRef, node_scheduling_manager: NodeSchedulingManagerRef, region_migration_manager: RegionMigrationManagerRef, selector_ctx: SelectorContext, + persistent_ctx: PersistentContext, ) -> Self { Self { @@ -151,24 +185,27 @@ impl Context { } impl Context { - /// Builds the drain plan when it is not available yet. + /// Rebuilds the drain plan from the authoritative table routes. /// - /// The plan is derived from the authoritative table routes, so a restarted procedure - /// rebuilds it instead of reading a persisted copy. - pub(crate) async fn ensure_plan(&mut self) -> Result<()> { - if self.volatile_ctx.plan.is_some() { - return Ok(()); + /// Builds are bounded by [`MAX_RESCAN_ROUNDS`]: a datanode whose leader regions keep + /// coming back must not be rescanned forever. + pub(crate) async fn rebuild_plan(&mut self) -> Result<()> { + let node_id = self.persistent_ctx.node_id; + if self.volatile_ctx.rescan_round >= MAX_RESCAN_ROUNDS { + return error::DrainDatanodeRescanExceededSnafu { + node_id, + max_rounds: MAX_RESCAN_ROUNDS, + } + .fail(); } - let node_id = self.persistent_ctx.node_id; + self.volatile_ctx.rescan_round += 1; let plan = utils::build_datanode_plan(&self.table_metadata_manager, node_id).await?; self.volatile_ctx.plan = Some(plan); Ok(()) } -} -impl Context { /// Removes `region_id` from the current plan. A later scan may add it back. pub(crate) fn forget_leader_region(&mut self, region_id: RegionId) { if let Some(plan) = self.volatile_ctx.plan.as_mut() { @@ -185,7 +222,7 @@ impl VolatileContext { .map_or(&[], |plan| plan.leader_regions.as_slice()) } - /// Returns the follower regions collected by the plan. + /// Returns the follower regions collected by the last plan build. pub(crate) fn follower_regions(&self) -> &[RegionId] { self.plan .as_ref() @@ -230,8 +267,11 @@ pub(crate) trait State: Sync + Send + Debug { /// The procedure draining a datanode. /// -/// It fences the datanode and collects the leader and follower regions placed on it. -/// Migrating the leader regions off the datanode lands in a follow-up. +/// It fences the datanode, migrates every leader region off it through +/// [`RegionMigrationProcedure`](crate::procedure::region_migration::RegionMigrationProcedure) +/// sub-procedures, and rescans for regions that still need a new migration. Migrations +/// already registered in the tracker continue independently and may outlive this +/// procedure. Follower regions are collected but not removed yet. pub(crate) struct DrainDatanodeProcedure { state: Box, context: Context, @@ -272,13 +312,17 @@ impl DrainDatanodeProcedure { .insert_running_procedure(&DrainDatanodeProcedureTask { node_id: persistent_ctx.node_id, target_node_id: persistent_ctx.target_node_id, + parallelism: persistent_ctx.parallelism, + strategy: persistent_ctx.strategy, + timeout: None, }) .into_iter() .collect(); + let context = ctx_factory(persistent_ctx); Ok(Self { state, - context: ctx_factory(persistent_ctx), + context, _guards: guards, }) } @@ -338,25 +382,42 @@ impl Procedure for DrainDatanodeProcedure { #[cfg(test)] mod tests { - use super::*; + use std::time::Duration; + + use common_error::ext::ErrorExt; + use common_error::status_code::StatusCode; + use common_meta::state_store::KvStateStore; + use common_procedure::store::ProcedureMessage; + use common_procedure::store::state_store::StateStore; + use common_procedure::{Procedure, ProcedureId, ProcedureState, Status}; + use futures::TryStreamExt; + use store_api::storage::RegionId; + + use crate::error; use crate::key::NodeSchedulingRole; + use crate::procedure::drain_datanode::drain_datanode_start::DrainDatanodeStart; + use crate::procedure::drain_datanode::manager::DrainDatanodeProcedureTracker; use crate::procedure::drain_datanode::test_util::{ TestEnv, follower_region, leader_region, new_persistent_context, run_until_done, + spawn_mock_datanode, wait_for_terminal_state, + }; + use crate::procedure::drain_datanode::{ + DEFAULT_PARALLELISM, DrainDatanodeProcedure, DrainDatanodeProcedureTask, + DrainDatanodeStrategy, State, }; use crate::procedure::region_migration::test_util::TestingEnv; + use crate::service::mailbox::Channel; #[tokio::test] - async fn test_drain_datanode_fences_and_collects_the_plan() { + async fn test_drain_datanode_fences_and_completes_without_leaders() { let env = TestEnv::new().await; - let region_id = RegionId::new(1024, 1); - let follower_region_id = RegionId::new(1024, 2); - env.create_table(vec![ - follower_region(follower_region_id, 2, 1), - leader_region(region_id, 1), - ]) - .await; + env.create_table(vec![follower_region(RegionId::new(1024, 1), 2, 1)]) + .await; - let mut ctx = env.context(new_persistent_context(Some(2))); + let mut ctx = env.context(new_persistent_context( + None, + DrainDatanodeStrategy::FailFast, + )); let procedure_ctx = TestingEnv::procedure_context(); let mut state: Box = Box::new(DrainDatanodeStart); @@ -365,13 +426,13 @@ mod tests { .unwrap(); assert!(matches!(status, Status::Done { .. })); - assert_eq!(vec![region_id], ctx.volatile_ctx.leader_regions()); + assert_eq!(0, ctx.volatile_ctx.migrated); assert_eq!( - vec![follower_region_id], + vec![RegionId::new(1024, 1)], ctx.volatile_ctx.follower_regions() ); assert!( - !ctx.node_scheduling_manager + !env.node_scheduling_manager .is_schedulable(NodeSchedulingRole::Datanode, 1) .await .unwrap() @@ -381,7 +442,7 @@ mod tests { #[tokio::test] async fn test_drain_datanode_dump_roundtrip() { let env = TestEnv::new().await; - let persistent_ctx = new_persistent_context(Some(2)); + let persistent_ctx = new_persistent_context(Some(2), DrainDatanodeStrategy::TryAll); let procedure = DrainDatanodeProcedure::new(env.context(persistent_ctx), vec![]); let json = procedure.dump().unwrap(); @@ -395,7 +456,366 @@ mod tests { let ctx = &reloaded.context.persistent_ctx; assert_eq!(1, ctx.node_id); assert_eq!(Some(2), ctx.target_node_id); - // The plan is volatile, so it is not dumped. + assert_eq!(DEFAULT_PARALLELISM, ctx.parallelism); + assert_eq!(DrainDatanodeStrategy::TryAll, ctx.strategy); + // The plan and the progress of a run are volatile, so they are not dumped. assert!(!json.contains("leader_regions")); + assert!(!json.contains("rescan_round")); + } + + #[tokio::test] + async fn test_drain_datanode_default_strategy_continues_after_batch_failures() { + let env = TestEnv::new().await; + let regions: Vec<_> = (1..=5).map(|number| RegionId::new(1024, number)).collect(); + env.create_table( + regions + .iter() + .copied() + .map(|region_id| leader_region(region_id, 1)) + .collect(), + ) + .await; + // The pinned target is unschedulable, so every region is given up. + env.node_scheduling_manager + .mark_unschedulable(NodeSchedulingRole::Datanode, 3) + .await + .unwrap(); + + let mut persistent_ctx = new_persistent_context(Some(3), DrainDatanodeStrategy::default()); + persistent_ctx.parallelism = 2; + let mut ctx = env.context(persistent_ctx); + let procedure_ctx = TestingEnv::procedure_context(); + let mut state: Box = Box::new(DrainDatanodeStart); + + let error = run_until_done(&mut state, &mut ctx, &procedure_ctx) + .await + .unwrap_err(); + + let error::Error::DrainDatanodeAbort { failed_regions, .. } = error else { + panic!("expected the drain procedure to abort"); + }; + assert_eq!(regions, failed_regions); + assert_eq!(regions, ctx.volatile_ctx.failed_regions); + assert_eq!(0, ctx.volatile_ctx.migrated); + } + + #[tokio::test] + async fn test_drain_datanode_recovery_obeys_deadline() { + for expired in [true, false] { + let mut env = TestEnv::new().await; + let first_id = RegionId::new(1024, 1); + let second_id = RegionId::new(1024, 2); + env.create_table(vec![ + leader_region(first_id, 1), + leader_region(second_id, 1), + ]) + .await; + env.mock_datanodes().await; + let mut persistent_ctx = + new_persistent_context(Some(2), DrainDatanodeStrategy::FailFast); + if expired { + persistent_ctx.deadline_ms = 0; + } + let mut procedure = + DrainDatanodeProcedure::new(env.context(persistent_ctx.clone()), vec![]); + assert!( + procedure + .execute(&TestingEnv::procedure_context()) + .await + .unwrap() + .need_persist() + ); + let checkpoint = procedure.dump().unwrap(); + let json: serde_json::Value = serde_json::from_str(&checkpoint).unwrap(); + assert_eq!( + serde_json::json!({ + "persistent_ctx": { + "node_id": 1, "target_node_id": 2, + "parallelism": DEFAULT_PARALLELISM, "strategy": "FailFast", + "deadline_ms": persistent_ctx.deadline_ms, + }, + "state": {"drain_datanode_state": "CheckDrained"}, + }), + json + ); + drop(procedure); + let recovered = DrainDatanodeProcedure::from_json( + &checkpoint, + &DrainDatanodeProcedureTracker::default(), + |persistent_ctx| env.context(persistent_ctx), + ) + .unwrap(); + assert!(recovered.context.volatile_ctx.plan.is_none()); + drop(recovered); + let id = ProcedureId::random(); + let message = common_procedure::store::ProcedureMessage { + type_name: DrainDatanodeProcedure::TYPE_NAME.to_string(), + data: checkpoint, + parent_id: None, + subprocedure_recovery: Default::default(), + step: 0, + error: None, + context: Default::default(), + }; + let store = common_meta::state_store::KvStateStore::new(env.env.kv_backend()); + common_procedure::store::state_store::StateStore::put( + &store, + &format!("procedure/{id}/0000000000.step"), + serde_json::to_vec(&message).unwrap(), + ) + .await + .unwrap(); + let manager = env.drain_manager(); + manager.try_start().unwrap(); + env.procedure_manager().start().await.unwrap(); + let state = wait_for_terminal_state(env.procedure_manager(), id).await; + if expired { + let ProcedureState::Failed { error } = state else { + panic!("expected expired drain: {state:?}"); + }; + assert_eq!(StatusCode::Cancelled, error.status_code()); + } else { + assert!(state.is_done(), "{state:?}"); + } + let table_route = env + .table_metadata_manager + .table_route_manager() + .table_route_storage() + .get(1024) + .await + .unwrap() + .unwrap(); + let routes = table_route.region_routes().unwrap(); + assert_eq!(2, routes.len()); + for route in routes { + assert_eq!( + Some(if expired { 1 } else { 2 }), + route.leader_peer.as_ref().map(|peer| peer.id) + ); + } + } + } + + #[tokio::test] + async fn test_drain_datanode_recovered_migration_outlives_parent() { + let mut env = TestEnv::new().await; + let region_id = RegionId::new(1024, 1); + env.create_table(vec![leader_region(region_id, 1)]).await; + let procedure_ctx = TestingEnv::procedure_context(); + let mut parent = DrainDatanodeProcedure::new( + env.context(new_persistent_context( + Some(2), + DrainDatanodeStrategy::FailFast, + )), + vec![], + ); + assert!(parent.execute(&procedure_ctx).await.unwrap().need_persist()); + let parent_checkpoint = parent.dump().unwrap(); + let Status::Suspended { + subprocedures: mut children, + recovery, + .. + } = run_until_done(&mut parent.state, &mut parent.context, &procedure_ctx) + .await + .unwrap() + else { + panic!("expected a migration batch"); + }; + assert_eq!(1, children.len()); + let mut child = children.pop().unwrap(); + // Persist the first migration checkpoint, before it opens the candidate region. + assert!( + child + .procedure + .execute(&procedure_ctx) + .await + .unwrap() + .need_persist() + ); + let parent_id = ProcedureId::random(); + let child_id = child.id; + let store = KvStateStore::new(env.env.kv_backend()); + for (id, message) in [ + ( + parent_id, + ProcedureMessage { + type_name: parent.type_name().to_string(), + data: parent_checkpoint, + parent_id: None, + subprocedure_recovery: Default::default(), + step: 0, + error: None, + context: Default::default(), + }, + ), + ( + child_id, + ProcedureMessage { + type_name: child.procedure.type_name().to_string(), + data: child.procedure.dump().unwrap(), + parent_id: Some(parent_id), + subprocedure_recovery: recovery, + step: 0, + error: None, + context: child.context.clone(), + }, + ), + ] { + store + .put( + &format!("procedure/{id}/0000000000.step"), + serde_json::to_vec(&message).unwrap(), + ) + .await + .unwrap(); + } + drop(child); + drop(parent); + assert!(!env.region_migration_manager.tracker().contains(region_id)); + + // Queue mailbox requests without replying until the recovered parent has ended. + let mailbox_ctx = env.env.mailbox_context(); + let mut receivers = Vec::new(); + for node_id in [1, 2] { + let (tx, rx) = tokio::sync::mpsc::channel(8); + mailbox_ctx + .insert_heartbeat_response_receiver(Channel::Datanode(node_id), tx) + .await; + receivers.push(rx); + } + let mailbox = mailbox_ctx.mailbox().clone(); + env.region_migration_manager.try_start().unwrap(); + let drain_manager = env.drain_manager(); + drain_manager.try_start().unwrap(); + env.procedure_manager().start().await.unwrap(); + assert!( + wait_for_terminal_state(env.procedure_manager(), parent_id) + .await + .is_done() + ); + assert!(env.region_migration_manager.tracker().contains(region_id)); + assert!(matches!( + env.procedure_manager() + .procedure_state(child_id) + .await + .unwrap(), + Some(ProcedureState::Running) + )); + // No replacement migration was submitted by the recovering drain. + assert_eq!( + 2, + env.procedure_manager() + .list_procedures() + .await + .unwrap() + .len() + ); + + for rx in receivers { + spawn_mock_datanode(mailbox.clone(), rx); + } + assert!( + wait_for_terminal_state(env.procedure_manager(), child_id) + .await + .is_done() + ); + assert!( + !crate::procedure::drain_datanode::utils::leads_region( + &env.table_metadata_manager, + 1, + region_id, + ) + .await + .unwrap() + ); + // Completion must release the recovered guard and clean its own checkpoints. + tokio::time::timeout(Duration::from_secs(5), async { + loop { + let mut records = store.walk_top_down("procedure/").await.unwrap(); + if records.try_next().await.unwrap().is_none() + && !env.region_migration_manager.tracker().contains(region_id) + { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("recovered migration did not clean up"); + env.procedure_manager().stop().await.unwrap(); + } + + #[tokio::test] + async fn test_drain_datanode_migrates_leader_region() { + assert_drain_migrates_leader_regions(1).await; + } + + #[tokio::test] + async fn test_drain_datanode_migrates_multiple_batches() { + assert_drain_migrates_leader_regions(3).await; + } + + async fn assert_drain_migrates_leader_regions(region_count: u32) { + let mut env = TestEnv::new().await; + let regions: Vec<_> = (1..=region_count) + .map(|number| RegionId::new(1024, number)) + .collect(); + env.create_table( + regions + .iter() + .copied() + .map(|region_id| leader_region(region_id, 1)) + .collect(), + ) + .await; + + env.mock_datanodes().await; + + env.procedure_manager().start().await.unwrap(); + let manager = env.drain_manager(); + let procedure_id = manager + .submit_procedure( + common_procedure::ProcedureContext::default(), + DrainDatanodeProcedureTask { + node_id: 1, + target_node_id: Some(2), + parallelism: 1, + strategy: DrainDatanodeStrategy::FailFast, + timeout: None, + }, + ) + .await + .unwrap(); + + let state = wait_for_terminal_state(env.procedure_manager(), procedure_id).await; + assert!( + matches!(state, ProcedureState::Done { .. }), + "the drain procedure did not complete successfully: {state:?}" + ); + + let table_route = env + .table_metadata_manager + .table_route_manager() + .table_route_storage() + .get(1024) + .await + .unwrap() + .unwrap(); + assert_eq!( + region_count as usize, + table_route.region_routes().unwrap().len() + ); + assert!( + !env.node_scheduling_manager + .is_schedulable(NodeSchedulingRole::Datanode, 1) + .await + .unwrap() + ); + for region_route in table_route.region_routes().unwrap() { + assert_eq!( + Some(2), + region_route.leader_peer.as_ref().map(|peer| peer.id) + ); + } } } diff --git a/src/meta-srv/src/procedure/drain_datanode/check_drained.rs b/src/meta-srv/src/procedure/drain_datanode/check_drained.rs new file mode 100644 index 00000000000..28db9bcab49 --- /dev/null +++ b/src/meta-srv/src/procedure/drain_datanode/check_drained.rs @@ -0,0 +1,233 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::any::Any; + +use common_procedure::{Context as ProcedureContext, Status}; +use common_telemetry::info; +use serde::{Deserialize, Serialize}; + +use crate::error::Result; +use crate::procedure::drain_datanode::dispatch::Dispatch; +use crate::procedure::drain_datanode::drain_datanode_end::DrainDatanodeEnd; +use crate::procedure::drain_datanode::{Context, State, utils}; + +/// Rescans the datanode for regions that still need a migration procedure. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub(crate) struct CheckDrained; + +#[async_trait::async_trait] +#[typetag::serde] +impl State for CheckDrained { + async fn next( + &mut self, + ctx: &mut Context, + _procedure_ctx: &ProcedureContext, + ) -> Result<(Box, Status)> { + utils::ensure_drain_deadline(ctx)?; + + let node_id = ctx.persistent_ctx.node_id; + // Reconcile against the table routes after settling all batches of this round. + ctx.rebuild_plan().await?; + // Regions given up by `DrainDatanodeStrategy::TryAll` are still led by the datanode; + // they must not trigger another rescan round. + // Migrations registered by recovery or another caller own their lifecycle. They + // need neither a new child nor a restored parent link, even if still in progress. + let tracker = ctx.region_migration_manager.tracker(); + let volatile_ctx = &mut ctx.volatile_ctx; + if let Some(plan) = volatile_ctx.plan.as_mut() { + plan.leader_regions.retain(|region_id| { + !volatile_ctx.failed_regions.contains(region_id) && !tracker.contains(*region_id) + }); + } + + let migrated = ctx.volatile_ctx.migrated; + let rescan_round = ctx.volatile_ctx.rescan_round; + let remaining_leader_regions = ctx.volatile_ctx.leader_regions().len(); + if remaining_leader_regions == 0 { + let follower_regions = ctx.volatile_ctx.follower_regions().len(); + info!( + "No more migrations to dispatch for datanode {}, migrated regions: {}, follower regions: {}", + node_id, migrated, follower_regions + ); + return Ok((Box::new(DrainDatanodeEnd), Status::executing(false))); + } + + info!( + "Datanode {} still leads {} regions, rescan round: {}", + node_id, remaining_leader_regions, rescan_round + ); + + Ok((Box::new(Dispatch), Status::executing(false))) + } + + fn as_any(&self) -> &dyn Any { + self + } +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use common_meta::peer::Peer; + use common_procedure::Status; + use store_api::storage::RegionId; + + use crate::error; + use crate::procedure::drain_datanode::check_drained::CheckDrained; + use crate::procedure::drain_datanode::dispatch::Dispatch; + use crate::procedure::drain_datanode::drain_datanode_end::DrainDatanodeEnd; + use crate::procedure::drain_datanode::test_util::{ + TestEnv, leader_region, leader_regions, new_persistent_context, + }; + use crate::procedure::drain_datanode::{DrainDatanodeStrategy, MAX_RESCAN_ROUNDS, State}; + use crate::procedure::region_migration::manager::{ + RegionMigrationProcedureTask, RegionMigrationTriggerReason, + }; + use crate::procedure::region_migration::test_util::TestingEnv; + + #[tokio::test] + async fn test_drain_datanode_check_drained_excludes_running_migrations() { + for has_pending_region in [false, true] { + let env = TestEnv::new().await; + let running = RegionId::new(1024, 1); + let pending = RegionId::new(1024, 2); + let failed = RegionId::new(1024, 3); + let mut routes = vec![leader_region(running, 1), leader_region(failed, 1)]; + if has_pending_region { + routes.push(leader_region(pending, 1)); + } + env.create_table(routes).await; + let guard = env + .region_migration_manager + .tracker() + .insert_running_procedure(&RegionMigrationProcedureTask::new( + running, + Peer::empty(1), + Peer::empty(2), + Duration::from_secs(10), + RegionMigrationTriggerReason::Manual, + )) + .unwrap(); + let mut ctx = env.context(new_persistent_context( + Some(2), + DrainDatanodeStrategy::TryAll, + )); + ctx.volatile_ctx.failed_regions.push(failed); + let procedure_ctx = TestingEnv::procedure_context(); + + let (next, status) = CheckDrained.next(&mut ctx, &procedure_ctx).await.unwrap(); + assert!(matches!(status, Status::Executing { persist: false, .. })); + assert_eq!(vec![failed], ctx.volatile_ctx.failed_regions); + assert_eq!(0, ctx.volatile_ctx.migrated); + if has_pending_region { + assert!(next.as_any().is::()); + assert_eq!(vec![pending], leader_regions(&ctx)); + } else { + assert!(next.as_any().is::()); + assert!(leader_regions(&ctx).is_empty()); + } + + // Once the migration releases its guard, an unchanged route needs a new child. + drop(guard); + let (next, _) = CheckDrained.next(&mut ctx, &procedure_ctx).await.unwrap(); + assert!(next.as_any().is::()); + let expected = if has_pending_region { + vec![running, pending] + } else { + vec![running] + }; + assert_eq!(expected, leader_regions(&ctx)); + } + } + + #[tokio::test] + async fn test_drain_datanode_check_drained_transitions_to_end() { + let env = TestEnv::new().await; + + let mut ctx = env.context(new_persistent_context( + Some(2), + DrainDatanodeStrategy::FailFast, + )); + let procedure_ctx = TestingEnv::procedure_context(); + + let (next, status) = CheckDrained.next(&mut ctx, &procedure_ctx).await.unwrap(); + + assert!(next.as_any().downcast_ref::().is_some()); + assert!(matches!(status, Status::Executing { persist: false, .. })); + } + + #[tokio::test] + async fn test_drain_datanode_check_drained_rescan_requeues() { + let env = TestEnv::new().await; + let region_id = RegionId::new(1024, 1); + env.create_table(vec![leader_region(region_id, 1)]).await; + + // The plan of the procedure is stale: a region still on the datanode is not listed. + let mut ctx = env.context(new_persistent_context( + None, + DrainDatanodeStrategy::FailFast, + )); + let procedure_ctx = TestingEnv::procedure_context(); + + let (next, status) = CheckDrained.next(&mut ctx, &procedure_ctx).await.unwrap(); + + assert!(matches!(status, Status::Executing { .. })); + assert!(next.as_any().downcast_ref::().is_some()); + assert_eq!(vec![region_id], leader_regions(&ctx)); + assert_eq!(1, ctx.volatile_ctx.rescan_round); + } + + #[tokio::test] + async fn test_drain_datanode_check_drained_exceeded_rounds() { + let env = TestEnv::new().await; + let region_id = RegionId::new(1024, 1); + env.create_table(vec![leader_region(region_id, 1)]).await; + + let mut ctx = env.context(new_persistent_context( + None, + DrainDatanodeStrategy::FailFast, + )); + ctx.volatile_ctx.rescan_round = MAX_RESCAN_ROUNDS; + let procedure_ctx = TestingEnv::procedure_context(); + + let error = CheckDrained + .next(&mut ctx, &procedure_ctx) + .await + .unwrap_err(); + + assert!(matches!( + error, + error::Error::DrainDatanodeRescanExceeded { .. } + )); + assert!(!error.is_retryable()); + } + + #[tokio::test] + async fn test_drain_datanode_check_drained_ignores_failed_regions() { + let env = TestEnv::new().await; + let region_id = RegionId::new(1024, 1); + env.create_table(vec![leader_region(region_id, 1)]).await; + + let mut ctx = env.context(new_persistent_context(None, DrainDatanodeStrategy::TryAll)); + ctx.volatile_ctx.failed_regions = vec![region_id]; + let procedure_ctx = TestingEnv::procedure_context(); + + let (next, status) = CheckDrained.next(&mut ctx, &procedure_ctx).await.unwrap(); + + assert!(matches!(status, Status::Executing { .. })); + assert!(next.as_any().downcast_ref::().is_some()); + } +} diff --git a/src/meta-srv/src/procedure/drain_datanode/dispatch.rs b/src/meta-srv/src/procedure/drain_datanode/dispatch.rs new file mode 100644 index 00000000000..e5f89a6ffb1 --- /dev/null +++ b/src/meta-srv/src/procedure/drain_datanode/dispatch.rs @@ -0,0 +1,629 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::any::Any; + +use common_procedure::{Context as ProcedureContext, Status, SubprocedureRecovery}; +use common_telemetry::{error, info, warn}; +use serde::{Deserialize, Serialize}; +use store_api::storage::RegionId; + +use crate::error::{self, Result}; +use crate::metrics::{METRIC_META_REGION_MIGRATION_DATANODES, METRIC_META_REGION_MIGRATION_FAIL}; +use crate::procedure::drain_datanode::check_drained::CheckDrained; +use crate::procedure::drain_datanode::utils; +use crate::procedure::drain_datanode::{Context, DrainDatanodeStrategy, State, SubprocedureMeta}; + +/// Migrates the leader regions off the drained datanode in batches. +/// +/// The state first settles the previous batch, then dispatches the next one. Sub-procedures +/// are tracked in the volatile context, and the suspension is not persisted, so a metasrv +/// restart rebuilds the plan by rescanning the datanode. Recovered migrations run +/// independently and are excluded by the migration tracker; regions that already left +/// the datanode are skipped by the route check in +/// [`Context::build_migration_procedure`]. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub(crate) struct Dispatch; + +#[async_trait::async_trait] +#[typetag::serde] +impl State for Dispatch { + async fn next( + &mut self, + ctx: &mut Context, + procedure_ctx: &ProcedureContext, + ) -> Result<(Box, Status)> { + let node_id = ctx.persistent_ctx.node_id; + + // Collect the previous batch's results before dispatching more migrations. + if !ctx.volatile_ctx.inflight_subprocedures.is_empty() { + ctx.wait_for_inflight_subprocedures(procedure_ctx, ctx.persistent_ctx.strategy) + .await?; + return Ok((Box::new(Dispatch), Status::executing(false))); + } + + utils::ensure_drain_deadline(ctx)?; + + let parallelism = ctx.persistent_ctx.parallelism; + let regions: Vec = ctx + .volatile_ctx + .leader_regions() + .iter() + .take(parallelism) + .copied() + .collect(); + if regions.is_empty() { + return Ok((Box::new(CheckDrained), Status::executing(false))); + } + + let mut children = Vec::with_capacity(regions.len()); + let mut inflight = Vec::with_capacity(regions.len()); + let mut targets = Vec::with_capacity(regions.len()); + for region_id in regions { + match ctx + .build_migration_procedure(procedure_ctx, region_id) + .await + { + Ok(Some((child, target_node_id))) => { + info!( + "Dispatching migration of region {} off datanode {}, procedure id: {}", + region_id, node_id, child.id + ); + inflight.push(SubprocedureMeta { + region_id, + procedure_id: child.id, + }); + targets.push(target_node_id); + children.push(child); + } + Ok(None) => { + info!( + "Skipping region {}, it is no longer led by datanode {}", + region_id, node_id + ); + ctx.forget_leader_region(region_id); + } + Err(error::Error::MigrationRunning { .. }) => { + // Another procedure is migrating the region already, for example a manual + // migration submitted before this batch. Ignore it for this round; the + // rescan adds it back when it is still led by the datanode. + warn!( + "A migration of region {} off datanode {} is already running, \ + skipping the region", + region_id, node_id + ); + ctx.forget_leader_region(region_id); + } + Err( + e @ (error::Error::TableRouteNotFound { .. } + | error::Error::TableInfoNotFound { .. }), + ) => { + warn!( + e; + "Skipping region {} while draining datanode {}: table metadata is missing", + region_id, node_id + ); + ctx.forget_leader_region(region_id); + } + Err(e) => { + // Another migration may finish between target selection and the + // migration manager's route check, releasing its tracker guard. + if matches!( + e, + error::Error::RegionMigrated { .. } + | error::Error::LeaderPeerChanged { .. } + ) && !utils::leads_region(&ctx.table_metadata_manager, node_id, region_id) + .await? + { + info!( + "Skipping region {}, it is no longer led by datanode {}", + region_id, node_id + ); + ctx.forget_leader_region(region_id); + continue; + } + + match ctx.persistent_ctx.strategy { + // Nothing is committed yet, so the batch can be retried as a whole. + DrainDatanodeStrategy::FailFast => return Err(e), + DrainDatanodeStrategy::TryAll => { + METRIC_META_REGION_MIGRATION_FAIL.inc(); + error!( + e; + "Gave up draining region {} off datanode {}", + region_id, + node_id + ); + ctx.forget_leader_region(region_id); + ctx.volatile_ctx.failed_regions.push(region_id); + } + } + } + } + } + ctx.volatile_ctx.inflight_subprocedures = inflight; + + if children.is_empty() { + // The whole batch was handled without a sub-procedure, so keep draining the + // remaining leaders instead of spending a rescan round on them. + return Ok((Box::new(Dispatch), Status::executing(false))); + } + + for target_node_id in targets { + METRIC_META_REGION_MIGRATION_DATANODES + .with_label_values(&["src", &node_id.to_string()]) + .inc(); + METRIC_META_REGION_MIGRATION_DATANODES + .with_label_values(&["desc", &target_node_id.to_string()]) + .inc(); + } + + // Recovered children continue independently. The parent rebuilds its plan and + // consults the migration tracker instead of restoring the batch's child links. + Ok(( + Box::new(Dispatch), + Status::Suspended { + subprocedures: children, + recovery: SubprocedureRecovery::Detach, + persist: false, + }, + )) + } + + fn as_any(&self) -> &dyn Any { + self + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::sync::Arc; + use std::time::Duration; + + use crate::procedure::test_util::{procedure_context_with_receivers, procedure_state_receiver}; + use api::v1::meta::heartbeat_request::NodeWorkloads; + use common_meta::cluster::NodeInfo; + use common_meta::error::Result as MetaResult; + use common_meta::key::table_route::TableRouteKey; + use common_meta::key::{MetadataKey, TableMetadataManager, TableMetadataManagerRef}; + use common_meta::kv_backend::test_util::MockKvBackendBuilder; + use common_meta::peer::{Peer, PeerDiscovery, PeerDiscoveryRef}; + use common_meta::rpc::store::BatchGetResponse; + use common_procedure::{ProcedureId, ProcedureState, Status}; + use store_api::storage::RegionId; + + use crate::error; + use crate::key::NodeSchedulingRole; + use crate::procedure::drain_datanode::check_drained::CheckDrained; + use crate::procedure::drain_datanode::dispatch::Dispatch; + use crate::procedure::drain_datanode::drain_datanode_start::DrainDatanodeStart; + use crate::procedure::drain_datanode::test_util::{ + TestEnv, follower_region, leader_region, leader_regions, new_persistent_context, + run_until_done, seed_plan, + }; + use crate::procedure::drain_datanode::{DrainDatanodeStrategy, State, SubprocedureMeta}; + use crate::procedure::region_migration::manager::{ + RegionMigrationManager, RegionMigrationProcedureTask, RegionMigrationTriggerReason, + }; + use crate::procedure::region_migration::test_util::TestingEnv; + + #[tokio::test] + async fn test_drain_datanode_dispatch_transitions_to_check_drained() { + let env = TestEnv::new().await; + // Only a follower region: there is no leader region to migrate. + env.create_table(vec![follower_region(RegionId::new(1024, 1), 2, 1)]) + .await; + + let mut ctx = env.context(new_persistent_context( + Some(2), + DrainDatanodeStrategy::FailFast, + )); + let procedure_ctx = TestingEnv::procedure_context(); + seed_plan(&mut ctx, vec![]); + let (next, status) = Dispatch.next(&mut ctx, &procedure_ctx).await.unwrap(); + + assert!(next.as_any().downcast_ref::().is_some()); + assert!(matches!(status, Status::Executing { persist: false, .. })); + } + + #[tokio::test] + async fn test_drain_datanode_dispatches_migration() { + let env = TestEnv::new().await; + let region_id = RegionId::new(1024, 1); + env.create_table(vec![leader_region(region_id, 1)]).await; + + let mut ctx = env.context(new_persistent_context( + Some(2), + DrainDatanodeStrategy::FailFast, + )); + let procedure_ctx = TestingEnv::procedure_context(); + + let (mut next, _) = DrainDatanodeStart + .next(&mut ctx, &procedure_ctx) + .await + .unwrap(); + let status = run_until_done(&mut next, &mut ctx, &procedure_ctx) + .await + .unwrap(); + + let Status::Suspended { + subprocedures, + recovery, + persist, + } = status + else { + panic!("expected the dispatch state to suspend the procedure"); + }; + // Recovery resumes children independently and rebuilds the parent's plan. + assert!(!persist); + assert_eq!(common_procedure::SubprocedureRecovery::Detach, recovery); + assert_eq!(1, subprocedures.len()); + assert_eq!(1, ctx.volatile_ctx.inflight_subprocedures.len()); + assert_eq!( + region_id, + ctx.volatile_ctx.inflight_subprocedures[0].region_id + ); + // The region is only removed after its migration is settled. + assert_eq!(vec![region_id], leader_regions(&ctx)); + assert!(next.as_any().downcast_ref::().is_some()); + } + + #[tokio::test] + async fn test_drain_datanode_dispatch_skips_region_moved_away() { + let env = TestEnv::new().await; + let region_id = RegionId::new(1024, 1); + // The plan lists the region, but the table route says another datanode leads it. + env.create_table(vec![leader_region(region_id, 2)]).await; + + let mut ctx = env.context(new_persistent_context( + Some(3), + DrainDatanodeStrategy::FailFast, + )); + seed_plan(&mut ctx, vec![region_id]); + let procedure_ctx = TestingEnv::procedure_context(); + + let (next, status) = Dispatch.next(&mut ctx, &procedure_ctx).await.unwrap(); + + assert!(matches!(status, Status::Executing { .. })); + assert!(next.as_any().downcast_ref::().is_some()); + assert!(leader_regions(&ctx).is_empty()); + assert!(ctx.volatile_ctx.inflight_subprocedures.is_empty()); + } + + /// Changes the route after the drain's first read, during target discovery. + struct ChangeRouteDuringDiscovery { + inner: PeerDiscoveryRef, + metadata: TableMetadataManagerRef, + region_id: RegionId, + leader_id: Option, + } + + #[async_trait::async_trait] + impl PeerDiscovery for ChangeRouteDuringDiscovery { + async fn active_frontends(&self) -> MetaResult> { + self.inner.active_frontends().await + } + + async fn active_datanodes( + &self, + filter: Option fn(&'a NodeWorkloads) -> bool>, + ) -> MetaResult> { + let storage = self.metadata.table_route_manager().table_route_storage(); + let table_id = self.region_id.table_id(); + let current = storage.get_with_raw_bytes(table_id).await?.unwrap(); + if let Some(leader_id) = self.leader_id { + let updated = current.update(vec![leader_region(self.region_id, leader_id)])?; + let (txn, _) = storage.build_update_txn(table_id, ¤t, &updated)?; + assert!(self.metadata.kv_backend().txn(txn).await?.succeeded); + } else { + let table_info = self + .metadata + .table_info_manager() + .get(table_id) + .await? + .unwrap(); + self.metadata + .delete_table_metadata( + table_id, + &table_info.table_name(), + ¤t, + &HashMap::new(), + None, + ) + .await?; + } + self.inner.active_datanodes(filter).await + } + + async fn active_flownodes( + &self, + filter: Option fn(&'a NodeWorkloads) -> bool>, + ) -> MetaResult> { + self.inner.active_flownodes(filter).await + } + } + + #[tokio::test] + async fn test_drop_after_route_read_fail_fast() { + assert_drop_after_route_read_is_skipped(DrainDatanodeStrategy::FailFast).await; + } + + #[tokio::test] + async fn test_drop_after_route_read_try_all() { + assert_drop_after_route_read_is_skipped(DrainDatanodeStrategy::TryAll).await; + } + + async fn assert_drop_after_route_read_is_skipped(strategy: DrainDatanodeStrategy) { + let env = TestEnv::new().await; + let region_id = RegionId::new(1024, 1); + env.create_table(vec![leader_region(region_id, 1)]).await; + let metadata = env.table_metadata_manager.clone(); + // Only intercept the migration manager's reads. Return its route snapshot + // after deleting the table, so its following table-info read finds nothing. + let backend = MockKvBackendBuilder::default() + // This fixture has no node-address metadata to remap the route. + .batch_get_fn(Arc::new(|_| Ok(BatchGetResponse { kvs: vec![] }))) + .range_async_fn(Arc::new(move |request| { + let metadata = metadata.clone(); + Box::pin(async move { + let response = metadata.kv_backend().range(request.clone()).await?; + if request.key == TableRouteKey::new(region_id.table_id()).to_bytes() { + let route = metadata + .table_route_manager() + .table_route_storage() + .get_with_raw_bytes(region_id.table_id()) + .await? + .unwrap(); + let info = metadata + .table_info_manager() + .get(region_id.table_id()) + .await? + .unwrap(); + metadata + .delete_table_metadata( + region_id.table_id(), + &info.table_name(), + &route, + &HashMap::new(), + None, + ) + .await?; + } + Ok(response) + }) + })) + .build() + .unwrap(); + let factory = env + .env + .context_factory_with_metadata(Arc::new(TableMetadataManager::new(Arc::new(backend)))); + let manager = Arc::new(RegionMigrationManager::new( + env.procedure_manager().clone(), + factory, + )); + let mut ctx = env.context(new_persistent_context(Some(2), strategy)); + ctx.region_migration_manager = manager.clone(); + seed_plan(&mut ctx, vec![region_id]); + + let (next, status) = Dispatch + .next(&mut ctx, &TestingEnv::procedure_context()) + .await + .unwrap(); + + assert!(next.as_any().is::()); + assert!(matches!(status, Status::Executing { persist: false, .. })); + assert!(leader_regions(&ctx).is_empty()); + assert!(ctx.volatile_ctx.failed_regions.is_empty()); + assert!(ctx.volatile_ctx.inflight_subprocedures.is_empty()); + assert_eq!(0, ctx.volatile_ctx.migrated); + assert!(!manager.tracker().contains(region_id)); + assert!( + env.table_metadata_manager + .table_info_manager() + .get(region_id.table_id()) + .await + .unwrap() + .is_none() + ); + } + + #[tokio::test] + async fn test_drain_datanode_dispatch_skips_region_moved_or_dropped_during_discovery() { + for strategy in [ + DrainDatanodeStrategy::FailFast, + DrainDatanodeStrategy::TryAll, + ] { + // Target 2 yields RegionMigrated; leader 3 yields LeaderPeerChanged; + // dropping the table yields TableRouteNotFound. + for leader_id in [Some(2), Some(3), None] { + let env = TestEnv::new().await; + let region_id = RegionId::new(1024, 1); + env.create_table(vec![leader_region(region_id, 1)]).await; + let mut ctx = env.context(new_persistent_context(Some(2), strategy)); + ctx.selector_ctx.peer_discovery = Arc::new(ChangeRouteDuringDiscovery { + inner: ctx.selector_ctx.peer_discovery.clone(), + metadata: env.table_metadata_manager.clone(), + region_id, + leader_id, + }); + seed_plan(&mut ctx, vec![region_id]); + + let (next, status) = Dispatch + .next(&mut ctx, &TestingEnv::procedure_context()) + .await + .unwrap(); + + assert!(next.as_any().is::()); + assert!(matches!(status, Status::Executing { persist: false, .. })); + assert!(leader_regions(&ctx).is_empty()); + assert!(ctx.volatile_ctx.failed_regions.is_empty()); + assert!(ctx.volatile_ctx.inflight_subprocedures.is_empty()); + assert_eq!(0, ctx.volatile_ctx.migrated); + assert!(!env.region_migration_manager.tracker().contains(region_id)); + } + } + } + + #[tokio::test] + async fn test_drain_datanode_skips_region_whose_migration_is_running() { + let env = TestEnv::new().await; + let region_id = RegionId::new(1024, 1); + env.create_table(vec![leader_region(region_id, 1)]).await; + + let mut ctx = env.context(new_persistent_context( + Some(2), + DrainDatanodeStrategy::FailFast, + )); + seed_plan(&mut ctx, vec![region_id]); + let procedure_ctx = TestingEnv::procedure_context(); + + // A manual migration of the same region is already running. Building it keeps the + // region in the running set of the migration manager until it is dropped. + let _child = env + .region_migration_manager + .build_subprocedure( + common_procedure::ProcedureContext::default(), + RegionMigrationProcedureTask::new( + region_id, + Peer::empty(1), + Peer::empty(2), + Duration::from_secs(10), + RegionMigrationTriggerReason::Manual, + ), + ) + .await + .expect("the manual migration should be buildable"); + + let (next, status) = Dispatch.next(&mut ctx, &procedure_ctx).await.unwrap(); + + // The region is ignored for this round instead of failing the drain: it leaves the plan + // and is neither migrated nor given up. The rescan re-adds it when the datanode still + // leads it. + assert!(matches!(status, Status::Executing { .. })); + assert!(next.as_any().downcast_ref::().is_some()); + assert!(leader_regions(&ctx).is_empty()); + assert!(ctx.volatile_ctx.failed_regions.is_empty()); + assert!(ctx.volatile_ctx.inflight_subprocedures.is_empty()); + } + + #[tokio::test] + async fn test_drain_datanode_fail_fast_on_unschedulable_target() { + let env = TestEnv::new().await; + let region_id = RegionId::new(1024, 1); + env.create_table(vec![leader_region(region_id, 1)]).await; + env.node_scheduling_manager + .mark_unschedulable(NodeSchedulingRole::Datanode, 3) + .await + .unwrap(); + + let mut ctx = env.context(new_persistent_context( + Some(3), + DrainDatanodeStrategy::FailFast, + )); + let procedure_ctx = TestingEnv::procedure_context(); + + let (mut state, _) = DrainDatanodeStart + .next(&mut ctx, &procedure_ctx) + .await + .unwrap(); + let error = run_until_done(&mut state, &mut ctx, &procedure_ctx) + .await + .unwrap_err(); + + assert!(matches!(error, error::Error::DrainDatanodeAbort { .. })); + assert!(!error.is_retryable()); + assert_eq!(vec![region_id], leader_regions(&ctx)); + } + + #[tokio::test] + async fn test_drain_datanode_waits_for_migrated_batch() { + let env = TestEnv::new().await; + let region_id = RegionId::new(1024, 1); + env.create_table(vec![leader_region(region_id, 1)]).await; + + let mut ctx = env.context(new_persistent_context( + Some(2), + DrainDatanodeStrategy::FailFast, + )); + seed_plan(&mut ctx, vec![region_id]); + ctx.volatile_ctx.inflight_subprocedures = vec![SubprocedureMeta { + region_id, + procedure_id: ProcedureId::random(), + }]; + let procedure_ctx = procedure_context_with_receivers(HashMap::from([( + ctx.volatile_ctx.inflight_subprocedures[0].procedure_id, + procedure_state_receiver(ProcedureState::Done { output: None }), + )])); + + let (mut next, status) = Dispatch.next(&mut ctx, &procedure_ctx).await.unwrap(); + + assert!(matches!(status, Status::Executing { .. })); + assert!(next.as_any().downcast_ref::().is_some()); + assert_eq!(1, ctx.volatile_ctx.migrated); + assert!(leader_regions(&ctx).is_empty()); + assert!(ctx.volatile_ctx.inflight_subprocedures.is_empty()); + // The next step checks the remaining regions. + let (next, _) = next.next(&mut ctx, &procedure_ctx).await.unwrap(); + assert!(next.as_any().downcast_ref::().is_some()); + } + + #[tokio::test] + async fn test_drain_datanode_exceeded_deadline() { + let env = TestEnv::new().await; + let region_id = RegionId::new(1024, 1); + env.create_table(vec![leader_region(region_id, 1)]).await; + + let mut persistent_ctx = new_persistent_context(Some(2), DrainDatanodeStrategy::FailFast); + persistent_ctx.deadline_ms = 0; + let mut ctx = env.context(persistent_ctx); + seed_plan(&mut ctx, vec![region_id]); + let procedure_ctx = TestingEnv::procedure_context(); + + let error = Dispatch.next(&mut ctx, &procedure_ctx).await.unwrap_err(); + + assert!(matches!(error, error::Error::ExceededDeadline { .. })); + assert!(!error.is_retryable()); + assert_eq!(vec![region_id], leader_regions(&ctx)); + } + + #[tokio::test] + async fn test_drain_datanode_deadline_does_not_interrupt_batch() { + let env = TestEnv::new().await; + let region_id = RegionId::new(1024, 1); + env.create_table(vec![leader_region(region_id, 1)]).await; + + let mut persistent_ctx = new_persistent_context(Some(2), DrainDatanodeStrategy::FailFast); + persistent_ctx.deadline_ms = 0; + let mut ctx = env.context(persistent_ctx); + seed_plan(&mut ctx, vec![region_id]); + ctx.volatile_ctx.inflight_subprocedures = vec![SubprocedureMeta { + region_id, + procedure_id: ProcedureId::random(), + }]; + let procedure_ctx = procedure_context_with_receivers(HashMap::from([( + ctx.volatile_ctx.inflight_subprocedures[0].procedure_id, + procedure_state_receiver(ProcedureState::Done { output: None }), + )])); + + // The in-flight batch is settled before the deadline is checked. + let (_, status) = Dispatch.next(&mut ctx, &procedure_ctx).await.unwrap(); + assert!(matches!(status, Status::Executing { .. })); + assert_eq!(1, ctx.volatile_ctx.migrated); + + // The next step observes the exceeded deadline. + let error = Dispatch.next(&mut ctx, &procedure_ctx).await.unwrap_err(); + assert!(matches!(error, error::Error::ExceededDeadline { .. })); + } +} diff --git a/src/meta-srv/src/procedure/drain_datanode/drain_datanode_end.rs b/src/meta-srv/src/procedure/drain_datanode/drain_datanode_end.rs index b2880a1c948..5819699fd72 100644 --- a/src/meta-srv/src/procedure/drain_datanode/drain_datanode_end.rs +++ b/src/meta-srv/src/procedure/drain_datanode/drain_datanode_end.rs @@ -15,12 +15,17 @@ use std::any::Any; use common_procedure::{Context as ProcedureContext, Status}; +use common_telemetry::{error, info}; use serde::{Deserialize, Serialize}; -use crate::error::Result; +use crate::error::{self, Result}; use crate::procedure::drain_datanode::{Context, State}; -/// Finishes the procedure. +/// Finishes the procedure and reports the outcome. +/// +/// The procedure state carries no output: only its status is observable through the +/// procedure API, so a drain that left regions behind ends as `Failed` with the failed +/// regions in the error instead of as `Done` with a report. #[derive(Debug, Clone, Serialize, Deserialize)] pub(crate) struct DrainDatanodeEnd; @@ -29,9 +34,33 @@ pub(crate) struct DrainDatanodeEnd; impl State for DrainDatanodeEnd { async fn next( &mut self, - _ctx: &mut Context, + ctx: &mut Context, _procedure_ctx: &ProcedureContext, ) -> Result<(Box, Status)> { + let node_id = ctx.persistent_ctx.node_id; + let migrated = ctx.volatile_ctx.migrated; + let follower_regions = ctx.volatile_ctx.follower_regions().len(); + let failed_regions = ctx.volatile_ctx.failed_regions.clone(); + + if !failed_regions.is_empty() { + error!( + "Drained datanode {} except {} regions, failed regions: {:?}", + node_id, + failed_regions.len(), + failed_regions + ); + return error::DrainDatanodeAbortSnafu { + node_id, + failed_regions, + } + .fail(); + } + + info!( + "Drained datanode {}, migrated leader regions: {}, follower regions left: {}", + node_id, migrated, follower_regions + ); + Ok((Box::new(DrainDatanodeEnd), Status::done())) } @@ -42,14 +71,24 @@ impl State for DrainDatanodeEnd { #[cfg(test)] mod tests { - use super::*; + use common_procedure::Status; + use store_api::storage::RegionId; + + use crate::error; + use crate::procedure::drain_datanode::drain_datanode_end::DrainDatanodeEnd; use crate::procedure::drain_datanode::test_util::{TestEnv, new_persistent_context}; + use crate::procedure::drain_datanode::{DrainDatanodeStrategy, State}; use crate::procedure::region_migration::test_util::TestingEnv; #[tokio::test] async fn test_drain_datanode_end_transitions_to_done() { let env = TestEnv::new().await; - let mut ctx = env.context(new_persistent_context(Some(2))); + + let mut ctx = env.context(new_persistent_context( + Some(2), + DrainDatanodeStrategy::TryAll, + )); + ctx.volatile_ctx.migrated = 1; let procedure_ctx = TestingEnv::procedure_context(); let (next, status) = DrainDatanodeEnd @@ -62,17 +101,25 @@ mod tests { } #[tokio::test] - async fn test_drain_datanode_end_finishes() { + async fn test_drain_datanode_end_aborts_when_regions_are_left() { let env = TestEnv::new().await; - let mut ctx = env.context(new_persistent_context(Some(2))); + let region_id = RegionId::new(1024, 1); + + let mut ctx = env.context(new_persistent_context( + Some(2), + DrainDatanodeStrategy::TryAll, + )); + ctx.volatile_ctx.failed_regions = vec![region_id]; let procedure_ctx = TestingEnv::procedure_context(); - let (next, status) = DrainDatanodeEnd + let error = DrainDatanodeEnd .next(&mut ctx, &procedure_ctx) .await - .unwrap(); + .unwrap_err(); - assert!(next.as_any().downcast_ref::().is_some()); - assert!(matches!(status, Status::Done { .. })); + let error::Error::DrainDatanodeAbort { failed_regions, .. } = error else { + panic!("expected the drain procedure to abort"); + }; + assert_eq!(vec![region_id], failed_regions); } } diff --git a/src/meta-srv/src/procedure/drain_datanode/drain_datanode_start.rs b/src/meta-srv/src/procedure/drain_datanode/drain_datanode_start.rs index 4cd28c7853f..42ee8f2c437 100644 --- a/src/meta-srv/src/procedure/drain_datanode/drain_datanode_start.rs +++ b/src/meta-srv/src/procedure/drain_datanode/drain_datanode_start.rs @@ -20,10 +20,10 @@ use serde::{Deserialize, Serialize}; use crate::error::Result; use crate::key::NodeSchedulingRole; -use crate::procedure::drain_datanode::drain_datanode_end::DrainDatanodeEnd; +use crate::procedure::drain_datanode::check_drained::CheckDrained; use crate::procedure::drain_datanode::{Context, State}; -/// Fences the datanode and builds the drain plan. +/// Fences the datanode before scanning its regions. #[derive(Debug, Clone, Serialize, Deserialize)] pub(crate) struct DrainDatanodeStart; @@ -44,15 +44,11 @@ impl State for DrainDatanodeStart { info!("Marked datanode {} as unschedulable", node_id); } - ctx.ensure_plan().await?; - info!( - "Planned drain of datanode {}, leader regions: {}, follower regions: {}", - node_id, - ctx.volatile_ctx.leader_regions().len(), - ctx.volatile_ctx.follower_regions().len() - ); - - Ok((Box::new(DrainDatanodeEnd), Status::executing(true))) + // The only persisted status of the procedure: persisting it registers the procedure, + // so that a metasrv restart resumes it. The plan and the progress of a run stay in + // the volatile context, and a resumed procedure rebuilds the plan by rescanning the + // datanode. + Ok((Box::new(CheckDrained), Status::executing(true))) } fn as_any(&self) -> &dyn Any { @@ -62,20 +58,28 @@ impl State for DrainDatanodeStart { #[cfg(test)] mod tests { - use super::*; + use common_procedure::Status; + use store_api::storage::RegionId; + + use crate::key::NodeSchedulingRole; + use crate::procedure::drain_datanode::check_drained::CheckDrained; + use crate::procedure::drain_datanode::drain_datanode_start::DrainDatanodeStart; use crate::procedure::drain_datanode::test_util::{ TestEnv, leader_region, new_persistent_context, }; + use crate::procedure::drain_datanode::{DrainDatanodeStrategy, State}; use crate::procedure::region_migration::test_util::TestingEnv; - use store_api::storage::RegionId; #[tokio::test] - async fn test_drain_datanode_start_transitions_to_end() { + async fn test_drain_datanode_start_transitions_to_check_drained() { let env = TestEnv::new().await; env.create_table(vec![leader_region(RegionId::new(1024, 1), 1)]) .await; - let mut ctx = env.context(new_persistent_context(Some(2))); + let mut ctx = env.context(new_persistent_context( + Some(2), + DrainDatanodeStrategy::FailFast, + )); let procedure_ctx = TestingEnv::procedure_context(); let (next, status) = DrainDatanodeStart @@ -83,15 +87,12 @@ mod tests { .await .unwrap(); - assert!(next.as_any().downcast_ref::().is_some()); + assert!(next.as_any().downcast_ref::().is_some()); // Registering the procedure is the only persisted step of the run. assert!(matches!(status, Status::Executing { persist: true, .. })); - assert_eq!( - vec![RegionId::new(1024, 1)], - ctx.volatile_ctx.leader_regions() - ); + assert!(ctx.volatile_ctx.plan.is_none()); assert!( - !ctx.node_scheduling_manager + !env.node_scheduling_manager .is_schedulable(NodeSchedulingRole::Datanode, 1) .await .unwrap() diff --git a/src/meta-srv/src/procedure/drain_datanode/manager.rs b/src/meta-srv/src/procedure/drain_datanode/manager.rs index f262cae2ed1..7c227610424 100644 --- a/src/meta-srv/src/procedure/drain_datanode/manager.rs +++ b/src/meta-srv/src/procedure/drain_datanode/manager.rs @@ -20,14 +20,15 @@ use common_meta::DatanodeId; use common_meta::key::TableMetadataManagerRef; use common_procedure::{ProcedureContext, ProcedureId, ProcedureManagerRef, ProcedureWithId}; use common_telemetry::info; +use common_time::util::current_time_millis; use snafu::ResultExt; use crate::error::{self, Result}; use crate::metasrv::SelectorContext; use crate::node_scheduling::NodeSchedulingManagerRef; use crate::procedure::drain_datanode::{ - Context as DrainDatanodeContext, DrainDatanodeProcedure, DrainDatanodeProcedureTask, - PersistentContext, + Context as DrainDatanodeContext, DEFAULT_DRAIN_TIMEOUT, DrainDatanodeProcedure, + DrainDatanodeProcedureTask, MAX_PARALLELISM, PersistentContext, }; use crate::procedure::region_migration::manager::RegionMigrationManagerRef; @@ -80,15 +81,16 @@ impl Drop for DrainDatanodeProcedureGuard { /// Manages the [`DrainDatanodeProcedure`]s. /// -/// It registers the procedure loader and submits drain tasks. A datanode is drained by at -/// most one procedure at a time: [`DrainDatanodeManager::submit_procedure`] rejects a second -/// drain through the [`DrainDatanodeProcedureTracker`]. +/// It registers the procedure loader and submits drain tasks. Procedures of the same +/// datanode are serialized by the procedure lock declared in +/// [`DrainDatanodeProcedure::lock_key`]. pub struct DrainDatanodeManager { procedure_manager: ProcedureManagerRef, table_metadata_manager: TableMetadataManagerRef, node_scheduling_manager: NodeSchedulingManagerRef, region_migration_manager: RegionMigrationManagerRef, selector_ctx: SelectorContext, + tracker: DrainDatanodeProcedureTracker, } @@ -107,6 +109,7 @@ impl DrainDatanodeManager { node_scheduling_manager, region_migration_manager, selector_ctx, + tracker: DrainDatanodeProcedureTracker::default(), } } @@ -140,6 +143,19 @@ impl DrainDatanodeManager { } .fail(); } + if task.parallelism == 0 || task.parallelism > MAX_PARALLELISM { + return error::InvalidArgumentsSnafu { + err_msg: format!("the parallelism must be in 1..={}", MAX_PARALLELISM), + } + .fail(); + } + let timeout = task.timeout.unwrap_or(DEFAULT_DRAIN_TIMEOUT); + if timeout.is_zero() { + return error::InvalidArgumentsSnafu { + err_msg: "the timeout must be greater than zero", + } + .fail(); + } let Some(guard) = self.tracker.insert_running_procedure(&task) else { return error::DrainDatanodeRunningSnafu { @@ -151,14 +167,19 @@ impl DrainDatanodeManager { let persistent_ctx = PersistentContext { node_id: task.node_id, target_node_id: task.target_node_id, + parallelism: task.parallelism, + strategy: task.strategy, + // The budget starts when the task is submitted, so neither a slow start nor a + // restart extends it. + deadline_ms: current_time_millis() + timeout.as_millis() as i64, }; let procedure = DrainDatanodeProcedure::new(self.context(persistent_ctx), vec![guard]); let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); let procedure_id = procedure_with_id.id; info!( - "Starting drain datanode procedure {procedure_id} for datanode {}", - task.node_id + "Starting drain datanode procedure {procedure_id} for datanode {}, timeout: {:?}", + task.node_id, timeout ); self.procedure_manager .submit(procedure_with_id) @@ -204,26 +225,71 @@ impl DrainDatanodeManager { #[cfg(test)] mod tests { - use super::*; - use crate::procedure::drain_datanode::test_util::{TestEnv, new_persistent_context}; + use std::time::Duration; + use common_procedure::Procedure; + use crate::error; + use crate::procedure::drain_datanode::manager::DrainDatanodeProcedureTracker; + use crate::procedure::drain_datanode::test_util::{TestEnv, new_persistent_context}; + use crate::procedure::drain_datanode::{ + DEFAULT_PARALLELISM, DrainDatanodeProcedure, DrainDatanodeProcedureTask, + DrainDatanodeStrategy, + }; + #[tokio::test] - async fn test_submit_procedure_rejects_same_target() { + async fn test_submit_procedure_rejects_invalid_task() { let env = TestEnv::new().await; let manager = env.drain_manager(); + // The pinned target must differ from the drained datanode. let error = manager .submit_procedure( common_procedure::ProcedureContext::default(), DrainDatanodeProcedureTask { node_id: 1, target_node_id: Some(1), + parallelism: DEFAULT_PARALLELISM, + strategy: DrainDatanodeStrategy::FailFast, + timeout: None, }, ) .await .unwrap_err(); + assert!(matches!(error, error::Error::InvalidArguments { .. })); + // The parallelism must be in range. + for parallelism in [0, crate::procedure::drain_datanode::MAX_PARALLELISM + 1] { + let error = manager + .submit_procedure( + common_procedure::ProcedureContext::default(), + DrainDatanodeProcedureTask { + node_id: 1, + target_node_id: Some(2), + parallelism, + strategy: DrainDatanodeStrategy::FailFast, + timeout: None, + }, + ) + .await + .unwrap_err(); + assert!(matches!(error, error::Error::InvalidArguments { .. })); + } + + // The timeout must be greater than zero. + let error = manager + .submit_procedure( + common_procedure::ProcedureContext::default(), + DrainDatanodeProcedureTask { + node_id: 1, + target_node_id: Some(2), + parallelism: DEFAULT_PARALLELISM, + strategy: DrainDatanodeStrategy::FailFast, + timeout: Some(Duration::ZERO), + }, + ) + .await + .unwrap_err(); assert!(matches!(error, error::Error::InvalidArguments { .. })); } @@ -234,6 +300,9 @@ mod tests { let task = DrainDatanodeProcedureTask { node_id: 1, target_node_id: Some(2), + parallelism: DEFAULT_PARALLELISM, + strategy: DrainDatanodeStrategy::FailFast, + timeout: None, }; // A drain of datanode 1 is running, for example one recovered after a restart. @@ -249,10 +318,13 @@ mod tests { } #[tokio::test] - async fn test_drain_datanode_tracker_takes_the_place_back_on_reload() { + async fn test_drain_datanode_takes_the_place_back_on_reload() { let env = TestEnv::new().await; - let procedure = - DrainDatanodeProcedure::new(env.context(new_persistent_context(Some(2))), vec![]); + let context = env.context(new_persistent_context( + Some(2), + DrainDatanodeStrategy::FailFast, + )); + let procedure = DrainDatanodeProcedure::new(context, vec![]); let json = procedure.dump().unwrap(); let tracker = DrainDatanodeProcedureTracker::default(); @@ -264,6 +336,9 @@ mod tests { let task = DrainDatanodeProcedureTask { node_id: 1, target_node_id: Some(2), + parallelism: DEFAULT_PARALLELISM, + strategy: DrainDatanodeStrategy::FailFast, + timeout: None, }; assert!(tracker.insert_running_procedure(&task).is_none()); diff --git a/src/meta-srv/src/procedure/drain_datanode/migration.rs b/src/meta-srv/src/procedure/drain_datanode/migration.rs index 1767f7494ad..594106fc2e3 100644 --- a/src/meta-srv/src/procedure/drain_datanode/migration.rs +++ b/src/meta-srv/src/procedure/drain_datanode/migration.rs @@ -188,7 +188,10 @@ mod tests { let region_id = RegionId::new(1024, 1); env.create_table(vec![leader_region(region_id, leader_id)]) .await; - let ctx = env.context(new_persistent_context(Some(2))); + let ctx = env.context(new_persistent_context( + Some(2), + DrainDatanodeStrategy::FailFast, + )); let procedure_ctx = crate::procedure::region_migration::test_util::TestingEnv::procedure_context(); let child = ctx @@ -219,7 +222,10 @@ mod tests { let region_id = RegionId::new(1024, 1); env.create_table(vec![leader_region(region_id, 1)]).await; - let mut ctx = env.context(new_persistent_context(Some(2))); + let mut ctx = env.context(new_persistent_context( + Some(2), + DrainDatanodeStrategy::FailFast, + )); seed_plan(&mut ctx, vec![region_id]); ctx.volatile_ctx.inflight_subprocedures = vec![SubprocedureMeta { region_id, @@ -245,7 +251,10 @@ mod tests { let region_id = RegionId::new(1024, 1); env.create_table(vec![leader_region(region_id, 1)]).await; - let mut ctx = env.context(new_persistent_context(Some(2))); + let mut ctx = env.context(new_persistent_context( + Some(2), + DrainDatanodeStrategy::FailFast, + )); seed_plan(&mut ctx, vec![region_id]); ctx.volatile_ctx.inflight_subprocedures = vec![SubprocedureMeta { region_id, @@ -276,7 +285,10 @@ mod tests { let region_id = RegionId::new(1024, 1); env.create_table(vec![leader_region(region_id, 1)]).await; - let mut ctx = env.context(new_persistent_context(Some(2))); + let mut ctx = env.context(new_persistent_context( + Some(2), + DrainDatanodeStrategy::FailFast, + )); seed_plan(&mut ctx, vec![region_id]); ctx.volatile_ctx.inflight_subprocedures = vec![SubprocedureMeta { region_id, @@ -306,7 +318,10 @@ mod tests { // raced with the migration leaves behind. env.create_table(vec![leader_region(region_id, 2)]).await; - let mut ctx = env.context(new_persistent_context(Some(2))); + let mut ctx = env.context(new_persistent_context( + Some(2), + DrainDatanodeStrategy::FailFast, + )); seed_plan(&mut ctx, vec![region_id]); ctx.volatile_ctx.inflight_subprocedures = vec![SubprocedureMeta { region_id, @@ -334,7 +349,10 @@ mod tests { let region_id = RegionId::new(1024, 1); env.create_table(vec![leader_region(region_id, 1)]).await; - let mut ctx = env.context(new_persistent_context(Some(2))); + let mut ctx = env.context(new_persistent_context( + Some(2), + DrainDatanodeStrategy::FailFast, + )); seed_plan(&mut ctx, vec![region_id]); ctx.volatile_ctx.inflight_subprocedures = vec![SubprocedureMeta { region_id, @@ -367,7 +385,10 @@ mod tests { leader_region(succeeded_region, 1), ]) .await; - let mut ctx = env.context(new_persistent_context(Some(2))); + let mut ctx = env.context(new_persistent_context( + Some(2), + DrainDatanodeStrategy::FailFast, + )); seed_plan(&mut ctx, vec![failed_region, succeeded_region]); let failed_id = ProcedureId::random(); let pending_id = ProcedureId::random(); @@ -427,7 +448,10 @@ mod tests { leader_region(succeeded_region, 1), ]) .await; - let mut ctx = env.context(new_persistent_context(Some(2))); + let mut ctx = env.context(new_persistent_context( + Some(2), + DrainDatanodeStrategy::FailFast, + )); seed_plan(&mut ctx, vec![failed_region, succeeded_region]); let failed_id = ProcedureId::random(); let succeeded_id = ProcedureId::random(); diff --git a/src/meta-srv/src/procedure/drain_datanode/test_util.rs b/src/meta-srv/src/procedure/drain_datanode/test_util.rs index aebdb16e0f6..bcbaca4c957 100644 --- a/src/meta-srv/src/procedure/drain_datanode/test_util.rs +++ b/src/meta-srv/src/procedure/drain_datanode/test_util.rs @@ -12,34 +12,53 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::metasrv::SelectorContext; -use crate::procedure::drain_datanode::DrainDatanodePlan; -use crate::procedure::drain_datanode::manager::DrainDatanodeManager; -use crate::procedure::region_migration::manager::{ - RegionMigrationManager, RegionMigrationManagerRef, -}; -use common_meta::affinity::NoopAffinity; -use common_procedure::ProcedureManagerRef; use std::sync::Arc; +use std::time::Duration; +use api::v1::meta::MailboxMessage; +use api::v1::meta::mailbox_message::Payload; use common_meta::DatanodeId; +use common_meta::affinity::NoopAffinity; +use common_meta::instruction::{ + DowngradeRegionReply, DowngradeRegionsReply, FlushRegionReply, Instruction, InstructionReply, + SimpleReply, UpgradeRegionReply, UpgradeRegionsReply, +}; use common_meta::key::TableMetadataManagerRef; use common_meta::key::test_utils::new_test_table_info; use common_meta::peer::Peer; use common_meta::rpc::router::{Region, RegionRoute}; -use common_procedure::{Context as ProcedureContext, Status}; +use common_procedure::{ + Context as ExecContext, ProcedureId, ProcedureManagerRef, ProcedureState, Status, +}; +use common_time::util::current_time_millis; use store_api::storage::RegionId; use crate::error::Result; +use crate::metasrv::SelectorContext; use crate::node_scheduling::NodeSchedulingManagerRef; -use crate::procedure::drain_datanode::{Context, PersistentContext, State}; +use crate::procedure::drain_datanode::manager::DrainDatanodeManager; +use crate::procedure::drain_datanode::{ + Context, DEFAULT_PARALLELISM, DrainDatanodePlan, DrainDatanodeStrategy, PersistentContext, + State, +}; +use crate::procedure::region_migration::manager::{ + RegionMigrationManager, RegionMigrationManagerRef, +}; use crate::procedure::region_migration::test_util::TestingEnv; +use crate::procedure::test_util::{MockHeartbeatReceiver, new_open_region_reply}; +use crate::service::mailbox::{Channel, MailboxRef}; use crate::test_util::{create_meta_peer_client, new_node_scheduling_discovery, put_datanodes}; -pub(super) fn new_persistent_context(target_node_id: Option) -> PersistentContext { +pub(super) fn new_persistent_context( + target_node_id: Option, + strategy: DrainDatanodeStrategy, +) -> PersistentContext { PersistentContext { node_id: 1, target_node_id, + parallelism: DEFAULT_PARALLELISM, + strategy, + deadline_ms: current_time_millis() + 60 * 1000, } } @@ -152,14 +171,127 @@ impl TestEnv { self.selector_ctx.clone(), ) } + + pub(super) async fn mock_datanodes(&mut self) { + let mailbox_ctx = self.env.mailbox_context(); + for node_id in [1, 2] { + let (tx, rx) = tokio::sync::mpsc::channel(8); + mailbox_ctx + .insert_heartbeat_response_receiver(Channel::Datanode(node_id), tx) + .await; + spawn_mock_datanode(mailbox_ctx.mailbox().clone(), rx); + } + } +} + +/// Spawns a mock datanode that answers every mailbox message of a migration. +pub(super) fn spawn_mock_datanode(mailbox: MailboxRef, mut rx: MockHeartbeatReceiver) { + common_runtime::spawn_global(async move { + while let Some(Ok(response)) = rx.recv().await { + let message = response.mailbox_message.unwrap(); + if let Some(reply) = mock_instruction_reply(&message) { + mailbox.on_recv(message.id, Ok(reply)).await.unwrap(); + } + } + }); +} + +/// Returns the reply of `message`, or `None` if the instruction needs no reply. +fn mock_instruction_reply(message: &MailboxMessage) -> Option { + let Some(Payload::Json(payload)) = message.payload.as_ref() else { + panic!("unexpected mailbox payload of message {}", message.id); + }; + let instruction: Instruction = serde_json::from_str(payload).unwrap(); + let reply = match instruction { + Instruction::OpenRegions(_) => { + return Some(new_open_region_reply(message.id, true, None)); + } + Instruction::CloseRegions(_) => InstructionReply::CloseRegions(SimpleReply { + result: true, + error: None, + }), + Instruction::FlushRegions(flush) => { + InstructionReply::FlushRegions(FlushRegionReply::from_results( + flush + .region_ids + .into_iter() + .map(|region_id| (region_id, Ok(()))) + .collect(), + )) + } + Instruction::DowngradeRegions(regions) => { + InstructionReply::DowngradeRegions(DowngradeRegionsReply::new( + regions + .into_iter() + .map(|region| DowngradeRegionReply { + region_id: region.region_id, + last_entry_id: Some(1), + metadata_last_entry_id: Some(1), + exists: true, + error: None, + }) + .collect(), + )) + } + Instruction::UpgradeRegions(regions) => { + InstructionReply::UpgradeRegions(UpgradeRegionsReply::new( + regions + .into_iter() + .map(|region| UpgradeRegionReply { + region_id: region.region_id, + ready: true, + exists: true, + error: None, + }) + .collect(), + )) + } + // One-way instruction: the meta server does not await a reply. + Instruction::InvalidateCaches(_) => return None, + other => panic!("unexpected instruction: {other:?}"), + }; + + Some(MailboxMessage { + id: message.id, + subject: "mock".to_string(), + from: "datanode".to_string(), + to: "meta".to_string(), + timestamp_millis: current_time_millis(), + payload: Some(Payload::Json(serde_json::to_string(&reply).unwrap())), + header: None, + }) +} + +pub(super) async fn wait_for_terminal_state( + procedure_manager: &ProcedureManagerRef, + procedure_id: ProcedureId, +) -> ProcedureState { + for _ in 0..500 { + if let Some(state) = procedure_manager + .procedure_state(procedure_id) + .await + .unwrap() + && matches!( + state, + ProcedureState::Done { .. } + | ProcedureState::Failed { .. } + | ProcedureState::Poisoned { .. } + ) + { + return state; + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + + panic!("the drain procedure did not reach a terminal state in 10 seconds"); } pub(super) async fn run_until_done( state: &mut Box, ctx: &mut Context, - procedure_ctx: &ProcedureContext, + procedure_ctx: &ExecContext, ) -> Result { - for _ in 0..10 { + for _ in 0..50 { let (next, status) = state.next(ctx, procedure_ctx).await?; *state = next; if !matches!(status, Status::Executing { .. }) { @@ -167,10 +299,10 @@ pub(super) async fn run_until_done( } } - panic!("the drain procedure did not finish in 10 steps"); + panic!("the drain procedure did not finish in 50 steps"); } -/// Seeds the volatile plan for migration batch tests. +/// Seeds the volatile plan, as if `CheckDrained` had built it. pub(super) fn seed_plan(ctx: &mut Context, leader_regions: Vec) { ctx.volatile_ctx.plan = Some(DrainDatanodePlan { leader_regions, diff --git a/src/meta-srv/src/procedure/drain_datanode/utils.rs b/src/meta-srv/src/procedure/drain_datanode/utils.rs index 8f5dfa70af2..45e2e6f43ec 100644 --- a/src/meta-srv/src/procedure/drain_datanode/utils.rs +++ b/src/meta-srv/src/procedure/drain_datanode/utils.rs @@ -19,6 +19,7 @@ use common_meta::key::TableMetadataManagerRef; use common_meta::key::table_route::TableRouteValue; use common_meta::peer::Peer; use common_meta::rpc::router::region_distribution; +use common_time::util::current_time_millis; use futures::TryStreamExt; use snafu::ResultExt; use store_api::storage::RegionId; @@ -116,6 +117,26 @@ pub(crate) async fn leads_region( Ok(region_route.leader_peer.as_ref().map(|peer| peer.id) == Some(node_id)) } +/// Returns an error if the drain exceeded its deadline. +pub(crate) fn ensure_drain_deadline(ctx: &Context) -> Result<()> { + if current_time_millis() < ctx.persistent_ctx.deadline_ms { + return Ok(()); + } + + let remaining_leader_regions = ctx + .volatile_ctx + .plan + .as_ref() + .map_or(0, |plan| plan.leader_regions.len()); + error::ExceededDeadlineSnafu { + operation: format!( + "Drain datanode {}, migrated: {}, remaining leader regions: {}", + ctx.persistent_ctx.node_id, ctx.volatile_ctx.migrated, remaining_leader_regions + ), + } + .fail() +} + /// Selects the migration destination of `region_id` off the drained datanode. /// /// Returns the pinned target when the procedure has one, otherwise the schedulable @@ -150,14 +171,16 @@ pub(crate) async fn select_migration_target( #[cfg(test)] mod tests { - use crate::error; - use crate::key::NodeSchedulingRole; - use crate::procedure::drain_datanode::test_util::new_persistent_context; - use crate::procedure::drain_datanode::test_util::{TestEnv, follower_region, leader_region}; - use crate::procedure::drain_datanode::utils; use common_meta::peer::Peer; use store_api::storage::RegionId; + use crate::error; + use crate::key::NodeSchedulingRole; + use crate::procedure::drain_datanode::test_util::{ + TestEnv, follower_region, leader_region, new_persistent_context, + }; + use crate::procedure::drain_datanode::{DrainDatanodeStrategy, utils}; + #[tokio::test] async fn test_build_datanode_plan() { let env = TestEnv::new().await; @@ -190,8 +213,8 @@ mod tests { async fn test_build_datanode_plan_reads_the_table_route() { let env = TestEnv::new().await; let region_id = RegionId::new(1024, 1); - // The datanode is a follower of the region, so it is not a leader even though the - // datanode table index of the datanode contains the table. + // The datanode is listed as a follower of the region, so it is not a leader even + // though the table index of the datanode contains the table. env.create_table(vec![follower_region(region_id, 2, 1)]) .await; @@ -202,12 +225,16 @@ mod tests { assert!(plan.leader_regions.is_empty()); assert_eq!(vec![region_id], plan.follower_regions); } + #[tokio::test] async fn test_select_migration_target() { let env = TestEnv::new().await; let region_id = RegionId::new(1024, 1); - let mut ctx = env.context(new_persistent_context(None)); + let mut ctx = env.context(new_persistent_context( + None, + DrainDatanodeStrategy::FailFast, + )); let peer = utils::select_migration_target(&ctx, region_id, &[]) .await .unwrap(); diff --git a/src/meta-srv/src/procedure/region_migration/test_util.rs b/src/meta-srv/src/procedure/region_migration/test_util.rs index 44180dc6f8c..ecee2eae0c5 100644 --- a/src/meta-srv/src/procedure/region_migration/test_util.rs +++ b/src/meta-srv/src/procedure/region_migration/test_util.rs @@ -137,6 +137,17 @@ impl TestingEnv { } } + /// Returns a context factory using the supplied metadata manager. + pub(crate) fn context_factory_with_metadata( + &self, + table_metadata_manager: TableMetadataManagerRef, + ) -> DefaultContextFactory { + DefaultContextFactory { + table_metadata_manager, + ..self.context_factory() + } + } + /// Returns the mutable [MailboxContext]. pub fn mailbox_context(&mut self) -> &mut MailboxContext { &mut self.mailbox_ctx From ec59131d39ae96d9e15f6b8992b88a50167eac5f Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Fri, 30 Oct 2026 00:23:47 +0000 Subject: [PATCH 144/239] feat(ci): support optional Docker targets for query regression (#10028) * feat(ci): support optional Docker targets for query regression Signed-off-by: WenyXu * fix(ci): enable Docker when provisioning image regression runners Signed-off-by: WenyXu * feat(ci): publish precompiled query regression helpers with pin update PRs Signed-off-by: WenyXu * fix(ci): retry helper image publication on registry timeouts Signed-off-by: WenyXu * fix(ci): publish helper artifacts with hosted Skopeo job Signed-off-by: WenyXu * refactor(ci): update helper image through repository variable Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- .../actions/github-script-tests/action.yml | 1 + .../scripts/build-query-regression-helpers.sh | 43 +++++ .../publish-query-regression-helpers.sh | 28 +++ .github/scripts/query-regression-helpers.py | 82 +++++++++ .github/scripts/query-regression-run.py | 149 +++++++++++++++- .../update-query-regression-helpers-image.sh | 14 ++ .../publish-query-regression-helpers.yml | 166 ++++++++++++++++++ .github/workflows/query-regression.yml | 85 ++++++++- .../workflows/release-dev-builder-images.yaml | 10 ++ docker/query-regression-helpers/Dockerfile | 7 + tests/perf/README.md | 99 ++++++++++- tests/perf/test_query_regression_helpers.py | 155 ++++++++++++++++ ...ery_regression_runner_compaction_toctou.py | 131 +++++++++++++- 13 files changed, 950 insertions(+), 20 deletions(-) create mode 100644 .github/scripts/build-query-regression-helpers.sh create mode 100644 .github/scripts/publish-query-regression-helpers.sh create mode 100644 .github/scripts/query-regression-helpers.py create mode 100644 .github/scripts/update-query-regression-helpers-image.sh create mode 100644 .github/workflows/publish-query-regression-helpers.yml create mode 100644 docker/query-regression-helpers/Dockerfile create mode 100644 tests/perf/test_query_regression_helpers.py diff --git a/.github/actions/github-script-tests/action.yml b/.github/actions/github-script-tests/action.yml index 5721aead096..2d2a310a9ef 100644 --- a/.github/actions/github-script-tests/action.yml +++ b/.github/actions/github-script-tests/action.yml @@ -44,6 +44,7 @@ runs: python3 tests/perf/test_query_regression_runner_otlp_trace_load.py python3 tests/perf/test_query_regression_summary_otlp.py python3 tests/perf/test_query_regression_case_selection.py + python3 tests/perf/test_query_regression_helpers.py python3 tests/perf/test_query_regression_nightly_refs.py python3 tests/perf/test_query_regression_slash.py python3 tests/perf/test_aliyun_ecs_runner_scripts.py diff --git a/.github/scripts/build-query-regression-helpers.sh b/.github/scripts/build-query-regression-helpers.sh new file mode 100644 index 00000000000..6b6cdb6fcde --- /dev/null +++ b/.github/scripts/build-query-regression-helpers.sh @@ -0,0 +1,43 @@ +#!/usr/bin/env bash +set -euo pipefail +: "${CARGO_TARGET_DIR:?CARGO_TARGET_DIR must be configured by the runner}" +revision="$(git rev-parse HEAD)" +tag="query-regression-helpers:${revision}" +staging="$(mktemp -d)" +trap 'rm -rf -- "$staging"' EXIT +cargo build --profile nightly -p cmd --features dev-tools \ + --bin query_perf_fixture --bin query_regression_runner +for binary in query_perf_fixture query_regression_runner; do + "$CARGO_TARGET_DIR/nightly/$binary" --help >/dev/null + cp "$CARGO_TARGET_DIR/nightly/$binary" "$staging/$binary" +done +mkdir -p helpers-image helpers-validation +docker build --provenance=false --build-arg "SOURCE_REVISION=$revision" \ + --file docker/query-regression-helpers/Dockerfile --tag "$tag" "$staging" +container="$(docker create "$tag" /query_perf_fixture --help)" +trap 'docker rm --force --volumes "$container"; rm -rf -- "$staging"' EXIT +mkdir -p "$staging/extracted" +for binary in query_perf_fixture query_regression_runner; do + docker cp "$container:/$binary" "$staging/extracted/$binary" + chmod +x "$staging/extracted/$binary" + "$staging/extracted/$binary" --help >/dev/null + cmp "$staging/$binary" "$staging/extracted/$binary" +done +# Hold the revision and cases constant: compare source-built versus extracted +# fixture normalization twice. Keep the raw outputs as workflow artifacts. +printf '%s\n' "$revision" > helpers-validation/source-revision.txt +for iteration in 1 2; do + while IFS= read -r case; do + name="$(basename "$(dirname "$case")")" + for mode in source image; do + binary="$staging/query_perf_fixture" + [[ "$mode" != image ]] || binary="$staging/extracted/query_perf_fixture" + "$binary" plan --case "$case" > "helpers-validation/${name}-${iteration}-${mode}.json" + done + cmp "helpers-validation/${name}-${iteration}-source.json" \ + "helpers-validation/${name}-${iteration}-image.json" + done < <(find tests/perf/query_cases -name case.toml -type f | sort) +done +(cd "$staging" && sha256sum query_perf_fixture query_regression_runner) > helpers-image/SHA256SUMS +printf '%s\n' "$revision" > helpers-image/source-revision.txt +docker save "$tag" --output helpers-image/image.tar diff --git a/.github/scripts/publish-query-regression-helpers.sh b/.github/scripts/publish-query-regression-helpers.sh new file mode 100644 index 00000000000..b867455f498 --- /dev/null +++ b/.github/scripts/publish-query-regression-helpers.sh @@ -0,0 +1,28 @@ +#!/usr/bin/env bash +set -euo pipefail +repository=greptime-registry.cn-hangzhou.cr.aliyuncs.com/tools/query-regression-helpers +revision="$(cat helpers-image/source-revision.txt)" +[[ "$revision" == "$GITHUB_SHA" ]] +tag="${repository}:${revision}-${GITHUB_RUN_ID}-${GITHUB_RUN_ATTEMPT}" +# Match release-dev-builder-images.yaml: publish ACR images with Skopeo from +# a hosted runner. A failed publish can reuse the completed build artifact. +mkdir -p helpers-validation +docker run --rm \ + -v "${DOCKER_CONFIG:-$HOME/.docker}:/root/.docker:ro" \ + -v "$PWD/helpers-image:/image:ro" \ + -v "$PWD/helpers-validation:/result" \ + -e REGISTRY_AUTH_FILE=/root/.docker/config.json \ + quay.io/skopeo/stable:latest --command-timeout 20m copy --retry-times 3 \ + --digestfile /result/digest.txt \ + docker-archive:/image/image.tar "docker://$tag" +digest="$(cat helpers-validation/digest.txt)" +[[ "$digest" =~ ^sha256:[0-9a-f]{64}$ ]] +image="$repository@$digest" +# The ECS build tested these exact bytes. Verify them again from the registry; +# hosted runners need not share the ECS runtime library versions. +python3 .github/scripts/query-regression-helpers.py "$image" helpers-validation/extracted \ + --checksums helpers-image/SHA256SUMS +jq -e --arg revision "$revision" '.source_revision == $revision' \ + helpers-validation/extracted/helpers-image.json >/dev/null +printf 'image=%s\n' "$image" >> "$GITHUB_OUTPUT" +printf 'Published helpers: `%s` (source `%s`)\n' "$image" "$revision" >> "$GITHUB_STEP_SUMMARY" diff --git a/.github/scripts/query-regression-helpers.py b/.github/scripts/query-regression-helpers.py new file mode 100644 index 00000000000..7f1e22a5c83 --- /dev/null +++ b/.github/scripts/query-regression-helpers.py @@ -0,0 +1,82 @@ +#!/usr/bin/env python3 +# 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. + +"""Extract digest-pinned query regression tools from an artifact image.""" + +import argparse +import hashlib +import json +import os +import re +import subprocess +from pathlib import Path + +BINARIES = ("query_perf_fixture", "query_regression_runner") + + +def resolve_image(value): + if value == "pinned": + value = os.environ.get("QUERY_REGRESSION_HELPERS_IMAGE", "") + if not re.fullmatch(r"[a-zA-Z0-9][a-zA-Z0-9._/:@-]*@sha256:[0-9a-f]{64}", value): + raise ValueError("helpers_image requires an image@sha256 digest; set QUERY_REGRESSION_HELPERS_IMAGE by publishing a verified image before using 'pinned'") + return value + + +def run(*args): + return subprocess.check_output(args, text=True).strip() + + +def extract(image, destination, checksums=None): + expected = None + if checksums is not None: + entries = [line.split() for line in checksums.read_text().splitlines()] + expected = {name: digest for digest, name in entries} + if set(expected) != set(BINARIES) or len(entries) != len(BINARIES): + raise ValueError("Checksums must contain exactly the two helper binaries") + image = resolve_image(image) + destination.mkdir(parents=True, exist_ok=True) + run("docker", "pull", image) + metadata = json.loads(run("docker", "image", "inspect", image))[0] + revision = (metadata.get("Config", {}).get("Labels") or {}).get("org.opencontainers.image.revision", "") + if not re.fullmatch(r"[0-9a-f]{40}", revision): + raise ValueError("Helper image is missing its source revision label") + container = run("docker", "create", image, "/query_perf_fixture", "--help") + try: + for binary in BINARIES: + target = destination / binary + run("docker", "cp", f"{container}:/{binary}", str(target)) + target.chmod(0o755) + if expected is None: + run(str(target.resolve()), "--help") + else: + with target.open("rb") as file: + actual = hashlib.file_digest(file, "sha256").hexdigest() + if actual != expected[binary]: + raise ValueError(f"Published helper checksum mismatch: {binary}") + finally: + run("docker", "rm", "--force", "--volumes", container) + provenance = {"image": image, "image_id": metadata["Id"], "source_revision": revision} + (destination / "helpers-image.json").write_text(json.dumps(provenance, indent=2) + "\n") + print(json.dumps(provenance)) + return provenance + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("image") + parser.add_argument("destination", type=Path) + parser.add_argument("--checksums", type=Path, help="Verify build checksums instead of executing on this host") + args = parser.parse_args() + extract(args.image, args.destination, args.checksums) diff --git a/.github/scripts/query-regression-run.py b/.github/scripts/query-regression-run.py index d488cee1594..e59dcdfb709 100644 --- a/.github/scripts/query-regression-run.py +++ b/.github/scripts/query-regression-run.py @@ -21,10 +21,13 @@ import argparse import json import os import re +import shlex import signal import socket import subprocess +import sys import time +import uuid import urllib.request from dataclasses import dataclass from pathlib import Path @@ -130,6 +133,7 @@ class RunTarget: datanode_http_port: int datanode_data_dir: Path frontend_config: Path | None = None + containerized: bool = False def allocate_ports(n: int) -> list[int]: @@ -151,6 +155,7 @@ def make_target( root: Path, ports: list[int], frontend_config: Path | None = None, + containerized: bool = False, ) -> RunTarget: work_dir = root / name return RunTarget( @@ -167,6 +172,7 @@ def make_target( datanode_http_port=ports[3], datanode_data_dir=work_dir / "datanode-0" / "data", frontend_config=frontend_config, + containerized=containerized, ) @@ -258,7 +264,11 @@ def stop_component( procs: dict[tuple[str, str], subprocess.Popen[bytes]], ) -> None: proc = procs.pop((target.name, name), None) - if proc is None or proc.poll() is not None: + if proc is None: + return + if proc.poll() is not None: + if target.containerized and proc.returncode != 0: + raise RuntimeError(f"Docker target {target.name}/{name} exited with {proc.returncode}") return try: os.killpg(proc.pid, signal.SIGTERM) @@ -272,6 +282,8 @@ def stop_component( except ProcessLookupError: return proc.wait(timeout=20) + if target.containerized and proc.returncode not in (0, 128 + signal.SIGTERM): + raise RuntimeError(f"Docker target {target.name}/{name} cleanup failed: {proc.returncode}") def restart_component( @@ -284,9 +296,15 @@ def restart_component( def stop_all(targets: list[RunTarget], procs: dict[tuple[str, str], subprocess.Popen[bytes]]) -> None: + errors = [] for target in reversed(targets): for name in ("frontend", "datanode", "metasrv"): - stop_component(target, name, procs) + try: + stop_component(target, name, procs) + except Exception as err: + errors.append(err) + if errors: + raise RuntimeError(f"component cleanup failed: {errors}") from errors[0] def load_plan(fixture_generator: Path, case_path: Path) -> dict[str, Any]: @@ -313,8 +331,8 @@ def run_direct_case( ) -> int: ports = allocate_ports(16) targets = [ - make_target("base", base_bin, work_dir, ports[:8]), - make_target("candidate", candidate_bin, work_dir, ports[8:]), + make_target("base", base_bin, work_dir, ports[:8], containerized=bool(getattr(args, "base_image", ""))), + make_target("candidate", candidate_bin, work_dir, ports[8:], containerized=bool(getattr(args, "candidate_image", ""))), ] for target in targets: if target.work_dir.exists() and any(target.work_dir.iterdir()): @@ -407,6 +425,7 @@ def run_remote_case( work_dir, ports[:8], work_dir / "base" / "frontend-prom-store.toml", + containerized=bool(getattr(args, "base_image", "")), ), make_target( "candidate", @@ -414,6 +433,7 @@ def run_remote_case( work_dir, ports[8:], work_dir / "candidate" / "frontend-prom-store.toml", + containerized=bool(getattr(args, "candidate_image", "")), ), ] for target in targets: @@ -495,8 +515,8 @@ def run_otlp_case( raise ValueError("--otelgen-bin is required for otlp_trace_load") ports = allocate_ports(16) targets = [ - make_target("base", base_bin, work_dir, ports[:8]), - make_target("candidate", candidate_bin, work_dir, ports[8:]), + make_target("base", base_bin, work_dir, ports[:8], containerized=bool(getattr(args, "base_image", ""))), + make_target("candidate", candidate_bin, work_dir, ports[8:], containerized=bool(getattr(args, "candidate_image", ""))), ] for target in targets: if target.work_dir.exists() and any(target.work_dir.iterdir()): @@ -634,6 +654,106 @@ def write_failed_report(work_dir: Path, case_path: Path, error: Exception) -> No ) +def docker_target(image: str, root: Path, command: list[str]) -> int: + """Expose an image as a binary, including read-bench subcommands. + + The caller terminates a process group. Keep Docker clients in separate + sessions so cleanup can stop the container before SST inspection starts. + """ + name = f"query-regression-{uuid.uuid4().hex}" + + def interrupted(signum, _frame): + raise SystemExit(128 + signum) + + previous = {sig: signal.signal(sig, interrupted) for sig in (signal.SIGTERM, signal.SIGINT)} + try: + subprocess.run([ + "docker", "create", "--name", name, "--network", "host", + "--user", f"{os.getuid()}:{os.getgid()}", + "--volume", f"{root}:{root}", "--workdir", str(root), + "--entrypoint", "/greptime/bin/greptime", image, *command, + ], check=True, stdout=subprocess.DEVNULL, start_new_session=True) + attached = subprocess.run( + ["docker", "start", "--attach", name], check=False, start_new_session=True, + ) + if attached.returncode != 0: + return attached.returncode + # docker start's status is not the contract for the container's exit code. + result = subprocess.run( + ["docker", "inspect", "--format", "{{.State.ExitCode}}", name], + check=True, capture_output=True, text=True, start_new_session=True, + ) + return int(result.stdout.strip()) + finally: + for sig in previous: + signal.signal(sig, signal.SIG_IGN) + try: + subprocess.run( + ["docker", "stop", "--time", "5", name], check=False, + stdout=subprocess.DEVNULL, timeout=8, start_new_session=True, + ) + except subprocess.TimeoutExpired: + # Force-removal below still waits for the container to stop. + pass + finally: + try: + subprocess.run( + ["docker", "rm", "--force", "--volumes", name], check=True, + stdout=subprocess.DEVNULL, timeout=5, start_new_session=True, + ) + finally: + for sig, handler in previous.items(): + signal.signal(sig, handler) + + +def prepare_targets(args: argparse.Namespace) -> dict[str, Any]: + """Resolve images once and create binary-compatible launchers for all cases.""" + metadata = {} + images = {} + for name in ("base", "candidate"): + image = getattr(args, f"{name}_image") + ref = getattr(args, f"{name}_ref") + if not image: + metadata[name] = {"kind": "source", "ref": os.environ.get(f"VERIFIED_{name.upper()}_SHA") or ref} + continue + if sys.platform != "linux": + raise ValueError("Docker targets require Linux host networking and a local Docker daemon") + if image.startswith("-") or any(c.isspace() for c in image): + raise ValueError(f"invalid {name} image reference: {image!r}") + if image not in images: + result = subprocess.run( + ["docker", "image", "inspect", image], check=False, capture_output=True, text=True, + ) + if result.returncode != 0: + subprocess.run(["docker", "pull", image], check=True) + result = subprocess.run( + ["docker", "image", "inspect", image], check=True, capture_output=True, text=True, + ) + images[image] = json.loads(result.stdout)[0] + info = images[image] + image_id = info["Id"] + if not re.fullmatch(r"sha256:[0-9a-f]{64}", image_id): + raise ValueError(f"invalid Docker image ID: {image_id!r}") + args.work_dir = args.work_dir.resolve() + if ":" in str(args.work_dir): + raise ValueError("Docker target work directory cannot contain ':'") + launchers = args.work_dir / ".targets" + launchers.mkdir(parents=True, exist_ok=True) + launcher = launchers / name + command = [sys.executable, str(Path(__file__).resolve()), + "docker-target", image_id, str(args.work_dir)] + launcher.write_text("#!/bin/sh\nexec " + shlex.join(command) + ' "$@"\n', encoding="utf-8") + launcher.chmod(0o755) + setattr(args, f"{name}_bin", launcher) + setattr(args, f"{name}_ref", f"{image} ({image_id})") + metadata[name] = { + "kind": "docker", "image": image, "image_id": image_id, + "repo_digests": info.get("RepoDigests", []), + "network": "host", "user": f"{os.getuid()}:{os.getgid()}", + } + return metadata + + def main() -> int: parser = argparse.ArgumentParser() parser.add_argument("--cases", action="append", help="'all', 'heavy', or comma/space separated case paths") @@ -663,6 +783,8 @@ def main() -> int: parser.add_argument("--base-ref", default=os.environ.get("BASE_REF", "")) parser.add_argument("--candidate-ref", default=os.environ.get("CANDIDATE_REF", "")) parser.add_argument("--github-output", default=os.environ.get("GITHUB_OUTPUT")) + parser.add_argument("--base-image", default=os.environ.get("BASE_IMAGE", "")) + parser.add_argument("--candidate-image", default=os.environ.get("CANDIDATE_IMAGE", "")) args = parser.parse_args() try: cases = split_cases(args.cases or [os.environ.get("CASE_PATHS", "all")]) @@ -671,6 +793,12 @@ def main() -> int: append_github_output(args.github_output, 1) return 0 + preparation_error = None + metadata = {} + try: + metadata = prepare_targets(args) + except Exception as err: + preparation_error = err reports: list[Path] = [] status = 0 for case in cases: @@ -678,11 +806,18 @@ def main() -> int: work_dir = args.work_dir / case_slug(case_path) reports.append(work_dir / "query-regression-report.json") try: + if preparation_error is not None: + raise preparation_error case_status = run_case(args, case_path, work_dir) except Exception as err: # noqa: BLE001 - preserve an aggregate report for the summary step print(f"error: query regression case {case_path}: {err}", flush=True) write_failed_report(work_dir, case_path, err) case_status = 1 + report_path = reports[-1] + if (args.base_image or args.candidate_image) and report_path.exists(): + report = json.loads(report_path.read_text(encoding="utf-8")) + report["execution_targets"] = metadata + report_path.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8") if case_status != 0: status = case_status or 1 @@ -695,4 +830,6 @@ def main() -> int: if __name__ == "__main__": + if len(sys.argv) > 1 and sys.argv[1] == "docker-target": + raise SystemExit(docker_target(sys.argv[2], Path(sys.argv[3]), sys.argv[4:])) raise SystemExit(main()) diff --git a/.github/scripts/update-query-regression-helpers-image.sh b/.github/scripts/update-query-regression-helpers-image.sh new file mode 100644 index 00000000000..9ef9122da74 --- /dev/null +++ b/.github/scripts/update-query-regression-helpers-image.sh @@ -0,0 +1,14 @@ +#!/usr/bin/env bash +set -euo pipefail +: "${HELPERS_IMAGE:?Verified image digest is required}" +: "${GITHUB_REPOSITORY:?Repository is required}" +repository=greptime-registry.cn-hangzhou.cr.aliyuncs.com/tools/query-regression-helpers +[[ "$HELPERS_IMAGE" == "$repository@sha256:"* ]] +[[ "${HELPERS_IMAGE#"$repository@"}" =~ ^sha256:[0-9a-f]{64}$ ]] +# Like QUERY_REGRESSION_ECS_IMAGE_ID, this is CI configuration, not a source +# change. Create on first publication, then update only after image verification. +gh variable set QUERY_REGRESSION_HELPERS_IMAGE \ + --repo "$GITHUB_REPOSITORY" --body "$HELPERS_IMAGE" +gh api "repos/$GITHUB_REPOSITORY/actions/variables/QUERY_REGRESSION_HELPERS_IMAGE" \ + | jq -e --arg expected "$HELPERS_IMAGE" '.value == $expected' >/dev/null +printf 'QUERY_REGRESSION_HELPERS_IMAGE updated to `%s`\n' "$HELPERS_IMAGE" diff --git a/.github/workflows/publish-query-regression-helpers.yml b/.github/workflows/publish-query-regression-helpers.yml new file mode 100644 index 00000000000..eca58cff96d --- /dev/null +++ b/.github/workflows/publish-query-regression-helpers.yml @@ -0,0 +1,166 @@ +name: Publish query regression helpers + +on: + workflow_dispatch: + workflow_call: + push: + branches: [main] + paths: + - src/cmd/src/bin/query_perf_fixture.rs + - src/cmd/src/bin/query_perf_fixture/** + - src/cmd/src/bin/query_regression_runner.rs + - src/cmd/src/bin/query_regression_runner/** + - docker/query-regression-helpers/** + - .github/workflows/publish-query-regression-helpers.yml + - .github/scripts/publish-query-regression-helpers.sh + - .github/scripts/build-query-regression-helpers.sh + - .github/scripts/update-query-regression-helpers-image.sh + - Cargo.lock + +permissions: + contents: read + +# Serialize publication and updates of the repository image variable. +concurrency: + group: publish-query-regression-helpers + cancel-in-progress: false + +jobs: + provision: + if: ${{ github.repository == 'GreptimeTeam/greptimedb' }} + runs-on: ubuntu-latest + timeout-minutes: 45 + outputs: + label: ${{ steps.provision.outputs.label }} + instance_id: ${{ steps.provision.outputs.instance_id }} + runner_name: ${{ steps.provision.outputs.runner_name }} + steps: + - name: Checkout trusted provisioning scripts + uses: actions/checkout@v4 + with: + repository: ${{ github.repository }} + ref: ${{ github.sha }} + persist-credentials: false + + - name: Provision ECS runner + id: provision + uses: ./.github/actions/aliyun-ecs-create + with: + access-key-id: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_ID }} + access-key-secret: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_SECRET }} + github-token: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + region-id: ${{ vars.ALIYUN_ECS_REGION_ID }} + vswitch-id: ${{ vars.ALIYUN_ECS_VSWITCH_ID }} + security-group-id: ${{ vars.ALIYUN_ECS_SECURITY_GROUP_ID }} + instance-type: ${{ vars.ALIYUN_ECS_INSTANCE_TYPE }} + # Preserve query regression's original resource and janitor budgets. + system-disk-gib: '80' + ttl-hours: '4' + enable-docker: 'true' + resource-group-id: ${{ vars.ALIYUN_ECS_RESOURCE_GROUP_ID }} + image-id: ${{ vars.QUERY_REGRESSION_ECS_IMAGE_ID }} + runner-uid: ${{ vars.QUERY_REGRESSION_RUNNER_UID || '1001' }} + runner-gid: ${{ vars.QUERY_REGRESSION_RUNNER_GID || '1001' }} + + build: + needs: provision + runs-on: ${{ needs.provision.outputs.label }} + timeout-minutes: 150 + env: + CARGO_HOME: /home/runner/.cargo + CARGO_TARGET_DIR: /home/runner/query-regression-target + RUSTUP_HOME: /opt/rustup + RUSTUP_TOOLCHAIN: nightly-2026-03-21 + RUSTUP_AUTO_INSTALL: "0" + CARGO_INCREMENTAL: "0" + RUSTFLAGS: -D warnings -C link-arg=-fuse-ld=mold + steps: + - uses: actions/checkout@v4 + with: + persist-credentials: false + - name: Build and validate helpers + run: bash .github/scripts/build-query-regression-helpers.sh + - name: Upload reusable helper image + uses: actions/upload-artifact@v4 + with: + name: helpers-image + path: helpers-image/ + compression-level: 1 + retention-days: 7 + if-no-files-found: error + - name: Upload build validation evidence + if: always() + uses: actions/upload-artifact@v4 + with: + name: helpers-build-validation + path: helpers-validation/ + if-no-files-found: warn + + publish: + needs: build + runs-on: ubuntu-latest + timeout-minutes: 45 + outputs: + image: ${{ steps.publish.outputs.image }} + steps: + - uses: actions/checkout@v4 + with: + persist-credentials: false + - uses: actions/download-artifact@v4 + with: + name: helpers-image + path: helpers-image/ + - name: Login to registry + uses: docker/login-action@v3 + with: + registry: greptime-registry.cn-hangzhou.cr.aliyuncs.com + username: ${{ secrets.ALICLOUD_USERNAME }} + password: ${{ secrets.ALICLOUD_PASSWORD }} + - name: Publish and verify helpers + id: publish + run: bash .github/scripts/publish-query-regression-helpers.sh + - name: Upload publication validation evidence + if: always() + uses: actions/upload-artifact@v4 + with: + name: helpers-publish-validation + path: | + helpers-validation/digest.txt + helpers-validation/extracted/helpers-image.json + if-no-files-found: warn + + update-variable: + needs: publish + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + with: + persist-credentials: false + - name: Update helpers image repository variable + env: + GH_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + HELPERS_IMAGE: ${{ needs.publish.outputs.image }} + run: bash .github/scripts/update-query-regression-helpers-image.sh + + teardown: + if: ${{ always() && needs.provision.outputs.instance_id != '' }} + needs: [provision, build] + runs-on: ubuntu-latest + timeout-minutes: 15 + steps: + - name: Checkout trusted teardown scripts + uses: actions/checkout@v4 + with: + repository: ${{ github.repository }} + ref: ${{ github.sha }} + persist-credentials: false + + - name: Teardown ECS runner + uses: ./.github/actions/aliyun-ecs-delete + with: + access-key-id: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_ID }} + access-key-secret: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_SECRET }} + github-token: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + region-id: ${{ vars.ALIYUN_ECS_REGION_ID }} + instance-id: ${{ needs.provision.outputs.instance_id }} + runner-name: ${{ needs.provision.outputs.runner_name }} diff --git a/.github/workflows/query-regression.yml b/.github/workflows/query-regression.yml index 6543281409d..d65d6666658 100644 --- a/.github/workflows/query-regression.yml +++ b/.github/workflows/query-regression.yml @@ -26,6 +26,21 @@ on: description: Candidate ref/sha to build; reusable release calls must provide an immutable full SHA required: true type: string + base_image: + description: Base Docker image tag/digest (empty builds base_ref from source) + required: false + type: string + default: "" + candidate_image: + description: Candidate Docker image tag/digest (empty builds candidate_ref from source) + required: false + type: string + default: "" + helpers_image: + description: Helpers image@sha256 or pinned from repository variable (empty builds from source) + required: false + type: string + default: "" cargo_profile: description: Cargo profile to build base and candidate binaries with required: false @@ -80,6 +95,21 @@ on: description: Runner HTTP timeout seconds required: true default: "300" + base_image: + description: Base Docker image tag/digest (empty builds base_ref from source) + required: false + type: string + default: "" + candidate_image: + description: Candidate Docker image tag/digest (empty builds candidate_ref from source) + required: false + type: string + default: "" + helpers_image: + description: Helpers image@sha256 or pinned from repository variable (empty builds from source) + required: false + type: string + default: "" cargo_profile: description: Cargo profile to build base and candidate binaries with required: true @@ -132,6 +162,7 @@ jobs: python3 tests/perf/test_query_regression_runner_otlp_trace_load.py python3 tests/perf/test_query_regression_summary_otlp.py python3 tests/perf/test_query_regression_case_selection.py + python3 tests/perf/test_query_regression_helpers.py python3 tests/perf/test_query_regression_nightly_refs.py python3 tests/perf/test_query_regression_slash.py python3 tests/perf/test_aliyun_ecs_runner_scripts.py @@ -171,6 +202,7 @@ jobs: # Preserve query regression's original resource and janitor budgets. system-disk-gib: '80' ttl-hours: '4' + enable-docker: ${{ inputs.base_image != '' || inputs.candidate_image != '' || inputs.helpers_image != '' }} resource-group-id: ${{ vars.ALIYUN_ECS_RESOURCE_GROUP_ID }} image-id: ${{ vars.QUERY_REGRESSION_ECS_IMAGE_ID }} runner-uid: ${{ vars.QUERY_REGRESSION_RUNNER_UID || '1001' }} @@ -214,6 +246,10 @@ jobs: PR_NUMBER: ${{ inputs.pr_number }} HEAD_REPO: ${{ inputs.head_repo }} BASE_REPO: ${{ inputs.base_repo }} + HELPERS_IMAGE: ${{ inputs.helpers_image }} + QUERY_REGRESSION_HELPERS_IMAGE: ${{ vars.QUERY_REGRESSION_HELPERS_IMAGE }} + BASE_IMAGE: ${{ inputs.base_image }} + CANDIDATE_IMAGE: ${{ inputs.candidate_image }} BASE_REF: ${{ inputs.base_ref }} CANDIDATE_REF: ${{ inputs.candidate_ref || github.ref }} EVENT_MERGE_SHA: ${{ inputs.pr_number != '' && inputs.candidate_ref || '' }} @@ -224,6 +260,16 @@ jobs: EXPECTED_RUNNER_UID: ${{ vars.QUERY_REGRESSION_RUNNER_UID || '1001' }} EXPECTED_RUNNER_GID: ${{ vars.QUERY_REGRESSION_RUNNER_GID || '1001' }} steps: + - name: Validate image target admission + run: | + if [[ "${PR_ADMISSION}" == "true" && ( -n "${BASE_IMAGE}" || -n "${CANDIDATE_IMAGE}" || -n "${HELPERS_IMAGE}" ) ]]; then + echo "PR admission requires source targets for verified commit attribution." >&2 + exit 1 + fi + if [[ -n "${BASE_IMAGE}" || -n "${CANDIDATE_IMAGE}" || -n "${HELPERS_IMAGE}" ]]; then + docker info >/dev/null + fi + - name: Report provisioned ECS runner if: ${{ needs.provision.outputs.instance_id != '' }} shell: bash @@ -418,6 +464,7 @@ jobs: > "${GITHUB_WORKSPACE}/query-regression-trusted-scripts/query-regression-summary.py" - name: Verify runner image tools + if: ${{ inputs.base_image == '' || inputs.candidate_image == '' || inputs.helpers_image == '' }} shell: bash run: | set -euo pipefail @@ -534,6 +581,7 @@ jobs: fi - name: Prepare query regression cache + if: ${{ inputs.base_image == '' || inputs.candidate_image == '' || inputs.helpers_image == '' }} shell: bash working-directory: src run: | @@ -694,6 +742,7 @@ jobs: sccache --show-stats - name: Build base greptime + if: ${{ inputs.base_image == '' }} working-directory: src run: | set -euo pipefail @@ -728,14 +777,27 @@ jobs: git reset --hard "${VERIFIED_CANDIDATE_SHA}" git clean -ffdx + - name: Extract precompiled query regression helpers + if: ${{ inputs.helpers_image != '' }} + working-directory: src + run: | + python3 .github/scripts/query-regression-helpers.py "$HELPERS_IMAGE" \ + "$GITHUB_WORKSPACE/query-regression-bins/candidate" + cat "$GITHUB_WORKSPACE/query-regression-bins/candidate/helpers-image.json" >> "$GITHUB_STEP_SUMMARY" + - name: Build candidate greptime and query regression helpers + if: ${{ inputs.candidate_image == '' || inputs.helpers_image == '' }} working-directory: src run: | set -euo pipefail SECONDS=0 - cargo build --profile "${CARGO_PROFILE}" -p cmd --bin greptime - cargo build --profile "${CARGO_PROFILE}" -p cmd --bin query_perf_fixture --features dev-tools - cargo build --profile "${CARGO_PROFILE}" -p cmd --bin query_regression_runner --features dev-tools + if [[ -z "${CANDIDATE_IMAGE}" ]]; then + cargo build --profile "${CARGO_PROFILE}" -p cmd --bin greptime + fi + if [[ -z "${HELPERS_IMAGE}" ]]; then + cargo build --profile "${CARGO_PROFILE}" -p cmd --bin query_perf_fixture --features dev-tools + cargo build --profile "${CARGO_PROFILE}" -p cmd --bin query_regression_runner --features dev-tools + fi candidate_build_elapsed="${SECONDS}" printf 'Candidate greptime and query regression helper cargo builds elapsed: %s seconds\n' "${candidate_build_elapsed}" printf -- '- Candidate greptime and query regression helper cargo builds: %s seconds\n' "${candidate_build_elapsed}" >> "${GITHUB_STEP_SUMMARY}" @@ -745,12 +807,16 @@ jobs: target_dir="debug" fi mkdir -p "${GITHUB_WORKSPACE}/query-regression-bins/candidate" - cp "${CARGO_TARGET_DIR}/${target_dir}/greptime" \ - "${GITHUB_WORKSPACE}/query-regression-bins/candidate/greptime" - cp "${CARGO_TARGET_DIR}/${target_dir}/query_perf_fixture" \ - "${GITHUB_WORKSPACE}/query-regression-bins/candidate/query_perf_fixture" - cp "${CARGO_TARGET_DIR}/${target_dir}/query_regression_runner" \ - "${GITHUB_WORKSPACE}/query-regression-bins/candidate/query_regression_runner" + if [[ -z "${CANDIDATE_IMAGE}" ]]; then + cp "${CARGO_TARGET_DIR}/${target_dir}/greptime" \ + "${GITHUB_WORKSPACE}/query-regression-bins/candidate/greptime" + fi + if [[ -z "${HELPERS_IMAGE}" ]]; then + cp "${CARGO_TARGET_DIR}/${target_dir}/query_perf_fixture" \ + "${GITHUB_WORKSPACE}/query-regression-bins/candidate/query_perf_fixture" + cp "${CARGO_TARGET_DIR}/${target_dir}/query_regression_runner" \ + "${GITHUB_WORKSPACE}/query-regression-bins/candidate/query_regression_runner" + fi - name: Run query regression id: run @@ -792,6 +858,7 @@ jobs: query-regression-work/**/logs/** query-regression-work/**/otelgen/** query-regression-summary.md + query-regression-bins/candidate/helpers-image.json machine-telemetry.log if-no-files-found: warn retention-days: 7 diff --git a/.github/workflows/release-dev-builder-images.yaml b/.github/workflows/release-dev-builder-images.yaml index 2e1041c1ff4..9f74035f8c9 100644 --- a/.github/workflows/release-dev-builder-images.yaml +++ b/.github/workflows/release-dev-builder-images.yaml @@ -35,6 +35,11 @@ on: description: Rebuild the query-regression ECS runner image (and update QUERY_REGRESSION_ECS_IMAGE_ID + RUNNER_IMAGE_EPOCH) required: false default: false + release_query_regression_helpers_image: + type: boolean + description: Publish precompiled query regression helpers and update the image repository variable + required: false + default: false update_dev_builder_image_tag: type: boolean description: Update the DEV_BUILDER_IMAGE_TAG in Makefile and create a PR @@ -451,3 +456,8 @@ jobs: echo "## Lockstep updates" >> "$GITHUB_STEP_SUMMARY" echo "- QUERY_REGRESSION_ECS_IMAGE_ID → \`${image_id}\`" >> "$GITHUB_STEP_SUMMARY" echo "- RUNNER_IMAGE_EPOCH → \`${new_epoch}\` (PR on \`${BRANCH}\`, pending review/merge)" >> "$GITHUB_STEP_SUMMARY" + + release-query-regression-helpers: + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && inputs.release_query_regression_helpers_image }} + uses: ./.github/workflows/publish-query-regression-helpers.yml + secrets: inherit diff --git a/docker/query-regression-helpers/Dockerfile b/docker/query-regression-helpers/Dockerfile new file mode 100644 index 00000000000..d1a93171695 --- /dev/null +++ b/docker/query-regression-helpers/Dockerfile @@ -0,0 +1,7 @@ +# Artifact image: binaries run on the query-regression ECS host after extraction. +# Build on the same ECS image as the consumer to preserve the runtime ABI. +FROM scratch +ARG SOURCE_REVISION +LABEL org.opencontainers.image.source="https://github.com/GreptimeTeam/greptimedb" +LABEL org.opencontainers.image.revision=$SOURCE_REVISION +COPY query_perf_fixture query_regression_runner / diff --git a/tests/perf/README.md b/tests/perf/README.md index 0361f598a3b..a7e698fe047 100644 --- a/tests/perf/README.md +++ b/tests/perf/README.md @@ -495,18 +495,68 @@ query_regression_runner measure --case --fixture-generator ## GitHub Actions `.github/workflows/query-regression.yml` provides an opt-in CI entrypoint for -query regression runs. It builds its own binaries for now: +query regression runs. By default, it builds its own binaries: - base `greptime` from the PR base commit, or `workflow_dispatch` `base_ref` - candidate `greptime`, `query_perf_fixture`, and `query_regression_runner` from the PR merge ref/current candidate checkout - outer driver and summary formatter from the candidate checkout -The workflow builds base and candidate `greptime` as normal release-equivalent -binaries. Candidate `query_perf_fixture` and `query_regression_runner` are the +For source targets, the workflow builds `greptime` as a normal release-equivalent +binary. Candidate `query_perf_fixture` and `query_regression_runner` are the extra head-side helpers; `finalize-remote` uses candidate `greptime datanode parquetbench/scanbench` as the read-bench tool against each target's data directory. +### Docker image targets + +`base_image` and `candidate_image` are independent optional inputs to manual +and reusable query-regression workflows. Leave an input empty to build that +side from source, or set it to a GreptimeDB image tag/digest to skip its server +build. All four source/image combinations are supported. PR slash-command runs +continue to use verified source revisions; image overrides are rejected for PR +admission so reports cannot attribute an image to a different commit. + +The candidate checkout still supplies cases and builds `query_perf_fixture` +and `query_regression_runner`, even with two image targets. `base_ref` and +`candidate_ref` remain required for reusable calls. Remote-write read-bench +commands use the candidate target, including its image when configured. + +The local driver accepts `--base-image` / `--candidate-image` (or `BASE_IMAGE` / +`CANDIDATE_IMAGE`). An image overrides that side's `--*-bin`. For example: + +```bash +python3 .github/scripts/query-regression-run.py \ + --candidate-src . \ + --base-image greptime/greptimedb:v1.3.0 \ + --candidate-bin /absolute/path/to/candidate/greptime \ + --fixture-generator /absolute/path/to/query_perf_fixture \ + --runner /absolute/path/to/query_regression_runner \ + --cases tests/perf/query_cases/smoke_direct_sst/case.toml \ + --summary-script .github/scripts/query-regression-summary.py +``` + +Image targets require Linux and a local Docker daemon accessible to the runner +user; the workflow checks access before building. For image targets, ECS +provisioning enables the existing `enable-docker` option: root cloud-init starts +the image-provided Docker daemon, grants the runner Docker group access, and +verifies access before starting the runner. Source-only runs keep this option +disabled. No ECS image rebuild is needed. Manually managed hosts must provision +Docker access and restart the runner service before selecting image targets. +Images must expose the GreptimeDB binary at `/greptime/bin/greptime` and support +the commands and data formats used by the selected cases. Images are pulled +when absent locally, then pinned to their inspected image ID for the entire +run; use a registry digest for repeatability across hosts. Reports include +requested images, resolved IDs and repository digests alongside source commit +identities. + +Containers use host networking and the runner's UID/GID. The run work directory +is bind-mounted at the same absolute path, preserving generated configuration, +SST materialization, logs and post-stop storage inspection. No additional CPU +or memory limits are imposed. Keep host resources and workloads comparable; +container execution and image build options can still affect timing. Use a +fresh work directory for each run. Container cleanup completes before the +lifecycle proceeds to storage inspection or a component restart. + The workflow runs when an allowlisted repository admin comments `/query-regression` on a non-draft PR. It does not rerun on pushes, ready-for-review, or reopen events. `/query-regression` runs the nine routine @@ -565,3 +615,46 @@ Additional SQL optimizer cases: `LIMIT`. - `sql_join_filter_order`: two direct-SST tables joined on a shared tag with time filters, aggregate ordering, and `LIMIT`. + +### Precompiled query-regression helpers + +`helpers_image` is optional in Query Regression (dispatch and reusable calls). +The empty default still builds both helpers from the candidate source. Set it to +`pinned` to use the `QUERY_REGRESSION_HELPERS_IMAGE` repository variable, or provide an explicit +`image@sha256:...` reference. With both database image inputs and a helper image, +the run skips all Cargo builds and build-cache preparation. PR comment admission +continues to require source builds. + +The helper image is an artifact image containing `query_perf_fixture` and +`query_regression_runner`. They are extracted and run on the host, so their child +processes can use the same database launchers and paths. Images are built on the +query-regression ECS runner; alternate runner hosts must have compatible Linux +x86_64 libraries. Extraction smoke-tests both executables and fails without a +source-build fallback. `helpers-image.json` in the report artifact records the +image digest, image ID, and source revision. Case files and orchestration still +come from the candidate ref; choose a compatible helper revision when changing +the case DSL. + +The **Publish query regression helpers** workflow builds and pushes to +`greptime-registry.cn-hangzhou.cr.aliyuncs.com/tools/query-regression-helpers`. +It runs manually or on main when helper sources, packaging, or Cargo.lock change. +After verifying the published digest, it creates or updates the +`QUERY_REGRESSION_HELPERS_IMAGE` Actions repository variable, using +`GH_PERSONAL_ACCESS_TOKEN` with Actions-variable write permission (the same +credential used to update the ECS runner image variable). No version-update PR +is needed. Subsequent runs with `helpers_image: pinned` use that digest; runs +with an explicit digest keep their chosen version, and empty inputs still build +from source. To roll back, set the variable to a previously verified digest. +No mutable `latest` tag is consumed. Local extraction with `pinned` reads the +same name from the environment. +The existing **Release builder images** workflow can also invoke publication via +`release_query_regression_helpers_image`, including for initial branch testing. + +Helper publication separates the ECS build from the hosted registry upload. +The ECS job smoke-tests the packaged binaries and compares fixture plans twice, +then retains the image archive and binary checksums as `helpers-image` for seven +days. The hosted job uses Skopeo (as in builder-image publishing), reads back the +published digest, and checks both binaries against those checksums. If upload +fails, rerun failed jobs to reuse the completed build artifact without compiling +again. Runtime compatibility is checked on ECS; the hosted publishing job only +compares bytes because its system libraries can differ. diff --git a/tests/perf/test_query_regression_helpers.py b/tests/perf/test_query_regression_helpers.py new file mode 100644 index 00000000000..329f34e0759 --- /dev/null +++ b/tests/perf/test_query_regression_helpers.py @@ -0,0 +1,155 @@ +#!/usr/bin/env python3 +# 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. + +"""Checks for helper image extraction and unchanged source-build defaults.""" + +import hashlib +import importlib.util +import json +import os +import subprocess +import tempfile +import unittest +from pathlib import Path +from unittest.mock import patch + +ROOT = Path(__file__).resolve().parents[2] +spec = importlib.util.spec_from_file_location("helpers", ROOT / ".github/scripts/query-regression-helpers.py") +helpers = importlib.util.module_from_spec(spec) +spec.loader.exec_module(helpers) +IMAGE = "registry/tools/helpers@sha256:" + "a" * 64 +REVISION = "b" * 40 + + +class HelpersTest(unittest.TestCase): + def test_repository_variable_and_explicit_digest(self): + with patch.dict(os.environ, {"QUERY_REGRESSION_HELPERS_IMAGE": IMAGE}): + self.assertEqual(helpers.resolve_image("pinned"), IMAGE) + explicit = IMAGE.replace("a" * 64, "b" * 64) + self.assertEqual(helpers.resolve_image(explicit), explicit) + for value in ("", "registry/helpers:latest", IMAGE + "\nmalicious", "--help"): + with self.subTest(value=value), self.assertRaises(ValueError): + helpers.resolve_image(value) + for variables in ({}, {"QUERY_REGRESSION_HELPERS_IMAGE": "registry/helpers:latest"}): + with patch.dict(os.environ, variables, clear=True), self.assertRaises(ValueError): + helpers.resolve_image("pinned") + + def test_extraction_provenance_and_cleanup_on_failure(self): + for fail in (False, True): + with self.subTest(fail=fail), tempfile.TemporaryDirectory() as tmp: + destination = Path(tmp) + calls = [] + + def run(*args): + calls.append(args) + if args[:3] == ("docker", "image", "inspect"): + return json.dumps([{"Id": "sha256:image", "Config": {"Labels": { + "org.opencontainers.image.revision": REVISION}}}]) + if args[:2] == ("docker", "create"): + return "container" + if args[:2] == ("docker", "cp"): + if fail: + raise subprocess.CalledProcessError(1, args) + Path(args[-1]).write_bytes(b"binary") + return "" + + with patch.object(helpers, "run", side_effect=run): + if fail: + with self.assertRaises(subprocess.CalledProcessError): + helpers.extract(IMAGE, destination) + self.assertFalse((destination / "helpers-image.json").exists()) + else: + self.assertEqual(helpers.extract(IMAGE, destination), { + "image": IMAGE, "image_id": "sha256:image", "source_revision": REVISION}) + for binary in helpers.BINARIES: + self.assertIn((str(destination / binary), "--help"), calls) + self.assertEqual(calls[-1], ("docker", "rm", "--force", "--volumes", "container")) + + def test_missing_revision_fails_before_container_creation(self): + with tempfile.TemporaryDirectory() as tmp, patch.object( + helpers, "run", side_effect=["", '[{"Config": {}}]'] + ) as run: + with self.assertRaises(ValueError): + helpers.extract(IMAGE, Path(tmp)) + self.assertEqual(run.call_count, 2) + + def test_published_checksums_and_mismatch_cleanup(self): + for valid in (True, False): + with self.subTest(valid=valid), tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + checksum = hashlib.sha256(b"binary" if valid else b"wrong").hexdigest() + checksums = root / "SHA256SUMS" + checksums.write_text("".join(f"{checksum} {name}\n" for name in helpers.BINARIES)) + calls = [] + + def run(*args): + calls.append(args) + if args[:3] == ("docker", "image", "inspect"): + return json.dumps([{"Id": "sha256:image", "Config": {"Labels": { + "org.opencontainers.image.revision": REVISION}}}]) + if args[:2] == ("docker", "create"): + return "container" + if args[:2] == ("docker", "cp"): + Path(args[-1]).write_bytes(b"binary") + return "" + + with patch.object(helpers, "run", side_effect=run): + if valid: + helpers.extract(IMAGE, root / "bins", checksums) + else: + with self.assertRaisesRegex(ValueError, "checksum mismatch"): + helpers.extract(IMAGE, root / "bins", checksums) + self.assertFalse((root / "bins/helpers-image.json").exists()) + self.assertTrue(all(call[0] == "docker" for call in calls)) + self.assertEqual(calls[-1], ("docker", "rm", "--force", "--volumes", "container")) + + def test_checksum_manifest_requires_both_helpers(self): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + checksums = root / "SHA256SUMS" + checksums.write_text("0" * 64 + " query_perf_fixture\n") + with patch.object(helpers, "run") as run, self.assertRaises(ValueError): + helpers.extract(IMAGE, root, checksums) + run.assert_not_called() + + def test_candidate_build_matrix(self): + # Execute the actual workflow shell with recorder commands. Changing + # only image inputs must preserve source defaults and skip opted-in builds. + import os + workflow = (ROOT / ".github/workflows/query-regression.yml").read_text() + step = workflow.split(" - name: Build candidate greptime and query regression helpers\n", 1)[1] + step = step.split(" - name:", 1)[0].split(" run: |\n", 1)[1] + script = "\n".join(line[10:] for line in step.splitlines()) + for candidate in ("", "server-image" ): + for image in ("", IMAGE): + with self.subTest(candidate=candidate, helpers=image), tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + log = root / "commands" + for command in ("cargo", "cp", "sccache"): + path = root / command + path.write_text('#!/bin/sh\nprintf "%s\\n" "' + command + ' $*" >> "$COMMAND_LOG"\n') + path.chmod(0o755) + env = dict(os.environ, PATH=tmp + os.pathsep + os.environ["PATH"], + COMMAND_LOG=str(log), CANDIDATE_IMAGE=candidate, HELPERS_IMAGE=image, + CARGO_PROFILE="nightly", CARGO_TARGET_DIR=str(root / "target"), + GITHUB_WORKSPACE=tmp, GITHUB_STEP_SUMMARY=str(root / "summary")) + subprocess.run(["bash", "-c", script], env=env, check=True, capture_output=True) + builds = [line for line in log.read_text().splitlines() if line.startswith("cargo ")] + expected = ([] if candidate else ["greptime"]) + ([] if image else list(helpers.BINARIES)) + self.assertEqual([line.split("--bin ")[1].split()[0] for line in builds], expected) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/perf/test_query_regression_runner_compaction_toctou.py b/tests/perf/test_query_regression_runner_compaction_toctou.py index 58f7f31b7fa..120e74d22f4 100644 --- a/tests/perf/test_query_regression_runner_compaction_toctou.py +++ b/tests/perf/test_query_regression_runner_compaction_toctou.py @@ -13,15 +13,18 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Regression coverage for the outer remote-write lifecycle.""" +"""Regression coverage for remote-write and Docker target lifecycles.""" import importlib.util import json import sys +import signal +import subprocess import tempfile import unittest from pathlib import Path -from unittest.mock import patch +from types import SimpleNamespace +from unittest.mock import Mock, patch RUNNER_PATH = Path(__file__).resolve().parents[2] / ".github/scripts/query-regression-run.py" @@ -110,5 +113,129 @@ class RemoteWriteLifecycleTest(unittest.TestCase): self.assertIn("stop:candidate:frontend", events) +class DockerTargetTest(unittest.TestCase): + image_id = "sha256:" + "a" * 64 + + def test_independent_target_selection_and_image_identity(self) -> None: + for base_image, candidate_image in (("", ""), ("image:v1", ""), ("", "image:v1"), ("image:v1", "image:v1")): + with self.subTest(base=base_image, candidate=candidate_image), tempfile.TemporaryDirectory() as tmp: + args = SimpleNamespace( + base_image=base_image, candidate_image=candidate_image, + base_bin=Path("base"), candidate_bin=Path("candidate"), + base_ref="base-sha", candidate_ref="candidate-sha", work_dir=Path(tmp), + ) + result = SimpleNamespace(returncode=0, stdout=json.dumps([ + {"Id": self.image_id, "RepoDigests": ["image@sha256:abc"]}, + ])) + with patch.object(runner.subprocess, "run", return_value=result) as run: + metadata = runner.prepare_targets(args) + self.assertEqual(run.call_count, int(bool(base_image or candidate_image))) + for name, image in (("base", base_image), ("candidate", candidate_image)): + if image: + launcher = getattr(args, f"{name}_bin") + self.assertTrue(launcher.stat().st_mode & 0o111) + self.assertIn(self.image_id, launcher.read_text()) + self.assertIn('"$@"', launcher.read_text()) + self.assertEqual(metadata[name]["image"], image) + self.assertEqual(metadata[name]["image_id"], self.image_id) + else: + self.assertEqual(getattr(args, f"{name}_bin"), Path(name)) + self.assertEqual(metadata[name], {"kind": "source", "ref": f"{name}-sha"}) + + def test_default_source_run_preserves_report_and_refs(self) -> None: + payload = '{"status":"ok","targets":[],"thresholds":[]}\n' + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + def run_case(args, _case_path, work_dir): + self.assertEqual((args.base_image, args.candidate_image), ("", "")) + self.assertEqual((args.base_ref, args.candidate_ref), ("main", "feature")) + self.assertEqual(args.cargo_profile, "nightly") + work_dir.mkdir(parents=True) + (work_dir / "query-regression-report.json").write_text(payload) + return 0 + argv = ["runner", "--cases", "case.toml", "--work-dir", str(root), + "--base-ref", "main", "--candidate-ref", "feature"] + with (patch.dict(runner.os.environ, {}, clear=True), + patch.object(sys, "argv", argv), + patch.object(runner, "run_case", side_effect=run_case), + patch.object(runner, "write_summary", return_value=0), + patch.object(runner.subprocess, "run") as subprocess_run): + self.assertEqual(runner.main(), 0) + subprocess_run.assert_not_called() + reports = list(root.rglob("query-regression-report.json")) + self.assertEqual(len(reports), 1) + self.assertEqual(reports[0].read_bytes(), payload.encode()) + + def test_missing_image_is_pulled_and_inspected(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + args = SimpleNamespace(base_image="image:v1", candidate_image="", base_ref="", candidate_ref="", work_dir=Path(tmp)) + results = [SimpleNamespace(returncode=1), SimpleNamespace(returncode=0), + SimpleNamespace(stdout=json.dumps([{"Id": self.image_id}]))] + with patch.object(runner.subprocess, "run", side_effect=results) as run: + runner.prepare_targets(args) + self.assertEqual([call.args[0][:2] for call in run.call_args_list], + [["docker", "image"], ["docker", "pull"], ["docker", "image"]]) + + def test_container_exit_code_and_cleanup(self) -> None: + commands = [] + def run(command, **kwargs): + commands.append(command) + self.assertTrue(kwargs["start_new_session"]) + return SimpleNamespace(returncode=0, stdout="42\n") + with patch.object(runner.subprocess, "run", side_effect=run): + status = runner.docker_target(self.image_id, Path("/tmp/work space"), ["datanode", "scanbench", "--help"]) + self.assertEqual(status, 42) + create = commands[0] + self.assertIn("/tmp/work space:/tmp/work space", create) + self.assertEqual(create[-4:], [self.image_id, "datanode", "scanbench", "--help"]) + self.assertEqual([cmd[1] for cmd in commands], ["create", "start", "inspect", "stop", "rm"]) + self.assertEqual(commands[-1][-1], create[3]) + + def test_signal_and_failed_start_both_remove_container(self) -> None: + for interruption in (True, False): + commands = [] + def run(command, **kwargs): + commands.append(command) + if command[1] == "start": + if interruption: + signal.getsignal(signal.SIGTERM)(signal.SIGTERM, None) + raise subprocess.CalledProcessError(1, command) + return SimpleNamespace(returncode=0) + previous = signal.getsignal(signal.SIGTERM) + with self.subTest(interruption=interruption), patch.object(runner.subprocess, "run", side_effect=run): + with self.assertRaises(SystemExit if interruption else subprocess.CalledProcessError): + runner.docker_target(self.image_id, Path("/tmp/work"), ["metasrv", "start"]) + self.assertEqual([cmd[1] for cmd in commands], ["create", "start", "stop", "rm"]) + self.assertEqual(signal.getsignal(signal.SIGTERM), previous) + + + def test_stop_timeout_still_force_removes(self) -> None: + commands = [] + def run(command, **kwargs): + commands.append(command) + if command[1] == "stop": + raise subprocess.TimeoutExpired(command, 8) + return SimpleNamespace(returncode=0, stdout="0") + with patch.object(runner.subprocess, "run", side_effect=run): + self.assertEqual(runner.docker_target(self.image_id, Path("/tmp/work"), ["--version"]), 0) + self.assertEqual(commands[-1][1], "rm") + + def test_failed_container_cleanup_is_not_ignored(self) -> None: + target = runner.make_target("base", Path("launcher"), Path("/tmp/work"), list(range(8)), containerized=True) + proc = Mock() + proc.poll.return_value = None + proc.returncode = 1 + with patch.object(runner.os, "killpg"): + with self.assertRaisesRegex(RuntimeError, "cleanup failed"): + runner.stop_component(target, "datanode", {("base", "datanode"): proc}) + + def test_cleanup_failure_does_not_skip_other_components(self) -> None: + target = runner.make_target("base", Path("launcher"), Path("/tmp/work"), list(range(8)), containerized=True) + with patch.object(runner, "stop_component", side_effect=[RuntimeError("failed"), None, None]) as stop: + with self.assertRaisesRegex(RuntimeError, "component cleanup failed"): + runner.stop_all([target], {}) + self.assertEqual([call.args[1] for call in stop.call_args_list], ["frontend", "datanode", "metasrv"]) + + if __name__ == "__main__": unittest.main() From 8a38b901352c89729eb3e9be2c4282c2e749383b Mon Sep 17 00:00:00 2001 From: zhaiAohan <1573612521@qq.com> Date: Fri, 30 Oct 2026 04:06:01 +0000 Subject: [PATCH 145/239] feat: add relational region engine skeleton and registration (#9953) * feat: add relational region engine skeleton and registration Add a read-only relational RegionEngine (create/open/drop lifecycle) and register it as the 4th engine in build_store_engines alongside mito/metric/file. Add a datanode-level test asserting all 4 engines are registered. Signed-off-by: tian1220A <1573612521@qq.com> * feat: add relational connection and idempotent region lifecycle Move ConnectionInfo into relational-engine (DSN parsing + redacted rendering), place the RelationalOptions/RELATIONAL_META_KEY serialization contract in table::requests, make CREATE idempotent by comparing the persisted manifest definition against the incoming request (RegionExists on mismatch), and drop the empty EngineConfig placeholder. Signed-off-by: tian1220A <1573612521@qq.com> * refactor: address review comments on relational region engine - remove connection.rs (only used by its own tests); it will be reintroduced when connection handling is wired up - drop RelationalOptions::version (not validated or used to select a deserialization format) - validate an idempotent create against the existing region definition (region id / schema / primary key / options) and return Unexpected on mismatch, with a reason field on RegionExists - use serde.workspace and remove the empty [features] section - cargo fmt Signed-off-by: tian1220A <1573612521@qq.com> * fix: adapt relational engine to latest main - implement RegionEngine::set_region_role_with_lease (no-op lease, the engine takes no writes so no lease is needed) - pass node_id to build_object_store_manager in the registration test - drop stale url dependency from the relational-engine lock entry Signed-off-by: tian1220A <1573612521@qq.com> * ci: register relational-engine in scope rules Signed-off-by: tian1220A <1573612521@qq.com> * fix: remove table self-dependency from Cargo.lock Signed-off-by: tian1220A <1573612521@qq.com> --------- Signed-off-by: tian1220A <1573612521@qq.com> --- .github/ci/scope-rules.toml | 1 + Cargo.lock | 23 + Cargo.toml | 2 + src/common/catalog/src/consts.rs | 1 + src/datanode/Cargo.toml | 1 + src/datanode/src/datanode.rs | 56 ++ src/datanode/src/region_server.rs | 10 +- src/relational-engine/Cargo.toml | 28 + src/relational-engine/src/engine.rs | 642 ++++++++++++++++++++++ src/relational-engine/src/error.rs | 181 ++++++ src/relational-engine/src/lib.rs | 21 + src/relational-engine/src/manifest.rs | 195 +++++++ src/relational-engine/src/region.rs | 312 +++++++++++ src/relational-engine/src/test_util.rs | 67 +++ src/store-api/src/metric_engine_consts.rs | 2 + src/store-api/src/region_engine.rs | 10 +- src/table/src/requests.rs | 13 + 17 files changed, 1561 insertions(+), 4 deletions(-) create mode 100644 src/relational-engine/Cargo.toml create mode 100644 src/relational-engine/src/engine.rs create mode 100644 src/relational-engine/src/error.rs create mode 100644 src/relational-engine/src/lib.rs create mode 100644 src/relational-engine/src/manifest.rs create mode 100644 src/relational-engine/src/region.rs create mode 100644 src/relational-engine/src/test_util.rs diff --git a/.github/ci/scope-rules.toml b/.github/ci/scope-rules.toml index 0c9920d5711..05ee736d625 100644 --- a/.github/ci/scope-rules.toml +++ b/.github/ci/scope-rules.toml @@ -141,6 +141,7 @@ engine = [ "src/mito-codec/**", "src/metric-engine/**", "src/file-engine/**", + "src/relational-engine/**", "src/log-store/**", "src/common/wal/**", ] diff --git a/Cargo.lock b/Cargo.lock index 3b7ab8897ec..6b7bd35e546 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4567,6 +4567,7 @@ dependencies = [ "prost 0.14.1", "query", "rand 0.9.4", + "relational-engine", "reqwest 0.13.4", "serde", "serde_json", @@ -12460,6 +12461,28 @@ dependencies = [ "memchr", ] +[[package]] +name = "relational-engine" +version = "1.4.0-alpha.0" +dependencies = [ + "api", + "async-trait", + "common-catalog", + "common-error", + "common-macro", + "common-telemetry", + "common-test-util", + "common-time", + "datatypes", + "object-store", + "serde", + "serde_json", + "snafu 0.8.6", + "store-api", + "table", + "tokio", +] + [[package]] name = "relative-path" version = "1.9.3" diff --git a/Cargo.toml b/Cargo.toml index e1271bdecd1..a6f77b1e9de 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -64,6 +64,7 @@ members = [ "src/promql", "src/puffin", "src/query", + "src/relational-engine", "src/standalone", "src/servers", "src/session", @@ -355,6 +356,7 @@ plugins = { path = "src/plugins" } promql = { path = "src/promql" } puffin = { path = "src/puffin" } query = { path = "src/query" } +relational-engine = { path = "src/relational-engine" } servers = { path = "src/servers" } session = { path = "src/session" } sql = { path = "src/sql" } diff --git a/src/common/catalog/src/consts.rs b/src/common/catalog/src/consts.rs index d999f871487..27680491ff0 100644 --- a/src/common/catalog/src/consts.rs +++ b/src/common/catalog/src/consts.rs @@ -152,6 +152,7 @@ pub fn default_engine() -> &'static str { } pub const FILE_ENGINE: &str = "file"; +pub const RELATIONAL_ENGINE: &str = "relational"; pub const SEMANTIC_TYPE_PRIMARY_KEY: &str = "TAG"; pub const SEMANTIC_TYPE_FIELD: &str = "FIELD"; diff --git a/src/datanode/Cargo.toml b/src/datanode/Cargo.toml index 45cdc783197..93bb6e89ed8 100644 --- a/src/datanode/Cargo.toml +++ b/src/datanode/Cargo.toml @@ -62,6 +62,7 @@ prometheus.workspace = true prost.workspace = true query.workspace = true rand.workspace = true +relational-engine.workspace = true reqwest.workspace = true serde.workspace = true serde_json.workspace = true diff --git a/src/datanode/src/datanode.rs b/src/datanode/src/datanode.rs index 6ebe7ffe05d..9a797f49c95 100644 --- a/src/datanode/src/datanode.rs +++ b/src/datanode/src/datanode.rs @@ -61,6 +61,7 @@ use object_store::util::normalize_dir; use query::QueryEngineFactory; use query::dummy_catalog::TableProviderFactoryRef; use query::region_query::RegionQueryHandlerFactoryRef; +use relational_engine::engine::RelationalEngine; use servers::server::ServerHandlers; use snafu::{OptionExt, ResultExt, ensure}; use store_api::logstore::LogStore; @@ -710,10 +711,14 @@ impl DatanodeBuilder { self.local_file_access.clone(), ); + let relational_engine = + RelationalEngine::new(object_store_manager.default_object_store().clone()); + Ok(vec![ Arc::new(mito_engine) as _, Arc::new(metric_engine) as _, Arc::new(file_engine) as _, + Arc::new(relational_engine) as _, ]) } @@ -1146,7 +1151,9 @@ mod tests { use common_config::Configurable; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; + use common_meta::cache::{SchemaCacheRef, TableSchemaCacheRef}; use common_meta::key::RegionRoleSet; + use common_meta::key::SchemaMetadataManager; use common_meta::key::datanode_table::DatanodeTableManager; use common_meta::kv_backend::KvBackendRef; use common_meta::kv_backend::memory::MemoryKvBackend; @@ -1159,6 +1166,7 @@ mod tests { use log_store::error::Error as LogStoreError; use meta_client::client::MetaClientBuilder; use mito2::engine::MITO_ENGINE_NAME; + use mito2::sst::file_ref::FileReferenceManager; use mito2::test_util::{CreateRequestBuilder, build_rows, rows_schema}; use object_store::ObjectStore; use object_store::config::{S3Config, S3Connection}; @@ -2006,4 +2014,52 @@ mod tests { }; validate_object_store_wal_config(&derived).unwrap(); } + + #[tokio::test] + async fn test_build_store_engines_registers_relational_engine() { + common_telemetry::init_default_ut_logging(); + let kv_backend = Arc::new(MemoryKvBackend::new()); + let cache_registry = Arc::new(build_datanode_layered_cache_registry(kv_backend.clone())); + let schema_cache: SchemaCacheRef = cache_registry.get().unwrap(); + let table_id_schema_cache: TableSchemaCacheRef = cache_registry.get().unwrap(); + let schema_metadata_manager = Arc::new(SchemaMetadataManager::new( + table_id_schema_cache, + schema_cache, + )); + let file_ref_manager = Arc::new(FileReferenceManager::with_gc_enabled(Some(0), false)); + + let data_home = create_temp_dir("build-store-engines"); + let mut opts = DatanodeOptions { + node_id: Some(0), + wal: DatanodeWalConfig::Noop, + enable_telemetry: false, + ..Default::default() + }; + opts.storage.data_home = data_home.path().to_str().unwrap().to_string(); + + let object_store_manager = + DatanodeBuilder::build_object_store_manager(&opts.storage, opts.node_id.unwrap()) + .await + .unwrap(); + + let mut builder = DatanodeBuilder::new(opts, Plugins::default(), kv_backend); + let engines = builder + .build_store_engines( + object_store_manager, + schema_metadata_manager, + file_ref_manager, + Plugins::default(), + ) + .await + .unwrap(); + + let names: Vec<&str> = engines.iter().map(|engine| engine.name()).collect(); + assert_eq!(engines.len(), 4, "expected 4 engines, got: {names:?}"); + for expected in [MITO_ENGINE_NAME, "metric", "file", "relational"] { + assert!( + names.contains(&expected), + "engine {expected:?} not registered, got: {names:?}" + ); + } + } } diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index 9d4c4b375fb..4ae350002c9 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -77,7 +77,7 @@ use session::context::{ }; use snafu::{IntoError, OptionExt, ResultExt, ensure}; use store_api::metric_engine_consts::{ - FILE_ENGINE_NAME, LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME, + FILE_ENGINE_NAME, LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME, RELATIONAL_ENGINE_NAME, }; use store_api::region_engine::{ RegionEngineRef, RegionManifestInfo, RegionRole, RegionStatistic, RemapManifestsRequest, @@ -2074,6 +2074,9 @@ impl RegionServerInner { RegionAttribute::File => { // do nothing } + RegionAttribute::Relational => { + // do nothing + } } } RegionChange::Deregisters => { @@ -2281,7 +2284,7 @@ fn is_metric_engine(engine: &str) -> bool { /// Records failed open attempts without using arbitrary engine names as metric labels. fn record_region_open_failures(engine: &str, status_code: StatusCode, count: usize) { let engine = match engine { - MITO_ENGINE_NAME | METRIC_ENGINE_NAME | FILE_ENGINE_NAME => engine, + MITO_ENGINE_NAME | METRIC_ENGINE_NAME | FILE_ENGINE_NAME | RELATIONAL_ENGINE_NAME => engine, _ => "unknown", }; crate::metrics::REGION_OPEN_FAILURES_TOTAL @@ -2330,6 +2333,7 @@ fn parse_region_attribute( Ok(RegionAttribute::Metric { physical }) } FILE_ENGINE_NAME => Ok(RegionAttribute::File), + RELATIONAL_ENGINE_NAME => Ok(RegionAttribute::Relational), _ => error::UnexpectedSnafu { violated: format!("Unknown engine: {}", engine), } @@ -2342,6 +2346,7 @@ enum RegionAttribute { Mito, Metric { physical: bool }, File, + Relational, } impl RegionAttribute { @@ -2350,6 +2355,7 @@ impl RegionAttribute { RegionAttribute::Mito => MITO_ENGINE_NAME, RegionAttribute::Metric { .. } => METRIC_ENGINE_NAME, RegionAttribute::File => FILE_ENGINE_NAME, + RegionAttribute::Relational => RELATIONAL_ENGINE_NAME, } } } diff --git a/src/relational-engine/Cargo.toml b/src/relational-engine/Cargo.toml new file mode 100644 index 00000000000..0d9a4666aa0 --- /dev/null +++ b/src/relational-engine/Cargo.toml @@ -0,0 +1,28 @@ +[package] +name = "relational-engine" +version.workspace = true +edition.workspace = true +license.workspace = true + +[lints] +workspace = true + +[dependencies] +api.workspace = true +async-trait.workspace = true +common-catalog.workspace = true +common-error.workspace = true +common-macro.workspace = true +common-telemetry.workspace = true +object-store.workspace = true +serde.workspace = true +serde_json.workspace = true +snafu.workspace = true +store-api.workspace = true +table.workspace = true +tokio.workspace = true + +[dev-dependencies] +common-test-util.workspace = true +common-time.workspace = true +datatypes.workspace = true diff --git a/src/relational-engine/src/engine.rs b/src/relational-engine/src/engine.rs new file mode 100644 index 00000000000..3d84c10971d --- /dev/null +++ b/src/relational-engine/src/engine.rs @@ -0,0 +1,642 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::any::Any; +use std::collections::HashMap; +use std::sync::{Arc, RwLock}; +use std::time::Instant; + +use api::region::RegionResponse; +use async_trait::async_trait; +use common_catalog::consts::RELATIONAL_ENGINE; +use common_error::ext::BoxedError; +use common_telemetry::{error, info}; +use object_store::ObjectStore; +use snafu::{OptionExt, ensure}; +use store_api::metadata::RegionMetadataRef; +use store_api::region_engine::{ + RegionEngine, RegionRole, RegionScannerRef, RegionStatistic, RemapManifestsRequest, + RemapManifestsResponse, SetRegionRoleStateResponse, SetRegionRoleStateSuccess, + SettableRegionRoleState, SyncRegionFromRequest, SyncRegionFromResponse, +}; +use store_api::region_request::{ + AffectedRows, RegionCloseRequest, RegionCreateRequest, RegionDropRequest, RegionOpenRequest, + RegionRequest, RegionRequirements, +}; +use store_api::storage::{RegionId, ScanRequest, SequenceNumber}; +use tokio::sync::Mutex; + +use crate::error::{ + RegionNotFoundSnafu, Result as EngineResult, UnexpectedEngineSnafu, UnsupportedSnafu, +}; +use crate::manifest::RelationalRegionManifest; +use crate::region::{RelationalRegion, RelationalRegionRef}; + +pub struct RelationalEngine { + inner: EngineInnerRef, +} + +impl RelationalEngine { + pub fn new(object_store: ObjectStore) -> Self { + Self { + inner: Arc::new(EngineInner::new(object_store)), + } + } +} + +#[async_trait] +impl RegionEngine for RelationalEngine { + fn name(&self) -> &str { + RELATIONAL_ENGINE + } + + async fn handle_request( + &self, + region_id: RegionId, + request: RegionRequest, + ) -> Result { + self.inner + .handle_request(region_id, request) + .await + .map_err(BoxedError::new) + } + + async fn handle_query( + &self, + _region_id: RegionId, + _request: ScanRequest, + ) -> Result { + Err(BoxedError::new( + UnsupportedSnafu { operation: "scan" }.build(), + )) + } + + async fn get_metadata(&self, region_id: RegionId) -> Result { + self.inner + .get_region(region_id) + .await + .map(|r| r.metadata()) + .context(RegionNotFoundSnafu { region_id }) + .map_err(BoxedError::new) + } + + async fn stop(&self) -> Result<(), BoxedError> { + self.inner.stop().await.map_err(BoxedError::new) + } + + fn region_statistic(&self, _: RegionId) -> Option { + None + } + + async fn get_committed_sequence(&self, _: RegionId) -> Result { + Ok(Default::default()) + } + + fn set_region_role(&self, region_id: RegionId, role: RegionRole) -> Result<(), BoxedError> { + self.inner + .set_region_role(region_id, role) + .map_err(BoxedError::new) + } + + /// Regions of this engine take no writes, so the lease is not needed. + fn set_region_role_with_lease( + &self, + region_id: RegionId, + role: RegionRole, + _sent_at: Instant, + _deadline: Instant, + ) -> Result { + self.set_region_role(region_id, role).map(|_| true) + } + + async fn set_region_role_state_gracefully( + &self, + region_id: RegionId, + _region_role_state: SettableRegionRoleState, + ) -> Result { + let exists = self.inner.get_region(region_id).await.is_some(); + + if exists { + Ok(SetRegionRoleStateResponse::success( + SetRegionRoleStateSuccess::relational(), + )) + } else { + Ok(SetRegionRoleStateResponse::NotFound) + } + } + + async fn sync_region( + &self, + _region_id: RegionId, + _request: SyncRegionFromRequest, + ) -> Result { + // Relational engine doesn't need to sync region manifest. + Ok(SyncRegionFromResponse::NotSupported) + } + + async fn remap_manifests( + &self, + _request: RemapManifestsRequest, + ) -> Result { + Err(BoxedError::new( + UnsupportedSnafu { + operation: "remap_manifests", + } + .build(), + )) + } + + fn role(&self, region_id: RegionId) -> Option { + self.inner.state(region_id) + } + + fn as_any(&self) -> &dyn Any { + self + } +} + +struct EngineInner { + /// All regions opened by the engine. + /// + /// Writing to `regions` should also hold the `region_mutex`. + regions: RwLock>, + + /// Region mutex is used to protect the operations such as creating/opening/closing + /// a region, to avoid things like opening the same region simultaneously. + region_mutex: Mutex<()>, + + object_store: ObjectStore, +} + +type EngineInnerRef = Arc; + +fn ensure_region_requirements( + requirements: RegionRequirements, + object_store: &ObjectStore, +) -> EngineResult<()> { + if !requirements.object_storage { + return Ok(()); + } + + ensure!( + object_store::util::is_object_storage(object_store), + UnsupportedSnafu { + operation: "create or open region with object storage requirement on non-object storage" + } + ); + + Ok(()) +} + +impl EngineInner { + fn new(object_store: ObjectStore) -> Self { + Self { + regions: RwLock::new(HashMap::new()), + region_mutex: Mutex::new(()), + object_store, + } + } + + async fn handle_request( + &self, + region_id: RegionId, + request: RegionRequest, + ) -> EngineResult { + let result = match request { + RegionRequest::Create(req) => self.handle_create(region_id, req).await, + RegionRequest::Drop(req) => self.handle_drop(region_id, req).await, + RegionRequest::Open(req) => self.handle_open(region_id, req).await, + RegionRequest::Close(req) => self.handle_close(region_id, req).await, + _ => UnsupportedSnafu { + operation: request.to_string(), + } + .fail(), + }; + result.map(RegionResponse::new) + } + + async fn stop(&self) -> EngineResult<()> { + let _lock = self.region_mutex.lock().await; + self.regions.write().unwrap().clear(); + Ok(()) + } + + fn set_region_role(&self, _region_id: RegionId, _region_role: RegionRole) -> EngineResult<()> { + // TODO(zhongzc): Improve the semantics and implementation of this API. + Ok(()) + } + + fn state(&self, region_id: RegionId) -> Option { + if self.regions.read().unwrap().get(®ion_id).is_some() { + Some(RegionRole::Leader) + } else { + None + } + } +} + +impl EngineInner { + async fn handle_create( + &self, + region_id: RegionId, + request: RegionCreateRequest, + ) -> EngineResult { + ensure!( + request.engine == RELATIONAL_ENGINE, + UnexpectedEngineSnafu { + engine: request.engine + } + ); + + // If the region is already open, validate the definition before treating the + // create as idempotent: a mismatched create must fail, not silently succeed. + if let Some(region) = self.get_region(region_id).await { + region.check_create_request(region_id, &request)?; + return Ok(0); + } + + info!("Try to create region, region_id: {}", region_id); + + let _lock = self.region_mutex.lock().await; + // Check again after acquiring the lock + if let Some(region) = self.get_region(region_id).await { + region.check_create_request(region_id, &request)?; + return Ok(0); + } + + ensure_region_requirements(request.requirements, &self.object_store)?; + + // The manifest is the source of truth for whether a region already + // exists. The in-memory `regions` map only tracks regions opened in + // this process, so after a restart it is empty even though the + // manifest was already persisted. Recover the region from the manifest + // instead of failing with `ManifestExists`, keeping `CREATE EXTERNAL + // TABLE` idempotent. + let region_dir = RelationalRegion::region_dir(&request.table_dir, region_id); + if RelationalRegionManifest::exists(®ion_dir, &self.object_store).await? { + let manifest = + RelationalRegionManifest::load(region_id, ®ion_dir, &self.object_store).await?; + let region = RelationalRegion::from_manifest(®ion_dir, manifest)?; + // The persisted definition must match the incoming request, otherwise a + // retried CREATE with a different schema would silently reuse the old + // manifest. Reject the mismatch instead. + region.check_create_request(region_id, &request)?; + self.regions.write().unwrap().insert(region_id, region); + info!( + "Region recovered from existing manifest, region_id: {}", + region_id + ); + return Ok(0); + } + + let res = RelationalRegion::create(region_id, request, &self.object_store).await; + let region = res.inspect_err(|err| { + error!( + err; + "Failed to create region, region_id: {}", + region_id + ); + })?; + self.regions.write().unwrap().insert(region_id, region); + + info!("A new region is created, region_id: {}", region_id); + Ok(0) + } + + async fn handle_open( + &self, + region_id: RegionId, + request: RegionOpenRequest, + ) -> EngineResult { + if self.exists(region_id).await { + return Ok(0); + } + + info!("Try to open region, region_id: {}", region_id); + + let _lock = self.region_mutex.lock().await; + // Check again after acquiring the lock + if self.exists(region_id).await { + return Ok(0); + } + + ensure_region_requirements(request.requirements, &self.object_store)?; + + let res = RelationalRegion::open(region_id, request, &self.object_store).await; + let region = res.inspect_err(|err| { + error!( + err; + "Failed to open region, region_id: {}", + region_id + ); + })?; + self.regions.write().unwrap().insert(region_id, region); + + info!("Region opened, region_id: {}", region_id); + Ok(0) + } + + async fn handle_close( + &self, + region_id: RegionId, + _request: RegionCloseRequest, + ) -> EngineResult { + let _lock = self.region_mutex.lock().await; + + let mut regions = self.regions.write().unwrap(); + if regions.remove(®ion_id).is_some() { + info!("Region closed, region_id: {}", region_id); + } + + Ok(0) + } + + async fn handle_drop( + &self, + region_id: RegionId, + _request: RegionDropRequest, + ) -> EngineResult { + // Serialize against create/open/close so the region can't be closed + // between the existence check and the manifest deletion. A region that + // is no longer in memory (e.g. closed concurrently) is reported as not + // found instead of silently returning success with a leaked manifest. + let _lock = self.region_mutex.lock().await; + + let region = self.get_region(region_id).await; + let Some(region) = region else { + return RegionNotFoundSnafu { region_id }.fail(); + }; + + info!("Try to drop region, region_id: {}", region_id); + + let res = RelationalRegion::drop(®ion, &self.object_store).await; + res.inspect_err(|err| { + error!( + err; + "Failed to drop region, region_id: {}", + region_id + ); + })?; + + let _ = self.regions.write().unwrap().remove(®ion_id); + + info!("Region dropped, region_id: {}", region_id); + Ok(0) + } + + async fn get_region(&self, region_id: RegionId) -> Option { + self.regions.read().unwrap().get(®ion_id).cloned() + } + + async fn exists(&self, region_id: RegionId) -> bool { + self.regions.read().unwrap().contains_key(®ion_id) + } +} + +#[cfg(test)] +mod tests { + use std::assert_matches; + + use object_store::ObjectStore; + use object_store::services::{Fs, S3}; + use store_api::region_request::{ + PathType, RegionCreateRequest, RegionDropRequest, RegionRequirements, + }; + use store_api::storage::RegionId; + + use super::{EngineInner, ensure_region_requirements}; + use crate::error::Error; + use crate::test_util::{new_test_column_metadata, new_test_object_store, new_test_options}; + + fn build_fs_object_store() -> ObjectStore { + ObjectStore::new(Fs::default().root("/tmp")).unwrap() + } + + fn build_s3_object_store() -> ObjectStore { + ObjectStore::new( + S3::default() + .bucket("test-bucket") + .region("us-east-1") + .disable_ec2_metadata(), + ) + .unwrap() + } + + #[test] + fn test_empty_region_requirements_are_supported() { + ensure_region_requirements(RegionRequirements::empty(), &build_fs_object_store()).unwrap(); + } + + #[test] + fn test_object_storage_region_requirement_rejects_fs_object_store() { + let err = ensure_region_requirements( + RegionRequirements::object_storage(), + &build_fs_object_store(), + ) + .unwrap_err(); + + assert!(matches!(err, Error::Unsupported { .. })); + } + + #[test] + fn test_object_storage_region_requirement_accepts_s3_object_store() { + ensure_region_requirements( + RegionRequirements::object_storage(), + &build_s3_object_store(), + ) + .unwrap(); + } + + #[tokio::test] + async fn test_drop_returns_not_found_when_region_absent() { + let (_dir, object_store) = new_test_object_store("test_drop_absent"); + let inner = EngineInner::new(object_store); + + let err = inner + .handle_drop( + RegionId::new(1, 0), + RegionDropRequest { + fast_path: false, + force: false, + partial_drop: false, + }, + ) + .await + .unwrap_err(); + + assert_matches!(err, Error::RegionNotFound { .. }); + } + + #[tokio::test] + async fn test_drop_removes_manifest() { + let (_dir, object_store) = new_test_object_store("test_drop_removes_manifest"); + let inner = EngineInner::new(object_store.clone()); + + let region_id = RegionId::new(1, 0); + let request = RegionCreateRequest { + engine: "relational".to_string(), + column_metadatas: new_test_column_metadata(), + primary_key: vec![1], + options: new_test_options(), + table_dir: "drop_dir/".to_string(), + path_type: PathType::Bare, + partition_expr_json: Some("".to_string()), + requirements: Default::default(), + }; + + inner.handle_create(region_id, request).await.unwrap(); + + let manifest_path = "drop_dir/1_0000000000/manifest/_relational_manifest"; + assert!(object_store.exists(manifest_path).await.unwrap()); + + inner + .handle_drop( + region_id, + RegionDropRequest { + fast_path: false, + force: false, + partial_drop: false, + }, + ) + .await + .unwrap(); + + assert!(!object_store.exists(manifest_path).await.unwrap()); + assert!(!inner.exists(region_id).await); + } + + #[tokio::test] + async fn test_create_is_idempotent_across_engine_rebuild() { + let (_dir, object_store) = new_test_object_store("test_create_idempotent"); + let region_id = RegionId::new(1, 0); + let request = RegionCreateRequest { + engine: "relational".to_string(), + column_metadatas: new_test_column_metadata(), + primary_key: vec![1], + options: new_test_options(), + table_dir: "idempotent_dir/".to_string(), + path_type: PathType::Bare, + partition_expr_json: Some("".to_string()), + requirements: Default::default(), + }; + + // First create persists the manifest to object storage. + { + let inner = EngineInner::new(object_store.clone()); + inner + .handle_create(region_id, request.clone()) + .await + .unwrap(); + assert!(inner.exists(region_id).await); + } + + // Rebuild the engine on the same object store: the in-memory region + // map is empty, but the manifest is still persisted. Retrying the + // same create must recover from the manifest instead of failing. + let rebuilt = EngineInner::new(object_store.clone()); + let affected = rebuilt.handle_create(region_id, request).await.unwrap(); + assert_eq!(affected, 0); + assert!(rebuilt.exists(region_id).await); + + let manifest_path = "idempotent_dir/1_0000000000/manifest/_relational_manifest"; + assert!(object_store.exists(manifest_path).await.unwrap()); + } + + #[tokio::test] + async fn test_create_rejects_mismatched_definition() { + let (_dir, object_store) = new_test_object_store("test_create_mismatch"); + let region_id = RegionId::new(1, 0); + + let request = RegionCreateRequest { + engine: "relational".to_string(), + column_metadatas: new_test_column_metadata(), + primary_key: vec![1], + options: new_test_options(), + table_dir: "mismatch_dir/".to_string(), + path_type: PathType::Bare, + partition_expr_json: Some("".to_string()), + requirements: Default::default(), + }; + + // First create persists the manifest. + let inner = EngineInner::new(object_store.clone()); + inner.handle_create(region_id, request).await.unwrap(); + + // Rebuild the engine and retry the create with a different primary key: + // the persisted definition no longer matches, so this must fail instead of + // silently reusing the old manifest. + let rebuilt = EngineInner::new(object_store.clone()); + let mismatched = RegionCreateRequest { + engine: "relational".to_string(), + column_metadatas: new_test_column_metadata(), + primary_key: vec![2], + options: new_test_options(), + table_dir: "mismatch_dir/".to_string(), + path_type: PathType::Bare, + partition_expr_json: Some("".to_string()), + requirements: Default::default(), + }; + let err = rebuilt + .handle_create(region_id, mismatched) + .await + .unwrap_err(); + assert_matches!(err, Error::RegionExists { .. }); + } + + #[tokio::test] + async fn test_create_rejects_mismatched_definition_in_memory() { + let (_dir, object_store) = new_test_object_store("test_create_mismatch_in_memory"); + let inner = EngineInner::new(object_store); + let region_id = RegionId::new(1, 0); + + let request = RegionCreateRequest { + engine: "relational".to_string(), + column_metadatas: new_test_column_metadata(), + primary_key: vec![1], + options: new_test_options(), + table_dir: "mismatch_in_memory_dir/".to_string(), + path_type: PathType::Bare, + partition_expr_json: Some("".to_string()), + requirements: Default::default(), + }; + + inner + .handle_create(region_id, request.clone()) + .await + .unwrap(); + + // An identical create is idempotent while the region stays in memory. + assert_eq!(inner.handle_create(region_id, request).await.unwrap(), 0); + + // A mismatched create must fail, not silently return success. + let mismatched = RegionCreateRequest { + engine: "relational".to_string(), + column_metadatas: new_test_column_metadata(), + primary_key: vec![2], + options: new_test_options(), + table_dir: "mismatch_in_memory_dir/".to_string(), + path_type: PathType::Bare, + partition_expr_json: Some("".to_string()), + requirements: Default::default(), + }; + let err = inner + .handle_create(region_id, mismatched) + .await + .unwrap_err(); + assert_matches!( + err, + Error::RegionExists { reason, .. } if reason.contains("primary key") + ); + } +} diff --git a/src/relational-engine/src/error.rs b/src/relational-engine/src/error.rs new file mode 100644 index 00000000000..b688bed44ae --- /dev/null +++ b/src/relational-engine/src/error.rs @@ -0,0 +1,181 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::any::Any; + +use common_error::ext::{ErrorExt, RetryHint}; +use common_error::status_code::StatusCode; +use common_macro::stack_trace_debug; +use object_store::error::retry_hint_from_opendal_error; +use serde_json::error::Error as JsonError; +use snafu::{Location, Snafu}; +use store_api::storage::RegionId; + +#[derive(Snafu)] +#[snafu(visibility(pub))] +#[stack_trace_debug] +pub enum Error { + #[snafu(display("Unsupported operation: {}", operation))] + Unsupported { + operation: String, + #[snafu(implicit)] + location: Location, + }, + + #[snafu(display("Unexpected engine: {}", engine))] + UnexpectedEngine { + engine: String, + #[snafu(implicit)] + location: Location, + }, + + #[snafu(display("Invalid region metadata"))] + InvalidMetadata { + source: store_api::metadata::MetadataError, + #[snafu(implicit)] + location: Location, + }, + + #[snafu(display("Region not found, region_id: {}", region_id))] + RegionNotFound { + region_id: RegionId, + #[snafu(implicit)] + location: Location, + }, + + #[snafu(display( + "Region already exists with a different definition, region_id: {}, reason: {}", + region_id, + reason + ))] + RegionExists { + region_id: RegionId, + reason: String, + #[snafu(implicit)] + location: Location, + }, + + #[snafu(display("Failed to check object from path: {}", path))] + CheckObject { + path: String, + #[snafu(implicit)] + location: Location, + #[snafu(source)] + error: object_store::Error, + }, + + #[snafu(display("Fail to encode object into json"))] + EncodeJson { + #[snafu(implicit)] + location: Location, + #[snafu(source)] + error: JsonError, + }, + + #[snafu(display("Fail to decode object from json"))] + DecodeJson { + #[snafu(implicit)] + location: Location, + #[snafu(source)] + error: JsonError, + }, + + #[snafu(display("Failed to store region manifest, region_id: {}", region_id))] + StoreRegionManifest { + #[snafu(source)] + error: object_store::Error, + region_id: RegionId, + #[snafu(implicit)] + location: Location, + }, + + #[snafu(display("Failed to load region manifest, region_id: {}", region_id))] + LoadRegionManifest { + #[snafu(source)] + error: object_store::Error, + region_id: RegionId, + #[snafu(implicit)] + location: Location, + }, + + #[snafu(display("Failed to delete region manifest, region_id: {},", region_id))] + DeleteRegionManifest { + #[snafu(source)] + error: object_store::Error, + region_id: RegionId, + #[snafu(implicit)] + location: Location, + }, + + #[snafu(display("Manifest already exists: {}", path))] + ManifestExists { + path: String, + #[snafu(implicit)] + location: Location, + }, + + #[snafu(display("Missing required field: {}", name))] + MissingRequiredField { + name: String, + #[snafu(implicit)] + location: Location, + }, +} + +pub type Result = std::result::Result; + +impl ErrorExt for Error { + fn status_code(&self) -> StatusCode { + use Error::*; + + match self { + // Unlike file-engine, an unsupported operation (e.g. scan) must surface as + // `Unimplemented` (StatusCode::Unsupported) so the client gets a clear + // "not implemented" signal instead of a generic "invalid arguments". + Unsupported { .. } => StatusCode::Unsupported, + + InvalidMetadata { .. } | MissingRequiredField { .. } => StatusCode::InvalidArguments, + + RegionNotFound { .. } => StatusCode::RegionNotFound, + + RegionExists { .. } => StatusCode::Unexpected, + + CheckObject { .. } + | StoreRegionManifest { .. } + | LoadRegionManifest { .. } + | DeleteRegionManifest { .. } => StatusCode::StorageUnavailable, + + EncodeJson { .. } + | DecodeJson { .. } + | ManifestExists { .. } + | UnexpectedEngine { .. } => StatusCode::Unexpected, + } + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn retry_hint(&self) -> RetryHint { + use Error::*; + + match self { + CheckObject { error, .. } + | StoreRegionManifest { error, .. } + | LoadRegionManifest { error, .. } + | DeleteRegionManifest { error, .. } => retry_hint_from_opendal_error(error), + _ => RetryHint::NonRetryable, + } + } +} diff --git a/src/relational-engine/src/lib.rs b/src/relational-engine/src/lib.rs new file mode 100644 index 00000000000..40c98529dd9 --- /dev/null +++ b/src/relational-engine/src/lib.rs @@ -0,0 +1,21 @@ +// 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. + +pub mod engine; +pub mod error; +pub(crate) mod manifest; +pub(crate) mod region; + +#[cfg(test)] +pub(crate) mod test_util; diff --git a/src/relational-engine/src/manifest.rs b/src/relational-engine/src/manifest.rs new file mode 100644 index 00000000000..ed2ddcdfda0 --- /dev/null +++ b/src/relational-engine/src/manifest.rs @@ -0,0 +1,195 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::HashMap; +use std::sync::Arc; + +use object_store::ObjectStore; +use serde::{Deserialize, Serialize}; +use snafu::{OptionExt, ResultExt, ensure}; +use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder, RegionMetadataRef}; +use store_api::storage::{ColumnId, RegionId}; +use table::requests::{RELATIONAL_META_KEY, RelationalOptions}; + +use crate::error::{ + CheckObjectSnafu, DecodeJsonSnafu, DeleteRegionManifestSnafu, EncodeJsonSnafu, + InvalidMetadataSnafu, LoadRegionManifestSnafu, ManifestExistsSnafu, MissingRequiredFieldSnafu, + Result, StoreRegionManifestSnafu, +}; + +#[inline] +fn region_manifest_path(region_dir: &str) -> String { + format!("{region_dir}manifest/_relational_manifest") +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct RelationalRegionManifest { + pub region_id: RegionId, + pub column_metadatas: Vec, + pub primary_key: Vec, + pub options: HashMap, +} + +impl RelationalRegionManifest { + pub async fn store(&self, region_dir: &str, object_store: &ObjectStore) -> Result<()> { + let path = ®ion_manifest_path(region_dir); + let exist = object_store + .exists(path) + .await + .context(CheckObjectSnafu { path })?; + ensure!(!exist, ManifestExistsSnafu { path }); + + let bs = self.encode()?; + object_store + .write(path, bs) + .await + .context(StoreRegionManifestSnafu { + region_id: self.region_id, + })?; + + Ok(()) + } + + /// Checks whether a region manifest has already been persisted under + /// `region_dir`. This is the source of truth for region existence and is + /// used to make `CREATE EXTERNAL TABLE` idempotent across restarts. + pub async fn exists(region_dir: &str, object_store: &ObjectStore) -> Result { + let path = ®ion_manifest_path(region_dir); + object_store + .exists(path) + .await + .context(CheckObjectSnafu { path }) + } + + pub async fn load( + region_id: RegionId, + region_dir: &str, + object_store: &ObjectStore, + ) -> Result { + let path = ®ion_manifest_path(region_dir); + let bs = object_store + .read(path) + .await + .context(LoadRegionManifestSnafu { region_id })? + .to_vec(); + Self::decode(bs.as_slice()) + } + + pub async fn delete( + region_id: RegionId, + region_dir: &str, + object_store: &ObjectStore, + ) -> Result<()> { + let path = ®ion_manifest_path(region_dir); + object_store + .delete(path) + .await + .context(DeleteRegionManifestSnafu { region_id }) + } + + pub fn metadata(&self) -> Result { + let mut builder = RegionMetadataBuilder::new(self.region_id); + for column in &self.column_metadatas { + builder.push_column_metadata(column.clone()); + } + builder.primary_key(self.primary_key.clone()); + let metadata = builder + .build_without_validation() + .context(InvalidMetadataSnafu)?; + + Ok(Arc::new(metadata)) + } + + pub fn relational_options(&self) -> Result { + let encoded_opts = self.get_option(RELATIONAL_META_KEY)?; + serde_json::from_str(&encoded_opts).context(DecodeJsonSnafu) + } + + fn encode(&self) -> Result> { + serde_json::to_vec(self).context(EncodeJsonSnafu) + } + + fn decode(src: &[u8]) -> Result { + serde_json::from_slice(src).context(DecodeJsonSnafu) + } + + fn get_option(&self, name: &str) -> Result { + self.options + .get(name) + .cloned() + .context(MissingRequiredFieldSnafu { name }) + } +} + +#[cfg(test)] +mod tests { + use std::assert_matches; + + use store_api::storage::RegionId; + + use super::*; + use crate::error::Error; + use crate::test_util::{new_test_column_metadata, new_test_object_store, new_test_options}; + + #[tokio::test] + async fn test_manifest_round_trip() { + let (_dir, object_store) = new_test_object_store("test_manifest_round_trip"); + + let region_id = RegionId::new(1, 0); + let manifest = RelationalRegionManifest { + region_id, + column_metadatas: new_test_column_metadata(), + primary_key: vec![1], + options: new_test_options(), + }; + let region_dir = "manifest_dir/1_0000000000/"; + + manifest.store(region_dir, &object_store).await.unwrap(); + + // Options and schema are persisted verbatim and recovered on load. + let loaded = RelationalRegionManifest::load(region_id, region_dir, &object_store) + .await + .unwrap(); + assert_eq!(loaded.options, manifest.options); + assert_eq!(loaded.region_id, region_id); + assert_eq!(loaded.primary_key, vec![1]); + assert_eq!(loaded.column_metadatas.len(), 3); + assert_eq!( + loaded.relational_options().unwrap(), + manifest.relational_options().unwrap() + ); + + // Storing the same manifest twice must fail. + let err = manifest.store(region_dir, &object_store).await.unwrap_err(); + assert_matches!(err, Error::ManifestExists { .. }); + } + + #[tokio::test] + async fn test_metadata_rebuild() { + let manifest = RelationalRegionManifest { + region_id: RegionId::new(2, 0), + column_metadatas: new_test_column_metadata(), + primary_key: vec![1], + options: new_test_options(), + }; + + let metadata = manifest.metadata().unwrap(); + assert_eq!(metadata.region_id, RegionId::new(2, 0)); + assert_eq!(metadata.primary_key, vec![1]); + assert_eq!(metadata.column_metadatas.len(), 3); + assert_eq!(metadata.column_metadatas[0].column_id, 0); + assert_eq!(metadata.column_metadatas[1].column_id, 1); + assert_eq!(metadata.column_metadatas[2].column_id, 2); + } +} diff --git a/src/relational-engine/src/region.rs b/src/relational-engine/src/region.rs new file mode 100644 index 00000000000..38dd2c369a0 --- /dev/null +++ b/src/relational-engine/src/region.rs @@ -0,0 +1,312 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::HashMap; +use std::sync::Arc; + +use object_store::ObjectStore; +use store_api::metadata::RegionMetadataRef; +use store_api::path_utils::region_name; +use store_api::region_request::{RegionCreateRequest, RegionOpenRequest}; +use store_api::storage::RegionId; + +use crate::error::{RegionExistsSnafu, Result}; +use crate::manifest::RelationalRegionManifest; + +#[derive(Debug)] +pub struct RelationalRegion { + pub(crate) region_dir: String, + pub(crate) metadata: RegionMetadataRef, + /// Frozen connection / remote table options. Kept alongside the metadata so an + /// idempotent `Create` can be validated against the existing definition. + options: HashMap, +} + +pub type RelationalRegionRef = Arc; + +impl RelationalRegion { + /// Computes the region directory under a table's data home. + pub(crate) fn region_dir(table_dir: &str, region_id: RegionId) -> String { + object_store::util::join_dir( + table_dir, + ®ion_name(region_id.table_id(), region_id.region_sequence()), + ) + } + + pub async fn create( + region_id: RegionId, + request: RegionCreateRequest, + object_store: &ObjectStore, + ) -> Result { + let options = request.options.clone(); + let manifest = RelationalRegionManifest { + region_id, + column_metadatas: request.column_metadatas.clone(), + primary_key: request.primary_key.clone(), + options: request.options, + }; + + let region_dir = Self::region_dir(&request.table_dir, region_id); + // Validate the options carry the relational metadata (connection + remote table) + // before persisting, so a malformed region fails fast here rather than on open. + manifest.relational_options()?; + let metadata = manifest.metadata()?; + + manifest.store(®ion_dir, object_store).await?; + + Ok(Arc::new(Self { + region_dir, + metadata, + options, + })) + } + + pub async fn open( + region_id: RegionId, + request: RegionOpenRequest, + object_store: &ObjectStore, + ) -> Result { + let region_dir = Self::region_dir(&request.table_dir, region_id); + Self::load(region_id, ®ion_dir, object_store).await + } + + /// Loads a region from its persisted manifest. Used both by `open` and to + /// recover a region whose manifest already exists (idempotent create). + pub(crate) async fn load( + region_id: RegionId, + region_dir: &str, + object_store: &ObjectStore, + ) -> Result { + let manifest = RelationalRegionManifest::load(region_id, region_dir, object_store).await?; + Self::from_manifest(region_dir, manifest) + } + + /// Builds a region from an already-loaded manifest, validating the + /// relational options and rebuilding the region metadata. + pub(crate) fn from_manifest( + region_dir: &str, + manifest: RelationalRegionManifest, + ) -> Result { + manifest.relational_options()?; + let metadata = manifest.metadata()?; + let options = manifest.options; + + Ok(Arc::new(Self { + region_dir: region_dir.to_string(), + metadata, + options, + })) + } + + pub async fn drop(&self, object_store: &ObjectStore) -> Result<()> { + RelationalRegionManifest::delete(self.metadata.region_id, &self.region_dir, object_store) + .await + } + + pub fn metadata(&self) -> RegionMetadataRef { + self.metadata.clone() + } + + /// Validates an idempotent `Create` against the already-open region, mirroring + /// mito2's `check_recovered_region`. A create whose definition (region id, schema, + /// primary key, or options) matches the existing region succeeds; a mismatched + /// create must fail instead of silently returning success. + pub(crate) fn check_create_request( + &self, + region_id: RegionId, + request: &RegionCreateRequest, + ) -> Result<()> { + if self.metadata.region_id != region_id { + return RegionExistsSnafu { + region_id, + reason: format!( + "existing metadata has different region id {}", + self.metadata.region_id + ), + } + .fail(); + } + if self.metadata.column_metadatas != request.column_metadatas { + return RegionExistsSnafu { + region_id, + reason: "existing metadata has different schema".to_string(), + } + .fail(); + } + if self.metadata.primary_key != request.primary_key { + return RegionExistsSnafu { + region_id, + reason: "existing metadata has different primary key".to_string(), + } + .fail(); + } + if self.options != request.options { + return RegionExistsSnafu { + region_id, + reason: "existing metadata has different options".to_string(), + } + .fail(); + } + + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use std::assert_matches; + use std::collections::HashMap; + + use store_api::region_request::{PathType, RegionCreateRequest, RegionOpenRequest}; + use store_api::storage::RegionId; + + use super::*; + use crate::error::Error; + use crate::test_util::{new_test_column_metadata, new_test_object_store, new_test_options}; + + #[tokio::test] + async fn test_create_region() { + let (_dir, object_store) = new_test_object_store("test_create_region"); + + let request = RegionCreateRequest { + engine: "relational".to_string(), + column_metadatas: new_test_column_metadata(), + primary_key: vec![1], + options: new_test_options(), + table_dir: "create_region_dir/".to_string(), + path_type: PathType::Bare, + partition_expr_json: Some("".to_string()), + requirements: Default::default(), + }; + let region_id = RegionId::new(1, 0); + + let region = RelationalRegion::create(region_id, request.clone(), &object_store) + .await + .unwrap(); + + assert_eq!(region.region_dir, "create_region_dir/1_0000000000/"); + assert_eq!(region.metadata.region_id, region_id); + assert_eq!(region.metadata.primary_key, vec![1]); + assert_eq!(region.metadata.column_metadatas.len(), 3); + + assert!( + object_store + .exists("create_region_dir/1_0000000000/manifest/_relational_manifest") + .await + .unwrap() + ); + + // Creating the same region twice must fail: the manifest already exists. + let err = RelationalRegion::create(region_id, request, &object_store) + .await + .unwrap_err(); + assert_matches!(err, Error::ManifestExists { .. }); + } + + #[tokio::test] + async fn test_open_region() { + let (_dir, object_store) = new_test_object_store("test_open_region"); + + let region_dir = "open_region_dir/".to_string(); + let request = RegionCreateRequest { + engine: "relational".to_string(), + column_metadatas: new_test_column_metadata(), + primary_key: vec![1], + options: new_test_options(), + table_dir: region_dir.clone(), + path_type: PathType::Bare, + partition_expr_json: Some("".to_string()), + requirements: Default::default(), + }; + let region_id = RegionId::new(1, 0); + + RelationalRegion::create(region_id, request, &object_store) + .await + .unwrap(); + + let request = RegionOpenRequest { + engine: "relational".to_string(), + table_dir: region_dir, + path_type: PathType::Bare, + options: HashMap::default(), + skip_wal_replay: false, + checkpoint: None, + requirements: Default::default(), + }; + + let region = RelationalRegion::open(region_id, request, &object_store) + .await + .unwrap(); + + assert_eq!(region.region_dir, "open_region_dir/1_0000000000/"); + assert_eq!(region.metadata.region_id, region_id); + assert_eq!(region.metadata.primary_key, vec![1]); + assert_eq!(region.metadata.column_metadatas.len(), 3); + } + + #[tokio::test] + async fn test_drop_region() { + let (_dir, object_store) = new_test_object_store("test_drop_region"); + + let region_dir = "drop_region_dir/".to_string(); + let request = RegionCreateRequest { + engine: "relational".to_string(), + column_metadatas: new_test_column_metadata(), + primary_key: vec![1], + options: new_test_options(), + table_dir: region_dir.clone(), + path_type: PathType::Bare, + partition_expr_json: Some("".to_string()), + requirements: Default::default(), + }; + let region_id = RegionId::new(1, 0); + + let region = RelationalRegion::create(region_id, request, &object_store) + .await + .unwrap(); + + assert!( + object_store + .exists("drop_region_dir/1_0000000000/manifest/_relational_manifest") + .await + .unwrap() + ); + + RelationalRegion::drop(®ion, &object_store) + .await + .unwrap(); + + assert!( + !object_store + .exists("drop_region_dir/1_0000000000/manifest/_relational_manifest") + .await + .unwrap() + ); + + // Re-opening a dropped region must fail. + let request = RegionOpenRequest { + engine: "relational".to_string(), + table_dir: region_dir, + path_type: PathType::Bare, + options: HashMap::default(), + skip_wal_replay: false, + checkpoint: None, + requirements: Default::default(), + }; + let err = RelationalRegion::open(region_id, request, &object_store) + .await + .unwrap_err(); + assert_matches!(err, Error::LoadRegionManifest { .. }); + } +} diff --git a/src/relational-engine/src/test_util.rs b/src/relational-engine/src/test_util.rs new file mode 100644 index 00000000000..b71512f7c57 --- /dev/null +++ b/src/relational-engine/src/test_util.rs @@ -0,0 +1,67 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::HashMap; + +use api::v1::SemanticType; +use common_test_util::temp_dir::{TempDir, create_temp_dir}; +use common_time::timestamp::TimeUnit; +use datatypes::prelude::ConcreteDataType; +use datatypes::schema::ColumnSchema; +use object_store::ObjectStore; +use object_store::services::Fs; +use store_api::metadata::ColumnMetadata; +use table::requests::{RELATIONAL_META_KEY, RelationalOptions}; + +pub(crate) fn new_test_object_store(prefix: &str) -> (TempDir, ObjectStore) { + let dir = create_temp_dir(prefix); + let store_dir = dir.path().to_string_lossy(); + let builder = Fs::default().root(&store_dir); + (dir, ObjectStore::new(builder).unwrap()) +} + +pub(crate) fn new_test_column_metadata() -> Vec { + vec![ + ColumnMetadata { + column_schema: ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_datatype(TimeUnit::Millisecond), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: 0, + }, + ColumnMetadata { + column_schema: ColumnSchema::new("host", ConcreteDataType::string_datatype(), false), + semantic_type: SemanticType::Tag, + column_id: 1, + }, + ColumnMetadata { + column_schema: ColumnSchema::new("val", ConcreteDataType::int64_datatype(), false), + semantic_type: SemanticType::Field, + column_id: 2, + }, + ] +} + +pub(crate) fn new_test_options() -> HashMap { + HashMap::from([( + RELATIONAL_META_KEY.to_string(), + serde_json::to_string(&RelationalOptions { + connection: "mysql://user:***@127.0.0.1:3306/db".to_string(), + remote_table: "remote_table".to_string(), + }) + .unwrap(), + )]) +} diff --git a/src/store-api/src/metric_engine_consts.rs b/src/store-api/src/metric_engine_consts.rs index 16d7a67773b..7c7792068fb 100644 --- a/src/store-api/src/metric_engine_consts.rs +++ b/src/store-api/src/metric_engine_consts.rs @@ -41,6 +41,8 @@ pub const METRIC_ENGINE_NAME: &str = "metric"; pub const FILE_ENGINE_NAME: &str = "file"; +pub const RELATIONAL_ENGINE_NAME: &str = "relational"; + /// Metadata key present in the `CREATE TABLE ... WITH ()` clause. This key is /// used to identify the table is a physical metric table. E.g.: /// ```sql diff --git a/src/store-api/src/region_engine.rs b/src/store-api/src/region_engine.rs index fb3c7f01b90..002ffb1691c 100644 --- a/src/store-api/src/region_engine.rs +++ b/src/store-api/src/region_engine.rs @@ -87,6 +87,7 @@ pub struct SetRegionRoleStateRequest { #[derive(Debug, PartialEq, Eq)] pub enum SetRegionRoleStateSuccess { File, + Relational, Mito { last_entry_id: entry::Id, }, @@ -102,6 +103,11 @@ impl SetRegionRoleStateSuccess { Self::File } + /// Returns a [SetRegionRoleStateSuccess::Relational]. + pub fn relational() -> Self { + Self::Relational + } + /// Returns a [SetRegionRoleStateSuccess::Mito] with the `last_entry_id`. pub fn mito(last_entry_id: entry::Id) -> Self { SetRegionRoleStateSuccess::Mito { last_entry_id } @@ -120,7 +126,7 @@ impl SetRegionRoleStateSuccess { /// Returns the last entry id of the region. pub fn last_entry_id(&self) -> Option { match self { - SetRegionRoleStateSuccess::File => None, + SetRegionRoleStateSuccess::File | SetRegionRoleStateSuccess::Relational => None, SetRegionRoleStateSuccess::Mito { last_entry_id } => Some(*last_entry_id), SetRegionRoleStateSuccess::Metric { last_entry_id, .. } => Some(*last_entry_id), } @@ -129,7 +135,7 @@ impl SetRegionRoleStateSuccess { /// Returns the last entry id of the metadata of the region. pub fn metadata_last_entry_id(&self) -> Option { match self { - SetRegionRoleStateSuccess::File => None, + SetRegionRoleStateSuccess::File | SetRegionRoleStateSuccess::Relational => None, SetRegionRoleStateSuccess::Mito { .. } => None, SetRegionRoleStateSuccess::Metric { metadata_last_entry_id, diff --git a/src/table/src/requests.rs b/src/table/src/requests.rs index 2674c2d0a1e..4cfded64fb5 100644 --- a/src/table/src/requests.rs +++ b/src/table/src/requests.rs @@ -54,10 +54,23 @@ mod semantic; pub use semantic::*; pub const FILE_TABLE_META_KEY: &str = "__private.file_table_meta"; +pub const RELATIONAL_META_KEY: &str = "__private.relational_meta"; pub const FILE_TABLE_LOCATION_KEY: &str = "location"; pub const FILE_TABLE_PATTERN_KEY: &str = "pattern"; pub const FILE_TABLE_FORMAT_KEY: &str = "format"; +/// Options captured at `CREATE EXTERNAL TABLE ... ENGINE=relational WITH (...)` time +/// and persisted in the region manifest. +/// +/// `connection` is persisted in **plaintext** so the datanode can reconnect to the +/// remote database when serving queries on the region. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct RelationalOptions { + /// Full DSN, e.g. `mysql://user:password@host:3306/db`. + pub connection: String, + pub remote_table: String, +} + pub const TABLE_DATA_MODEL: &str = "table_data_model"; pub const TABLE_DATA_MODEL_TRACE_V1: &str = "greptime_trace_v1"; /// Table data model used by the JSON2-based OTLP trace pipeline. From c035670ec03be0c0ec9f725c3f7d86645efdee8c Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Fri, 30 Oct 2026 07:25:55 +0000 Subject: [PATCH 146/239] feat: share Puffin metadata and index page loads (#9792) Signed-off-by: WenyXu --- Cargo.lock | 2 + src/common/cache/src/paged.rs | 17 +- src/common/cache/src/paged/tests.rs | 6 + src/index/src/bloom_filter/error.rs | 26 + src/index/src/inverted_index/error.rs | 26 + src/mito2/src/cache.rs | 83 ++- src/mito2/src/cache/index.rs | 666 +++++++++++------- .../src/cache/index/bloom_filter_index.rs | 264 +++++-- src/mito2/src/cache/index/inverted_index.rs | 228 ++++-- src/mito2/src/metrics.rs | 168 +++++ .../src/sst/index/inverted_index/creator.rs | 3 +- src/puffin/Cargo.toml | 2 + src/puffin/src/puffin_manager/cache.rs | 216 +++++- .../fs_puffin_manager/reader.rs | 35 +- src/puffin/src/puffin_manager/tests.rs | 165 +++++ 15 files changed, 1464 insertions(+), 443 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 6b7bd35e546..a87b99721ff 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11908,6 +11908,7 @@ dependencies = [ "bitflags 2.12.1", "bytes", "common-base", + "common-cache", "common-error", "common-macro", "common-runtime", @@ -11925,6 +11926,7 @@ dependencies = [ "snafu 0.8.6", "tokio", "tokio-util", + "tracing", "uuid", ] diff --git a/src/common/cache/src/paged.rs b/src/common/cache/src/paged.rs index f2e8a659ddf..3131329dc04 100644 --- a/src/common/cache/src/paged.rs +++ b/src/common/cache/src/paged.rs @@ -31,7 +31,7 @@ use common_base::range_read::RangeReader; use futures::channel::{mpsc, oneshot}; use futures::{StreamExt, future}; -use crate::metrics::{CacheMetrics, CacheSizeMetric}; +use crate::metrics::{CacheMetrics, CacheSizeMetric, CacheSizeRegistration}; use crate::page::PageLayout; use crate::{Cache, RemovalCause}; @@ -41,8 +41,6 @@ pub enum PageEvent { Lookup { hits: usize, misses: usize }, /// One successful native page loader, excluding waiters and hits. Loaded { bytes: u64 }, - /// Backend weight, not process RSS; sampled after fetch/maintenance. - Residency { usage: u64, capacity: u64 }, /// Backend eviction, excluding unadmitted values released by their owners. Evicted { cause: RemovalCause }, } @@ -51,6 +49,7 @@ pub enum PageEvent { #[derive(Clone)] pub struct PageCacheObserver { pub bytes: CacheSizeMetric, + pub capacity: CacheSizeMetric, pub loads: Option>, pub events: Arc, } @@ -78,6 +77,7 @@ pub struct PagedCache { cache: Cache<(K, u64), Arc>, page_size: u64, observer: PageCacheObserver, + _capacity_registration: CacheSizeRegistration, } impl PagedCache { @@ -102,7 +102,9 @@ impl PagedCache { if let Some(metrics) = &observer.loads { cache = cache.with_metrics(metrics.clone()); } + let capacity_registration = observer.capacity.register(move || capacity); Self { + _capacity_registration: capacity_registration, cache, page_size, observer, @@ -139,7 +141,6 @@ impl PagedCache { /// Applies pending cache admission and eviction work. pub async fn run_pending_tasks(&self) { self.cache.run_pending_tasks().await; - self.observe_residency(); } /// Fetches bytes in input order. The reader must address the same immutable @@ -329,7 +330,6 @@ impl PagedCache { } output.push(assembled.freeze()); } - self.observe_residency(); Ok((output, stats)) } @@ -339,13 +339,6 @@ impl PagedCache { }); Arc::new(CachedPage { file_size, bytes }) } - - fn observe_residency(&self) { - (self.observer.events)(PageEvent::Residency { - usage: self.cache.weighted_size(), - capacity: self.cache.capacity(), - }); - } } struct CachedPage { diff --git a/src/common/cache/src/paged/tests.rs b/src/common/cache/src/paged/tests.rs index 556455ebfb1..7d729f8733f 100644 --- a/src/common/cache/src/paged/tests.rs +++ b/src/common/cache/src/paged/tests.rs @@ -100,6 +100,12 @@ fn cache() -> (PagedCache, Arc) { let events = observation.clone(); let observer = PageCacheObserver { bytes: observation.bytes.bind(&[]), + capacity: crate::metrics::CacheSizeMetrics::new( + prometheus::Opts::new("test_page_capacity_bytes", "test"), + &[], + ) + .unwrap() + .bind(&[]), loads: None, events: Arc::new(move |event| { if let PageEvent::Loaded { .. } = event { diff --git a/src/index/src/bloom_filter/error.rs b/src/index/src/bloom_filter/error.rs index 93b2ab19933..f248fc07b0f 100644 --- a/src/index/src/bloom_filter/error.rs +++ b/src/index/src/bloom_filter/error.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::any::Any; +use std::sync::Arc; use common_error::ext::{BoxedError, ErrorExt}; use common_error::status_code::StatusCode; @@ -68,6 +69,13 @@ pub enum Error { location: Location, }, + #[snafu(display("Shared index metadata read failed"))] + SharedRead { + source: Arc, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("External error"))] External { source: BoxedError, @@ -89,6 +97,7 @@ impl ErrorExt for Error { Intermediate { source, .. } => source.status_code(), External { source, .. } => source.status_code(), + SharedRead { source, .. } => source.status_code(), } } @@ -98,3 +107,20 @@ impl ErrorExt for Error { } pub type Result = std::result::Result; + +#[cfg(test)] +mod tests { + use snafu::IntoError; + + use super::*; + + #[test] + fn shared_metadata_error_preserves_status() { + let source = FileSizeTooSmallSnafu { size: 0_u64 }.build(); + let status = source.status_code(); + let retry = source.retry_hint(); + let error = SharedReadSnafu.into_error(Arc::new(source)); + assert_eq!(error.status_code(), status); + assert_eq!(error.retry_hint(), retry); + } +} diff --git a/src/index/src/inverted_index/error.rs b/src/index/src/inverted_index/error.rs index f9007fbbe40..dd8f8b6016d 100644 --- a/src/index/src/inverted_index/error.rs +++ b/src/index/src/inverted_index/error.rs @@ -14,6 +14,7 @@ use std::any::Any; use std::io::Error as IoError; +use std::sync::Arc; use common_error::ext::{BoxedError, ErrorExt}; use common_error::status_code::StatusCode; @@ -221,6 +222,13 @@ pub enum Error { expected_row_count: usize, }, + #[snafu(display("Shared index metadata read failed"))] + SharedRead { + source: Arc, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("External error"))] External { source: BoxedError, @@ -269,6 +277,7 @@ impl ErrorExt for Error { Intermediate { source, .. } => source.status_code(), External { source, .. } => source.status_code(), + SharedRead { source, .. } => source.status_code(), } } @@ -278,3 +287,20 @@ impl ErrorExt for Error { } pub type Result = std::result::Result; + +#[cfg(test)] +mod tests { + use snafu::IntoError; + + use super::*; + + #[test] + fn shared_metadata_error_preserves_status() { + let source = CommonIoSnafu.into_error(std::io::Error::other("read failed")); + let status = source.status_code(); + let retry = source.retry_hint(); + let error = SharedReadSnafu.into_error(Arc::new(source)); + assert_eq!(error.status_code(), status); + assert_eq!(error.retry_hint(), retry); + } +} diff --git a/src/mito2/src/cache.rs b/src/mito2/src/cache.rs index 2d63a09329f..1b3645eff62 100644 --- a/src/mito2/src/cache.rs +++ b/src/mito2/src/cache.rs @@ -1522,7 +1522,7 @@ impl CacheManager { } if let Some(cache) = &self.puffin_metadata_cache { - cache.remove(&file_id.to_string()); + cache.remove(&file_id.to_string()).await; } if let Some(write_cache) = &self.write_cache { @@ -1809,7 +1809,8 @@ impl CacheManagerBuilder { let prefilter_result_cache = (self.prefilter_result_cache_size != 0) .then(|| new_prefilter_result_cache(self.prefilter_result_cache_size)); let puffin_metadata_cache = - PuffinMetadataCache::new(self.puffin_metadata_size, CACHE_BYTES.gauges()); + PuffinMetadataCache::new(self.puffin_metadata_size, &CACHE_BYTES) + .with_load_metrics(crate::metrics::cache_loading_metrics("puffin_metadata")); let selector_result_cache = (self.selector_result_cache_size != 0).then(|| { Cache::builder() .max_capacity(self.selector_result_cache_size) @@ -3604,11 +3605,15 @@ mod tests { Arc::new(::index::target::IndexTarget::ColumnId(column_id)), Tag::Skipping, ); - bloom_cache.put_metadata(bloom_key.clone(), Arc::new(BloomFilterMeta::default())); - inverted_cache.put_metadata( - (index_id.file_id(), index_id.version), - Arc::new(InvertedIndexMetas::default()), - ); + bloom_cache + .put_metadata(bloom_key.clone(), Arc::new(BloomFilterMeta::default())) + .await; + inverted_cache + .put_metadata( + (index_id.file_id(), index_id.version), + Arc::new(InvertedIndexMetas::default()), + ) + .await; let predicate = PredicateKey::new_bloom(Arc::new(BTreeMap::new())); let selection = Arc::new(RowGroupSelection::default()); result_cache.put(predicate.clone(), index_id.file_id(), selection); @@ -3617,58 +3622,84 @@ mod tests { blobs: Vec::new(), properties: HashMap::new(), }); - puffin_metadata_cache.put_metadata(file_id_str.clone(), metadata); + puffin_metadata_cache + .put_metadata(file_id_str.clone(), metadata) + .await; - assert!(bloom_cache.get_metadata(bloom_key.clone()).is_some()); + assert!(bloom_cache.get_metadata(bloom_key.clone()).await.is_some()); assert!( inverted_cache .get_metadata((index_id.file_id(), index_id.version)) + .await .is_some() ); assert!(result_cache.get(&predicate, index_id.file_id()).is_some()); - assert!(puffin_metadata_cache.get_metadata(&file_id_str).is_some()); + assert!( + puffin_metadata_cache + .get_metadata(&file_id_str) + .await + .is_some() + ); cache.evict_puffin_cache(index_id).await; - assert!(bloom_cache.get_metadata(bloom_key.clone()).is_none()); + assert!(bloom_cache.get_metadata(bloom_key.clone()).await.is_none()); assert!( inverted_cache .get_metadata((index_id.file_id(), index_id.version)) + .await .is_none() ); assert!(result_cache.get(&predicate, index_id.file_id()).is_none()); - assert!(puffin_metadata_cache.get_metadata(&file_id_str).is_none()); + assert!( + puffin_metadata_cache + .get_metadata(&file_id_str) + .await + .is_none() + ); // Refill caches and evict via CacheStrategy to ensure delegation works. - bloom_cache.put_metadata(bloom_key.clone(), Arc::new(BloomFilterMeta::default())); - inverted_cache.put_metadata( - (index_id.file_id(), index_id.version), - Arc::new(InvertedIndexMetas::default()), - ); + bloom_cache + .put_metadata(bloom_key.clone(), Arc::new(BloomFilterMeta::default())) + .await; + inverted_cache + .put_metadata( + (index_id.file_id(), index_id.version), + Arc::new(InvertedIndexMetas::default()), + ) + .await; result_cache.put( predicate.clone(), index_id.file_id(), Arc::new(RowGroupSelection::default()), ); - puffin_metadata_cache.put_metadata( - file_id_str.clone(), - Arc::new(FileMetadata { - blobs: Vec::new(), - properties: HashMap::new(), - }), - ); + puffin_metadata_cache + .put_metadata( + file_id_str.clone(), + Arc::new(FileMetadata { + blobs: Vec::new(), + properties: HashMap::new(), + }), + ) + .await; let strategy = CacheStrategy::EnableAll(cache.clone()); strategy.evict_puffin_cache(index_id).await; - assert!(bloom_cache.get_metadata(bloom_key.clone()).is_none()); + assert!(bloom_cache.get_metadata(bloom_key.clone()).await.is_none()); assert!( inverted_cache .get_metadata((index_id.file_id(), index_id.version)) + .await .is_none() ); assert!(result_cache.get(&predicate, index_id.file_id()).is_none()); - assert!(puffin_metadata_cache.get_metadata(&file_id_str).is_none()); + assert!( + puffin_metadata_cache + .get_metadata(&file_id_str) + .await + .is_none() + ); } fn wide_region_metadata(column_count: u32) -> RegionMetadata { diff --git a/src/mito2/src/cache/index.rs b/src/mito2/src/cache/index.rs index ab813c0353d..5f9208f006f 100644 --- a/src/mito2/src/cache/index.rs +++ b/src/mito2/src/cache/index.rs @@ -16,16 +16,18 @@ pub mod bloom_filter_index; pub mod inverted_index; pub mod result_cache; -use std::collections::{HashMap, HashSet}; use std::future::Future; use std::hash::Hash; +use std::io; use std::ops::Range; use std::sync::Arc; +use async_trait::async_trait; use bytes::Bytes; -use object_store::Buffer; +use common_base::range_read::{Metadata, RangeReader}; +use common_cache::paged::PagedCache; -use crate::metrics::{CACHE_BYTES, CACHE_HIT, CACHE_MISS}; +use crate::metrics::CACHE_BYTES; /// Metrics for index metadata. const INDEX_METADATA_TYPE: &str = "index_metadata"; @@ -45,65 +47,18 @@ pub struct IndexCacheMetrics { pub page_bytes: u64, } -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] -pub struct PageKey { - page_id: u64, -} - -impl PageKey { - /// Converts an offset to a page ID based on the page size. - fn calculate_page_id(offset: u64, page_size: u64) -> u64 { - offset / page_size - } - - /// Calculates the total number of pages that a given size spans, starting from a specific offset. - fn calculate_page_count(offset: u64, size: u32, page_size: u64) -> u32 { - let start_page = Self::calculate_page_id(offset, page_size); - let end_page = Self::calculate_page_id(offset + (size as u64) - 1, page_size); - (end_page + 1 - start_page) as u32 - } - - /// Calculates the byte range for data retrieval based on the specified offset and size. - /// - /// This function determines the starting and ending byte positions required for reading data. - /// For example, with an offset of 5000 and a size of 5000, using a PAGE_SIZE of 4096, - /// the resulting byte range will be 904..5904. This indicates that: - /// - The reader will first access fixed-size pages [4096, 8192) and [8192, 12288). - /// - To read the range [5000..10000), it only needs to fetch bytes within the range [904, 5904) across two pages. - fn calculate_range(offset: u64, size: u32, page_size: u64) -> Range { - let start = (offset % page_size) as usize; - let end = start + size as usize; - start..end - } - - /// Generates a iterator of `IndexKey` for the pages that a given offset and size span. - fn generate_page_keys(offset: u64, size: u32, page_size: u64) -> impl Iterator { - let start_page = Self::calculate_page_id(offset, page_size); - let total_pages = Self::calculate_page_count(offset, size, page_size); - (0..total_pages).map(move |i| Self { - page_id: start_page + i as u64, - }) - } -} - /// Cache for index metadata and content. pub struct IndexCache { /// Cache for index metadata - index_metadata: moka::sync::Cache>, + index_metadata: common_cache::Cache>, /// Cache for index content. - index: moka::sync::Cache<(K, PageKey), Bytes>, - // Page size for index content. - page_size: u64, - - /// Weighter for metadata. - weight_of_metadata: fn(&K, &Arc) -> u32, - /// Weighter for content. - weight_of_content: fn(&(K, PageKey), &Bytes) -> u32, + index: PagedCache, + index_type: &'static str, } impl IndexCache where - K: Hash + Eq + Send + Sync + 'static, + K: Hash + Eq + Clone + Send + Sync + 'static, M: Send + Sync + 'static, { pub fn new_with_weighter( @@ -111,40 +66,34 @@ where index_content_cap: u64, page_size: u64, index_type: &'static str, + metadata_type: &'static str, weight_of_metadata: fn(&K, &Arc) -> u32, - weight_of_content: fn(&(K, PageKey), &Bytes) -> u32, + weight_of_content: fn(&(K, u64), &Bytes) -> u32, ) -> Self { common_telemetry::debug!( "Building IndexCache with metadata size: {index_metadata_cap}, content size: {index_content_cap}, page size: {page_size}, index type: {index_type}" ); - let index_metadata = moka::sync::CacheBuilder::new(index_metadata_cap) - .name(&format!("index_metadata_{}", index_type)) - .weigher(weight_of_metadata) - .eviction_listener(move |k, v, _cause| { - let size = weight_of_metadata(&k, &v); - CACHE_BYTES - .with_label_values(&[INDEX_METADATA_TYPE]) - .sub(size.into()); - }) - .support_invalidation_closures() - .build(); - let index_cache = moka::sync::CacheBuilder::new(index_content_cap) - .name(&format!("index_content_{}", index_type)) - .weigher(weight_of_content) - .eviction_listener(move |k, v, _cause| { - let size = weight_of_content(&k, &v); - CACHE_BYTES - .with_label_values(&[INDEX_CONTENT_TYPE]) - .sub(size.into()); - }) - .support_invalidation_closures() - .build(); + let index_metadata = common_cache::Cache::new( + &format!("index_metadata_{index_type}"), + index_metadata_cap, + weight_of_metadata, + |_, _, _| {}, + ) + .with_size_metric(CACHE_BYTES.bind(&[INDEX_METADATA_TYPE])) + .with_metrics(crate::metrics::cache_loading_metrics(metadata_type)); + let index_cache = PagedCache::new( + &format!("index_content_{index_type}"), + index_content_cap, + page_size, + weight_of_content, + crate::metrics::cache_page_observer(index_type, INDEX_CONTENT_TYPE, |hits, misses| { + (hits, misses) + }), + ); Self { index_metadata, index: index_cache, - page_size, - weight_of_content, - weight_of_metadata, + index_type, } } } @@ -154,19 +103,31 @@ where K: Hash + Eq + Clone + Send + Sync + 'static, M: Send + Sync + 'static, { - pub fn get_metadata(&self, key: K) -> Option> { - self.index_metadata.get(&key) + /// Looks up completed metadata without starting or waiting for a load. + pub async fn get_metadata(&self, key: K) -> Option> { + self.index_metadata.get(&key).await } - pub fn put_metadata(&self, key: K, metadata: Arc) { - CACHE_BYTES - .with_label_values(&[INDEX_METADATA_TYPE]) - .add((self.weight_of_metadata)(&key, &metadata).into()); - self.index_metadata.insert(key, metadata) + /// Inserts completed metadata for cache eviction and accounting tests. + #[cfg(test)] + pub async fn put_metadata(&self, key: K, metadata: Arc) { + self.index_metadata.insert(key, metadata).await; } - /// Gets given range of index data from cache, and loads from source if the file - /// is not already cached. + /// Shares loads after the caller's initial lookup missed. Native request counts + /// describe these submissions, including waiters, rather than completed hits. + async fn get_or_load_metadata(&self, key: K, load: F) -> Result, Arc> + where + F: Future, E>>, + E: Send + Sync + 'static, + { + if self.index_metadata.capacity() == 0 { + return load.await.map_err(Arc::new); + } + self.index_metadata.try_get_with(&key, load).await + } + + /// Fetches one byte range through fixed-size pages. async fn get_or_load( &self, key: K, @@ -174,157 +135,52 @@ where offset: u64, size: u32, load: F, - ) -> Result<(Vec, IndexCacheMetrics), E> + ) -> Result<(Bytes, IndexCacheMetrics), Arc> where - F: Fn(Vec>) -> Fut, - Fut: Future, E>>, - E: std::error::Error, + F: Fn(Vec>) -> Fut + Send + Sync, + Fut: Future, E>> + Send, + E: std::error::Error + Send + Sync + 'static, { - let mut metrics = IndexCacheMetrics::default(); - let page_keys = - PageKey::generate_page_keys(offset, size, self.page_size).collect::>(); - // Size is 0, return empty data. - if page_keys.is_empty() { - return Ok((Vec::new(), metrics)); - } - metrics.num_pages = page_keys.len(); - let mut data = Vec::with_capacity(page_keys.len()); - data.resize(page_keys.len(), Bytes::new()); - let mut cache_miss_range = vec![]; - let mut cache_miss_idx = vec![]; - let last_index = page_keys.len() - 1; - // TODO: Avoid copy as much as possible. - for (i, page_key) in page_keys.iter().enumerate() { - match self.get_page(key.clone(), *page_key) { - Some(page) => { - CACHE_HIT.with_label_values(&[INDEX_CONTENT_TYPE]).inc(); - metrics.cache_hit += 1; - metrics.page_bytes += page.len() as u64; - data[i] = page; - } - None => { - CACHE_MISS.with_label_values(&[INDEX_CONTENT_TYPE]).inc(); - metrics.cache_miss += 1; - let base_offset = page_key.page_id * self.page_size; - let pruned_size = if i == last_index { - prune_size(page_keys.iter(), file_size, self.page_size) - } else { - self.page_size - }; - cache_miss_range.push(base_offset..base_offset + pruned_size); - cache_miss_idx.push(i); - } - } - } - if !cache_miss_range.is_empty() { - let pages = load(cache_miss_range).await?; - for (i, page) in cache_miss_idx.into_iter().zip(pages) { - let page_key = page_keys[i]; - metrics.page_bytes += page.len() as u64; - data[i] = page.clone(); - self.put_page(key.clone(), page_key, page.clone()); - } - } - let buffer = Buffer::from_iter(data); - Ok(( - buffer - .slice(PageKey::calculate_range(offset, size, self.page_size)) - .to_vec(), - metrics, - )) + let end = offset.checked_add(u64::from(size)).ok_or_else(|| { + Arc::new(io::Error::new( + io::ErrorKind::InvalidInput, + "Index range overflow", + )) + })?; + let (pages, metrics) = self + .get_or_load_vec(key, file_size, std::slice::from_ref(&(offset..end)), load) + .await?; + Ok((pages.into_iter().next().unwrap_or_default(), metrics)) } - /// Like [`IndexCache::get_or_load`] for several ranges, but loads the missing pages of - /// all ranges with a single `load` call so a cold read costs one round trip instead of - /// one per range. + /// Shares missing page loads while retaining legacy unique-page metrics. async fn get_or_load_vec( &self, key: K, file_size: u64, ranges: &[Range], load: F, - ) -> Result<(Vec, IndexCacheMetrics), E> + ) -> Result<(Vec, IndexCacheMetrics), Arc> where - F: FnOnce(Vec>) -> Fut, - Fut: Future, E>>, - E: std::error::Error, + F: Fn(Vec>) -> Fut + Send + Sync, + Fut: Future, E>> + Send, + E: std::error::Error + Send + Sync + 'static, { - let mut metrics = IndexCacheMetrics::default(); - let mut pages: HashMap = HashMap::new(); - let mut missing: Vec = Vec::new(); - let mut seen = HashSet::new(); - for range in ranges.iter().filter(|r| r.end > r.start) { - let size = (range.end - range.start) as u32; - for page_key in PageKey::generate_page_keys(range.start, size, self.page_size) { - if !seen.insert(page_key.page_id) { - continue; - } - metrics.num_pages += 1; - match self.get_page(key.clone(), page_key) { - Some(page) => { - CACHE_HIT.with_label_values(&[INDEX_CONTENT_TYPE]).inc(); - metrics.cache_hit += 1; - metrics.page_bytes += page.len() as u64; - pages.insert(page_key.page_id, page); - } - None => { - CACHE_MISS.with_label_values(&[INDEX_CONTENT_TYPE]).inc(); - metrics.cache_miss += 1; - missing.push(page_key.page_id); - } - } - } - } - - if !missing.is_empty() { - missing.sort_unstable(); - let load_ranges = missing - .iter() - .map(|page_id| { - let start = page_id * self.page_size; - start..(start + self.page_size).min(file_size) - }) - .collect(); - let loaded = load(load_ranges).await?; - debug_assert_eq!(loaded.len(), missing.len()); - for (page_id, page) in missing.into_iter().zip(loaded) { - metrics.page_bytes += page.len() as u64; - self.put_page(key.clone(), PageKey { page_id }, page.clone()); - pages.insert(page_id, page); - } - } - - let result = ranges - .iter() - .map(|range| { - let size = (range.end - range.start) as u32; - if size == 0 { - return Bytes::new(); - } - let buffer = Buffer::from_iter( - PageKey::generate_page_keys(range.start, size, self.page_size) - .map(|p| pages[&p.page_id].clone()), - ); - buffer - .slice(PageKey::calculate_range(range.start, size, self.page_size)) - .to_bytes() - }) - .collect(); - Ok((result, metrics)) - } - - fn get_page(&self, key: K, page_key: PageKey) -> Option { - self.index.get(&(key, page_key)) - } - - fn put_page(&self, key: K, page_key: PageKey, value: Bytes) { - // Clones the value to ensure it doesn't reference a larger buffer. - let value = Bytes::from(value.to_vec()); - let key = (key, page_key); - CACHE_BYTES - .with_label_values(&[INDEX_CONTENT_TYPE]) - .add((self.weight_of_content)(&key, &value).into()); - self.index.insert(key, value); + let reader = IndexPageReader { + load: &load, + file_size, + }; + let (pages, stats) = self.index.fetch(key, file_size, ranges, &reader).await?; + crate::metrics::cache_page_fetch_metrics(self.index_type, &stats); + Ok(( + pages, + IndexCacheMetrics { + cache_hit: stats.cache_hit, + cache_miss: stats.cache_miss, + num_pages: stats.cache_hit + stats.cache_miss, + page_bytes: stats.aligned_bytes, + }, + )) } /// Invalidates all cache entries whose keys satisfy `predicate`. @@ -336,24 +192,354 @@ where let metadata_predicate = Arc::clone(&predicate); self.index_metadata - .invalidate_entries_if(move |key, _| metadata_predicate(key)) + .invalidate_if(move |key, _| metadata_predicate(key)) .expect("cache should support invalidation closures"); self.index - .invalidate_entries_if(move |(key, _), _| predicate(key)) + .invalidate_if(move |key| predicate(key)) .expect("cache should support invalidation closures"); } } -/// Prunes the size of the last page based on the indexes. -/// We have following cases: -/// 1. The rest file size is less than the page size, read to the end of the file. -/// 2. Otherwise, read the page size. -fn prune_size<'a>( - indexes: impl Iterator, +/// Adapts borrowed index readers to the common range-reader interface. +struct IndexPageReader<'a, F> { + load: &'a F, file_size: u64, - page_size: u64, -) -> u64 { - let last_page_start = indexes.last().map(|i| i.page_id * page_size).unwrap_or(0); - page_size.min(file_size - last_page_start) +} + +#[async_trait] +impl RangeReader for IndexPageReader<'_, F> +where + F: Fn(Vec>) -> Fut + Send + Sync, + Fut: Future, E>> + Send, + E: std::error::Error + Send + Sync + 'static, +{ + async fn metadata(&self) -> io::Result { + Ok(Metadata { + content_length: self.file_size, + }) + } + + async fn read(&self, range: Range) -> io::Result { + let mut reads = self.read_vec(&[range]).await?; + if reads.len() != 1 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "Index reader returned an invalid result count", + )); + } + Ok(reads.remove(0)) + } + + async fn read_vec(&self, ranges: &[Range]) -> io::Result> { + (self.load)(ranges.to_vec()) + .await + .map_err(|error| io::Error::other(Arc::new(error))) + } +} + +/// Finds a typed shared error through the range-reader I/O wrappers. +fn find_shared_page_source( + mut error: &io::Error, +) -> Option> { + loop { + let source = error.get_ref()?; + if let Some(source) = source.downcast_ref::>() { + return Some(source.clone()); + } + if let Some(nested) = source.downcast_ref::>() { + error = nested; + } else { + error = source.downcast_ref::()?; + } + } +} + +#[cfg(test)] +mod tests { + use tokio::sync::Notify; + + use super::*; + + fn metadata_cache() -> IndexCache { + IndexCache::new_with_weighter( + 1024, + 1024, + 4, + "test", + "test_index_metadata", + |_, _| 8, + |_, bytes| bytes.len() as u32, + ) + } + + #[tokio::test] + async fn metadata_shares_same_key_without_blocking_other_keys() { + let cache = metadata_cache(); + let release = Notify::new(); + let mut owner = Box::pin(cache.get_or_load_metadata(1, async { + release.notified().await; + Ok::<_, std::io::Error>(Arc::new(11)) + })); + assert!(futures::poll!(&mut owner).is_pending()); + let mut waiter = Box::pin(cache.get_or_load_metadata(1, async { + Err::, _>(std::io::Error::other("unexpected duplicate load")) + })); + assert!(futures::poll!(&mut waiter).is_pending()); + let other = cache + .get_or_load_metadata(2, async { Ok::<_, std::io::Error>(Arc::new(33)) }) + .await + .unwrap(); + assert_eq!(*other, 33); + release.notify_one(); + let owner = owner.await.unwrap(); + let waiter = waiter.await.unwrap(); + assert!(Arc::ptr_eq(&owner, &waiter)); + assert_eq!(*cache.get_metadata(1).await.unwrap(), 11); + } + + #[tokio::test] + async fn metadata_failure_is_shared_but_not_cached() { + let cache = metadata_cache(); + let release = Notify::new(); + let mut owner = Box::pin(cache.get_or_load_metadata(1, async { + release.notified().await; + Err::, _>(std::io::Error::other("metadata read failed")) + })); + assert!(futures::poll!(&mut owner).is_pending()); + let mut waiter = Box::pin( + cache.get_or_load_metadata(1, async { Ok::<_, std::io::Error>(Arc::new(22)) }), + ); + assert!(futures::poll!(&mut waiter).is_pending()); + release.notify_one(); + let error = owner.await.unwrap_err(); + let shared = waiter.await.unwrap_err(); + assert!(Arc::ptr_eq(&error, &shared)); + assert!(cache.get_metadata(1).await.is_none()); + let retried = cache + .get_or_load_metadata(1, async { Ok::<_, std::io::Error>(Arc::new(33)) }) + .await + .unwrap(); + assert_eq!(*retried, 33); + } + + #[tokio::test] + async fn cancelled_metadata_loader_allows_waiter_to_take_over() { + let cache = metadata_cache(); + let mut owner = Box::pin(cache.get_or_load_metadata( + 1, + std::future::pending::, std::io::Error>>(), + )); + assert!(futures::poll!(&mut owner).is_pending()); + let mut waiter = Box::pin( + cache.get_or_load_metadata(1, async { Ok::<_, std::io::Error>(Arc::new(22)) }), + ); + assert!(futures::poll!(&mut waiter).is_pending()); + drop(owner); + let value = tokio::time::timeout(std::time::Duration::from_secs(1), waiter) + .await + .unwrap() + .unwrap(); + assert_eq!(*value, 22); + } + + #[tokio::test] + async fn metadata_invalidation_releases_charge_with_live_reader() { + let cache = metadata_cache(); + let value = Arc::new(11); + cache.index_metadata.insert(1, value.clone()).await; + cache.put_metadata(2, Arc::new(22)).await; + cache.index_metadata.run_pending_tasks().await; + assert_eq!(cache.index_metadata.entry_count(), 2); + cache.invalidate_if(|key| *key == 1); + assert!(cache.get_metadata(1).await.is_none()); + assert_eq!(*cache.get_metadata(2).await.unwrap(), 22); + cache.index_metadata.run_pending_tasks().await; + assert_eq!(cache.index_metadata.entry_count(), 1); + assert_eq!(*value, 11); + drop(value); + assert_eq!(cache.index_metadata.entry_count(), 1); + } + + #[tokio::test] + async fn disabled_metadata_cache_does_not_share_loads() { + let cache = IndexCache::new_with_weighter( + 0, + 1024, + 4, + "test_disabled", + "test_disabled_index_metadata", + |_: &u64, _: &Arc| 8, + |_, bytes| bytes.len() as u32, + ); + let mut owner = Box::pin(cache.get_or_load_metadata( + 1, + std::future::pending::, std::io::Error>>(), + )); + assert!(futures::poll!(&mut owner).is_pending()); + let mut independent = Box::pin( + cache.get_or_load_metadata(1, async { Ok::<_, std::io::Error>(Arc::new(22)) }), + ); + assert!(futures::poll!(&mut independent).is_ready()); + } + + #[tokio::test] + async fn metadata_metrics_distinguish_requests_from_loads() { + let metrics = crate::metrics::cache_loading_metrics("test_index_metadata_metrics"); + let mut cache = metadata_cache(); + cache.index_metadata = cache.index_metadata.with_metrics(metrics.clone()); + let release = Notify::new(); + let mut owner = Box::pin(cache.get_or_load_metadata(1, async { + release.notified().await; + Ok::<_, std::io::Error>(Arc::new(11)) + })); + let mut waiter = Box::pin(cache.get_or_load_metadata(1, async { + Err::, _>(std::io::Error::other("duplicate metadata load")) + })); + assert!(futures::poll!(&mut owner).is_pending()); + assert!(futures::poll!(&mut waiter).is_pending()); + assert_eq!(metrics.requests.active.get(), 2); + assert_eq!(metrics.loads.active.get(), 1); + release.notify_one(); + let (owner, waiter) = tokio::join!(owner, waiter); + assert!(Arc::ptr_eq(&owner.unwrap(), &waiter.unwrap())); + assert_eq!(metrics.requests.succeeded.get(), 2); + assert_eq!(metrics.loads.succeeded.get(), 1); + assert_eq!(metrics.requests.active.get(), 0); + assert_eq!(metrics.loads.active.get(), 0); + } + #[tokio::test] + async fn overlapping_page_requests_share_bytes_and_preserve_legacy_metrics() { + let cache = metadata_cache(); + let release = Notify::new(); + let reads = std::sync::Mutex::new(Vec::new()); + let load = |ranges: Vec>| { + reads.lock().unwrap().push(ranges.clone()); + let release = &release; + async move { + release.notified().await; + Ok::<_, io::Error>( + ranges + .iter() + .map(|range| { + Bytes::from(range.clone().map(|v| v as u8).collect::>()) + }) + .collect(), + ) + } + }; + let first_ranges = std::slice::from_ref(&(0..8)); + let second_ranges = std::slice::from_ref(&(4..12)); + let mut first = Box::pin(cache.get_or_load_vec(1, 12, first_ranges, &load)); + assert!(futures::poll!(&mut first).is_pending()); + let mut second = Box::pin(cache.get_or_load_vec(1, 12, second_ranges, &load)); + assert!(futures::poll!(&mut second).is_pending()); + release.notify_waiters(); + let (first, second) = tokio::join!(first, second); + let (first, first_metrics) = first.unwrap(); + let (second, second_metrics) = second.unwrap(); + assert_eq!(first, vec![Bytes::from_static(&[0, 1, 2, 3, 4, 5, 6, 7])]); + assert_eq!( + second, + vec![Bytes::from_static(&[4, 5, 6, 7, 8, 9, 10, 11])] + ); + for metrics in [first_metrics, second_metrics] { + assert_eq!( + ( + metrics.cache_hit, + metrics.cache_miss, + metrics.num_pages, + metrics.page_bytes + ), + (0, 2, 2, 8) + ); + } + let batches = reads.lock().unwrap().clone(); + assert_eq!(batches, vec![vec![0..4, 4..8], vec![8..12]]); + let mut physical = batches.into_iter().flatten().collect::>(); + physical.sort_by_key(|range| range.start); + assert!(physical.windows(2).all(|pair| pair[0].end <= pair[1].start)); + assert_eq!( + physical + .iter() + .map(|range| range.end - range.start) + .sum::(), + 12 + ); + let (hot, metrics) = cache + .get_or_load_vec(1, 12, &[1..6, 4..12], |_| async { + Err::, _>(io::Error::other("unexpected hot read")) + }) + .await + .unwrap(); + assert_eq!(hot[0], Bytes::from_static(&[1, 2, 3, 4, 5])); + assert_eq!( + ( + metrics.cache_hit, + metrics.cache_miss, + metrics.num_pages, + metrics.page_bytes + ), + (3, 0, 3, 12) + ); + cache.invalidate_if(|key| *key == 1); + assert!(!cache.index.contains(&1, 0).await); + } + + #[tokio::test] + async fn index_pages_preserve_shared_typed_read_errors_and_retry() { + use common_error::ext::ErrorExt; + use index::inverted_index::error::{Error, IndexNotFoundSnafu}; + let cache = metadata_cache(); + let error = cache + .get_or_load_vec(1, 8, std::slice::from_ref(&(0..8)), |_| async { + Err::, _>(IndexNotFoundSnafu { name: "missing" }.build()) + }) + .await + .unwrap_err(); + let retained = find_shared_page_source::(&error).unwrap(); + assert_eq!( + retained.status_code(), + common_error::status_code::StatusCode::InvalidArguments + ); + assert!(!cache.index.contains(&1, 0).await); + let (pages, _) = cache + .get_or_load_vec(1, 8, std::slice::from_ref(&(0..8)), |ranges| async move { + Ok::<_, io::Error>( + ranges + .iter() + .map(|r| Bytes::from(vec![7; (r.end - r.start) as usize])) + .collect(), + ) + }) + .await + .unwrap(); + assert_eq!(pages, vec![Bytes::from(vec![7; 8])]); + } + #[tokio::test] + async fn single_page_range_retains_shared_bytes() { + let cache = metadata_cache(); + let (first, _) = cache + .get_or_load(1, 8, 1, 2, |ranges| async move { + Ok::<_, io::Error>( + ranges + .into_iter() + .map(|range| Bytes::from(range.map(|n| n as u8).collect::>())) + .collect(), + ) + }) + .await + .unwrap(); + let (second, metrics) = cache + .get_or_load(1, 8, 1, 2, |_| async { + Err::, _>(io::Error::other("unexpected read on cache hit")) + }) + .await + .unwrap(); + assert_eq!(&first[..], &[1, 2]); + assert_eq!(first, second); + assert_eq!(first.as_ptr(), second.as_ptr()); + assert_eq!(metrics.cache_hit, 1); + assert_eq!(metrics.cache_miss, 0); + } } diff --git a/src/mito2/src/cache/index/bloom_filter_index.rs b/src/mito2/src/cache/index/bloom_filter_index.rs index 0e81817a754..1bd8b0fbae6 100644 --- a/src/mito2/src/cache/index/bloom_filter_index.rs +++ b/src/mito2/src/cache/index/bloom_filter_index.rs @@ -19,12 +19,13 @@ use std::time::Instant; use api::v1::index::{BloomFilterLoc, BloomFilterMeta}; use async_trait::async_trait; use bytes::Bytes; -use index::bloom_filter::error::Result; +use index::bloom_filter::error::{Error, IoSnafu, Result, SharedReadSnafu}; use index::bloom_filter::reader::{BloomFilterReadMetrics, BloomFilterReader}; use index::target::IndexTarget; +use snafu::IntoError; use store_api::storage::{FileId, IndexVersion}; -use crate::cache::index::{INDEX_METADATA_TYPE, IndexCache, PageKey}; +use crate::cache::index::{INDEX_METADATA_TYPE, IndexCache}; use crate::metrics::{CACHE_HIT, CACHE_MISS}; const INDEX_TYPE_BLOOM_FILTER_INDEX: &str = "bloom_filter_index"; @@ -50,6 +51,7 @@ impl BloomFilterIndexCache { index_content_cap, page_size, INDEX_TYPE_BLOOM_FILTER_INDEX, + "bloom_filter_index_metadata", bloom_filter_index_metadata_weight, bloom_filter_index_content_weight, ) @@ -76,7 +78,7 @@ fn bloom_filter_index_metadata_weight(k: &BloomFilterIndexKey, meta: &Arc u32 { +fn bloom_filter_index_content_weight((k, _): &(BloomFilterIndexKey, u64), v: &Bytes) -> u32 { (k.0.as_bytes().len() + k.2.approx_mem_size() + v.len()) as u32 } @@ -138,7 +140,8 @@ impl BloomFilterReader for CachedBloomFilterIndexBl size, move |ranges| async move { inner.read_vec(&ranges, None).await }, ) - .await?; + .await + .map_err(map_page_error)?; if let Some(m) = metrics { m.total_ranges += cache_metrics.num_pages; @@ -150,7 +153,7 @@ impl BloomFilterReader for CachedBloomFilterIndexBl } } - Ok(result.into()) + Ok(result) } async fn read_vec( @@ -174,7 +177,8 @@ impl BloomFilterReader for CachedBloomFilterIndexBl ranges, move |ranges| async move { inner.read_vec(&ranges, None).await }, ) - .await?; + .await + .map_err(map_page_error)?; if let Some(m) = metrics { m.total_ranges += total_cache_metrics.num_pages; m.total_bytes += total_cache_metrics.page_bytes; @@ -193,31 +197,55 @@ impl BloomFilterReader for CachedBloomFilterIndexBl &self, metrics: Option<&mut BloomFilterReadMetrics>, ) -> Result> { - if let Some(cached) = self.cache.get_metadata(( + let key = ( self.file_id, self.index_version, self.target.clone(), self.tag, - )) { + ); + if let Some(cached) = self.cache.get_metadata(key.clone()).await { CACHE_HIT.with_label_values(&[INDEX_METADATA_TYPE]).inc(); if let Some(m) = metrics { m.cache_hit += 1; } - Ok(cached) - } else { - let meta = self.inner.metadata(metrics).await?; - self.cache.put_metadata( - ( - self.file_id, - self.index_version, - self.target.clone(), - self.tag, - ), - meta.clone(), - ); - CACHE_MISS.with_label_values(&[INDEX_METADATA_TYPE]).inc(); - Ok(meta) + return Ok(cached); } + + let start = metrics.as_ref().map(|_| Instant::now()); + let collect_metrics = metrics.is_some(); + let mut loaded = BloomFilterReadMetrics::default(); + let result = self + .cache + .get_or_load_metadata(key, async { + self.inner + .metadata(collect_metrics.then_some(&mut loaded)) + .await + }) + .await; + if let Some(m) = metrics { + // Only the selected loader performs I/O. Waiters retain zero I/O + // counts. fetch_elapsed is per-caller elapsed (including wait and error), + // not physical I/O time summed across callers. + m.total_bytes += loaded.total_bytes; + m.total_ranges += loaded.total_ranges; + m.cache_hit += loaded.cache_hit; + m.cache_miss += loaded.cache_miss; + if let Some(start) = start { + m.fetch_elapsed += start.elapsed(); + } + } + let meta = result.map_err(|source| SharedReadSnafu.into_error(source))?; + CACHE_MISS.with_label_values(&[INDEX_METADATA_TYPE]).inc(); + Ok(meta) + } +} + +/// Recovers the typed source retained by the range-reader adapter. +fn map_page_error(error: Arc) -> Error { + if let Some(source) = crate::cache::index::find_shared_page_source::(&error) { + SharedReadSnafu.into_error(source) + } else { + IoSnafu.into_error(std::io::Error::other(error)) } } @@ -230,8 +258,109 @@ mod test { const FUZZ_REPEAT_TIMES: usize = 100; - #[test] - fn json_paths_have_independent_cache_entries() { + struct MetadataReader { + release: tokio::sync::Notify, + fail: std::sync::atomic::AtomicBool, + } + + #[async_trait] + impl BloomFilterReader for MetadataReader { + async fn range_read( + &self, + _: u64, + _: u32, + _: Option<&mut BloomFilterReadMetrics>, + ) -> Result { + unreachable!("metadata test does not read content pages") + } + + async fn read_vec( + &self, + _: &[Range], + _: Option<&mut BloomFilterReadMetrics>, + ) -> Result> { + unreachable!("metadata test does not read content pages") + } + + async fn metadata( + &self, + metrics: Option<&mut BloomFilterReadMetrics>, + ) -> Result> { + self.release.notified().await; + if let Some(metrics) = metrics { + metrics.total_bytes += 8192; + metrics.total_ranges += 1; + // The wrapper must replace, not add, this physical-fetch duration. + metrics.fetch_elapsed += std::time::Duration::from_secs(3600); + } + if self.fail.load(std::sync::atomic::Ordering::Relaxed) { + return Err(IoSnafu.into_error(std::io::Error::other("metadata failed"))); + } + Ok(Arc::new(BloomFilterMeta::default())) + } + } + + #[tokio::test] + async fn metadata_waiters_do_not_duplicate_io_metrics() { + for fail in [false, true] { + let cache = Arc::new(BloomFilterIndexCache::new(4096, 4096, 1024)); + let reader = CachedBloomFilterIndexBlobReader::new( + FileId::random(), + 0, + IndexTarget::ColumnId(42), + Tag::Skipping, + 8192, + MetadataReader { + release: tokio::sync::Notify::new(), + fail: std::sync::atomic::AtomicBool::new(fail), + }, + cache, + ); + let mut owner_metrics = BloomFilterReadMetrics::default(); + let mut waiter_metrics = BloomFilterReadMetrics::default(); + let mut owner = Box::pin(reader.metadata(Some(&mut owner_metrics))); + assert!(futures::poll!(&mut owner).is_pending()); + let mut waiter = Box::pin(reader.metadata(Some(&mut waiter_metrics))); + assert!(futures::poll!(&mut waiter).is_pending()); + reader.inner.release.notify_one(); + let owner_result = owner.await; + let waiter_result = waiter.await; + assert_eq!(owner_result.is_err(), fail); + assert_eq!(waiter_result.is_err(), fail); + assert_eq!( + (owner_metrics.total_bytes, owner_metrics.total_ranges), + (8192, 1) + ); + assert_eq!( + (waiter_metrics.total_bytes, waiter_metrics.total_ranges), + (0, 0) + ); + assert_eq!((owner_metrics.cache_hit, waiter_metrics.cache_hit), (0, 0)); + for metrics in [&owner_metrics, &waiter_metrics] { + // Per-caller elapsed includes waiting, including on error; it is + // not a sum of physical I/O durations across concurrent callers. + assert!(metrics.fetch_elapsed > std::time::Duration::ZERO); + assert!(metrics.fetch_elapsed < std::time::Duration::from_secs(3600)); + } + if fail { + let Error::SharedRead { source: owner, .. } = owner_result.unwrap_err() else { + panic!("owner must preserve the typed shared read error"); + }; + let Error::SharedRead { source: waiter, .. } = waiter_result.unwrap_err() else { + panic!("waiter must preserve the typed shared read error"); + }; + assert!(Arc::ptr_eq(&owner, &waiter)); + } else { + let mut hot_metrics = BloomFilterReadMetrics::default(); + reader.metadata(Some(&mut hot_metrics)).await.unwrap(); + assert_eq!(hot_metrics.cache_hit, 1); + assert_eq!(hot_metrics.total_bytes, 0); + } + } + } + + #[tokio::test] + async fn json_paths_have_independent_cache_entries() { let cache = BloomFilterIndexCache::new(1024 * 1024, 1024 * 1024, 1024); let file_id = FileId::random(); let first = ( @@ -246,25 +375,35 @@ mod test { Arc::new(IndexTarget::new_json_path(1, vec!["a.b".into()]).unwrap()), Tag::Skipping, ); - cache.put_metadata( - first.clone(), - Arc::new(BloomFilterMeta { - row_count: 1, - ..Default::default() - }), + cache + .put_metadata( + first.clone(), + Arc::new(BloomFilterMeta { + row_count: 1, + ..Default::default() + }), + ) + .await; + cache + .put_metadata( + second.clone(), + Arc::new(BloomFilterMeta { + row_count: 2, + ..Default::default() + }), + ) + .await; + assert_eq!( + cache.get_metadata(first.clone()).await.unwrap().row_count, + 1 ); - cache.put_metadata( - second.clone(), - Arc::new(BloomFilterMeta { - row_count: 2, - ..Default::default() - }), + assert_eq!( + cache.get_metadata(second.clone()).await.unwrap().row_count, + 2 ); - assert_eq!(cache.get_metadata(first.clone()).unwrap().row_count, 1); - assert_eq!(cache.get_metadata(second.clone()).unwrap().row_count, 2); cache.invalidate_file(file_id); - assert!(cache.get_metadata(first).is_none()); - assert!(cache.get_metadata(second).is_none()); + assert!(cache.get_metadata(first).await.is_none()); + assert!(cache.get_metadata(second).await.is_none()); } #[test] @@ -317,42 +456,29 @@ mod test { #[test] fn fuzz_index_calculation() { + use common_cache::page::PageLayout; + use std::num::NonZeroU64; + let mut rng = rand::rng(); let mut data = vec![0u8; 1024 * 1024]; rng.fill_bytes(&mut data); - for _ in 0..FUZZ_REPEAT_TIMES { let offset = rng.random_range(0..data.len() as u64); - let size = rng.random_range(0..data.len() as u32 - offset as u32); - let page_size: usize = rng.random_range(1..1024); - - let indexes = - PageKey::generate_page_keys(offset, size, page_size as u64).collect::>(); - let page_num = indexes.len(); + let size = rng.random_range(0..data.len() as u64 - offset); + let page_size = NonZeroU64::new(rng.random_range(1_u64..1024)).unwrap(); + let layout = PageLayout::new(page_size, data.len() as u64); + let requested = offset..offset + size; let mut read = Vec::with_capacity(size as usize); - for key in indexes.into_iter() { - let start = key.page_id as usize * page_size; - let page = if start + page_size < data.len() { - &data[start..start + page_size] - } else { - &data[start..] - }; - read.extend_from_slice(page); + for id in layout.page_ids(requested.clone()).unwrap() { + let range = layout.page_range(id).unwrap(); + let page = &data[range.start as usize..range.end as usize]; + let slice = layout + .request_slice(id, requested.clone()) + .unwrap() + .unwrap(); + read.extend_from_slice(&page[slice]); } - let expected_range = offset as usize..(offset + size as u64 as u64) as usize; - let read = read[PageKey::calculate_range(offset, size, page_size as u64)].to_vec(); - assert_eq!( - read, - data.get(expected_range).unwrap(), - "fuzz_read_index failed, offset: {}, size: {}, page_size: {}\nread len: {}, expected len: {}\nrange: {:?}, page num: {}", - offset, - size, - page_size, - read.len(), - size as usize, - PageKey::calculate_range(offset, size, page_size as u64), - page_num - ); + assert_eq!(read, data[offset as usize..(offset + size) as usize]); } } } diff --git a/src/mito2/src/cache/index/inverted_index.rs b/src/mito2/src/cache/index/inverted_index.rs index 675ed736110..b3a0c6af955 100644 --- a/src/mito2/src/cache/index/inverted_index.rs +++ b/src/mito2/src/cache/index/inverted_index.rs @@ -19,12 +19,13 @@ use std::time::Instant; use api::v1::index::InvertedIndexMetas; use async_trait::async_trait; use bytes::Bytes; -use index::inverted_index::error::Result; +use index::inverted_index::error::{CommonIoSnafu, Error, Result, SharedReadSnafu}; use index::inverted_index::format::reader::{InvertedIndexReadMetrics, InvertedIndexReader}; use prost::Message; +use snafu::IntoError; use store_api::storage::{FileId, IndexVersion}; -use crate::cache::index::{INDEX_METADATA_TYPE, IndexCache, PageKey}; +use crate::cache::index::{INDEX_METADATA_TYPE, IndexCache}; use crate::metrics::{CACHE_HIT, CACHE_MISS}; const INDEX_TYPE_INVERTED_INDEX: &str = "inverted_index"; @@ -41,6 +42,7 @@ impl InvertedIndexCache { index_content_cap, page_size, INDEX_TYPE_INVERTED_INDEX, + "inverted_index_metadata", inverted_index_metadata_weight, inverted_index_content_weight, ) @@ -58,7 +60,7 @@ fn inverted_index_metadata_weight(k: &(FileId, IndexVersion), v: &Arc u32 { +fn inverted_index_content_weight((k, _): &((FileId, IndexVersion), u64), v: &Bytes) -> u32 { (k.0.as_bytes().len() + size_of::() + v.len()) as u32 } @@ -110,7 +112,8 @@ impl InvertedIndexReader for CachedInvertedIndexBlobRead size, move |ranges| async move { inner.read_vec(&ranges, None).await }, ) - .await?; + .await + .map_err(map_page_error)?; if let Some(m) = metrics { m.total_bytes += cache_metrics.page_bytes; @@ -120,7 +123,7 @@ impl InvertedIndexReader for CachedInvertedIndexBlobRead m.fetch_elapsed += start.unwrap().elapsed(); } - Ok(result) + Ok(result.to_vec()) } async fn read_vec<'a>( @@ -139,7 +142,8 @@ impl InvertedIndexReader for CachedInvertedIndexBlobRead ranges, move |ranges| async move { inner.read_vec(&ranges, None).await }, ) - .await?; + .await + .map_err(map_page_error)?; if let Some(m) = metrics { m.total_bytes += total_cache_metrics.page_bytes; m.total_ranges += total_cache_metrics.num_pages; @@ -155,19 +159,50 @@ impl InvertedIndexReader for CachedInvertedIndexBlobRead &self, metrics: Option<&'a mut InvertedIndexReadMetrics>, ) -> Result> { - if let Some(cached) = self.cache.get_metadata((self.file_id, self.index_version)) { + let key = (self.file_id, self.index_version); + if let Some(cached) = self.cache.get_metadata(key).await { CACHE_HIT.with_label_values(&[INDEX_METADATA_TYPE]).inc(); if let Some(m) = metrics { m.cache_hit += 1; } - Ok(cached) - } else { - let meta = self.inner.metadata(metrics).await?; - self.cache - .put_metadata((self.file_id, self.index_version), meta.clone()); - CACHE_MISS.with_label_values(&[INDEX_METADATA_TYPE]).inc(); - Ok(meta) + return Ok(cached); } + + let start = metrics.as_ref().map(|_| Instant::now()); + let collect_metrics = metrics.is_some(); + let mut loaded = InvertedIndexReadMetrics::default(); + let result = self + .cache + .get_or_load_metadata(key, async { + self.inner + .metadata(collect_metrics.then_some(&mut loaded)) + .await + }) + .await; + if let Some(m) = metrics { + // Only the selected loader performs I/O. Waiters retain zero I/O + // counts. fetch_elapsed is per-caller elapsed (including wait and error), + // not physical I/O time summed across callers. + m.total_bytes += loaded.total_bytes; + m.total_ranges += loaded.total_ranges; + m.cache_hit += loaded.cache_hit; + m.cache_miss += loaded.cache_miss; + if let Some(start) = start { + m.fetch_elapsed += start.elapsed(); + } + } + let meta = result.map_err(|source| SharedReadSnafu.into_error(source))?; + CACHE_MISS.with_label_values(&[INDEX_METADATA_TYPE]).inc(); + Ok(meta) + } +} + +/// Recovers the typed source retained by the range-reader adapter. +fn map_page_error(error: Arc) -> Error { + if let Some(source) = crate::cache::index::find_shared_page_source::(&error) { + SharedReadSnafu.into_error(source) + } else { + CommonIoSnafu.into_error(std::io::Error::other(error)) } } @@ -190,46 +225,131 @@ mod test { // Repeat times for following little fuzz tests. const FUZZ_REPEAT_TIMES: usize = 100; + struct MetadataReader { + release: tokio::sync::Notify, + fail: std::sync::atomic::AtomicBool, + } + + #[async_trait] + impl InvertedIndexReader for MetadataReader { + async fn range_read<'a>( + &self, + _: u64, + _: u32, + _: Option<&'a mut InvertedIndexReadMetrics>, + ) -> Result> { + unreachable!("metadata test does not read content pages") + } + + async fn read_vec<'a>( + &self, + _: &[Range], + _: Option<&'a mut InvertedIndexReadMetrics>, + ) -> Result> { + unreachable!("metadata test does not read content pages") + } + + async fn metadata<'a>( + &self, + metrics: Option<&'a mut InvertedIndexReadMetrics>, + ) -> Result> { + self.release.notified().await; + if let Some(metrics) = metrics { + metrics.total_bytes += 8192; + metrics.total_ranges += 1; + // The wrapper must replace, not add, this physical-fetch duration. + metrics.fetch_elapsed += std::time::Duration::from_secs(3600); + } + if self.fail.load(std::sync::atomic::Ordering::Relaxed) { + return Err(CommonIoSnafu.into_error(std::io::Error::other("metadata failed"))); + } + Ok(Arc::new(InvertedIndexMetas::default())) + } + } + + #[tokio::test] + async fn metadata_waiters_do_not_duplicate_io_metrics() { + for fail in [false, true] { + let cache = Arc::new(InvertedIndexCache::new(4096, 4096, 1024)); + let reader = CachedInvertedIndexBlobReader::new( + FileId::random(), + 0, + 8192, + MetadataReader { + release: tokio::sync::Notify::new(), + fail: std::sync::atomic::AtomicBool::new(fail), + }, + cache, + ); + let mut owner_metrics = InvertedIndexReadMetrics::default(); + let mut waiter_metrics = InvertedIndexReadMetrics::default(); + let mut owner = Box::pin(reader.metadata(Some(&mut owner_metrics))); + assert!(futures::poll!(&mut owner).is_pending()); + let mut waiter = Box::pin(reader.metadata(Some(&mut waiter_metrics))); + assert!(futures::poll!(&mut waiter).is_pending()); + reader.inner.release.notify_one(); + let owner_result = owner.await; + let waiter_result = waiter.await; + assert_eq!(owner_result.is_err(), fail); + assert_eq!(waiter_result.is_err(), fail); + assert_eq!( + (owner_metrics.total_bytes, owner_metrics.total_ranges), + (8192, 1) + ); + assert_eq!( + (waiter_metrics.total_bytes, waiter_metrics.total_ranges), + (0, 0) + ); + assert_eq!((owner_metrics.cache_hit, waiter_metrics.cache_hit), (0, 0)); + for metrics in [&owner_metrics, &waiter_metrics] { + // Per-caller elapsed includes waiting, including on error; it is + // not a sum of physical I/O durations across concurrent callers. + assert!(metrics.fetch_elapsed > std::time::Duration::ZERO); + assert!(metrics.fetch_elapsed < std::time::Duration::from_secs(3600)); + } + if fail { + let Error::SharedRead { source: owner, .. } = owner_result.unwrap_err() else { + panic!("owner must preserve the typed shared read error"); + }; + let Error::SharedRead { source: waiter, .. } = waiter_result.unwrap_err() else { + panic!("waiter must preserve the typed shared read error"); + }; + assert!(Arc::ptr_eq(&owner, &waiter)); + } else { + let mut hot_metrics = InvertedIndexReadMetrics::default(); + reader.metadata(Some(&mut hot_metrics)).await.unwrap(); + assert_eq!(hot_metrics.cache_hit, 1); + assert_eq!(hot_metrics.total_bytes, 0); + } + } + } + // Fuzz test for index data page key #[test] fn fuzz_index_calculation() { - // randomly generate a large u8 array + use common_cache::page::PageLayout; + use std::num::NonZeroU64; + let mut rng = rand::rng(); let mut data = vec![0u8; 1024 * 1024]; rng.fill_bytes(&mut data); - for _ in 0..FUZZ_REPEAT_TIMES { let offset = rng.random_range(0..data.len() as u64); - let size = rng.random_range(0..data.len() as u32 - offset as u32); - let page_size: usize = rng.random_range(1..1024); - - let indexes = - PageKey::generate_page_keys(offset, size, page_size as u64).collect::>(); - let page_num = indexes.len(); + let size = rng.random_range(0..data.len() as u64 - offset); + let page_size = NonZeroU64::new(rng.random_range(1_u64..1024)).unwrap(); + let layout = PageLayout::new(page_size, data.len() as u64); + let requested = offset..offset + size; let mut read = Vec::with_capacity(size as usize); - for key in indexes.into_iter() { - let start = key.page_id as usize * page_size; - let page = if start + page_size < data.len() { - &data[start..start + page_size] - } else { - &data[start..] - }; - read.extend_from_slice(page); - } - let expected_range = offset as usize..(offset + size as u64 as u64) as usize; - let read = read[PageKey::calculate_range(offset, size, page_size as u64)].to_vec(); - if read != data.get(expected_range).unwrap() { - panic!( - "fuzz_read_index failed, offset: {}, size: {}, page_size: {}\nread len: {}, expected len: {}\nrange: {:?}, page num: {}", - offset, - size, - page_size, - read.len(), - size as usize, - PageKey::calculate_range(offset, size, page_size as u64), - page_num - ); + for id in layout.page_ids(requested.clone()).unwrap() { + let range = layout.page_range(id).unwrap(); + let page = &data[range.start as usize..range.end as usize]; + let slice = layout + .request_slice(id, requested.clone()) + .unwrap() + .unwrap(); + read.extend_from_slice(&page[slice]); } + assert_eq!(read, data[offset as usize..(offset + size) as usize]); } } @@ -460,7 +580,7 @@ mod test { ) .await .unwrap(); - assert_eq!(read, expected); + assert_eq!(read.as_ref(), expected.as_slice()); } } @@ -489,10 +609,13 @@ mod test { .collect::>(); let loads = AtomicUsize::new(0); + let physical_ranges = std::sync::Mutex::new(Vec::new()); let load = |pages: Vec>| { loads.fetch_add(1, Ordering::Relaxed); - // Pages are requested once each, in file order. - assert!(pages.windows(2).all(|w| w[0].end <= w[1].start)); + physical_ranges + .lock() + .unwrap() + .extend(pages.iter().cloned()); let data = &data; async move { Ok::<_, std::io::Error>( @@ -512,8 +635,15 @@ mod test { .await .unwrap(); assert_eq!(cold, expected); - let cold_loads = usize::from(cold_metrics.cache_miss > 0); - assert_eq!(loads.load(Ordering::Relaxed), cold_loads); + let cold_loads = loads.load(Ordering::Relaxed); + assert_eq!(cold_loads, usize::from(cold_metrics.cache_miss > 0)); + let mut observed = physical_ranges.lock().unwrap().clone(); + observed.sort_by_key(|r| r.start); + assert!(observed.windows(2).all(|w| w[0].end <= w[1].start)); + assert_eq!( + observed.iter().map(|r| r.end - r.start).sum::(), + cold_metrics.page_bytes + ); let (warm, metrics) = cache .get_or_load_vec(key, file_size, &ranges, load) diff --git a/src/mito2/src/metrics.rs b/src/mito2/src/metrics.rs index acb6a9da4c4..b88b7e5231c 100644 --- a/src/mito2/src/metrics.rs +++ b/src/mito2/src/metrics.rs @@ -738,3 +738,171 @@ impl StagerNotifier for StagerMetrics { self.recycle_cache_bytes.sub(size as i64); } } + +lazy_static! { + pub static ref CACHE_PAGE_LOADED_BYTES: IntCounterVec = register_int_counter_vec!( + "greptime_mito_cache_page_loaded_bytes_total", + "Bytes produced by successful page loaders, excluding waiters", + &[TYPE_LABEL] + ) + .unwrap(); + /// Configured capacity summed across live page caches in each category. + pub static ref CACHE_CAPACITY_BYTES: common_cache::metrics::CacheSizeMetrics = { + let metrics = common_cache::metrics::CacheSizeMetrics::new( + prometheus::Opts::new("greptime_mito_cache_capacity_bytes", "Configured capacity of live page caches in bytes"), + &[TYPE_LABEL], + ).unwrap(); + prometheus::register(Box::new(metrics.clone())).unwrap(); + metrics + }; +} + +/// Binds shared page/cache I/O metrics while retaining the business hit/miss labels. +pub(crate) fn cache_page_observer( + cache_name_label: &'static str, + cache_category_label: &'static str, + lookup_counts: fn(usize, usize) -> (usize, usize), +) -> common_cache::paged::PageCacheObserver { + use common_cache::paged::PageEvent; + let hits = CACHE_HIT.with_label_values(&[cache_category_label]); + let misses = CACHE_MISS.with_label_values(&[cache_category_label]); + let loaded = CACHE_PAGE_LOADED_BYTES.with_label_values(&[cache_name_label]); + common_cache::paged::PageCacheObserver { + bytes: CACHE_BYTES.bind(&[cache_category_label]), + capacity: CACHE_CAPACITY_BYTES.bind(&[cache_category_label]), + loads: Some(cache_loading_metrics(cache_name_label)), + events: std::sync::Arc::new(move |event| match event { + PageEvent::Lookup { + hits: count, + misses: missing, + } => { + let (count, missing) = lookup_counts(count, missing); + hits.inc_by(count as u64); + misses.inc_by(missing as u64); + } + PageEvent::Loaded { bytes } => loaded.inc_by(bytes), + PageEvent::Evicted { cause } => CACHE_EVICTION + .with_label_values(&[cache_category_label, cause.as_str()]) + .inc(), + }), + } +} + +lazy_static! { + pub static ref CACHE_PAGE_BYTES: IntCounterVec = register_int_counter_vec!( + "greptime_mito_cache_page_bytes_total", + "Logical page fetch bytes by kind; not physical backend bytes", + &[TYPE_LABEL, "kind"] + ) + .unwrap(); +} + +/// Records logical byte amplification independently of physical batch reads. +pub(crate) fn cache_page_fetch_metrics( + cache_type: &'static str, + stats: &common_cache::paged::PageFetchStats, +) { + for (kind, value) in [ + ("requested", stats.requested_bytes), + ("aligned", stats.aligned_bytes), + ("unique_requested", stats.unique_requested_bytes), + ("cached", stats.cached_bytes), + ] { + CACHE_PAGE_BYTES + .with_label_values(&[cache_type, kind]) + .inc_by(value); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use common_cache::paged::{PageEvent, PageFetchStats}; + + #[test] + fn page_cache_capacity_sums_live_instances_by_category() { + use bytes::Bytes; + use common_cache::paged::PagedCache; + use prometheus::core::Collector; + + let category = "test_shared_page_category"; + let sample = || { + CACHE_CAPACITY_BYTES + .collect() + .into_iter() + .flat_map(|family| family.get_metric().to_vec()) + .find(|metric| { + metric + .get_label() + .iter() + .any(|label| label.value() == category) + }) + .unwrap() + .get_gauge() + .as_ref() + .unwrap() + .value() + }; + let make_cache = |name, capacity| { + PagedCache::::new( + name, + capacity, + 64, + |_, bytes: &Bytes| bytes.len() as u32, + cache_page_observer(name, category, |hits, misses| (hits, misses)), + ) + }; + let first = make_cache("test_first_pages", 128); + let second = make_cache("test_second_pages", 256); + assert_eq!(sample(), 384.0); + drop(first); + assert_eq!(sample(), 256.0); + drop(second); + assert_eq!(sample(), 0.0); + } + + #[test] + fn test_page_observer_retains_business_lookup_units() { + let page = cache_page_observer("test_page_unit", "test_page_unit", |hits, misses| { + (hits, misses) + }); + let request = + cache_page_observer("test_request_unit", "test_request_unit", |hits, misses| { + (usize::from(hits > 0), usize::from(misses > 0)) + }); + for observer in [&page, &request] { + (observer.events)(PageEvent::Lookup { hits: 3, misses: 2 }); + (observer.events)(PageEvent::Loaded { bytes: 4096 }); + } + assert_eq!(CACHE_HIT.with_label_values(&["test_page_unit"]).get(), 3); + assert_eq!(CACHE_MISS.with_label_values(&["test_page_unit"]).get(), 2); + assert_eq!(CACHE_HIT.with_label_values(&["test_request_unit"]).get(), 1); + assert_eq!( + CACHE_MISS.with_label_values(&["test_request_unit"]).get(), + 1 + ); + cache_page_fetch_metrics( + "test_page_unit", + &PageFetchStats { + requested_bytes: 7, + aligned_bytes: 12, + unique_requested_bytes: 6, + cached_bytes: 3, + ..Default::default() + }, + ); + for (kind, expected) in [ + ("requested", 7), + ("aligned", 12), + ("unique_requested", 6), + ("cached", 3), + ] { + assert_eq!( + CACHE_PAGE_BYTES + .with_label_values(&["test_page_unit", kind]) + .get(), + expected + ); + } + } +} diff --git a/src/mito2/src/sst/index/inverted_index/creator.rs b/src/mito2/src/sst/index/inverted_index/creator.rs index 239fb608c1d..8c12eba6348 100644 --- a/src/mito2/src/sst/index/inverted_index/creator.rs +++ b/src/mito2/src/sst/index/inverted_index/creator.rs @@ -631,8 +631,7 @@ mod tests { move |expr| { let _d = &d; let cache = Arc::new(InvertedIndexCache::new(10, 10, 100)); - let puffin_metadata_cache = - Arc::new(PuffinMetadataCache::new(10, CACHE_BYTES.gauges())); + let puffin_metadata_cache = Arc::new(PuffinMetadataCache::new(10, &CACHE_BYTES)); let applier = InvertedIndexApplierBuilder::new( table_dir.clone(), PathType::Bare, diff --git a/src/puffin/Cargo.toml b/src/puffin/Cargo.toml index a51ca99d444..42d6a02ce92 100644 --- a/src/puffin/Cargo.toml +++ b/src/puffin/Cargo.toml @@ -16,6 +16,7 @@ base64.workspace = true bitflags.workspace = true bytes.workspace = true common-base.workspace = true +common-cache.workspace = true common-error.workspace = true common-macro.workspace = true common-runtime.workspace = true @@ -33,4 +34,5 @@ sha2 = "0.10.8" snafu.workspace = true tokio.workspace = true tokio-util.workspace = true +tracing.workspace = true uuid.workspace = true diff --git a/src/puffin/src/puffin_manager/cache.rs b/src/puffin/src/puffin_manager/cache.rs index fc6df95d027..bfc4269ef11 100644 --- a/src/puffin/src/puffin_manager/cache.rs +++ b/src/puffin/src/puffin_manager/cache.rs @@ -12,11 +12,18 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::future::Future; use std::sync::Arc; -use prometheus::IntGaugeVec; +use common_cache::Cache; +use common_cache::metrics::CacheMetrics; +use common_cache::metrics::CacheSizeMetrics; +use common_telemetry::tracing::{self, Instrument}; +use snafu::ResultExt; +use crate::error::{CacheGetSnafu, Result}; use crate::file_metadata::FileMetadata; + /// Metrics for index metadata. const PUFFIN_METADATA_TYPE: &str = "puffin_metadata"; @@ -24,8 +31,7 @@ pub type PuffinMetadataCacheRef = Arc; /// A cache for storing the metadata of the index files. pub struct PuffinMetadataCache { - cache: moka::sync::Cache>, - cache_bytes: IntGaugeVec, + cache: Cache>, } fn puffin_metadata_weight(k: &str, v: &Arc) -> u32 { @@ -33,39 +39,199 @@ fn puffin_metadata_weight(k: &str, v: &Arc) -> u32 { } impl PuffinMetadataCache { - pub fn new(capacity: u64, cache_bytes: &'static IntGaugeVec) -> Self { + pub fn new(capacity: u64, cache_bytes: &CacheSizeMetrics) -> Self { common_telemetry::debug!("Building PuffinMetadataCache with capacity: {capacity}"); Self { - cache: moka::sync::CacheBuilder::new(capacity) - .name("puffin_metadata") - .weigher(|k: &String, v| puffin_metadata_weight(k, v)) - .eviction_listener(move |k, v, _cause| { - let size = puffin_metadata_weight(&k, &v); - cache_bytes - .with_label_values(&[PUFFIN_METADATA_TYPE]) - .sub(size.into()); - }) - .build(), - cache_bytes: cache_bytes.clone(), + cache: Cache::new( + "puffin_metadata", + capacity, + |k: &String, v: &Arc| puffin_metadata_weight(k, v), + |_, _, _| {}, + ) + .with_size_metric(cache_bytes.bind(&[PUFFIN_METADATA_TYPE])), } } + /// Attaches native loader metrics without changing existing cache byte accounting. + pub fn with_load_metrics(mut self, metrics: Arc) -> Self { + self.cache = self.cache.with_metrics(metrics); + self + } + /// Gets the metadata from the cache. - pub fn get_metadata(&self, file_id: &str) -> Option> { - self.cache.get(file_id) + pub async fn get_metadata(&self, file_id: &str) -> Option> { + self.cache.get(file_id).await + } + + /// Shares one metadata load among concurrent readers of the same file. + #[tracing::instrument(name = "puffin_metadata_lookup", skip_all)] + pub async fn try_get_with(&self, file_id: &str, load: F) -> Result> + where + F: Future>>, + { + self.cache + .try_get_with(file_id, async { + let metadata = load + .instrument(tracing::info_span!("puffin_metadata_load")) + .await?; + Ok(metadata) + }) + .await + .context(CacheGetSnafu) } /// Puts the metadata into the cache. - pub fn put_metadata(&self, file_id: String, metadata: Arc) { - let size = puffin_metadata_weight(&file_id, &metadata); - self.cache_bytes - .with_label_values(&[PUFFIN_METADATA_TYPE]) - .add(size.into()); - self.cache.insert(file_id, metadata); + pub async fn put_metadata(&self, file_id: String, metadata: Arc) { + self.cache.insert(file_id, metadata).await; } /// Removes the metadata of the given file from the cache, if present. - pub fn remove(&self, file_id: &str) { - self.cache.invalidate(file_id); + pub async fn remove(&self, file_id: &str) { + self.cache.invalidate(file_id).await; + } +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use prometheus::Opts; + use prometheus::core::Collector; + + use super::*; + use crate::error::InvalidPuffinFooterSnafu; + + fn new_cache() -> (PuffinMetadataCache, CacheSizeMetrics) { + let metrics = + CacheSizeMetrics::new(Opts::new("test_metadata_bytes", "test"), &["type"]).unwrap(); + (PuffinMetadataCache::new(1024 * 1024, &metrics), metrics) + } + + fn sample_bytes(metrics: &CacheSizeMetrics) -> i64 { + metrics.collect()[0].get_metric()[0] + .get_gauge() + .as_ref() + .unwrap() + .value() as i64 + } + + fn metadata() -> Arc { + Arc::new(FileMetadata { + blobs: Vec::new(), + properties: Default::default(), + }) + } + + // Hold both requests at the same cold-load boundary. Only coalescing differs. + #[tokio::test] + async fn compare_separate_lookup_with_coalesced_load() { + for _ in 0..3 { + for coalesce in [false, true] { + let (cache, _) = new_cache(); + let calls = AtomicUsize::new(0); + let gate = tokio::sync::Notify::new(); + let request = || async { + let load = async { + calls.fetch_add(1, Ordering::Relaxed); + gate.notified().await; + Ok(metadata()) + }; + if coalesce { + cache.try_get_with("file", load).await + } else if let Some(value) = cache.get_metadata("file").await { + Ok(value) + } else { + let value = load.await?; + cache.put_metadata("file".to_string(), value.clone()).await; + Ok(value) + } + }; + let mut first = Box::pin(request()); + let mut second = Box::pin(request()); + assert!(futures::poll!(&mut first).is_pending()); + assert!(futures::poll!(&mut second).is_pending()); + assert_eq!(calls.load(Ordering::Relaxed), if coalesce { 1 } else { 2 }); + gate.notify_waiters(); + let (first, second) = tokio::join!(first, second); + assert_eq!(first.unwrap().as_ref(), second.unwrap().as_ref()); + } + } + } + + #[tokio::test] + async fn concurrent_loads_share_result_and_accounting() { + let (cache, gauge) = new_cache(); + let calls = AtomicUsize::new(0); + let gate = tokio::sync::Notify::new(); + let load = || async { + calls.fetch_add(1, Ordering::Relaxed); + gate.notified().await; + Ok(metadata()) + }; + let mut first = Box::pin(cache.try_get_with("file", load())); + let mut second = Box::pin(cache.try_get_with("file", load())); + assert!(futures::poll!(&mut first).is_pending()); + assert!(futures::poll!(&mut second).is_pending()); + assert_eq!(calls.load(Ordering::Relaxed), 1); + gate.notify_waiters(); + let (first, second) = tokio::join!(first, second); + let first = first.unwrap(); + assert!(Arc::ptr_eq(&first, &second.unwrap())); + cache.cache.run_pending_tasks().await; + assert_eq!( + sample_bytes(&gauge), + i64::from(puffin_metadata_weight("file", &first)) + ); + let warm = cache + .try_get_with("file", async { panic!("warm read loaded") }) + .await + .unwrap(); + assert!(Arc::ptr_eq(&first, &warm)); + cache.remove("file").await; + cache.cache.run_pending_tasks().await; + assert!(cache.get_metadata("file").await.is_none()); + assert_eq!(sample_bytes(&gauge), 0); + } + + #[tokio::test] + async fn shared_failure_can_retry() { + let (cache, gauge) = new_cache(); + let calls = AtomicUsize::new(0); + let gate = tokio::sync::Notify::new(); + let load = || async { + calls.fetch_add(1, Ordering::Relaxed); + gate.notified().await; + InvalidPuffinFooterSnafu.fail() + }; + let mut first = Box::pin(cache.try_get_with("file", load())); + let mut second = Box::pin(cache.try_get_with("file", load())); + assert!(futures::poll!(&mut first).is_pending()); + assert!(futures::poll!(&mut second).is_pending()); + gate.notify_waiters(); + let (first, second) = tokio::join!(first, second); + assert!(first.is_err()); + assert!(second.is_err()); + assert_eq!(calls.load(Ordering::Relaxed), 1); + assert_eq!(sample_bytes(&gauge), 0); + assert!( + cache + .try_get_with("file", async { Ok(metadata()) }) + .await + .is_ok() + ); + } + + #[tokio::test] + async fn cancelled_loader_allows_waiter_to_load() { + let (cache, _) = new_cache(); + let mut first = Box::pin(cache.try_get_with("file", std::future::pending())); + assert!(futures::poll!(&mut first).is_pending()); + let mut second = Box::pin(cache.try_get_with("file", async { Ok(metadata()) })); + assert!(futures::poll!(&mut second).is_pending()); + drop(first); + tokio::time::timeout(std::time::Duration::from_secs(5), second) + .await + .unwrap() + .unwrap(); } } diff --git a/src/puffin/src/puffin_manager/fs_puffin_manager/reader.rs b/src/puffin/src/puffin_manager/fs_puffin_manager/reader.rs index b35b03875e4..77834ca2671 100644 --- a/src/puffin/src/puffin_manager/fs_puffin_manager/reader.rs +++ b/src/puffin/src/puffin_manager/fs_puffin_manager/reader.rs @@ -102,7 +102,8 @@ where async fn blob(&self, key: &str) -> Result> { let mut file = self.puffin_reader().await?; - let blob_metadata = self.get_blob_metadata(key, &mut file).await?; + let metadata = self.get_puffin_file_metadata(&mut file).await?; + let blob_metadata = Self::get_blob_metadata(key, &metadata)?; let blob = if blob_metadata.compression_codec.is_none() { // If the blob is not compressed, we can directly read it from the puffin file. Either::L(RandomReadBlob { @@ -132,7 +133,8 @@ where async fn dir(&self, key: &str) -> Result<(GuardWithMetadata, DirMetrics)> { let mut file = self.puffin_reader().await?; - let blob_metadata = self.get_blob_metadata(key, &mut file).await?; + let metadata = self.get_puffin_file_metadata(&mut file).await?; + let blob_metadata = Self::get_blob_metadata(key, &metadata)?; let (dir, metrics) = self .stager .get_dir( @@ -144,6 +146,7 @@ where let blob_metadata = blob_metadata.clone(); Box::pin(Self::init_dir_to_stager( file, + metadata, blob_metadata, handle, writer_provider, @@ -166,26 +169,18 @@ where &self, reader: &mut PuffinFileReader, ) -> Result> { - let id = self.handle.to_string(); - if let Some(cache) = self.puffin_file_metadata_cache.as_ref() - && let Some(metadata) = cache.get_metadata(&id) - { - return Ok(metadata); + if let Some(cache) = &self.puffin_file_metadata_cache { + cache + .try_get_with(&self.handle.to_string(), async { + reader.metadata().await.map(Arc::new) + }) + .await + } else { + reader.metadata().await.map(Arc::new) } - - let metadata = Arc::new(reader.metadata().await?); - if let Some(cache) = self.puffin_file_metadata_cache.as_ref() { - cache.put_metadata(id, metadata.clone()); - } - Ok(metadata) } - async fn get_blob_metadata( - &self, - key: &str, - file: &mut PuffinFileReader, - ) -> Result { - let metadata = self.get_puffin_file_metadata(file).await?; + fn get_blob_metadata(key: &str, metadata: &FileMetadata) -> Result { let blob_metadata = metadata .blobs .iter() @@ -218,12 +213,12 @@ where async fn init_dir_to_stager( mut file: PuffinFileReader, + puffin_metadata: Arc, blob_metadata: BlobMetadata, handle: F::FileHandle, writer_provider: DirWriterProviderRef, accessor: F, ) -> Result { - let puffin_metadata = file.metadata().await?; let reader = file.blob_reader(&blob_metadata)?; let meta = reader.metadata().await.context(MetadataSnafu)?; let buf = reader diff --git a/src/puffin/src/puffin_manager/tests.rs b/src/puffin/src/puffin_manager/tests.rs index f1ee9fabd70..a09c5696b41 100644 --- a/src/puffin/src/puffin_manager/tests.rs +++ b/src/puffin/src/puffin_manager/tests.rs @@ -383,3 +383,168 @@ async fn check_dir( assert_eq!(buf, *raw_data); } } + +#[tokio::test] +async fn test_metadata_cache_shared_by_independent_readers() { + use common_cache::metrics::{CacheMetrics, OperationMetrics}; + use prometheus::{Histogram, HistogramOpts, IntCounter, IntGauge, Opts}; + + use crate::puffin_manager::cache::PuffinMetadataCache; + + let operation = |name: &str| OperationMetrics { + started: IntCounter::new(format!("{name}_started"), "test").unwrap(), + succeeded: IntCounter::new(format!("{name}_succeeded"), "test").unwrap(), + failed: IntCounter::new(format!("{name}_failed"), "test").unwrap(), + cancelled: IntCounter::new(format!("{name}_cancelled"), "test").unwrap(), + active: IntGauge::new(format!("{name}_active"), "test").unwrap(), + elapsed: Histogram::with_opts(HistogramOpts::new(format!("{name}_seconds"), "test")) + .unwrap(), + }; + let metrics = Arc::new(CacheMetrics { + requests: operation("requests"), + loads: operation("loads"), + }); + let bytes = common_cache::metrics::CacheSizeMetrics::new( + Opts::new("metadata_bytes", "test"), + &["type"], + ) + .unwrap(); + let cache = + Arc::new(PuffinMetadataCache::new(1024 * 1024, &bytes).with_load_metrics(metrics.clone())); + let (_dir, stager) = new_bounded_stager("shared_metadata", 1024).await; + let accessor = Arc::new(MockFileAccessor::new("shared_metadata")); + let manager = + FsPuffinManager::new(stager, accessor).with_puffin_metadata_cache(Some(cache.clone())); + let file = "file".to_string(); + let mut writer = manager.writer(&file).await.unwrap(); + put_blob("blob", b"payload", None, &mut writer).await; + writer.finish().await.unwrap(); + + let requests = (0..8).map(|_| async { + manager + .reader(&file) + .await + .unwrap() + .metadata() + .await + .unwrap() + }); + let values = futures::future::join_all(requests).await; + assert!(values.iter().all(|value| Arc::ptr_eq(value, &values[0]))); + assert_eq!(8, metrics.requests.succeeded.get()); + assert_eq!(1, metrics.loads.succeeded.get()); + assert_eq!(0, metrics.loads.active.get()); + assert_eq!(1, values[0].blobs.len()); + assert_eq!("blob", values[0].blobs[0].blob_type); + + cache.remove(&file).await; + let reloaded = manager + .reader(&file) + .await + .unwrap() + .metadata() + .await + .unwrap(); + assert_eq!(values[0].as_ref(), reloaded.as_ref()); + assert!(!Arc::ptr_eq(&values[0], &reloaded)); + assert_eq!(2, metrics.loads.succeeded.get()); +} + +// After warming the metadata cache, reject footer reads while allowing blob reads. +struct FooterGuardAccessor { + inner: MockFileAccessor, + blob_end: Arc, +} + +struct FooterGuardReader { + inner: common_base::range_read::FileReader, + blob_end: Arc, +} + +#[async_trait::async_trait] +impl crate::puffin_manager::file_accessor::PuffinFileAccessor for FooterGuardAccessor { + type Reader = FooterGuardReader; + type Writer = + ::Writer; + type FileHandle = String; + + async fn reader(&self, handle: &String) -> crate::error::Result { + Ok(FooterGuardReader { + inner: self.inner.reader(handle).await?, + blob_end: self.blob_end.clone(), + }) + } + + async fn writer(&self, handle: &String) -> crate::error::Result { + self.inner.writer(handle).await + } +} + +#[async_trait::async_trait] +impl RangeReader for FooterGuardReader { + async fn metadata(&self) -> std::io::Result { + self.inner.metadata().await + } + + async fn read(&self, range: std::ops::Range) -> std::io::Result { + if range.end > self.blob_end.load(std::sync::atomic::Ordering::Relaxed) { + return Err(std::io::Error::other( + "unexpected footer read after cache warmup", + )); + } + self.inner.read(range).await + } +} + +impl common_base::range_read::SizeAwareRangeReader for FooterGuardReader { + fn with_file_size_hint(&mut self, size: u64) { + self.inner.with_file_size_hint(size); + } +} + +#[tokio::test] +async fn test_cached_metadata_used_when_staging_directory() { + use std::sync::atomic::{AtomicU64, Ordering}; + + use prometheus::Opts; + + use crate::puffin_manager::cache::PuffinMetadataCache; + + let bytes = common_cache::metrics::CacheSizeMetrics::new( + Opts::new("directory_metadata_bytes", "test"), + &["type"], + ) + .unwrap(); + let cache = Arc::new(PuffinMetadataCache::new(1024 * 1024, &bytes)); + let (_dir, stager) = new_bounded_stager("cached_directory", 1024 * 1024).await; + let blob_end = Arc::new(AtomicU64::new(u64::MAX)); + let accessor = Arc::new(FooterGuardAccessor { + inner: MockFileAccessor::new("cached_directory"), + blob_end: blob_end.clone(), + }); + let manager = + FsPuffinManager::new(stager.clone(), accessor).with_puffin_metadata_cache(Some(cache)); + let file = "directory".to_string(); + let mut writer = manager.writer(&file).await.unwrap(); + let files: &[(&str, &[u8])] = &[("a.txt", b"hello"), ("b.txt", b"world")]; + put_dir("dir", files, None, &mut writer).await; + writer.finish().await.unwrap(); + let metadata = manager + .reader(&file) + .await + .unwrap() + .metadata() + .await + .unwrap(); + blob_end.store( + metadata + .blobs + .iter() + .map(|blob| u64::try_from(blob.offset + blob.length).unwrap()) + .max() + .unwrap(), + Ordering::Relaxed, + ); + let reader = manager.reader(&file).await.unwrap(); + check_dir(&file, "dir", files, &stager, &reader).await; +} From f2d74650d9c0f548c965410ac7fadf454f6b2219 Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Fri, 30 Oct 2026 08:56:34 +0000 Subject: [PATCH 147/239] fix: optimize Prometheus remote write v2 ingestion (#9928) * fix: optimize Prometheus remote write v2 ingestion Signed-off-by: shuiyisong * fix: correct remote write v2 memory accounting and response counts Signed-off-by: shuiyisong --------- Signed-off-by: shuiyisong --- Cargo.lock | 1 + src/mito2/src/region_write_ctx.rs | 1 + src/operator/AGENTS.md | 2 + src/operator/src/bulk_insert.rs | 27 +- src/operator/src/insert.rs | 5 +- .../src/req_convert/common/partitioner.rs | 126 +++- src/operator/src/write_dispatcher.rs | 159 ++++- .../src/write_dispatcher/attribution.rs | 144 +++++ src/operator/src/write_dispatcher/tests.rs | 96 ++- src/servers/AGENTS.md | 2 + src/servers/benches/flush_batch_physical.rs | 5 +- src/servers/benches/prom_decode.rs | 91 +++ src/servers/src/batcher/logical_table.rs | 5 + .../src/batcher/logical_table/batch.rs | 19 +- .../batcher/logical_table/batch_convert.rs | 3 + .../batcher/logical_table/pending_worker.rs | 14 +- .../src/batcher/logical_table/region_write.rs | 38 +- src/servers/src/batcher/table.rs | 13 +- src/servers/src/batcher/table/batch.rs | 9 + .../src/batcher/table/pending_batch.rs | 2 + src/servers/src/http/memory_limit.rs | 19 +- src/servers/src/http/prom_store.rs | 598 ++++++++++++------ src/servers/src/prom_remote_write/mod.rs | 19 +- src/servers/src/prom_remote_write/v2.rs | 580 ++++++++++++++--- src/servers/src/prom_store.rs | 51 +- src/servers/src/request_memory_limiter.rs | 21 +- src/servers/src/row_writer.rs | 4 + src/servers/tests/http/prom_store_test.rs | 8 +- src/session/Cargo.toml | 1 + src/session/src/context.rs | 13 + src/session/src/lib.rs | 1 + src/session/src/write_progress.rs | 146 +++++ tests-integration/src/tests/instance_test.rs | 29 +- 33 files changed, 1898 insertions(+), 354 deletions(-) create mode 100644 src/operator/src/write_dispatcher/attribution.rs create mode 100644 src/session/src/write_progress.rs diff --git a/Cargo.lock b/Cargo.lock index a87b99721ff..a5e8e050a72 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -13848,6 +13848,7 @@ dependencies = [ "prost 0.14.1", "snafu 0.8.6", "sql", + "tokio", "uuid", ] diff --git a/src/mito2/src/region_write_ctx.rs b/src/mito2/src/region_write_ctx.rs index d37c42545a0..43083fa7744 100644 --- a/src/mito2/src/region_write_ctx.rs +++ b/src/mito2/src/region_write_ctx.rs @@ -398,6 +398,7 @@ impl RegionWriteCtx { } self.has_skip_wal_writes |= bulk.num_rows() > 0 && self.is_skip_wal_write(skip_wal); + // WriteAttribution requires the input count on success, before any deduplication. self.bulk_notifiers .push(WriteNotify::new(sender, bulk.num_rows())); diff --git a/src/operator/AGENTS.md b/src/operator/AGENTS.md index 4d31f44de96..d7d2bb1177e 100644 --- a/src/operator/AGENTS.md +++ b/src/operator/AGENTS.md @@ -41,6 +41,8 @@ cargo sqlness bare -t Check standalone and distributed routing when changing region dispatch. Write retries operate on prepared region requests in `write_dispatcher`; keep admission and Flow side effects outside retries, and preserve successful regions. +Process-local write attribution must follow the same row masks through batching +and retry repartitioning; only confirmed storage responses advance progress. Write admission uses fallible `write_meter!` calls in `Inserter::do_request` and `handle_bulk_insert`, before flow side effects and datanode dispatch. Bulk diff --git a/src/operator/src/bulk_insert.rs b/src/operator/src/bulk_insert.rs index d4a1ed0497c..3144daec3a0 100644 --- a/src/operator/src/bulk_insert.rs +++ b/src/operator/src/bulk_insert.rs @@ -41,7 +41,7 @@ use crate::error::Result; use crate::insert::Inserter; use crate::req_convert::TablePartitioner; use crate::req_convert::insert; -use crate::write_dispatcher::WriteAttempt; +use crate::write_dispatcher::{WriteAttempt, WriteAttribution}; use crate::{error, metrics}; impl Inserter { @@ -54,26 +54,31 @@ impl Inserter { table_info: TableInfoRef, batch: RecordBatch, ctx: QueryContextRef, + attribution: WriteAttribution, ) -> Result { let requests = TablePartitioner::try_new(&self.partition_manager, &table_info) .await .context(error::InvalidPartitionSnafu)? - .partition_bulk_insert_requests(batch, None, ctx.skip_wal())?; + .partition_attributed_bulk_requests(batch, ctx.skip_wal(), attribution)?; let writes = requests .into_iter() - .map(|request| RegionRequest { - header: Some(RegionRequestHeader { - dbname: ctx.get_db_string(), - tracing_context: TracingContext::from_current_span().to_w3c(), - ..Default::default() - }), - body: Some(region_request::Body::BulkInsert(request)), + .map(|(request, attribution)| { + ( + RegionRequest { + header: Some(RegionRequestHeader { + dbname: ctx.get_db_string(), + tracing_context: TracingContext::from_current_span().to_w3c(), + ..Default::default() + }), + body: Some(region_request::Body::BulkInsert(request)), + }, + attribution, + ) }) .collect(); - let writes = self.write_dispatcher.prepare(writes).await?; Ok(self .write_dispatcher - .execute(writes, None, None) + .dispatch_attributed(writes, None) .await .context(error::RequestInsertsSnafu)? .affected_rows) diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index 88cb565c779..3dabb963bdd 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -861,8 +861,11 @@ impl Inserter { // Dropping the caller stops retries without abandoning in-flight writes. let cancellation = CancellationToken::new(); let _guard = cancellation.clone().drop_guard(); + let progress = ctx.write_progress().cloned(); let affected_rows = common_runtime::spawn_global(async move { - dispatcher.execute(writes, None, Some(&cancellation)).await + dispatcher + .execute_with_progress(writes, None, Some(&cancellation), progress) + .await }) .await .context(JoinTaskSnafu)? diff --git a/src/operator/src/req_convert/common/partitioner.rs b/src/operator/src/req_convert/common/partitioner.rs index 6df73d7d1b9..05ad400abf0 100644 --- a/src/operator/src/req_convert/common/partitioner.rs +++ b/src/operator/src/req_convert/common/partitioner.rs @@ -37,6 +37,7 @@ use store_api::storage::{RegionId, RegionNumber, TableId}; use table::metadata::TableInfo; use crate::error::{self, Result, SplitDeleteSnafu, SplitInsertSnafu}; +use crate::write_dispatcher::{WriteAttribution, write_row_count}; /// Loads table routing rules and converts rows to region write requests. pub struct Partitioner<'a> { @@ -83,6 +84,101 @@ impl<'a> Partitioner<'a> { } impl TablePartitioner { + pub(crate) fn rebuild_attributed_write( + &self, + request: RegionRequest, + table_info: &TableInfo, + attribution: WriteAttribution, + ) -> Result> { + if !attribution.is_tracked() { + return Ok(self + .rebuild_write_requests(request, table_info)? + .into_iter() + .map(|request| (request, WriteAttribution::default())) + .collect()); + } + let RegionRequest { header, body } = request; + if let Some(Body::BulkInsert(bulk)) = body { + return Ok(self + .rebuild_bulk_insert_request(bulk, table_info, attribution)? + .into_iter() + .map(|(body, attribution)| { + ( + RegionRequest { + header: header.clone(), + body: Some(body), + }, + attribution, + ) + }) + .collect()); + } + self.rebuild_write_requests(RegionRequest { header, body }, table_info)? + .into_iter() + .map(|request| { + let attribution = attribution + .for_rows(write_row_count(&request)) + .map_err(BoxedError::new) + .context(error::ExternalSnafu)?; + Ok((request, attribution)) + }) + .collect() + } + + pub fn partition_attributed_bulk_requests( + &self, + batch: RecordBatch, + skip_wal: bool, + attribution: WriteAttribution, + ) -> Result> { + if !attribution.is_tracked() { + return Ok(self + .partition_bulk_insert_requests(batch, None, skip_wal)? + .into_iter() + .map(|request| (request, WriteAttribution::default())) + .collect()); + } + let masks = self + .rule + .split_record_batch(&batch) + .context(SplitInsertSnafu)?; + self.build_attributed_bulk_requests(batch, masks, None, skip_wal, attribution) + } + + fn build_attributed_bulk_requests( + &self, + batch: RecordBatch, + masks: HashMap, + original_ipc: Option, + skip_wal: bool, + attribution: WriteAttribution, + ) -> Result> { + let mut by_region = masks + .iter() + .filter(|(_, mask)| !mask.select_none()) + .map(|(region, mask)| { + Ok(( + *region, + attribution + .filter(mask) + .map_err(BoxedError::new) + .context(error::ExternalSnafu)?, + )) + }) + .collect::>>()?; + self.build_bulk_insert_requests(batch, masks, original_ipc, skip_wal)? + .into_iter() + .map(|request| { + let attribution = by_region + .remove(&RegionId::from_u64(request.region_id).region_number()) + .context(error::UnexpectedSnafu { + violated: "Missing write ownership after partitioning", + })?; + Ok((request, attribution)) + }) + .collect() + } + /// Reads a table's partition rule and expression versions into one routing snapshot. pub async fn try_new( partition_manager: &PartitionRuleManager, @@ -119,7 +215,11 @@ impl TablePartitioner { })? { Body::Inserts(inserts) => self.rebuild_insert_requests(inserts)?, Body::Deletes(deletes) => self.rebuild_delete_requests(deletes)?, - Body::BulkInsert(bulk) => self.rebuild_bulk_insert_request(bulk, table_info)?, + Body::BulkInsert(bulk) => self + .rebuild_bulk_insert_request(bulk, table_info, WriteAttribution::default())? + .into_iter() + .map(|(body, _)| body) + .collect(), _ => { return error::UnexpectedSnafu { violated: "Only writes can be rebuilt", @@ -278,10 +378,11 @@ impl TablePartitioner { &self, bulk: BulkInsertRequest, table_info: &TableInfo, - ) -> Result> { + attribution: WriteAttribution, + ) -> Result> { let region_id = RegionId::from_u64(bulk.region_id); if self.is_partition_expr_unchanged(region_id, bulk.partition_expr_version) { - return Ok(vec![Body::BulkInsert(bulk)]); + return Ok(vec![(Body::BulkInsert(bulk), attribution)]); } let bulk_insert_request::Body::ArrowIpc(ipc) = bulk.body.context(error::UnexpectedSnafu { @@ -296,10 +397,23 @@ impl TablePartitioner { .rule .split_record_batch(routing_batch.as_ref().unwrap_or(&batch)) .context(SplitInsertSnafu)?; - Ok(self - .build_bulk_insert_requests(batch, masks, Some(ipc), bulk.skip_wal)? + let writes = if attribution.is_tracked() { + self.build_attributed_bulk_requests( + batch, + masks, + Some(ipc), + bulk.skip_wal, + attribution, + )? + } else { + self.build_bulk_insert_requests(batch, masks, Some(ipc), bulk.skip_wal)? + .into_iter() + .map(|request| (request, WriteAttribution::default())) + .collect() + }; + Ok(writes .into_iter() - .map(Body::BulkInsert) + .map(|(request, attribution)| (Body::BulkInsert(request), attribution)) .collect()) } diff --git a/src/operator/src/write_dispatcher.rs b/src/operator/src/write_dispatcher.rs index 2d12feb74c5..b28eff9d550 100644 --- a/src/operator/src/write_dispatcher.rs +++ b/src/operator/src/write_dispatcher.rs @@ -12,7 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. +mod attribution; mod options; +pub use attribution::WriteAttribution; #[cfg(test)] mod tests; @@ -35,6 +37,7 @@ use common_meta::rpc::router; use futures::future; use partition::manager::{PartitionRuleManager, PartitionRuleManagerRef}; use prometheus::Histogram; +use session::write_progress::WriteProgress; use snafu::{IntoError, OptionExt, ResultExt, ensure}; use store_api::storage::RegionId; use table::table_name::TableName; @@ -110,6 +113,35 @@ impl WriteDispatcher { self.dispatch_inner(requests, Some(histogram)).await } + /// Dispatches shared batches with row ownership retained across retries. + pub async fn dispatch_attributed( + &self, + requests: Vec<(RegionRequest, WriteAttribution)>, + histogram: Option, + ) -> Result { + let regions = requests + .iter() + .map(|(request, _)| write_region_id(request)) + .collect::>>()?; + let peers = self + .resolve_peers(®ions) + .await + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + let writes = requests + .into_iter() + .zip(regions) + .map(|((request, attribution), region)| { + Ok((get_peer(&peers, region)?, request, attribution)) + }) + .collect::>>()?; + let attempts = future::join_all(writes.into_iter().map(|(peer, request, attribution)| { + self.send_attributed(peer, request, attribution, histogram.as_ref()) + })) + .await; + self.finish(attempts, histogram.as_ref(), None).await + } + /// Resolves routes before writes are detached from the caller. pub(crate) async fn prepare( &self, @@ -178,11 +210,23 @@ impl WriteDispatcher { writes: Vec<(PeerRef, RegionRequest)>, histogram: Option, cancellation: Option<&CancellationToken>, + ) -> Result { + self.execute_with_progress(writes, histogram, cancellation, None) + .await + } + + pub(crate) async fn execute_with_progress( + &self, + writes: Vec<(PeerRef, RegionRequest)>, + histogram: Option, + cancellation: Option<&CancellationToken>, + progress: Option>, ) -> Result { let histogram = histogram.as_ref(); + let progress = progress.as_ref(); let attempts = future::join_all(writes.into_iter().map(|(peer, request)| async move { - if self.retry.is_some() { - self.send_batch(peer, request, histogram).await + if self.retry.is_some() || progress.is_some() { + self.send_batch(peer, request, histogram, progress).await } else { vec![self.send(peer, request, histogram).await] } @@ -196,9 +240,20 @@ impl WriteDispatcher { /// Sends the first attempt immediately so bulk encoding can overlap other sends. pub(crate) async fn send( + &self, + peer: PeerRef, + request: RegionRequest, + histogram: Option<&Histogram>, + ) -> WriteAttempt { + self.send_attributed(peer, request, WriteAttribution::default(), histogram) + .await + } + + async fn send_attributed( &self, peer: PeerRef, mut request: RegionRequest, + attribution: WriteAttribution, histogram: Option<&Histogram>, ) -> WriteAttempt { if let Some(retry) = &self.retry { @@ -210,6 +265,7 @@ impl WriteDispatcher { RoutedWrite { peer, request: Arc::new(request), + attribution, }, histogram, ) @@ -218,7 +274,7 @@ impl WriteDispatcher { let _timer = histogram.map(Histogram::start_timer); let datanode = self.node_manager.datanode(&peer).await; match datanode.handle(request).await { - Ok(response) => WriteAttempt::Success(response), + Ok(response) => acknowledged(response, &attribution), Err(error) => WriteAttempt::Failed { error, request: None, @@ -267,6 +323,7 @@ impl WriteDispatcher { peer: PeerRef, mut request: RegionRequest, histogram: Option<&Histogram>, + progress: Option<&Arc>, ) -> Vec { if !matches!(&request.body, Some(Body::Inserts(_) | Body::Deletes(_))) { return vec![self.send(peer, request, histogram).await]; @@ -274,6 +331,40 @@ impl WriteDispatcher { // Upgrade datanodes first and roll back frontends first: old datanodes // cannot report which entries succeeded, so retrying a batch may replay writes. api::region::set_accept_partial_result(request.header.get_or_insert_default()); + if self.retry.is_none() { + // Standalone consumes owned requests; retaining an Arc prevents local dispatch. + let requester = self.node_manager.datanode(&peer).await; + let _timer = histogram.map(Histogram::start_timer); + return match requester.handle(request).await { + Ok(response) => { + if let Some(progress) = progress { + progress.add_written(response.affected_rows); + } + vec![WriteAttempt::Success(response)] + } + Err(Error::PartialRegion { response }) => { + if let Some(progress) = progress { + progress.add_written(response.affected_rows); + } + let mut attempts = vec![WriteAttempt::Success(RegionResponse { + affected_rows: response.affected_rows, + extensions: response.extensions, + metadata: Vec::new(), + })]; + attempts.extend(response.failures.into_iter().map(|(_, error)| { + WriteAttempt::Failed { + error, + request: None, + } + })); + attempts + } + Err(error) => vec![WriteAttempt::Failed { + error, + request: None, + }], + }; + } let request = Arc::new(request); let requester = self.node_manager.datanode(&peer).await; let result = { @@ -281,7 +372,12 @@ impl WriteDispatcher { requester.handle(Arc::clone(&request)).await }; let error = match result { - Ok(response) => return vec![WriteAttempt::Success(response)], + Ok(response) => { + if let Some(progress) = progress { + progress.add_written(response.affected_rows); + } + return vec![WriteAttempt::Success(response)]; + } Err(error) => error, }; let request = match Arc::try_unwrap(request) { @@ -299,6 +395,9 @@ impl WriteDispatcher { let requests = split_request(request); match error { Error::PartialRegion { response } => { + if let Some(progress) = progress { + progress.add_written(response.affected_rows); + } let mut failures: HashMap<_, _> = response.failures.into_iter().collect(); let mut attempts = vec![WriteAttempt::Success(RegionResponse { affected_rows: response.affected_rows, @@ -309,6 +408,7 @@ impl WriteDispatcher { if let Some(error) = failures.remove(&index) { let request = error.retry_hint().is_retryable().then(|| RoutedWrite { peer: Arc::clone(&peer), + attribution: direct_attribution(&request, progress.cloned()), request: Arc::new(request), }); attempts.push(WriteAttempt::Failed { error, request }); @@ -324,6 +424,7 @@ impl WriteDispatcher { .map(|request| { let request = error.retry_hint().is_retryable().then(|| RoutedWrite { peer: Arc::clone(&peer), + attribution: direct_attribution(&request, progress.cloned()), request: Arc::new(request), }); WriteAttempt::Failed { @@ -363,6 +464,7 @@ struct WriteRetry { pub(crate) struct RoutedWrite { peer: PeerRef, request: Arc, + attribution: WriteAttribution, } impl WriteRetry { @@ -370,7 +472,7 @@ impl WriteRetry { let _timer = histogram.map(Histogram::start_timer); let requester = self.node_manager.datanode(&write.peer).await; match requester.handle(Arc::clone(&write.request)).await { - Ok(response) => WriteAttempt::Success(response), + Ok(response) => acknowledged(response, &write.attribution), Err(error) => { let request = error.retry_hint().is_retryable().then_some(write); WriteAttempt::Failed { error, request } @@ -544,7 +646,7 @@ impl WriteRetry { err_msg: "Refreshed write has no partition snapshot", })?; let requests = partitioner - .rebuild_write_requests(request, table_info) + .rebuild_attributed_write(request, table_info, write.attribution) .map_err(BoxedError::new) .context(ExternalSnafu)?; rebuilt_requests.extend(requests); @@ -552,10 +654,13 @@ impl WriteRetry { self.prepare(rebuilt_requests).await } - async fn prepare(&self, requests: Vec) -> Result> { + async fn prepare( + &self, + requests: Vec<(RegionRequest, WriteAttribution)>, + ) -> Result> { let regions = requests .iter() - .map(write_region_id) + .map(|(request, _)| write_region_id(request)) .collect::>>()?; let peers = resolve_peers(&self.partition_manager, ®ions) .await @@ -564,10 +669,11 @@ impl WriteRetry { requests .into_iter() .zip(regions) - .map(|(request, region)| { + .map(|((request, attribution), region)| { Ok(RoutedWrite { peer: get_peer(&peers, region)?, request: Arc::new(request), + attribution, }) }) .collect() @@ -599,6 +705,41 @@ impl WriteRetry { } } +fn acknowledged(response: RegionResponse, attribution: &WriteAttribution) -> WriteAttempt { + match attribution.record(response.affected_rows) { + Ok(()) => WriteAttempt::Success(response), + Err(error) => WriteAttempt::Failed { + error, + request: None, + }, + } +} + +pub(crate) fn write_row_count(request: &RegionRequest) -> usize { + match &request.body { + Some(Body::Inserts(inserts)) => inserts + .requests + .iter() + .filter_map(|request| request.rows.as_ref()) + .map(|rows| rows.rows.len()) + .sum(), + Some(Body::Deletes(deletes)) => deletes + .requests + .iter() + .filter_map(|request| request.rows.as_ref()) + .map(|rows| rows.rows.len()) + .sum(), + _ => 0, + } +} + +fn direct_attribution( + request: &RegionRequest, + progress: Option>, +) -> WriteAttribution { + WriteAttribution::new(write_row_count(request), progress) +} + async fn resolve_peers( partition_manager: &PartitionRuleManager, regions: &[RegionId], diff --git a/src/operator/src/write_dispatcher/attribution.rs b/src/operator/src/write_dispatcher/attribution.rs new file mode 100644 index 00000000000..9fa23d36370 --- /dev/null +++ b/src/operator/src/write_dispatcher/attribution.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. + +use std::sync::Arc; + +use common_meta::error::{Result, UnexpectedSnafu}; +use partition::partition::RegionMask; +use session::write_progress::WriteProgress; +use snafu::ensure; + +/// Row ownership in payload order. Kept in process memory, never encoded in an RPC. +#[derive(Clone, Debug, Default)] +pub struct WriteAttribution { + spans: Vec<(usize, Option>)>, +} + +impl WriteAttribution { + pub fn new(rows: usize, progress: Option>) -> Self { + Self { + spans: vec![(rows, progress)], + } + } + + pub fn extend(&mut self, other: Self) { + self.spans.extend(other.spans); + } + + pub fn row_count(&self) -> usize { + self.spans.iter().map(|(rows, _)| rows).sum() + } + + pub fn is_tracked(&self) -> bool { + self.spans.iter().any(|(_, progress)| progress.is_some()) + } + + /// Selects ownership using the exact mask applied to the Arrow payload. + pub fn filter(&self, mask: &RegionMask) -> Result { + ensure!( + self.row_count() == mask.array().len(), + UnexpectedSnafu { + err_msg: "Write ownership does not match the partition mask", + } + ); + if !self.is_tracked() { + return Ok(Self::default()); + } + let mut offset = 0; + let mut selected = Self::default(); + for (rows, progress) in &self.spans { + let count = mask.array().slice(offset, *rows).true_count(); + if count > 0 { + selected.spans.push((count, progress.clone())); + } + offset += rows; + } + Ok(selected) + } + + /// Successful writes must acknowledge every input row, including duplicates, as + /// Mito's `RegionWriteCtx::push_bulk` does. A partial count cannot identify owners. + pub(crate) fn record(&self, affected_rows: usize) -> Result<()> { + if !self.is_tracked() { + return Ok(()); + } + ensure!( + affected_rows == self.row_count(), + UnexpectedSnafu { + err_msg: format!( + "Write affected {affected_rows} rows, expected {}; cannot attribute the result", + self.row_count() + ), + } + ); + for (rows, progress) in &self.spans { + if let Some(progress) = progress { + progress.add_written(*rows); + } + } + Ok(()) + } + + /// Direct row writes have one owner even when a refreshed route splits them. + pub(crate) fn for_rows(&self, rows: usize) -> Result { + ensure!( + self.spans.len() <= 1, + UnexpectedSnafu { + err_msg: "Cannot repartition shared ownership without a row mask", + } + ); + Ok(Self::new( + rows, + self.spans.first().and_then(|(_, p)| p.clone()), + )) + } +} + +#[cfg(test)] +mod tests { + use datatypes::arrow::array::BooleanArray; + + use super::*; + + #[test] + fn ownership_survives_filtering_and_repartition() { + let first = Arc::new(WriteProgress::default()); + let second = Arc::new(WriteProgress::default()); + let mut owners = WriteAttribution::new(3, Some(first.clone())); + owners.extend(WriteAttribution::new(2, Some(second.clone()))); + let initial = owners + .filter(&RegionMask::new( + BooleanArray::from(vec![true, false, true, false, true]), + 3, + )) + .unwrap(); + let retry = initial + .filter(&RegionMask::new( + BooleanArray::from(vec![false, true, true]), + 2, + )) + .unwrap(); + assert!(retry.record(1).is_err()); + assert_eq!(first.written(), 0); + retry.record(2).unwrap(); + assert_eq!(first.written(), 1); + assert_eq!(second.written(), 1); + assert!(owners.for_rows(5).is_err()); + assert!( + owners + .filter(&RegionMask::new(BooleanArray::from(vec![true]), 1)) + .is_err() + ); + } +} diff --git a/src/operator/src/write_dispatcher/tests.rs b/src/operator/src/write_dispatcher/tests.rs index 2607a3f194d..52954355cc6 100644 --- a/src/operator/src/write_dispatcher/tests.rs +++ b/src/operator/src/write_dispatcher/tests.rs @@ -416,6 +416,59 @@ fn build_insert_request(regions: &[u32]) -> RegionRequest { } } +#[tokio::test] +async fn test_direct_partial_progress_with_and_without_retries() { + for max_retries in [0, 1] { + let server = RegionTestServer::start().await; + server.fail_region(1, [RetryHint::Retryable; 2]); + let dispatcher = server.create_dispatcher(max_retries).await; + let mut request = build_insert_request(&[1, 2]); + if let Some(Body::Inserts(inserts)) = &mut request.body { + for insert in &mut inserts.requests { + insert.rows = Some(Rows { + rows: vec![Row::default()], + ..Default::default() + }); + } + } + let writes = dispatcher.prepare(vec![request]).await.unwrap(); + let progress = Arc::new(WriteProgress::default()); + assert!( + dispatcher + .execute_with_progress(writes, None, None, Some(progress.clone())) + .await + .is_err() + ); + assert_eq!(progress.written(), 1); + assert_eq!(server.attempt_counts()[&RegionId::new(1, 2).as_u64()], 1); + } +} + +#[tokio::test] +async fn test_shared_batch_partial_progress_credits_only_its_owners() { + for max_retries in [0, 1] { + let server = RegionTestServer::start().await; + server.fail_region(1, [RetryHint::Retryable; 2]); + let dispatcher = server.create_dispatcher(max_retries).await; + let first = Arc::new(WriteProgress::default()); + let second = Arc::new(WriteProgress::default()); + let writes = vec![ + ( + build_partitioned_request(WriteKind::BulkInsert, 1, None, &[70]), + WriteAttribution::new(1, Some(first.clone())), + ), + ( + build_partitioned_request(WriteKind::BulkInsert, 2, None, &[20]), + WriteAttribution::new(1, Some(second.clone())), + ), + ]; + assert!(dispatcher.dispatch_attributed(writes, None).await.is_err()); + assert_eq!(first.written(), 0); + assert_eq!(second.written(), 1); + assert_eq!(server.attempt_counts()[&RegionId::new(1, 2).as_u64()], 1); + } +} + #[tokio::test] async fn test_write_retry_disabled_and_enabled_ablation() { common_telemetry::init_default_ut_logging(); @@ -628,6 +681,7 @@ async fn test_should_retry_requires_a_retained_request_and_budget() { }; let retryable = || { failure(Some(RoutedWrite { + attribution: Default::default(), peer: Arc::new(Peer::new(1, "peer")), request: Arc::new(build_insert_request(&[1])), })) @@ -1037,10 +1091,18 @@ fn request_parts(request: &RegionRequest) -> (u32, Option, Vec) { #[tokio::test] async fn test_write_retry_rebuilds_only_failed_rows_after_route_refresh() { - for kind in [WriteKind::Insert, WriteKind::Delete, WriteKind::BulkInsert] { + for (kind, fail_repartitioned) in [ + (WriteKind::Insert, false), + (WriteKind::Delete, false), + (WriteKind::BulkInsert, false), + (WriteKind::BulkInsert, true), + ] { let old_server = RegionTestServer::start().await; let new_server = RegionTestServer::start().await; old_server.fail_region(2, [RetryHint::RetryAfterRouteRefresh]); + if fail_repartitioned { + new_server.fail_region(5, [RetryHint::NonRetryable]); + } let backend = test_util::prepare_mocked_backend().await; test_util::create_partition_rule_manager(backend.clone()).await; update_table_leader(&backend, &old_server.peer).await; @@ -1101,14 +1163,30 @@ async fn test_write_retry_rebuilds_only_failed_rows_after_route_refresh() { catalog_manager, invalidator.clone(), ); - let response = dispatcher - .dispatch(vec![ - build_partitioned_request(kind, 1, old_versions[&1], &[70]), - build_partitioned_request(kind, 2, old_versions[&2], &[20, 40]), - ]) - .await - .unwrap(); - assert_eq!(response.affected_rows, 3, "{kind:?}"); + let first = Arc::new(WriteProgress::default()); + let second = Arc::new(WriteProgress::default()); + let requests = vec![ + build_partitioned_request(kind, 1, old_versions[&1], &[70]), + build_partitioned_request(kind, 2, old_versions[&2], &[20, 40]), + ]; + let response = if matches!(kind, WriteKind::BulkInsert) { + let mut shared = WriteAttribution::new(1, Some(first.clone())); + shared.extend(WriteAttribution::new(1, Some(second.clone()))); + let owners = [WriteAttribution::new(1, Some(first.clone())), shared]; + let response = dispatcher + .dispatch_attributed(requests.into_iter().zip(owners).collect(), None) + .await; + assert_eq!(first.written(), 2); + assert_eq!(second.written(), if fail_repartitioned { 0 } else { 1 }); + response + } else { + dispatcher.dispatch(requests).await + }; + if fail_repartitioned { + assert!(response.is_err()); + } else { + assert_eq!(response.unwrap().affected_rows, 3, "{kind:?}"); + } assert_eq!( old_server.attempt_counts(), HashMap::from([ diff --git a/src/servers/AGENTS.md b/src/servers/AGENTS.md index 6ee7675313c..9556c1a9909 100644 --- a/src/servers/AGENTS.md +++ b/src/servers/AGENTS.md @@ -43,6 +43,8 @@ cargo nextest run -p servers Keep protocol translation here and permissions/database behavior in frontend. Async table batching returns storage-completion futures with its early acknowledgements. Keep completion reporting independent of the caller's lifetime. +Remote-write v2 keeps decoder reservations and per-request progress through queued +writes; error responses drain those writes before reporting partial counts. Logical-table physical flushes retain partition columns in Arrow RPC payloads and pass all region writes to the operator write dispatcher together; the dispatcher owns peer resolution and retries. See `batcher/logical_table/region_write.rs`. diff --git a/src/servers/benches/flush_batch_physical.rs b/src/servers/benches/flush_batch_physical.rs index d879486c2d3..e4933a9b371 100644 --- a/src/servers/benches/flush_batch_physical.rs +++ b/src/servers/benches/flush_batch_physical.rs @@ -101,7 +101,10 @@ struct BenchWriteDispatcher; #[async_trait] impl PhysicalFlushWriteDispatcher for BenchWriteDispatcher { - async fn dispatch(&self, _requests: Vec) -> error::Result { + async fn dispatch( + &self, + _requests: Vec<(RegionRequest, operator::write_dispatcher::WriteAttribution)>, + ) -> error::Result { Ok(0) } } diff --git a/src/servers/benches/prom_decode.rs b/src/servers/benches/prom_decode.rs index bbca7a4269f..62888550597 100644 --- a/src/servers/benches/prom_decode.rs +++ b/src/servers/benches/prom_decode.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::collections::HashMap; +use std::hash::{DefaultHasher, Hash, Hasher}; use std::time::Duration; use api::greptime_proto::io::prometheus::write::v2 as write_v2; @@ -303,6 +304,95 @@ fn bench_prom_v1_vs_v2_decode(c: &mut Criterion) { } } +/// Exercises the production v2 decoder, including validation and row allocation. +#[allow(clippy::print_stderr)] +fn bench_prom_v2_histograms(c: &mut Criterion) { + use write_v2::histogram::{Count, ZeroCount}; + + let mut group = c.benchmark_group("prom_rw_v2_histograms"); + group.sample_size(50); + group.measurement_time(Duration::from_secs(5)); + + for (name, points_per_series, buckets, float, custom) in [ + ("int_1k_series_1_point_8_buckets", 1, 8, false, false), + ("int_1k_series_10_points_8_buckets", 10, 8, false, false), + ("int_1k_series_1_point_128_buckets", 1, 128, false, false), + ("float_1k_series_1_point_8_buckets", 1, 8, true, false), + ("float_1k_series_10_points_8_buckets", 10, 8, true, false), + ("custom_1k_series_1_point_8_buckets", 1, 8, false, true), + ] { + let (_, mut request) = generate_sample_workload(1_000, points_per_series, 8); + for series in &mut request.timeseries { + series.histograms = series + .samples + .drain(..) + .map(|sample| { + let mut histogram = write_v2::Histogram { + schema: if custom { -53 } else { 0 }, + timestamp: sample.timestamp, + start_timestamp: sample.timestamp - 1_000, + sum: buckets as f64, + positive_spans: vec![write_v2::BucketSpan { + offset: 0, + length: buckets, + }], + ..Default::default() + }; + if float { + histogram.count = Some(Count::CountFloat(buckets as f64)); + histogram.zero_count = Some(ZeroCount::ZeroCountFloat(0.0)); + histogram.positive_counts = vec![1.0; buckets as usize]; + } else { + histogram.count = Some(Count::CountInt(buckets as u64)); + histogram.zero_count = Some(ZeroCount::ZeroCountInt(0)); + histogram.positive_deltas = vec![0; buckets as usize]; + histogram.positive_deltas[0] = 1; + } + if custom { + histogram.custom_values = (1..buckets).map(f64::from).collect(); + } + histogram + }) + .collect(); + } + let bytes = request.encode_to_vec(); + let expected_rows = (1_000 * points_per_series) as u64; + let (samples, mut histograms, sample_count, histogram_count) = + remote_write_v2::decode_uncompressed_write_requests(&bytes).unwrap(); + assert!(samples.is_empty()); + assert_eq!(sample_count, 0); + assert_eq!(histogram_count, expected_rows); + assert_eq!( + histograms + .iter() + .map(|insert| insert.rows.as_ref().unwrap().rows.len() as u64) + .sum::(), + expected_rows + ); + // Stable table order makes before/after output fingerprints comparable. + histograms.sort_unstable_by(|a, b| a.table_name.cmp(&b.table_name)); + let mut fingerprint = DefaultHasher::new(); + for insert in &histograms { + insert.encode_to_vec().hash(&mut fingerprint); + } + eprintln!( + "{name}: {} bytes, {expected_rows} rows, fingerprint {:016x}", + bytes.len(), + fingerprint.finish() + ); + drop(histograms); + group.throughput(Throughput::Elements(expected_rows)); + group.bench_with_input(BenchmarkId::from_parameter(name), &bytes, |b, bytes| { + b.iter(|| { + black_box( + remote_write_v2::decode_uncompressed_write_requests(black_box(bytes)).unwrap(), + ); + }); + }); + } + group.finish(); +} + /// Benchmark comparing UTF-8 string validation (`decode_string`) vs /// direct byte-level Prometheus label name validation (`decode_label_name`). fn bench_label_name_validation(c: &mut Criterion) { @@ -395,6 +485,7 @@ criterion_group!( benches, bench_decode_prom_request, bench_prom_v1_vs_v2_decode, + bench_prom_v2_histograms, bench_label_name_validation, bench_utf8_validation ); diff --git a/src/servers/src/batcher/logical_table.rs b/src/servers/src/batcher/logical_table.rs index f34b0ae406e..1391d85ee3f 100644 --- a/src/servers/src/batcher/logical_table.rs +++ b/src/servers/src/batcher/logical_table.rs @@ -316,9 +316,11 @@ impl LogicalTablePendingRowsBatcher { }; let (response_tx, response_rx) = oneshot::channel(); + let progress = ctx.write_progress().cloned(); let batch_key = batch_key_from_ctx(&ctx); let mut cmd = Some(WorkerCommand::Submit { + _completion: progress.as_ref().map(|progress| progress.start()), table_batches, total_rows, ctx, @@ -369,6 +371,9 @@ impl LogicalTablePendingRowsBatcher { .context(error::SubmitBatchSnafu) .map(|()| (total_rows as u64, prepared)) } else { + if let Some(progress) = progress { + progress.mark_accepted(); + } Ok((total_rows as u64, prepared)) } } diff --git a/src/servers/src/batcher/logical_table/batch.rs b/src/servers/src/batcher/logical_table/batch.rs index 1ce90f39d01..b6200b1d7c9 100644 --- a/src/servers/src/batcher/logical_table/batch.rs +++ b/src/servers/src/batcher/logical_table/batch.rs @@ -237,6 +237,13 @@ pub async fn flush_batch_physical( // 4. Concatenate all modified batches (all share the same physical schema) let combined_batch = batch_convert::concat_modified_batches(&modified_batches)?; + let mut attribution = operator::write_dispatcher::WriteAttribution::default(); + for batch in table_batches.iter().flat_map(|table| &table.batches) { + attribution.extend(operator::write_dispatcher::WriteAttribution::new( + batch.batch.num_rows(), + batch.progress.clone(), + )); + } // 5. Split by physical partition rule and send to regions let physical_table_id = physical_table_info.table_id(); @@ -244,6 +251,7 @@ pub async fn flush_batch_physical( combined_batch, physical_table_id, partition_rule.as_ref(), + attribution, )?; let region_writes = region_write::encode_region_write_requests( @@ -428,6 +436,7 @@ mod tests { let limiter = RequestLimiter::try_new(1).unwrap(); let (response_tx, response_rx) = oneshot::channel(); let waiter = FlushWaiter { + _completion: None, response_tx, _permit: limiter.acquire().await.unwrap(), }; @@ -679,7 +688,10 @@ mod tests { #[async_trait] impl PhysicalFlushWriteDispatcher for MockFlushWriteDispatcher { - async fn dispatch(&self, _requests: Vec) -> error::Result { + async fn dispatch( + &self, + _requests: Vec<(RegionRequest, operator::write_dispatcher::WriteAttribution)>, + ) -> error::Result { self.writes.fetch_add(1, Ordering::SeqCst); if self.fail { return Err(Error::Internal { @@ -706,7 +718,10 @@ mod tests { #[async_trait] impl PhysicalFlushWriteDispatcher for AffectedRowsFlushWriteDispatcher { - async fn dispatch(&self, _requests: Vec) -> error::Result { + async fn dispatch( + &self, + _requests: Vec<(RegionRequest, operator::write_dispatcher::WriteAttribution)>, + ) -> error::Result { Ok(self.affected_rows) } } diff --git a/src/servers/src/batcher/logical_table/batch_convert.rs b/src/servers/src/batcher/logical_table/batch_convert.rs index 65204ec038c..322b12488e8 100644 --- a/src/servers/src/batcher/logical_table/batch_convert.rs +++ b/src/servers/src/batcher/logical_table/batch_convert.rs @@ -32,6 +32,8 @@ use crate::metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED; /// An aligned logical record batch and its timestamp column index. #[derive(Debug, Clone)] pub struct RecordBatchWithTsIdx { + pub(in crate::batcher::logical_table) progress: + Option>, /// The aligned logical record batch. pub(in crate::batcher::logical_table) batch: RecordBatch, /// The timestamp column index in `batch`. @@ -63,6 +65,7 @@ impl RecordBatchWithTsIdx { ); Ok(Self { + progress: None, batch, timestamp_index, }) diff --git a/src/servers/src/batcher/logical_table/pending_worker.rs b/src/servers/src/batcher/logical_table/pending_worker.rs index ff363c7637c..9f9c689e6b0 100644 --- a/src/servers/src/batcher/logical_table/pending_worker.rs +++ b/src/servers/src/batcher/logical_table/pending_worker.rs @@ -43,6 +43,8 @@ pub(in crate::batcher::logical_table) struct PendingBatch { } pub(in crate::batcher::logical_table) struct FlushWaiter { + pub(in crate::batcher::logical_table) _completion: + Option, pub(in crate::batcher::logical_table) response_tx: oneshot::Sender>>, pub(in crate::batcher::logical_table) _permit: Arc, @@ -56,6 +58,7 @@ pub(in crate::batcher::logical_table) struct PendingWorker { /// Commands accepted by a logical-table batching worker. pub(crate) enum WorkerCommand { Submit { + _completion: Option, table_batches: Vec<(String, u32, RecordBatchWithTsIdx)>, total_rows: usize, ctx: QueryContextRef, @@ -122,20 +125,22 @@ pub(crate) fn start_worker( tokio::select! { cmd = rx.recv() => { match cmd { - Some(WorkerCommand::Submit { table_batches, total_rows, ctx, response_tx, _permit }) => { + Some(WorkerCommand::Submit { table_batches, total_rows, ctx, response_tx, _permit, _completion }) => { let submitted_at = tokio::time::Instant::now(); idle_timer.as_mut().reset(submitted_at + worker_idle_timeout); pending_flush.submit( - FlushWaiter { response_tx, _permit }, + FlushWaiter { response_tx, _permit, _completion }, total_rows, ); + let progress = ctx.write_progress().cloned(); let pending_batch = batch.get_or_insert_with(||{ PENDING_BATCHES.inc(); PendingBatch::new(ctx) }); - for (table_name, table_id, record_batch) in table_batches { + for (table_name, table_id, mut record_batch) in table_batches { + record_batch.progress = progress.clone(); pending_batch.add_table_batch(table_name, table_id, record_batch); } @@ -316,6 +321,7 @@ mod tests { PendingCore::new(TimingFlushPolicy::try_new(Duration::from_secs(10), 1).unwrap()); pending_flush.submit( FlushWaiter { + _completion: None, response_tx, _permit: Arc::new(permit), }, @@ -362,6 +368,7 @@ mod tests { let (response_tx, mut response_rx) = oneshot::channel(); pending_flush.submit( FlushWaiter { + _completion: None, response_tx, _permit: Arc::new(semaphore.clone().acquire_owned().await.unwrap()), }, @@ -614,6 +621,7 @@ mod tests { let permit = Arc::new(Semaphore::new(1)).acquire_owned().await.unwrap(); worker_tx .send(WorkerCommand::Submit { + _completion: None, table_batches: vec![( "cpu".to_string(), 42, diff --git a/src/servers/src/batcher/logical_table/region_write.rs b/src/servers/src/batcher/logical_table/region_write.rs index 9511d21a998..e64cced9d54 100644 --- a/src/servers/src/batcher/logical_table/region_write.rs +++ b/src/servers/src/batcher/logical_table/region_write.rs @@ -25,7 +25,7 @@ use catalog::CatalogManagerRef; use common_grpc::error::Error as GrpcError; use common_grpc::flight; use common_telemetry::tracing_context::TracingContext; -use operator::write_dispatcher::WriteDispatcherRef; +use operator::write_dispatcher::{WriteAttribution, WriteDispatcherRef}; use partition::manager::PartitionRuleManagerRef; use partition::partition::PartitionRuleRef; use snafu::ResultExt; @@ -68,7 +68,7 @@ pub trait PhysicalFlushPartitionProvider: Send + Sync { #[async_trait] pub trait PhysicalFlushWriteDispatcher: Send + Sync { /// Dispatches one physical flush together, sharing routing refreshes across its writes. - async fn dispatch(&self, requests: Vec) -> Result; + async fn dispatch(&self, requests: Vec<(RegionRequest, WriteAttribution)>) -> Result; } #[derive(Clone)] @@ -126,12 +126,14 @@ pub(crate) struct WriteDispatcherPhysicalFlushAdapter { #[async_trait] impl PhysicalFlushWriteDispatcher for WriteDispatcherPhysicalFlushAdapter { - async fn dispatch(&self, requests: Vec) -> Result { + async fn dispatch(&self, requests: Vec<(RegionRequest, WriteAttribution)>) -> Result { self.write_dispatcher - .dispatch_with_metrics( + .dispatch_attributed( requests, - metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED - .with_label_values(&["flush_write_region"]), + Some( + metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED + .with_label_values(&["flush_write_region"]), + ), ) .await .context(error::CommonMetaSnafu) @@ -141,6 +143,7 @@ impl PhysicalFlushWriteDispatcher for WriteDispatcherPhysicalFlushAdapter { /// Arrow rows selected for one physical region. pub(crate) struct PlannedRegionBatch { + pub(crate) attribution: WriteAttribution, pub(crate) region_id: RegionId, pub(crate) batch: RecordBatch, } @@ -169,6 +172,7 @@ pub(crate) fn plan_region_batch( physical_table_id: TableId, region_number: u32, mask: &partition::partition::RegionMask, + attribution: &WriteAttribution, ) -> Result> { if mask.select_none() { return Ok(None); @@ -189,6 +193,7 @@ pub(crate) fn plan_region_batch( } Ok(Some(PlannedRegionBatch { + attribution: attribution.filter(mask).context(error::CommonMetaSnafu)?, region_id: RegionId::new(physical_table_id, region_number), batch: region_batch, })) @@ -199,6 +204,7 @@ pub(crate) fn plan_region_batches( combined_batch: RecordBatch, physical_table_id: TableId, partition_rule: &dyn partition::partition::PartitionRule, + attribution: WriteAttribution, ) -> Result> { let region_masks = split_combined_batch_by_region(&combined_batch, partition_rule)?; @@ -206,9 +212,13 @@ pub(crate) fn plan_region_batches( // Mito's sparse bulk conversion projects the stored fields and encoded primary key. let mut planned_batches = Vec::new(); for (region_number, mask) in region_masks { - if let Some(planned_batch) = - plan_region_batch(&combined_batch, physical_table_id, region_number, &mask)? - { + if let Some(planned_batch) = plan_region_batch( + &combined_batch, + physical_table_id, + region_number, + &mask, + &attribution, + )? { planned_batches.push(planned_batch); } } @@ -221,7 +231,7 @@ pub(crate) fn encode_region_write_requests( planned_batches: Vec, skip_wal: bool, partition_versions: &HashMap>, -) -> Result> { +) -> Result> { let mut region_writes = Vec::with_capacity(planned_batches.len()); for planned in planned_batches { let region_id = planned.region_id; @@ -256,7 +266,7 @@ pub(crate) fn encode_region_write_requests( })), }; - region_writes.push(request); + region_writes.push((request, planned.attribution)); } Ok(region_writes) @@ -334,6 +344,7 @@ mod tests { &TwoRegionPartitionRule { partition_columns: vec!["host".to_string()], }, + operator::write_dispatcher::WriteAttribution::new(2, None), ) .unwrap(); planned_batches.sort_by_key(|planned| planned.region_id.region_number()); @@ -352,7 +363,7 @@ mod tests { &HashMap::from([(1, Some(7)), (2, Some(8))]), ) .unwrap(); - for (index, write) in writes.into_iter().enumerate() { + for (index, (write, _)) in writes.into_iter().enumerate() { let Some(region_request::Body::BulkInsert(request)) = write.body else { panic!("expected bulk") }; @@ -426,6 +437,7 @@ mod tests { fn check_encode_region_write_requests(skip_wal: bool) { let planned_batch = PlannedRegionBatch { + attribution: operator::write_dispatcher::WriteAttribution::new(1, None), region_id: RegionId::new(1024, 1), batch: RecordBatch::try_new( Arc::new(ArrowSchema::new(vec![ @@ -453,7 +465,7 @@ mod tests { .unwrap(); assert_eq!(1, writes.len()); - let Some(region_request::Body::BulkInsert(request)) = &writes[0].body else { + let Some(region_request::Body::BulkInsert(request)) = &writes[0].0.body else { panic!("expected bulk insert request"); }; assert_eq!(RegionId::new(1024, 1).as_u64(), request.region_id); diff --git a/src/servers/src/batcher/table.rs b/src/servers/src/batcher/table.rs index daae05251f8..b2228ecb057 100644 --- a/src/servers/src/batcher/table.rs +++ b/src/servers/src/batcher/table.rs @@ -188,7 +188,9 @@ impl PendingRowsBatcher for TablePendingRowsBatcher { } let key = batch_key_from_ctx(&table_info.name, &ctx); let (response_tx, response_rx) = oneshot::channel(); + let progress = ctx.write_progress().cloned(); let pending = PendingBatch { + _completion: progress.as_ref().map(|progress| progress.start()), table_info, batch, ctx, @@ -231,6 +233,9 @@ impl PendingRowsBatcher for TablePendingRowsBatcher { if self.pending_rows_batch_sync { completion.await?; } else { + if let Some(progress) = progress { + progress.mark_accepted(); + } output .meta .write_completions @@ -379,7 +384,10 @@ mod tests { ); Arc::get_mut(&mut batcher).unwrap().pending_rows_batch_sync = sync; let table = Arc::new(new_test_table_info(1, "ack", [0].into_iter())); - let ctx = QueryContext::arc(); + let progress = Arc::new(session::write_progress::WriteProgress::default()); + let mut ctx = QueryContext::with("greptime", "public"); + ctx.set_write_progress(progress.clone()); + let ctx = Arc::new(ctx); let key = batch_key_from_ctx(&table.name, &ctx); // Hold the worker command to control completion independently of scheduling. let (_, receiver) = batcher.workers.get_or_create(key, 1).await; @@ -414,6 +422,8 @@ mod tests { assert!(futures::poll!(pending.as_mut()).is_pending()); completion = Some(pending); } + assert_eq!(progress.accepted(), !sync); + assert_eq!(progress.written(), 0); // Early acknowledgement must not release capacity before completion. let acquire = batcher.acquire(); tokio::pin!(acquire); @@ -827,6 +837,7 @@ mod tests { let (response_tx, response_rx) = oneshot::channel(); self.tx .send(WorkerCommand::Submit(PendingBatch { + _completion: None, table_info, batch, ctx: QueryContext::arc(), diff --git a/src/servers/src/batcher/table/batch.rs b/src/servers/src/batcher/table/batch.rs index e9ab339a6ac..42dfd52f82e 100644 --- a/src/servers/src/batcher/table/batch.rs +++ b/src/servers/src/batcher/table/batch.rs @@ -21,6 +21,7 @@ use common_telemetry::error; use operator::error::{ComputeArrowSnafu, Result, UnexpectedSnafu}; use operator::insert::Inserter; use operator::metrics::DIST_INGEST_ROW_COUNT; +use operator::write_dispatcher::WriteAttribution; use snafu::ResultExt; use table::metadata::TableInfoRef; @@ -72,11 +73,19 @@ async fn send_batch(batch: &Batch, inserter: &Inserter) -> Result<(TableInfoRef, .build() })?; let combined = combine_batches(batch.submissions.iter().map(|pending| &pending.batch))?; + let mut attribution = WriteAttribution::default(); + for pending in &batch.submissions { + attribution.extend(WriteAttribution::new( + pending.batch.num_rows(), + pending.ctx.write_progress().cloned(), + )); + } let affected_rows = inserter .flush_bulk_batch( first.table_info.clone(), combined.clone(), first.ctx.clone(), + attribution, ) .await?; validate_affected_rows(batch.total_rows, affected_rows)?; diff --git a/src/servers/src/batcher/table/pending_batch.rs b/src/servers/src/batcher/table/pending_batch.rs index 4f82865da55..580407e6b5f 100644 --- a/src/servers/src/batcher/table/pending_batch.rs +++ b/src/servers/src/batcher/table/pending_batch.rs @@ -27,6 +27,7 @@ pub(in crate::batcher::table) struct PendingBatch { pub ctx: QueryContextRef, pub response_tx: oneshot::Sender>>, pub _permit: Arc, + pub _completion: Option, } pub(in crate::batcher::table) fn notify_batches( @@ -59,6 +60,7 @@ mod tests { let (response_tx, response_rx) = oneshot::channel(); let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); batches.push(PendingBatch { + _completion: None, table_info: Arc::new(new_test_table_info(1, "t", [0].into_iter())), batch: RecordBatch::try_new( schema, diff --git a/src/servers/src/http/memory_limit.rs b/src/servers/src/http/memory_limit.rs index 41905616f44..75eda067223 100644 --- a/src/servers/src/http/memory_limit.rs +++ b/src/servers/src/http/memory_limit.rs @@ -25,7 +25,7 @@ //! would otherwise be admitted for free), and a client can understate the //! header while sending a much larger body. //! -//! Permits acquired while streaming are held by [`BodyMemoryAccounting`], +//! Upfront and streaming permits are held by [`BodyMemoryAccounting`], //! which is also inserted into the request extensions so they stay alive until //! the whole request is finished (extractors collect the body into owned //! buffers that outlive the body stream itself). @@ -67,7 +67,7 @@ pub async fn memory_limit_middleware( .and_then(|v| v.parse::().ok()) .unwrap_or(0); - let _guard = match limiter.acquire(content_length).await { + let guard = match limiter.acquire(content_length).await { Ok(guard) => guard, Err(e) => { return ( @@ -87,6 +87,7 @@ pub async fn memory_limit_middleware( // Account whatever is actually streamed beyond the upfront reservation // (chunked transfer has none). let accounting = BodyMemoryAccounting::default(); + accounting.hold(guard); // The middleware keeps its own handle so the acquired permits stay held // for the whole request: the body wrapper and the extensions are dropped // once the extractors finish collecting the body, before the handler is @@ -168,13 +169,13 @@ pub async fn decoded_body_accounting_middleware( quota_exceeded_response(&retained_accounting, response) } -/// Holds the memory guards acquired while a request body is streamed. Shared +/// Holds the memory guards acquired upfront and while a request body is streamed. Shared /// between the [`AccountedBody`] wrapper, the request extensions and the /// middleware itself, so the permits are only released when the request /// (including any collected body buffers and the handler still using them) /// is finished. #[derive(Clone, Default)] -struct BodyMemoryAccounting { +pub(crate) struct BodyMemoryAccounting { guards: Arc>>>, /// Set when an incremental charge hit the quota, so the rejection can be /// rewritten into a 429 after the middleware regains control. @@ -182,6 +183,16 @@ struct BodyMemoryAccounting { } impl BodyMemoryAccounting { + /// Includes permit rounding and all frames collected so far. + pub(crate) fn reserved_bytes(&self) -> u64 { + self.guards + .lock() + .unwrap() + .iter() + .map(MemoryGuard::granted_bytes) + .sum() + } + fn hold(&self, guard: MemoryGuard) { self.guards.lock().unwrap().push(guard); } diff --git a/src/servers/src/http/prom_store.rs b/src/servers/src/http/prom_store.rs index 05d9725f947..5a1be58502e 100644 --- a/src/servers/src/http/prom_store.rs +++ b/src/servers/src/http/prom_store.rs @@ -20,7 +20,7 @@ use async_trait::async_trait; use axum::Extension; use axum::body::Bytes; use axum::extract::{Query, State}; -use axum::http::{HeaderMap, HeaderValue, StatusCode, header}; +use axum::http::{Extensions, HeaderMap, HeaderValue, StatusCode, header}; use axum::response::IntoResponse; use axum_extra::TypedHeader; use common_catalog::consts::DEFAULT_SCHEMA_NAME; @@ -46,6 +46,7 @@ use crate::http::extractor::PipelineInfo; use crate::http::header::{ CONTENT_TYPE_PROTOBUF_STR, GREPTIME_DB_HEADER_METRICS, write_cost_header_map, }; +use crate::http::memory_limit::BodyMemoryAccounting; use crate::prom_remote_write::decode::PromSeriesProcessor; use crate::prom_remote_write::v2::decode_remote_write_v2; use crate::prom_remote_write::validation::PromValidationMode; @@ -109,7 +110,7 @@ impl Default for RemoteWriteQuery { pub async fn remote_write( State(state): State, Query(params): Query, - Extension(query_ctx): Extension, + (Extension(query_ctx), extensions): (Extension, Extensions), content_type: Option>, pipeline_info: PipelineInfo, content_encoding: TypedHeader, @@ -126,7 +127,20 @@ pub async fn remote_write( { return Ok(response); } - remote_write_v2(state, params, query_ctx, pipeline_info, is_zstd, body).await + // Read after Bytes has collected every frame, including chunked-body reservations. + let body_reserved_bytes = extensions + .get::() + .map_or(0, BodyMemoryAccounting::reserved_bytes); + remote_write_v2( + state, + params, + query_ctx, + pipeline_info, + is_zstd, + body, + body_reserved_bytes, + ) + .await } RemoteWriteProto::Unsupported(content_type) => Ok(( StatusCode::UNSUPPORTED_MEDIA_TYPE, @@ -235,6 +249,7 @@ async fn remote_write_v2( pipeline_info: PipelineInfo, is_zstd: bool, body: Bytes, + body_reserved_bytes: u64, ) -> Result { let PromStoreState { prom_store_handler, @@ -256,20 +271,38 @@ async fn remote_write_v2( let (db, mut query_ctx, _timer) = prepare_remote_write_context(¶ms, query_ctx, REMOTE_WRITE_V2_VERSION); - let req = match decode_remote_write_v2(is_zstd, body, &memory_limiter).await { - Ok(req) => req, - Err(error) => return Ok(remote_write_v2_error_response(error, 0, 0, 0)), - }; + let req = + match decode_remote_write_v2(is_zstd, body, &memory_limiter, body_reserved_bytes).await { + Ok(req) => req, + Err(error) => return Ok(remote_write_v2_error_response(error, 0, 0, 0)), + }; // The v2 per-series metadata upgrades the written tables' semantic options // (metric type/unit, declared quality) at auto-create time. if let Some(index) = req.semantic_index.encode(&query_ctx.current_schema()) { query_ctx.set_extension(SEMANTIC_PER_TABLE_INDEX_KEY, index); } let query_ctx = Arc::new(query_ctx); - let sample_count = req.sample_count; - let histogram_count = req.histogram_count; - let sample_batches = into_prom_write_batches(req.samples, query_ctx.clone()); - let histogram_batches = into_prom_write_batches(req.histograms, query_ctx); + let (sample_count, histogram_count) = (req.sample_count, req.histogram_count); + let mut sample_batches = into_prom_write_batches(req.samples, query_ctx.clone()); + let mut histogram_batches = into_prom_write_batches(req.histograms, query_ctx); + for (batches, histogram) in [(&mut sample_batches, false), (&mut histogram_batches, true)] { + let counter = if histogram { + crate::metrics::PROM_STORE_REMOTE_WRITE_HISTOGRAMS + .with_label_values(&[&db, REMOTE_WRITE_V2_VERSION]) + } else { + crate::metrics::PROM_STORE_REMOTE_WRITE_SAMPLES + .with_label_values(&[&db, REMOTE_WRITE_V2_VERSION]) + }; + for (ctx, _) in batches { + let counter = counter.clone(); + let progress = + session::write_progress::WriteProgress::with_reservation(req.memory_guard.clone()) + .with_observer(move |rows| counter.inc_by(rows)); + let mut fork = ctx.fork(); + fork.set_write_progress(Arc::new(progress)); + *ctx = Arc::new(fork); + } + } let outcome = match write_prometheus_v2_rows_with_progress( prom_store_handler, pending_rows_batcher, @@ -281,8 +314,6 @@ async fn remote_write_v2( { Ok(outcome) => outcome, Err(error) => { - record_remote_write_samples(&db, REMOTE_WRITE_V2_VERSION, error.samples_written); - record_remote_write_histograms(&db, REMOTE_WRITE_V2_VERSION, error.histograms_written); return Ok(remote_write_v2_error_response( error.error, error.samples_written, @@ -291,11 +322,8 @@ async fn remote_write_v2( )); } }; - debug_assert_eq!(outcome.samples_written, sample_count); - debug_assert_eq!(outcome.histograms_written, histogram_count); - record_remote_write_samples(&db, REMOTE_WRITE_V2_VERSION, outcome.samples_written); - record_remote_write_histograms(&db, REMOTE_WRITE_V2_VERSION, outcome.histograms_written); - + debug_assert!(outcome.samples_written <= sample_count); + debug_assert!(outcome.histograms_written <= histogram_count); let mut headers = write_cost_header_map(outcome.write_cost); append_remote_write_v2_written_headers( &mut headers, @@ -304,7 +332,12 @@ async fn remote_write_v2( 0, ); - Ok((StatusCode::NO_CONTENT, headers).into_response()) + let status = if outcome.accepted { + StatusCode::ACCEPTED + } else { + StatusCode::NO_CONTENT + }; + Ok((status, headers).into_response()) } fn vm_proto_version_response(params: &RemoteWriteQuery) -> Option { @@ -356,6 +389,7 @@ struct PromWriteError { } struct PromWriteV2Outcome { + accepted: bool, write_cost: usize, samples_written: u64, histograms_written: u64, @@ -524,187 +558,120 @@ async fn write_prometheus_v2_rows_with_progress( sample_batches: Vec, histogram_batches: Vec, ) -> std::result::Result { - if histogram_batches.is_empty() { - return write_prometheus_rows_with_progress( - prom_store_handler, - pending_rows_batcher, - prom_store_with_metric_engine, - sample_batches, - ) - .await - .map(|outcome| PromWriteV2Outcome { - write_cost: outcome.write_cost, - samples_written: outcome.rows_written, - histograms_written: 0, - }) - .map_err(|error| PromWriteV2Error { - error: error.error, - samples_written: error.rows_written, - histograms_written: 0, - }); - } - let sample_batch_count = sample_batches.len(); let mut batches = sample_batches; batches.extend(histogram_batches); - - if prom_store_with_metric_engine && let Some(batcher) = pending_rows_batcher { - // Same ordering as the v1 path: preflight (which may redirect the - // contexts) before the bulk eligibility decision, and the fallback - // reuses the prepared batches without re-running hooks or admission. - preflight_prometheus_rows(&prom_store_handler, &mut batches) - .await - .map_err(|error| PromWriteV2Error { - error, - samples_written: 0, - histograms_written: 0, - })?; - // Destinations bound to another physical table must stay on the - // ordinary insert path, which routes per destination; time index - // units need no check — the bulk encode converts each request to - // its destination's unit. - if batcher.accepts_bulk_destinations(batches.iter()).await { - return write_batched_prometheus_v2_rows_with_progress( - prom_store_handler, - batcher.as_ref(), - prom_store_with_metric_engine, - sample_batch_count, - batches, - ) - .await; - } - return write_prepared_prometheus_v2_rows_with_progress( - prom_store_handler, - batches, - sample_batch_count, - prom_store_with_metric_engine, - ) - .await; - } - - let row_counts = batches - .iter() - .map(|(_, request)| prom_write_row_count(request)) - .collect::>(); - let batch_count = batches.len(); - let outputs = prom_store_handler - .write_all(batches, prom_store_with_metric_engine) + preflight_prometheus_rows(&prom_store_handler, &mut batches) .await .map_err(|error| PromWriteV2Error { error, samples_written: 0, histograms_written: 0, })?; - - let mut write_cost = 0; - let mut samples_written = 0; - let mut histograms_written = 0; - let mut output_count = 0; - for (index, (output, rows)) in outputs.into_iter().zip(row_counts).enumerate() { - let output = output.map_err(|error| PromWriteV2Error { - error, - samples_written, - histograms_written, - })?; - write_cost += output.meta.cost; - if index < sample_batch_count { - samples_written += rows; - } else { - histograms_written += rows; - } - output_count += 1; - } - if output_count != batch_count { - return Err(PromWriteV2Error { - error: incomplete_prom_write_error(), - samples_written, - histograms_written, - }); - } - - Ok(PromWriteV2Outcome { - write_cost, - samples_written, - histograms_written, - }) + let batcher = if prom_store_with_metric_engine + && let Some(batcher) = pending_rows_batcher.as_ref() + && batcher.accepts_bulk_destinations(batches.iter()).await + { + Some(batcher.as_ref() as &dyn PromWriteBatcher) + } else { + None + }; + write_batched_prometheus_v2_rows_with_progress( + prom_store_handler, + batcher, + prom_store_with_metric_engine, + sample_batch_count, + batches, + ) + .await } -/// Writes already-preflighted batches through the ordinary (prepared) write -/// path with v2 partial-progress accounting; the pre_write hooks and -/// admission have already run. -async fn write_prepared_prometheus_v2_rows_with_progress( +async fn write_batched_prometheus_v2_rows_with_progress( prom_store_handler: PromStoreProtocolHandlerRef, - batches: Vec, - sample_batch_count: usize, + batcher: Option<&dyn PromWriteBatcher>, prom_store_with_metric_engine: bool, + sample_batch_count: usize, + mut batches: Vec, ) -> std::result::Result { + let progress = batches + .iter_mut() + .map(|(ctx, _)| { + if ctx.write_progress().is_none() { + let mut fork = ctx.fork(); + fork.set_write_progress( + Arc::new(session::write_progress::WriteProgress::default()), + ); + *ctx = Arc::new(fork); + } + ctx.write_progress().cloned().unwrap_or_default() + }) + .collect::>(); + let counts = || { + let samples = progress[..sample_batch_count] + .iter() + .map(|p| p.written()) + .sum(); + let histograms = progress[sample_batch_count..] + .iter() + .map(|p| p.written()) + .sum(); + (samples, histograms) + }; let mut write_cost = 0; - let mut samples_written = 0; - let mut histograms_written = 0; - for (index, (ctx, request)) in batches.into_iter().enumerate() { - let rows = prom_write_row_count(&request); - let output = prom_store_handler - .write_prepared(request, ctx, prom_store_with_metric_engine) - .await - .map_err(|error| PromWriteV2Error { - error, - samples_written, - histograms_written, - })?; - write_cost += output.meta.cost; - if index < sample_batch_count { - samples_written += rows; + let mut samples_accepted = 0; + let mut histograms_accepted = 0; + for (index, (ctx, requests)) in batches.into_iter().enumerate() { + let result = if index < sample_batch_count + && let Some(batcher) = batcher + { + batcher + .submit(requests, ctx) + .await + .map(|rows| (rows, 0, false)) } else { - histograms_written += rows; + prom_store_handler + .write_prepared(requests, ctx, prom_store_with_metric_engine) + .await + .map(|output| { + let pending = !output.meta.write_completions.is_empty(); + let (rows, cost) = output.extract_rows_and_cost(); + (rows as u64, cost, pending) + }) + }; + match result { + Ok((rows, cost, pending)) => { + write_cost += cost; + if index < sample_batch_count { + samples_accepted += rows; + } else { + histograms_accepted += rows; + } + if pending { + progress[index].mark_accepted(); + } + if !progress[index].accepted() { + progress[index].confirm(rows); + } + } + Err(error) => { + for progress in &progress { + progress.wait().await; + } + let (samples_written, histograms_written) = counts(); + return Err(PromWriteV2Error { + error, + samples_written, + histograms_written, + }); + } } } Ok(PromWriteV2Outcome { write_cost, - samples_written, - histograms_written, - }) -} - -async fn write_batched_prometheus_v2_rows_with_progress( - prom_store_handler: PromStoreProtocolHandlerRef, - batcher: &B, - prom_store_with_metric_engine: bool, - sample_batch_count: usize, - batches: Vec, -) -> std::result::Result { - let mut samples_written = 0; - let mut histograms_written = 0; - let mut write_cost = 0; - let mut batches = batches.into_iter(); - for (ctx, requests) in batches.by_ref().take(sample_batch_count) { - let rows = batcher - .submit(requests, ctx) - .await - .map_err(|error| PromWriteV2Error { - error, - samples_written, - histograms_written, - })?; - samples_written += rows; - } - for (ctx, requests) in batches { - let rows = prom_write_row_count(&requests); - let output = prom_store_handler - .write_prepared(requests, ctx, prom_store_with_metric_engine) - .await - .map_err(|error| PromWriteV2Error { - error, - samples_written, - histograms_written, - })?; - write_cost += output.meta.cost; - histograms_written += rows; - } - - Ok(PromWriteV2Outcome { - write_cost, - samples_written, - histograms_written, + // PRW counts accepted rows on success; errors above report confirmed storage writes. + samples_written: samples_accepted, + histograms_written: histograms_accepted, + accepted: progress.iter().any(|p| p.accepted()), }) } @@ -732,15 +699,6 @@ fn record_remote_write_samples(db: &str, version: &str, rows: u64) { .inc_by(rows); } -fn record_remote_write_histograms(db: &str, version: &str, rows: u64) { - if rows == 0 { - return; - } - crate::metrics::PROM_STORE_REMOTE_WRITE_HISTOGRAMS - .with_label_values(&[db, version]) - .inc_by(rows); -} - fn remote_write_v2_error_response( error: error::Error, samples: u64, @@ -872,7 +830,8 @@ async fn decode_remote_read_request( ) -> Result { // Holds the memory permits for the decompressed bytes until the protobuf // decoding is finished. - let buf = snappy_decompress_limited(&body[..], MAX_DECOMPRESSED_REQUEST_SIZE, limiter).await?; + let buf = + snappy_decompress_limited(&body[..], MAX_DECOMPRESSED_REQUEST_SIZE, limiter, 0).await?; ReadRequest::decode(&buf[..]).context(error::DecodePromRemoteRequestSnafu) } @@ -880,15 +839,22 @@ async fn decode_remote_read_request( #[cfg(test)] mod tests { use std::sync::Mutex; + use std::time::Duration; use api::prom_store::remote::ReadRequest; use api::v1::{Row, RowInsertRequest, Rows}; use async_trait::async_trait; + use axum::body::Body; + use axum::http::Request; + use axum::{Router, middleware, routing}; + use common_memory_manager::OnExhaustedPolicy; use common_query::Output; use pipeline::GreptimePipelineParams; use session::context::{QueryContext, QueryContextRef}; + use tower::ServiceExt; use super::*; + use crate::http::memory_limit::memory_limit_middleware; use crate::prom_remote_write::validation::PromValidationMode; use crate::query_handler::PromStoreProtocolHandler; @@ -986,9 +952,14 @@ mod tests { preflight_prometheus_rows(&handler, &mut batches) .await .unwrap(); - let Ok(outcome) = - write_batched_prometheus_v2_rows_with_progress(handler, &batcher, true, 1, batches) - .await + let Ok(outcome) = write_batched_prometheus_v2_rows_with_progress( + handler, + Some(&batcher), + true, + 1, + batches, + ) + .await else { panic!("mixed remote write should succeed") }; @@ -1021,6 +992,37 @@ mod tests { ) } + #[tokio::test] + async fn test_async_v2_counts_acceptance_without_confirming_storage() { + let handler: PromStoreProtocolHandlerRef = Arc::new(RecordingPromStoreHandler { + events: Default::default(), + }); + let mut batches = vec![ + test_prom_write_batch("queued"), + test_prom_write_batch("queued"), + ]; + let mut progress = Vec::new(); + for (ctx, _) in &mut batches { + let tracker = Arc::new(session::write_progress::WriteProgress::default()); + Arc::get_mut(ctx) + .unwrap() + .set_write_progress(tracker.clone()); + progress.push(tracker); + } + preflight_prometheus_rows(&handler, &mut batches) + .await + .unwrap(); + let Ok(outcome) = + write_batched_prometheus_v2_rows_with_progress(handler, None, true, 1, batches).await + else { + panic!("queued writes should be accepted") + }; + assert!(outcome.accepted); + assert_eq!(outcome.samples_written, 1); + assert_eq!(outcome.histograms_written, 1); + assert!(progress.iter().all(|tracker| tracker.written() == 0)); + } + fn record_write_event(events: &Mutex>, phase: &str, request: &RowInsertRequests) { events.lock().unwrap().push(format!( "{phase}:{}", @@ -1028,6 +1030,60 @@ mod tests { )); } + #[tokio::test] + async fn test_failed_request_waits_for_earlier_queued_writes() { + struct DelayedBatcher( + Mutex< + Option<( + Arc, + session::write_progress::PendingWrite, + )>, + >, + ); + #[async_trait] + impl PromWriteBatcher for DelayedBatcher { + async fn submit( + &self, + request: RowInsertRequests, + ctx: QueryContextRef, + ) -> Result { + let progress = ctx.write_progress().unwrap().clone(); + let pending = progress.start(); + progress.mark_accepted(); + *self.0.lock().unwrap() = Some((progress, pending)); + Ok(prom_write_row_count(&request)) + } + } + let handler: PromStoreProtocolHandlerRef = Arc::new(RecordingPromStoreHandler { + events: Default::default(), + }); + let batcher = DelayedBatcher(Mutex::new(None)); + let mut batches = vec![ + test_prom_write_batch("queued"), + test_prom_write_batch("partial"), + ]; + preflight_prometheus_rows(&handler, &mut batches) + .await + .unwrap(); + let write = write_batched_prometheus_v2_rows_with_progress( + handler, + Some(&batcher), + true, + 1, + batches, + ); + tokio::pin!(write); + assert!(futures::poll!(write.as_mut()).is_pending()); + let (progress, pending) = batcher.0.lock().unwrap().take().unwrap(); + progress.add_written(1); + drop(pending); + let Err(error) = write.await else { + panic!("the second batch must fail") + }; + assert_eq!(error.samples_written, 1); + assert_eq!(error.histograms_written, 1); + } + struct RecordingPromWriteBatcher { events: Arc>>, } @@ -1064,7 +1120,20 @@ mod tests { ) -> Result { assert_eq!(ctx.write_rows_to_admit("greptime", "public", 1), 0); record_write_event(&self.events, "direct", &request); - Ok(Output::new_with_affected_rows(0)) + match request.inserts[0].table_name.as_str() { + "partial" => { + ctx.write_progress().unwrap().add_written(1); + return Err(InternalSnafu { + err_msg: "one region failed", + } + .build()); + } + "queued" => ctx.write_progress().unwrap().mark_accepted(), + _ => {} + } + Ok(Output::new_with_affected_rows( + prom_write_row_count(&request) as usize, + )) } async fn write( @@ -1093,6 +1162,146 @@ mod tests { } } + #[tokio::test] + async fn test_remote_write_v2_reports_confirmed_partial_and_queued_writes() { + for (name, expected_status, expected_count) in [ + ("confirmed", StatusCode::NO_CONTENT, "2"), + ("partial", StatusCode::INTERNAL_SERVER_ERROR, "1"), + ("queued", StatusCode::ACCEPTED, "2"), + ] { + let request = crate::prom_remote_write::v2::test_util::request_with_labels_and_samples( + vec![("__name__", name)], + vec![api::greptime_proto::io::prometheus::write::v2::Sample::default(); 2], + ); + let body = + Bytes::from(crate::prom_store::snappy_compress(&request.encode_to_vec()).unwrap()); + let mut state = test_state(); + state.prom_store_handler = Arc::new(RecordingPromStoreHandler { + events: Default::default(), + }); + let response = remote_write_v2( + state, + RemoteWriteQuery::default(), + QueryContext::with("greptime", "public"), + pipeline_info(None), + false, + body, + 0, + ) + .await + .unwrap(); + assert_eq!(response.status(), expected_status); + assert_eq!( + response.headers()[REMOTE_WRITE_V2_SAMPLES_WRITTEN_HEADER], + expected_count + ); + assert_eq!( + response.headers()[REMOTE_WRITE_V2_HISTOGRAMS_WRITTEN_HEADER], + "0" + ); + assert_eq!( + response.headers()[REMOTE_WRITE_V2_EXEMPLARS_WRITTEN_HEADER], + "0" + ); + } + } + + fn memory_limited_prom_write_app(limiter: ServerMemoryLimiter) -> Router { + let mut state = test_state(); + state.memory_limiter = limiter.clone(); + Router::new() + .route("/write", routing::post(remote_write)) + .with_state(state) + .layer(Extension(QueryContext::with("greptime", "public"))) + .layer(middleware::from_fn_with_state( + limiter, + memory_limit_middleware, + )) + } + + #[tokio::test(start_paused = true)] + async fn test_remote_write_v2_accounts_for_body_when_rejecting_or_waiting() { + let wire = api::greptime_proto::io::prometheus::write::v2::Request { + symbols: vec![ + String::new(), + format!("{:?}", (0..2048).collect::>()), + ], + timeseries: Vec::new(), + } + .encode_to_vec(); + let decoded_bytes = wire.len().div_ceil(1024) * 1024; + + for encoding in [DEFAULT_ENCODING, VM_ENCODING] { + let body = Bytes::from(if encoding == DEFAULT_ENCODING { + crate::prom_store::snappy_compress(&wire).unwrap() + } else { + zstd::stream::encode_all(wire.as_slice(), 3).unwrap() + }); + for chunked in [false, true] { + // Small frames make their rounded permits exceed the rounded total body length. + let chunk_size = if chunked { 513 } else { body.len() }; + let frames = body + .chunks(chunk_size) + .map(Bytes::copy_from_slice) + .collect::>(); + let body_reserved: usize = + frames.iter().map(|f| f.len().div_ceil(1024) * 1024).sum(); + if chunked { + assert!(body_reserved > body.len().div_ceil(1024) * 1024); + } + // Reject during decompression or scratch allocation; otherwise wait for contention. + for decoder_quota in [decoded_bytes - 1024, decoded_bytes, decoded_bytes + 1024] { + let limit = (body_reserved + decoder_quota) as u64; + let limiter = ServerMemoryLimiter::new(limit, OnExhaustedPolicy::default()); + let fits = decoder_quota > decoded_bytes; + let occupied = fits.then(|| limiter.try_acquire(1024).unwrap()); + let app = memory_limited_prom_write_app(limiter.clone()); + let mut request = Request::builder() + .uri("/write") + .method("POST") + .header( + header::CONTENT_TYPE, + "application/x-protobuf;proto=io.prometheus.write.v2.Request", + ) + .header(header::CONTENT_ENCODING, encoding); + let request_body = if chunked { + Body::from_stream(futures::stream::iter( + frames.clone().into_iter().map(Ok::<_, std::io::Error>), + )) + } else { + request = request.header(header::CONTENT_LENGTH, body.len().to_string()); + Body::from(body.clone()) + }; + let started = tokio::time::Instant::now(); + let response = app.oneshot(request.body(request_body).unwrap()); + tokio::pin!(response); + if let Some(occupied) = occupied { + assert!(futures::poll!(response.as_mut()).is_pending()); + // A failed attempt releases decoder permits, retaining only the HTTP body. + assert_eq!(limiter.used_bytes(), body_reserved as u64 + 1024); + drop(occupied); + } + let response = response.await.unwrap(); + assert_eq!( + response.status(), + if fits { + StatusCode::NO_CONTENT + } else { + StatusCode::TOO_MANY_REQUESTS + }, + "encoding={encoding}, chunked={chunked}, decoder_quota={decoder_quota}" + ); + assert_eq!(started.elapsed(), Duration::ZERO); + assert_eq!( + response.headers()[REMOTE_WRITE_V2_SAMPLES_WRITTEN_HEADER], + "0" + ); + assert_eq!(limiter.used_bytes(), 0); + } + } + } + } + #[tokio::test] async fn test_remote_write_v2_ignores_pipeline() { let request = api::greptime_proto::io::prometheus::write::v2::Request { @@ -1109,6 +1318,7 @@ mod tests { pipeline_info(Some("pipeline")), false, body, + 0, ) .await .unwrap(); diff --git a/src/servers/src/prom_remote_write/mod.rs b/src/servers/src/prom_remote_write/mod.rs index 06729795b5b..a54b50ad55b 100644 --- a/src/servers/src/prom_remote_write/mod.rs +++ b/src/servers/src/prom_remote_write/mod.rs @@ -70,12 +70,25 @@ pub(crate) async fn decompress_remote_write_body( is_zstd: bool, body: &[u8], limiter: &ServerMemoryLimiter, + body_reserved_bytes: u64, ) -> crate::error::Result { let decompress = |is_zstd: bool| async move { if is_zstd { - zstd_decompress_limited(body, MAX_DECOMPRESSED_REQUEST_SIZE, limiter).await + zstd_decompress_limited( + body, + MAX_DECOMPRESSED_REQUEST_SIZE, + limiter, + body_reserved_bytes, + ) + .await } else { - snappy_decompress_limited(body, MAX_DECOMPRESSED_REQUEST_SIZE, limiter).await + snappy_decompress_limited( + body, + MAX_DECOMPRESSED_REQUEST_SIZE, + limiter, + body_reserved_bytes, + ) + .await } }; @@ -112,7 +125,7 @@ pub async fn decode_remote_write_request( // Holds the memory permits for the decompressed bytes; the caller must // keep them alive as long as the returned `TablesBuilder`, which retains // the decompressed buffer as its raw data. - let buf = decompress_remote_write_body(is_zstd, &body[..], limiter).await?; + let buf = decompress_remote_write_body(is_zstd, &body[..], limiter, 0).await?; // Decompression copied the payload out, so the compressed body is no longer needed. drop(body); let (data, guards) = buf.into_parts(); diff --git a/src/servers/src/prom_remote_write/v2.rs b/src/servers/src/prom_remote_write/v2.rs index 9178748bc77..3c36370fa2a 100644 --- a/src/servers/src/prom_remote_write/v2.rs +++ b/src/servers/src/prom_remote_write/v2.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::collections::hash_map::Entry; +use std::sync::Arc; use ahash::{HashMap, HashMapExt, HashSet, HashSetExt}; #[cfg(test)] @@ -29,7 +30,8 @@ use api::v1::ColumnSchema; use api::v1::value::ValueData; use api::v1::{ColumnDataType, RowInsertRequest, Rows, SemanticType, Value}; use bytes::{Buf, Bytes}; -use common_grpc::precision::Precision; +use common_error::ext::BoxedError; +use common_memory_manager::{MemoryGuard, OnExhaustedPolicy}; #[cfg(test)] use common_query::native_histogram::*; use common_query::native_histogram::{ @@ -41,7 +43,7 @@ use prost::encoding::{ DecodeContext, WireType, decode_key, decode_varint, message, skip_field, uint32, }; use prost::{DecodeError, Message}; -use snafu::{OptionExt, ResultExt, ensure}; +use snafu::{IntoError, OptionExt, ResultExt, ensure}; use table::requests::{ METADATA_QUALITY_DECLARED, SEMANTIC_METRIC_METADATA_QUALITY, SEMANTIC_METRIC_TYPE, SEMANTIC_METRIC_UNIT, @@ -49,7 +51,6 @@ use table::requests::{ use crate::error::{self, Result}; use crate::prom_remote_write::row_builder::PromCtx; -use crate::prom_remote_write::validation::validate_label_name; use crate::prom_remote_write::{REMOTE_WRITE_V2_VERSION, decompress_remote_write_body}; #[allow(deprecated)] use crate::prom_store::{ @@ -57,6 +58,7 @@ use crate::prom_store::{ PHYSICAL_TABLE_LABEL_ALT, SCHEMA_LABEL, }; use crate::request_memory_limiter::ServerMemoryLimiter; +use crate::request_memory_metrics::RequestMemoryMetrics; use crate::row_writer::{self, TableData}; use crate::semantic::{ METRIC_TYPE_COUNTER, METRIC_TYPE_GAUGE, METRIC_TYPE_GAUGE_HISTOGRAM, METRIC_TYPE_HISTOGRAM, @@ -126,12 +128,132 @@ pub(crate) struct RemoteWriteV2WriteRequests { /// Per-table semantic metadata from the series' inline `Metadata`, folded /// into table options at auto-create time. pub semantic_index: SemanticIndexes, + pub memory_guard: Arc>, +} + +struct ConversionBudget<'a> { + limiter: &'a ServerMemoryLimiter, + guard: MemoryGuard, + used_bytes: u64, + held_bytes: u64, +} + +impl<'a> ConversionBudget<'a> { + fn new(limiter: &'a ServerMemoryLimiter, held_bytes: u64) -> Result { + Ok(Self { + limiter, + used_bytes: 0, + held_bytes, + guard: limiter.try_acquire(0).context(error::InternalSnafu { + err_msg: "cannot create an empty memory reservation", + })?, + }) + } + + fn reserve(&mut self, bytes: Option) -> Result<()> { + let required = bytes + .and_then(|bytes| u64::try_from(bytes).ok()) + .and_then(|bytes| self.used_bytes.checked_add(bytes)); + if let Some(required) = required { + let granted = self.guard.granted_bytes(); + let additional = required.saturating_sub(granted); + if additional == 0 { + self.used_bytes = required; + return Ok(()); + } + // Amortize quota updates; cap the chunk size at 1/64 of a finite limit. + let limit = self.limiter.limit_bytes(); + let (chunk, available) = if limit == 0 { + (64 * 1024, u64::MAX - granted) + } else { + ( + (64 * 1024).min(limit / 64), + limit.saturating_sub(self.limiter.used_bytes()), + ) + }; + let reservation = additional.max(chunk).min(available).max(additional); + if self.guard.try_acquire_additional(reservation) + || (reservation > additional && self.guard.try_acquire_additional(additional)) + { + self.used_bytes = required; + return Ok(()); + } + } + // The caller drops this attempt before waiting for its entire required quota. + Err(error::MemoryLimitExceededSnafu.into_error( + common_memory_manager::Error::MemoryLimitExceeded { + requested_bytes: required + .and_then(|bytes| self.held_bytes.checked_add(bytes)) + .unwrap_or(u64::MAX), + limit_bytes: self.limiter.limit_bytes(), + }, + )) + } } pub(crate) async fn decode_remote_write_v2( is_zstd: bool, body: Bytes, limiter: &ServerMemoryLimiter, + body_reserved_bytes: u64, +) -> Result { + let mut requested_bytes = 0; + let decode = async { + loop { + match decode_remote_write_v2_once(is_zstd, &body, limiter, body_reserved_bytes).await { + Err(error::Error::MemoryLimitExceeded { + source: + common_memory_manager::Error::MemoryLimitExceeded { + requested_bytes: required, + limit_bytes, + }, + .. + }) if matches!(limiter.policy(), OnExhaustedPolicy::Wait { .. }) + && required <= limit_bytes => + { + limiter.check_request_fits(required, body_reserved_bytes)?; + requested_bytes = required; + // The attempt is gone; HTTP body permits remain held until the handler returns. + tokio::task::yield_now().await; + drop(limiter.acquire(required).await?); + } + result => return result, + } + } + }; + let result = match limiter.policy() { + OnExhaustedPolicy::Wait { timeout } => { + // Another request can take the released permits before our next attempt. + tokio::time::timeout(timeout, decode) + .await + .unwrap_or_else(|_| { + Err(error::MemoryLimitExceededSnafu.into_error( + common_memory_manager::Error::MemoryAcquireTimeout { + requested_bytes, + waited: timeout, + }, + )) + }) + } + OnExhaustedPolicy::Fail => decode.await, + }; + result.map_err(|error| match error { + error::Error::MemoryLimitExceeded { + source: source @ common_memory_manager::Error::MemoryAcquireTimeout { .. }, + .. + } => { + // Prometheus retries 5xx; 429 may discard a request delayed by other writes. + error::OtherSnafu.into_error(BoxedError::new(source)) + } + error => error, + }) +} + +async fn decode_remote_write_v2_once( + is_zstd: bool, + body: &[u8], + limiter: &ServerMemoryLimiter, + body_reserved_bytes: u64, ) -> Result { let decode_timer = crate::metrics::METRIC_HTTP_PROM_STORE_CODEC_ELAPSED .with_label_values(&["decode", REMOTE_WRITE_V2_VERSION]) @@ -139,19 +261,68 @@ pub(crate) async fn decode_remote_write_v2( // Holds the memory permits for the decompressed bytes until the protobuf // decoding and conversion below are finished. - let buf = decompress_remote_write_body(is_zstd, &body[..], limiter).await?; - // Decompression copied the payload out, so the compressed body is no longer needed. - drop(body); + let buf = decompress_remote_write_body(is_zstd, body, limiter, body_reserved_bytes).await?; + let mut scratch = ConversionBudget::new(limiter, buf.reserved_bytes())?; + scratch.reserve(Some( + conversion_scratch_bytes(&buf).context(error::DecodePromRemoteRequestSnafu)?, + ))?; + let budget = ConversionBudget::new( + limiter, + buf.reserved_bytes() + scratch.guard.granted_bytes(), + )?; let request = BorrowedRequest::decode(&buf).context(error::DecodePromRemoteRequestSnafu)?; drop(decode_timer); let _convert_timer = crate::metrics::METRIC_HTTP_PROM_STORE_CODEC_ELAPSED .with_label_values(&["convert", REMOTE_WRITE_V2_VERSION]) .start_timer(); - convert_remote_write_v2(request) + convert_remote_write_v2(request, budget) } -fn convert_remote_write_v2(request: BorrowedRequest<'_>) -> Result { +/// Bounds borrowed-vector capacity and protobuf scratch without charging symbol contents or samples. +fn conversion_scratch_bytes(mut buf: &[u8]) -> std::result::Result { + // A one-byte varint can become an i64; Vec growth can double its capacity. + const PROTOBUF_SCRATCH_EXPANSION: usize = 2 * std::mem::size_of::(); + // Cover the minimum capacities of the two borrowed vectors. + let mut bytes = Some(8 * std::mem::size_of::<&[u8]>()); + while buf.has_remaining() { + let (tag, wire_type) = decode_key(&mut buf)?; + match tag { + 4 | 5 => { + let value = take_length_delimited(wire_type, &mut buf)?; + bytes = bytes.and_then(|n| n.checked_add(2 * std::mem::size_of::<&[u8]>())); + if tag == 5 { + let mut series = value; + while series.has_remaining() { + let before = series.len(); + let (tag, wire_type) = decode_key(&mut series)?; + skip_field(wire_type, tag, &mut series, DecodeContext::default())?; + if matches!( + tag, + TIME_SERIES_LABELS_REFS_TAG + | TIME_SERIES_HISTOGRAMS_TAG + | TIME_SERIES_EXEMPLARS_TAG + ) { + bytes = bytes.and_then(|n| { + n.checked_add( + (before - series.len()) + .checked_mul(PROTOBUF_SCRATCH_EXPANSION)?, + ) + }); + } + } + } + } + _ => skip_field(wire_type, tag, &mut buf, DecodeContext::default())?, + } + } + bytes.ok_or_else(|| DecodeError::new("remote write v2 scratch size overflows usize")) +} + +fn convert_remote_write_v2( + request: BorrowedRequest<'_>, + mut budget: ConversionBudget<'_>, +) -> Result { ensure!( request.symbols.first().copied() == Some(""), error::InvalidPromRemoteRequestSnafu { @@ -178,9 +349,23 @@ fn convert_remote_write_v2(request: BorrowedRequest<'_>) -> Result) -> Result 0 { + &sample_tables + } else { + &histogram_tables + } + .get(&prom_ctx) + .and_then(|tables| tables.get(&table_name)); + let old_columns = existing.map_or(2, TableData::num_columns); + let old_rows = existing.map_or(0, TableData::num_rows); + let new_columns = tags + .iter() + .filter(|(name, _)| existing.is_none_or(|table| table.column_index(name).is_none())) + .count(); + let row_count = counts.samples.checked_add(counts.histograms); + let tag_bytes = tags + .iter() + .try_fold(0usize, |bytes, (_, value)| bytes.checked_add(value.len())); + budget.reserve((|| { + let rows = row_count?; + let width = old_columns.checked_add(new_columns)?; + // Include row capacity, the tag template, and NULL padding of older rows. + let cells = rows + .checked_add(1)? + .checked_mul(width)? + .checked_add(old_rows.checked_mul(new_columns)?)?; + let values = cells.checked_mul(2 * std::mem::size_of::())?; + let strings = tag_bytes?.checked_mul(rows.checked_add(1)?)?; + let row_slots = rows.checked_mul(2 * std::mem::size_of::())?; + let schema = new_columns + .checked_mul(512)? + .checked_add(if existing.is_none() { 8192 } else { 0 })?; + let histograms = if counts.histograms > 0 { + // A one-byte bucket delta expands into a Value plus decoding scratch. + series + .len() + .checked_mul(2 * std::mem::size_of::() + 32)? + .checked_add( + counts + .histograms + .checked_mul(18 * std::mem::size_of::())?, + )? + } else { + 0 + }; + values + .checked_add(strings)? + .checked_add(row_slots)? + .checked_add(schema)? + .checked_add(histograms) + })())?; let (writer, row_count) = if counts.samples > 0 { ( SeriesWriter::Samples(get_or_create_table_data( @@ -249,6 +484,7 @@ fn convert_remote_write_v2(request: BorrowedRequest<'_>) -> Result Result<()> { let mut sample_row_template = None; + let mut histogram_row_template = None; while buf.has_remaining() { let (tag, wire_type) = decode_key(&mut buf).context(error::DecodePromRemoteRequestSnafu)?; @@ -480,18 +717,54 @@ fn decode_series_leaves( } })?; if let Some(SeriesWriter::Histograms(table_data)) = &mut writer { - if rows_remaining == 0 { - write_native_histogram( + // Validate every payload, including subsequent histograms in a series. + let value = encode_native_histogram(&scratch.histogram).map_err(|error| { + error::InvalidPromRemoteRequestSnafu { + msg: format!("remote write v2 {error}"), + } + .build() + })?; + if histogram_row_template.is_none() { + let timestamp_index = table_data.ensure_column_by_name( + greptime_timestamp(), + ColumnDataType::TimestampMillisecond, + SemanticType::Timestamp, + )?; + let value_index = match table_data.column_index(greptime_native_histogram()) + { + Some(index) => index, + None => { + let schema = native_histogram_column_schema().map_err(|error| { + error::InvalidPromRemoteRequestSnafu { + msg: format!("remote write v2 {error}"), + } + .build() + })?; + table_data.ensure_column(schema)? + } + }; + let mut row = table_data.alloc_one_row(); + row_writer::write_tags( table_data, - &scratch.histogram, std::mem::take(&mut tags).into_iter(), + &mut row, )?; - } else { - write_native_histogram( - table_data, - &scratch.histogram, - tags.iter().cloned(), - )?; + histogram_row_template = Some((row, timestamp_index, value_index)); + } + if let Some((template, timestamp_index, value_index)) = + histogram_row_template.as_mut() + { + // Clone only tags, never the previous histogram's bucket lists. + let mut row = if rows_remaining == 0 { + std::mem::take(template) + } else { + template.clone() + }; + row[*timestamp_index].value_data = Some( + ValueData::TimestampMillisecondValue(scratch.histogram.timestamp), + ); + row[*value_index].value_data = Some(value); + table_data.add_row(row); } } } @@ -584,47 +857,6 @@ fn get_or_create_table_data( } } -fn write_native_histogram<'a>( - table_data: &mut TableData, - histogram: &Histogram, - tags: impl Iterator, -) -> Result<()> { - let value = encode_native_histogram(histogram).map_err(|error| { - error::InvalidPromRemoteRequestSnafu { - msg: format!("remote write v2 {error}"), - } - .build() - })?; - let column_schema = native_histogram_column_schema().map_err(|error| { - error::InvalidPromRemoteRequestSnafu { - msg: format!("remote write v2 {error}"), - } - .build() - })?; - - // Persist both int and float families into the logical table schema. Only one - // family is populated per row; the other is written as NULL so PromQL can - // infer the original histogram flavor without a separate type column. - let mut row = table_data.alloc_one_row(); - row_writer::write_ts_to_millis( - table_data, - greptime_timestamp(), - Some(histogram.timestamp), - Precision::Millisecond, - &mut row, - )?; - row_writer::write_by_schema( - table_data, - std::iter::once((column_schema, Some(value))), - &mut row, - )?; - - row_writer::write_tags(table_data, tags, &mut row)?; - table_data.add_row(row); - - Ok(()) -} - fn ensure_no_internal_histogram_labels(tags: &PromTags<'_>) -> Result<()> { // The histogram field column is generated from the protobuf payload. for (name, _) in tags { @@ -695,7 +927,7 @@ fn resolve_series_labels<'a>( fn validate_label(name: &str) -> Result<()> { ensure!( - validate_label_name(name.as_bytes()), + !name.is_empty(), error::InvalidPromRemoteRequestSnafu { msg: format!("remote write v2 invalid label name `{name}`"), } @@ -818,7 +1050,7 @@ pub mod test_util { .build() .unwrap() .block_on(async move { - decompress_remote_write_body(is_zstd, &body[..], &ServerMemoryLimiter::default()) + decompress_remote_write_body(is_zstd, &body[..], &ServerMemoryLimiter::default(), 0) .await .map(|buf| buf.to_vec()) }) @@ -843,6 +1075,7 @@ pub mod test_util { is_zstd, body, &ServerMemoryLimiter::default(), + 0, ))?; Ok(( requests.samples.all_req().collect(), @@ -855,9 +1088,15 @@ pub mod test_util { pub fn decode_uncompressed_write_requests( body: &[u8], ) -> Result<(Vec, Vec, u64, u64)> { + let limiter = ServerMemoryLimiter::default(); + let mut scratch = super::ConversionBudget::new(&limiter, 0)?; + scratch.reserve(Some( + super::conversion_scratch_bytes(body).context(error::DecodePromRemoteRequestSnafu)?, + ))?; + let budget = super::ConversionBudget::new(&limiter, scratch.guard.granted_bytes())?; let request = super::BorrowedRequest::decode(body).context(error::DecodePromRemoteRequestSnafu)?; - let requests = super::convert_remote_write_v2(request)?; + let requests = super::convert_remote_write_v2(request, budget)?; Ok(( requests.samples.all_req().collect(), requests.histograms.all_req().collect(), @@ -1362,6 +1601,206 @@ mod tests { ); } + #[tokio::test] + async fn test_expanded_rows_are_charged_before_allocation_and_released() { + use common_memory_manager::OnExhaustedPolicy; + let label = "x".repeat(64 * 1024); + let mut request = request_with_sample(vec![(METRIC_NAME_LABEL, "metric"), ("tag", &label)]); + request.timeseries[0].samples = vec![Sample::default(); 10_000]; + let limiter = ServerMemoryLimiter::new(8 * 1024 * 1024, OnExhaustedPolicy::Fail); + let body = + Bytes::from(crate::prom_store::snappy_compress(&request.encode_to_vec()).unwrap()); + let error = match decode_remote_write_v2(false, body, &limiter, 0).await { + Ok(_) => panic!("interned labels must not expand past the quota"), + Err(error) => error, + }; + assert!(matches!(error, error::Error::MemoryLimitExceeded { .. })); + assert_eq!(limiter.used_bytes(), 0); + + request.timeseries[0].samples.truncate(1); + let body = + Bytes::from(crate::prom_store::snappy_compress(&request.encode_to_vec()).unwrap()); + let decoded = decode_remote_write_v2(false, body, &limiter, 0) + .await + .unwrap(); + assert!(limiter.used_bytes() >= label.len() as u64); + let progress = + session::write_progress::WriteProgress::with_reservation(decoded.memory_guard.clone()); + drop(decoded); + assert!(limiter.used_bytes() >= label.len() as u64); + drop(progress); + assert_eq!(limiter.used_bytes(), 0); + let mut budget = ConversionBudget::new(&limiter, 0).unwrap(); + assert!(budget.reserve(usize::MAX.checked_mul(2)).is_err()); + } + + #[tokio::test] + async fn test_large_symbols_do_not_reserve_body_multiples_or_retain_scratch() { + let limiter = ServerMemoryLimiter::new(64 * 1024 * 1024, OnExhaustedPolicy::Fail); + let mut request = request_with_sample(vec![(METRIC_NAME_LABEL, "metric")]); + let body = + Bytes::from(crate::prom_store::snappy_compress(&request.encode_to_vec()).unwrap()); + let decoded = decode_remote_write_v2(false, body, &limiter, 0) + .await + .unwrap(); + let row_bytes = limiter.used_bytes(); + drop(decoded); + + request + .symbols + .extend((0..4096).map(|i| format!("unused_{i}"))); + let help_ref = request.symbols.len() as u32; + request.symbols.push("x".repeat(4 * 1024 * 1024)); + request.timeseries[0].metadata = Some(Metadata { + help_ref, + ..Default::default() + }); + let body = + Bytes::from(crate::prom_store::snappy_compress(&request.encode_to_vec()).unwrap()); + let decoded = decode_remote_write_v2(false, body, &limiter, 0) + .await + .unwrap(); + assert_eq!(decoded.sample_count, 1); + assert_eq!(limiter.used_bytes(), row_bytes); + drop(decoded); + assert_eq!(limiter.used_bytes(), 0); + } + + #[tokio::test(start_paused = true)] + async fn test_conversion_wait_timeout_is_retryable_but_oversized_requests_do_not_wait() { + use common_error::ext::{ErrorExt, RetryHint}; + + let limiter = ServerMemoryLimiter::new(64 * 1024, OnExhaustedPolicy::default()); + let occupied = limiter.try_acquire(60 * 1024).unwrap(); + let mut request = request_with_sample(vec![(METRIC_NAME_LABEL, "metric")]); + let body = + Bytes::from(crate::prom_store::snappy_compress(&request.encode_to_vec()).unwrap()); + let error = decode_remote_write_v2(false, body, &limiter, 0) + .await + .err() + .unwrap(); + assert_eq!(error.retry_hint(), RetryHint::Retryable); + assert!(crate::error::status_code_to_http_status(&error.status_code()).is_server_error()); + assert_eq!(limiter.used_bytes(), occupied.granted_bytes()); + drop(occupied); + + request.timeseries[0].samples = vec![Sample::default(); 10_000]; + let body = + Bytes::from(crate::prom_store::snappy_compress(&request.encode_to_vec()).unwrap()); + let started = tokio::time::Instant::now(); + let error = decode_remote_write_v2(false, body, &limiter, 0) + .await + .err() + .unwrap(); + assert!(matches!(error, error::Error::MemoryLimitExceeded { .. })); + assert_eq!(started.elapsed(), std::time::Duration::ZERO); + assert_eq!(limiter.used_bytes(), 0); + } + + #[test] + fn test_chunked_reservations_respect_small_and_shared_quotas() { + let limiter = + ServerMemoryLimiter::new(8192, common_memory_manager::OnExhaustedPolicy::Fail); + let mut first = ConversionBudget::new(&limiter, 0).unwrap(); + let mut second = ConversionBudget::new(&limiter, 0).unwrap(); + first.reserve(Some(3000)).unwrap(); + second.reserve(Some(3000)).unwrap(); + assert!(limiter.used_bytes() <= 8192); + assert!(first.reserve(Some(8192)).is_err()); + assert!(first.guard.granted_bytes() >= first.used_bytes); + drop(second); + first.reserve(Some(3000)).unwrap(); + assert!(first.guard.granted_bytes() >= first.used_bytes); + first.used_bytes = u64::MAX; + assert!(first.reserve(Some(1)).is_err()); + drop(first); + assert_eq!(limiter.used_bytes(), 0); + } + + #[test] + fn test_schema_widening_reserves_null_padding() { + let mut request = request_with_sample(vec![(METRIC_NAME_LABEL, "metric")]); + let series = request.timeseries.pop().unwrap(); + for index in 0..200 { + let mut series = series.clone(); + series.labels_refs.extend([request.symbols.len() as u32, 0]); + request.symbols.push(format!("label_{index}")); + request.timeseries.push(series); + } + let body = request.encode_to_vec(); + let limiter = ServerMemoryLimiter::default(); + let budget = ConversionBudget::new(&limiter, 0).unwrap(); + let decoded = + convert_remote_write_v2(BorrowedRequest::decode(&body).unwrap(), budget).unwrap(); + let minimum = 200 * 202 * std::mem::size_of::(); + assert!(limiter.used_bytes() >= minimum as u64); + let rows = decoded.samples.all_req().next().unwrap().rows.unwrap(); + assert_eq!(rows.schema.len(), 202); + assert!(rows.rows.iter().all(|row| row.values.len() == 202)); + drop(decoded.memory_guard); + assert_eq!(limiter.used_bytes(), 0); + + let limiter = ServerMemoryLimiter::new( + minimum as u64 / 2, + common_memory_manager::OnExhaustedPolicy::Fail, + ); + let budget = ConversionBudget::new(&limiter, 0).unwrap(); + assert!(convert_remote_write_v2(BorrowedRequest::decode(&body).unwrap(), budget).is_err()); + assert_eq!(limiter.used_bytes(), 0); + } + + #[test] + fn test_utf8_labels_and_histogram_row_templates() { + let labels = vec![ + (METRIC_NAME_LABEL, "metric"), + ("service.name", "api"), + ("区域", "西"), + ("has-dash", "value"), + ]; + let sample = request_with_sample(labels.clone()); + let (requests, _, _, _) = test_util::write_requests(sample).unwrap(); + let sample_rows = requests[0].rows.as_ref().unwrap(); + assert_eq!(sample_rows.schema[3].column_name, "区域"); + assert_eq!( + sample_rows.rows[0].values[3].value_data, + Some(ValueData::StringValue("西".into())) + ); + + let mut request = request_with_sample(labels); + request.timeseries[0].samples.clear(); + request.timeseries[0].histograms = vec![ + Histogram { + timestamp: i64::MIN, + sum: 1.0, + ..Default::default() + }, + Histogram { + timestamp: i64::MAX, + sum: 2.0, + ..Default::default() + }, + ]; + let (_, requests, _, count) = test_util::write_requests(request.clone()).unwrap(); + assert_eq!(count, 2); + let rows = requests[0].rows.as_ref().unwrap(); + for (index, timestamp) in [i64::MIN, i64::MAX].into_iter().enumerate() { + assert_eq!( + rows.rows[index].values[0].value_data, + Some(ValueData::TimestampMillisecondValue(timestamp)) + ); + assert_eq!( + rows.rows[index].values[1].value_data, + Some(encode_native_histogram(&request.timeseries[0].histograms[index]).unwrap()) + ); + assert_eq!( + rows.rows[index].values[3].value_data, + Some(ValueData::StringValue("西".into())) + ); + } + request.timeseries[0].histograms[1].schema = 99; + assert!(test_util::write_requests(request).is_err()); + } + #[test] fn test_into_context_req_rejects_invalid_requests() { let mut cases = Vec::new(); @@ -1412,24 +1851,6 @@ mod tests { "invalid label name", )); - cases.push(( - "invalid label name", - request_with_sample(vec![(METRIC_NAME_LABEL, "metric"), ("has-dash", "api")]), - "invalid label name", - )); - - cases.push(( - "dotted label name", - request_with_sample(vec![(METRIC_NAME_LABEL, "metric"), ("service.name", "api")]), - "invalid label name", - )); - - cases.push(( - "non-ascii label name", - request_with_sample(vec![(METRIC_NAME_LABEL, "metric"), ("区域", "api")]), - "invalid label name", - )); - cases.push(( "empty metric name", request_with_sample(vec![(METRIC_NAME_LABEL, "")]), @@ -2055,6 +2476,7 @@ mod tests { is_zstd, body, &ServerMemoryLimiter::default(), + 0, )) } diff --git a/src/servers/src/prom_store.rs b/src/servers/src/prom_store.rs index d45cabceaf7..cc83b35e6fc 100644 --- a/src/servers/src/prom_store.rs +++ b/src/servers/src/prom_store.rs @@ -737,6 +737,10 @@ impl std::fmt::Debug for ChargedBuffer { } impl ChargedBuffer { + pub(crate) fn reserved_bytes(&self) -> u64 { + self.guards.iter().map(MemoryGuard::granted_bytes).sum() + } + /// Creates a buffer with the permits charged for it. pub(crate) fn new(data: Vec, guards: Vec>) -> Self { Self { data, guards } @@ -759,6 +763,7 @@ pub(crate) async fn snappy_decompress_limited( buf: &[u8], max_decompressed: usize, limiter: &ServerMemoryLimiter, + held_bytes: u64, ) -> Result { let decoded_len = snap::raw::decompress_len(buf).context(error::DecompressSnappyPromRemoteRequestSnafu)?; @@ -771,6 +776,7 @@ pub(crate) async fn snappy_decompress_limited( ); // Reserve the declared size before the allocation happens so concurrent // decompressions cannot exceed the aggregate quota. + limiter.check_request_fits(decoded_len as u64, held_bytes)?; let guard = limiter.acquire(decoded_len as u64).await?; let mut decoder = Decoder::new(); let data = decoder @@ -789,6 +795,7 @@ pub(crate) async fn zstd_decompress_limited( buf: &[u8], max_decompressed: usize, limiter: &ServerMemoryLimiter, + mut held_bytes: u64, ) -> Result { use std::io::Read; @@ -816,7 +823,10 @@ pub(crate) async fn zstd_decompress_limited( ); // Charge each chunk before it is appended so the aggregate quota // bounds the memory while it is being materialized. - guards.push(limiter.acquire(n as u64).await?); + limiter.check_request_fits(n as u64, held_bytes)?; + let guard = limiter.acquire(n as u64).await?; + held_bytes += guard.granted_bytes(); + guards.push(guard); data.extend_from_slice(&chunk[..n]); } Ok(ChargedBuffer { data, guards }) @@ -2283,7 +2293,7 @@ mod tests { let bomb = [0xffu8, 0xff, 0xff, 0xff, 0x0f]; let limiter = ServerMemoryLimiter::default(); - let err = snappy_decompress_limited(&bomb, MAX_DECOMPRESSED_REQUEST_SIZE, &limiter) + let err = snappy_decompress_limited(&bomb, MAX_DECOMPRESSED_REQUEST_SIZE, &limiter, 0) .await .unwrap_err(); assert!(matches!( @@ -2301,7 +2311,7 @@ mod tests { let limiter = ServerMemoryLimiter::default(); let decoded = - snappy_decompress_limited(&compressed, MAX_DECOMPRESSED_REQUEST_SIZE, &limiter) + snappy_decompress_limited(&compressed, MAX_DECOMPRESSED_REQUEST_SIZE, &limiter, 0) .await .unwrap(); assert_eq!(&decoded[..], &raw[..]); @@ -2315,7 +2325,7 @@ mod tests { let limiter = ServerMemoryLimiter::new(128 * 1024, OnExhaustedPolicy::Fail); let decoded = - snappy_decompress_limited(&compressed, MAX_DECOMPRESSED_REQUEST_SIZE, &limiter) + snappy_decompress_limited(&compressed, MAX_DECOMPRESSED_REQUEST_SIZE, &limiter, 0) .await .unwrap(); // Rounding to the 1 KiB permit granularity. @@ -2333,9 +2343,10 @@ mod tests { // Quota smaller than the decoded size. let limiter = ServerMemoryLimiter::new(1024, OnExhaustedPolicy::Fail); - let err = snappy_decompress_limited(&compressed, MAX_DECOMPRESSED_REQUEST_SIZE, &limiter) - .await - .unwrap_err(); + let err = + snappy_decompress_limited(&compressed, MAX_DECOMPRESSED_REQUEST_SIZE, &limiter, 0) + .await + .unwrap_err(); assert!(matches!( err, crate::error::Error::MemoryLimitExceeded { .. } @@ -2351,7 +2362,7 @@ mod tests { // Exactly at the cap: accepted. let raw = vec![0u8; small_cap]; let compressed = zstd::stream::encode_all(&raw[..], 3).unwrap(); - let decoded = zstd_decompress_limited(&compressed, small_cap, &limiter) + let decoded = zstd_decompress_limited(&compressed, small_cap, &limiter, 0) .await .unwrap(); assert_eq!(decoded.len(), small_cap); @@ -2359,7 +2370,7 @@ mod tests { // One byte past the cap: rejected. let raw = vec![0u8; small_cap + 1]; let compressed = zstd::stream::encode_all(&raw[..], 3).unwrap(); - let err = zstd_decompress_limited(&compressed, small_cap, &limiter) + let err = zstd_decompress_limited(&compressed, small_cap, &limiter, 0) .await .unwrap_err(); assert!(matches!( @@ -2368,7 +2379,7 @@ mod tests { )); } - #[tokio::test] + #[tokio::test(start_paused = true)] async fn test_zstd_charges_incrementally() { // 64 KiB wire that expands to 2 MiB must be charged chunk by chunk // against a quota that only fits the expanded size. @@ -2377,13 +2388,27 @@ mod tests { assert!(compressed.len() < 16 * 1024); let limiter = ServerMemoryLimiter::new(4 * 1024 * 1024, OnExhaustedPolicy::Fail); - let decoded = zstd_decompress_limited(&compressed, MAX_DECOMPRESSED_REQUEST_SIZE, &limiter) - .await - .unwrap(); + let decoded = + zstd_decompress_limited(&compressed, MAX_DECOMPRESSED_REQUEST_SIZE, &limiter, 0) + .await + .unwrap(); assert_eq!(decoded.len(), raw.len()); assert!(limiter.used_bytes() >= raw.len() as u64); drop(decoded); assert_eq!(0, limiter.used_bytes()); + // The next chunk cannot fit alongside this request's earlier chunks. + let limiter = ServerMemoryLimiter::new( + (2 * ZSTD_DECOMPRESS_CHUNK) as u64, + OnExhaustedPolicy::default(), + ); + let started = tokio::time::Instant::now(); + let error = + zstd_decompress_limited(&compressed, MAX_DECOMPRESSED_REQUEST_SIZE, &limiter, 0) + .await + .unwrap_err(); + assert!(matches!(error, error::Error::MemoryLimitExceeded { .. })); + assert_eq!(started.elapsed(), std::time::Duration::ZERO); + assert_eq!(limiter.used_bytes(), 0); } #[test] diff --git a/src/servers/src/request_memory_limiter.rs b/src/servers/src/request_memory_limiter.rs index ac072f51b2e..e2be0a08151 100644 --- a/src/servers/src/request_memory_limiter.rs +++ b/src/servers/src/request_memory_limiter.rs @@ -17,7 +17,7 @@ use std::sync::Arc; use common_memory_manager::{MemoryGuard, MemoryManager, OnExhaustedPolicy, PermitGranularity}; -use snafu::ResultExt; +use snafu::{IntoError, ResultExt}; use crate::error::{MemoryLimitExceededSnafu, Result}; use crate::request_memory_metrics::RequestMemoryMetrics; @@ -69,6 +69,25 @@ impl ServerMemoryLimiter { self.manager.try_acquire(bytes) } + /// Rejects requests whose additional bytes cannot fit alongside their own held permits. + pub(crate) fn check_request_fits(&self, bytes: u64, held_bytes: u64) -> Result<()> { + let requested_bytes = held_bytes.checked_add(bytes); + let limit_bytes = self.limit_bytes(); + match requested_bytes { + Some(bytes) if limit_bytes == 0 || bytes <= limit_bytes => Ok(()), + _ => Err(MemoryLimitExceededSnafu.into_error( + common_memory_manager::Error::MemoryLimitExceeded { + requested_bytes: requested_bytes.unwrap_or(u64::MAX), + limit_bytes, + }, + )), + } + } + + pub(crate) fn policy(&self) -> OnExhaustedPolicy { + self.policy + } + /// Returns total memory limit in bytes (0 if unlimited). pub fn limit_bytes(&self) -> u64 { self.manager.limit_bytes() diff --git a/src/servers/src/row_writer.rs b/src/servers/src/row_writer.rs index 780a19df3f8..74ef6cedc0d 100644 --- a/src/servers/src/row_writer.rs +++ b/src/servers/src/row_writer.rs @@ -65,6 +65,10 @@ impl TableData { self.schema.len() } + pub(crate) fn column_index(&self, name: &str) -> Option { + self.column_indexes.get(name).copied() + } + #[inline] pub fn num_rows(&self) -> usize { self.rows.len() diff --git a/src/servers/tests/http/prom_store_test.rs b/src/servers/tests/http/prom_store_test.rs index e4ef6219b30..47d7004940c 100644 --- a/src/servers/tests/http/prom_store_test.rs +++ b/src/servers/tests/http/prom_store_test.rs @@ -97,6 +97,12 @@ impl PromStoreProtocolHandler for DummyInstance { .fail(); } + let rows = request + .inserts + .iter() + .filter_map(|insert| insert.rows.as_ref()) + .map(|rows| rows.rows.len()) + .sum(); let _ = self .write_tx .send(RemoteWriteCapture { @@ -108,7 +114,7 @@ impl PromStoreProtocolHandler for DummyInstance { }) .await; - Ok(Output::new_with_affected_rows(0)) + Ok(Output::new_with_affected_rows(rows)) } async fn write( diff --git a/src/session/Cargo.toml b/src/session/Cargo.toml index 93c46ff7ff7..1cd6a88567c 100644 --- a/src/session/Cargo.toml +++ b/src/session/Cargo.toml @@ -25,6 +25,7 @@ derive_builder.workspace = true derive_more.workspace = true snafu.workspace = true sql.workspace = true +tokio.workspace = true uuid.workspace = true [dev-dependencies] diff --git a/src/session/src/context.rs b/src/session/src/context.rs index 02ee621e0e4..03643331ddf 100644 --- a/src/session/src/context.rs +++ b/src/session/src/context.rs @@ -43,6 +43,7 @@ use crate::query_id::QueryId; use crate::session_config::{ InvalidConfigValueSnafu, PGByteaOutputValue, PGDateOrder, PGDateTimeStyle, PGIntervalStyle, }; +use crate::write_progress::WriteProgress; use crate::{MutableInner, ReadPreference}; pub type QueryContextRef = Arc; @@ -95,6 +96,9 @@ pub struct QueryContext { /// Process-local admission for one database in this request. Never sent over the wire. #[builder(setter(skip))] admitted_write: Option<(String, String)>, + /// Process-local storage acknowledgements; never accepted from a wire context. + #[builder(setter(skip))] + write_progress: Option>, /// Process id for managing on-going queries #[builder(default)] process_id: u32, @@ -289,6 +293,14 @@ impl From<&QueryContext> for api::v1::QueryContext { } impl QueryContext { + pub fn write_progress(&self) -> Option<&Arc> { + self.write_progress.as_ref() + } + + pub fn set_write_progress(&mut self, progress: Arc) { + self.write_progress = Some(progress); + } + /// Forks this context with an independent snapshot of mutable session data. /// /// Unlike [`Clone`], changes to the schema, user, timezone, and other fields @@ -763,6 +775,7 @@ impl QueryContextBuilder { channel, batching_enabled: self.batching_enabled.unwrap_or_default(), admitted_write: None, + write_progress: None, logical_batching_enabled: self.logical_batching_enabled.unwrap_or_default(), process_id: self.process_id.unwrap_or_default(), conn_info: self.conn_info.unwrap_or_default(), diff --git a/src/session/src/lib.rs b/src/session/src/lib.rs index af0a4918e4b..fe4e7bdcc01 100644 --- a/src/session/src/lib.rs +++ b/src/session/src/lib.rs @@ -19,6 +19,7 @@ pub mod query_id; pub mod query_options; pub mod session_config; pub mod table_name; +pub mod write_progress; use std::collections::{HashMap, VecDeque}; use std::net::SocketAddr; diff --git a/src/session/src/write_progress.rs b/src/session/src/write_progress.rs new file mode 100644 index 00000000000..e6f7fc511d4 --- /dev/null +++ b/src/session/src/write_progress.rs @@ -0,0 +1,146 @@ +// 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. + +//! Process-local acknowledgements for a finite write request. + +use std::fmt; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}; +use tokio::sync::Notify; + +/// Counts only confirmed storage writes, including partial successes. +#[derive(Default)] +pub struct WriteProgress { + written: AtomicU64, + accepted: AtomicBool, + // A queued write must keep its decoder's memory reservation alive. + _reservation: Option>, + on_write: Option>, + pending: AtomicUsize, + completed: Notify, +} + +/// Keeps a queued submission pending until the worker reports its result. +pub struct PendingWrite(Arc); + +impl Drop for PendingWrite { + fn drop(&mut self) { + if self.0.pending.fetch_sub(1, Ordering::AcqRel) == 1 { + self.0.completed.notify_waiters(); + } + } +} + +impl fmt::Debug for WriteProgress { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("WriteProgress") + .field("written", &self.written()) + .field("accepted", &self.accepted()) + .finish_non_exhaustive() + } +} + +impl WriteProgress { + pub fn start(self: &Arc) -> PendingWrite { + self.pending.fetch_add(1, Ordering::Relaxed); + PendingWrite(self.clone()) + } + + /// Waits for previously queued submissions before returning a failed request's counts. + pub async fn wait(&self) { + loop { + let notified = self.completed.notified(); + tokio::pin!(notified); + notified.as_mut().enable(); + if self.pending.load(Ordering::Acquire) == 0 { + return; + } + notified.await; + } + } + + pub fn with_reservation(reservation: Arc) -> Self { + Self { + _reservation: Some(reservation), + ..Self::default() + } + } + + pub fn add_written(&self, rows: usize) { + self.written.fetch_add(rows as u64, Ordering::Relaxed); + self.observe(rows as u64); + } + + pub fn with_observer(mut self, observer: impl Fn(u64) + Send + Sync + 'static) -> Self { + self.on_write = Some(Box::new(observer)); + self + } + + /// A synchronous whole-request acknowledgement also covers handlers without region tracking. + pub fn confirm(&self, rows: u64) { + let previous = self.written.fetch_max(rows, Ordering::Relaxed); + self.observe(rows.saturating_sub(previous)); + } + + fn observe(&self, rows: u64) { + if rows > 0 + && let Some(observer) = &self.on_write + { + observer(rows); + } + } + + pub fn written(&self) -> u64 { + self.written.load(Ordering::Relaxed) + } + + /// Records an acknowledgement returned before storage completion. + pub fn mark_accepted(&self) { + self.accepted.store(true, Ordering::Relaxed); + } + + pub fn accepted(&self) -> bool { + self.accepted.load(Ordering::Relaxed) + } +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::AtomicUsize; + + use super::*; + + #[test] + fn confirmations_do_not_double_count_and_contexts_keep_local_ownership() { + let observed = Arc::new(AtomicUsize::new(0)); + let counter = observed.clone(); + let progress = Arc::new(WriteProgress::default().with_observer(move |rows| { + counter.fetch_add(rows as usize, Ordering::Relaxed); + })); + let mut context = crate::context::QueryContext::with("greptime", "public"); + context.set_write_progress(progress.clone()); + let fork = context.with_write_admission(); + assert!(Arc::ptr_eq(fork.write_progress().unwrap(), &progress)); + let wire = api::v1::QueryContext::from(context); + let received = crate::context::QueryContext::try_from(wire).unwrap(); + assert!(received.write_progress().is_none()); + progress.add_written(2); + progress.confirm(3); + progress.confirm(3); + progress.mark_accepted(); + assert_eq!(progress.written(), 3); + assert_eq!(observed.load(Ordering::Relaxed), 3); + assert!(progress.accepted()); + } +} diff --git a/tests-integration/src/tests/instance_test.rs b/tests-integration/src/tests/instance_test.rs index 4f6cd5a3fab..c1f886c6ed7 100644 --- a/tests-integration/src/tests/instance_test.rs +++ b/tests-integration/src/tests/instance_test.rs @@ -4279,13 +4279,23 @@ async fn test_histogram_ingestion_storage_lifecycle(#[case] metric_engine: bool) let mut router = make_router(context.frontend()); for round in 0..2 { let timestamp = 1_000 + round * 10_000; - let mut labels = vec![("__name__", "lifecycle_histogram"), ("job", "api")]; + let mut labels = vec![ + ("__name__", "lifecycle_histogram"), + ("job", "api"), + ("service.name", "api"), + ("区域", "西"), + ]; if round == 1 { // Exercise schema evolution after the first SST already exists. labels.push(("zone", "east")); } let sample = remote_write_v2::request_with_labels_and_samples( - vec![("__name__", "lifecycle_sample"), ("job", "api")], + vec![ + ("__name__", "lifecycle_sample"), + ("job", "api"), + ("service.name", "api"), + ("区域", "西"), + ], vec![Sample { value: 42.0, timestamp, @@ -4310,7 +4320,7 @@ async fn test_histogram_ingestion_storage_lifecycle(#[case] metric_engine: bool) }], ); // Create the scalar table first, then add histogram storage to its physical table. - for request in [sample, histograms] { + for (request, samples, histograms) in [(sample, "1", "0"), (histograms, "0", "1")] { let response = router .clone() .oneshot( @@ -4327,6 +4337,14 @@ async fn test_histogram_ingestion_storage_lifecycle(#[case] metric_engine: bool) ) .await .unwrap(); + assert_eq!( + response.headers()["X-Prometheus-Remote-Write-Samples-Written"], + samples + ); + assert_eq!( + response.headers()["X-Prometheus-Remote-Write-Histograms-Written"], + histograms + ); assert_eq!( response.status(), 204, @@ -4336,6 +4354,11 @@ async fn test_histogram_ingestion_storage_lifecycle(#[case] metric_engine: bool) } let expected = snapshot(context.frontend()).await; + for table in &expected { + assert!( + table.contains("service.name") && table.contains("区域") && table.contains("西") + ); + } assert!(expected[1].contains("9007199254740993"), "{}", expected[1]); // Rebuild before flushing, including after adding fields to the physical // schema: recovery must retain the columns filled into bulk writes. From 00b2c57c0afb47886110db8ba64bc42bf198bdf3 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Sat, 31 Oct 2026 15:38:08 +0000 Subject: [PATCH 148/239] docs: link the observability database definition from README (#10061) * docs: link the observability database definition from README Signed-off-by: Dennis Zhuang * ci: add the new required checks to the Docs CI placeholders The main ruleset now requires five more checks. Their workflows skip docs-only PRs via paths-ignore, so without placeholders those PRs can never satisfy the required checks. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- .github/workflows/docs.yml | 5 +++++ README.md | 4 ++-- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index c9a6e8caaf3..fa8bc7af2d9 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -74,5 +74,10 @@ jobs: - coverage - Sqlness Test (Flat format) - Compatibility Test (recent releases) + - GitHub Script Tests (Rust CI) + - GitHub Script Tests (Integration CI) + - Check for conflict + - Check Unused Dependencies + - Build GreptimeDB binaries (ubuntu-latest) steps: - run: 'echo "No action required"' diff --git a/README.md b/README.md index 96058cc79b7..5204dab9513 100644 --- a/README.md +++ b/README.md @@ -31,7 +31,7 @@ One engine, on your infrastructure. ## Introduction -**GreptimeDB** is an open-source observability database. Metrics, logs, and traces run on one columnar engine over object storage and share one [table model](https://docs.greptime.com/user-guide/concepts/data-model/): tags, timestamp, and fields. When signals carry common identifiers such as service, host, or trace ID, you can correlate them in SQL without moving data between databases. +**GreptimeDB** is an open-source [observability database](https://docs.greptime.com/user-guide/concepts/why-greptimedb/#what-an-observability-database-is). Metrics, logs, and traces run on one columnar engine over object storage and share one [table model](https://docs.greptime.com/user-guide/concepts/data-model/): tags, timestamp, and fields. When signals carry common identifiers such as service, host, or trace ID, you can correlate them in SQL without moving data between databases. Ingest through OpenTelemetry, Prometheus Remote Write, Loki Push, or Elasticsearch Bulk. Use SQL across observability data and PromQL for metrics. Migrate ingestion one signal at a time without rebuilding your collectors. @@ -202,7 +202,7 @@ See the [Contribution Guidelines](CONTRIBUTING.md) for the full developer workfl GreptimeDB is generally available, with stable APIs and regular releases. It runs in production at scale — [OceanBase Cloud](https://greptime.com/blogs/2025-07-22-user-case-obcloud-log-management-greptimedb) operates 80+ GreptimeDB clusters managing 300 TB of logs, cutting log storage cost by 60%+ after migrating from Grafana Loki. See more in [case studies](https://greptime.com/blogs/?category=Use%20Case). -Release lines and support windows are in the [version reference](https://docs.greptime.com/nightly/reference/about-greptimedb-version). For where the project is going, read the [v1.0 highlights](https://greptime.com/blogs/2025-11-05-greptimedb-v1-highlights) and the [2026 roadmap](https://greptime.com/blogs/2026-02-11-greptimedb-roadmap-2026). +Release lines and support windows are in the [version reference](https://docs.greptime.com/nightly/reference/about-greptimedb-version). For where the project is going, read the [v1.0 highlights](https://greptime.com/blogs/2025-11-05-greptimedb-v1-highlights), the [2026 roadmap](https://greptime.com/blogs/2026-02-11-greptimedb-roadmap-2026), and [what we mean by an observability database](https://greptime.com/blogs/2026-08-26-what-is-an-observability-database). ## Community From c40e86c6533eac94b28652f0255f7cd07ea86c15 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Sat, 31 Oct 2026 16:25:49 +0000 Subject: [PATCH 149/239] fix(datanode): pick a node epoch above the recorded one (#9984) * fix(datanode): pick a node epoch above the recorded one The metasrv fences a datanode whose node epoch is lower than the highest one it has recorded for the node id. A datanode used its start time as its epoch, so after a restart on a host whose clock is behind the previous process's, or after the clock stepped back, the new process picked a lower epoch and was fenced for good: no region lease, no keep-alive. Before it creates the heartbeat task, the datanode now reads the recorded epoch through its metadata backend and uses one more than it if its start time is not higher. The node info still reports the local start time. A failed read fails the startup, like the other metadata reads. Signed-off-by: jeremyhi * test(datanode): tighten the node epoch tests Use a lower record that leaves room below the start time, so the test fails if the start time is ignored, and a nonzero node id, so it fails if the node id is ignored. Clarify the doc comment of the epoch choice. This assumes the previous process of a node id has stopped before a new one starts. Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- src/datanode/src/heartbeat.rs | 99 +++++++++++++++++++++++++-- src/meta-srv/src/service/heartbeat.rs | 14 ++++ 2 files changed, 108 insertions(+), 5 deletions(-) diff --git a/src/datanode/src/heartbeat.rs b/src/datanode/src/heartbeat.rs index f8d92212728..3bf1f6b5365 100644 --- a/src/datanode/src/heartbeat.rs +++ b/src/datanode/src/heartbeat.rs @@ -32,6 +32,8 @@ use common_meta::heartbeat::handler::{ }; use common_meta::heartbeat::mailbox::{HeartbeatMailbox, MailboxRef}; use common_meta::heartbeat::utils::outgoing_message_to_mailbox_message; +use common_meta::key::MetadataValue; +use common_meta::key::datanode_epoch::{DatanodeEpochKey, DatanodeEpochValue}; use common_meta::kv_backend::KvBackendRef; use common_stat::ResourceStatRef; use common_telemetry::{debug, error, info, trace, warn}; @@ -47,7 +49,9 @@ use tokio::time::Instant; use self::handler::RegionHeartbeatResponseHandler; use crate::alive_keeper::{CountdownTaskHandlerExtRef, RegionAliveKeeper}; use crate::config::DatanodeOptions; -use crate::error::{self, MetaClientInitSnafu, RegionEngineNotFoundSnafu, Result}; +use crate::error::{ + self, GetMetadataSnafu, MetaClientInitSnafu, RegionEngineNotFoundSnafu, Result, +}; use crate::event_listener::RegionServerEventReceiver; use crate::metrics::{self, HEARTBEAT_RECV_COUNT, HEARTBEAT_SENT_COUNT}; use crate::region_server::RegionServer; @@ -55,11 +59,29 @@ use crate::region_server::RegionServer; pub(crate) mod handler; pub(crate) mod task_tracker; +/// Returns the epoch of this process: its start time, or one more than the epoch the metasrv has +/// recorded for the node when the record is not lower than the start time, as after a restart on +/// a slower clock or after the clock stepped back. The metasrv fences an epoch lower than the +/// recorded one. +async fn node_epoch_above_record( + kv_backend: &KvBackendRef, + node_id: u64, + start_time_ms: u64, +) -> Result { + let key = DatanodeEpochKey::new(node_id).to_bytes(); + let Some(kv) = kv_backend.get(&key).await.context(GetMetadataSnafu)? else { + return Ok(start_time_ms); + }; + let recorded = DatanodeEpochValue::try_from_raw_value(&kv.value).context(GetMetadataSnafu)?; + Ok(start_time_ms.max(recorded.node_epoch + 1)) +} + /// The datanode heartbeat task which sending `[HeartbeatRequest]` to Metasrv periodically in background. pub struct HeartbeatTask { node_id: u64, workload_types: Vec, node_epoch: u64, + start_time_ms: u64, peer_addr: String, running: Arc, meta_client: MetaClientRef, @@ -87,6 +109,9 @@ impl HeartbeatTask { plugins: Plugins, resource_stat: ResourceStatRef, ) -> Result { + let node_id = opts.node_id.unwrap_or(0); + let start_time_ms = common_time::util::current_time_millis() as u64; + let node_epoch = node_epoch_above_record(&kv_backend, node_id, start_time_ms).await?; let countdown_task_handler_ext = plugins.get::(); let region_alive_keeper = Arc::new(RegionAliveKeeper::new( region_server.clone(), @@ -105,10 +130,10 @@ impl HeartbeatTask { ])); Ok(Self { - node_id: opts.node_id.unwrap_or(0), + node_id, workload_types: opts.workload_types.clone(), - // We use datanode's start time millis as the node's epoch. - node_epoch: common_time::util::current_time_millis() as u64, + node_epoch, + start_time_ms, peer_addr: addrs::resolve_addr(&opts.grpc.bind_addr, Some(&opts.grpc.server_addr)), running: Arc::new(AtomicBool::new(false)), meta_client, @@ -227,6 +252,7 @@ impl HeartbeatTask { } let node_id = self.node_id; let node_epoch = self.node_epoch; + let start_time_ms = self.start_time_ms; let addr = &self.peer_addr; let meta_client = self.meta_client.clone(); @@ -301,7 +327,7 @@ impl HeartbeatTask { info: Some(NodeInfo { version: build_info.version.to_string(), git_commit: build_info.commit_short.to_string(), - start_time_ms: node_epoch, + start_time_ms, total_cpu_millicores, total_memory_bytes, cpu_usage_millicores: 0, @@ -471,3 +497,66 @@ impl HeartbeatTask { Ok(()) } } + +#[cfg(test)] +mod tests { + use common_meta::cache_invalidator::DummyCacheInvalidator; + use common_meta::kv_backend::memory::MemoryKvBackend; + use common_meta::rpc::store::PutRequest; + use common_stat::ResourceStatImpl; + use meta_client::client::MetaClientBuilder; + + use super::*; + use crate::tests::mock_region_server; + + const NODE_ID: u64 = 7; + + async fn kv_backend_with_record(node_epoch: Option) -> KvBackendRef { + let kv_backend: KvBackendRef = Arc::new(MemoryKvBackend::new()); + if let Some(node_epoch) = node_epoch { + let key = DatanodeEpochKey::new(NODE_ID).to_bytes(); + let value = DatanodeEpochValue { node_epoch } + .try_as_raw_value() + .unwrap(); + let req = PutRequest::new().with_key(key).with_value(value); + kv_backend.put(req).await.unwrap(); + } + kv_backend + } + + #[tokio::test] + async fn test_node_epoch_above_record() { + for (record, expected) in [ + (None, 1000), + (Some(500), 1000), + (Some(1000), 1001), + (Some(5000), 5001), + ] { + let kv_backend = kv_backend_with_record(record).await; + let node_epoch = node_epoch_above_record(&kv_backend, NODE_ID, 1000).await; + assert_eq!(node_epoch.unwrap(), expected); + } + } + + #[tokio::test] + async fn test_start_time_stays_local_above_a_record() { + let now = common_time::util::current_time_millis() as u64; + let record = now + 60_000; + let task = HeartbeatTask::try_new( + &DatanodeOptions { + node_id: Some(NODE_ID), + ..Default::default() + }, + mock_region_server(), + Arc::new(MetaClientBuilder::datanode_default_options(NODE_ID).build()), + kv_backend_with_record(Some(record)).await, + Arc::new(DummyCacheInvalidator), + Plugins::default(), + Arc::new(ResourceStatImpl::default()), + ) + .await + .unwrap(); + assert_eq!(task.node_epoch, record + 1); + assert!(task.start_time_ms >= now && task.start_time_ms < record); + } +} diff --git a/src/meta-srv/src/service/heartbeat.rs b/src/meta-srv/src/service/heartbeat.rs index c4f4978e746..67ccb9df49d 100644 --- a/src/meta-srv/src/service/heartbeat.rs +++ b/src/meta-srv/src/service/heartbeat.rs @@ -1311,6 +1311,20 @@ mod tests { assert!(accept_node_epoch(&cached, 1, e2).await.unwrap()); } + #[tokio::test] + async fn test_epoch_above_a_record_ahead_of_the_clock_is_accepted() { + let group = new_group(); + let ctx = TestEnv::new().ctx(); + // The previous process of the node ran on a faster clock. + let recorded = current_time_millis() as u64 + 60_000; + record(&ctx, recorded).await; + + let mut restarted = Stream::open(&group, &ctx).await; + let res = restarted.send(heartbeat(recorded + 1, "a2")).await; + assert!(res.region_lease.is_some()); + assert_eq!(lease_addr(&ctx).await.as_deref(), Some("a2")); + } + #[tokio::test] async fn test_epoch_far_ahead_is_not_recorded() { let kv_backend = MemoryKvBackend::new(); From a57ab18eaa1b6823163cb4ea64173680b659ffdd Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Mon, 2 Nov 2026 03:09:13 +0000 Subject: [PATCH 150/239] feat(ingest): report confirmed progress for batched writes (#9921) * feat(ingest): track confirmed writes through batching and retries Batched writes can partially reach storage before another partition fails. Callers need the successful row count and the failure together to report accurate ingestion results. Add WriteProgress and a confirmed-write context, carry row ownership through partitioning and retry rebuilds, and retain completions through ordinary and logical-table batchers. Existing callers retain the completion/result wrappers. Tests cover partial progress, request ownership, batching and retries. This supplies the accounting required by Remote Write exemplar written-count headers. Signed-off-by: shuiyisong * fix(servers): coalesce flow notifications Signed-off-by: shuiyisong * perf(ingest): reduce write progress overhead Signed-off-by: shuiyisong * fix: preserve retry hints across error wrappers Signed-off-by: shuiyisong * fix(ingest): drain confirmed writes on cancellation Signed-off-by: shuiyisong * perf(batcher): reduce row owner allocations Signed-off-by: shuiyisong * fix: sqlness Signed-off-by: shuiyisong * fix(batcher): restore per-flush metrics Signed-off-by: shuiyisong * revert: scope retry-hint fix to storage writes Signed-off-by: shuiyisong * chore: comments Signed-off-by: shuiyisong * fix(ingest): restore flush metrics and benchmark Signed-off-by: shuiyisong * fix(ingest): propagate sync batch write errors Signed-off-by: shuiyisong * fix: preserve peer batching for confirmed writes Signed-off-by: shuiyisong * perf: avoid unnecessary row owner copies Signed-off-by: shuiyisong * perf: reuse IPC buffers in tracked retries Signed-off-by: shuiyisong * style: order write dispatcher imports Signed-off-by: shuiyisong * fix: rebase main Signed-off-by: shuiyisong --------- Signed-off-by: shuiyisong --- src/common/error/src/mock.rs | 20 +- src/common/query/src/error.rs | 42 +- src/common/query/src/lib.rs | 68 +- src/datanode/src/region_server.rs | 61 +- src/operator/src/bulk_insert.rs | 33 + src/operator/src/insert.rs | 303 ++++++++- .../src/req_convert/common/partitioner.rs | 394 +++++++++-- src/operator/src/write_dispatcher.rs | 118 +++- src/operator/src/write_dispatcher/progress.rs | 333 +++++++++ src/operator/src/write_dispatcher/tests.rs | 639 +++++++++++++++++- src/servers/benches/flush_batch_physical.rs | 25 +- src/servers/src/batcher/logical_table.rs | 279 +++++++- .../src/batcher/logical_table/batch.rs | 537 +++++++++++++-- .../batcher/logical_table/batch_convert.rs | 2 + .../batcher/logical_table/pending_worker.rs | 32 +- .../src/batcher/logical_table/region_write.rs | 233 ++++--- src/servers/src/batcher/table.rs | 133 +++- src/servers/src/batcher/table/batch.rs | 84 +-- .../src/batcher/table/flow_notifier.rs | 234 ++++++- .../src/batcher/table/pending_batch.rs | 27 +- src/servers/src/error.rs | 46 +- src/servers/src/metrics.rs | 4 +- src/session/src/context.rs | 22 + 23 files changed, 3284 insertions(+), 385 deletions(-) create mode 100644 src/operator/src/write_dispatcher/progress.rs diff --git a/src/common/error/src/mock.rs b/src/common/error/src/mock.rs index 572e11deabb..499135d4812 100644 --- a/src/common/error/src/mock.rs +++ b/src/common/error/src/mock.rs @@ -17,26 +17,38 @@ use std::any::Any; use std::fmt; -use crate::ext::{ErrorExt, StackError}; +use crate::ext::{ErrorExt, RetryHint, StackError}; use crate::status_code::StatusCode; /// A mock error mainly for test. #[derive(Debug)] pub struct MockError { pub code: StatusCode, + retry_hint: RetryHint, source: Option>, } impl MockError { /// Create a new [MockError] without backtrace. pub fn new(code: StatusCode) -> MockError { - MockError { code, source: None } + MockError { + code, + retry_hint: RetryHint::NonRetryable, + source: None, + } + } + + /// Sets the retry hint independently of the status code. + pub fn with_retry_hint(mut self, retry_hint: RetryHint) -> Self { + self.retry_hint = retry_hint; + self } /// Create a new [MockError] with source. pub fn with_source(source: MockError) -> MockError { MockError { code: source.code, + retry_hint: source.retry_hint, source: Some(Box::new(source)), } } @@ -62,6 +74,10 @@ impl ErrorExt for MockError { fn as_any(&self) -> &dyn Any { self } + + fn retry_hint(&self) -> RetryHint { + self.retry_hint + } } impl StackError for MockError { diff --git a/src/common/query/src/error.rs b/src/common/query/src/error.rs index 4c53be86b6e..511f0a82db3 100644 --- a/src/common/query/src/error.rs +++ b/src/common/query/src/error.rs @@ -13,9 +13,10 @@ // limitations under the License. use std::any::Any; +use std::sync::Arc; use arrow::error::ArrowError; -use common_error::ext::{BoxedError, ErrorExt}; +use common_error::ext::{BoxedError, ErrorExt, RetryHint}; use common_error::status_code::StatusCode; use common_macro::stack_trace_debug; use datafusion_common::DataFusionError; @@ -29,6 +30,12 @@ use snafu::{Location, Snafu}; #[snafu(visibility(pub))] #[stack_trace_debug] pub enum Error { + #[snafu(display("Storage write did not complete: {source}"))] + StorageWrite { + source: Arc, + #[snafu(implicit)] + location: Location, + }, #[snafu(display("Unsupported input datatypes {:?} in function {}", datatypes, function))] UnsupportedInputDataType { function: String, @@ -218,6 +225,7 @@ pub type Result = std::result::Result; impl ErrorExt for Error { fn status_code(&self) -> StatusCode { match self { + Error::StorageWrite { source, .. } => source.status_code(), Error::InvalidInputState { .. } | Error::ToScalarValue { .. } | Error::GetScalarVector { .. } @@ -257,6 +265,13 @@ impl ErrorExt for Error { fn as_any(&self) -> &dyn Any { self } + + fn retry_hint(&self) -> RetryHint { + match self { + Error::StorageWrite { source, .. } => source.retry_hint(), + _ => RetryHint::NonRetryable, + } + } } impl From for DataFusionError { @@ -301,9 +316,34 @@ mod tests { use std::sync::Arc; use common_error::ext::PlainError; + use common_error::mock::MockError; use super::*; + #[tokio::test] + async fn test_write_completion_preserves_retry_hint() { + for hint in [ + RetryHint::NonRetryable, + RetryHint::Retryable, + RetryHint::RetryAfterRouteRefresh, + ] { + let source = + || BoxedError::new(MockError::new(StatusCode::RegionBusy).with_retry_hint(hint)); + let progress = crate::WriteProgress::failed(source()); + let error = progress.clone().into_result().unwrap_err(); + assert_eq!(error.status_code(), StatusCode::RegionBusy); + assert_eq!(error.retry_hint(), hint); + assert_eq!(progress.error.unwrap().retry_hint(), hint); + + let error = crate::WriteCompletion::new(std::future::ready(Err(source()))) + .wait() + .await + .unwrap_err(); + assert_eq!(error.status_code(), StatusCode::RegionBusy); + assert_eq!(error.retry_hint(), hint); + } + } + #[test] fn test_datafusion_status_code_preserves_external_errors_through_wrappers() { let boxed_error = |status| { diff --git a/src/common/query/src/lib.rs b/src/common/query/src/lib.rs index 8151dfc5731..c68e4c76470 100644 --- a/src/common/query/src/lib.rs +++ b/src/common/query/src/lib.rs @@ -38,6 +38,7 @@ use common_recordbatch::{ }; use datafusion::physical_plan::ExecutionPlan; use serde::{Deserialize, Serialize}; +use snafu::IntoError; use sqlparser_derive::{Visit, VisitMut}; /// new Output struct with output data(previously Output) and output meta @@ -77,23 +78,68 @@ pub struct OutputMeta { /// May exist for query output. One can retrieve execution metrics from this plan. pub plan: Option>, pub cost: OutputCost, - /// Process-local storage results attached by asynchronous ordinary-table batching. + /// Process-local storage results retained for asynchronous batching or confirmed writes. /// Callers needing confirmed writes must await these. Dropping them does not cancel writes. pub write_completions: Vec, } -/// Actual storage completion of an already accepted asynchronous write. -pub struct WriteCompletion(Pin> + Send>>); +/// Storage outcome of an already accepted write, whether synchronous or asynchronous. +pub struct WriteCompletion(Pin + Send>>); + +/// Acknowledged rows for one original input, even when another part failed. +/// Failed or lost acknowledgements do not establish how many rows were persisted. +#[derive(Debug, Default, Clone)] +pub struct WriteProgress { + pub affected_rows: usize, + pub error: Option>, +} + +impl WriteProgress { + pub fn failed(error: BoxedError) -> Self { + Self { + affected_rows: 0, + error: Some(Arc::new(error)), + } + } + + pub fn merge(&mut self, other: Self) { + self.affected_rows += other.affected_rows; + if self.error.is_none() { + self.error = other.error; + } + } + + pub fn into_result(self) -> Result { + match self.error { + Some(error) => Err(BoxedError::new(error::StorageWriteSnafu.into_error(error))), + None => Ok(self.affected_rows), + } + } +} impl WriteCompletion { /// Wraps the storage result without changing how the write is scheduled. + /// Success reports zero affected rows; use [`Self::with_progress`] for counted writes. pub fn new(completion: impl Future> + Send + 'static) -> Self { + Self::with_progress(async move { + match completion.await { + Ok(()) => WriteProgress::default(), + Err(error) => WriteProgress::failed(error), + } + }) + } + + pub fn with_progress(completion: impl Future + Send + 'static) -> Self { Self(Box::pin(completion)) } + pub async fn progress(self) -> WriteProgress { + self.0.await + } + /// Waits for storage success or failure, independently of queue admission. pub async fn wait(self) -> Result<(), BoxedError> { - self.0.await + self.0.await.into_result().map(|_| ()) } } @@ -104,6 +150,20 @@ impl Debug for WriteCompletion { } impl Output { + /// Observes every dispatched completion before returning any failure. + pub async fn write_progress(self) -> WriteProgress { + if self.meta.write_completions.is_empty() { + return WriteProgress { + affected_rows: self.extract_rows_and_cost().0, + error: None, + }; + } + let mut progress = WriteProgress::default(); + for completion in self.meta.write_completions { + progress.merge(completion.progress().await); + } + progress + } pub fn new_with_affected_rows(affected_rows: OutputRows) -> Self { Self { data: OutputData::AffectedRows(affected_rows), diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index 4ae350002c9..9c60972fbc3 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -652,17 +652,20 @@ impl RegionServer { let first_region_id = first.region_id; let request_type = first.request.request_type(); - // Upstream batching keeps requests from the same engine together. - // Ordinary execution reports engine lookup failures per entry. + // Ordinary execution handles mixed engines and unavailable regions per entry. let Ok(CurrentEngine::Engine(engine)) = self.inner.get_engine(first_region_id, &RegionChange::None) else { return Ok(Either::Right(requests)); }; if engine.name() != METRIC_ENGINE_NAME - || !requests - .iter() - .all(|request| matches!(request.request, RegionRequest::Put(_))) + || !requests.iter().all(|request| { + matches!(request.request, RegionRequest::Put(_)) + && matches!( + self.inner.region_map.get(&request.region_id).as_deref(), + Some(RegionEngineWithStatus::Ready(engine)) if engine.name() == METRIC_ENGINE_NAME + ) + }) { return Ok(Either::Right(requests)); } @@ -3630,9 +3633,11 @@ mod tests { let mut env = TestEnv::with_prefix("region-server-metric-results").await; let mito = env.create_engine(MitoConfig::default()).await; let metric = Arc::new( - MetricEngine::try_new(mito, metric_engine::config::EngineConfig::default()).unwrap(), + MetricEngine::try_new(mito.clone(), metric_engine::config::EngineConfig::default()) + .unwrap(), ); let mut server = mock_region_server(); + server.register_engine(Arc::new(mito)); server.register_engine(metric.clone()); let mut inserts = Vec::new(); for table in [10, 20] { @@ -3731,6 +3736,50 @@ mod tests { .iter() .all(|(_, error)| error.retry_hint() == RetryHint::RetryAfterRouteRefresh) ); + + let regular_region = RegionId::new(30, 0); + let mut create = CreateRequestBuilder::new().build(); + create.table_dir = "metric-results/30/".to_string(); + let regular = InsertRequest { + region_id: regular_region.as_u64(), + rows: Some(Rows { + schema: rows_schema(&create), + rows: build_rows(0, 3), + }), + ..Default::default() + }; + server + .handle_request(regular_region, RegionRequest::Create(create)) + .await + .unwrap(); + // A metric entry must not route its Mito2 sibling through the metric engine. + for (metric_insert, metric_failed) in inserts.iter().zip([false, true]) { + for metric_first in [true, false] { + let requests = if metric_first { + vec![metric_insert.clone(), regular.clone()] + } else { + vec![regular.clone(), metric_insert.clone()] + }; + let response = server + .handle_requests_in_parallel(region_request::Body::Inserts(InsertRequests { + requests, + })) + .await + .unwrap(); + assert_eq!( + response.affected_rows, + 3 + usize::from(!metric_failed), + "{:?}", + response.failures + ); + assert_eq!(response.failures.len(), usize::from(metric_failed)); + if let Some((index, error)) = response.failures.first() { + assert_eq!(*index, usize::from(!metric_first)); + assert_eq!(error.status_code(), StatusCode::RegionNotReady); + assert_eq!(error.retry_hint(), RetryHint::RetryAfterRouteRefresh); + } + } + } } #[tokio::test] diff --git a/src/operator/src/bulk_insert.rs b/src/operator/src/bulk_insert.rs index 3144daec3a0..d17cdcf6743 100644 --- a/src/operator/src/bulk_insert.rs +++ b/src/operator/src/bulk_insert.rs @@ -45,6 +45,39 @@ use crate::write_dispatcher::{WriteAttempt, WriteAttribution}; use crate::{error, metrics}; impl Inserter { + /// Writes a coalesced batch while retaining each original input's acknowledgements. + pub async fn flush_bulk_batch_tracked( + &self, + table_info: TableInfoRef, + batch: RecordBatch, + ctx: QueryContextRef, + owners: Vec, + inputs: usize, + ) -> Result> { + let requests = TablePartitioner::try_new(&self.partition_manager, &table_info) + .await + .context(error::InvalidPartitionSnafu)? + .partition_tracked_bulk_insert_requests(batch, owners, ctx.skip_wal())?; + let requests = requests + .into_iter() + .map(|(request, owners)| crate::write_dispatcher::TrackedWrite { + request: RegionRequest { + header: Some(RegionRequestHeader { + dbname: ctx.get_db_string(), + tracing_context: TracingContext::from_current_span().to_w3c(), + ..Default::default() + }), + body: Some(region_request::Body::BulkInsert(request)), + }, + owners, + }) + .collect(); + Ok(self + .write_dispatcher + .dispatch_tracked(requests, inputs, None) + .await) + } + /// Routes and writes a prepared table batch, returning the affected row count. /// /// Callers must exclude instant-TTL tables and handle Flow notifications after diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index 3dabb963bdd..5ecc01457a1 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -575,10 +575,31 @@ impl Inserter { }; let mut affected_rows = 0; let mut meta = OutputMeta::new_with_cost(write_cost as _); + let mut first_error = None; for result in results { - let output = result?; - affected_rows += output.extract_rows_and_cost().0; - meta.write_completions.extend(output.meta.write_completions); + match result { + Ok(output) => { + affected_rows += output.extract_rows_and_cost().0; + meta.write_completions.extend(output.meta.write_completions); + } + Err(error) if ctx.confirmed_writes() => { + meta.write_completions + .push(common_query::WriteCompletion::with_progress(async move { + common_query::WriteProgress::failed(common_error::ext::BoxedError::new( + error, + )) + })); + } + Err(error) => { + first_error.get_or_insert(error); + } + } + } + if let Some(error) = first_error { + for completion in meta.write_completions { + let _ = completion.wait().await; + } + return Err(error); } Ok(Output::new(OutputData::AffectedRows(affected_rows), meta)) } @@ -856,6 +877,54 @@ impl Inserter { )?; let dispatcher = self.write_dispatcher.clone(); + if ctx.confirmed_writes() { + let writes = normal_requests + .requests + .into_iter() + .map(|request| { + let rows = request.rows.as_ref().map_or(0, |rows| rows.rows.len()); + crate::write_dispatcher::TrackedWrite { + request: request_factory.build_insert(RegionInsertRequests { + requests: vec![request], + }), + owners: vec![0; rows], + } + }) + .collect(); + let (writes, progress) = dispatcher.prepare_tracked(writes, 1).await; + let cancellation = CancellationToken::new(); + let _guard = cancellation.clone().drop_guard(); + let counter = ctx.write_progress().cloned(); + let progress = common_runtime::spawn_global(async move { + let progress = dispatcher + .execute_tracked(writes, progress, None, Some(&cancellation)) + .await; + if let Some(counter) = counter { + counter.add_written(progress.iter().map(|part| part.affected_rows).sum()); + } + progress + }) + .await + .context(JoinTaskSnafu)? + .into_iter() + .fold( + common_query::WriteProgress::default(), + |mut result, part| { + result.merge(part); + result + }, + ); + crate::metrics::DIST_INGEST_ROW_COUNT + .with_label_values(&[ctx.get_db_string().as_str()]) + .inc_by(progress.affected_rows as u64); + let mut meta = OutputMeta::new_with_cost(write_cost as _); + let rows = progress.affected_rows; + meta.write_completions + .push(common_query::WriteCompletion::with_progress(async move { + progress + })); + return Ok(Output::new(OutputData::AffectedRows(rows), meta)); + } let request = request_factory.build_insert(normal_requests); let writes = dispatcher.prepare(vec![request]).await?; // Dropping the caller stops retries without abandoning in-flight writes. @@ -2320,16 +2389,26 @@ fn release_mirror_pending_rows(mirror_pending_rows: &AtomicU64, rows: u64) { #[cfg(test)] mod tests { + use std::future::Future; use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; + use std::task::Poll; use api::helper::ColumnDataTypeWrapper; + use api::region::RegionResponse; use api::v1::flow::FlowResponse; use api::v1::helper::{field_column_schema, time_index_column_schema}; use api::v1::{Row, RowInsertRequest, Rows, Value}; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; + use common_error::ext::{ErrorExt, RetryHint}; + use common_error::mock::MockError; + use common_error::status_code::StatusCode; use common_meta::cache::new_table_flownode_set_cache; + use common_meta::cache_invalidator::DummyCacheInvalidator; use common_meta::ddl::test_util::datanode_handler::NaiveDatanodeHandler; + use common_meta::node_manager::{ + Datanode, DatanodeManager, DatanodeRef, FlownodeManager, FlownodeRef, + }; use common_meta::test_util::{MockDatanodeManager, MockFlownodeHandler, MockFlownodeManager}; use common_query::native_histogram::NATIVE_HISTOGRAM_FIELD; use common_query::prelude::{greptime_native_histogram, set_default_prefix}; @@ -2340,11 +2419,13 @@ mod tests { use table::TableRef; use table::dist_table::DummyDataSource; use table::metadata::{TableInfoBuilder, TableMetaBuilder, TableType}; + use tokio::sync::{Semaphore, mpsc}; use crate::insert::*; use crate::test_util::{ create_partition_rule_manager, new_test_table_info, prepare_mocked_backend, }; + use crate::write_dispatcher::{TrackedWrite, WriteRetryOptions}; fn make_table_ref_with_schema( ts_name: &str, @@ -2732,6 +2813,222 @@ mod tests { assert!(table_is_native_histogram(&table)); } + #[derive(Clone)] + struct GatedDatanodeHandler { + started: mpsc::UnboundedSender, + completed: mpsc::UnboundedSender, + release: Arc, + } + + #[async_trait::async_trait] + impl Datanode for GatedDatanodeHandler { + async fn handle_request( + &self, + request: Arc, + ) -> common_meta::error::Result { + let Some(api::v1::region::region_request::Body::Inserts(requests)) = &request.body + else { + panic!("expected inserts"); + }; + let region = RegionId::from(requests.requests[0].region_id).region_number(); + self.started.send(region).unwrap(); + self.release.acquire().await.unwrap().forget(); + self.completed.send(region).unwrap(); + if region == 2 { + Err(common_meta::error::Error::retry_later(MockError::new( + StatusCode::RegionBusy, + ))) + } else { + Ok(RegionResponse::new(1)) + } + } + + async fn handle_query( + &self, + _request: common_query::request::QueryRequest, + ) -> common_meta::error::Result { + unreachable!() + } + } + + #[async_trait::async_trait] + impl DatanodeManager for GatedDatanodeHandler { + async fn datanode(&self, _peer: &Peer) -> DatanodeRef { + Arc::new(self.clone()) + } + } + + #[async_trait::async_trait] + impl FlownodeManager for GatedDatanodeHandler { + async fn flownode(&self, _peer: &Peer) -> FlownodeRef { + unreachable!() + } + } + + #[tokio::test] + async fn test_insert_cancellation_drains_inflight_writes() { + for confirmed in [false, true] { + for max_retries in [0, 2] { + let mut inserter = batcher_test_inserter().await; + let (started, mut starts) = mpsc::unbounded_channel(); + let (completed, mut completions) = mpsc::unbounded_channel(); + let release = Arc::new(Semaphore::new(0)); + let node_manager = Arc::new(GatedDatanodeHandler { + started, + completed, + release: release.clone(), + }); + inserter.node_manager = node_manager.clone(); + let dispatcher = Arc::new( + WriteDispatcher::new(node_manager, inserter.partition_manager.clone()) + .with_retry( + WriteRetryOptions { + max_retries, + initial_backoff: Duration::ZERO, + max_backoff: Duration::ZERO, + ..Default::default() + }, + inserter.catalog_manager.clone(), + Arc::new(DummyCacheInvalidator), + ), + ); + let inserter = inserter.with_write_dispatcher(dispatcher); + let table_infos = HashMap::from_iter([( + 1, + Arc::new(new_test_table_info(1, "table_1", [1].into_iter())), + )]); + let counter = Arc::new(session::write_progress::WriteProgress::default()); + let mut ctx = QueryContext::arc().fork(); + ctx.set_write_progress(counter.clone()); + let ctx = Arc::new(ctx); + let ctx = if confirmed { + Arc::new(ctx.with_confirmed_writes()) + } else { + ctx + }; + let requests = InstantAndNormalInsertRequests { + normal_requests: RegionInsertRequests { + requests: [1, 2] + .map(|region| RegionInsertRequest { + region_id: RegionId::new(1, region).as_u64(), + rows: Some(Rows { + schema: vec![], + rows: vec![Row { values: vec![] }], + }), + ..Default::default() + }) + .to_vec(), + }, + instant_requests: RegionInsertRequests::default(), + }; + let caller = + tokio::spawn( + async move { inserter.do_request(requests, &table_infos, &ctx).await }, + ); + for _ in 0..2 { + tokio::time::timeout(Duration::from_secs(5), starts.recv()) + .await + .unwrap() + .unwrap(); + } + caller.abort(); + assert!(caller.await.unwrap_err().is_cancelled()); + release.add_permits(2); + // Closing the channel proves the detached dispatcher has finished. + let mut completed = tokio::time::timeout(Duration::from_secs(5), async { + let mut regions = Vec::new(); + while let Some(region) = completions.recv().await { + regions.push(region); + } + regions + }) + .await + .unwrap(); + completed.sort_unstable(); + assert_eq!( + completed, + vec![1, 2], + "confirmed={confirmed}, retries={max_retries}" + ); + assert!( + starts.try_recv().is_err(), + "cancelled caller started another retry" + ); + assert_eq!(counter.written(), 1); + } + } + } + + #[tokio::test] + async fn test_tracked_cancellation_stops_backoff() { + let inserter = batcher_test_inserter().await; + let (started, mut starts) = mpsc::unbounded_channel(); + let (completed, mut completions) = mpsc::unbounded_channel(); + let node_manager = Arc::new(GatedDatanodeHandler { + started, + completed, + release: Arc::new(Semaphore::new(2)), + }); + let dispatcher = WriteDispatcher::new(node_manager, inserter.partition_manager.clone()) + .with_retry( + WriteRetryOptions { + max_retries: 2, + initial_backoff: Duration::from_secs(60), + max_backoff: Duration::from_secs(60), + ..Default::default() + }, + inserter.catalog_manager.clone(), + Arc::new(DummyCacheInvalidator), + ); + let factory = RegionRequestFactory::new(RegionRequestHeader::default()); + let requests = [1, 2] + .into_iter() + .map(|region| TrackedWrite { + request: factory.build_insert(RegionInsertRequests { + requests: vec![RegionInsertRequest { + region_id: RegionId::new(1, region).as_u64(), + ..Default::default() + }], + }), + owners: vec![(region - 1) as usize], + }) + .collect(); + let (writes, progress) = dispatcher.prepare_tracked(requests, 2).await; + let cancellation = CancellationToken::new(); + let execute = dispatcher.execute_tracked(writes, progress, None, Some(&cancellation)); + tokio::pin!(execute); + // Wait for both sends to complete before cancelling the backoff. + tokio::time::timeout( + Duration::from_secs(1), + future::poll_fn(|cx| { + assert!(execute.as_mut().poll(cx).is_pending()); + if completions.len() == 2 { + Poll::Ready(()) + } else { + Poll::Pending + } + }), + ) + .await + .unwrap(); + for region in [1, 2] { + assert_eq!(starts.try_recv().unwrap(), region); + assert_eq!(completions.try_recv().unwrap(), region); + } + cancellation.cancel(); + let progress = tokio::time::timeout(Duration::from_secs(1), execute) + .await + .unwrap(); + assert_eq!(progress[0].affected_rows, 1); + assert!(progress[0].error.is_none()); + assert_eq!(progress[1].affected_rows, 0); + assert_eq!( + progress[1].error.as_ref().unwrap().retry_hint(), + RetryHint::Retryable + ); + assert!(starts.try_recv().is_err()); + } + // Keep global meter registration in one test, isolated by nextest's per-test process. #[tokio::test] async fn test_write_meter_admission() { diff --git a/src/operator/src/req_convert/common/partitioner.rs b/src/operator/src/req_convert/common/partitioner.rs index 05ad400abf0..48900836b3d 100644 --- a/src/operator/src/req_convert/common/partitioner.rs +++ b/src/operator/src/req_convert/common/partitioner.rs @@ -179,6 +179,150 @@ impl TablePartitioner { .collect() } + /// Keeps process-local row ownership aligned when a failed write is repartitioned. + pub(crate) fn rebuild_tracked_write_requests( + &self, + request: RegionRequest, + table_info: &TableInfo, + owners: Option>, + ) -> Result>)>> { + let Some(owners) = owners else { + return Ok(self + .rebuild_write_requests(request, table_info)? + .into_iter() + .map(|request| (request, None)) + .collect()); + }; + let RegionRequest { header, body } = request; + if let Some(Body::BulkInsert(bulk)) = body { + let region_id = RegionId::from_u64(bulk.region_id); + if self.is_partition_expr_unchanged(region_id, bulk.partition_expr_version) { + return Ok(vec![( + RegionRequest { + header, + body: Some(Body::BulkInsert(bulk)), + }, + Some(owners), + )]); + } + let bulk_insert_request::Body::ArrowIpc(ipc) = + bulk.body.context(error::UnexpectedSnafu { + violated: "Cannot rebuild a bulk insert without Arrow IPC data", + })?; + let batch = FlightDecoder::try_from_schema_bytes(&ipc.schema) + .and_then(|mut decoder| { + decoder.try_decode_record_batch(&ipc.data_header, &ipc.payload) + }) + .map_err(BoxedError::new) + .context(error::ExternalSnafu)?; + let routing = self.restore_partition_tags(&batch, table_info)?; + let masks = self + .rule + .split_record_batch(routing.as_ref().unwrap_or(&batch)) + .context(SplitInsertSnafu)?; + return self + .build_tracked_bulk_insert_requests(batch, masks, owners, Some(ipc), bulk.skip_wal) + .map(|requests| { + requests + .into_iter() + .map(|(request, owners)| { + ( + RegionRequest { + header: header.clone(), + body: Some(Body::BulkInsert(request)), + }, + Some(owners), + ) + }) + .collect() + }); + } + let owner = owners.first().copied().context(error::UnexpectedSnafu { + violated: "Missing row ownership", + })?; + ensure!( + owners.iter().all(|value| *value == owner), + error::UnexpectedSnafu { + violated: "Mixed row ownership requires Arrow batching" + } + ); + self.rebuild_write_requests(RegionRequest { header, body }, table_info)? + .into_iter() + .map(|request| { + let count = match &request.body { + Some(Body::Inserts(inserts)) => inserts + .requests + .iter() + .filter_map(|r| r.rows.as_ref()) + .map(|r| r.rows.len()) + .sum(), + _ => { + return error::UnexpectedSnafu { + violated: "Unsupported tracked write", + } + .fail(); + } + }; + Ok((request, Some(vec![owner; count]))) + }) + .collect() + } + + pub(crate) fn partition_tracked_bulk_insert_requests( + &self, + batch: RecordBatch, + owners: Vec, + skip_wal: bool, + ) -> Result)>> { + let masks = self + .rule + .split_record_batch(&batch) + .context(SplitInsertSnafu)?; + self.build_tracked_bulk_insert_requests(batch, masks, owners, None, skip_wal) + } + + fn build_tracked_bulk_insert_requests( + &self, + batch: RecordBatch, + masks: HashMap, + owners: Vec, + original_ipc: Option, + skip_wal: bool, + ) -> Result)>> { + ensure!( + owners.len() == batch.num_rows(), + error::UnexpectedSnafu { + violated: "Write ownership does not match row count" + } + ); + let mut selected = HashMap::new(); + for (region, mask) in &masks { + let mut selected_owners = Vec::with_capacity(mask.selected_rows()); + selected_owners.extend( + mask.array() + .values() + .set_indices() + .map(|index| owners[index]), + ); + selected.insert( + RegionId::new(self.table_id, *region).as_u64(), + selected_owners, + ); + } + self.build_bulk_insert_requests(batch, masks, original_ipc, skip_wal)? + .into_iter() + .map(|request| { + let owners = + selected + .remove(&request.region_id) + .context(error::UnexpectedSnafu { + violated: "Missing partition ownership", + })?; + Ok((request, owners)) + }) + .collect() + } + /// Reads a table's partition rule and expression versions into one routing snapshot. pub async fn try_new( partition_manager: &PartitionRuleManager, @@ -530,7 +674,7 @@ mod tests { use common_meta::kv_backend::test_util::MockKvBackendBuilder; use common_meta::rpc::router::{Region, RegionRoute}; use common_meta::rpc::store::RangeResponse; - use datatypes::arrow::array::{Int32Array, StringArray}; + use datatypes::arrow::array::{BooleanArray, Int32Array, StringArray}; use datatypes::arrow::datatypes::{DataType, Field, Schema}; use datatypes::arrow::record_batch::RecordBatch; use datatypes::data_type::ConcreteDataType; @@ -1073,81 +1217,108 @@ mod tests { .find_table_partition_rule(&table_info) .await .unwrap(); - let request = build_bulk_request(99, Some(0), vec![1, 20, 70]); - let header = request.header.clone(); - let rerouted = rebuild_with_cached_table(&manager, &cache, request) + let partitioner = TablePartitioner::try_new(&manager, &table_info) .await .unwrap(); - assert_eq!(rerouted.len(), 3); - let mut actual = Vec::new(); - for request in rerouted { - assert_eq!(request.header, header); - let Some(Body::BulkInsert(bulk)) = request.body else { - panic!("expected bulk insert") - }; - let region = RegionId::from_u64(bulk.region_id); - assert_eq!(region.table_id(), 1); - assert!(bulk.skip_wal); - assert_eq!(bulk.aligned_schema_version, None); - assert_eq!( - bulk.partition_expr_version.map(|v| v.value), - versions[®ion.region_number()] - ); - let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = bulk.body else { - panic!("expected Arrow IPC") - }; - let batch = FlightDecoder::try_from_schema_bytes(&ipc.schema) - .unwrap() - .try_decode_record_batch(&ipc.data_header, &ipc.payload) + for tracked in [false, true] { + let request = build_bulk_request(99, Some(0), vec![1, 20, 70]); + let header = request.header.clone(); + let rerouted = partitioner + .rebuild_tracked_write_requests( + request, + &table_info, + tracked.then(|| vec![2, 0, 2]), + ) .unwrap(); - assert_eq!(batch.num_rows(), 1); - let values = batch - .column(0) - .as_any() - .downcast_ref::() - .unwrap(); - actual.push((region.region_number(), values.value(0))); - } - actual.sort_unstable(); - assert_eq!(actual, [(1, 70), (2, 20), (3, 1)]); + assert_eq!(rerouted.len(), 3); + let mut actual = Vec::new(); + for (request, owners) in rerouted { + assert_eq!(request.header, header); + let Some(Body::BulkInsert(bulk)) = request.body else { + panic!("expected bulk insert") + }; + let region = RegionId::from_u64(bulk.region_id); + assert_eq!(region.table_id(), 1); + assert!(bulk.skip_wal); + assert_eq!(bulk.aligned_schema_version, None); + assert_eq!( + bulk.partition_expr_version.map(|v| v.value), + versions[®ion.region_number()] + ); + assert_eq!( + owners, + tracked.then(|| vec![if region.region_number() == 2 { 0 } else { 2 }]) + ); + let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = bulk.body else { + panic!("expected Arrow IPC") + }; + let batch = FlightDecoder::try_from_schema_bytes(&ipc.schema) + .unwrap() + .try_decode_record_batch(&ipc.data_header, &ipc.payload) + .unwrap(); + assert_eq!(batch.num_rows(), 1); + let values = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + actual.push((region.region_number(), values.value(0))); + } + actual.sort_unstable(); + assert_eq!(actual, [(1, 70), (2, 20), (3, 1)]); - // A leader-only change retains the original IPC buffers and schema hint. - let request = build_bulk_request(3, versions[&3], vec![1]); - let Some(Body::BulkInsert(bulk)) = request.body.as_ref() else { - unreachable!() - }; - let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = &bulk.body else { - unreachable!() - }; - let pointers = ( - ipc.schema.as_ptr(), - ipc.data_header.as_ptr(), - ipc.payload.as_ptr(), - ); - let rerouted = rebuild_with_cached_table(&manager, &cache, request) - .await - .unwrap(); - assert_eq!(rerouted.len(), 1); - let Some(Body::BulkInsert(bulk)) = rerouted[0].body.as_ref() else { - unreachable!() - }; - assert_eq!( - bulk.aligned_schema_version, - Some(AlignedSchemaVersion { - schema_version: 123 - }) - ); - let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = &bulk.body else { - unreachable!() - }; - assert_eq!( - ( - ipc.schema.as_ptr(), - ipc.data_header.as_ptr(), - ipc.payload.as_ptr() - ), - pointers - ); + // Both unchanged routes and reroutes selecting all rows reuse the original IPC. + for (region, version, unchanged) in [(99, Some(0), false), (3, versions[&3], true)] { + let request = build_bulk_request(region, version, vec![1, 2, 3]); + let header = request.header.clone(); + let Some(Body::BulkInsert(bulk)) = request.body.as_ref() else { + unreachable!() + }; + let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = &bulk.body else { + unreachable!() + }; + let pointers = ( + ipc.schema.as_ptr(), + ipc.data_header.as_ptr(), + ipc.payload.as_ptr(), + ); + let rerouted = partitioner + .rebuild_tracked_write_requests( + request, + &table_info, + tracked.then(|| vec![2, 0, 2]), + ) + .unwrap(); + assert_eq!(rerouted.len(), 1); + let (request, owners) = &rerouted[0]; + assert_eq!(request.header, header); + assert_eq!(owners, &tracked.then(|| vec![2, 0, 2])); + let Some(Body::BulkInsert(bulk)) = request.body.as_ref() else { + unreachable!() + }; + assert_eq!(bulk.region_id, RegionId::new(1, 3).as_u64()); + assert!(bulk.skip_wal); + assert_eq!(bulk.partition_expr_version.map(|v| v.value), versions[&3]); + assert_eq!( + bulk.aligned_schema_version, + unchanged.then_some(AlignedSchemaVersion { + schema_version: 123 + }) + ); + let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = &bulk.body else { + unreachable!() + }; + assert_eq!( + ( + ipc.schema.as_ptr(), + ipc.data_header.as_ptr(), + ipc.payload.as_ptr() + ), + pointers, + "tracked={tracked}, unchanged={unchanged}" + ); + } + } } #[tokio::test] @@ -1199,4 +1370,83 @@ mod tests { .context(error::InvalidPartitionSnafu)?; partitioner.rebuild_write_requests(request, &table_info) } + + #[tokio::test] + async fn test_tracked_bulk_owners_preserve_payload_alignment() { + let (manager, cache) = create_fixture().await; + let table_info = cache.get(1).await.unwrap().unwrap(); + let partitioner = TablePartitioner::try_new(&manager, &table_info) + .await + .unwrap(); + for len in [0, 1, 7, 64, 65, 257] { + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])), + vec![Arc::new(Int32Array::from_iter_values(0..len as i32))], + ) + .unwrap(); + let owners = (0..len) + .map(|i| [2, 2, 0, 1, 0, 0][i % 6]) + .collect::>(); + for offset in [0, 1, 7, 65] { + for pattern in 0..3 { + let values = (0..offset + len) + .map(|i| match pattern { + 0 => false, + 1 => true, + _ => i % 3 != 0, + }) + .collect::>(); + let expected = (0..len).filter(|i| values[offset + i]).collect::>(); + let mask = BooleanArray::from(values).slice(offset, len); + let requests = partitioner + .build_tracked_bulk_insert_requests( + batch.clone(), + [(1, RegionMask::from(mask))].into_iter().collect(), + owners.clone(), + None, + false, + ) + .unwrap(); + if expected.is_empty() { + assert!(requests.is_empty()); + continue; + } + assert_eq!(requests.len(), 1); + let (request, selected_owners) = &requests[0]; + let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = &request.body else { + panic!("expected Arrow IPC") + }; + let payload = FlightDecoder::try_from_schema_bytes(&ipc.schema) + .unwrap() + .try_decode_record_batch(&ipc.data_header, &ipc.payload) + .unwrap(); + let actual_rows = payload + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + actual_rows + .values() + .iter() + .map(|x| *x as usize) + .collect::>(), + expected + ); + assert_eq!( + *selected_owners, + expected.iter().map(|i| owners[*i]).collect::>() + ); + } + } + let wrong = partitioner.build_tracked_bulk_insert_requests( + batch, + HashMap::new(), + vec![0; len + 1], + None, + false, + ); + assert!(wrong.is_err()); + } + } } diff --git a/src/operator/src/write_dispatcher.rs b/src/operator/src/write_dispatcher.rs index b28eff9d550..d3b6334814a 100644 --- a/src/operator/src/write_dispatcher.rs +++ b/src/operator/src/write_dispatcher.rs @@ -14,7 +14,7 @@ mod attribution; mod options; -pub use attribution::WriteAttribution; +mod progress; #[cfg(test)] mod tests; @@ -43,7 +43,9 @@ use store_api::storage::RegionId; use table::table_name::TableName; use tokio_util::sync::CancellationToken; +pub use self::attribution::WriteAttribution; pub use self::options::WriteRetryOptions; +pub use self::progress::TrackedWrite; use crate::error::{self, FindRegionLeaderSnafu, RequestRegionSnafu}; use crate::req_convert::TablePartitioner; @@ -226,14 +228,16 @@ impl WriteDispatcher { let progress = progress.as_ref(); let attempts = future::join_all(writes.into_iter().map(|(peer, request)| async move { if self.retry.is_some() || progress.is_some() { - self.send_batch(peer, request, histogram, progress).await + self.send_batch(peer, request, Vec::new(), histogram, progress) + .await } else { - vec![self.send(peer, request, histogram).await] + vec![(Vec::new(), self.send(peer, request, histogram).await)] } })) .await .into_iter() .flatten() + .map(|(_, attempt)| attempt) .collect(); self.finish(attempts, histogram, cancellation).await } @@ -266,6 +270,7 @@ impl WriteDispatcher { peer, request: Arc::new(request), attribution, + owners: None, }, histogram, ) @@ -318,15 +323,20 @@ impl WriteDispatcher { resolve_peers(&self.partition_manager, regions).await } + /// `owners` has one vector per request entry for tracked sends, or is empty otherwise. async fn send_batch( &self, peer: PeerRef, mut request: RegionRequest, + owners: Vec>, histogram: Option<&Histogram>, progress: Option<&Arc>, - ) -> Vec { + ) -> Vec<(Vec, WriteAttempt)> { if !matches!(&request.body, Some(Body::Inserts(_) | Body::Deletes(_))) { - return vec![self.send(peer, request, histogram).await]; + return vec![( + owners.into_iter().flatten().collect(), + self.send(peer, request, histogram).await, + )]; } // Upgrade datanodes first and roll back frontends first: old datanodes // cannot report which entries succeeded, so retrying a batch may replay writes. @@ -363,7 +373,10 @@ impl WriteDispatcher { error, request: None, }], - }; + } + .into_iter() + .map(|attempt| (Vec::new(), attempt)) + .collect(); } let request = Arc::new(request); let requester = self.node_manager.datanode(&peer).await; @@ -376,44 +389,64 @@ impl WriteDispatcher { if let Some(progress) = progress { progress.add_written(response.affected_rows); } - return vec![WriteAttempt::Success(response)]; + return vec![( + owners.into_iter().flatten().collect(), + WriteAttempt::Success(response), + )]; } Err(error) => error, }; let request = match Arc::try_unwrap(request) { Ok(request) => request, Err(_) => { - return vec![WriteAttempt::Failed { - error: UnexpectedSnafu { - err_msg: "Region request is still shared after its RPC completed", - } - .build(), - request: None, - }]; + return vec![( + owners.into_iter().flatten().collect(), + WriteAttempt::Failed { + error: UnexpectedSnafu { + err_msg: "Region request is still shared after its RPC completed", + } + .build(), + request: None, + }, + )]; } }; let requests = split_request(request); + let mut owners = owners.into_iter(); match error { Error::PartialRegion { response } => { if let Some(progress) = progress { progress.add_written(response.affected_rows); } let mut failures: HashMap<_, _> = response.failures.into_iter().collect(); - let mut attempts = vec![WriteAttempt::Success(RegionResponse { - affected_rows: response.affected_rows, - extensions: response.extensions, - metadata: Vec::new(), - })]; + let mut attempts = Vec::new(); + let mut success_owners = Vec::new(); for (index, request) in requests.into_iter().enumerate() { + let owners = owners.next(); if let Some(error) = failures.remove(&index) { + let failed_owners = owners.clone().unwrap_or_default(); let request = error.retry_hint().is_retryable().then(|| RoutedWrite { peer: Arc::clone(&peer), attribution: direct_attribution(&request, progress.cloned()), request: Arc::new(request), + owners, }); - attempts.push(WriteAttempt::Failed { error, request }); + attempts.push((failed_owners, WriteAttempt::Failed { error, request })); + } else { + success_owners.extend(owners.unwrap_or_default()); } } + attempts.insert( + 0, + ( + success_owners, + WriteAttempt::Success(RegionResponse { + affected_rows: response.affected_rows, + extensions: response.extensions, + metadata: Vec::new(), + }), + ), + ); attempts } error => { @@ -422,15 +455,21 @@ impl WriteDispatcher { requests .into_iter() .map(|request| { + let owners = owners.next(); + let failed_owners = owners.clone().unwrap_or_default(); let request = error.retry_hint().is_retryable().then(|| RoutedWrite { peer: Arc::clone(&peer), attribution: direct_attribution(&request, progress.cloned()), request: Arc::new(request), + owners, }); - WriteAttempt::Failed { - error: BatchRequestSnafu.into_error(Arc::clone(&error)), - request, - } + ( + failed_owners, + WriteAttempt::Failed { + error: BatchRequestSnafu.into_error(Arc::clone(&error)), + request, + }, + ) }) .collect() } @@ -465,6 +504,7 @@ pub(crate) struct RoutedWrite { peer: PeerRef, request: Arc, attribution: WriteAttribution, + owners: Option>, } impl WriteRetry { @@ -645,13 +685,30 @@ impl WriteRetry { partitioners.get(&table_id).context(UnexpectedSnafu { err_msg: "Refreshed write has no partition snapshot", })?; - let requests = partitioner - .rebuild_attributed_write(request, table_info, write.attribution) - .map_err(BoxedError::new) - .context(ExternalSnafu)?; - rebuilt_requests.extend(requests); + if let Some(owners) = write.owners { + let requests = partitioner + .rebuild_tracked_write_requests(request, table_info, Some(owners)) + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + rebuilt_requests.extend( + requests + .into_iter() + .map(|(request, owners)| ((request, WriteAttribution::default()), owners)), + ); + } else { + let requests = partitioner + .rebuild_attributed_write(request, table_info, write.attribution) + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + rebuilt_requests.extend(requests.into_iter().map(|request| (request, None))); + } } - self.prepare(rebuilt_requests).await + let (requests, owners): (Vec<_>, Vec<_>) = rebuilt_requests.into_iter().unzip(); + let mut writes = self.prepare(requests).await?; + for (write, owners) in writes.iter_mut().zip(owners) { + write.owners = owners; + } + Ok(writes) } async fn prepare( @@ -674,6 +731,7 @@ impl WriteRetry { peer: get_peer(&peers, region)?, request: Arc::new(request), attribution, + owners: None, }) }) .collect() diff --git a/src/operator/src/write_dispatcher/progress.rs b/src/operator/src/write_dispatcher/progress.rs new file mode 100644 index 00000000000..76d49bc2311 --- /dev/null +++ b/src/operator/src/write_dispatcher/progress.rs @@ -0,0 +1,333 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; + +use api::v1::region::{RegionRequest, region_request::Body}; +use backon::{BackoffBuilder, ExponentialBuilder}; +use common_error::ext::{BoxedError, ErrorExt, RetryHint}; +use common_meta::error::{Error, UnexpectedSnafu}; +use common_meta::peer::PeerRef; +use common_query::WriteProgress; +use futures::future; +use prometheus::Histogram; +use tokio_util::sync::CancellationToken; + +use crate::write_dispatcher::{ + RoutedWrite, WriteAttempt, WriteAttribution, WriteDispatcher, get_peer, write_region_id, +}; + +/// A region payload and its original-input indices, in payload row order. +/// Ownership stays local; it is never added to the schema or sent to storage. +pub struct TrackedWrite { + pub request: RegionRequest, + pub owners: Vec, +} + +/// A wire request retaining one row-owner vector per region entry. +struct TrackedBatch { + request: RegionRequest, + owners: Vec>, +} + +impl WriteDispatcher { + /// Collects acknowledgements per input across shared flushes and repartition retries. + /// Records each datanode send attempt, including retries, in the optional histogram. + pub async fn dispatch_tracked( + &self, + requests: Vec, + inputs: usize, + histogram: Option<&Histogram>, + ) -> Vec { + let (writes, progress) = self.prepare_tracked(requests, inputs).await; + self.execute_tracked(writes, progress, histogram, None) + .await + } + + /// Resolves routes before tracked writes are detached from the caller. + pub(crate) async fn prepare_tracked( + &self, + requests: Vec, + inputs: usize, + ) -> (Vec, Vec) { + let mut progress = vec![WriteProgress::default(); inputs]; + if requests + .iter() + .any(|request| request.owners.iter().any(|owner| *owner >= inputs)) + { + let error = UnexpectedSnafu { + err_msg: "Write owner exceeds original input count", + } + .build(); + return ( + Vec::new(), + vec![WriteProgress::failed(BoxedError::new(error)); inputs], + ); + } + let regions = match requests + .iter() + .map(|write| write_region_id(&write.request)) + .collect::>>() + { + Ok(regions) => regions, + Err(error) => { + return ( + Vec::new(), + vec![WriteProgress::failed(BoxedError::new(error)); inputs], + ); + } + }; + let peers = match self.resolve_peers(®ions).await { + Ok(peers) => peers, + Err(error) => { + return ( + Vec::new(), + vec![WriteProgress::failed(BoxedError::new(error)); inputs], + ); + } + }; + let mut pending = Vec::with_capacity(requests.len()); + for (write, region) in requests.into_iter().zip(regions) { + match get_peer(&peers, region) { + Ok(peer) => pending.push(RoutedWrite { + peer, + request: Arc::new(write.request), + attribution: WriteAttribution::default(), + owners: Some(write.owners), + }), + Err(error) => record_failure(&mut progress, &write.owners, error), + } + } + (pending, progress) + } + + /// Drains in-flight writes while cancellation stops subsequent retries. + pub(crate) async fn execute_tracked( + &self, + mut pending: Vec, + mut progress: Vec, + histogram: Option<&Histogram>, + cancellation: Option<&CancellationToken>, + ) -> Vec { + let mut backoff = self.retry.as_ref().map(|retry| { + ExponentialBuilder::default() + .with_min_delay(retry.options.initial_backoff) + .with_max_delay(retry.options.max_backoff) + .with_factor(retry.options.backoff_multiplier) + .with_max_times(retry.options.max_retries) + .with_jitter() + .build() + }); + let mut failures = Vec::new(); + loop { + failures.clear(); + // Partial responses can be attributed directly when all rows belong to one input. + let attempts = if progress.len() == 1 + && self.retry.is_some() + && pending.len() > 1 + && pending + .iter() + .any(|write| matches!(write.request.body, Some(Body::Inserts(_)))) + { + self.send_tracked_batches(pending, histogram).await + } else { + future::join_all( + pending + .into_iter() + .map(|write| self.send_tracked_write(write, histogram)), + ) + .await + }; + pending = Vec::new(); + let mut refresh = Vec::new(); + let mut terminal = false; + for (owners, attempt) in attempts { + match attempt { + WriteAttempt::Success(response) if response.affected_rows == owners.len() => { + for run in owners.chunk_by(|a, b| a == b) { + progress[run[0]].affected_rows += run.len(); + } + } + WriteAttempt::Success(response) => { + terminal = true; + record_failure(&mut progress, &owners, UnexpectedSnafu { + err_msg: format!("Write acknowledged {} rows, expected {}; cannot attribute this acknowledgement", response.affected_rows, owners.len()), + }.build()); + } + WriteAttempt::Failed { + error, + request: Some(write), + } => { + if error.retry_hint() == RetryHint::RetryAfterRouteRefresh { + refresh.push(write); + } else { + pending.push(write); + } + failures.push((owners, error)); + } + WriteAttempt::Failed { + error, + request: None, + } => { + terminal = true; + record_failure(&mut progress, &owners, error); + } + } + } + if failures.is_empty() { + break; + } + let delay = (!terminal) + .then(|| backoff.as_mut().and_then(Iterator::next)) + .flatten(); + let (Some(delay), Some(retry)) = (delay, &self.retry) else { + break; + }; + if let Some(cancellation) = cancellation { + tokio::select! { + biased; + _ = cancellation.cancelled() => break, + _ = tokio::time::sleep(delay) => {} + } + } else { + tokio::time::sleep(delay).await; + } + if cancellation.is_some_and(CancellationToken::is_cancelled) { + break; + } + if !refresh.is_empty() { + let owners = refresh + .iter() + .flat_map(|write| write.owners.iter().flatten().copied()) + .collect::>(); + match retry.reroute_failed_writes(refresh).await { + Ok(writes) => pending.extend(writes), + Err(error) => record_failure(&mut progress, &owners, error), + } + } + if pending.is_empty() || cancellation.is_some_and(CancellationToken::is_cancelled) { + break; + } + } + for (owners, error) in failures { + record_failure(&mut progress, &owners, error); + } + progress + } + + async fn send_tracked_write( + &self, + mut write: RoutedWrite, + histogram: Option<&Histogram>, + ) -> (Vec, WriteAttempt) { + let owners = write.owners.take().unwrap_or_default(); + let mut attempt = if let Some(retry) = &self.retry { + retry.send(write, histogram).await + } else { + let _timer = histogram.map(Histogram::start_timer); + let datanode = self.node_manager.datanode(&write.peer).await; + match datanode.handle(write.request).await { + Ok(response) => WriteAttempt::Success(response), + Err(error) => WriteAttempt::Failed { + error, + request: None, + }, + } + }; + if let WriteAttempt::Failed { + request: Some(write), + .. + } = &mut attempt + { + write.owners = Some(owners.clone()); + } + (owners, attempt) + } + + async fn send_tracked_batches( + &self, + writes: Vec, + histogram: Option<&Histogram>, + ) -> Vec<(Vec, WriteAttempt)> { + let mut batches: HashMap> = HashMap::new(); + let mut singles = Vec::new(); + let mut attempts = Vec::new(); + for write in writes { + if !matches!(write.request.body, Some(Body::Inserts(_))) { + singles.push(write); + continue; + } + let owners = write.owners.unwrap_or_default(); + let mut request = match Arc::try_unwrap(write.request) { + Ok(request) => request, + Err(_) => { + attempts.push(( + owners, + WriteAttempt::Failed { + error: UnexpectedSnafu { + err_msg: "Region request is still shared before batching", + } + .build(), + request: None, + }, + )); + continue; + } + }; + let batches = batches.entry(write.peer).or_default(); + if let Some(batch) = batches + .iter_mut() + .find(|batch| batch.request.header == request.header) + && let (Some(Body::Inserts(existing)), Some(Body::Inserts(next))) = + (&mut batch.request.body, &mut request.body) + { + existing.requests.append(&mut next.requests); + batch.owners.push(owners); + } else { + batches.push(TrackedBatch { + request, + owners: vec![owners], + }); + } + } + let batches = future::join_all(batches.into_iter().flat_map(|(peer, batches)| { + batches.into_iter().map(move |batch| { + self.send_batch( + Arc::clone(&peer), + batch.request, + batch.owners, + histogram, + None, + ) + }) + })); + let singles = future::join_all( + singles + .into_iter() + .map(|write| self.send_tracked_write(write, histogram)), + ); + let (batches, singles) = future::join(batches, singles).await; + attempts.extend(batches.into_iter().flatten()); + attempts.extend(singles); + attempts + } +} + +fn record_failure(progress: &mut [WriteProgress], owners: &[usize], error: Error) { + let error = Arc::new(BoxedError::new(error)); + for owner in owners.iter().copied().collect::>() { + progress[owner].error.get_or_insert_with(|| error.clone()); + } +} diff --git a/src/operator/src/write_dispatcher/tests.rs b/src/operator/src/write_dispatcher/tests.rs index 52954355cc6..cb66bfd1dbc 100644 --- a/src/operator/src/write_dispatcher/tests.rs +++ b/src/operator/src/write_dispatcher/tests.rs @@ -79,6 +79,7 @@ use crate::test_util; #[derive(Clone, Default)] struct RecordingRegionService { request_sizes: Arc>>, + response_rows: Arc>>, rpc_rejections: Arc>>, requests: Arc>>, rejections: Arc>>>, @@ -147,6 +148,9 @@ impl Region for RecordingRegionService { .unwrap() .status_code = StatusCode::PartialFailure as u32; } + if let Some(rows) = *self.response_rows.lock().unwrap() { + response.affected_rows = rows as _; + } Ok(Response::new(response)) } } @@ -157,6 +161,13 @@ impl RecordingRegionService { request: RegionRequest, ) -> std::result::Result, Status> { let region_id = write_region_id(&request).unwrap().as_u64(); + let affected_rows = match &request.body { + Some(Body::Inserts(inserts)) => inserts.requests[0] + .rows + .as_ref() + .map_or(1, |rows| rows.rows.len()), + _ => 1, + }; self.requests.lock().unwrap().push(request); let release = self.release.lock().unwrap().clone(); if let Some(release) = release { @@ -178,7 +189,7 @@ impl RecordingRegionService { ..Default::default() }), }), - affected_rows: 1, + affected_rows: affected_rows as _, ..Default::default() })) } @@ -190,6 +201,7 @@ struct RecordingCacheInvalidator { inner: Option, route_update: Mutex>, rename_and_recreate: Mutex>, + cancellation: Option, } #[async_trait::async_trait] @@ -240,7 +252,11 @@ impl CacheInvalidator for RecordingCacheInvalidator { .await?; } if let Some(inner) = &self.inner { - inner.invalidate(ctx, caches).await + inner.invalidate(ctx, caches).await?; + if let Some(cancellation) = &self.cancellation { + cancellation.cancel(); + } + Ok(()) } else { UnexpectedSnafu { err_msg: "stop after recording route refresh", @@ -684,6 +700,7 @@ async fn test_should_retry_requires_a_retained_request_and_budget() { attribution: Default::default(), peer: Arc::new(Peer::new(1, "peer")), request: Arc::new(build_insert_request(&[1])), + owners: None, })) }; for (name, attempts, retries, expected) in [ @@ -745,18 +762,46 @@ async fn test_write_retry_cancellation_stops_backoff() { #[tokio::test] async fn test_write_retry_cancellation_drains_inflight_requests() { + for tracked in [false, true] { + check_write_retry_cancellation_drains_inflight_requests(tracked).await; + } +} + +async fn check_write_retry_cancellation_drains_inflight_requests(tracked: bool) { let server = RegionTestServer::start().await; server.fail_region(2, [RetryHint::RetryAfterRouteRefresh]); let release = Arc::new(Semaphore::new(0)); *server.service.release.lock().unwrap() = Some(release.clone()); let dispatcher = server.create_dispatcher(2).await; - let writes = dispatcher - .prepare(vec![build_insert_request(&[1, 2])]) - .await - .unwrap(); let cancellation = CancellationToken::new(); let guard = cancellation.clone().drop_guard(); - let execute = dispatcher.execute(writes, None, Some(&cancellation)); + let execute = async { + if tracked { + let requests = split_request(build_insert_request(&[1, 2])) + .into_iter() + .map(|request| TrackedWrite { + request, + owners: vec![0], + }) + .collect(); + let (writes, progress) = dispatcher.prepare_tracked(requests, 1).await; + let mut progress = dispatcher + .execute_tracked(writes, progress, None, Some(&cancellation)) + .await; + assert_eq!(progress[0].affected_rows, 1); + progress.pop().unwrap().into_result() + } else { + let writes = dispatcher + .prepare(vec![build_insert_request(&[1, 2])]) + .await + .unwrap(); + dispatcher + .execute(writes, None, Some(&cancellation)) + .await + .map(|response| response.affected_rows) + .map_err(BoxedError::new) + } + }; tokio::pin!(execute); tokio::time::timeout( Duration::from_secs(1), @@ -791,9 +836,74 @@ async fn test_write_retry_cancellation_drains_inflight_requests() { (RegionId::new(1, 2).as_u64(), 1), ]) ); + assert_eq!(*server.service.request_sizes.lock().unwrap(), [2]); assert!(server.invalidator.calls.lock().unwrap().is_empty()); } +#[tokio::test] +async fn test_tracked_cancellation_during_route_refresh_stops_retries() { + for inputs in [1, 2] { + check_tracked_cancellation_during_route_refresh_stops_retries(inputs).await; + } +} + +async fn check_tracked_cancellation_during_route_refresh_stops_retries(inputs: usize) { + let server = RegionTestServer::start().await; + server.fail_region(2, [RetryHint::RetryAfterRouteRefresh]); + let backend = test_util::prepare_mocked_backend().await; + test_util::create_partition_rule_manager(backend.clone()).await; + update_table_leader(&backend, &server.peer).await; + let (partition_manager, _, registry, catalog_manager) = create_routing_caches(backend); + let cancellation = CancellationToken::new(); + let invalidator = Arc::new(RecordingCacheInvalidator { + inner: Some(registry), + cancellation: Some(cancellation.clone()), + ..Default::default() + }); + let dispatcher = WriteDispatcher::new(Arc::new(NodeClients::default()), partition_manager) + .with_retry(build_retry_options(2), catalog_manager, invalidator.clone()); + let (writes, progress) = dispatcher + .prepare_tracked( + vec![ + TrackedWrite { + request: build_partitioned_request(WriteKind::Insert, 1, None, &[70]), + owners: vec![0], + }, + TrackedWrite { + request: build_partitioned_request(WriteKind::Insert, 2, None, &[20]), + owners: vec![inputs - 1], + }, + ], + inputs, + ) + .await; + let progress = dispatcher + .execute_tracked(writes, progress, None, Some(&cancellation)) + .await; + assert!(cancellation.is_cancelled()); + assert_eq!(invalidator.calls.lock().unwrap().len(), 2); + assert_eq!(progress[0].affected_rows, 1); + assert_eq!(progress[0].error.is_none(), inputs == 2); + if inputs == 2 { + assert_eq!(progress[1].affected_rows, 0); + } + assert_eq!( + progress[inputs - 1].error.as_ref().unwrap().retry_hint(), + RetryHint::RetryAfterRouteRefresh + ); + assert_eq!( + *server.service.request_sizes.lock().unwrap(), + if inputs == 1 { vec![2] } else { vec![1, 1] } + ); + assert_eq!( + server.attempt_counts(), + HashMap::from([ + (RegionId::new(1, 1).as_u64(), 1), + (RegionId::new(1, 2).as_u64(), 1), + ]) + ); +} + #[tokio::test] async fn test_write_retry_refreshes_routes_once_per_failed_round() { let server = RegionTestServer::start().await; @@ -1091,6 +1201,162 @@ fn request_parts(request: &RegionRequest) -> (u32, Option, Vec) { #[tokio::test] async fn test_write_retry_rebuilds_only_failed_rows_after_route_refresh() { + for (kind, tracked) in [ + (WriteKind::Insert, false), + (WriteKind::Insert, true), + (WriteKind::Delete, false), + (WriteKind::BulkInsert, false), + (WriteKind::BulkInsert, true), + ] { + let old_server = RegionTestServer::start().await; + let new_server = RegionTestServer::start().await; + old_server.fail_region(2, [RetryHint::RetryAfterRouteRefresh]); + let backend = test_util::prepare_mocked_backend().await; + test_util::create_partition_rule_manager(backend.clone()).await; + update_table_leader(&backend, &old_server.peer).await; + let (partition_manager, table_info_cache, registry, catalog_manager) = + create_routing_caches(backend.clone()); + let table_info = table_info_cache.get(1).await.unwrap().unwrap(); + let (_, old_versions) = partition_manager + .find_table_partition_rule(&table_info) + .await + .unwrap(); + let manager = TableRouteManager::new(backend.clone()); + let current = manager + .table_route_storage() + .get_with_raw_bytes(1) + .await + .unwrap() + .unwrap(); + let old_routes = current.region_routes().unwrap(); + let replaced = old_routes + .iter() + .find(|route| route.region.id.region_number() == 2) + .unwrap(); + let mut routes = old_routes + .iter() + .filter(|route| route.region.id.region_number() != 2) + .cloned() + .collect::>(); + for (region, lower, upper) in [(4, 10, 30), (5, 30, 50)] { + let mut route = replaced.clone(); + route.region.id = RegionId::new(1, region); + route.leader_peer = Some(new_server.peer.clone()); + route.region.partition_expr = PartitionExpr::new( + Operand::Expr(PartitionExpr::new( + Operand::Column("a".to_string()), + RestrictedOp::GtEq, + Operand::Value(datatypes::value::Value::Int32(lower)), + )), + RestrictedOp::And, + Operand::Expr(PartitionExpr::new( + Operand::Column("a".to_string()), + RestrictedOp::Lt, + Operand::Value(datatypes::value::Value::Int32(upper)), + )), + ) + .as_json_str() + .unwrap(); + routes.push(route); + } + let invalidator = Arc::new(RecordingCacheInvalidator { + calls: Mutex::default(), + inner: Some(registry), + route_update: Mutex::new(Some((backend, current.update(routes).unwrap()))), + ..Default::default() + }); + let clients = Arc::new(NodeClients::new(ChannelConfig::default())); + let dispatcher = WriteDispatcher::new(clients, partition_manager.clone()).with_retry( + build_retry_options(2), + catalog_manager, + invalidator.clone(), + ); + if tracked { + let inputs = if matches!(kind, WriteKind::Insert) { + 1 + } else { + 2 + }; + new_server.fail_region(5, [RetryHint::NonRetryable]); + let progress = dispatcher + .dispatch_tracked( + vec![ + TrackedWrite { + request: build_partitioned_request(kind, 1, old_versions[&1], &[70]), + owners: vec![0], + }, + TrackedWrite { + request: build_partitioned_request( + kind, + 2, + old_versions[&2], + &[20, 40], + ), + owners: vec![0, inputs - 1], + }, + ], + inputs, + None, + ) + .await; + assert_eq!(progress[0].affected_rows, 2); + assert_eq!(progress[0].error.is_none(), inputs == 2); + if inputs == 2 { + assert_eq!(progress[1].affected_rows, 0); + assert!(progress[1].error.is_some()); + } else { + assert_eq!(*old_server.service.request_sizes.lock().unwrap(), [2]); + assert_eq!(*new_server.service.request_sizes.lock().unwrap(), [2]); + } + } else { + let response = dispatcher + .dispatch(vec![ + build_partitioned_request(kind, 1, old_versions[&1], &[70]), + build_partitioned_request(kind, 2, old_versions[&2], &[20, 40]), + ]) + .await + .unwrap(); + assert_eq!(response.affected_rows, 3, "{kind:?}"); + } + assert_eq!( + old_server.attempt_counts(), + HashMap::from([ + (RegionId::new(1, 1).as_u64(), 1), + (RegionId::new(1, 2).as_u64(), 1), + ]), + "{kind:?}" + ); + let (_, versions) = partition_manager + .find_table_partition_rule(&table_info) + .await + .unwrap(); + let mut writes = new_server + .service + .requests + .lock() + .unwrap() + .iter() + .map(request_parts) + .collect::>(); + writes.sort_by_key(|write| write.0); + assert_eq!( + writes, + vec![(4, versions[&4], vec![20]), (5, versions[&5], vec![40])], + "{kind:?}" + ); + assert_eq!( + *invalidator.calls.lock().unwrap(), + vec![ + vec![CacheIdent::TableId(1)], + vec![CacheIdent::TableName(TableName::new( + "greptime", "public", "table_1" + ))], + ] + ); + } +} +#[tokio::test] +async fn test_write_retry_rebuilds_only_failed_rows_after_route_refresh_attributed() { for (kind, fail_repartitioned) in [ (WriteKind::Insert, false), (WriteKind::Delete, false), @@ -1979,3 +2245,362 @@ async fn test_with_retry_can_disable_previously_enabled_retries() { ); assert_eq!(server.attempt_counts()[&RegionId::new(1, 1).as_u64()], 1); } + +#[tokio::test] +async fn test_tracked_row_writes_preserve_peer_batching() { + for max_retries in [0, 2] { + for hint in [ + None, + Some(RetryHint::Retryable), + Some(RetryHint::NonRetryable), + ] { + let server = RegionTestServer::start().await; + if let Some(hint) = hint { + server.fail_region(1, [hint]); + } + let dispatcher = server.create_dispatcher(max_retries).await; + let histogram = Histogram::with_opts(prometheus::HistogramOpts::new( + "tracked_batch_attempts", + "Tracked batch send attempts", + )) + .unwrap(); + // Repeated region IDs must still be distinguished by their entry indices. + let writes = [(1, 2), (1, 2), (2, 3)] + .into_iter() + .enumerate() + .map(|(index, (region, rows))| TrackedWrite { + request: build_partitioned_request( + WriteKind::Insert, + region, + None, + &vec![20 + index as i32; rows], + ), + owners: vec![0; rows], + }) + .collect(); + let progress = dispatcher + .dispatch_tracked(writes, 1, Some(&histogram)) + .await; + let retried = max_retries > 0 && hint == Some(RetryHint::Retryable); + let failed = hint.is_some() && !retried; + assert_eq!(progress[0].affected_rows, 7 - 2 * usize::from(failed)); + assert_eq!(progress[0].error.is_some(), failed); + assert_eq!( + server.attempt_counts(), + HashMap::from([ + (RegionId::new(1, 1).as_u64(), 2 + usize::from(retried)), + (RegionId::new(1, 2).as_u64(), 1), + ]) + ); + let mut expected_sizes = if max_retries == 0 { + vec![1, 1, 1] + } else { + vec![3] + }; + if retried { + expected_sizes.push(1); + let values = server + .service + .requests + .lock() + .unwrap() + .iter() + .map(|request| request_parts(request).2[0]) + .collect::>(); + assert_eq!(values, [20, 21, 22, 20]); + } + assert_eq!(histogram.get_sample_count(), expected_sizes.len() as u64); + assert_eq!( + *server.service.request_sizes.lock().unwrap(), + expected_sizes + ); + } + } +} + +#[tokio::test] +async fn test_tracked_batches_preserve_peers_and_headers() { + let first = RegionTestServer::start().await; + let second = RegionTestServer::start().await; + let dispatcher = first.create_dispatcher(2).await; + let requests = [1, 1, 1, 2, 2] + .into_iter() + .enumerate() + .map(|(index, region)| { + let mut request = build_partitioned_request(WriteKind::Insert, region, None, &[20]); + if index == 2 { + request.header = Some(RegionRequestHeader { + dbname: "greptime.other".to_string(), + ..Default::default() + }); + } + TrackedWrite { + request, + owners: vec![0], + } + }) + .collect(); + let (mut writes, progress) = dispatcher.prepare_tracked(requests, 1).await; + for write in &mut writes { + if write_region_id(&write.request).unwrap().region_number() == 2 { + write.peer = Arc::new(second.peer.clone()); + } + } + let progress = dispatcher + .execute_tracked(writes, progress, None, None) + .await; + assert_eq!(progress[0].affected_rows, 5); + assert!(progress[0].error.is_none()); + let mut sizes = first.service.request_sizes.lock().unwrap().clone(); + sizes.sort_unstable(); + assert_eq!(sizes, [1, 2]); + assert_eq!(*second.service.request_sizes.lock().unwrap(), [2]); + assert_eq!( + first + .service + .requests + .lock() + .unwrap() + .iter() + .filter(|request| { request.header.as_ref().unwrap().dbname == "greptime.other" }) + .count(), + 1 + ); +} + +#[tokio::test] +async fn test_tracked_batches_reject_mismatched_acknowledgements() { + for (partial, rows) in [(false, 0), (false, 1), (false, 3), (true, 0), (true, 2)] { + let server = RegionTestServer::start().await; + *server.service.response_rows.lock().unwrap() = Some(rows); + if partial { + server.fail_region(2, [RetryHint::Retryable]); + } + let dispatcher = server.create_dispatcher(2).await; + let requests = [1, 2] + .into_iter() + .map(|region| TrackedWrite { + request: build_partitioned_request(WriteKind::Insert, region, None, &[20]), + owners: vec![0], + }) + .collect(); + let progress = dispatcher.dispatch_tracked(requests, 1, None).await; + assert_eq!(progress[0].affected_rows, 0); + assert!( + progress[0] + .error + .as_ref() + .unwrap() + .to_string() + .contains("cannot attribute") + ); + assert_eq!(*server.service.request_sizes.lock().unwrap(), [2]); + } +} + +#[tokio::test] +async fn test_tracked_batches_retry_whole_rpc_failures() { + for failures in [1, 3] { + let server = RegionTestServer::start().await; + server + .service + .rpc_rejections + .lock() + .unwrap() + .extend((0..failures).map(|_| Status::unavailable("acknowledgement lost"))); + let dispatcher = server.create_dispatcher(2).await; + let requests = [1, 2] + .into_iter() + .map(|region| TrackedWrite { + request: build_partitioned_request(WriteKind::Insert, region, None, &[20]), + owners: vec![0], + }) + .collect(); + let progress = dispatcher.dispatch_tracked(requests, 1, None).await; + assert_eq!(progress[0].affected_rows, if failures == 1 { 2 } else { 0 }); + assert_eq!(progress[0].error.is_some(), failures == 3); + assert_eq!( + *server.service.request_sizes.lock().unwrap(), + vec![2; (failures + 1).min(3)] + ); + } +} + +#[tokio::test] +async fn test_tracked_batches_leave_other_write_kinds_separate() { + for kind in [WriteKind::BulkInsert, WriteKind::Delete] { + let server = RegionTestServer::start().await; + let dispatcher = server.create_dispatcher(2).await; + let histogram = Histogram::with_opts(prometheus::HistogramOpts::new( + "tracked_other_attempts", + "Tracked non-insert attempts", + )) + .unwrap(); + let requests = [1, 2] + .into_iter() + .map(|region| TrackedWrite { + request: build_partitioned_request(kind, region, None, &[20]), + owners: vec![0], + }) + .collect(); + let progress = dispatcher + .dispatch_tracked(requests, 1, Some(&histogram)) + .await; + assert_eq!(progress[0].affected_rows, 2); + assert!(progress[0].error.is_none()); + assert_eq!(histogram.get_sample_count(), 2); + assert_eq!(server.service.requests.lock().unwrap().len(), 2); + } +} + +#[tokio::test] +async fn test_tracked_dispatch_records_each_send_attempt() { + for (max_retries, failures, expected_attempts) in + [(0, 0, 1), (0, 1, 1), (2, 0, 1), (2, 1, 2), (2, 3, 3)] + { + let server = RegionTestServer::start().await; + server.fail_region(2, std::iter::repeat_n(RetryHint::Retryable, failures)); + let dispatcher = server.create_dispatcher(max_retries).await; + let histogram = Histogram::with_opts(prometheus::HistogramOpts::new( + "tracked_write_attempts", + "Tracked write attempt time", + )) + .unwrap(); + let progress = dispatcher + .dispatch_tracked( + vec![ + TrackedWrite { + request: build_insert_request(&[1]), + owners: vec![0], + }, + TrackedWrite { + request: build_partitioned_request( + WriteKind::Insert, + 2, + None, + &[20, 21, 22], + ), + owners: vec![1, 0, 1], + }, + ], + 2, + Some(&histogram), + ) + .await; + + assert_eq!(histogram.get_sample_count(), (1 + expected_attempts) as u64); + assert_eq!(server.attempt_counts()[&RegionId::new(1, 1).as_u64()], 1); + assert_eq!( + server.attempt_counts()[&RegionId::new(1, 2).as_u64()], + expected_attempts + ); + let failed = failures > max_retries; + assert_eq!(progress[0].affected_rows, 1 + usize::from(!failed)); + assert_eq!(progress[0].error.is_some(), failed); + assert_eq!(progress[1].affected_rows, 2 * usize::from(!failed)); + assert_eq!(progress[1].error.is_some(), failed); + } +} + +#[tokio::test] +async fn test_exemplar_progress_retains_success_and_unknown_acknowledgements() { + for retries in [0, 1] { + let server = RegionTestServer::start().await; + server.service.rejections.lock().unwrap().insert( + RegionId::new(1, 2).as_u64(), + (0..=retries) + .map(|_| Status::unavailable("acknowledgement lost")) + .collect(), + ); + let dispatcher = server.create_dispatcher(retries).await; + let progress = dispatcher + .dispatch_tracked( + vec![ + TrackedWrite { + request: build_insert_request(&[1]), + owners: vec![0], + }, + TrackedWrite { + request: build_insert_request(&[2]), + owners: vec![0], + }, + TrackedWrite { + request: build_insert_request(&[1]), + owners: vec![1], + }, + ], + 2, + None, + ) + .await; + assert_eq!(progress[0].affected_rows, 1); + assert!(progress[0].error.is_some()); + assert_eq!(progress[1].affected_rows, 1); + assert!(progress[1].error.is_none()); + assert_eq!(server.attempt_counts()[&RegionId::new(1, 1).as_u64()], 2); + assert_eq!( + server.attempt_counts()[&RegionId::new(1, 2).as_u64()], + retries + 1 + ); + } +} + +#[tokio::test] +async fn test_tracked_progress_counts_noncontiguous_owners() { + #[derive(Clone)] + struct CountRows; + + #[async_trait::async_trait] + impl MockDatanodeHandler for CountRows { + async fn handle( + &self, + _: &Peer, + request: RegionRequest, + ) -> Result { + let Some(Body::Inserts(inserts)) = request.body else { + panic!("expected row inserts") + }; + Ok(api::region::RegionResponse::new( + inserts + .requests + .into_iter() + .map(|r| r.rows.unwrap().rows.len()) + .sum(), + )) + } + + async fn handle_query( + &self, + _: &Peer, + _: QueryRequest, + ) -> Result { + unreachable!() + } + } + + let manager = + test_util::create_partition_rule_manager(test_util::prepare_mocked_backend().await).await; + let dispatcher = WriteDispatcher::new(Arc::new(MockDatanodeManager::new(CountRows)), manager); + for owners in [ + vec![], + vec![0], + vec![0; 1024], + vec![0, 0, 2, 2, 1, 0, 0], + (0..129).map(|i| i % 3).collect(), + ] { + let request = + build_partitioned_request(WriteKind::Insert, 1, None, &vec![70; owners.len()]); + let mut expected = [0; 4]; + for owner in &owners { + expected[*owner] += 1; + } + let progress = dispatcher + .dispatch_tracked(vec![TrackedWrite { request, owners }], 4, None) + .await; + assert!(progress.iter().all(|p| p.error.is_none())); + assert_eq!( + progress.iter().map(|p| p.affected_rows).collect::>(), + expected + ); + } +} diff --git a/src/servers/benches/flush_batch_physical.rs b/src/servers/benches/flush_batch_physical.rs index e4933a9b371..7ce73c5f8c6 100644 --- a/src/servers/benches/flush_batch_physical.rs +++ b/src/servers/benches/flush_batch_physical.rs @@ -21,9 +21,11 @@ use arrow::datatypes::{DataType as ArrowDataType, Field, Schema as ArrowSchema, use arrow::record_batch::RecordBatch; use async_trait::async_trait; use catalog::error::Result as CatalogResult; +use common_query::WriteProgress; use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema as DtColumnSchema, Schema as DtSchema}; +use operator::write_dispatcher::TrackedWrite; use partition::error::Result as PartitionResult; use partition::partition::{PartitionRule, PartitionRuleRef, RegionMask}; use servers::batcher::logical_table::{ @@ -101,11 +103,20 @@ struct BenchWriteDispatcher; #[async_trait] impl PhysicalFlushWriteDispatcher for BenchWriteDispatcher { - async fn dispatch( + async fn dispatch(&self, _requests: Vec) -> error::Result { + unreachable!() + } + + async fn dispatch_tracked( &self, - _requests: Vec<(RegionRequest, operator::write_dispatcher::WriteAttribution)>, - ) -> error::Result { - Ok(0) + requests: Vec, + inputs: usize, + ) -> Vec { + let mut progress = vec![WriteProgress::default(); inputs]; + for owner in requests.into_iter().flat_map(|write| write.owners) { + progress[owner].affected_rows += 1; + } + progress } } @@ -232,7 +243,7 @@ fn bench_flush_batch_physical(c: &mut Criterion) { |b, batches| { b.iter(|| { rt.block_on(async { - flush_batch_physical( + let affected_rows = flush_batch_physical( batches, "phy", &ctx, @@ -242,6 +253,7 @@ fn bench_flush_batch_physical(c: &mut Criterion) { ) .await .unwrap(); + assert_eq!(affected_rows, num_tables * rows_per_table); }); }); }, @@ -259,7 +271,7 @@ fn bench_flush_batch_physical(c: &mut Criterion) { |b, batches| { b.iter(|| { rt.block_on(async { - flush_batch_physical( + let affected_rows = flush_batch_physical( batches, "phy", &ctx, @@ -269,6 +281,7 @@ fn bench_flush_batch_physical(c: &mut Criterion) { ) .await .unwrap(); + assert_eq!(affected_rows, num_tables * rows_per_table); }); }); }, diff --git a/src/servers/src/batcher/logical_table.rs b/src/servers/src/batcher/logical_table.rs index 1391d85ee3f..26277c43406 100644 --- a/src/servers/src/batcher/logical_table.rs +++ b/src/servers/src/batcher/logical_table.rs @@ -33,9 +33,11 @@ use common_batcher::flush_limiter::FlushLimiter; use common_batcher::flush_policy::timing::TimingFlushPolicy; use common_batcher::request_limiter::RequestLimiter; use common_batcher::worker_registry::WorkerRegistry; +use common_error::ext::BoxedError; use common_meta::cache::TableFlownodeSetCacheRef; use common_meta::node_manager::NodeManagerRef; use common_query::prelude::GREPTIME_PHYSICAL_TABLE; +use common_query::{Output, WriteCompletion, WriteProgress}; use common_time::timestamp::TimeUnit; use meter_core::data::MeterRecord; use meter_macros::write_meter; @@ -270,6 +272,9 @@ impl LogicalTablePendingRowsBatcher { /// Submits with request-level accounting after schema preparation and before /// queue admission. Acknowledgement follows the global batching policy. + /// In synchronous mode, waits for storage even when confirmed writes are enabled. + /// Use [`Self::submit_output_with`] with confirmed writes to retain partial row + /// counts on failure. pub async fn submit_with( &self, requests: RowInsertRequests, @@ -279,6 +284,33 @@ impl LogicalTablePendingRowsBatcher { where F: Future> + Send, { + let (output, prepared) = self + .submit_output_with(requests, ctx, after_prepare) + .await?; + let rows = output.extract_rows_and_cost().0 as u64; + if self.pending_rows_batch_sync { + output + .write_progress() + .await + .into_result() + .context(error::ExecuteGrpcQuerySnafu)?; + } + Ok((rows, prepared)) + } + + /// Retains completion even when the batching policy acknowledges queue admission. + /// With confirmed writes, callers must await [`Output::write_progress`] to observe + /// storage errors and acknowledged rows, including in synchronous batching mode. + pub async fn submit_output_with( + &self, + requests: RowInsertRequests, + ctx: QueryContextRef, + after_prepare: impl FnOnce(RowInsertRequests) -> F, + ) -> Result<(Output, T)> + where + F: Future> + Send, + { + let confirmed = ctx.confirmed_writes(); let (table_batches, total_rows) = { let _timer = PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED .with_label_values(&["submit_build_and_align"]) @@ -287,7 +319,7 @@ impl LogicalTablePendingRowsBatcher { }; let prepared = after_prepare(requests).await?; if total_rows == 0 { - return Ok((0, prepared)); + return Ok((Output::new_with_affected_rows(0), prepared)); } // Flushes dispatch directly to datanodes, so admit once before enqueueing. @@ -358,24 +390,30 @@ impl LogicalTablePendingRowsBatcher { } } - if self.pending_rows_batch_sync { - let result = { - let _timer = PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED - .with_label_values(&["submit_wait_flush_result"]) - .start_timer(); - response_rx - .await - .map_err(|_| error::BatcherChannelClosedSnafu.build())? - }; - result - .context(error::SubmitBatchSnafu) - .map(|()| (total_rows as u64, prepared)) + let completion = async move { + let _timer = PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED + .with_label_values(&["submit_wait_flush_result"]) + .start_timer(); + response_rx.await.unwrap_or_else(|_| { + WriteProgress::failed(BoxedError::new(error::BatcherChannelClosedSnafu.build())) + }) + }; + let mut output = Output::new_with_affected_rows(total_rows); + if self.pending_rows_batch_sync && !confirmed { + completion + .await + .into_result() + .context(error::ExecuteGrpcQuerySnafu)?; } else { if let Some(progress) = progress { progress.mark_accepted(); } - Ok((total_rows as u64, prepared)) + output + .meta + .write_completions + .push(WriteCompletion::with_progress(completion)); } + Ok((output, prepared)) } } @@ -430,16 +468,229 @@ impl Drop for LogicalTablePendingRowsBatcher { #[cfg(test)] mod tests { + use std::future::ready; + use std::num::NonZeroUsize; use std::sync::Arc; + use std::time::Duration; + use api::v1::value::ValueData; + use api::v1::{ + ColumnDataType, ColumnSchema, Row, RowInsertRequest, RowInsertRequests, Rows, SemanticType, + Value, + }; use arrow::array::{StringArray, TimestampMillisecondArray}; use arrow::datatypes::{DataType as ArrowDataType, Field, Schema as ArrowSchema}; use arrow::record_batch::RecordBatch; + use catalog::memory::MemoryCatalogManager; + use common_error::ext::{BoxedError, ErrorExt}; + use common_error::status_code::StatusCode; + use common_meta::test_util::MockDatanodeManager; + use common_query::WriteProgress; use common_query::prelude::greptime_timestamp; + use datatypes::prelude::ConcreteDataType; + use datatypes::schema::{ColumnSchema as DtColumnSchema, Schema as DtSchema}; + use operator::test_util::{create_partition_rule_manager, prepare_mocked_backend}; + use operator::write_dispatcher::WriteDispatcher; + use session::context::{QueryContext, QueryContextRef}; + use table::dist_table::DistTable; + use table::test_util::table_info::test_table_info; use crate::batcher::logical_table::batch_convert::{RecordBatchWithTsIdx, TableBatch}; use crate::batcher::logical_table::flow_notifier::extract_timestamps; + use crate::batcher::logical_table::pending_worker::WorkerCommand; use crate::batcher::logical_table::test_util::{mock_aligned_tag_batch, mock_tag_batch}; + use crate::batcher::logical_table::{ + LogicalTablePendingRowsBatcher, PendingRowsSchemaAlterer, batch_key_from_ctx, + }; + use crate::batcher::test_util::mock_table_flownode_cache; + use crate::error; + + struct UnexpectedSchemaAlterer; + + #[async_trait::async_trait] + impl PendingRowsSchemaAlterer for UnexpectedSchemaAlterer { + async fn create_tables_if_missing_batch( + &self, + _: &str, + _: &str, + _: &[(&str, &[ColumnSchema])], + _: bool, + _: QueryContextRef, + ) -> error::Result<()> { + unreachable!() + } + + async fn add_missing_prom_tag_columns_batch( + &self, + _: &str, + _: &str, + _: &[(&str, &[String])], + _: QueryContextRef, + ) -> error::Result<()> { + unreachable!() + } + } + + async fn submission_batcher(sync: bool) -> Arc { + let schema = DtSchema::try_new(vec![ + DtColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + DtColumnSchema::new("value", ConcreteDataType::float64_datatype(), true), + ]) + .unwrap(); + let table = test_table_info(1, "cpu", "public", "greptime", Arc::new(schema)); + let catalog = MemoryCatalogManager::new_with_table(DistTable::table(Arc::new(table))); + let partitions = create_partition_rule_manager(prepare_mocked_backend().await).await; + let nodes = Arc::new(MockDatanodeManager::new(())); + let dispatcher = Arc::new(WriteDispatcher::new(nodes.clone(), partitions.clone())); + let mut batcher = LogicalTablePendingRowsBatcher::try_new( + partitions, + nodes, + dispatcher, + catalog, + mock_table_flownode_cache(1, vec![]).await, + true, + Arc::new(UnexpectedSchemaAlterer), + Duration::from_secs(1), + 2, + 1, + 1, + 1, + NonZeroUsize::new(1).unwrap(), + ) + .unwrap(); + Arc::get_mut(&mut batcher).unwrap().pending_rows_batch_sync = sync; + batcher + } + + fn submission_requests() -> RowInsertRequests { + RowInsertRequests { + inserts: vec![RowInsertRequest { + table_name: "cpu".to_string(), + rows: Some(Rows { + schema: vec![ColumnSchema { + column_name: "ts".to_string(), + datatype: ColumnDataType::TimestampMillisecond as i32, + semantic_type: SemanticType::Timestamp as i32, + ..Default::default() + }], + rows: [1000, 2000] + .into_iter() + .map(|ts| Row { + values: vec![Value { + value_data: Some(ValueData::TimestampMillisecondValue(ts)), + }], + }) + .collect(), + }), + }], + } + } + + fn flush_progress(affected_rows: usize) -> WriteProgress { + WriteProgress { + affected_rows, + error: (affected_rows < 2).then(|| { + Arc::new(BoxedError::new( + error::InvalidQuerySnafu { + reason: "storage write failed", + } + .build(), + )) + }), + } + } + + #[tokio::test] + async fn test_submit_with_preserves_sync_acknowledgement() { + for confirmed in [false, true] { + for sync in [false, true] { + for affected_rows in [0, 1, 2] { + let batcher = submission_batcher(sync).await; + let ctx = QueryContext::arc(); + let ctx = if confirmed { + Arc::new(ctx.with_confirmed_writes()) + } else { + ctx + }; + let (_, receiver) = batcher + .workers + .get_or_create(batch_key_from_ctx(&ctx), 1) + .await; + let mut receiver = receiver.unwrap(); + let submit = batcher.submit_with(submission_requests(), ctx, |_| ready(Ok(7))); + tokio::pin!(submit); + let mut early_result = None; + // Poll submission first to detect early acknowledgement without timing assumptions. + let command = tokio::select! { + biased; + result = &mut submit => { + assert!(result.is_ok(), "{result:?}"); + early_result = Some(result); + receiver.try_recv().unwrap() + } + command = receiver.recv() => command.unwrap(), + }; + assert_eq!(early_result.is_none(), sync); + let WorkerCommand::Submit { response_tx, .. } = command else { + panic!("expected submission") + }; + assert_eq!( + response_tx.send(flush_progress(affected_rows)).is_ok(), + sync + ); + let result = match early_result { + Some(result) => result, + None => submit.await, + }; + if sync && affected_rows < 2 { + assert_eq!( + result.unwrap_err().status_code(), + StatusCode::InvalidArguments + ); + } else { + assert_eq!(result.unwrap(), (2, 7)); + } + } + } + } + } + + #[tokio::test] + async fn test_submit_output_with_preserves_confirmed_partial_progress() { + for sync in [false, true] { + let batcher = submission_batcher(sync).await; + let ctx = Arc::new(QueryContext::arc().with_confirmed_writes()); + let (_, receiver) = batcher + .workers + .get_or_create(batch_key_from_ctx(&ctx), 1) + .await; + let mut receiver = receiver.unwrap(); + let (output, prepared) = batcher + .submit_output_with(submission_requests(), ctx, |_| ready(Ok(7))) + .await + .unwrap(); + assert_eq!(prepared, 7); + assert_eq!(output.meta.write_completions.len(), 1); + let WorkerCommand::Submit { response_tx, .. } = receiver.try_recv().unwrap() else { + panic!("expected submission") + }; + let progress = output.write_progress(); + tokio::pin!(progress); + assert!(futures::poll!(progress.as_mut()).is_pending()); + response_tx.send(flush_progress(1)).unwrap(); + let progress = progress.await; + assert_eq!(progress.affected_rows, 1); + assert_eq!( + progress.error.unwrap().status_code(), + StatusCode::InvalidArguments + ); + } + } #[test] fn test_extract_timestamps_appends_non_null_batches_in_order() { diff --git a/src/servers/src/batcher/logical_table/batch.rs b/src/servers/src/batcher/logical_table/batch.rs index b6200b1d7c9..fc07ef5ca41 100644 --- a/src/servers/src/batcher/logical_table/batch.rs +++ b/src/servers/src/batcher/logical_table/batch.rs @@ -13,17 +13,18 @@ // limitations under the License. use std::collections::HashSet; -use std::sync::Arc; use std::time::Instant; use catalog::CatalogManagerRef; use common_batcher::flush_limiter::FlushLimiter; -use common_query::prelude; +use common_error::ext::BoxedError; +use common_query::{WriteProgress, prelude}; use common_telemetry::{debug, warn}; +use operator::write_dispatcher::TrackedWrite; use operator::write_dispatcher::WriteDispatcherRef; use partition::manager::PartitionRuleManagerRef; use session::context::QueryContextRef; -use snafu::OptionExt; +use snafu::{IntoError, OptionExt, ensure}; use crate::batcher::flow_notifier::FlowNotifier; use crate::batcher::logical_table; @@ -129,13 +130,21 @@ pub(crate) async fn flush_batch( .extension(logical_table::PHYSICAL_TABLE_KEY) .unwrap_or(prelude::GREPTIME_PHYSICAL_TABLE) .to_string(); - let result = flush_batch_physical( + let result = flush_batch_physical_tracked( &table_batches, &physical_table_name, &ctx, partition_manager, write_dispatcher, catalog_manager, + waiters.len().max( + table_batches + .iter() + .flat_map(|table| &table.batches) + .map(|batch| batch.owner + 1) + .max() + .unwrap_or(0), + ), ) .await; @@ -148,15 +157,42 @@ pub(crate) async fn flush_batch( ); match result { - Ok(affected_rows) => { - metrics::FLUSH_TOTAL.inc(); - metrics::FLUSH_ROWS.observe(total_row_count as f64); - operator::metrics::DIST_INGEST_ROW_COUNT - .with_label_values(&[db_string.as_str()]) - .inc_by(affected_rows as u64); - - notify_waiters(waiters, Ok(())); - flow_notifier::enqueue_flow_notifications(table_batches, &flow_notification_tx); + Ok(progress) => { + if progress.iter().any(|result| result.error.is_some()) { + metrics::FLUSH_FAILURES.inc(); + } else { + metrics::FLUSH_TOTAL.inc(); + metrics::FLUSH_ROWS.observe(total_row_count as f64); + } + let affected_rows = progress.iter().map(|result| result.affected_rows).sum(); + metrics::FLUSH_DROPPED_ROWS + .inc_by(total_row_count.saturating_sub(affected_rows) as u64); + let successful = progress + .iter() + .enumerate() + .filter_map(|(owner, result)| result.error.is_none().then_some(owner)) + .collect::>(); + for (waiter, result) in waiters.into_iter().zip(progress) { + operator::metrics::DIST_INGEST_ROW_COUNT + .with_label_values(&[db_string.as_str()]) + .inc_by(result.affected_rows as u64); + let _ = waiter.response_tx.send(result); + } + let accepted = table_batches + .into_iter() + .filter_map(|mut table| { + table + .batches + .retain(|batch| successful.contains(&batch.owner)); + table.row_count = table + .batches + .iter() + .map(|batch| batch.batch.num_rows()) + .sum(); + (!table.batches.is_empty()).then_some(table) + }) + .collect(); + flow_notifier::enqueue_flow_notifications(accepted, &flow_notification_tx); } Err(err) => { metrics::FLUSH_FAILURES.inc(); @@ -184,6 +220,40 @@ pub async fn flush_batch_physical( write_dispatcher: &(impl PhysicalFlushWriteDispatcher + ?Sized), catalog_manager: &(impl PhysicalFlushCatalogProvider + ?Sized), ) -> Result { + let inputs = table_batches + .iter() + .flat_map(|table| &table.batches) + .map(|batch| batch.owner + 1) + .max() + .unwrap_or(1); + let progress = flush_batch_physical_tracked( + table_batches, + physical_table_name, + ctx, + partition_manager, + write_dispatcher, + catalog_manager, + inputs, + ) + .await?; + let mut result = WriteProgress::default(); + for part in progress { + result.merge(part); + } + result + .into_result() + .map_err(|source| error::ExecuteGrpcQuerySnafu.into_error(source)) +} + +async fn flush_batch_physical_tracked( + table_batches: &[TableBatch], + physical_table_name: &str, + ctx: &QueryContextRef, + partition_manager: &(impl PhysicalFlushPartitionProvider + ?Sized), + write_dispatcher: &(impl PhysicalFlushWriteDispatcher + ?Sized), + catalog_manager: &(impl PhysicalFlushCatalogProvider + ?Sized), + inputs: usize, +) -> Result> { // 1. Resolve the physical table and get column ID mapping let physical_table = { let _timer = metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED @@ -237,42 +307,82 @@ pub async fn flush_batch_physical( // 4. Concatenate all modified batches (all share the same physical schema) let combined_batch = batch_convert::concat_modified_batches(&modified_batches)?; - let mut attribution = operator::write_dispatcher::WriteAttribution::default(); - for batch in table_batches.iter().flat_map(|table| &table.batches) { - attribution.extend(operator::write_dispatcher::WriteAttribution::new( - batch.batch.num_rows(), - batch.progress.clone(), - )); - } - // 5. Split by physical partition rule and send to regions let physical_table_id = physical_table_info.table_id(); - let planned_batches = region_write::plan_region_batches( - combined_batch, - physical_table_id, - partition_rule.as_ref(), - attribution, - )?; - - let region_writes = region_write::encode_region_write_requests( - planned_batches, - ctx.skip_wal(), - &partition_versions, - )?; - write_dispatcher.dispatch(region_writes).await + let row_count = table_batches + .iter() + .flat_map(|table| &table.batches) + .map(|batch| batch.batch.num_rows()) + .sum::(); + ensure!( + row_count == combined_batch.num_rows(), + error::InternalSnafu { + err_msg: "Physical transformation changed row ownership" + } + ); + let masks = + region_write::split_combined_batch_by_region(&combined_batch, partition_rule.as_ref())?; + let mut writes = Vec::new(); + for (region, mask) in masks { + let mut selected = Vec::with_capacity(mask.selected_rows()); + let bits = mask.array().values(); + let mut offset = 0; + // Filtering preserves batch order, and each original batch has one owner. + for batch in table_batches.iter().flat_map(|table| &table.batches) { + let num_rows = batch.batch.num_rows(); + let selected_rows = if num_rows == 1 { + // Avoid bitmap range setup for single-row submissions. + usize::from(bits.value(offset)) + } else { + // `offset` is mask-relative; the raw buffer also needs the mask's slice offset. + bits.inner() + .count_set_bits_offset(bits.offset() + offset, num_rows) + }; + selected.extend(std::iter::repeat_n(batch.owner, selected_rows)); + offset += num_rows; + } + if let Some(batch) = + region_write::plan_region_batch(&combined_batch, physical_table_id, region, &mask)? + { + let request = region_write::encode_region_write_request( + batch, + ctx.skip_wal(), + &partition_versions, + )?; + writes.push(TrackedWrite { + request, + owners: selected, + }); + } + } + let progress = write_dispatcher.dispatch_tracked(writes, inputs).await; + // A submission can contribute several logical batches; credit its counter once. + let counters = table_batches + .iter() + .flat_map(|table| &table.batches) + .filter_map(|batch| { + batch + .progress + .as_ref() + .map(|counter| (batch.owner, counter)) + }) + .collect::>(); + for (owner, counter) in counters { + counter.add_written(progress[owner].affected_rows); + } + Ok(progress) } pub(in crate::batcher::logical_table) fn notify_waiters( waiters: Vec, result: Result<()>, ) { - let shared_result = result.map_err(Arc::new); + let progress = match result { + Ok(()) => WriteProgress::default(), + Err(error) => WriteProgress::failed(BoxedError::new(error)), + }; for waiter in waiters { - let _ = waiter.response_tx.send(match &shared_result { - Ok(()) => Ok(()), - Err(error) => Err(Arc::clone(error)), - }); - // waiter._permit is dropped here, releasing the inflight semaphore slot + let _ = waiter.response_tx.send(progress.clone()); } } @@ -316,7 +426,7 @@ mod tests { use crate::batcher::test_util::mock_table_flownode_cache; use crate::error; use crate::error::Error; - use crate::metrics::FLOW_NOTIFICATION_DROPPED; + use crate::metrics::{self, FLOW_NOTIFICATION_DROPPED}; #[tokio::test] async fn test_flush_batch_notifies_flownode_after_successful_physical_write() { @@ -480,7 +590,7 @@ mod tests { } #[tokio::test] - async fn test_flush_batch_physical_returns_actual_affected_rows() { + async fn test_flush_batch_physical_rejects_unattributable_affected_rows() { let table_batches = vec![TableBatch { table_name: "t1".to_string(), table_id: 11, @@ -489,7 +599,7 @@ mod tests { }]; let ctx = session::context::QueryContext::arc(); - let affected_rows = flush_batch_physical( + let error = flush_batch_physical( &table_batches, "phy", &ctx, @@ -502,9 +612,334 @@ mod tests { }, ) .await - .unwrap(); + .unwrap_err(); - assert_eq!(7, affected_rows); + assert!(format!("{error:?}").contains("Unattributable")); + } + + #[tokio::test] + async fn test_exemplar_coalesced_logical_flush_retains_each_requests_progress() { + use common_error::ext::BoxedError; + use common_query::WriteProgress; + use operator::write_dispatcher::TrackedWrite; + + struct ProgressDispatcher([(usize, bool); 2]); + #[async_trait] + impl PhysicalFlushWriteDispatcher for ProgressDispatcher { + async fn dispatch(&self, _: Vec) -> error::Result { + unreachable!() + } + async fn dispatch_tracked( + &self, + requests: Vec, + inputs: usize, + ) -> Vec { + assert_eq!(inputs, 2); + assert_eq!( + requests + .iter() + .flat_map(|r| &r.owners) + .copied() + .collect::>(), + vec![0, 1, 1] + ); + self.0 + .iter() + .map(|&(affected_rows, failed)| WriteProgress { + affected_rows, + error: failed.then(|| { + Arc::new(BoxedError::new( + error::InternalSnafu { + err_msg: "lost acknowledgement", + } + .build(), + )) + }), + }) + .collect() + } + } + let ctx = session::context::QueryContext::arc(); + let ingested = operator::metrics::DIST_INGEST_ROW_COUNT + .with_label_values(&[ctx.get_db_string().as_str()]); + for (expected_progress, successful_flush, expected_timestamps) in [ + ([(1, false), (2, false)], true, vec![1000; 3]), + ([(1, false), (1, true)], false, vec![1000]), + ([(0, true), (0, true)], false, vec![]), + ] { + let limiter = RequestLimiter::try_new(2).unwrap(); + let mut waiters = Vec::new(); + let mut receivers = Vec::new(); + let counters = + [0, 1].map(|_| Arc::new(session::write_progress::WriteProgress::default())); + for counter in &counters { + let (response_tx, response_rx) = oneshot::channel(); + waiters.push(FlushWaiter { + _completion: Some(counter.start()), + response_tx, + _permit: limiter.acquire().await.unwrap(), + }); + receivers.push(response_rx); + } + let batches = [0, 1, 1] + .into_iter() + .map(|owner| { + let mut batch = mock_aligned_tag_batch("tag1", "host-1", 1000, 1.0); + batch.owner = owner; + batch.progress = Some(counters[owner].clone()); + batch + }) + .collect(); + let (notifier, mut notifications) = + FlowNotifier::try_new(1, FLOW_NOTIFICATION_DROPPED.clone()).unwrap(); + let successes = metrics::FLUSH_TOTAL.get(); + let failures = metrics::FLUSH_FAILURES.get(); + let samples = metrics::FLUSH_ROWS.get_sample_count(); + let sum = metrics::FLUSH_ROWS.get_sample_sum(); + let dropped = metrics::FLUSH_DROPPED_ROWS.get(); + let written = ingested.get(); + flush_batch( + Batch { + table_batches: vec![TableBatch { + table_name: "cpu".into(), + table_id: 42, + batches, + row_count: 3, + }], + total_row_count: 3, + db_string: ctx.get_db_string(), + ctx: ctx.clone(), + waiters, + }, + &MockFlushPartitionProvider { + partition_rule_calls: Arc::new(AtomicUsize::new(0)), + }, + &ProgressDispatcher(expected_progress), + &MockFlushCatalogProvider { + table: Some(mock_physical_table_metadata(1024)), + }, + notifier, + ) + .await; + let expected_written: usize = expected_progress.iter().map(|(rows, _)| rows).sum(); + assert_eq!( + metrics::FLUSH_TOTAL.get(), + successes + u64::from(successful_flush) + ); + assert_eq!( + metrics::FLUSH_FAILURES.get(), + failures + u64::from(!successful_flush) + ); + assert_eq!( + metrics::FLUSH_ROWS.get_sample_count(), + samples + u64::from(successful_flush) + ); + assert_eq!( + metrics::FLUSH_ROWS.get_sample_sum(), + sum + if successful_flush { 3.0 } else { 0.0 } + ); + assert_eq!( + metrics::FLUSH_DROPPED_ROWS.get(), + dropped + (3 - expected_written) as u64 + ); + assert_eq!(ingested.get(), written + expected_written as u64); + for (counter, (affected_rows, _)) in counters.iter().zip(expected_progress) { + counter.wait().await; + assert_eq!(counter.written(), affected_rows as u64); + } + for (receiver, (affected_rows, failed)) in receivers.into_iter().zip(expected_progress) + { + let result = receiver.await.unwrap(); + assert_eq!(result.affected_rows, affected_rows); + assert_eq!(result.error.is_some(), failed); + } + if !expected_timestamps.is_empty() { + let notification = notifications.try_recv().unwrap(); + assert_eq!(notification.table_id, 42); + assert_eq!(notification.timestamps, expected_timestamps); + } + assert!(notifications.try_recv().is_err()); + } + } + + #[tokio::test] + async fn test_flush_batch_physical_preserves_owners_with_sliced_masks() { + use std::sync::Mutex; + + use api::v1::region::{bulk_insert_request, region_request}; + use arrow::array::{BooleanArray, Float64Array, StringArray, TimestampMillisecondArray}; + use common_grpc::flight::FlightDecoder; + use common_query::WriteProgress; + use operator::write_dispatcher::TrackedWrite; + + #[derive(Clone)] + struct MaskPartitions { + offset: usize, + pattern: usize, + } + + impl PartitionRule for MaskPartitions { + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn partition_columns(&self) -> &[String] { + &[] + } + + fn find_region(&self, _: &[datatypes::prelude::Value]) -> PartitionResult { + unreachable!() + } + + fn split_record_batch( + &self, + batch: &RecordBatch, + ) -> PartitionResult> { + Ok([false, true] + .into_iter() + .enumerate() + .map(|(region, selected)| { + let values = std::iter::repeat_n(true, self.offset) + .chain((0..batch.num_rows()).map(|row| { + let matches = match self.pattern { + 0 => false, + 1 => true, + _ => row % 3 != 0, + }; + matches == selected + })) + .chain([true; 8]) + .collect::>(); + let mask = BooleanArray::from(values).slice(self.offset, batch.num_rows()); + (region as u32, RegionMask::from(mask)) + }) + .collect()) + } + } + + #[async_trait] + impl PhysicalFlushPartitionProvider for MaskPartitions { + async fn find_table_partition_rule( + &self, + _: &table::metadata::TableInfo, + ) -> PartitionResult<(PartitionRuleRef, HashMap>)> { + Ok((Arc::new(self.clone()), HashMap::new())) + } + } + + #[derive(Default)] + struct RecordingDispatcher(Mutex>); + + #[async_trait] + impl PhysicalFlushWriteDispatcher for RecordingDispatcher { + async fn dispatch(&self, _: Vec) -> error::Result { + unreachable!() + } + + async fn dispatch_tracked( + &self, + requests: Vec, + inputs: usize, + ) -> Vec { + let mut progress = vec![WriteProgress::default(); inputs]; + for owner in requests.iter().flat_map(|request| &request.owners) { + progress[*owner].affected_rows += 1; + } + *self.0.lock().unwrap() = requests; + progress + } + } + + let ctx = session::context::QueryContext::arc(); + let catalog = MockFlushCatalogProvider { + table: Some(mock_physical_table_metadata(1024)), + }; + for lengths in [[0; 7], [1; 7], [0, 3, 65, 1, 127, 0, 64]] { + let mut expected = Vec::new(); + let tables = lengths + .chunks(4) + .enumerate() + .map(|(table, lengths)| { + let batches = lengths + .iter() + .enumerate() + .map(|(index, &rows)| { + let owner = [2, 0, 1, 0, 2, 1, 0][table * 4 + index]; + let start = expected.len(); + expected.extend(std::iter::repeat_n(owner, rows)); + let mut batch = mock_aligned_tag_batch("tag1", "host-1", 0, 1.0); + batch.batch = RecordBatch::try_new( + batch.batch.schema(), + vec![ + Arc::new(TimestampMillisecondArray::from_iter_values( + (start..start + rows).map(|row| row as i64), + )), + Arc::new(Float64Array::from(vec![1.0; rows])), + Arc::new(StringArray::from(vec!["host-1"; rows])), + ], + ) + .unwrap(); + batch.owner = owner; + batch + }) + .collect(); + TableBatch { + table_name: format!("logical_{table}"), + table_id: 100 + table as u32, + batches, + row_count: lengths.iter().sum(), + } + }) + .collect::>(); + for offset in [0, 1, 7, 65] { + for pattern in 0..3 { + let dispatcher = RecordingDispatcher::default(); + let affected = flush_batch_physical( + &tables, + "phy", + &ctx, + &MaskPartitions { offset, pattern }, + &dispatcher, + &catalog, + ) + .await + .unwrap(); + assert_eq!(affected, expected.len()); + let mut actual = Vec::new(); + for write in dispatcher.0.into_inner().unwrap() { + let Some(region_request::Body::BulkInsert(request)) = write.request.body + else { + panic!("expected bulk request") + }; + let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = request.body else { + panic!("expected Arrow IPC") + }; + let batch = FlightDecoder::try_from_schema_bytes(&ipc.schema) + .unwrap() + .try_decode_record_batch(&ipc.data_header, &ipc.payload) + .unwrap(); + assert_eq!(batch.num_rows(), write.owners.len()); + let timestamps = batch + .column_by_name("greptime_timestamp") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + assert!(timestamps.values().windows(2).all(|pair| pair[0] < pair[1])); + actual.extend(timestamps.values().iter().copied().zip(write.owners)); + } + actual.sort_unstable(); + assert_eq!( + actual, + expected + .iter() + .enumerate() + .map(|(row, owner)| (row as i64, *owner)) + .collect::>(), + ); + } + } + } } #[tokio::test] @@ -688,17 +1123,14 @@ mod tests { #[async_trait] impl PhysicalFlushWriteDispatcher for MockFlushWriteDispatcher { - async fn dispatch( - &self, - _requests: Vec<(RegionRequest, operator::write_dispatcher::WriteAttribution)>, - ) -> error::Result { + async fn dispatch(&self, _requests: Vec) -> error::Result { self.writes.fetch_add(1, Ordering::SeqCst); if self.fail { return Err(Error::Internal { err_msg: "physical write failed".to_string(), }); } - Ok(0) + Ok(1) } } @@ -718,10 +1150,7 @@ mod tests { #[async_trait] impl PhysicalFlushWriteDispatcher for AffectedRowsFlushWriteDispatcher { - async fn dispatch( - &self, - _requests: Vec<(RegionRequest, operator::write_dispatcher::WriteAttribution)>, - ) -> error::Result { + async fn dispatch(&self, _requests: Vec) -> error::Result { Ok(self.affected_rows) } } diff --git a/src/servers/src/batcher/logical_table/batch_convert.rs b/src/servers/src/batcher/logical_table/batch_convert.rs index 322b12488e8..5711fa190a9 100644 --- a/src/servers/src/batcher/logical_table/batch_convert.rs +++ b/src/servers/src/batcher/logical_table/batch_convert.rs @@ -34,6 +34,7 @@ use crate::metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED; pub struct RecordBatchWithTsIdx { pub(in crate::batcher::logical_table) progress: Option>, + pub(in crate::batcher::logical_table) owner: usize, /// The aligned logical record batch. pub(in crate::batcher::logical_table) batch: RecordBatch, /// The timestamp column index in `batch`. @@ -66,6 +67,7 @@ impl RecordBatchWithTsIdx { Ok(Self { progress: None, + owner: 0, batch, timestamp_index, }) diff --git a/src/servers/src/batcher/logical_table/pending_worker.rs b/src/servers/src/batcher/logical_table/pending_worker.rs index 9f9c689e6b0..c803c3de49f 100644 --- a/src/servers/src/batcher/logical_table/pending_worker.rs +++ b/src/servers/src/batcher/logical_table/pending_worker.rs @@ -22,6 +22,7 @@ use common_batcher::flush_policy::FlushTrigger; use common_batcher::flush_policy::timing::TimingFlushPolicy; use common_batcher::pending_worker::PendingWorker as PendingCore; use common_batcher::worker_registry::WorkerRegistry; +use common_query::WriteProgress; use common_telemetry::debug; use operator::write_dispatcher::WriteDispatcherRef; use partition::manager::PartitionRuleManagerRef; @@ -33,10 +34,10 @@ use crate::batcher::flow_notifier::FlowNotifier; use crate::batcher::logical_table::BatchKey; use crate::batcher::logical_table::batch::{self, Batch}; use crate::batcher::logical_table::batch_convert::{RecordBatchWithTsIdx, TableBatch}; -use crate::error::Error; use crate::metrics::{PENDING_BATCHES, PENDING_ROWS, PENDING_WORKERS}; pub(in crate::batcher::logical_table) struct PendingBatch { + inputs: usize, pub(in crate::batcher::logical_table) tables: HashMap, pub(in crate::batcher::logical_table) db_string: String, pub(in crate::batcher::logical_table) ctx: QueryContextRef, @@ -45,8 +46,7 @@ pub(in crate::batcher::logical_table) struct PendingBatch { pub(in crate::batcher::logical_table) struct FlushWaiter { pub(in crate::batcher::logical_table) _completion: Option, - pub(in crate::batcher::logical_table) response_tx: - oneshot::Sender>>, + pub(in crate::batcher::logical_table) response_tx: oneshot::Sender, pub(in crate::batcher::logical_table) _permit: Arc, } @@ -62,7 +62,7 @@ pub(crate) enum WorkerCommand { table_batches: Vec<(String, u32, RecordBatchWithTsIdx)>, total_rows: usize, ctx: QueryContextRef, - response_tx: oneshot::Sender>>, + response_tx: oneshot::Sender, _permit: Arc, }, #[cfg(test)] @@ -73,6 +73,7 @@ impl PendingBatch { pub(in crate::batcher::logical_table) fn new(ctx: QueryContextRef) -> Self { let db_string = ctx.get_db_string(); Self { + inputs: 0, tables: HashMap::new(), db_string, ctx, @@ -141,9 +142,11 @@ pub(crate) fn start_worker( for (table_name, table_id, mut record_batch) in table_batches { record_batch.progress = progress.clone(); + record_batch.owner = pending_batch.inputs; pending_batch.add_table_batch(table_name, table_id, record_batch); } + pending_batch.inputs += 1; PENDING_ROWS.add(total_rows as i64); if let Some(flush) = drain_batch(&mut batch, &mut pending_flush, Some(FlushTrigger::Submission)) { @@ -279,6 +282,7 @@ pub(in crate::batcher::logical_table) fn drain_batch( #[cfg(test)] mod tests { + use common_query::WriteProgress; use std::collections::HashMap; use std::sync::Arc; use std::time::Duration; @@ -309,7 +313,6 @@ mod tests { use crate::batcher::logical_table::test_util::{ ConcurrentMockNodeManager, mock_aligned_tag_batch, }; - use crate::error::Error; use crate::metrics::FLOW_NOTIFICATION_DROPPED; #[tokio::test] @@ -328,6 +331,7 @@ mod tests { 1, ); let mut batch = Some(PendingBatch { + inputs: 1, tables: HashMap::from([( 42, TableBatch { @@ -616,7 +620,7 @@ mod tests { worker_tx: &mpsc::Sender, total_rows: usize, timestamp: i64, - ) -> oneshot::Receiver>> { + ) -> oneshot::Receiver { let (response_tx, response_rx) = oneshot::channel(); let permit = Arc::new(Semaphore::new(1)).acquire_owned().await.unwrap(); worker_tx @@ -648,9 +652,9 @@ mod tests { } async fn receive_mock_flush_result( - response_rx: oneshot::Receiver>>, + response_rx: oneshot::Receiver, context: &str, - ) -> std::result::Result<(), Arc> { + ) -> WriteProgress { // Under paused time the timeout auto-advances the clock and fires // deterministically if the flush never completes. tokio::time::timeout(WORKER_TEST_TIMEOUT, response_rx) @@ -659,14 +663,12 @@ mod tests { .expect("flush result channel closed without a result") } - fn assert_missing_physical_table(result: std::result::Result<(), Arc>) { - let err = result.expect_err("the empty catalog should make the flush fail"); + fn assert_missing_physical_table(result: WriteProgress) { + let err = result + .error + .expect("the empty catalog should make the flush fail"); assert!( - matches!( - err.as_ref(), - Error::Internal { err_msg } - if err_msg.contains("not found during pending flush") - ), + err.to_string().contains("not found during pending flush"), "unexpected flush error: {err}" ); } diff --git a/src/servers/src/batcher/logical_table/region_write.rs b/src/servers/src/batcher/logical_table/region_write.rs index e64cced9d54..1875cc088d9 100644 --- a/src/servers/src/batcher/logical_table/region_write.rs +++ b/src/servers/src/batcher/logical_table/region_write.rs @@ -22,10 +22,13 @@ use arrow::compute; use arrow::record_batch::RecordBatch; use async_trait::async_trait; use catalog::CatalogManagerRef; +use common_error::ext::BoxedError; use common_grpc::error::Error as GrpcError; use common_grpc::flight; +use common_query::WriteProgress; use common_telemetry::tracing_context::TracingContext; -use operator::write_dispatcher::{WriteAttribution, WriteDispatcherRef}; +use operator::write_dispatcher::TrackedWrite; +use operator::write_dispatcher::WriteDispatcherRef; use partition::manager::PartitionRuleManagerRef; use partition::partition::PartitionRuleRef; use snafu::ResultExt; @@ -68,7 +71,43 @@ pub trait PhysicalFlushPartitionProvider: Send + Sync { #[async_trait] pub trait PhysicalFlushWriteDispatcher: Send + Sync { /// Dispatches one physical flush together, sharing routing refreshes across its writes. - async fn dispatch(&self, requests: Vec<(RegionRequest, WriteAttribution)>) -> Result; + async fn dispatch(&self, requests: Vec) -> Result; + + async fn dispatch_tracked( + &self, + requests: Vec, + inputs: usize, + ) -> Vec { + let mut progress = vec![WriteProgress::default(); inputs]; + let total = requests.iter().map(|r| r.owners.len()).sum::(); + let owners = requests + .iter() + .flat_map(|r| r.owners.iter().copied()) + .collect::>(); + match self + .dispatch(requests.into_iter().map(|r| r.request).collect()) + .await + { + Ok(rows) if rows == total => { + for owner in owners { + progress[owner].affected_rows += 1; + } + } + Ok(_) => { + return vec![ + WriteProgress::failed(BoxedError::new( + error::InternalSnafu { + err_msg: "Unattributable physical write acknowledgement" + } + .build() + )); + inputs + ]; + } + Err(error) => return vec![WriteProgress::failed(BoxedError::new(error)); inputs], + } + progress + } } #[derive(Clone)] @@ -126,16 +165,30 @@ pub(crate) struct WriteDispatcherPhysicalFlushAdapter { #[async_trait] impl PhysicalFlushWriteDispatcher for WriteDispatcherPhysicalFlushAdapter { - async fn dispatch(&self, requests: Vec<(RegionRequest, WriteAttribution)>) -> Result { + async fn dispatch_tracked( + &self, + requests: Vec, + inputs: usize, + ) -> Vec { self.write_dispatcher - .dispatch_attributed( + .dispatch_tracked( requests, + inputs, Some( - metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED + &metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED .with_label_values(&["flush_write_region"]), ), ) .await + } + async fn dispatch(&self, requests: Vec) -> Result { + self.write_dispatcher + .dispatch_with_metrics( + requests, + metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED + .with_label_values(&["flush_write_region"]), + ) + .await .context(error::CommonMetaSnafu) .map(|response| response.affected_rows) } @@ -143,7 +196,6 @@ impl PhysicalFlushWriteDispatcher for WriteDispatcherPhysicalFlushAdapter { /// Arrow rows selected for one physical region. pub(crate) struct PlannedRegionBatch { - pub(crate) attribution: WriteAttribution, pub(crate) region_id: RegionId, pub(crate) batch: RecordBatch, } @@ -172,7 +224,6 @@ pub(crate) fn plan_region_batch( physical_table_id: TableId, region_number: u32, mask: &partition::partition::RegionMask, - attribution: &WriteAttribution, ) -> Result> { if mask.select_none() { return Ok(None); @@ -193,18 +244,17 @@ pub(crate) fn plan_region_batch( } Ok(Some(PlannedRegionBatch { - attribution: attribution.filter(mask).context(error::CommonMetaSnafu)?, region_id: RegionId::new(physical_table_id, region_number), batch: region_batch, })) } /// Partitions a physical flush while retaining columns needed for rerouting. +#[cfg(test)] pub(crate) fn plan_region_batches( combined_batch: RecordBatch, physical_table_id: TableId, partition_rule: &dyn partition::partition::PartitionRule, - attribution: WriteAttribution, ) -> Result> { let region_masks = split_combined_batch_by_region(&combined_batch, partition_rule)?; @@ -212,13 +262,9 @@ pub(crate) fn plan_region_batches( // Mito's sparse bulk conversion projects the stored fields and encoded primary key. let mut planned_batches = Vec::new(); for (region_number, mask) in region_masks { - if let Some(planned_batch) = plan_region_batch( - &combined_batch, - physical_table_id, - region_number, - &mask, - &attribution, - )? { + if let Some(planned_batch) = + plan_region_batch(&combined_batch, physical_table_id, region_number, &mask)? + { planned_batches.push(planned_batch); } } @@ -226,50 +272,43 @@ pub(crate) fn plan_region_batches( Ok(planned_batches) } -/// Encodes each region batch with the partition version used to select its rows. -pub(crate) fn encode_region_write_requests( - planned_batches: Vec, +/// Encodes a region batch with the partition version used to select its rows. +pub(crate) fn encode_region_write_request( + planned: PlannedRegionBatch, skip_wal: bool, partition_versions: &HashMap>, -) -> Result> { - let mut region_writes = Vec::with_capacity(planned_batches.len()); - for planned in planned_batches { - let region_id = planned.region_id; - let (schema_bytes, data_header, payload) = { - let _timer = metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED - .with_label_values(&["flush_physical_encode_ipc"]) - .start_timer(); - flight::record_batch_to_ipc(planned.batch).map_err(map_ipc_error)? - }; +) -> Result { + let region_id = planned.region_id; + let (schema_bytes, data_header, payload) = { + let _timer = metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED + .with_label_values(&["flush_physical_encode_ipc"]) + .start_timer(); + flight::record_batch_to_ipc(planned.batch).map_err(map_ipc_error)? + }; - let request = RegionRequest { - header: Some(RegionRequestHeader { - tracing_context: TracingContext::from_current_span().to_w3c(), - ..Default::default() - }), - body: Some(region_request::Body::BulkInsert(BulkInsertRequest { - skip_wal, - region_id: region_id.as_u64(), - partition_expr_version: partition_versions - .get(®ion_id.region_number()) - .copied() - .flatten() - .map(|value| PartitionExprVersion { value }), - // Set aligned_schema_version to None so that datanode will check the batch schema again to see if any - // column is missing. - aligned_schema_version: None, - body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc { - schema: schema_bytes, - data_header, - payload, - })), + Ok(RegionRequest { + header: Some(RegionRequestHeader { + tracing_context: TracingContext::from_current_span().to_w3c(), + ..Default::default() + }), + body: Some(region_request::Body::BulkInsert(BulkInsertRequest { + skip_wal, + region_id: region_id.as_u64(), + partition_expr_version: partition_versions + .get(®ion_id.region_number()) + .copied() + .flatten() + .map(|value| PartitionExprVersion { value }), + // Set aligned_schema_version to None so that datanode will check the batch schema again to see if any + // column is missing. + aligned_schema_version: None, + body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc { + schema: schema_bytes, + data_header, + payload, })), - }; - - region_writes.push((request, planned.attribution)); - } - - Ok(region_writes) + })), + }) } pub(in crate::batcher::logical_table) fn map_ipc_error(error: GrpcError) -> Error { @@ -287,19 +326,65 @@ mod tests { use std::collections::{HashMap, HashSet}; use std::sync::Arc; + use api::region::RegionResponse; use api::v1::PartitionExprVersion; - use api::v1::region::{bulk_insert_request, region_request}; + use api::v1::region::{ + InsertRequest, InsertRequests, RegionRequest, bulk_insert_request, region_request, + }; use arrow::array::{Array, BinaryArray, BooleanArray, StringArray, TimestampMillisecondArray}; use arrow::datatypes::{DataType as ArrowDataType, Field, Schema as ArrowSchema}; use arrow::record_batch::RecordBatch; use common_grpc::flight::FlightDecoder; + use common_meta::ddl::test_util::datanode_handler::DatanodeWatcher; + use common_meta::test_util::MockDatanodeManager; + use operator::test_util::{create_partition_rule_manager, prepare_mocked_backend}; + use operator::write_dispatcher::{TrackedWrite, WriteDispatcher}; use partition::partition::{PartitionRule, RegionMask}; use store_api::storage::RegionId; use crate::batcher::logical_table::batch_convert::{self, RecordBatchWithTsIdx, TableBatch}; use crate::batcher::logical_table::region_write::{ - PlannedRegionBatch, encode_region_write_requests, plan_region_batches, + PhysicalFlushWriteDispatcher, PlannedRegionBatch, WriteDispatcherPhysicalFlushAdapter, + encode_region_write_request, plan_region_batches, }; + use crate::metrics; + + #[tokio::test] + async fn test_tracked_dispatch_records_flush_write_region_metric() { + let partitions = create_partition_rule_manager(prepare_mocked_backend().await).await; + let (sender, _receiver) = tokio::sync::mpsc::channel(1); + let handler = DatanodeWatcher::new(sender).with_handler(|_, _| Ok(RegionResponse::new(1))); + let dispatcher = WriteDispatcherPhysicalFlushAdapter { + write_dispatcher: Arc::new(WriteDispatcher::new( + Arc::new(MockDatanodeManager::new(handler)), + partitions, + )), + }; + let histogram = metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED + .with_label_values(&["flush_write_region"]); + let before = histogram.get_sample_count(); + let progress = dispatcher + .dispatch_tracked( + vec![TrackedWrite { + request: RegionRequest { + body: Some(region_request::Body::Inserts(InsertRequests { + requests: vec![InsertRequest { + region_id: RegionId::new(1, 1).as_u64(), + ..Default::default() + }], + })), + ..Default::default() + }, + owners: vec![0], + }], + 1, + ) + .await; + + assert_eq!(progress[0].affected_rows, 1); + assert!(progress[0].error.is_none()); + assert_eq!(histogram.get_sample_count(), before + 1); + } #[test] fn test_plan_region_batches_preserves_partition_columns_through_encoding() { @@ -344,7 +429,6 @@ mod tests { &TwoRegionPartitionRule { partition_columns: vec!["host".to_string()], }, - operator::write_dispatcher::WriteAttribution::new(2, None), ) .unwrap(); planned_batches.sort_by_key(|planned| planned.region_id.region_number()); @@ -357,13 +441,9 @@ mod tests { assert_eq!(1, planned_batches[1].num_rows()); assert_eq!(4, planned_batches[1].batch.num_columns()); - let writes = encode_region_write_requests( - planned_batches, - false, - &HashMap::from([(1, Some(7)), (2, Some(8))]), - ) - .unwrap(); - for (index, (write, _)) in writes.into_iter().enumerate() { + let partition_versions = HashMap::from([(1, Some(7)), (2, Some(8))]); + for (index, planned) in planned_batches.into_iter().enumerate() { + let write = encode_region_write_request(planned, false, &partition_versions).unwrap(); let Some(region_request::Body::BulkInsert(request)) = write.body else { panic!("expected bulk") }; @@ -394,9 +474,9 @@ mod tests { } #[test] - fn test_encode_region_write_requests_builds_bulk_insert_requests() { - check_encode_region_write_requests(false); - check_encode_region_write_requests(true); + fn test_encode_region_write_request_builds_bulk_insert_request() { + check_encode_region_write_request(false); + check_encode_region_write_request(true); } struct TwoRegionPartitionRule { @@ -435,9 +515,8 @@ mod tests { } } - fn check_encode_region_write_requests(skip_wal: bool) { + fn check_encode_region_write_request(skip_wal: bool) { let planned_batch = PlannedRegionBatch { - attribution: operator::write_dispatcher::WriteAttribution::new(1, None), region_id: RegionId::new(1024, 1), batch: RecordBatch::try_new( Arc::new(ArrowSchema::new(vec![ @@ -457,15 +536,11 @@ mod tests { ) .unwrap(), }; - let writes = encode_region_write_requests( - vec![planned_batch], - skip_wal, - &HashMap::from([(1, Some(7))]), - ) - .unwrap(); + let write = + encode_region_write_request(planned_batch, skip_wal, &HashMap::from([(1, Some(7))])) + .unwrap(); - assert_eq!(1, writes.len()); - let Some(region_request::Body::BulkInsert(request)) = &writes[0].0.body else { + let Some(region_request::Body::BulkInsert(request)) = &write.body else { panic!("expected bulk insert request"); }; assert_eq!(RegionId::new(1024, 1).as_u64(), request.region_id); diff --git a/src/servers/src/batcher/table.rs b/src/servers/src/batcher/table.rs index b2228ecb057..7840fd355c4 100644 --- a/src/servers/src/batcher/table.rs +++ b/src/servers/src/batcher/table.rs @@ -31,9 +31,9 @@ use common_batcher::flush_policy::timing::TimingFlushPolicy; use common_batcher::request_limiter::RequestLimiter; use common_batcher::worker_registry::WorkerRegistry; use common_error::ext::BoxedError; -use common_query::{Output, WriteCompletion}; +use common_query::{Output, WriteCompletion, WriteProgress}; use operator::batcher::PendingRowsBatcher; -use operator::error::{BatchFlushSnafu, Result, UnexpectedSnafu}; +use operator::error::{ExternalSnafu, Result, UnexpectedSnafu}; use operator::insert::Inserter; use session::context::QueryContextRef; use snafu::ResultExt; @@ -175,6 +175,7 @@ impl PendingRowsBatcher for TablePendingRowsBatcher { permit: Arc, ) -> Result { let total_rows = batch.num_rows(); + let confirmed = ctx.confirmed_writes(); if total_rows == 0 { return Ok(Output::new_with_affected_rows(0)); } @@ -218,20 +219,13 @@ impl PendingRowsBatcher for TablePendingRowsBatcher { .fail(); } let completion = async move { - response_rx - .await - .map_err(|_| { - UnexpectedSnafu { - violated: "batch worker stopped before reporting its write result" - .to_string(), - } - .build() - })? - .context(BatchFlushSnafu) + response_rx.await.unwrap_or_else(|_| WriteProgress::failed(BoxedError::new( + UnexpectedSnafu { violated: "batch worker stopped before acknowledging its write; storage outcome unknown".to_string() }.build() + ))) }; let mut output = Output::new_with_affected_rows(total_rows); - if self.pending_rows_batch_sync { - completion.await?; + if self.pending_rows_batch_sync && !confirmed { + completion.await.into_result().context(ExternalSnafu)?; } else { if let Some(progress) = progress { progress.mark_accepted(); @@ -239,9 +233,7 @@ impl PendingRowsBatcher for TablePendingRowsBatcher { output .meta .write_completions - .push(WriteCompletion::new(async move { - completion.await.map_err(BoxedError::new) - })); + .push(WriteCompletion::with_progress(completion)); } Ok(output) } @@ -281,7 +273,6 @@ mod tests { use datatypes::value::Value as DtValue; use datatypes::vectors::{Int32Vector, TimestampMillisecondVector, VectorRef}; use operator::batcher::PendingRowsBatcher; - use operator::error::Error; use operator::insert::Inserter; use operator::metrics::DIST_INGEST_ROW_COUNT; use operator::req_convert::insert::rows_to_record_batch; @@ -695,6 +686,99 @@ mod tests { } } + #[tokio::test] + async fn test_exemplar_coalesced_table_flush_retains_each_requests_progress() { + #[derive(Clone)] + struct PartialHandler; + #[async_trait::async_trait] + impl MockDatanodeHandler for PartialHandler { + async fn handle(&self, _: &Peer, request: RegionRequest) -> MetaResult { + let Some(region_request::Body::BulkInsert(request)) = request.body else { + panic!("expected bulk") + }; + let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = request.body else { + panic!("expected Arrow") + }; + let batch = FlightDecoder::try_from_schema_bytes(&ipc.schema) + .unwrap() + .try_decode_record_batch(&ipc.data_header, &ipc.payload) + .unwrap(); + let values = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + if values.values().contains(&20) { + return common_meta::error::UnexpectedSnafu { + err_msg: "lost acknowledgement", + } + .fail(); + } + Ok(RegionResponse::new(batch.num_rows())) + } + async fn handle_query( + &self, + _: &Peer, + _: QueryRequest, + ) -> MetaResult { + unreachable!() + } + } + for sync in [false, true] { + let partitions = create_partition_rule_manager(prepare_mocked_backend().await).await; + let inserter = Arc::new(Inserter::new( + MemoryCatalogManager::new(), + partitions, + Arc::new(MockDatanodeManager::new(PartialHandler)), + mock_table_flownode_cache(1, vec![]).await, + true, + )); + let mut batcher = TablePendingRowsBatcher::try_new( + Duration::from_secs(3600), + 3, + 1, + 4, + 4, + NonZeroUsize::new(16).unwrap(), + inserter, + ) + .unwrap(); + Arc::get_mut(&mut batcher).unwrap().pending_rows_batch_sync = sync; + let table = Arc::new(new_test_table_info(1, "table_1", [1, 2, 3].into_iter())); + let mut first = rows(1); + first.rows.extend(rows(20).rows); + let first = rows_to_record_batch(&first, &table).unwrap(); + let second = rows_to_record_batch(&rows(2), &table).unwrap(); + let counters = + [0, 1].map(|_| Arc::new(session::write_progress::WriteProgress::default())); + let contexts = counters.each_ref().map(|counter| { + let mut ctx = QueryContext::arc().with_confirmed_writes(); + ctx.set_write_progress(counter.clone()); + Arc::new(ctx) + }); + let first_permit = batcher.acquire().await.unwrap(); + let second_permit = batcher.acquire().await.unwrap(); + let (first, second) = tokio::join!( + batcher.submit(table.clone(), first, contexts[0].clone(), first_permit), + batcher.submit(table, second, contexts[1].clone(), second_permit), + ); + let first = timeout(Duration::from_secs(5), first.unwrap().write_progress()) + .await + .unwrap(); + let second = timeout(Duration::from_secs(5), second.unwrap().write_progress()) + .await + .unwrap(); + assert_eq!(first.affected_rows, 1); + assert!(first.error.is_some()); + assert_eq!(second.affected_rows, 1); + assert!(second.error.is_none()); + for counter in counters { + counter.wait().await; + assert_eq!(counter.written(), 1); + } + } + } + #[test] fn test_batch_key_uses_name_and_wal_policy() { let ctx = QueryContext::arc(); @@ -810,7 +894,7 @@ mod tests { &self, count: usize, changed_schema: bool, - ) -> oneshot::Receiver>> { + ) -> oneshot::Receiver { self.submit_with_table(count, changed_schema, self.table.clone()) .await } @@ -820,7 +904,7 @@ mod tests { count: usize, changed_schema: bool, table_info: TableInfoRef, - ) -> oneshot::Receiver>> { + ) -> oneshot::Receiver { let mut input = rows(1); input.rows = vec![input.rows[0].clone(); count]; let mut batch = rows_to_record_batch(&input, &self.table).unwrap(); @@ -888,11 +972,13 @@ mod tests { .await .unwrap() .unwrap() + .into_result() .unwrap(); timeout(WORKER_TIMEOUT, second) .await .unwrap() .unwrap() + .into_result() .unwrap(); worker.shutdown.send(()).unwrap(); worker.stopped().await; @@ -917,11 +1003,13 @@ mod tests { .await .unwrap() .unwrap() + .into_result() .unwrap(); timeout(WORKER_TIMEOUT, second) .await .unwrap() .unwrap() + .into_result() .unwrap(); worker.shutdown.send(()).unwrap(); worker.stopped().await; @@ -956,6 +1044,7 @@ mod tests { .await .unwrap() .unwrap() + .into_result() .unwrap(); if recreated { // The recreated table has no route in this fixture; its error belongs only to it. @@ -964,6 +1053,7 @@ mod tests { .await .unwrap() .unwrap() + .into_result() .is_err() ); } else { @@ -973,6 +1063,7 @@ mod tests { .await .unwrap() .unwrap() + .into_result() .unwrap(); } worker.shutdown.send(()).unwrap(); @@ -994,6 +1085,7 @@ mod tests { .await .unwrap() .unwrap() + .into_result() .unwrap(); worker.stopped().await; assert!(matches!( @@ -1005,6 +1097,7 @@ mod tests { .await .unwrap() .unwrap() + .into_result() .unwrap(); } } diff --git a/src/servers/src/batcher/table/batch.rs b/src/servers/src/batcher/table/batch.rs index 42dfd52f82e..babdf9e04bc 100644 --- a/src/servers/src/batcher/table/batch.rs +++ b/src/servers/src/batcher/table/batch.rs @@ -17,13 +17,12 @@ use std::time::Instant; use arrow::compute::concat_batches; use arrow::record_batch::RecordBatch; +use common_query::WriteProgress; use common_telemetry::error; use operator::error::{ComputeArrowSnafu, Result, UnexpectedSnafu}; use operator::insert::Inserter; use operator::metrics::DIST_INGEST_ROW_COUNT; -use operator::write_dispatcher::WriteAttribution; use snafu::ResultExt; -use table::metadata::TableInfoRef; use crate::batcher::table::flow_notifier::FlowNotifier; use crate::batcher::table::metrics::{ @@ -46,15 +45,37 @@ pub(in crate::batcher::table) async fn flush_batch( let result = send_batch(&batch, &inserter).await; FLUSH_ELAPSED.observe(started.elapsed().as_secs_f64()); match result { - Ok((table, combined)) => { - FLUSH_TOTAL.inc(); - FLUSH_ROWS.observe(batch.total_rows as f64); - DIST_INGEST_ROW_COUNT - .with_label_values(&[batch.submissions[0].ctx.get_db_string().as_str()]) - .inc_by(batch.total_rows as u64); - notify_batches(batch.submissions, Ok(())); - // Admission is best-effort and only occurs after affected-row validation. - notifier.notify(table, &combined); + Ok(progress) => { + if progress.iter().any(|result| result.error.is_some()) { + FLUSH_FAILURES.inc(); + } else { + FLUSH_TOTAL.inc(); + FLUSH_ROWS.observe(batch.total_rows as f64); + } + if let Some(first) = batch.submissions.first() { + notifier.notify( + first.table_info.clone(), + batch + .submissions + .iter() + .zip(&progress) + .filter(|(_, result)| result.error.is_none()) + .map(|(pending, _)| &pending.batch), + ); + } + for (pending, result) in batch.submissions.into_iter().zip(progress) { + if let Some(counter) = pending.ctx.write_progress() { + counter.add_written(result.affected_rows); + } + DIST_INGEST_ROW_COUNT + .with_label_values(&[pending.ctx.get_db_string().as_str()]) + .inc_by(result.affected_rows as u64); + if result.error.is_some() { + FLUSH_DROPPED_ROWS + .inc_by((pending.batch.num_rows() - result.affected_rows) as u64); + } + let _ = pending.response_tx.send(result); + } } Err(error) => { error!(error; "Failed to flush table batch, rows: {}", batch.total_rows); @@ -65,7 +86,7 @@ pub(in crate::batcher::table) async fn flush_batch( } } -async fn send_batch(batch: &Batch, inserter: &Inserter) -> Result<(TableInfoRef, RecordBatch)> { +async fn send_batch(batch: &Batch, inserter: &Inserter) -> Result> { let first = batch.submissions.first().ok_or_else(|| { UnexpectedSnafu { violated: "cannot execute an empty flush".to_string(), @@ -73,23 +94,19 @@ async fn send_batch(batch: &Batch, inserter: &Inserter) -> Result<(TableInfoRef, .build() })?; let combined = combine_batches(batch.submissions.iter().map(|pending| &pending.batch))?; - let mut attribution = WriteAttribution::default(); - for pending in &batch.submissions { - attribution.extend(WriteAttribution::new( - pending.batch.num_rows(), - pending.ctx.write_progress().cloned(), - )); + let mut owners = Vec::with_capacity(batch.total_rows); + for (owner, pending) in batch.submissions.iter().enumerate() { + owners.extend(std::iter::repeat_n(owner, pending.batch.num_rows())); } - let affected_rows = inserter - .flush_bulk_batch( + inserter + .flush_bulk_batch_tracked( first.table_info.clone(), - combined.clone(), + combined, first.ctx.clone(), - attribution, + owners, + batch.submissions.len(), ) - .await?; - validate_affected_rows(batch.total_rows, affected_rows)?; - Ok((first.table_info.clone(), combined)) + .await } fn combine_batches<'a>(batches: impl IntoIterator) -> Result { @@ -110,15 +127,6 @@ fn combine_batches<'a>(batches: impl IntoIterator) -> Re concat_batches(&schema, batches).context(ComputeArrowSnafu) } -fn validate_affected_rows(expected: usize, actual: usize) -> Result<()> { - if expected != actual { - return UnexpectedSnafu { - violated: format!("batched write affected {actual} rows, expected {expected}; individual results cannot be attributed"), - }.fail(); - } - Ok(()) -} - #[cfg(test)] mod tests { use arrow::array::Int32Array; @@ -154,12 +162,4 @@ mod tests { let changed = RecordBatch::try_new(Arc::new(metadata), first.columns().to_vec()).unwrap(); assert!(combine_batches([&first, &changed]).is_err()); } - - #[test] - fn test_affected_rows_must_match_before_attribution() { - for (expected, actual, valid) in [(0, 0, true), (3, 3, true), (3, 2, false), (3, 4, false)] - { - assert_eq!(valid, validate_affected_rows(expected, actual).is_ok()); - } - } } diff --git a/src/servers/src/batcher/table/flow_notifier.rs b/src/servers/src/batcher/table/flow_notifier.rs index 306358bf509..6ecc8460857 100644 --- a/src/servers/src/batcher/table/flow_notifier.rs +++ b/src/servers/src/batcher/table/flow_notifier.rs @@ -56,18 +56,25 @@ impl FlowNotifier { Some(Self { notifier }) } + /// Enqueues one notification for batches from the same table. /// Never waits for queue capacity or delivery and never changes the write result. - pub fn notify(&self, table_info: TableInfoRef, batch: &RecordBatch) { + pub fn notify<'a>( + &self, + table_info: TableInfoRef, + batches: impl IntoIterator, + ) { let Some(timestamp_column) = table_info.meta.schema.timestamp_column() else { return; }; - let timestamps = match extract_timestamps(batch, ×tamp_column.name) { - Ok(timestamps) => timestamps, - Err(error) => { - error!(error; "Failed to extract flow notification timestamps, table_id: {}", table_info.table_id()); - return; + let mut timestamps = Vec::new(); + for batch in batches { + match extract_timestamps(batch, ×tamp_column.name) { + Ok(batch_timestamps) => timestamps.extend(batch_timestamps), + Err(error) => { + error!(error; "Failed to extract flow notification timestamps, table_id: {}", table_info.table_id()); + } } - }; + } if timestamps.is_empty() { return; } @@ -87,20 +94,220 @@ mod tests { use std::sync::Arc; use api::helper::ColumnDataTypeWrapper; - use api::v1::ColumnDataType; + use api::region::RegionResponse; + use api::v1::helper::{tag_column_schema, time_index_column_schema}; + use api::v1::region::{RegionRequest, bulk_insert_request, region_request}; + use api::v1::value::ValueData; + use api::v1::{ColumnDataType, Row, Rows, Value}; use arrow::array::{ - ArrayRef, TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray, - TimestampSecondArray, + ArrayRef, Int32Array, TimestampMicrosecondArray, TimestampMillisecondArray, + TimestampNanosecondArray, TimestampSecondArray, }; use arrow::record_batch::RecordBatch; + use catalog::memory::MemoryCatalogManager; + use common_grpc::flight::FlightDecoder; + use common_meta::error::Result as MetaResult; + use common_meta::peer::Peer; + use common_meta::test_util::{MockDatanodeHandler, MockDatanodeManager}; + use common_query::request::QueryRequest; + use common_recordbatch::SendableRecordBatchStream; use datatypes::data_type::ConcreteDataType; use datatypes::schema::{ColumnSchema, Schema}; - use operator::test_util::new_test_table_info; + use operator::insert::Inserter; + use operator::req_convert::insert::rows_to_record_batch; + use operator::test_util::{ + create_partition_rule_manager, new_test_table_info, prepare_mocked_backend, + }; + use session::context::QueryContext; + use tokio::sync::{Semaphore, oneshot}; use crate::batcher::flow_notifier::FlowNotifier as NotificationQueue; + use crate::batcher::table::batch::{Batch, flush_batch}; use crate::batcher::table::flow_notifier::{ FLOW_NOTIFICATION_DROPPED, FlowNotification, FlowNotifier, }; + use crate::batcher::table::metrics; + use crate::batcher::table::pending_batch::PendingBatch; + use crate::batcher::test_util::mock_table_flownode_cache; + + #[tokio::test] + async fn test_flush_coalesces_notifications_for_successful_submissions() { + #[derive(Clone)] + struct PartialHandler; + + #[async_trait::async_trait] + impl MockDatanodeHandler for PartialHandler { + async fn handle(&self, _: &Peer, request: RegionRequest) -> MetaResult { + let Some(region_request::Body::BulkInsert(request)) = request.body else { + panic!("expected bulk insert") + }; + let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = request.body else { + panic!("expected Arrow IPC") + }; + let batch = FlightDecoder::try_from_schema_bytes(&ipc.schema) + .unwrap() + .try_decode_record_batch(&ipc.data_header, &ipc.payload) + .unwrap(); + let values = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + if values.values().contains(&20) { + return common_meta::error::UnexpectedSnafu { + err_msg: "lost acknowledgement", + } + .fail(); + } + Ok(RegionResponse::new(batch.num_rows())) + } + + async fn handle_query( + &self, + _: &Peer, + _: QueryRequest, + ) -> MetaResult { + unreachable!() + } + } + + let partitions = create_partition_rule_manager(prepare_mocked_backend().await).await; + let inserter = Arc::new(Inserter::new( + MemoryCatalogManager::new(), + partitions, + Arc::new(MockDatanodeManager::new(PartialHandler)), + mock_table_flownode_cache(1, vec![]).await, + true, + )); + let table = Arc::new(new_test_table_info(1, "table_1", [1, 2, 3].into_iter())); + let cases = [ + ( + vec![vec![1], vec![2]], + vec![(1, false), (1, false)], + vec![60_000, 120_000], + true, + ), + ( + vec![vec![1, 20], vec![2], vec![3]], + vec![(1, true), (1, false), (1, false)], + vec![120_000, 180_000], + false, + ), + ( + vec![vec![20], vec![20]], + vec![(0, true), (0, true)], + vec![], + false, + ), + ( + vec![vec![1]; 100], + vec![(1, false); 100], + vec![60_000; 100], + true, + ), + ]; + let ctx = Arc::new(QueryContext::arc().with_confirmed_writes()); + let ingested = operator::metrics::DIST_INGEST_ROW_COUNT + .with_label_values(&[ctx.get_db_string().as_str()]); + for (inputs, expected_progress, expected_timestamps, successful_flush) in cases { + let permits = Arc::new(Semaphore::new(inputs.len())); + let mut submissions = Vec::new(); + let mut completions = Vec::new(); + let mut total_rows = 0; + for values in inputs { + let rows = Rows { + schema: vec![ + tag_column_schema("a", ColumnDataType::Int32), + time_index_column_schema("ts", ColumnDataType::TimestampMillisecond), + ], + rows: values + .into_iter() + .map(|value| Row { + values: vec![ + Value { + value_data: Some(ValueData::I32Value(value)), + }, + Value { + value_data: Some(ValueData::TimestampMillisecondValue( + i64::from(value) * 60_000, + )), + }, + ], + }) + .collect(), + }; + total_rows += rows.rows.len(); + let (response_tx, response_rx) = oneshot::channel(); + submissions.push(PendingBatch { + _completion: None, + table_info: table.clone(), + batch: rows_to_record_batch(&rows, &table).unwrap(), + ctx: ctx.clone(), + response_tx, + _permit: Arc::new(permits.clone().acquire_owned().await.unwrap()), + }); + completions.push(response_rx); + } + // Keep a one-slot queue undrained so per-submission notifications lose timestamps. + let (notifier, mut notifications) = + NotificationQueue::try_new(1, FLOW_NOTIFICATION_DROPPED.clone()).unwrap(); + let successes = metrics::FLUSH_TOTAL.get(); + let failures = metrics::FLUSH_FAILURES.get(); + let samples = metrics::FLUSH_ROWS.get_sample_count(); + let sum = metrics::FLUSH_ROWS.get_sample_sum(); + let dropped = metrics::FLUSH_DROPPED_ROWS.get(); + let written = ingested.get(); + flush_batch( + Batch { + submissions, + total_rows, + }, + inserter.clone(), + FlowNotifier { notifier }, + ) + .await; + + let expected_written: usize = expected_progress.iter().map(|(rows, _)| rows).sum(); + assert_eq!( + metrics::FLUSH_TOTAL.get(), + successes + u64::from(successful_flush) + ); + assert_eq!( + metrics::FLUSH_FAILURES.get(), + failures + u64::from(!successful_flush) + ); + assert_eq!( + metrics::FLUSH_ROWS.get_sample_count(), + samples + u64::from(successful_flush) + ); + assert_eq!( + metrics::FLUSH_ROWS.get_sample_sum(), + sum + if successful_flush { + total_rows as f64 + } else { + 0.0 + } + ); + assert_eq!( + metrics::FLUSH_DROPPED_ROWS.get(), + dropped + (total_rows - expected_written) as u64 + ); + assert_eq!(ingested.get(), written + expected_written as u64); + for (completion, (affected_rows, failed)) in + completions.into_iter().zip(expected_progress) + { + let progress = completion.await.unwrap(); + assert_eq!(progress.affected_rows, affected_rows); + assert_eq!(progress.error.is_some(), failed); + } + if !expected_timestamps.is_empty() { + let notification = notifications.try_recv().unwrap(); + assert_eq!(notification.table_id, 1); + assert_eq!(notification.timestamps, expected_timestamps); + } + assert!(notifications.try_recv().is_err()); + } + } #[test] fn test_full_and_closed_queues_do_not_block() { @@ -158,10 +365,11 @@ mod tests { .unwrap(); let (notifier, mut receiver) = NotificationQueue::try_new(1, FLOW_NOTIFICATION_DROPPED.clone()).unwrap(); - FlowNotifier { notifier }.notify(Arc::new(table), &batch); + FlowNotifier { notifier }.notify(Arc::new(table), [&batch, &batch]); let notification = receiver.try_recv().unwrap(); assert_eq!(notification.table_id, 1); - assert_eq!(notification.timestamps, vec![42]); + assert_eq!(notification.timestamps, vec![42, 42]); + assert!(receiver.try_recv().is_err()); } } } diff --git a/src/servers/src/batcher/table/pending_batch.rs b/src/servers/src/batcher/table/pending_batch.rs index 580407e6b5f..c4545f5992a 100644 --- a/src/servers/src/batcher/table/pending_batch.rs +++ b/src/servers/src/batcher/table/pending_batch.rs @@ -15,8 +15,11 @@ use std::sync::Arc; use arrow::record_batch::RecordBatch; +use common_error::ext::BoxedError; +use common_query::WriteProgress; use operator::error::Error; use session::context::QueryContextRef; +use snafu::IntoError; use table::metadata::TableInfoRef; use tokio::sync::{OwnedSemaphorePermit, oneshot}; @@ -25,7 +28,7 @@ pub(in crate::batcher::table) struct PendingBatch { pub table_info: TableInfoRef, pub batch: RecordBatch, pub ctx: QueryContextRef, - pub response_tx: oneshot::Sender>>, + pub response_tx: oneshot::Sender, pub _permit: Arc, pub _completion: Option, } @@ -35,7 +38,16 @@ pub(in crate::batcher::table) fn notify_batches( result: Result<(), Arc>, ) { for batch in batches { - let _ = batch.response_tx.send(result.clone()); + let progress = match &result { + Ok(()) => WriteProgress { + affected_rows: batch.batch.num_rows(), + error: None, + }, + Err(error) => WriteProgress::failed(BoxedError::new( + operator::error::BatchFlushSnafu.into_error(error.clone()), + )), + }; + let _ = batch.response_tx.send(progress); } } @@ -85,12 +97,11 @@ mod tests { notify_batches(batches, if fail { Err(error.clone()) } else { Ok(()) }); assert_eq!(3, semaphore.available_permits()); for receiver in receivers { - match receiver.await.unwrap() { - Ok(()) => assert!(!fail), - Err(actual) => { - assert!(fail); - assert!(Arc::ptr_eq(&error, &actual)); - } + let progress = receiver.await.unwrap(); + assert_eq!(progress.error.is_some(), fail); + assert_eq!(progress.affected_rows, if fail { 0 } else { 1 }); + if let Some(error) = progress.error { + assert!(error.to_string().contains("test flush failure")); } } } diff --git a/src/servers/src/error.rs b/src/servers/src/error.rs index 7421ac7bdf3..2e8f4cbefdc 100644 --- a/src/servers/src/error.rs +++ b/src/servers/src/error.rs @@ -15,7 +15,6 @@ use std::any::Any; use std::net::SocketAddr; use std::string::FromUtf8Error; -use std::sync::Arc; use axum::http::StatusCode as HttpStatusCode; use axum::response::{IntoResponse, Response}; @@ -735,9 +734,6 @@ pub enum Error { #[snafu(implicit)] location: Location, }, - - #[snafu(display("Failed to submit batch: {}", source))] - SubmitBatch { source: Arc }, } pub type Result = std::result::Result; @@ -876,7 +872,6 @@ impl ErrorExt for Error { Partition { source, .. } => source.status_code(), MetricEngine { source, .. } => source.status_code(), Operator { source, .. } => source.status_code(), - SubmitBatch { source, .. } => source.status_code(), } } @@ -911,7 +906,6 @@ impl ErrorExt for Error { Partition { source, .. } => source.retry_hint(), MetricEngine { source, .. } => source.retry_hint(), Operator { source, .. } => source.retry_hint(), - SubmitBatch { source, .. } => source.retry_hint(), MemoryLimitExceeded { source, .. } => source.retry_hint(), CollectRecordbatch { source, .. } => source.retry_hint(), @@ -1014,8 +1008,9 @@ pub fn status_code_to_http_status(status_code: &StatusCode) -> HttpStatusCode { #[cfg(test)] mod tests { - use common_error::GREPTIME_DB_HEADER_ERROR_CODE; use common_error::ext::PlainError; + use common_error::mock::MockError; + use common_error::{GREPTIME_DB_HEADER_ERROR_CODE, GREPTIME_DB_HEADER_ERROR_RETRY_HINT}; use super::*; @@ -1044,4 +1039,41 @@ mod tests { (StatusCode::RequestOutdated as u32).to_string() ); } + + #[tokio::test] + async fn test_storage_write_retry_hints_reach_tonic_status() { + for hint in [ + RetryHint::NonRetryable, + RetryHint::Retryable, + RetryHint::RetryAfterRouteRefresh, + ] { + let source = + BoxedError::new(MockError::new(StatusCode::RequestOutdated).with_retry_hint(hint)); + let completion = common_query::WriteCompletion::new(std::future::ready(Err(source))); + let error = Error::ExecuteGrpcQuery { + source: completion.wait().await.unwrap_err(), + location: Location::default(), + }; + let status: tonic::Status = error.into(); + assert_eq!(status.code(), tonic::Code::InvalidArgument); + assert_eq!( + status + .metadata() + .get(GREPTIME_DB_HEADER_ERROR_CODE) + .unwrap() + .to_str() + .unwrap(), + (StatusCode::RequestOutdated as u32).to_string() + ); + assert_eq!( + status + .metadata() + .get(GREPTIME_DB_HEADER_ERROR_RETRY_HINT) + .unwrap() + .to_str() + .unwrap(), + hint.as_str() + ); + } + } } diff --git a/src/servers/src/metrics.rs b/src/servers/src/metrics.rs index d928c60c741..963ab51b1e2 100644 --- a/src/servers/src/metrics.rs +++ b/src/servers/src/metrics.rs @@ -154,12 +154,12 @@ lazy_static! { .unwrap(); pub static ref FLUSH_TOTAL: IntCounter = register_int_counter!( "greptime_prom_store_flush_total", - "Total number of batch flushes" + "Total number of successful batch flushes" ) .unwrap(); pub static ref FLUSH_ROWS: Histogram = register_histogram!( "greptime_prom_store_flush_rows", - "Number of rows per flush", + "Number of rows per successful flush", vec![100.0, 1000.0, 10000.0, 50000.0, 100000.0, 500000.0] ) .unwrap(); diff --git a/src/session/src/context.rs b/src/session/src/context.rs index 03643331ddf..93890e7d535 100644 --- a/src/session/src/context.rs +++ b/src/session/src/context.rs @@ -99,6 +99,9 @@ pub struct QueryContext { /// Process-local storage acknowledgements; never accepted from a wire context. #[builder(setter(skip))] write_progress: Option>, + /// Preserve storage acknowledgements for protocol-level partial-write accounting. + #[builder(setter(skip))] + confirmed_writes: bool, /// Process id for managing on-going queries #[builder(default)] process_id: u32, @@ -301,6 +304,24 @@ impl QueryContext { self.write_progress = Some(progress); } + /// Forks this context to retain storage acknowledgements for partial write accounting. + /// + /// Insertion APIs returning `Output` may return `Ok` while carrying storage errors + /// in `Output.meta.write_completions`, even on synchronous write paths. Callers must + /// await `Output::write_progress()` and inspect both its acknowledged row count and + /// its `error` before treating the write as successful. + pub fn with_confirmed_writes(&self) -> Self { + let mut ctx = self.fork(); + ctx.confirmed_writes = true; + ctx + } + + /// Whether insertion outputs retain storage acknowledgements and may defer errors. + /// See [`Self::with_confirmed_writes`] for the caller's completion-handling contract. + pub fn confirmed_writes(&self) -> bool { + self.confirmed_writes + } + /// Forks this context with an independent snapshot of mutable session data. /// /// Unlike [`Clone`], changes to the schema, user, timezone, and other fields @@ -776,6 +797,7 @@ impl QueryContextBuilder { batching_enabled: self.batching_enabled.unwrap_or_default(), admitted_write: None, write_progress: None, + confirmed_writes: false, logical_batching_enabled: self.logical_batching_enabled.unwrap_or_default(), process_id: self.process_id.unwrap_or_default(), conn_info: self.conn_info.unwrap_or_default(), From 01b0076f3e7051618d46934f999eb863abf708c7 Mon Sep 17 00:00:00 2001 From: discord9 Date: Mon, 2 Nov 2026 04:15:19 +0000 Subject: [PATCH 151/239] fix(promql): collapse duplicate samples at one timestamp before range functions (#10069) * fix(promql): collapse duplicate samples at one timestamp before range functions A read path that does not deduplicate, such as an append-mode region, can hand PromQL two rows for one series at one timestamp. Every range function reads its samples from RangeManipulate, so the duplicate entries were treated as independent samples: count_over_time counted both, idelta and irate returned NULL because the last two entries shared a timestamp, and delta/increase/rate deflated because the average sampling interval divides the observed span by sample_count - 1. Prometheus sample identity is one sample per series per timestamp, so its range functions can never see this input. Restore that identity once, at the operator every range function shares, instead of special-casing each function's arithmetic. The collapse compares native ticks, never the millisecond payload, so distinct sub-millisecond samples that share a millisecond stay distinct as calculate_range documents and tests. Batches without duplicate timestamps take a single adjacency scan and are returned untouched; only when duplicates exist is a take built, keeping tag and value columns aligned. Keeping the last row of a group matches MergeMode::LastRow. The normalize test that fed three rows at one timestamp now uses distinct timestamps: it exercises histogram start-timestamp offsetting, and its repeated timestamp was incidental, not asserted behavior. Signed-off-by: discord9 * test(promql): cover duplicate samples end to end and keep the newer value The operator now keeps the first row of a tied timestamp group, which is the row a deduplicating read keeps: storage orders equal primary key and timestamp by sequence descending, and FlatDedupReader takes the first row of each group. The end-to-end case pins that choice by writing 100 and then 101 at the same timestamp: the window reports increase 6, while keeping the stale 100 would report 8. The sqlness case writes duplicate rows into an append_mode region, which is the read path that does not deduplicate, and asserts every range function plus the unaffected ordinary table in both the standalone and distributed environments. Signed-off-by: discord9 * docs(promql): qualify duplicate-timestamp conflict selection Signed-off-by: discord9 * perf(promql): avoid copying tags when collapsing duplicate samples Signed-off-by: discord9 --------- Signed-off-by: discord9 --- src/promql/src/extension_plan/normalize.rs | 10 +- .../src/extension_plan/range_manipulate.rs | 405 +++++++++++++++++- .../common/promql/duplicate_samples.result | 158 +++++++ .../common/promql/duplicate_samples.sql | 72 ++++ 4 files changed, 634 insertions(+), 11 deletions(-) create mode 100644 tests/cases/standalone/common/promql/duplicate_samples.result create mode 100644 tests/cases/standalone/common/promql/duplicate_samples.sql diff --git a/src/promql/src/extension_plan/normalize.rs b/src/promql/src/extension_plan/normalize.rs index b1ea6ca3e7c..5c076f48326 100644 --- a/src/promql/src/extension_plan/normalize.rs +++ b/src/promql/src/extension_plan/normalize.rs @@ -831,7 +831,7 @@ mod test { let batch = RecordBatch::try_new( schema.clone(), vec![ - Arc::new(TimestampMillisecondArray::from(vec![1_000; 3])), + Arc::new(TimestampMillisecondArray::from(vec![1_000, 1_001, 1_002])), histograms, ], ) @@ -843,11 +843,11 @@ mod test { let normalized = SeriesNormalize::new(1_000, TIME_INDEX_COLUMN, false, Vec::new(), logical_input); let range = RangeManipulate::new( - 2_000, - 2_000, + 2_002, + 2_002, 1, 1_000, - 1, + 3, TIME_INDEX_COLUMN.to_string(), vec!["value".to_string()], LogicalPlan::Extension(datafusion::logical_expr::Extension { @@ -909,7 +909,7 @@ mod test { .unwrap(); assert_eq!( timestamps.get(0).unwrap().to_data(), - TimestampMillisecondArray::from(vec![2_000; 3]).to_data() + TimestampMillisecondArray::from(vec![2_000, 2_001, 2_002]).to_data() ); } } diff --git a/src/promql/src/extension_plan/range_manipulate.rs b/src/promql/src/extension_plan/range_manipulate.rs index 38ecb7aa4db..7cc19e96aa7 100644 --- a/src/promql/src/extension_plan/range_manipulate.rs +++ b/src/promql/src/extension_plan/range_manipulate.rs @@ -682,9 +682,58 @@ impl RangeManipulateStream { // But they are not exactly the same, because we don't eager-evaluate on the data in this plan. // And the generated timestamp is not aligned to the step. It's expected to do later. pub fn manipulate(&self, input: RecordBatch) -> DataFusionResult> { + // Prometheus sample identity is one sample per series per timestamp, and this operator's + // input is one series per batch (the planner inserts a `SeriesDivide` for that). A read + // path that does not deduplicate can still deliver two rows for one timestamp, and every + // range function reads its samples from here, so collapse them once instead of + // special-casing each function. Comparing native ticks rather than the millisecond + // payload keeps distinct sub-millisecond samples apart. + // + // The first row of a tied group wins. A deduplicating scan also keeps the first row + // after ordering equal primary keys and timestamps by descending sequence, so the values + // agree when the scan presents those rows newest-first. This operator cannot establish + // that order itself: unordered scans and scans whose sequence values were overwritten + // can select a different tied value. The collapse fixes sample multiplicity, not global + // newest-write conflict resolution. + let (native_timestamps, _) = timestamp_array_to_primitive(input.column(self.time_index)) + .ok_or_else(|| { + DataFusionError::Execution("Time index column is not a timestamp".into()) + })?; + let native_values = native_timestamps.values().as_ref(); + let first_duplicate = if native_timestamps.null_count() == 0 { + native_values + .windows(2) + .position(|pair| pair[0] == pair[1]) + .map(|index| index + 1) + } else { + (1..input.num_rows()).find(|&i| { + !native_timestamps.is_null(i - 1) + && !native_timestamps.is_null(i) + && native_values[i - 1] == native_values[i] + }) + }; + let indices = first_duplicate.map(|first_duplicate| { + let mut indices = Vec::with_capacity(input.num_rows()); + indices.extend(0..first_duplicate); + indices.extend((first_duplicate + 1..input.num_rows()).filter(|&i| { + native_timestamps.is_null(i) + || native_timestamps.is_null(i - 1) + || native_values[i] != native_values[i - 1] + })); + Int64Array::from_iter_values(indices.into_iter().map(|index| index as i64)) + }); + // Only sample columns need collapsing. Tags below still come from row zero of the + // original batch, and the aligned time index is replaced below. In particular, do not + // build a RecordBatch with collapsed sample columns and uncollapsed tags: Arrow requires + // all columns to have the same length. + let ts_column = match indices.as_ref() { + Some(indices) => compute::take(input.column(self.time_index), indices, None)?, + None => input.column(self.time_index).clone(), + }; let mut other_columns = (0..input.columns().len()).collect::>(); + other_columns.remove(&self.time_index); // calculate the range - let (ranges, (start, end)) = self.calculate_range(&input)?; + let (ranges, (start, end)) = self.calculate_range_on_timestamps(&ts_column)?; // ignore this if all ranges are empty if ranges.iter().all(|(_, len)| *len == 0) { return Ok(None); @@ -694,9 +743,12 @@ impl RangeManipulateStream { let mut new_columns = input.columns().to_vec(); for index in self.field_columns.iter() { let _ = other_columns.remove(index); - let column = input.column(*index); + let column = match indices.as_ref() { + Some(indices) => compute::take(input.column(*index), indices, None)?, + None => input.column(*index).clone(), + }; let new_column = Arc::new( - RangeArray::from_ranges(column.clone(), ranges.clone()) + RangeArray::from_ranges(column, ranges.clone()) .map_err(|e| ArrowError::InvalidArgumentError(e.to_string()))? .into_dict(), ); @@ -707,7 +759,6 @@ impl RangeManipulateStream { // only when its own data type already is that ABI: the output schema fixes the // dictionary value type, so a millisecond column carrying a timezone must be // converted like any other native unit. - let ts_column = input.column(self.time_index); let timestamp_values: ArrayRef = if self.offset == 0 && matches!( ts_column.data_type(), @@ -718,7 +769,7 @@ impl RangeManipulateStream { // Shift in wide native precision before truncating toward zero, preserving // null validity. let scale = nanoseconds_per_native_tick(self.time_unit); - let (timestamps, _) = timestamp_array_to_primitive(ts_column).ok_or_else(|| { + let (timestamps, _) = timestamp_array_to_primitive(&ts_column).ok_or_else(|| { DataFusionError::Execution("Time index column is not a timestamp".into()) })?; // Single pass over the timestamps: an i128 intermediate holds any native @@ -796,12 +847,20 @@ impl RangeManipulateStream { /// Return values: /// - A vector of tuples where each tuple contains the start index and length of the range. /// - A tuple of the actual start/end timestamp used to calculate the range. + #[cfg(test)] #[allow(clippy::type_complexity)] fn calculate_range( &self, input: &RecordBatch, ) -> DataFusionResult<(Vec<(u32, u32)>, (i64, i64))> { - let ts_column = input.column(self.time_index); + self.calculate_range_on_timestamps(input.column(self.time_index)) + } + + #[allow(clippy::type_complexity)] + fn calculate_range_on_timestamps( + &self, + ts_column: &ArrayRef, + ) -> DataFusionResult<(Vec<(u32, u32)>, (i64, i64))> { let scale = nanoseconds_per_native_tick(self.time_unit); let (timestamps, _) = timestamp_array_to_primitive(ts_column).ok_or_else(|| { DataFusionError::Execution("Time index column is not a timestamp".into()) @@ -1170,6 +1229,340 @@ mod test { assert_eq!(result_literal, expected); } + #[tokio::test] + async fn duplicate_native_timestamps_collapse_in_real_plan() { + use datafusion::arrow::array::Float64Array; + + let schema = Arc::new(Schema::new(vec![ + Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new("value", DataType::Float64, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![ + 29_600, 29_600, 59_600, 59_600, + ])), + Arc::new(Float64Array::from(vec![100.0, 100.0, 102.0, 102.0])), + ], + ) + .unwrap(); + let plan = RangeManipulate::new( + 60_000, + 60_000, + 60_000, + 0, + 60_000, + TIME_INDEX_COLUMN.to_string(), + vec!["value".to_string()], + LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: schema.clone().to_dfschema_ref().unwrap(), + }), + ) + .unwrap(); + let input = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(), + ))); + let batches = datafusion::physical_plan::collect( + plan.to_execution_plan(input), + SessionContext::default().task_ctx(), + ) + .await + .unwrap(); + let out = &batches[0]; + let values = RangeArray::try_new( + out.column(1) + .as_any() + .downcast_ref::>() + .unwrap() + .clone(), + ) + .unwrap(); + let timestamps = RangeArray::try_new( + out.column(2) + .as_any() + .downcast_ref::>() + .unwrap() + .clone(), + ) + .unwrap(); + assert_eq!( + values + .get(0) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values(), + &[100.0, 102.0] + ); + assert_eq!( + timestamps + .get(0) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values(), + &[29_600, 59_600] + ); + } + + #[tokio::test] + async fn duplicate_timestamp_keeps_first_row_and_aligned_tags() { + use datafusion::arrow::array::{Float64Array, StringArray}; + + // Rows sharing a timestamp but carrying different values and tags show which row the + // collapse keeps, and that every column follows the same index instead of drifting apart. + let schema = Arc::new(Schema::new(vec![ + Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new("value", DataType::Float64, false), + Field::new("tag", DataType::Utf8, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![ + 5_000, 10_000, 10_000, 15_000, + ])), + Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0, 4.0])), + Arc::new(StringArray::from(vec!["a", "b", "c", "d"])), + ], + ) + .unwrap(); + let plan = RangeManipulate::new( + 20_000, + 20_000, + 20_000, + 0, + 20_000, + TIME_INDEX_COLUMN.to_string(), + vec!["value".to_string()], + LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: schema.clone().to_dfschema_ref().unwrap(), + }), + ) + .unwrap(); + let input = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(), + ))); + let batches = datafusion::physical_plan::collect( + plan.to_execution_plan(input), + SessionContext::default().task_ctx(), + ) + .await + .unwrap(); + let out = &batches[0]; + let values = RangeArray::try_new( + out.column(1) + .as_any() + .downcast_ref::>() + .unwrap() + .clone(), + ) + .unwrap(); + let timestamps = RangeArray::try_new( + out.column(3) + .as_any() + .downcast_ref::>() + .unwrap() + .clone(), + ) + .unwrap(); + // The first row of the tied pair is kept, matching what a deduplicating read returns, so + // 2.0 survives and 3.0 is dropped. + assert_eq!( + values + .get(0) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values(), + &[1.0, 2.0, 4.0] + ); + assert_eq!( + timestamps + .get(0) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values(), + &[5_000, 10_000, 15_000] + ); + // Tag columns are not range payloads, so the operator repeats the batch's first row for + // every output row: a batch holds one series, so every row's tag is that series' tag. + assert_eq!( + out.column(2) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .collect::>(), + vec![Some("a")] + ); + } + + #[test] + fn duplicate_sample_range_udfs_report_deduplicated_values() { + use crate::functions::{CountOverTime, Delta, IDelta, Increase, Rate}; + use datafusion::arrow::array::Float64Array; + use datafusion::physical_plan::ColumnarValue; + use datafusion_common::config::ConfigOptions; + use datafusion_expr::ScalarFunctionArgs; + let evaluate = |udf: datafusion::logical_expr::ScalarUDF, range_ms: Option| { + let ts = RangeArray::from_ranges( + Arc::new(TimestampMillisecondArray::from(vec![29_600, 59_600])), + [(0, 2)], + ) + .unwrap(); + let values = + RangeArray::from_ranges(Arc::new(Float64Array::from(vec![100.0, 102.0])), [(0, 2)]) + .unwrap(); + let mut args = vec![ + ColumnarValue::Array(Arc::new(ts.into_dict())), + ColumnarValue::Array(Arc::new(values.into_dict())), + ]; + let mut arg_fields = args + .iter() + .enumerate() + .map(|(i, arg)| { + Arc::new(Field::new( + if i == 0 { "ts" } else { "value" }, + arg.data_type(), + false, + )) + }) + .collect::>(); + if let Some(range_ms) = range_ms { + let eval_ts: ArrayRef = Arc::new(TimestampMillisecondArray::from(vec![60_000])); + args.push(ColumnarValue::Array(eval_ts)); + args.push(ColumnarValue::Scalar( + datafusion_common::ScalarValue::Int64(Some(range_ms)), + )); + arg_fields.push(Arc::new(Field::new("eval_ts", args[2].data_type(), false))); + arg_fields.push(Arc::new(Field::new("range", DataType::Int64, false))); + } + let result = udf + .invoke_with_args(ScalarFunctionArgs { + args, + arg_fields, + number_rows: 1, + return_field: Arc::new(Field::new("result", DataType::Float64, false)), + config_options: Arc::new(ConfigOptions::default()), + }) + .unwrap(); + let array = match result { + ColumnarValue::Array(a) => a, + ColumnarValue::Scalar(s) => s.to_array_of_size(1).unwrap(), + }; + array + .as_any() + .downcast_ref::() + .unwrap() + .value(0) + }; + assert_eq!(evaluate(CountOverTime::scalar_udf(), None), 2.0); + assert_eq!(evaluate(Delta::scalar_udf(), Some(60_000)), 4.0); + assert_eq!(evaluate(Increase::scalar_udf(), Some(60_000)), 4.0); + assert!((evaluate(Rate::scalar_udf(), Some(60_000)) - 4.0 / 60.0).abs() < 0.0001); + assert_eq!(evaluate(IDelta::::scalar_udf(), None), 2.0); + assert!((evaluate(IDelta::::scalar_udf(), None) - 2.0 / 30.0).abs() < 0.0001); + } + + #[tokio::test] + async fn distinct_submillisecond_timestamps_are_not_collapsed() { + let schema = Arc::new(Schema::new(vec![ + Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(TimeUnit::Microsecond, None), + false, + ), + Field::new("value", DataType::Float64, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMicrosecondArray::from(vec![1_000_001, 1_000_002])), + Arc::new(Float64Array::from(vec![10.0, 20.0])), + ], + ) + .unwrap(); + let plan = RangeManipulate::new( + 1_001, + 1_001, + 1, + 0, + 1, + TIME_INDEX_COLUMN.to_string(), + vec!["value".to_string()], + LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: schema.clone().to_dfschema_ref().unwrap(), + }), + ) + .unwrap(); + let input = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(), + ))); + let batches = datafusion::physical_plan::collect( + plan.to_execution_plan(input), + SessionContext::default().task_ctx(), + ) + .await + .unwrap(); + let output = &batches[0]; + let values = RangeArray::try_new( + output + .column(1) + .as_any() + .downcast_ref::>() + .unwrap() + .clone(), + ) + .unwrap(); + let timestamps = RangeArray::try_new( + output + .column(2) + .as_any() + .downcast_ref::>() + .unwrap() + .clone(), + ) + .unwrap(); + assert_eq!( + values + .get(0) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values(), + &[10.0, 20.0] + ); + assert_eq!( + timestamps + .get(0) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values(), + &[1_000, 1_000] + ); + } + #[tokio::test] async fn native_timestamps_preserve_range_membership_and_ms_payload() { for (unit, ticks_per_ms) in [ diff --git a/tests/cases/standalone/common/promql/duplicate_samples.result b/tests/cases/standalone/common/promql/duplicate_samples.result new file mode 100644 index 00000000000..e49ad1430fd --- /dev/null +++ b/tests/cases/standalone/common/promql/duplicate_samples.result @@ -0,0 +1,158 @@ +-- A duplicate row set for one series (same tags, same timestamp) is what an append-mode +-- region legitimately keeps, and what a retried write leaves behind. Prometheus sample +-- identity is (series, timestamp), so its range functions can never see two samples at one +-- timestamp; GreptimeDB must not let them see one either. +-- +-- Two samples per window, each written twice. The window is (0, 60s] with the first sample at +-- 29.6s and the last at 59.6s, so the pre-fix arithmetic is visible in the results below: +-- duplicates: n=4, average interval 10s, threshold 11s, leading gap clamped to 5s +-- factor = (30 + 5 + 0.4) / 30 = 1.18 => increase 2 * 1.18 = 2.36 +-- deduplicated: n=2, average interval 30s, threshold 33s, leading gap kept at 29.6s +-- factor = (30 + 29.6 + 0.4) / 30 = 2 => increase 2 * 2 = 4 +CREATE TABLE duplicate_samples ( + ts TIMESTAMP TIME INDEX, + greptime_value DOUBLE, + series STRING PRIMARY KEY +) engine=mito with('append_mode'='true'); + +Affected Rows: 0 + +INSERT INTO duplicate_samples VALUES + (29600, 100.0, 'a'), + (29600, 100.0, 'a'), + (59600, 102.0, 'a'), + (59600, 102.0, 'a'); + +Affected Rows: 4 + +-- Counting duplicate entries would report 4 samples for a window that holds two. +TQL EVAL (60, 60, '1m') count_over_time(duplicate_samples[1m]); + ++---------------------+-----------------------------------------------+--------+ +| ts | prom_count_over_time(ts_range,greptime_value) | series | ++---------------------+-----------------------------------------------+--------+ +| 1970-01-01T00:01:00 | 2.0 | a | ++---------------------+-----------------------------------------------+--------+ + +-- Both extrapolated functions must keep the same leading extrapolation. +TQL EVAL (60, 60, '1m') increase(duplicate_samples[1m]); + ++---------------------+--------------------------------------------------------+--------+ +| ts | prom_increase(ts_range,greptime_value,ts,Int64(60000)) | series | ++---------------------+--------------------------------------------------------+--------+ +| 1970-01-01T00:01:00 | 4.0 | a | ++---------------------+--------------------------------------------------------+--------+ + +TQL EVAL (60, 60, '1m') delta(duplicate_samples[1m]); + ++---------------------+-----------------------------------------------------+--------+ +| ts | prom_delta(ts_range,greptime_value,ts,Int64(60000)) | series | ++---------------------+-----------------------------------------------------+--------+ +| 1970-01-01T00:01:00 | 4.0 | a | ++---------------------+-----------------------------------------------------+--------+ + +TQL EVAL (60, 60, '1m') rate(duplicate_samples[1m]); + ++---------------------+----------------------------------------------------+--------+ +| ts | prom_rate(ts_range,greptime_value,ts,Int64(60000)) | series | ++---------------------+----------------------------------------------------+--------+ +| 1970-01-01T00:01:00 | 0.06666666666666667 | a | ++---------------------+----------------------------------------------------+--------+ + +-- The last two samples are 30s apart, not two entries of one timestamp. +TQL EVAL (60, 60, '1m') idelta(duplicate_samples[1m]); + ++---------------------+--------------------------------------+--------+ +| ts | prom_idelta(ts_range,greptime_value) | series | ++---------------------+--------------------------------------+--------+ +| 1970-01-01T00:01:00 | 2.0 | a | ++---------------------+--------------------------------------+--------+ + +TQL EVAL (60, 60, '1m') irate(duplicate_samples[1m]); + ++---------------------+-------------------------------------+--------+ +| ts | prom_irate(ts_range,greptime_value) | series | ++---------------------+-------------------------------------+--------+ +| 1970-01-01T00:01:00 | 0.06666666666666667 | a | ++---------------------+-------------------------------------+--------+ + +-- A duplicate row whose value differs from the first write is the same sample updated. One sample +-- remains for that timestamp, and it carries the value a deduplicating read would return, so the +-- later write wins: the window holds 101 and 104, giving an observed delta of 3 and an +-- extrapolation factor of 2, hence 6. Keeping the stale 100 would instead report 8. +CREATE TABLE duplicate_samples_conflict ( + ts TIMESTAMP TIME INDEX, + greptime_value DOUBLE, + series STRING PRIMARY KEY +) engine=mito with('append_mode'='true'); + +Affected Rows: 0 + +INSERT INTO duplicate_samples_conflict VALUES + (29600, 100.0, 'a'), + (29600, 101.0, 'a'), + (59600, 104.0, 'a'); + +Affected Rows: 3 + +TQL EVAL (60, 60, '1m') count_over_time(duplicate_samples_conflict[1m]); + ++---------------------+-----------------------------------------------+--------+ +| ts | prom_count_over_time(ts_range,greptime_value) | series | ++---------------------+-----------------------------------------------+--------+ +| 1970-01-01T00:01:00 | 2.0 | a | ++---------------------+-----------------------------------------------+--------+ + +TQL EVAL (60, 60, '1m') increase(duplicate_samples_conflict[1m]); + ++---------------------+--------------------------------------------------------+--------+ +| ts | prom_increase(ts_range,greptime_value,ts,Int64(60000)) | series | ++---------------------+--------------------------------------------------------+--------+ +| 1970-01-01T00:01:00 | 6.0 | a | ++---------------------+--------------------------------------------------------+--------+ + +-- A deduplicating region keeps its single row, so ordinary tables are unaffected. +CREATE TABLE duplicate_samples_deduped ( + ts TIMESTAMP TIME INDEX, + greptime_value DOUBLE, + series STRING PRIMARY KEY +); + +Affected Rows: 0 + +INSERT INTO duplicate_samples_deduped VALUES + (29600, 100.0, 'a'), + (29600, 100.0, 'a'), + (59600, 102.0, 'a'), + (59600, 102.0, 'a'); + +Affected Rows: 4 + +TQL EVAL (60, 60, '1m') count_over_time(duplicate_samples_deduped[1m]); + ++---------------------+-----------------------------------------------+--------+ +| ts | prom_count_over_time(ts_range,greptime_value) | series | ++---------------------+-----------------------------------------------+--------+ +| 1970-01-01T00:01:00 | 2.0 | a | ++---------------------+-----------------------------------------------+--------+ + +TQL EVAL (60, 60, '1m') increase(duplicate_samples_deduped[1m]); + ++---------------------+--------------------------------------------------------+--------+ +| ts | prom_increase(ts_range,greptime_value,ts,Int64(60000)) | series | ++---------------------+--------------------------------------------------------+--------+ +| 1970-01-01T00:01:00 | 4.0 | a | ++---------------------+--------------------------------------------------------+--------+ + +DROP TABLE duplicate_samples; + +Affected Rows: 0 + +DROP TABLE duplicate_samples_conflict; + +Affected Rows: 0 + +DROP TABLE duplicate_samples_deduped; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/promql/duplicate_samples.sql b/tests/cases/standalone/common/promql/duplicate_samples.sql new file mode 100644 index 00000000000..8eb42840e82 --- /dev/null +++ b/tests/cases/standalone/common/promql/duplicate_samples.sql @@ -0,0 +1,72 @@ +-- A duplicate row set for one series (same tags, same timestamp) is what an append-mode +-- region legitimately keeps, and what a retried write leaves behind. Prometheus sample +-- identity is (series, timestamp), so its range functions can never see two samples at one +-- timestamp; GreptimeDB must not let them see one either. +-- +-- Two samples per window, each written twice. The window is (0, 60s] with the first sample at +-- 29.6s and the last at 59.6s, so the pre-fix arithmetic is visible in the results below: +-- duplicates: n=4, average interval 10s, threshold 11s, leading gap clamped to 5s +-- factor = (30 + 5 + 0.4) / 30 = 1.18 => increase 2 * 1.18 = 2.36 +-- deduplicated: n=2, average interval 30s, threshold 33s, leading gap kept at 29.6s +-- factor = (30 + 29.6 + 0.4) / 30 = 2 => increase 2 * 2 = 4 +CREATE TABLE duplicate_samples ( + ts TIMESTAMP TIME INDEX, + greptime_value DOUBLE, + series STRING PRIMARY KEY +) engine=mito with('append_mode'='true'); + +INSERT INTO duplicate_samples VALUES + (29600, 100.0, 'a'), + (29600, 100.0, 'a'), + (59600, 102.0, 'a'), + (59600, 102.0, 'a'); + +-- Counting duplicate entries would report 4 samples for a window that holds two. +TQL EVAL (60, 60, '1m') count_over_time(duplicate_samples[1m]); + +-- Both extrapolated functions must keep the same leading extrapolation. +TQL EVAL (60, 60, '1m') increase(duplicate_samples[1m]); +TQL EVAL (60, 60, '1m') delta(duplicate_samples[1m]); +TQL EVAL (60, 60, '1m') rate(duplicate_samples[1m]); + +-- The last two samples are 30s apart, not two entries of one timestamp. +TQL EVAL (60, 60, '1m') idelta(duplicate_samples[1m]); +TQL EVAL (60, 60, '1m') irate(duplicate_samples[1m]); + +-- A duplicate row whose value differs from the first write is the same sample updated. One sample +-- remains for that timestamp, and it carries the value a deduplicating read would return, so the +-- later write wins: the window holds 101 and 104, giving an observed delta of 3 and an +-- extrapolation factor of 2, hence 6. Keeping the stale 100 would instead report 8. +CREATE TABLE duplicate_samples_conflict ( + ts TIMESTAMP TIME INDEX, + greptime_value DOUBLE, + series STRING PRIMARY KEY +) engine=mito with('append_mode'='true'); + +INSERT INTO duplicate_samples_conflict VALUES + (29600, 100.0, 'a'), + (29600, 101.0, 'a'), + (59600, 104.0, 'a'); + +TQL EVAL (60, 60, '1m') count_over_time(duplicate_samples_conflict[1m]); +TQL EVAL (60, 60, '1m') increase(duplicate_samples_conflict[1m]); + +-- A deduplicating region keeps its single row, so ordinary tables are unaffected. +CREATE TABLE duplicate_samples_deduped ( + ts TIMESTAMP TIME INDEX, + greptime_value DOUBLE, + series STRING PRIMARY KEY +); + +INSERT INTO duplicate_samples_deduped VALUES + (29600, 100.0, 'a'), + (29600, 100.0, 'a'), + (59600, 102.0, 'a'), + (59600, 102.0, 'a'); + +TQL EVAL (60, 60, '1m') count_over_time(duplicate_samples_deduped[1m]); +TQL EVAL (60, 60, '1m') increase(duplicate_samples_deduped[1m]); + +DROP TABLE duplicate_samples; +DROP TABLE duplicate_samples_conflict; +DROP TABLE duplicate_samples_deduped; From 94d71196cf63fa8783e9084b0aa96133a16bb569 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Mon, 2 Nov 2026 04:34:48 +0000 Subject: [PATCH 152/239] feat(mito2): track per-region flushed Parquet bytes (#10070) * feat(mito2): track per-region flushed Parquet bytes Signed-off-by: WenyXu * test(store-api): inline region statistic compatibility payload Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- src/metric-engine/src/utils.rs | 18 +++ src/mito2/src/engine/flush_test.rs | 103 ++++++++++++++++++ .../engine/object_store_wal_recovery_test.rs | 7 ++ src/mito2/src/engine/staging_test.rs | 17 +++ src/mito2/src/region.rs | 7 ++ src/mito2/src/worker/handle_flush.rs | 13 +++ src/store-api/src/region_engine.rs | 48 ++++++++ 7 files changed, 213 insertions(+) diff --git a/src/metric-engine/src/utils.rs b/src/metric-engine/src/utils.rs index 503b75a1cc5..64834a3f091 100644 --- a/src/metric-engine/src/utils.rs +++ b/src/metric-engine/src/utils.rs @@ -80,6 +80,8 @@ fn merge_region_statistic( metadata_manifest_version: metadata_stat.manifest.data_manifest_version(), }, written_bytes: metadata_stat.written_bytes + data_stat.written_bytes, + flushed_parquet_bytes: metadata_stat.flushed_parquet_bytes + + data_stat.flushed_parquet_bytes, query_cpu_time: data_stat.query_cpu_time, query_scanned_bytes: data_stat.query_scanned_bytes, data_topic_latest_entry_id: data_stat.data_topic_latest_entry_id, @@ -150,6 +152,22 @@ mod tests { assert_eq!(to_data_region_id(region_id), expected_region_id); } + #[test] + fn merge_region_statistic_sums_flushed_parquet_bytes() { + let metadata_stat = RegionStatistic { + flushed_parquet_bytes: 123, + ..Default::default() + }; + let data_stat = RegionStatistic { + flushed_parquet_bytes: 456, + ..Default::default() + }; + assert_eq!( + merge_region_statistic(&metadata_stat, &data_stat).flushed_parquet_bytes, + 579 + ); + } + #[test] fn merge_region_statistic_ignores_metadata_timestamps() { // Metadata rows are written at timestamp 0; folding them in would report diff --git a/src/mito2/src/engine/flush_test.rs b/src/mito2/src/engine/flush_test.rs index c7b4be14587..40b4cce0dbc 100644 --- a/src/mito2/src/engine/flush_test.rs +++ b/src/mito2/src/engine/flush_test.rs @@ -351,6 +351,109 @@ async fn test_manual_flush_with_format(flat_format: bool) { assert_eq!(expected, batches.pretty_print().unwrap()); } +#[tokio::test] +async fn test_flushed_parquet_bytes() { + for flat_format in [false, true] { + let mut env = TestEnv::with_prefix("flushed-parquet-bytes").await; + let engine = env + .create_engine(MitoConfig { + default_flat_format: flat_format, + schedule_compaction_after_edit: false, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new() + .insert_option("compaction.type", "twcs") + .insert_option("compaction.twcs.trigger_file_num", "100") + .build(); + let table_dir = request.table_dir.clone(); + let schema = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + let statistic = || engine.get_region_statistic(region_id).unwrap(); + assert_eq!(statistic().flushed_parquet_bytes, 0); + + let mut previous_bytes = 0; + for start in [0, 10] { + put_rows( + &engine, + region_id, + Rows { + schema: schema.clone(), + rows: build_rows(start, start + 3), + }, + ) + .await; + assert_eq!(statistic().flushed_parquet_bytes, previous_bytes); + flush_region(&engine, region_id, None).await; + let manifest = engine + .get_region(region_id) + .unwrap() + .manifest_ctx + .manifest() + .await; + let parquet_bytes: u64 = manifest.files.values().map(|file| file.file_size).sum(); + assert!(parquet_bytes > previous_bytes); + assert_eq!(statistic().flushed_parquet_bytes, parquet_bytes); + previous_bytes = parquet_bytes; + } + flush_region(&engine, region_id, None).await; + assert_eq!(statistic().flushed_parquet_bytes, previous_bytes); + + // A generic edit can add the same file metadata without producing a flush. + let file = engine + .get_region(region_id) + .unwrap() + .manifest_ctx + .manifest() + .await + .files + .values() + .next() + .unwrap() + .clone(); + engine + .edit_region( + region_id, + crate::manifest::action::RegionEdit { + files_to_add: vec![file], + files_to_remove: vec![], + timestamp_ms: None, + compaction_time_window: None, + flushed_entry_id: None, + flushed_sequence: None, + committed_sequence: None, + }, + ) + .await + .unwrap(); + assert_eq!(statistic().flushed_parquet_bytes, previous_bytes); + + engine + .handle_request( + region_id, + RegionRequest::Compact(RegionCompactRequest { + options: api::v1::region::compact_request::Options::StrictWindow( + api::v1::region::StrictWindow { + window_seconds: 3600, + }, + ), + ..Default::default() + }), + ) + .await + .unwrap(); + assert_eq!(statistic().flushed_parquet_bytes, previous_bytes); + assert_eq!(statistic().sst_num, 1); + reopen_region(&engine, region_id, table_dir, true, HashMap::new()).await; + assert!(statistic().sst_size > 0); + assert_eq!(statistic().flushed_parquet_bytes, 0); + } +} + #[tokio::test] async fn test_flush_engine() { test_flush_engine_with_format(false).await; diff --git a/src/mito2/src/engine/object_store_wal_recovery_test.rs b/src/mito2/src/engine/object_store_wal_recovery_test.rs index a9713d02fcc..0c028e067db 100644 --- a/src/mito2/src/engine/object_store_wal_recovery_test.rs +++ b/src/mito2/src/engine/object_store_wal_recovery_test.rs @@ -1862,6 +1862,13 @@ async fn test_flush_does_not_publish_a_frontier_for_a_lost_enqueued_backlog() { .await .manifest_flushed_entry_id ); + assert_eq!( + engine + .get_region_statistic(REGION_A) + .unwrap() + .flushed_parquet_bytes, + 0 + ); assert!(store.stop().await.is_err()); } diff --git a/src/mito2/src/engine/staging_test.rs b/src/mito2/src/engine/staging_test.rs index c29c4bbe3a3..08c272be6ab 100644 --- a/src/mito2/src/engine/staging_test.rs +++ b/src/mito2/src/engine/staging_test.rs @@ -812,6 +812,15 @@ async fn test_staging_exit_success_with_manifests_with_format(flat_format: bool) let manager = region.manifest_ctx.manifest_manager.read().await; let staging_manifest = manager.staging_manifest().unwrap(); assert_eq!(staging_manifest.files.len(), 3); + let parquet_bytes: u64 = staging_manifest + .files + .values() + .map(|file| file.file_size) + .sum(); + assert_eq!( + region.region_statistic().flushed_parquet_bytes, + parquet_bytes + ); assert_eq!( staging_manifest.metadata.partition_expr.as_ref().unwrap(), &partition_expr @@ -874,6 +883,9 @@ async fn test_staging_exit_success_with_manifests_with_format(flat_format: bool) assert_eq!(sst_entries.iter().filter(|e| e.visible).count(), 1); assert_eq!(sst_entries.iter().filter(|e| !e.visible).count(), 2); + let flushed_parquet_bytes = region.region_statistic().flushed_parquet_bytes; + assert!(flushed_parquet_bytes > 0); + // Streaming inspection must retain staging entries and their visibility. { use futures::StreamExt; @@ -991,6 +1003,11 @@ async fn test_staging_exit_success_with_manifests_with_format(flat_format: bool) "Region should no longer be in staging mode" ); + assert_eq!( + region.region_statistic().flushed_parquet_bytes, + flushed_parquet_bytes + ); + // Verify staging manifests have been cleared let staging_files_after = fs::read_dir(&staging_manifest_dir) .map(|entries| entries.collect::, _>>().unwrap_or_default()) diff --git a/src/mito2/src/region.rs b/src/mito2/src/region.rs index 819f32cf7e8..c13083fb9b7 100644 --- a/src/mito2/src/region.rs +++ b/src/mito2/src/region.rs @@ -198,6 +198,8 @@ pub struct MitoRegion { pub(crate) struct RegionStats { /// The total bytes written to the region. pub(crate) written_bytes: Arc, + /// Parquet bytes committed by flushes since this region instance was opened. + pub(crate) flushed_parquet_bytes: AtomicU64, /// The total query CPU time of the region in nanoseconds. pub(crate) query_cpu_time: Arc, /// The total scanned bytes of the region. @@ -208,6 +210,7 @@ impl RegionStats { pub(crate) fn new() -> Self { Self { written_bytes: Arc::new(AtomicU64::new(0)), + flushed_parquet_bytes: AtomicU64::new(0), query_cpu_time: Arc::new(AtomicU64::new(0)), query_scanned_bytes: Arc::new(AtomicU64::new(0)), } @@ -756,6 +759,10 @@ impl MitoRegion { data_topic_latest_entry_id: topic_latest_entry_id, metadata_topic_latest_entry_id: topic_latest_entry_id, written_bytes, + flushed_parquet_bytes: self + .region_stats + .flushed_parquet_bytes + .load(Ordering::Relaxed), query_cpu_time, query_scanned_bytes, min_timestamp: time_range.map(|(min, _)| min), diff --git a/src/mito2/src/worker/handle_flush.rs b/src/mito2/src/worker/handle_flush.rs index 4f26911077f..657559e12f3 100644 --- a/src/mito2/src/worker/handle_flush.rs +++ b/src/mito2/src/worker/handle_flush.rs @@ -417,6 +417,19 @@ impl RegionWorkerLoop { ); } + // The flush has committed its manifest, including in staging mode. Count only + // its Parquet output, before WAL cleanup (which can fail after the commit). + let flushed_parquet_bytes = request + .edit + .files_to_add + .iter() + .map(|file| file.file_size) + .sum(); + region + .region_stats + .flushed_parquet_bytes + .fetch_add(flushed_parquet_bytes, Ordering::Relaxed); + region.update_flush_millis(); // Update topic latest entry id as soon as possible after flush to make sure the prunable entry id is updated timely, // which is important for remote WAL pruning. diff --git a/src/store-api/src/region_engine.rs b/src/store-api/src/region_engine.rs index 002ffb1691c..5d7af812685 100644 --- a/src/store-api/src/region_engine.rs +++ b/src/store-api/src/region_engine.rs @@ -600,6 +600,12 @@ pub struct RegionStatistic { #[serde(default)] /// The total bytes written of the region since region opened. pub written_bytes: u64, + /// Total Parquet bytes committed by flushes since the region opened. + /// + /// Includes staging flushes, but excludes index files, compaction and imported SSTs. + /// Resets when the region instance is recreated; manifest replay does not increment it. + #[serde(default)] + pub flushed_parquet_bytes: u64, /// The total query CPU time of the region since region opened. /// /// Unit: nanoseconds. @@ -775,6 +781,48 @@ mod tests { use super::*; + #[test] + fn region_statistic_flush_bytes_compatibility() { + let legacy_bytes = br#"{ + "num_rows": 1, + "memtable_size": 2, + "wal_size": 3, + "manifest_size": 4, + "sst_size": 5, + "sst_num": 6, + "index_size": 7, + "manifest": { + "Mito": { + "manifest_version": 8, + "flushed_entry_id": 9, + "file_removed_cnt": 10 + } + }, + "written_bytes": 11 + }"#; + let mut statistic = RegionStatistic::deserialize_from_slice(legacy_bytes).unwrap(); + assert_eq!(statistic.flushed_parquet_bytes, 0); + statistic.flushed_parquet_bytes = 1234; + let encoded = statistic.serialize_to_vec().unwrap(); + let decoded = RegionStatistic::deserialize_from_slice(&encoded).unwrap(); + assert_eq!(decoded.flushed_parquet_bytes, 1234); + let legacy: serde_json::Value = serde_json::from_slice(legacy_bytes).unwrap(); + let mut current: serde_json::Value = serde_json::from_slice(&encoded).unwrap(); + for (key, value) in legacy.as_object().unwrap() { + assert_eq!(¤t[key], value); + } + current + .as_object_mut() + .unwrap() + .remove("flushed_parquet_bytes"); + assert_eq!( + serde_json::from_value::(current) + .unwrap() + .flushed_parquet_bytes, + 0 + ); + } + #[test] fn region_statistic_deserializes_without_query_stats() { let statistic: RegionStatistic = serde_json::from_value(json!({ From a09ebc70e6c95f067d9cbde2ed17b7f656e5ad62 Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Mon, 2 Nov 2026 08:57:05 +0000 Subject: [PATCH 153/239] feat(postgres): OTLP binary payload ingest via COPY FROM STDIN (metrics, traces, logs) (#10050) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit feat(postgres): OTLP binary payload ingest via COPY FROM STDIN COPY FROM STDIN WITH (FORMAT otlp [, signal '...'] [, ]) ingests one binary OTLP protobuf message over the postgres wire protocol — the first binary COPY format (wire format code 1) — for all three signals: - signal 'metrics' (default): ExportMetricsServiceRequest, completed through OpenTelemetryProtocolHandler::metrics with the same context preparation as the HTTP route (OtlpMetricCtx from the options); rows are accepted data points. - signal 'traces': ExportTraceServiceRequest through ::traces, with the trace table name and pipeline name/version/params from the options (an explicit pipeline name is required, exactly like the HTTP traces endpoint); rows are accepted spans. - signal 'logs': ExportLogsServiceRequest through ::logs, with the log table name, pipeline options and extract keys from the options, and the same OtlpLogDirect default as the HTTP route; rows are affected rows. The WITH options carry what the HTTP endpoints take as request headers, reusing the HTTP path's own parsing helpers (truthy, SelectInfo, GreptimePipelineParams, PipelineWay) so the two routes cannot drift; values already provided by the postgres session (database, user, permissions) are not repeated. Option semantics match the headers: promote_all_resource_attrs / resource_attrs (';'-separated include or exclude list) / promote_scope_attrs / translation_strategy for metrics; table_name / pipeline_name / pipeline_version / pipeline_params / extract_keys for traces and logs. Signal-inapplicable options are rejected. Like the influxdb format: table-less (a named table is rejected), the whole payload is buffered until CopyDone exactly like one OTLP/HTTP request body and bounded by the shared 64 MiB request limit, decoding uses the same protobuf types, and the server must be wired — PostgresServer::with_otlp_handler(handler, with_metric_engine) carries the same metric-engine flag the HTTP wiring passes from its options. The pipeline handler is derived from the protocol handler via trait upcasting, as the HTTP routes do. e2e: chunked binary metrics ingest; traces with table_name and pipeline_name options; logs on the default table; metrics promote_all_resource_attrs making a resource attribute queryable; invalid-payload failure with connection reuse. Signed-off-by: Ning Sun --- src/frontend/src/server.rs | 2 + src/servers/src/postgres.rs | 16 +- src/servers/src/postgres/copy_in/codec.rs | 7 + src/servers/src/postgres/copy_in/mod.rs | 231 +++++++++++- src/servers/src/postgres/copy_in/otlp.rs | 307 +++++++++++++++ src/servers/src/postgres/server.rs | 22 ++ tests-integration/src/test_util.rs | 9 +- .../src/tests/postgres_copy_in_test.rs | 354 ++++++++++++++++++ 8 files changed, 935 insertions(+), 13 deletions(-) create mode 100644 src/servers/src/postgres/copy_in/otlp.rs diff --git a/src/frontend/src/server.rs b/src/frontend/src/server.rs index d84f01fce95..2b79bddb3cd 100644 --- a/src/frontend/src/server.rs +++ b/src/frontend/src/server.rs @@ -409,6 +409,7 @@ where if opts.postgres.enable { // Init PosgresSQL Server + let with_metric_engine = effective_prom_store_options(&opts).with_metric_engine; let opts = &opts.postgres; let pg_addr = parse_addr(&opts.addr)?; @@ -430,6 +431,7 @@ where Some(self.instance.process_manager().clone()), ) .with_influxdb_handler(instance.clone()) + .with_otlp_handler(instance.clone(), with_metric_engine) .with_batching_enabled(postgres_batching), ) as Box; diff --git a/src/servers/src/postgres.rs b/src/servers/src/postgres.rs index fbef2e9ac4c..1d46ab9c010 100644 --- a/src/servers/src/postgres.rs +++ b/src/servers/src/postgres.rs @@ -44,7 +44,9 @@ use session::context::Channel; use self::auth_handler::{PgUserProviderSource, ZeroPidSecretKeyGenerator}; use self::handler::DefaultQueryParser; use crate::query_handler::sql::ServerSqlQueryHandlerRef; -use crate::query_handler::{CopyInHandlerRef, InfluxdbLineProtocolHandlerRef}; +use crate::query_handler::{ + CopyInHandlerRef, InfluxdbLineProtocolHandlerRef, OpenTelemetryProtocolHandlerRef, +}; pub(crate) struct GreptimeDBStartupParameters { version: String, @@ -86,6 +88,12 @@ pub struct PostgresServerHandlerInner { /// Absent on servers not wired for it; such connections reject the /// influxdb format when a copy starts. influxdb_handler: Option, + /// Executes the OTLP payload flushed by a table-less + /// `COPY FROM STDIN WITH (FORMAT otlp)` — the same handler the + /// OTLP/HTTP endpoints use — plus the with-metric-engine flag the + /// HTTP wiring passes from its options. + otlp_handler: Option, + otlp_with_metric_engine: bool, force_tls: bool, auth_driver: AuthDriver, UserInfoRef>, /// SCRAM selected; failures get the uniform metric label. @@ -106,6 +114,10 @@ pub(crate) struct MakePostgresServerHandler { copy_in_handler: CopyInHandlerRef, #[builder(default)] influxdb_handler: Option, + #[builder(default)] + otlp_handler: Option, + #[builder(default)] + otlp_with_metric_engine: bool, user_provider: Option, #[builder(default = "Arc::new(GreptimeDBStartupParameters::new())")] param_provider: Arc, @@ -157,6 +169,8 @@ impl MakePostgresServerHandler { .map(|p| Arc::new(PgUserProviderSource::new(p))), copy_in_handler: self.copy_in_handler.clone(), influxdb_handler: self.influxdb_handler.clone(), + otlp_handler: self.otlp_handler.clone(), + otlp_with_metric_engine: self.otlp_with_metric_engine, force_tls: self.force_tls, auth_driver: AuthDriver::new(self.param_provider.clone()) .with_pid_secret_key_generator(Arc::new(ZeroPidSecretKeyGenerator)), diff --git a/src/servers/src/postgres/copy_in/codec.rs b/src/servers/src/postgres/copy_in/codec.rs index a7672c619ef..9918061c273 100644 --- a/src/servers/src/postgres/copy_in/codec.rs +++ b/src/servers/src/postgres/copy_in/codec.rs @@ -55,6 +55,7 @@ use pgwire::error::PgWireResult; use crate::postgres::copy_in::csv::CsvCodec; use crate::postgres::copy_in::influxdb::InfluxdbOptions; +use crate::postgres::copy_in::otlp::OtlpOptions; /// A resolved COPY FROM STDIN format with its validated options. #[derive(Debug, Clone)] @@ -67,6 +68,12 @@ pub(super) enum CopyFormat { /// pipeline the HTTP write endpoints use (one COPY statement is one /// ingestion request; streaming ingest is not supported). Influxdb(InfluxdbOptions), + /// Table-less, **binary** OTLP: one `ExportMetricsServiceRequest`, + /// `ExportTraceServiceRequest` or `ExportLogsServiceRequest` + /// protobuf message, buffered whole and completed through the + /// matching OTLP pipeline the OTLP/HTTP endpoints use. Advertises + /// wire format code 1. + Otlp(OtlpOptions), } /// A parsed, insertable row: one gRPC value per targeted column, in COPY diff --git a/src/servers/src/postgres/copy_in/mod.rs b/src/servers/src/postgres/copy_in/mod.rs index 92528d4ab2f..064b65b8bb5 100644 --- a/src/servers/src/postgres/copy_in/mod.rs +++ b/src/servers/src/postgres/copy_in/mod.rs @@ -35,12 +35,14 @@ mod codec; mod csv; mod influxdb; mod options; +mod otlp; mod parse; use std::fmt::Debug; use std::sync::Arc; use crate::influxdb::InfluxdbRequest; +use crate::query_handler::PipelineHandlerRef; use api::helper::ColumnDataTypeWrapper; use api::v1::column_def::options_from_column_schema; use api::v1::{ @@ -48,6 +50,7 @@ use api::v1::{ }; use async_trait::async_trait; use codec::{CopyFormat, RecordParser}; +use common_catalog::consts::{TRACE_TABLE_NAME, TRACE_TABLE_NAME_SESSION_KEY}; use common_grpc::precision::Precision; use common_query::OutputData; use common_telemetry::{info, warn}; @@ -57,12 +60,15 @@ use datatypes::data_type::DataType; use datatypes::schema::{ColumnSchema, SchemaRef}; use futures::{Sink, SinkExt, stream}; use influxdb::InfluxdbOptions; +use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest; +use otlp::{OtlpOptions, OtlpSignal}; use pgwire::api::copy::CopyHandler; use pgwire::api::results::{CopyResponse, Response, Tag}; use pgwire::api::{ClientInfo, PgWireConnectionState}; use pgwire::error::{ErrorInfo, PgWireError, PgWireResult}; use pgwire::messages::PgWireBackendMessage; use pgwire::messages::copy::{CopyData, CopyDone, CopyFail}; +use prost::Message; use session::context::QueryContextRef; use table::metadata::TableInfo; @@ -84,11 +90,10 @@ use crate::postgres::utils::convert_err; /// the per-batch fixed costs negligible. pub(crate) const COPY_IN_FLUSH_ROWS: usize = 32768; -/// Upper bound of an influxdb COPY FROM STDIN payload. One COPY -/// statement is one bounded ingestion request (the whole payload is -/// buffered like one HTTP write body); a larger import is multiple COPY -/// statements. Streaming ingest may arrive later, with framing owned by -/// the shared line-protocol parser rather than the COPY layer. +/// Upper bound of a table-less COPY FROM STDIN payload (influxdb line +/// protocol, OTLP). One COPY statement is one bounded ingestion request +/// (the whole payload is buffered like one HTTP request body); a larger +/// import is multiple COPY statements. pub(crate) const COPY_IN_MAX_PAYLOAD_BYTES: usize = 64 * 1024 * 1024; /// Builds a format from the format-neutral option set. @@ -102,6 +107,8 @@ const REGISTERED_FORMATS: &[(&str, FormatBuilder)] = &[ ("csv", csv::build_csv_codec), ("influxdb", influxdb::build_influxdb_codec), ("influxdb_line", influxdb::build_influxdb_codec), + ("otlp", otlp::build_otlp_codec), + ("otlp_metrics", otlp::build_otlp_codec), ("text", build_unimplemented_text_codec), ("txt", build_unimplemented_text_codec), ("binary", build_unimplemented_binary_codec), @@ -268,6 +275,10 @@ enum Payload { data: Vec, precision: Option, }, + /// The whole buffered OTLP payload with its validated options, + /// decoded and completed as one request at `CopyDone` through the + /// matching OTLP pipeline. + Otlp { data: Vec, options: OtlpOptions }, } impl CopyInState { @@ -327,6 +338,22 @@ impl CopyInState { } } + /// State for the OTLP format: the whole binary payload is buffered + /// and completed once at `CopyDone` through the matching OTLP + /// pipeline. Binary: the advertised wire format code is 1. + fn new_otlp(options: &OtlpOptions, query_ctx: QueryContextRef) -> Self { + Self { + query_ctx, + payload: Payload::Otlp { + data: Vec::new(), + options: options.clone(), + }, + codec_name: "otlp", + codec_format_code: 1, + rows_written: 0, + } + } + /// The wire format code advertised to the client. pub(crate) fn format_code(&self) -> i8 { self.codec_format_code @@ -337,7 +364,7 @@ impl CopyInState { fn feed(&mut self, data: &[u8]) -> PgWireResult<()> { match &mut self.payload { Payload::Table { parser, rows, .. } => rows.extend(parser.feed(data)?), - Payload::Influxdb { data: buffered, .. } => { + Payload::Influxdb { data: buffered, .. } | Payload::Otlp { data: buffered, .. } => { if buffered.len() + data.len() > COPY_IN_MAX_PAYLOAD_BYTES { return Err(bad_copy_data(format!( "COPY FROM STDIN payload exceeds the {} limit; split the import into multiple COPY statements", @@ -356,7 +383,7 @@ impl CopyInState { match &mut self.payload { Payload::Table { parser, rows, .. } => rows.extend(parser.finish()?), // The whole payload is one request body; nothing to finalize. - Payload::Influxdb { .. } => {} + Payload::Influxdb { .. } | Payload::Otlp { .. } => {} } Ok(()) } @@ -365,7 +392,7 @@ impl CopyInState { match &self.payload { Payload::Table { rows, .. } => rows.len() >= COPY_IN_FLUSH_ROWS, // Completed once at CopyDone, not mid-stream. - Payload::Influxdb { .. } => false, + Payload::Influxdb { .. } | Payload::Otlp { .. } => false, } } @@ -395,6 +422,17 @@ impl CopyInState { }) } + /// Takes the buffered OTLP payload and its options, if any. + fn take_otlp(&mut self) -> Option<(OtlpOptions, Vec)> { + let Payload::Otlp { options, data } = &mut self.payload else { + return None; + }; + if data.is_empty() { + return None; + } + Some((options.clone(), std::mem::take(data))) + } + /// Takes the buffered payload as one line-protocol request body, if /// any — the same bytes an HTTP write of the same records would /// carry. @@ -558,11 +596,26 @@ impl PostgresServerHandlerInner { } CopyInState::new_influxdb(options, query_ctx) } + // OTLP forbids a target table too — the payload's resource + // metrics route to their own tables — and completes through + // the OTLP metrics pipeline; the server must be wired. + CopyFormat::Otlp(options) => { + if !stmt.table.is_empty() { + return Err(unsupported(format!( + "format otlp routes each record to its own table; \"COPY {} FROM STDIN\" must not name one", + stmt.table.join(".") + ))); + } + if self.otlp_handler.is_none() { + return Err(unsupported("this server is not wired for the otlp format")); + } + CopyInState::new_otlp(options, query_ctx) + } }; let num_columns = match &state.payload { Payload::Table { grpc_schema, .. } => grpc_schema.len(), // No fixed column set: the format routes per record. - Payload::Influxdb { .. } => 0, + Payload::Influxdb { .. } | Payload::Otlp { .. } => 0, }; let format_code = state.format_code(); @@ -585,6 +638,14 @@ impl PostgresServerHandlerInner { // timestamp alignment and inserts are identical to an HTTP // write of the same bytes. The context stays this connection's // postgres context. + if let Some((options, bytes)) = state.take_otlp() { + let Some(otlp_handler) = &self.otlp_handler else { + return Err(unsupported("this server is not wired for the otlp format")); + }; + return self + .flush_otlp(otlp_handler, options, bytes, query_ctx) + .await; + } if let Some(request) = state.take_lines()? { let Some(influxdb_handler) = &self.influxdb_handler else { return Err(unsupported( @@ -616,6 +677,158 @@ impl PostgresServerHandlerInner { } } } + + /// Completes an OTLP copy through the pipeline the matching + /// OTLP/HTTP endpoint uses — the payload is decoded with the same + /// protobuf types and handed to the same handler methods with the + /// same context preparation, so both ingestion routes are + /// identical. The context stays this connection's postgres context + /// apart from the signal-specific preparation below. + async fn flush_otlp( + &self, + otlp_handler: &crate::query_handler::OpenTelemetryProtocolHandlerRef, + options: OtlpOptions, + bytes: Vec, + query_ctx: QueryContextRef, + ) -> PgWireResult { + use opentelemetry_proto::tonic::collector::logs::v1::ExportLogsServiceRequest; + use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest; + use pipeline::{GreptimePipelineParams, PipelineWay, SelectInfo}; + use session::protocol_ctx::{MetricType, OtlpMetricCtx, ProtocolCtx}; + + // The HTTP routes derive the pipeline handler from the same + // protocol handler via trait upcasting; so does the copy. + let pipeline_handler: PipelineHandlerRef = otlp_handler.clone(); + match options.signal { + OtlpSignal::Metrics { + promote_all_resource_attrs, + resource_attrs, + promote_scope_attrs, + translation_strategy, + } => { + // Same context preparation as the HTTP metrics route. + let mut ctx = (*query_ctx).clone(); + ctx.set_protocol_ctx(ProtocolCtx::OtlpMetric(OtlpMetricCtx { + promote_all_resource_attrs, + resource_attrs, + promote_scope_attrs, + with_metric_engine: self.otlp_with_metric_engine, + // Set by the frontend from its config. + resource_info: false, + metric_type: MetricType::Init, + metric_translation_strategy: translation_strategy, + })); + let request = ExportMetricsServiceRequest::decode(bytes.as_slice()) + .map_err(|e| bad_copy_data(format!("invalid OTLP payload: {e}")))?; + let outcome = otlp_handler + .metrics(request, Arc::new(ctx)) + .await + .map_err(convert_err)?; + // Same failure semantics as the HTTP route: a payload + // whose points were all rejected is an error, not a + // successful `COPY 0`. + if outcome.accepted_data_points == 0 && outcome.rejected_data_points > 0 { + return Err(all_rejected_error( + "data points", + outcome.rejected_data_points, + outcome.error_message.as_deref(), + )); + } + // Rows are the accepted data points. + Ok(outcome.accepted_data_points.max(0) as u64) + } + OtlpSignal::Traces { + table_name, + pipeline_name, + pipeline_version, + pipeline_params, + } => { + let table_name = table_name.unwrap_or_else(|| TRACE_TABLE_NAME.to_string()); + // Same context preparation as the HTTP traces route. + let mut ctx = (*query_ctx).clone(); + ctx.set_extension(TRACE_TABLE_NAME_SESSION_KEY, &table_name); + let pipeline = PipelineWay::from_name_and_default( + pipeline_name.as_deref(), + pipeline_version.as_deref(), + None, + ) + .map_err(convert_err)?; + let request = ExportTraceServiceRequest::decode(bytes.as_slice()) + .map_err(|e| bad_copy_data(format!("invalid OTLP payload: {e}")))?; + let outcome = otlp_handler + .traces( + pipeline_handler, + request, + pipeline, + GreptimePipelineParams::from_params(pipeline_params.as_deref()), + table_name, + Arc::new(ctx), + ) + .await + .map_err(convert_err)?; + // Same failure semantics as the HTTP route. + if outcome.accepted_spans == 0 && outcome.rejected_spans > 0 { + return Err(all_rejected_error( + "spans", + outcome.rejected_spans as i64, + outcome.error_message.as_deref(), + )); + } + // Rows are the accepted spans. + Ok(outcome.accepted_spans as u64) + } + OtlpSignal::Logs { + table_name, + pipeline_name, + pipeline_version, + pipeline_params, + extract_keys, + } => { + let table_name = table_name.unwrap_or_else(|| "opentelemetry_logs".to_string()); + // Same select semantics as the HTTP extract-keys header. + let select_info = match extract_keys.as_deref() { + Some(keys) if !keys.is_empty() => SelectInfo::from(keys.to_string()), + _ => Default::default(), + }; + let pipeline = PipelineWay::from_name_and_default( + pipeline_name.as_deref(), + pipeline_version.as_deref(), + Some(PipelineWay::OtlpLogDirect(Box::new(select_info))), + ) + .map_err(convert_err)?; + let request = ExportLogsServiceRequest::decode(bytes.as_slice()) + .map_err(|e| bad_copy_data(format!("invalid OTLP payload: {e}")))?; + let outputs = otlp_handler + .logs( + pipeline_handler, + request, + pipeline, + GreptimePipelineParams::from_params(pipeline_params.as_deref()), + table_name, + query_ctx, + ) + .await + .map_err(convert_err)?; + // Rows are the affected rows across the outputs. + Ok(outputs + .iter() + .map(|o| match &o.data { + OutputData::AffectedRows(rows) => *rows as u64, + _ => 0, + }) + .sum()) + } + } + } +} + +/// The error for an OTLP payload whose points/spans were all rejected, +/// carrying the pipeline's reason when it provided one. +fn all_rejected_error(kind: &str, rejected: i64, reason: Option<&str>) -> PgWireError { + bad_copy_data(match reason { + Some(reason) => format!("OTLP ingestion rejected all {rejected} {kind}: {reason}"), + None => format!("OTLP ingestion rejected all {rejected} {kind}"), + }) } fn adjust_query_context( diff --git a/src/servers/src/postgres/copy_in/otlp.rs b/src/servers/src/postgres/copy_in/otlp.rs new file mode 100644 index 00000000000..df63b54562c --- /dev/null +++ b/src/servers/src/postgres/copy_in/otlp.rs @@ -0,0 +1,307 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! The `otlp` format for COPY FROM STDIN: one binary OTLP payload for +//! any of the three signals — `ExportMetricsServiceRequest`, +//! `ExportTraceServiceRequest` or `ExportLogsServiceRequest`. +//! +//! `COPY FROM STDIN WITH (FORMAT otlp [, signal 'traces'] [, ])` +//! is a GreptimeDB-specific, **binary** format (wire format code 1). +//! Like the influxdb format it is table-less and does not decode here: +//! the whole payload is buffered until `CopyDone` — exactly like an +//! OTLP/HTTP request body — then decoded and completed through the +//! frontend's `OpenTelemetryProtocolHandler`, the same path the +//! `/v1/otlp/v1/{metrics,traces,logs}` endpoints use, under the +//! connection's regular postgres query context. +//! +//! The `WITH` options carry what the HTTP endpoints take as request +//! headers; values already provided by the postgres session (database, +//! user, permissions) are not repeated. Option semantics, defaults and +//! parsing helpers are shared with the HTTP extractors (`truthy`, +//! `SelectInfo`, `GreptimePipelineParams`, `PipelineWay`), so the two +//! ingestion routes cannot drift. +//! +//! One COPY statement is one bounded ingestion request (see +//! [`crate::postgres::copy_in::COPY_IN_MAX_PAYLOAD_BYTES`]). + +use ahash::HashSet; + +use pgwire::error::PgWireResult; +use pipeline::truthy; +use session::protocol_ctx::OtlpMetricTranslationStrategy; + +use crate::postgres::copy_in::bad_copy_data; +use crate::postgres::copy_in::codec::CopyFormat; +use crate::postgres::copy_in::options::{CopyOptionSet, CopyOptionValue}; + +/// The format's validated options: the signal and its header-equivalent +/// configuration, carried to the flush that completes the copy. +#[derive(Debug, Clone)] +pub(super) struct OtlpOptions { + pub(super) signal: OtlpSignal, +} + +#[derive(Debug, Clone)] +pub(super) enum OtlpSignal { + /// Metric options mirroring the `x-greptime-otlp-metric-*` headers. + Metrics { + /// Persist all resource attributes; `resource_attrs` then acts + /// as an exclude list, otherwise as an include list. + promote_all_resource_attrs: bool, + resource_attrs: HashSet, + promote_scope_attrs: bool, + translation_strategy: OtlpMetricTranslationStrategy, + }, + /// Trace options mirroring `x-greptime-trace-table-name` and the + /// pipeline headers. The raw pipeline-params string is parsed by the + /// shared `GreptimePipelineParams::from_params` when the copy + /// completes, like the HTTP request does. + Traces { + table_name: Option, + pipeline_name: Option, + pipeline_version: Option, + pipeline_params: Option, + }, + /// Log options mirroring `x-greptime-log-table-name`, + /// `x-greptime-log-extract-keys` and the pipeline headers. + Logs { + table_name: Option, + pipeline_name: Option, + pipeline_version: Option, + pipeline_params: Option, + extract_keys: Option, + }, +} + +pub(super) fn build_otlp_codec(mut options: CopyOptionSet) -> PgWireResult { + let signal = options.take_string("signal")?; + let metric_options = ( + take_truthy(&mut options, "promote_all_resource_attrs")?, + options.take_string("resource_attrs")?, + take_truthy(&mut options, "promote_scope_attrs")?, + options.take_string("translation_strategy")?, + ); + let trace_log_options = ( + options.take_string("table_name")?, + options.take_string("pipeline_name")?, + options.take_string("pipeline_version")?, + options.take_string("pipeline_params")?, + ); + let extract_keys = options.take_string("extract_keys")?; + options.reject_leftovers("otlp")?; + + let unexpected = |option: &str| { + bad_copy_data(format!( + "COPY option \"{option}\" does not apply to the selected signal" + )) + }; + let signal = match signal.as_deref() { + None | Some("metrics") => { + if trace_log_options != (None, None, None, None) || extract_keys.is_some() { + return Err(unexpected("table_name/pipeline/extract_keys")); + } + let (promote_all, resource_attrs, promote_scope, strategy) = metric_options; + let promote_all = promote_all.unwrap_or(false); + let promote_scope = promote_scope.unwrap_or(false); + let translation_strategy = match strategy.as_deref() { + None => OtlpMetricTranslationStrategy::default(), + // Same values as the HTTP header (`VALUES`), parsed by + // the same `FromStr`. + Some(value) => value.parse().map_err(|_| { + bad_copy_data(format!( + "invalid translation_strategy \"{value}\", expected one of: {}", + OtlpMetricTranslationStrategy::VALUES.join(", ") + )) + })?, + }; + let resource_attrs = resource_attrs + // Same ';'-separated list and trimming as the header + // extractor. + .map(|s| s.split(';').map(|s| s.trim().to_string()).collect()) + .unwrap_or_default(); + OtlpSignal::Metrics { + promote_all_resource_attrs: promote_all, + resource_attrs, + promote_scope_attrs: promote_scope, + translation_strategy, + } + } + Some("traces") | Some("logs") => { + if metric_options != (None, None, None, None) { + return Err(unexpected("promote_*/resource_attrs/translation_strategy")); + } + let (table_name, pipeline_name, pipeline_version, pipeline_params) = trace_log_options; + if signal.as_deref() == Some("traces") { + if extract_keys.is_some() { + // `extract_keys` only applies to logs. + return Err(unexpected("extract_keys")); + } + OtlpSignal::Traces { + table_name, + pipeline_name, + pipeline_version, + pipeline_params, + } + } else { + OtlpSignal::Logs { + table_name, + pipeline_name, + pipeline_version, + pipeline_params, + extract_keys, + } + } + } + Some(unknown) => { + return Err(bad_copy_data(format!( + "invalid signal \"{unknown}\", expected one of: metrics, traces, logs" + ))); + } + }; + Ok(CopyFormat::Otlp(OtlpOptions { signal })) +} + +/// Takes a boolean option with the HTTP header's exact truthy +/// spellings (`true`, `1`, `yes`, `on`, `t`), plus the COPY forms +/// (bare flag or `TRUE`/`FALSE`), so both routes accept the same values. +fn take_truthy(options: &mut CopyOptionSet, key: &str) -> PgWireResult> { + match options.take_value(key) { + None => Ok(None), + Some(CopyOptionValue::Flag) => Ok(Some(true)), + Some(CopyOptionValue::Boolean(b)) => Ok(Some(b)), + Some(CopyOptionValue::String(v)) | Some(CopyOptionValue::Ident(v)) => Ok(Some(truthy(v))), + Some(other) => Err(bad_copy_data(format!( + "COPY option {key} requires a boolean value, got {}", + other.kind() + ))), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::postgres::copy_in::options::CopyOptionValue; + + fn opts(entries: &[(&str, CopyOptionValue)]) -> CopyOptionSet { + let mut set = CopyOptionSet::default(); + for (k, v) in entries { + set.insert((*k).into(), v.clone()).unwrap(); + } + set + } + + #[test] + fn test_signal_defaults_and_options() { + // Default signal is metrics. + assert!(matches!( + build_otlp_codec(CopyOptionSet::default()).unwrap(), + CopyFormat::Otlp(OtlpOptions { + signal: OtlpSignal::Metrics { .. } + }) + )); + + // Metric options parse like the headers. + let format = build_otlp_codec(opts(&[ + ("promote_all_resource_attrs", CopyOptionValue::Flag), + ( + "resource_attrs", + CopyOptionValue::String("service.name; host.id ".into()), + ), + ("promote_scope_attrs", CopyOptionValue::Flag), + ( + "translation_strategy", + CopyOptionValue::String("NoTranslation".into()), + ), + ])) + .unwrap(); + let CopyFormat::Otlp(OtlpOptions { + signal: + OtlpSignal::Metrics { + promote_all_resource_attrs, + resource_attrs, + promote_scope_attrs, + translation_strategy, + }, + }) = format + else { + panic!("expected metrics"); + }; + assert!(promote_all_resource_attrs && promote_scope_attrs); + assert_eq!( + resource_attrs, + ["service.name", "host.id"] + .into_iter() + .map(String::from) + .collect() + ); + assert!(matches!( + translation_strategy, + OtlpMetricTranslationStrategy::NoTranslation + )); + + // Traces and logs accept their header equivalents. + assert!(matches!( + build_otlp_codec(opts(&[( + "signal", + CopyOptionValue::String("traces".into()) + )])) + .unwrap(), + CopyFormat::Otlp(OtlpOptions { + signal: OtlpSignal::Traces { .. } + }) + )); + assert!(matches!( + build_otlp_codec(opts(&[ + ("signal", CopyOptionValue::String("logs".into())), + ("table_name", CopyOptionValue::String("my_logs".into())), + ("extract_keys", CopyOptionValue::String("a,b".into())), + ])) + .unwrap(), + CopyFormat::Otlp(OtlpOptions { + signal: OtlpSignal::Logs { .. } + }) + )); + } + + #[test] + fn test_rejects_unknown_and_mismatched_options() { + // Unknown signal. + let set = opts(&[("signal", CopyOptionValue::String("spans".into()))]); + assert!(build_otlp_codec(set).is_err()); + + // Unknown option. + let set = opts(&[("precision", CopyOptionValue::String("ns".into()))]); + assert!(build_otlp_codec(set).is_err()); + + // Signal-specific options do not cross. + let set = opts(&[ + ("signal", CopyOptionValue::String("traces".into())), + ("promote_scope_attrs", CopyOptionValue::Flag), + ]); + assert!(build_otlp_codec(set).is_err()); + let set = opts(&[ + ("signal", CopyOptionValue::String("metrics".into())), + ("table_name", CopyOptionValue::String("t".into())), + ]); + assert!(build_otlp_codec(set).is_err()); + let set = opts(&[ + ("signal", CopyOptionValue::String("traces".into())), + ( + "extract_keys", + CopyOptionValue::String("attributes.foo".into()), + ), + ]); + assert!(build_otlp_codec(set).is_err()); + } +} diff --git a/src/servers/src/postgres/server.rs b/src/servers/src/postgres/server.rs index 3919081ba62..8bdbce65a9c 100644 --- a/src/servers/src/postgres/server.rs +++ b/src/servers/src/postgres/server.rs @@ -31,6 +31,7 @@ use crate::error::Result; use crate::postgres::{MakePostgresServerHandler, MakePostgresServerHandlerBuilder}; use crate::query_handler::CopyInHandlerRef; use crate::query_handler::InfluxdbLineProtocolHandlerRef; +use crate::query_handler::OpenTelemetryProtocolHandlerRef; use crate::query_handler::sql::ServerSqlQueryHandlerRef; use crate::server::{AbortableStream, BaseTcpServer, Server}; use crate::tls::ReloadableTlsServerConfig; @@ -42,6 +43,8 @@ pub struct PostgresServer { keep_alive_secs: u64, batching_enabled: bool, influxdb_handler: Option, + otlp_handler: Option, + otlp_with_metric_engine: bool, bind_addr: Option, process_manager: Option, } @@ -80,6 +83,8 @@ impl PostgresServer { keep_alive_secs, batching_enabled: false, influxdb_handler: None, + otlp_handler: None, + otlp_with_metric_engine: false, bind_addr: None, process_manager, } @@ -100,6 +105,21 @@ impl PostgresServer { self } + /// Wires the OTLP handler used by table-less + /// `COPY FROM STDIN WITH (FORMAT otlp)` — the same handler the + /// OTLP/HTTP endpoints use, so the payload runs through the + /// identical pipeline. `with_metric_engine` carries the same flag + /// the HTTP wiring passes from its options. + pub fn with_otlp_handler( + mut self, + handler: OpenTelemetryProtocolHandlerRef, + with_metric_engine: bool, + ) -> Self { + self.otlp_handler = Some(handler); + self.otlp_with_metric_engine = with_metric_engine; + self + } + fn accept( &self, io_runtime: Runtime, @@ -115,6 +135,8 @@ impl PostgresServer { query_handler: make_handler.query_handler.clone(), copy_in_handler: make_handler.copy_in_handler.clone(), influxdb_handler: self.influxdb_handler.clone(), + otlp_handler: self.otlp_handler.clone(), + otlp_with_metric_engine: self.otlp_with_metric_engine, user_provider: make_handler.user_provider.clone(), param_provider: make_handler.param_provider.clone(), force_tls: make_handler.force_tls, diff --git a/tests-integration/src/test_util.rs b/tests-integration/src/test_util.rs index da8d8add687..6da25d659b6 100644 --- a/tests-integration/src/test_util.rs +++ b/tests-integration/src/test_util.rs @@ -1216,7 +1216,8 @@ pub async fn setup_pg_server_with_slow_query_threshold( None, None, ) - .with_influxdb_handler(fe_instance_ref.clone()), + .with_influxdb_handler(fe_instance_ref.clone()) + .with_otlp_handler(fe_instance_ref.clone(), true), ); pg_server @@ -1263,7 +1264,8 @@ pub async fn setup_pg_server_with_user_provider( user_provider, None, ) - .with_influxdb_handler(fe_instance_ref.clone()), + .with_influxdb_handler(fe_instance_ref.clone()) + .with_otlp_handler(fe_instance_ref.clone(), true), ); pg_server @@ -1322,7 +1324,8 @@ pub async fn setup_pg_server_with_prom_native_histogram( None, None, ) - .with_influxdb_handler(instance.fe_instance().clone()), + .with_influxdb_handler(instance.fe_instance().clone()) + .with_otlp_handler(instance.fe_instance().clone(), true), ); pg_server diff --git a/tests-integration/src/tests/postgres_copy_in_test.rs b/tests-integration/src/tests/postgres_copy_in_test.rs index 94b19c1315d..5d382330808 100644 --- a/tests-integration/src/tests/postgres_copy_in_test.rs +++ b/tests-integration/src/tests/postgres_copy_in_test.rs @@ -819,3 +819,357 @@ async fn test_postgres_copy_in_influxdb_invalid_utf8_and_trailing_record() { server.shutdown().await.unwrap(); } + +/// `COPY FROM STDIN WITH (FORMAT otlp)` ingests one binary OTLP +/// ExportMetricsServiceRequest through the same pipeline as the OTLP/HTTP +/// endpoint — and validates binary (format code 1) COPY support, including +/// payloads split across CopyData chunks. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_postgres_copy_in_otlp_binary() { + use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest; + use otel_arrow_rust::proto::opentelemetry::common::v1::any_value; + use otel_arrow_rust::proto::opentelemetry::common::v1::{AnyValue, KeyValue}; + use otel_arrow_rust::proto::opentelemetry::metrics::v1::number_data_point::Value; + use otel_arrow_rust::proto::opentelemetry::metrics::v1::{ + Gauge, Metric, NumberDataPoint, ResourceMetrics, ScopeMetrics, metric, + }; + use prost::Message; + + init_default_ut_logging(); + let (_guard, server) = setup_pg_server(StorageType::File, "pg_copy_in_otlp").await; + let port = server.bind_addr().unwrap().port(); + let client = connect(port).await; + + let request = ExportMetricsServiceRequest { + resource_metrics: vec![ResourceMetrics { + scope_metrics: vec![ScopeMetrics { + metrics: vec![Metric { + name: "copy_otlp_gauge".into(), + data: Some(metric::Data::Gauge(Gauge { + data_points: vec![ + NumberDataPoint { + time_unix_nano: 1_700_000_000_000_000_000, + value: Some(Value::AsDouble(1.5)), + attributes: vec![KeyValue { + key: "host".into(), + value: Some(AnyValue { + value: Some(any_value::Value::StringValue("a".into())), + }), + }], + ..Default::default() + }, + NumberDataPoint { + time_unix_nano: 1_700_000_001_000_000_000, + value: Some(Value::AsDouble(2.5)), + attributes: vec![KeyValue { + key: "host".into(), + value: Some(AnyValue { + value: Some(any_value::Value::StringValue("b".into())), + }), + }], + ..Default::default() + }, + ], + })), + ..Default::default() + }], + ..Default::default() + }], + ..Default::default() + }], + }; + let body = request.encode_to_vec(); + + // Feed the binary payload in small chunks: the message spans CopyData + // boundaries exactly like a large text payload would. + let rows = feed_copy_in( + &client, + "COPY FROM STDIN WITH (FORMAT otlp)", + body.chunks(9).map(Bytes::copy_from_slice).collect(), + ) + .await + .unwrap(); + assert_eq!(rows, 2); + + let messages = client + .simple_query("SELECT host, greptime_value FROM copy_otlp_gauge ORDER BY host") + .await + .unwrap(); + let rows: Vec<(String, String)> = messages + .iter() + .filter_map(|m| { + let tokio_postgres::SimpleQueryMessage::Row(row) = m else { + return None; + }; + Some(( + row.get(0).unwrap_or("").to_string(), + row.get(1).unwrap_or("").to_string(), + )) + }) + .collect(); + assert_eq!( + rows, + vec![("a".into(), "1.5".into()), ("b".into(), "2.5".into())] + ); + + // Garbage bytes fail the copy and the connection stays usable. + let error = feed_copy_in( + &client, + "COPY FROM STDIN WITH (FORMAT otlp)", + vec![Bytes::from_static(b"\xff\xff\xff\xff")], + ) + .await + .unwrap_err(); + assert!( + error + .as_db_error() + .map(|d| d.message().contains("invalid OTLP payload")) + .unwrap_or(false), + "unexpected error: {error}" + ); + let messages = client.simple_query("SELECT 1").await.unwrap(); + assert!(!messages.is_empty()); + + server.shutdown().await.unwrap(); +} + +/// `COPY FROM STDIN WITH (FORMAT otlp, signal 'traces')` ingests one +/// binary ExportTraceServiceRequest through the same pipeline as the +/// OTLP/HTTP traces endpoint, honoring the trace-table-name option. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_postgres_copy_in_otlp_traces() { + use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest; + use prost::Message; + use serde_json::json; + + init_default_ut_logging(); + let (_guard, server) = setup_pg_server(StorageType::File, "pg_copy_in_otlp_traces").await; + let port = server.bind_addr().unwrap().port(); + let client = connect(port).await; + + let request: ExportTraceServiceRequest = serde_json::from_value(json!({ + "resourceSpans": [{ + "resource": {"attributes": [ + {"key": "service.name", "value": {"stringValue": "copy-svc"}} + ]}, + "scopeSpans": [{"spans": [{ + "traceId": "c05d7a4ec8e1f231f02ed6e8da8655b4", + "spanId": "9630f2916e2f7909", + "name": "copy-op", + "kind": 2, + "startTimeUnixNano": "1736480942444376000", + "endTimeUnixNano": "1736480942444499000" + }]}] + }] + })) + .unwrap(); + + let rows = feed_copy_in( + &client, + "COPY FROM STDIN WITH (FORMAT otlp, signal 'traces', table_name 'copy_traces', pipeline_name 'greptime_trace_v2')", + vec![Bytes::from(request.encode_to_vec())], + ) + .await + .unwrap(); + assert_eq!(rows, 1); + + let messages = client + .simple_query("SELECT count(*) FROM copy_traces WHERE service_name = 'copy-svc'") + .await + .unwrap(); + let count: Vec = messages + .iter() + .filter_map(|m| match m { + tokio_postgres::SimpleQueryMessage::Row(row) => row.get(0).map(str::to_string), + _ => None, + }) + .collect(); + assert_eq!(count, vec!["1".to_string()]); + + server.shutdown().await.unwrap(); +} + +/// `COPY FROM STDIN WITH (FORMAT otlp, signal 'logs')` ingests one +/// binary ExportLogsServiceRequest through the same pipeline as the +/// OTLP/HTTP logs endpoint, with the default log table. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_postgres_copy_in_otlp_logs() { + use opentelemetry_proto::tonic::collector::logs::v1::ExportLogsServiceRequest; + use prost::Message; + use serde_json::json; + + init_default_ut_logging(); + let (_guard, server) = setup_pg_server(StorageType::File, "pg_copy_in_otlp_logs").await; + let port = server.bind_addr().unwrap().port(); + let client = connect(port).await; + + let request: ExportLogsServiceRequest = serde_json::from_value(json!({ + "resourceLogs": [{ + "resource": {"attributes": []}, + "scopeLogs": [{"logRecords": [{ + "timeUnixNano": "1736413568497632000", + "severityNumber": 9, + "severityText": "Info", + "body": {"stringValue": "copied via postgres"}, + "traceId": "f665100a612542b69cc362fe2ae9d3bf", + "spanId": "e58f01c4c69f4488" + }]}] + }] + })) + .unwrap(); + + let rows = feed_copy_in( + &client, + "COPY FROM STDIN WITH (FORMAT otlp, signal 'logs')", + vec![Bytes::from(request.encode_to_vec())], + ) + .await + .unwrap(); + assert_eq!(rows, 1); + + let messages = client + .simple_query("SELECT count(*) FROM opentelemetry_logs WHERE body = 'copied via postgres'") + .await + .unwrap(); + let count: Vec = messages + .iter() + .filter_map(|m| match m { + tokio_postgres::SimpleQueryMessage::Row(row) => row.get(0).map(str::to_string), + _ => None, + }) + .collect(); + assert_eq!(count, vec!["1".to_string()]); + + server.shutdown().await.unwrap(); +} + +/// A metrics COPY honors the header-equivalent options: promoting all +/// resource attributes adds them as columns. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_postgres_copy_in_otlp_metrics_options() { + use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest; + use prost::Message; + + init_default_ut_logging(); + let (_guard, server) = setup_pg_server(StorageType::File, "pg_copy_in_otlp_mopts").await; + let port = server.bind_addr().unwrap().port(); + let client = connect(port).await; + + let request = ExportMetricsServiceRequest { + resource_metrics: vec![otel_arrow_rust::proto::opentelemetry::metrics::v1::ResourceMetrics { + resource: Some(otel_arrow_rust::proto::opentelemetry::resource::v1::Resource { + attributes: vec![otel_arrow_rust::proto::opentelemetry::common::v1::KeyValue { + key: "deployment.environment".into(), + value: Some(otel_arrow_rust::proto::opentelemetry::common::v1::AnyValue { + value: Some( + otel_arrow_rust::proto::opentelemetry::common::v1::any_value::Value::StringValue("prod".into()), + ), + }), + }], + ..Default::default() + }), + scope_metrics: vec![otel_arrow_rust::proto::opentelemetry::metrics::v1::ScopeMetrics { + metrics: vec![otel_arrow_rust::proto::opentelemetry::metrics::v1::Metric { + name: "copy_opt_gauge".into(), + data: Some(otel_arrow_rust::proto::opentelemetry::metrics::v1::metric::Data::Gauge( + otel_arrow_rust::proto::opentelemetry::metrics::v1::Gauge { + data_points: vec![otel_arrow_rust::proto::opentelemetry::metrics::v1::NumberDataPoint { + time_unix_nano: 1_700_000_000_000_000_000, + value: Some(otel_arrow_rust::proto::opentelemetry::metrics::v1::number_data_point::Value::AsDouble(42.0)), + ..Default::default() + }], + }, + )), + ..Default::default() + }], + ..Default::default() + }], + ..Default::default() + }], + }; + + let rows = feed_copy_in( + &client, + "COPY FROM STDIN WITH (FORMAT otlp, signal 'metrics', promote_all_resource_attrs)", + vec![Bytes::from(request.encode_to_vec())], + ) + .await + .unwrap(); + assert_eq!(rows, 1); + + let messages = client + .simple_query("SELECT count(*) FROM copy_opt_gauge WHERE deployment_environment = 'prod'") + .await + .unwrap(); + let count: Vec = messages + .iter() + .filter_map(|m| match m { + tokio_postgres::SimpleQueryMessage::Row(row) => row.get(0).map(str::to_string), + _ => None, + }) + .collect(); + assert_eq!(count, vec!["1".to_string()]); + + server.shutdown().await.unwrap(); +} + +/// A payload whose data points are all rejected fails the copy with the +/// pipeline's reason — the same failure semantics as the HTTP route — +/// and the connection stays usable. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_postgres_copy_in_otlp_all_rejected() { + use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest; + use otel_arrow_rust::proto::opentelemetry::metrics::v1::{ + AggregationTemporality, Histogram, HistogramDataPoint, Metric, ResourceMetrics, + ScopeMetrics, metric, + }; + use prost::Message; + + init_default_ut_logging(); + let (_guard, server) = setup_pg_server(StorageType::File, "pg_copy_in_otlp_rej").await; + let port = server.bind_addr().unwrap().port(); + let client = connect(port).await; + + // A delta histogram cannot be stored as-is: its single point is + // rejected — the same trigger the OTLP suite uses for the HTTP + // route's failure outcome. + let request = ExportMetricsServiceRequest { + resource_metrics: vec![ResourceMetrics { + scope_metrics: vec![ScopeMetrics { + metrics: vec![Metric { + name: "copy_rejected_histogram".into(), + data: Some(metric::Data::Histogram(Histogram { + data_points: vec![HistogramDataPoint { + count: 1, + bucket_counts: vec![1], + explicit_bounds: vec![1.0, 2.0], + ..Default::default() + }], + aggregation_temporality: AggregationTemporality::Delta as i32, + })), + ..Default::default() + }], + ..Default::default() + }], + ..Default::default() + }], + }; + let error = feed_copy_in( + &client, + "COPY FROM STDIN WITH (FORMAT otlp)", + vec![Bytes::from(request.encode_to_vec())], + ) + .await + .unwrap_err(); + assert!( + error + .as_db_error() + .map(|d| d.message().contains("rejected all")) + .unwrap_or(false), + "unexpected error: {error}" + ); + let messages = client.simple_query("SELECT 1").await.unwrap(); + assert!(!messages.is_empty()); + + server.shutdown().await.unwrap(); +} From 295bed98ec1c0ebd229f123bd0669be87ae3b0ed Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Mon, 2 Nov 2026 10:52:31 +0000 Subject: [PATCH 154/239] feat: share Parquet byte page loads (#9793) Signed-off-by: WenyXu --- src/cmd/src/datanode/parquetbench.rs | 1 - src/common/cache/src/cache.rs | 33 +- src/common/cache/src/paged.rs | 16 +- src/mito2/src/cache.rs | 561 +---------- src/mito2/src/cache/index.rs | 19 +- .../src/cache/index/bloom_filter_index.rs | 2 +- src/mito2/src/cache/index/inverted_index.rs | 2 +- src/mito2/src/cache/parquet_page.rs | 353 +++++++ src/mito2/src/error.rs | 16 +- src/mito2/src/sst/file.rs | 2 +- src/mito2/src/sst/parquet.rs | 27 +- src/mito2/src/sst/parquet/push_decoder.rs | 887 ++++++++++++------ src/mito2/src/sst/parquet/reader.rs | 3 +- 13 files changed, 1096 insertions(+), 826 deletions(-) create mode 100644 src/mito2/src/cache/parquet_page.rs diff --git a/src/cmd/src/datanode/parquetbench.rs b/src/cmd/src/datanode/parquetbench.rs index a5c73f29527..7cfe01785a8 100644 --- a/src/cmd/src/datanode/parquetbench.rs +++ b/src/cmd/src/datanode/parquetbench.rs @@ -713,7 +713,6 @@ async fn run_direct_iteration( file_path.clone(), object_store.clone(), CacheStrategy::Disabled, - row_group_idx, None, ); let mut stream = build_sst_parquet_record_batch_stream( diff --git a/src/common/cache/src/cache.rs b/src/common/cache/src/cache.rs index 98efb890f29..42c4fa4bd24 100644 --- a/src/common/cache/src/cache.rs +++ b/src/common/cache/src/cache.rs @@ -15,6 +15,7 @@ use std::borrow::Borrow; use std::hash::Hash; use std::sync::Arc; +use std::time::Duration; use crate::metrics::{CacheMetrics, CacheSizeMetric, CacheSizeRegistration}; @@ -76,14 +77,32 @@ where W: Fn(&K, &V) -> u32 + Send + Sync + 'static, L: Fn(Arc, V, RemovalCause) + Send + Sync + 'static, { + Self::new_with_ttl(name, capacity, weigher, listener, None) + } + + /// Creates a cache with an optional write-based expiration time. + pub fn new_with_ttl( + name: &str, + capacity: u64, + weigher: W, + listener: L, + ttl: Option, + ) -> Self + where + W: Fn(&K, &V) -> u32 + Send + Sync + 'static, + L: Fn(Arc, V, RemovalCause) + Send + Sync + 'static, + { + let mut builder = moka::future::Cache::builder() + .name(name) + .max_capacity(capacity) + .support_invalidation_closures() + .weigher(weigher) + .eviction_listener(move |key, value, cause| listener(key, value, cause.into())); + if let Some(ttl) = ttl { + builder = builder.time_to_live(ttl); + } Self { - inner: moka::future::Cache::builder() - .name(name) - .max_capacity(capacity) - .support_invalidation_closures() - .weigher(weigher) - .eviction_listener(move |key, value, cause| listener(key, value, cause.into())) - .build(), + inner: builder.build(), capacity, size_registration: None, metrics: None, diff --git a/src/common/cache/src/paged.rs b/src/common/cache/src/paged.rs index 3131329dc04..6b986afc22a 100644 --- a/src/common/cache/src/paged.rs +++ b/src/common/cache/src/paged.rs @@ -25,6 +25,7 @@ use std::num::NonZeroU64; use std::ops::Range; use std::sync::Arc; use std::task::Poll; +use std::time::Duration; use bytes::{Bytes, BytesMut}; use common_base::range_read::RangeReader; @@ -88,15 +89,28 @@ impl PagedCache { page_size: u64, weigher: impl Fn(&(K, u64), &Bytes) -> u32 + Send + Sync + 'static, observer: PageCacheObserver, + ) -> Self { + Self::new_with_ttl(name, capacity, page_size, weigher, observer, None) + } + + /// Creates a page cache with optional expiration measured from insertion. + pub fn new_with_ttl( + name: &str, + capacity: u64, + page_size: u64, + weigher: impl Fn(&(K, u64), &Bytes) -> u32 + Send + Sync + 'static, + observer: PageCacheObserver, + ttl: Option, ) -> Self { let events = observer.events.clone(); - let mut cache = Cache::new( + let mut cache = Cache::new_with_ttl( name, capacity, move |key, value: &Arc| weigher(key, &value.bytes), move |_, _, cause| { events(PageEvent::Evicted { cause }); }, + ttl, ) .with_size_metric(observer.bytes.clone()); if let Some(metrics) = &observer.loads { diff --git a/src/mito2/src/cache.rs b/src/mito2/src/cache.rs index 1b3645eff62..7161e27b5c2 100644 --- a/src/mito2/src/cache.rs +++ b/src/mito2/src/cache.rs @@ -19,15 +19,14 @@ pub(crate) mod cache_size; pub(crate) mod file_cache; pub(crate) mod index; pub(crate) mod manifest_cache; +pub(crate) mod parquet_page; pub(crate) mod range_cache; #[cfg(test)] pub(crate) mod test_util; pub(crate) mod write_cache; -use std::collections::{BTreeMap, HashMap}; -use std::mem; -use std::ops::Range; -use std::sync::{Arc, RwLock, Weak}; +use std::sync::{Arc, Weak}; +use std::{io, mem}; use bytes::Bytes; use common_base::readable_size::ReadableSize; @@ -59,6 +58,7 @@ pub use write_cache::{WriteCacheUploadStoreWrapper, WriteCacheUploadStoreWrapper use crate::cache::cache_size::parquet_meta_size; use crate::cache::file_cache::{FileType, IndexKey}; use crate::cache::index::inverted_index::{InvertedIndexCache, InvertedIndexCacheRef}; +use crate::cache::parquet_page::ParquetPageCache; use crate::cache::write_cache::WriteCacheRef; use crate::error::{ CompressObjectSnafu, DecompressObjectSnafu, InvalidMetadataSnafu, InvalidParquetSnafu, @@ -74,14 +74,29 @@ use crate::sst::parquet::PARQUET_METADATA_KEY; use crate::sst::parquet::read_columns::ParquetReadColumns; use crate::sst::parquet::reader::MetadataCacheMetrics; +/// Finds a typed shared error through the range-reader I/O wrappers. +pub(crate) fn find_shared_page_source( + mut error: &io::Error, +) -> Option> { + loop { + let source = error.get_ref()?; + if let Some(source) = source.downcast_ref::>() { + return Some(source.clone()); + } + if let Some(nested) = source.downcast_ref::>() { + error = nested; + } else { + error = source.downcast_ref::()?; + } + } +} + /// Metrics type key for sst meta. const SST_META_TYPE: &str = "sst_meta"; /// Metrics type key for the optional decoded SST metadata acceleration tier. const SST_META_DECODED_TYPE: &str = "sst_meta_decoded"; /// Metrics type key for vector. const VECTOR_TYPE: &str = "vector"; -/// Metrics type key for pages. -const PAGE_TYPE: &str = "page"; /// Metrics type key for files on the local store. const FILE_TYPE: &str = "file"; /// Metrics type key for index files (puffin) on the local store. @@ -945,33 +960,11 @@ impl CacheStrategy { } } - /// Calls [CacheManager::get_page_ranges()]. - /// It returns None if the strategy is [CacheStrategy::Compaction] or [CacheStrategy::Disabled]. - pub fn get_page_ranges( - &self, - file_id: FileId, - row_group_idx: usize, - ranges: &[Range], - ) -> Option { + /// Returns the Parquet byte cache for normal query reads. + pub(crate) fn parquet_page_cache(&self) -> Option<&Arc> { match self { - CacheStrategy::EnableAll(cache_manager) => { - cache_manager.get_page_ranges(file_id, row_group_idx, ranges) - } - CacheStrategy::Compaction(_) | CacheStrategy::Disabled => None, - } - } - - /// Calls [CacheManager::put_page_ranges()]. - /// It does nothing if the strategy isn't [CacheStrategy::EnableAll]. - pub fn put_page_ranges( - &self, - file_id: FileId, - row_group_idx: usize, - ranges: &[Range], - pages: &[Bytes], - ) { - if let CacheStrategy::EnableAll(cache_manager) = self { - cache_manager.put_page_ranges(file_id, row_group_idx, ranges, pages); + Self::EnableAll(manager) => manager.page_cache.as_ref(), + Self::Compaction(_) | Self::Disabled => None, } } @@ -1152,7 +1145,7 @@ pub struct CacheManager { /// Cache for vectors. vector_cache: Option, /// Cache for SST byte ranges. - page_cache: Option>, + page_cache: Option>, /// A Cache for writing files to object stores. write_cache: Option, /// Cache for inverted index. @@ -1475,38 +1468,6 @@ impl CacheManager { } } - /// Gets cached byte fragments for the requested ranges. - pub fn get_page_ranges( - &self, - file_id: FileId, - row_group_idx: usize, - ranges: &[Range], - ) -> Option { - self.page_cache.as_ref().map(|page_cache| { - let lookup = page_cache.lookup(file_id, row_group_idx, ranges); - if lookup.cached_bytes > 0 { - CACHE_HIT.with_label_values(&[PAGE_TYPE]).inc(); - } - if !lookup.missing_ranges.is_empty() { - CACHE_MISS.with_label_values(&[PAGE_TYPE]).inc(); - } - lookup - }) - } - - /// Puts byte fragments into the page cache. - pub fn put_page_ranges( - &self, - file_id: FileId, - row_group_idx: usize, - ranges: &[Range], - pages: &[Bytes], - ) { - if let Some(cache) = &self.page_cache { - cache.insert_ranges(file_id, row_group_idx, ranges, pages); - } - } - /// Evicts every puffin-related cache entry for the given file. pub async fn evict_puffin_cache(&self, file_id: RegionIndexId) { if let Some(cache) = &self.bloom_filter_index_cache { @@ -1792,7 +1753,7 @@ impl CacheManagerBuilder { .build() }); let page_cache = - (self.page_cache_size != 0).then(|| PageRangeCache::new(self.page_cache_size)); + (self.page_cache_size != 0).then(|| ParquetPageCache::from_env(self.page_cache_size)); let inverted_index_cache = InvertedIndexCache::new( self.index_metadata_size, self.index_content_size, @@ -1866,10 +1827,6 @@ fn vector_cache_weight(_k: &(ConcreteDataType, Value), v: &VectorRef) -> u32 { (mem::size_of::() + mem::size_of::() + v.memory_size()) as u32 } -fn page_cache_weight(k: &PageFragmentKey, v: &Bytes) -> u32 { - (k.estimated_size() + mem::size_of::() + v.len()) as u32 -} - fn selector_result_cache_weight(k: &SelectorResultKey, v: &Arc) -> u32 { (mem::size_of_val(k) + v.estimated_size()) as u32 } @@ -1899,284 +1856,6 @@ impl SstMetaKey { } } -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] -struct PageFragmentGroupKey { - file_id: FileId, - row_group_idx: usize, -} - -/// Cache key for one byte fragment in an SST row group. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] -pub struct PageFragmentKey { - /// Id of the SST file. - file_id: FileId, - /// Index of the row group. - row_group_idx: usize, - /// Start offset of the cached byte fragment. - start: u64, - /// End offset of the cached byte fragment. - end: u64, -} - -impl PageFragmentKey { - fn new(file_id: FileId, row_group_idx: usize, range: &Range) -> PageFragmentKey { - PageFragmentKey { - file_id, - row_group_idx, - start: range.start, - end: range.end, - } - } - - fn group_key(&self) -> PageFragmentGroupKey { - PageFragmentGroupKey { - file_id: self.file_id, - row_group_idx: self.row_group_idx, - } - } - - /// Returns memory used by the key (estimated). - fn estimated_size(&self) -> usize { - mem::size_of::() - } -} - -/// One cached byte fragment that overlaps a requested range. -#[derive(Clone)] -pub struct PageRangePart { - /// Range covered by `bytes`. - pub range: Range, - /// Bytes for `range`. - pub bytes: Bytes, -} - -/// Result of looking up request ranges in the page range cache. -pub struct PageRangeLookup { - /// Cached fragments grouped by the original requested range index. - pub cached_parts: Vec>, - /// Ranges that are not covered by cached fragments and need fetching. - pub missing_ranges: Vec>, - /// Number of cached fragments used. - pub cached_range_count: usize, - /// Number of requested bytes served from cached fragments. - pub cached_bytes: u64, -} - -impl PageRangeLookup { - pub fn is_fully_cached(&self) -> bool { - self.missing_ranges.is_empty() - } -} - -type PageFragmentRangeIndex = BTreeMap<(u64, u64), PageFragmentKey>; -type PageFragmentIndex = HashMap; - -/// Byte-fragment cache for Parquet row-group reads. -pub struct PageRangeCache { - cache: Cache, - index: RwLock, -} - -impl PageRangeCache { - fn new(capacity: u64) -> Arc { - Arc::new_cyclic(|weak_cache: &std::sync::Weak| { - let cache = Cache::builder() - .max_capacity(capacity) - .weigher(page_cache_weight) - .eviction_listener({ - let weak_cache = weak_cache.clone(); - move |k, v, cause| { - let size = page_cache_weight(&k, &v); - CACHE_BYTES.with_label_values(&[PAGE_TYPE]).sub(size.into()); - CACHE_EVICTION - .with_label_values(&[PAGE_TYPE, removal_cause_str(cause)]) - .inc(); - - if let Some(cache) = weak_cache.upgrade() - && !matches!(cause, RemovalCause::Replaced) - { - cache.remove_index_entry(*k); - } - } - }) - .build(); - - PageRangeCache { - cache, - index: RwLock::new(HashMap::new()), - } - }) - } - - fn lookup( - &self, - file_id: FileId, - row_group_idx: usize, - ranges: &[Range], - ) -> PageRangeLookup { - let mut cached_parts = Vec::with_capacity(ranges.len()); - let mut missing_ranges = Vec::new(); - let mut cached_range_count = 0; - let mut cached_bytes = 0; - - for range in ranges { - if range.start >= range.end { - cached_parts.push(Vec::new()); - continue; - } - - let mut parts = Vec::new(); - let candidates = self.find_index_candidates(file_id, row_group_idx, range); - let mut stale_keys = Vec::new(); - - for fragment_key in candidates { - if let Some(bytes) = self.cache.get(&fragment_key) { - let part_start = range.start.max(fragment_key.start); - let part_end = range.end.min(fragment_key.end); - let slice_start = (part_start - fragment_key.start) as usize; - let slice_end = (part_end - fragment_key.start) as usize; - parts.push(PageRangePart { - range: part_start..part_end, - bytes: bytes.slice(slice_start..slice_end), - }); - } else { - stale_keys.push(fragment_key); - } - } - for key in stale_keys { - self.remove_uncached_index_entry(key); - } - - let mut cursor = range.start; - let mut compacted_parts: Vec = Vec::with_capacity(parts.len()); - for part in parts { - if part.range.end <= cursor { - continue; - } - - let part = if part.range.start < cursor { - let offset = (cursor - part.range.start) as usize; - PageRangePart { - range: cursor..part.range.end, - bytes: part.bytes.slice(offset..), - } - } else { - part - }; - - if cursor < part.range.start { - missing_ranges.push(cursor..part.range.start); - } - cached_bytes += part.range.end - part.range.start; - cached_range_count += 1; - cursor = part.range.end; - compacted_parts.push(part); - - if cursor >= range.end { - break; - } - } - - if cursor < range.end { - missing_ranges.push(cursor..range.end); - } - cached_parts.push(compacted_parts); - } - - PageRangeLookup { - cached_parts, - missing_ranges, - cached_range_count, - cached_bytes, - } - } - - fn insert_ranges( - &self, - file_id: FileId, - row_group_idx: usize, - ranges: &[Range], - pages: &[Bytes], - ) { - for (range, bytes) in ranges.iter().zip(pages) { - if range.start >= range.end || bytes.len() as u64 != range.end - range.start { - continue; - } - - let key = PageFragmentKey::new(file_id, row_group_idx, range); - let bytes = Bytes::copy_from_slice(bytes.as_ref()); - let size = page_cache_weight(&key, &bytes); - CACHE_BYTES.with_label_values(&[PAGE_TYPE]).add(size.into()); - self.cache.insert(key, bytes); - let mut index = self.index.write().unwrap(); - index - .entry(key.group_key()) - .or_default() - .insert((key.start, key.end), key); - } - } - - fn find_index_candidates( - &self, - file_id: FileId, - row_group_idx: usize, - range: &Range, - ) -> Vec { - let group_key = PageFragmentGroupKey { - file_id, - row_group_idx, - }; - let index = self.index.read().unwrap(); - index - .get(&group_key) - .map(|ranges| { - ranges - .range(..(range.end, 0)) - .filter_map(|(_, fragment_key)| { - (fragment_key.end > range.start).then_some(*fragment_key) - }) - .collect() - }) - .unwrap_or_default() - } - - fn remove_uncached_index_entry(&self, key: PageFragmentKey) { - let group_key = key.group_key(); - let mut index = self.index.write().unwrap(); - if self.cache.contains_key(&key) { - return; - } - - Self::remove_index_entry_locked(&mut index, group_key, key); - } - - fn remove_index_entry(&self, key: PageFragmentKey) { - let group_key = key.group_key(); - let mut index = self.index.write().unwrap(); - Self::remove_index_entry_locked(&mut index, group_key, key); - } - - fn remove_index_entry_locked( - index: &mut PageFragmentIndex, - group_key: PageFragmentGroupKey, - key: PageFragmentKey, - ) { - let Some(ranges) = index.get_mut(&group_key) else { - return; - }; - - let removed = ranges - .get(&(key.start, key.end)) - .is_some_and(|current| current == &key); - if removed { - ranges.remove(&(key.start, key.end)); - } - if ranges.is_empty() { - index.remove(&group_key); - } - } -} - /// Cache key for time series row selector result. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub struct SelectorResultKey { @@ -2369,6 +2048,24 @@ mod tests { use crate::read::read_columns::ReadColumns; use crate::sst::parquet::row_selection::RowGroupSelection; + #[test] + fn test_find_shared_page_source() { + let source = Arc::new( + UnexpectedSnafu { + reason: "page read failed", + } + .build(), + ); + let direct = io::Error::other(source.clone()); + let nested = io::Error::other(Arc::new(io::Error::other(direct))); + let retained = find_shared_page_source::(&nested).unwrap(); + assert!(Arc::ptr_eq(&source, &retained)); + assert!( + find_shared_page_source::(&io::Error::other("untyped error")) + .is_none() + ); + } + #[tokio::test] async fn test_disable_cache() { let cache = CacheManager::default(); @@ -3151,174 +2848,6 @@ mod tests { assert_eq!(vector, cached); } - #[test] - fn test_page_cache() { - let cache = CacheManager::builder().page_cache_size(1000).build(); - let file_id = FileId::random(); - let uncached = 0..10; - assert_eq!( - vec![0..10], - cache - .get_page_ranges(file_id, 0, std::slice::from_ref(&uncached)) - .unwrap() - .missing_ranges - ); - - let cached = 100..500; - cache.put_page_ranges( - file_id, - 0, - std::slice::from_ref(&cached), - &[Bytes::from(vec![7; 400])], - ); - - let subrange = 200..300; - let lookup = cache - .get_page_ranges(file_id, 0, std::slice::from_ref(&subrange)) - .unwrap(); - assert!(lookup.is_fully_cached()); - assert_eq!(100, lookup.cached_bytes); - assert_eq!(1, lookup.cached_parts.len()); - assert_eq!(200..300, lookup.cached_parts[0][0].range); - assert_eq!(100, lookup.cached_parts[0][0].bytes.len()); - - let overlapping = 400..600; - let lookup = cache - .get_page_ranges(file_id, 0, std::slice::from_ref(&overlapping)) - .unwrap(); - assert!(!lookup.is_fully_cached()); - assert_eq!(100, lookup.cached_bytes); - assert_eq!(vec![500..600], lookup.missing_ranges); - assert_eq!(400..500, lookup.cached_parts[0][0].range); - } - - #[test] - fn test_page_cache_detaches_fragment_bytes() { - let cache = PageRangeCache::new(1000); - let file_id = FileId::random(); - let backing = Bytes::from(vec![1; 1024]); - let page = backing.slice(512..522); - let page_ptr = page.as_ptr(); - let range = 0..10; - - cache.insert_ranges( - file_id, - 0, - std::slice::from_ref(&range), - std::slice::from_ref(&page), - ); - - let lookup = cache.lookup(file_id, 0, std::slice::from_ref(&range)); - assert!(lookup.is_fully_cached()); - assert_eq!(1, lookup.cached_parts[0].len()); - assert_eq!(&page[..], &lookup.cached_parts[0][0].bytes[..]); - assert_ne!(page_ptr, lookup.cached_parts[0][0].bytes.as_ptr()); - } - - #[test] - fn test_page_cache_replaces_fragment() { - let cache = PageRangeCache::new(1000); - let file_id = FileId::random(); - let range = 0..10; - - cache.insert_ranges( - file_id, - 0, - std::slice::from_ref(&range), - &[Bytes::from(vec![1; 10])], - ); - cache.insert_ranges( - file_id, - 0, - std::slice::from_ref(&range), - &[Bytes::from(vec![2; 10])], - ); - cache.cache.run_pending_tasks(); - assert_eq!( - vec![PageFragmentKey::new(file_id, 0, &range)], - cache.find_index_candidates(file_id, 0, &range) - ); - - let lookup = cache.lookup(file_id, 0, std::slice::from_ref(&range)); - assert!(lookup.is_fully_cached()); - assert_eq!(&vec![2; 10][..], &lookup.cached_parts[0][0].bytes[..]); - } - - #[test] - fn test_page_cache_retains_disjoint_inserts_for_same_row_group() { - let cache = PageRangeCache::new(1000); - let file_id = FileId::random(); - let range1 = 0..10; - let range2 = 20..30; - - cache.insert_ranges( - file_id, - 0, - std::slice::from_ref(&range1), - &[Bytes::from(vec![1; 10])], - ); - cache.insert_ranges( - file_id, - 0, - std::slice::from_ref(&range2), - &[Bytes::from(vec![2; 10])], - ); - - let lookup = cache.lookup(file_id, 0, &[range1, range2]); - assert!(lookup.is_fully_cached()); - assert_eq!(2, lookup.cached_range_count); - assert_eq!(&vec![1; 10][..], &lookup.cached_parts[0][0].bytes[..]); - assert_eq!(&vec![2; 10][..], &lookup.cached_parts[1][0].bytes[..]); - } - - #[test] - fn test_page_cache_fragment_eviction() { - let file_id = FileId::random(); - let range = 0..10; - let key = PageFragmentKey::new(file_id, 0, &range); - let page = Bytes::from(vec![1; 10]); - let cache = PageRangeCache::new(page_cache_weight(&key, &page) as u64); - - cache.insert_ranges( - file_id, - 0, - std::slice::from_ref(&range), - &[Bytes::from(vec![1; 10])], - ); - assert!( - cache - .lookup(file_id, 0, std::slice::from_ref(&range)) - .is_fully_cached() - ); - - cache.cache.invalidate(&key); - cache.cache.run_pending_tasks(); - assert!(cache.find_index_candidates(file_id, 0, &range).is_empty()); - - let lookup = cache.lookup(file_id, 0, std::slice::from_ref(&range)); - assert!(!lookup.is_fully_cached()); - assert_eq!(vec![0..10], lookup.missing_ranges); - } - - #[test] - fn test_page_cache_rejects_oversized_fragment() { - let cache = PageRangeCache::new(1); - let file_id = FileId::random(); - let range = 0..10; - - cache.insert_ranges( - file_id, - 0, - std::slice::from_ref(&range), - &[Bytes::from(vec![1; 10])], - ); - cache.cache.run_pending_tasks(); - - let lookup = cache.lookup(file_id, 0, std::slice::from_ref(&range)); - assert!(!lookup.is_fully_cached()); - assert_eq!(vec![0..10], lookup.missing_ranges); - } - #[test] fn test_selector_result_cache() { let cache = CacheManager::builder() diff --git a/src/mito2/src/cache/index.rs b/src/mito2/src/cache/index.rs index 5f9208f006f..ed2297709c1 100644 --- a/src/mito2/src/cache/index.rs +++ b/src/mito2/src/cache/index.rs @@ -238,23 +238,6 @@ where } } -/// Finds a typed shared error through the range-reader I/O wrappers. -fn find_shared_page_source( - mut error: &io::Error, -) -> Option> { - loop { - let source = error.get_ref()?; - if let Some(source) = source.downcast_ref::>() { - return Some(source.clone()); - } - if let Some(nested) = source.downcast_ref::>() { - error = nested; - } else { - error = source.downcast_ref::()?; - } - } -} - #[cfg(test)] mod tests { use tokio::sync::Notify; @@ -497,7 +480,7 @@ mod tests { }) .await .unwrap_err(); - let retained = find_shared_page_source::(&error).unwrap(); + let retained = crate::cache::find_shared_page_source::(&error).unwrap(); assert_eq!( retained.status_code(), common_error::status_code::StatusCode::InvalidArguments diff --git a/src/mito2/src/cache/index/bloom_filter_index.rs b/src/mito2/src/cache/index/bloom_filter_index.rs index 1bd8b0fbae6..80f5148185d 100644 --- a/src/mito2/src/cache/index/bloom_filter_index.rs +++ b/src/mito2/src/cache/index/bloom_filter_index.rs @@ -242,7 +242,7 @@ impl BloomFilterReader for CachedBloomFilterIndexBl /// Recovers the typed source retained by the range-reader adapter. fn map_page_error(error: Arc) -> Error { - if let Some(source) = crate::cache::index::find_shared_page_source::(&error) { + if let Some(source) = crate::cache::find_shared_page_source::(&error) { SharedReadSnafu.into_error(source) } else { IoSnafu.into_error(std::io::Error::other(error)) diff --git a/src/mito2/src/cache/index/inverted_index.rs b/src/mito2/src/cache/index/inverted_index.rs index b3a0c6af955..34b44e571bb 100644 --- a/src/mito2/src/cache/index/inverted_index.rs +++ b/src/mito2/src/cache/index/inverted_index.rs @@ -199,7 +199,7 @@ impl InvertedIndexReader for CachedInvertedIndexBlobRead /// Recovers the typed source retained by the range-reader adapter. fn map_page_error(error: Arc) -> Error { - if let Some(source) = crate::cache::index::find_shared_page_source::(&error) { + if let Some(source) = crate::cache::find_shared_page_source::(&error) { SharedReadSnafu.into_error(source) } else { CommonIoSnafu.into_error(std::io::Error::other(error)) diff --git a/src/mito2/src/cache/parquet_page.rs b/src/mito2/src/cache/parquet_page.rs new file mode 100644 index 00000000000..2d07fb1e68b --- /dev/null +++ b/src/mito2/src/cache/parquet_page.rs @@ -0,0 +1,353 @@ +// 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. + +//! Fixed byte blocks shared across row groups of an immutable Parquet file. + +use std::io; +use std::mem; +use std::ops::Range; +use std::sync::Arc; +use std::time::Duration; + +use bytes::Bytes; +use common_base::range_read::RangeReader; +use common_cache::paged::{PageFetchStats, PagedCache}; + +use crate::metrics; +use crate::sst::file::RegionFileId; + +/// Default byte blocks, independent of the Parquet format's encoded data-page boundaries. +pub(crate) const PAGE_SIZE: u64 = 64 * 1024; + +/// Internal tuning knob, in bytes; read only when constructing a cache. +const PAGE_SIZE_ENV: &str = "GREPTIME_PARQUET_PAGE_SIZE"; +/// Expire pages one hour after insertion; physical reclamation uses cache maintenance. +const PAGE_TTL: Duration = Duration::from_secs(3600); + +fn parse_page_size(value: Option<&str>) -> u64 { + match value { + None => PAGE_SIZE, + Some(value) => match value.parse::() { + Ok(size) if (1..=MAX_PAGE_SIZE).contains(&size) => size, + _ => { + common_telemetry::warn!( + "Invalid {PAGE_SIZE_ENV}: expected bytes in 1..={MAX_PAGE_SIZE}; using {PAGE_SIZE}" + ); + PAGE_SIZE + } + }, + } +} + +/// Largest block whose bytes and entry overhead fit the cache's u32 weight. +pub(crate) const MAX_PAGE_SIZE: u64 = u32::MAX as u64 + - (mem::size_of::<(RegionFileId, u64)>() + mem::size_of::() + mem::size_of::()) + as u64; + +/// Caches immutable Parquet file bytes in fixed-size blocks shared across row groups. +pub(crate) struct ParquetPageCache { + page_size: u64, + cache: PagedCache, +} + +impl ParquetPageCache { + /// Constructs the production cache, resolving the internal block-size override once. + pub(crate) fn from_env(capacity: u64) -> Arc { + let value = std::env::var(PAGE_SIZE_ENV).ok(); + Self::new(capacity, parse_page_size(value.as_deref())) + } + + /// Creates a cache with a total byte capacity and a byte-block size. + pub(crate) fn new(capacity: u64, page_size: u64) -> Arc { + Self::with_ttl(capacity, page_size, PAGE_TTL) + } + + fn with_ttl(capacity: u64, page_size: u64, ttl: Duration) -> Arc { + Arc::new(Self { + page_size, + cache: PagedCache::new_with_ttl( + "parquet_pages", + capacity, + page_size, + |_: &(RegionFileId, u64), bytes: &Bytes| { + (mem::size_of::<(RegionFileId, u64)>() + + mem::size_of::() + + mem::size_of::() + + bytes.len()) as u32 + }, + metrics::cache_page_observer("page", "page", |hits, misses| { + (usize::from(hits > 0), usize::from(misses > 0)) + }), + Some(ttl), + ), + }) + } + + /// Recovers the immutable file size from a requested block already in memory. + pub(crate) async fn known_file_size( + &self, + file: RegionFileId, + ranges: &[Range], + ) -> Option { + self.cache.known_file_size(&file, ranges).await + } + + /// Fetches ranges in input order, sharing missing blocks and returning [`PageFetchStats`]. + /// The reader must address the immutable file identified by `file`. + pub(crate) async fn fetch( + &self, + file: RegionFileId, + file_size: u64, + ranges: &[Range], + reader: &R, + ) -> Result<(Vec, PageFetchStats), Arc> { + // Also protect internal callers that supply an explicit block size. + if self.page_size == 0 || self.page_size > MAX_PAGE_SIZE { + return Err(Arc::new(io::Error::new( + io::ErrorKind::InvalidInput, + format!( + "Parquet cache page size {} is outside 1..={MAX_PAGE_SIZE} bytes", + self.page_size + ), + ))); + } + let result = self.cache.fetch(file, file_size, ranges, reader).await?; + metrics::cache_page_fetch_metrics("page", &result.1); + Ok(result) + } + + /// Checks whether a completed block is available without recording a lookup event. + #[cfg(test)] + pub(crate) async fn contains(&self, file: RegionFileId, page_id: u64) -> bool { + self.cache.contains(&file, page_id).await + } +} + +#[cfg(test)] +mod tests { + use super::*; + use async_trait::async_trait; + use common_base::range_read::Metadata; + use std::sync::Mutex; + use store_api::storage::{FileId, RegionId}; + + #[derive(Default)] + struct Reader(Mutex>>); + + #[async_trait] + impl RangeReader for Reader { + async fn metadata(&self) -> io::Result { + Ok(Metadata { + content_length: PAGE_SIZE * 2 + 10, + }) + } + async fn read(&self, range: Range) -> io::Result { + self.0.lock().unwrap().push(range.clone()); + tokio::task::yield_now().await; + Ok(Bytes::from(range.map(|n| n as u8).collect::>())) + } + } + + #[tokio::test] + async fn test_aligns_and_reuses_partial_overlaps_without_merging_reads() { + let cache = ParquetPageCache::new(1024 * 1024, PAGE_SIZE); + let file = RegionFileId::new(RegionId::new(1, 1), FileId::random()); + let size = PAGE_SIZE * 2 + 10; + let reader = Reader::default(); + let (result, stats) = cache + .fetch( + file, + size, + &[1..4, PAGE_SIZE - 2..PAGE_SIZE + 2, size - 3..size], + &reader, + ) + .await + .unwrap(); + assert_eq!(&result[0][..], &[1, 2, 3]); + assert_eq!(&result[1][..], &[254, 255, 0, 1]); + assert_eq!(&result[2][..], &[7, 8, 9]); + assert_eq!( + ( + stats.requested_bytes, + stats.unique_requested_bytes, + stats.aligned_bytes + ), + (10, 10, size) + ); + assert_eq!( + reader.0.lock().unwrap().as_slice(), + &[0..PAGE_SIZE, PAGE_SIZE..PAGE_SIZE * 2, PAGE_SIZE * 2..size] + ); + let (_, stats) = cache + .fetch(file, size, &[2..6, 4..8], &reader) + .await + .unwrap(); + assert_eq!( + (stats.cache_hit, stats.cache_miss, stats.cached_bytes), + (1, 0, 8) + ); + assert_eq!(reader.0.lock().unwrap().len(), 3); + } + + #[tokio::test] + async fn file_identity_includes_region_and_concurrent_reads_share() { + let cache = ParquetPageCache::new(1024 * 1024, PAGE_SIZE); + let file_id = FileId::random(); + let first = RegionFileId::new(RegionId::new(1, 1), file_id); + let second = RegionFileId::new(RegionId::new(2, 1), file_id); + let reader = Reader::default(); + let ranges = [0..4, 2..6]; + let (a, b) = tokio::join!( + cache.fetch(first, 100, &ranges, &reader), + cache.fetch(first, 100, &ranges, &reader) + ); + assert_eq!(a.unwrap().0, b.unwrap().0); + assert_eq!(reader.0.lock().unwrap().len(), 1); + cache.fetch(second, 100, &ranges, &reader).await.unwrap(); + assert_eq!(reader.0.lock().unwrap().len(), 2); + assert!(cache.contains(first, 0).await); + assert!(cache.contains(second, 0).await); + } + + #[tokio::test] + async fn configurable_blocks_preserve_cross_block_and_tail_reads() { + let file = RegionFileId::new(RegionId::new(1, 1), FileId::random()); + let size = PAGE_SIZE * 2 + 10; + // Include a non-power-of-two size: the layout only requires a nonzero divisor. + for page_size in [1, 37, 4096, PAGE_SIZE, PAGE_SIZE * 4] { + let cache = ParquetPageCache::new(16 * 1024 * 1024, page_size); + let reader = Reader::default(); + let boundary = page_size.min(size - 1); + let ranges = [ + 0..3, + boundary - 1..boundary + 1, + PAGE_SIZE - 2..PAGE_SIZE + 2, + size - 3..size, + size..size, + ]; + let expected: Vec<_> = ranges + .iter() + .map(|r| Bytes::from(r.clone().map(|n| n as u8).collect::>())) + .collect(); + let (actual, _) = cache.fetch(file, size, &ranges, &reader).await.unwrap(); + assert_eq!(actual, expected); + let reads = reader.0.lock().unwrap().len(); + assert!(reader.0.lock().unwrap().iter().all(|r| r.end <= size)); + let (cached, stats) = cache.fetch(file, size, &ranges, &reader).await.unwrap(); + assert_eq!(cached, expected); + assert_eq!(stats.cache_miss, 0); + assert_eq!(reader.0.lock().unwrap().len(), reads); + } + } + + #[tokio::test] + async fn builder_uses_fixed_block_size() { + use crate::cache::{CacheManager, CacheStrategy}; + + let strategy = CacheStrategy::EnableAll(Arc::new( + CacheManager::builder().page_cache_size(1024 * 1024).build(), + )); + let cache = strategy.parquet_page_cache().unwrap(); + let reader = Reader::default(); + let file = RegionFileId::new(RegionId::new(1, 1), FileId::random()); + cache + .fetch(file, PAGE_SIZE * 2, std::slice::from_ref(&(0..1)), &reader) + .await + .unwrap(); + assert_eq!( + reader.0.lock().unwrap().as_slice(), + std::slice::from_ref(&(0..PAGE_SIZE)) + ); + let disabled = CacheStrategy::EnableAll(Arc::new(CacheManager::builder().build())); + assert!(disabled.parquet_page_cache().is_none()); + } + + #[test] + fn block_size_override_uses_valid_bytes_or_default() { + for (value, expected) in [ + (None, PAGE_SIZE), + (Some(""), PAGE_SIZE), + (Some("64KiB"), PAGE_SIZE), + (Some("0"), PAGE_SIZE), + (Some("-1"), PAGE_SIZE), + (Some("37"), 37), + (Some("131072"), 131072), + ] { + assert_eq!(parse_page_size(value), expected); + } + assert_eq!( + parse_page_size(Some(&MAX_PAGE_SIZE.to_string())), + MAX_PAGE_SIZE + ); + assert_eq!( + parse_page_size(Some(&(MAX_PAGE_SIZE + 1).to_string())), + PAGE_SIZE + ); + } + + #[tokio::test] + async fn expired_pages_reload_without_file_invalidation() { + let cache = ParquetPageCache::with_ttl(1024 * 1024, PAGE_SIZE, Duration::from_millis(50)); + let retained = ParquetPageCache::new(1024 * 1024, PAGE_SIZE); + let file = RegionFileId::new(RegionId::new(1, 1), FileId::random()); + let reader = Reader::default(); + let ranges = [0..4, 4..8]; + let expected = cache.fetch(file, 100, &ranges, &reader).await.unwrap().0; + retained.fetch(file, 100, &ranges, &reader).await.unwrap(); + cache.cache.run_pending_tasks().await; + retained.cache.run_pending_tasks().await; + assert!(retained.contains(file, 0).await); + tokio::time::timeout(Duration::from_secs(5), async { + while cache.contains(file, 0).await { + tokio::time::sleep(Duration::from_millis(5)).await; + } + }) + .await + .unwrap(); + assert!(retained.contains(file, 0).await); + let (actual, stats) = cache.fetch(file, 100, &ranges, &reader).await.unwrap(); + assert_eq!(actual, expected); + assert_eq!(stats.cache_miss, 1); + assert_eq!(reader.0.lock().unwrap().len(), 3); + } + + #[tokio::test] + async fn invalid_block_size_is_rejected_before_io() { + let file = RegionFileId::new(RegionId::new(1, 1), FileId::random()); + for page_size in [0, MAX_PAGE_SIZE + 1, u64::MAX] { + let cache = ParquetPageCache::new(1024, page_size); + let reader = Reader::default(); + let error = cache + .fetch(file, 100, std::slice::from_ref(&(0..4)), &reader) + .await + .unwrap_err(); + assert_eq!(error.kind(), io::ErrorKind::InvalidInput); + assert!(reader.0.lock().unwrap().is_empty()); + } + } + + #[tokio::test] + async fn oversized_page_is_returned_without_retention() { + let cache = ParquetPageCache::new(1, PAGE_SIZE); + let file = RegionFileId::new(RegionId::new(1, 1), FileId::random()); + let reader = Reader::default(); + let result = cache + .fetch(file, 100, &[0..4, 4..8], &reader) + .await + .unwrap(); + assert_eq!(&result.0[0][..], &[0, 1, 2, 3]); + cache.cache.run_pending_tasks().await; + assert!(!cache.contains(file, 0).await); + } +} diff --git a/src/mito2/src/error.rs b/src/mito2/src/error.rs index 91402711136..d9340883d2e 100644 --- a/src/mito2/src/error.rs +++ b/src/mito2/src/error.rs @@ -205,6 +205,13 @@ pub enum Error { location: Location, }, + #[snafu(display("Failed to load Parquet page"))] + LoadParquetPage { + source: Arc, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Failed to write parquet file"))] WriteParquet { #[snafu(source)] @@ -1462,7 +1469,9 @@ impl Error { match self { Error::OpenDal { error, .. } => error.kind() == ErrorKind::NotFound, Error::ManifestDeltaNotFound { .. } => true, - Error::LoadSstMeta { source, .. } => source.is_object_not_found(), + Error::LoadSstMeta { source, .. } | Error::LoadParquetPage { source, .. } => { + source.is_object_not_found() + } _ => false, } } @@ -1585,7 +1594,9 @@ impl ErrorExt for Error { OpenRegion { source, .. } => source.status_code(), WriteParquet { .. } => StatusCode::StorageUnavailable, - WriteGroup { source, .. } | LoadSstMeta { source, .. } => source.status_code(), + WriteGroup { source, .. } + | LoadSstMeta { source, .. } + | LoadParquetPage { source, .. } => source.status_code(), InvalidBatch { .. } => StatusCode::InvalidArguments, InvalidRecordBatch { .. } => StatusCode::InvalidArguments, ConvertVector { source, .. } => source.status_code(), @@ -1741,6 +1752,7 @@ impl ErrorExt for Error { OpenRegion { source, .. } | LoadSstMeta { source, .. } + | LoadParquetPage { source, .. } | WriteGroup { source, .. } | FlushRegion { source, .. } | BuildIndexAsync { source, .. } diff --git a/src/mito2/src/sst/file.rs b/src/mito2/src/sst/file.rs index 6422819607b..337d22699ae 100644 --- a/src/mito2/src/sst/file.rs +++ b/src/mito2/src/sst/file.rs @@ -869,7 +869,7 @@ pub async fn delete_files( access_layer: &AccessLayerRef, cache_manager: &Option, ) -> crate::error::Result<()> { - // Remove meta of the file from cache. + // Remove cached metadata for these immutable files. if let Some(cache) = &cache_manager { for (file_id, _) in file_ids { cache.invalidate_range_results(*file_id); diff --git a/src/mito2/src/sst/parquet.rs b/src/mito2/src/sst/parquet.rs index fec4ede14df..381b1ba4647 100644 --- a/src/mito2/src/sst/parquet.rs +++ b/src/mito2/src/sst/parquet.rs @@ -706,22 +706,21 @@ mod tests { ranges }; - // Cache 4 row groups. + // Every byte block touched by the four row groups is retained. Blocks + // are shared across row groups, while file identity remains isolated. + let pages = cache.parquet_page_cache().unwrap(); for i in 0..4 { - let lookup = cache - .get_page_ranges(handle.file_id().file_id(), i, &get_ranges(i)) - .unwrap(); - assert!(lookup.is_fully_cached()); + for range in get_ranges(i) { + for page in range.start / crate::cache::parquet_page::PAGE_SIZE + ..=(range.end - 1) / crate::cache::parquet_page::PAGE_SIZE + { + assert!(pages.contains(handle.file_id(), page).await); + } + } } - let missing_range = 0..10; - let lookup = cache - .get_page_ranges( - handle.file_id().file_id(), - 5, - std::slice::from_ref(&missing_range), - ) - .unwrap(); - assert_eq!(vec![0..10], lookup.missing_ranges); + let other = + crate::sst::file::RegionFileId::new(handle.file_id().region_id(), FileId::random()); + assert!(!pages.contains(other, 0).await); } #[tokio::test] diff --git a/src/mito2/src/sst/parquet/push_decoder.rs b/src/mito2/src/sst/parquet/push_decoder.rs index 1a1446c13c2..9a53f2eee40 100644 --- a/src/mito2/src/sst/parquet/push_decoder.rs +++ b/src/mito2/src/sst/parquet/push_decoder.rs @@ -14,9 +14,16 @@ //! Push decoder stream implementation for SST parquet files. +use std::io; use std::ops::Range; +use std::sync::{Arc, Mutex}; +use std::time::Instant; -use bytes::{Bytes, BytesMut}; +use async_trait::async_trait; +use common_base::range_read::{Metadata, RangeReader}; +use tokio::sync::OnceCell; + +use bytes::Bytes; use common_recordbatch::QueryMemoryTracker; use datatypes::arrow::record_batch::RecordBatch; use futures::StreamExt; @@ -26,11 +33,13 @@ use parquet::DecodeResult; use parquet::arrow::ProjectionMask; use parquet::arrow::arrow_reader::{ArrowReaderMetadata, RowSelection}; use parquet::arrow::push_decoder::ParquetPushDecoderBuilder; -use snafu::{OptionExt, ResultExt, ensure}; +use snafu::{IntoError, OptionExt, ResultExt, ensure}; +use crate::cache::CacheStrategy; use crate::cache::file_cache::{FileType, IndexKey}; -use crate::cache::{CacheStrategy, PageRangePart}; -use crate::error::{OpenDalSnafu, ReadParquetSnafu, Result, UnexpectedSnafu}; +use crate::error::{ + Error, LoadParquetPageSnafu, OpenDalSnafu, ReadParquetSnafu, Result, UnexpectedSnafu, +}; use crate::metrics::{READ_STAGE_ELAPSED, READ_STAGE_FETCH_PAGES}; use crate::sst::file::RegionFileId; use crate::sst::parquet::helper::fetch_byte_ranges; @@ -50,8 +59,9 @@ pub struct SstParquetRangeFetcher { object_store: ObjectStore, /// Cache strategy for reading pages. cache_strategy: CacheStrategy, - /// Row group index for cache key. - row_group_idx: usize, + /// Immutable SST size, resolved once when no manifest hint is available. + file_size: OnceCell, + file_size_hint: u64, /// Optional metrics for tracking fetch operations. fetch_metrics: Option, /// Shared quota for decoding this fetcher's row group. @@ -65,7 +75,6 @@ impl SstParquetRangeFetcher { file_path: String, object_store: ObjectStore, cache_strategy: CacheStrategy, - row_group_idx: usize, fetch_metrics: Option, ) -> Self { Self { @@ -73,7 +82,8 @@ impl SstParquetRangeFetcher { file_path, object_store, cache_strategy, - row_group_idx, + file_size: OnceCell::new(), + file_size_hint: 0, fetch_metrics, memory_tracker: None, } @@ -85,240 +95,220 @@ impl SstParquetRangeFetcher { self } - /// Fetches byte ranges from page cache, write cache, or object store. - async fn fetch_bytes_with_cache(&self, ranges: Vec>) -> Result> { - let fetch_start = self - .fetch_metrics - .as_ref() - .map(|_| std::time::Instant::now()); + /// Supplies the immutable SST size from its manifest entry. + pub(crate) fn with_file_size(mut self, size: u64) -> Self { + self.file_size_hint = size; + self + } + + async fn resolve_file_size(&self, ranges: &[Range]) -> Result { + self.file_size + .get_or_try_init(|| async { + if self.file_size_hint != 0 { + return Ok(self.file_size_hint); + } + if let Some(cache) = self.cache_strategy.parquet_page_cache() + && let Some(size) = cache.known_file_size(self.region_file_id, ranges).await + { + return Ok(size); + } + if let Some(cache) = self.cache_strategy.write_cache() { + let key = IndexKey::new( + self.region_file_id.region_id(), + self.region_file_id.file_id(), + FileType::Parquet, + ); + let files = cache.file_cache(); + if files.get(key).await.is_some() { + // The write-cache index uses u32 sizes, so stat the actual + // file instead. Eviction or local I/O failure falls back below. + if let Ok(meta) = + files.local_store().stat(&files.cache_file_path(key)).await + { + return Ok(meta.content_length()); + } + } + } + Ok(self + .object_store + .stat(&self.file_path) + .await + .context(OpenDalSnafu)? + .content_length()) + }) + .await + .copied() + } + + /// Fetches decoder ranges through fixed-size pages and the existing local/remote hierarchy. + async fn fetch_bytes_with_cache(&self, ranges: &[Range]) -> Result> { + let start = self.fetch_metrics.as_ref().map(|_| Instant::now()); let _timer = READ_STAGE_FETCH_PAGES.start_timer(); - - let mut page_lookup = self.cache_strategy.get_page_ranges( - self.region_file_id.file_id(), - self.row_group_idx, - &ranges, - ); - if let Some(lookup) = &page_lookup - && lookup.cached_bytes > 0 - && let Some(metrics) = &self.fetch_metrics - { - let mut metrics_data = metrics.data.lock().unwrap(); - metrics_data.page_cache_hit += 1; - metrics_data.pages_to_fetch_mem += lookup.cached_range_count; - metrics_data.page_size_to_fetch_mem += lookup.cached_bytes; - metrics_data.page_size_needed += lookup.cached_bytes; - } - - // Fast path: all requested ranges can be assembled from cached fragments. - if page_lookup - .as_ref() - .map(|lookup| lookup.is_fully_cached()) - .unwrap_or(false) - { - let lookup = page_lookup.take().unwrap(); - if let Some(metrics) = &self.fetch_metrics - && let Some(start) = fetch_start - { - metrics.data.lock().unwrap().total_fetch_elapsed += start.elapsed(); + ensure!( + ranges.iter().all(|r| r.start <= r.end), + UnexpectedSnafu { + reason: "Reversed Parquet byte range" } - return assemble_ranges(&ranges, lookup.cached_parts, &[]); + ); + if ranges.iter().all(Range::is_empty) { + return Ok(vec![Bytes::new(); ranges.len()]); } + let reader = ParquetRangeReader { + fetcher: self, + local: self.fetch_metrics.as_ref().map(|_| Mutex::default()), + remote: self.fetch_metrics.as_ref().map(|_| Mutex::default()), + }; + let output = if let Some(cache) = self.cache_strategy.parquet_page_cache() { + let size = self.resolve_file_size(ranges).await?; + let (output, stats) = cache + .fetch(self.region_file_id, size, ranges, &reader) + .await + .map_err(|error| { + if let Some(source) = crate::cache::find_shared_page_source::(&error) { + LoadParquetPageSnafu.into_error(source) + } else { + ReadParquetSnafu { + path: &self.file_path, + } + .into_error(parquet::errors::ParquetError::External(Box::new(error))) + } + })?; + if let Some(metrics) = &self.fetch_metrics { + let mut data = metrics.data.lock().unwrap(); + data.page_cache_hit += usize::from(stats.cached_bytes > 0); + data.pages_to_fetch_mem += stats.cached_fragments; + data.page_size_to_fetch_mem += stats.cached_bytes; + } + output + } else { + reader.read_ranges(ranges).await? + }; + if let (Some(metrics), Some(local), Some(remote), Some(start)) = + (&self.fetch_metrics, &reader.local, &reader.remote, start) + { + let mut data = metrics.data.lock().unwrap(); + let local = local.lock().unwrap(); + let remote = remote.lock().unwrap(); + data.write_cache_hit += usize::from(local.ranges != 0); + data.cache_miss += usize::from(remote.ranges != 0); + data.pages_to_fetch_write_cache += local.ranges; + data.page_size_to_fetch_write_cache += local.bytes; + data.write_cache_fetch_elapsed += local.elapsed(); + data.pages_to_fetch_store += remote.ranges; + data.page_size_to_fetch_store += remote.bytes; + data.store_fetch_elapsed += remote.elapsed(); + data.page_size_needed += output.iter().map(|b| b.len() as u64).sum::(); + data.total_fetch_elapsed += start.elapsed(); + } + Ok(output) + } +} - let missing_ranges = page_lookup - .as_ref() - .map(|lookup| lookup.missing_ranges.clone()) - .unwrap_or_else(|| ranges.clone()); +/// Aggregates one fetch's physical reads without turning request hit counters +/// into per-I/O counters. Elapsed time spans concurrent reads rather than summing them. +#[derive(Default)] +struct StorageReadMetrics { + ranges: usize, + bytes: u64, + interval: Option<(Instant, Instant)>, +} - // Calculate total range size for metrics. - let (_, unaligned_size) = compute_total_range_size(&missing_ranges); +impl StorageReadMetrics { + fn record(&mut self, ranges: usize, bytes: u64, start: Instant) { + self.ranges += ranges; + self.bytes += bytes; + let end = Instant::now(); + self.interval = Some(match self.interval { + Some((first, last)) => (first.min(start), last.max(end)), + None => (start, end), + }); + } - // Check write cache. + fn elapsed(&self) -> std::time::Duration { + self.interval + .map(|(start, end)| end - start) + .unwrap_or_default() + } +} + +/// Reads one immutable SST through the local write cache or object store. +/// Both cached and uncached reads use the existing range merging. +struct ParquetRangeReader<'a> { + fetcher: &'a SstParquetRangeFetcher, + local: Option>, + remote: Option>, +} + +impl ParquetRangeReader<'_> { + async fn read_ranges(&self, ranges: &[Range]) -> Result> { + let fetcher = self.fetcher; let key = IndexKey::new( - self.region_file_id.region_id(), - self.region_file_id.file_id(), + fetcher.region_file_id.region_id(), + fetcher.region_file_id.file_id(), FileType::Parquet, ); - let fetch_write_cache_start = self - .fetch_metrics - .as_ref() - .map(|_| std::time::Instant::now()); - let write_cache_result = match self.cache_strategy.write_cache() { - Some(cache) => cache.file_cache().read_ranges(key, &missing_ranges).await, - None => None, + let bytes = if fetcher.fetch_metrics.is_none() { + 0 + } else { + compute_total_range_size(ranges).1 }; - - let fetched_pages = match write_cache_result { - Some(data) => { - if let Some(metrics) = &self.fetch_metrics { - let elapsed = fetch_write_cache_start - .map(|start| start.elapsed()) - .unwrap_or_default(); - let range_size_needed: u64 = - missing_ranges.iter().map(|r| r.end - r.start).sum(); - let mut metrics_data = metrics.data.lock().unwrap(); - metrics_data.write_cache_fetch_elapsed += elapsed; - metrics_data.write_cache_hit += 1; - metrics_data.pages_to_fetch_write_cache += missing_ranges.len(); - metrics_data.page_size_to_fetch_write_cache += unaligned_size; - metrics_data.page_size_needed += range_size_needed; - } - data + let start = self.local.as_ref().map(|_| Instant::now()); + if let Some(cache) = fetcher.cache_strategy.write_cache() + && let Some(data) = cache.file_cache().read_ranges(key, ranges).await + { + if let (Some(metrics), Some(start)) = (&self.local, start) { + metrics.lock().unwrap().record(ranges.len(), bytes, start); } - None => { - // Fetch data from object store. - let _timer = READ_STAGE_ELAPSED - .with_label_values(&["cache_miss_read"]) - .start_timer(); - - let start = self - .fetch_metrics - .as_ref() - .map(|_| std::time::Instant::now()); - let data = - fetch_byte_ranges(&self.file_path, self.object_store.clone(), &missing_ranges) - .await - .context(OpenDalSnafu)?; - - if let Some(metrics) = &self.fetch_metrics { - let elapsed = start.map(|start| start.elapsed()).unwrap_or_default(); - let range_size_needed: u64 = - missing_ranges.iter().map(|r| r.end - r.start).sum(); - let mut metrics_data = metrics.data.lock().unwrap(); - metrics_data.store_fetch_elapsed += elapsed; - metrics_data.cache_miss += 1; - metrics_data.pages_to_fetch_store += missing_ranges.len(); - metrics_data.page_size_to_fetch_store += unaligned_size; - metrics_data.page_size_needed += range_size_needed; - } - data - } - }; - ensure!( - fetched_pages.len() == missing_ranges.len(), - UnexpectedSnafu { - reason: format!( - "Invalid parquet range fetch: {} missing ranges but {} fetched byte ranges", - missing_ranges.len(), - fetched_pages.len() - ), - } - ); - - self.cache_strategy.put_page_ranges( - self.region_file_id.file_id(), - self.row_group_idx, - &missing_ranges, - &fetched_pages, - ); - - if let (Some(metrics), Some(start)) = (&self.fetch_metrics, fetch_start) { - metrics.data.lock().unwrap().total_fetch_elapsed += start.elapsed(); + return Ok(data); } - - if let Some(lookup) = page_lookup { - let fetched_parts = missing_ranges - .into_iter() - .zip(fetched_pages) - .map(|(range, bytes)| PageRangePart { range, bytes }) - .collect::>(); - return assemble_ranges(&ranges, lookup.cached_parts, &fetched_parts); + let _timer = READ_STAGE_ELAPSED + .with_label_values(&["cache_miss_read"]) + .start_timer(); + let start = self.remote.as_ref().map(|_| Instant::now()); + let data = fetch_byte_ranges(&fetcher.file_path, fetcher.object_store.clone(), ranges) + .await + .context(OpenDalSnafu)?; + if let (Some(metrics), Some(start)) = (&self.remote, start) { + metrics.lock().unwrap().record(ranges.len(), bytes, start); } - - Ok(fetched_pages) + Ok(data) } } -fn assemble_ranges( - ranges: &[Range], - cached_parts: Vec>, - fetched_parts: &[PageRangePart], -) -> Result> { - ensure!( - ranges.len() == cached_parts.len(), - UnexpectedSnafu { - reason: format!( - "Invalid parquet range assembly: {} requested ranges but {} cached part groups", - ranges.len(), - cached_parts.len() - ), - } - ); +#[async_trait] +impl RangeReader for ParquetRangeReader<'_> { + async fn metadata(&self) -> io::Result { + self.fetcher + .resolve_file_size(&[]) + .await + .map(|content_length| Metadata { content_length }) + .map_err(|error| io::Error::other(Arc::new(error))) + } - ranges - .iter() - .zip(cached_parts) - .map(|(range, mut parts)| { - parts.extend( - fetched_parts - .iter() - .filter_map(|part| overlapping_part(range, part)), - ); - assemble_range(range, parts) + async fn read_vec(&self, ranges: &[Range]) -> io::Result> { + self.read_ranges(ranges) + .await + .map_err(|error| io::Error::other(Arc::new(error))) + } + + async fn read(&self, range: Range) -> io::Result { + let mut data = self + .read_ranges(std::slice::from_ref(&range)) + .await + .map_err(|error| io::Error::other(Arc::new(error)))?; + if data.len() != 1 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "Invalid Parquet read response count", + )); + } + data.pop().ok_or_else(|| { + io::Error::new( + io::ErrorKind::UnexpectedEof, + "Missing Parquet read response", + ) }) - .collect() -} - -fn overlapping_part(range: &Range, part: &PageRangePart) -> Option { - let start = range.start.max(part.range.start); - let end = range.end.min(part.range.end); - if start >= end { - return None; } - - let slice_start = (start - part.range.start) as usize; - let slice_end = (end - part.range.start) as usize; - Some(PageRangePart { - range: start..end, - bytes: part.bytes.slice(slice_start..slice_end), - }) -} - -fn assemble_range(range: &Range, mut parts: Vec) -> Result { - if range.start >= range.end { - return Ok(Bytes::new()); - } - - parts.sort_unstable_by_key(|part| part.range.start); - if parts.len() == 1 && parts[0].range == *range { - return Ok(parts.pop().unwrap().bytes); - } - - let mut cursor = range.start; - let mut output = BytesMut::with_capacity((range.end - range.start) as usize); - for part in parts { - ensure!( - part.range.start <= cursor, - UnexpectedSnafu { - reason: format!( - "Missing cached parquet bytes for range {}..{}, next part starts at {}", - range.start, range.end, part.range.start - ), - } - ); - if part.range.end <= cursor { - continue; - } - - let slice_start = (cursor - part.range.start) as usize; - let slice_end = (part.range.end.min(range.end) - part.range.start) as usize; - output.extend_from_slice(&part.bytes.slice(slice_start..slice_end)); - cursor = part.range.end.min(range.end); - if cursor >= range.end { - break; - } - } - - ensure!( - cursor == range.end, - UnexpectedSnafu { - reason: format!( - "Missing cached parquet bytes for range {}..{}, assembled through {}", - range.start, range.end, cursor - ), - } - ); - - Ok(output.freeze()) } /// Builds a parquet record batch stream driven directly by [ParquetPushDecoderBuilder]. @@ -372,7 +362,7 @@ pub fn build_sst_parquet_record_batch_stream( loop { match decoder.try_decode().context(ReadParquetSnafu { path: &file_path })? { DecodeResult::NeedsData(ranges) => { - let data = fetcher.fetch_bytes_with_cache(ranges.clone()).await?; + let data = fetcher.fetch_bytes_with_cache(&ranges).await?; decoder.push_ranges(ranges, data) .context(ReadParquetSnafu { path: &file_path })?; } @@ -461,7 +451,6 @@ mod tests { } else { CacheStrategy::Disabled }, - row_group_idx, None, ) .with_memory_tracker(Some(tracker.clone())); @@ -520,51 +509,425 @@ mod tests { assert!(skipped.next().await.is_none()); assert_eq!(0, tracker.current()); } + #[tokio::test] + async fn large_local_file_size_is_lossless_and_missing_file_falls_back() { + use crate::cache::{CacheManager, file_cache::IndexValue}; + use crate::test_util::TestEnv; + use common_base::readable_size::ReadableSize; + use common_test_util::temp_dir::create_temp_dir; + use object_store::services::{Fs, Memory}; + use std::io::{Seek, SeekFrom, Write}; + use store_api::storage::{FileId, RegionId}; - #[test] - fn test_assemble_range_from_cached_subrange_and_fetched_tail() { - let cached_parts = vec![vec![PageRangePart { - range: 400..500, - bytes: Bytes::from(vec![1; 100]), - }]]; - let fetched_parts = vec![PageRangePart { - range: 500..600, - bytes: Bytes::from(vec![2; 100]), - }]; + let env = TestEnv::new().await; + let dir = create_temp_dir("large-page-cache"); + let local = ObjectStore::new(Fs::default().root(dir.path().to_str().unwrap())).unwrap(); + let write_cache = env.create_write_cache(local, ReadableSize::gb(8)).await; + let files = write_cache.file_cache(); + let id = RegionFileId::new(RegionId::new(1, 1), FileId::random()); + let key = IndexKey::new(id.region_id(), id.file_id(), FileType::Parquet); + let path = dir.path().join(files.cache_file_path(key)); + std::fs::create_dir_all(path.parent().unwrap()).unwrap(); + let mut file = std::fs::File::create(&path).unwrap(); + let size = 5_u64 * 1024 * 1024 * 1024; + let offset = 4_u64 * 1024 * 1024 * 1024 + 17; + // A sparse file reproduces the u32 overflow without writing 5 GiB. + file.set_len(size).unwrap(); + file.seek(SeekFrom::Start(offset)).unwrap(); + file.write_all(b"page").unwrap(); + drop(file); + files + .put( + key, + IndexValue { + file_size: size as u32, + }, + ) + .await; + let remote = ObjectStore::new(Memory::default()).unwrap(); + let strategy = CacheStrategy::EnableAll(Arc::new( + CacheManager::builder() + .page_cache_size(1024 * 1024) + .write_cache(Some(write_cache)) + .build(), + )); + let fetcher = SstParquetRangeFetcher::new( + id, + "remote.parquet".into(), + remote.clone(), + strategy.clone(), + None, + ); + let range = offset..offset + 4; + let ranges = std::slice::from_ref(&range); + assert_eq!(fetcher.resolve_file_size(ranges).await.unwrap(), size); + assert_eq!( + fetcher.fetch_bytes_with_cache(ranges).await.unwrap(), + vec![Bytes::from_static(b"page")] + ); - let requested = 400..600; - let output = assemble_ranges( - std::slice::from_ref(&requested), - cached_parts, - &fetched_parts, + // Keep the stale index entry but remove its file; a fresh fetcher must stat remote. + std::fs::remove_file(&path).unwrap(); + remote + .write("remote.parquet", b"remote".as_slice()) + .await + .unwrap(); + let fetcher = + SstParquetRangeFetcher::new(id, "remote.parquet".into(), remote, strategy, None); + assert_eq!(fetcher.resolve_file_size(&[]).await.unwrap(), 6); + assert_eq!( + fetcher.fetch_bytes_with_cache(&[0..3, 3..6]).await.unwrap(), + vec![Bytes::from_static(b"rem"), Bytes::from_static(b"ote")] + ); + } + + #[tokio::test] + async fn test_page_fetch_partial_hits_preserve_request_metrics() { + use crate::cache::CacheManager; + use crate::cache::parquet_page::PAGE_SIZE; + use object_store::services::Memory; + use std::sync::Arc; + use store_api::storage::{FileId, RegionId}; + + let store = ObjectStore::new(Memory::default()).unwrap(); + let size = PAGE_SIZE * 2 + 10; + store + .write( + "bytes.parquet", + (0..size).map(|n| n as u8).collect::>(), + ) + .await + .unwrap(); + let manager = Arc::new(CacheManager::builder().page_cache_size(1024 * 1024).build()); + let id = RegionFileId::new(RegionId::new(1, 1), FileId::random()); + let metrics = ParquetFetchMetrics::default(); + let fetcher = SstParquetRangeFetcher::new( + id, + "bytes.parquet".into(), + store.clone(), + CacheStrategy::EnableAll(manager), + Some(metrics.clone()), ) - .unwrap(); - assert_eq!(1, output.len()); - assert_eq!(vec![1; 100].as_slice(), &output[0][..100]); - assert_eq!(vec![2; 100].as_slice(), &output[0][100..]); + .with_file_size(size); + assert_eq!( + fetcher.fetch_bytes_with_cache(&[1..5, 1..1]).await.unwrap()[0], + Bytes::from_static(&[1, 2, 3, 4]) + ); + assert_eq!( + fetcher.fetch_bytes_with_cache(&[3..7, 1..1]).await.unwrap()[0], + Bytes::from_static(&[3, 4, 5, 6]) + ); + assert_eq!( + fetcher + .fetch_bytes_with_cache(&[PAGE_SIZE - 2..PAGE_SIZE + 2, 1..1]) + .await + .unwrap()[0], + Bytes::from_static(&[254, 255, 0, 1]) + ); + { + let data = metrics.data.lock().unwrap(); + assert_eq!(data.page_cache_hit, 2); + assert_eq!(data.cache_miss, 2); + assert_eq!(data.write_cache_hit, 0); + assert_eq!(data.page_size_needed, 12); + assert_eq!(data.page_size_to_fetch_mem, 6); + assert_eq!(data.page_size_to_fetch_store, PAGE_SIZE * 2); + } + // Completed pages serve another fetch without contacting the deleted backing object. + store.delete("bytes.parquet").await.unwrap(); + let result = fetcher + .fetch_bytes_with_cache(&[PAGE_SIZE + 1..PAGE_SIZE + 3, size..size]) + .await + .unwrap(); + assert_eq!(&result[0][..], &[1, 2]); + assert!(result[1].is_empty()); + let disabled = SstParquetRangeFetcher::new( + id, + "bytes.parquet".into(), + store, + CacheStrategy::Disabled, + None, + ); + assert!( + disabled + .fetch_bytes_with_cache(&[1..5, 1..1]) + .await + .is_err() + ); } - #[test] - fn test_assemble_range_returns_single_covering_part_without_copy() { - let bytes = Bytes::from_static(b"abcdef"); - let cached_parts = vec![vec![PageRangePart { - range: 10..16, - bytes: bytes.clone(), - }]]; + #[tokio::test] + async fn test_page_cache_coalesces_storage_reads() { + use std::sync::atomic::{AtomicUsize, Ordering}; - let requested = 10..16; - let output = assemble_ranges(std::slice::from_ref(&requested), cached_parts, &[]).unwrap(); - assert_eq!(bytes, output[0]); + use object_store::layers::mock; + use object_store::services::Memory; + use store_api::storage::{FileId, RegionId}; + + use crate::cache::CacheManager; + use crate::cache::parquet_page::PAGE_SIZE; + + struct CountingReader { + inner: mock::oio::Reader, + reads: Arc, + } + + impl mock::oio::Read for CountingReader { + async fn open( + &self, + range: mock::BytesRange, + ) -> object_store::Result<(mock::RpRead, Box)> + { + self.reads.fetch_add(1, Ordering::Relaxed); + self.inner.open(range).await + } + + async fn read( + &self, + range: mock::BytesRange, + ) -> object_store::Result<(mock::RpRead, object_store::Buffer)> { + self.reads.fetch_add(1, Ordering::Relaxed); + self.inner.read(range).await + } + } + + let reads = Arc::new(AtomicUsize::new(0)); + let layer = mock::MockLayerBuilder::default() + .reader_factory(Arc::new({ + let reads = reads.clone(); + move |_, _, inner| { + Box::new(CountingReader { + inner, + reads: reads.clone(), + }) + } + })) + .build() + .unwrap(); + let store = ObjectStore::new(Memory::default()).unwrap().layer(layer); + let bytes = Bytes::from( + (0..PAGE_SIZE * 128 + 13) + .map(|n| (n * 31 + n / PAGE_SIZE) as u8) + .collect::>(), + ); + store.write("bytes.parquet", bytes.clone()).await.unwrap(); + let manager = Arc::new( + CacheManager::builder() + .page_cache_size(16 * 1024 * 1024) + .build(), + ); + let fetcher = SstParquetRangeFetcher::new( + RegionFileId::new(RegionId::new(1, 1), FileId::random()), + "bytes.parquet".into(), + store, + CacheStrategy::EnableAll(manager), + None, + ) + .with_file_size(bytes.len() as u64); + + // A cold request spanning 65 adjacent cache pages opens one storage read. + let contiguous = 3..PAGE_SIZE * 65 - 3; + for _ in 0..2 { + assert_eq!( + fetcher + .fetch_bytes_with_cache(std::slice::from_ref(&contiguous)) + .await + .unwrap(), + vec![bytes.slice(contiguous.start as usize..contiguous.end as usize)] + ); + assert_eq!(reads.load(Ordering::Relaxed), 1); + } + + // Two distant cold pages stay separate; cached and duplicate ranges keep + // their requested order without adding storage reads. + let scattered = [ + PAGE_SIZE * 120 + 3..PAGE_SIZE * 120 + 17, + PAGE_SIZE * 80 + 7..PAGE_SIZE * 80 + 11, + 1..5, + PAGE_SIZE * 120 + 4..PAGE_SIZE * 120 + 8, + ]; + let expected: Vec<_> = scattered + .iter() + .map(|range| bytes.slice(range.start as usize..range.end as usize)) + .collect(); + for _ in 0..2 { + assert_eq!( + fetcher.fetch_bytes_with_cache(&scattered).await.unwrap(), + expected + ); + assert_eq!(reads.load(Ordering::Relaxed), 3); + } } - #[test] - fn test_assemble_range_clamps_overlapping_part_to_requested_end() { - let parts = vec![PageRangePart { - range: 0..10, - bytes: Bytes::from_static(b"0123456789"), - }]; + #[tokio::test] + async fn test_page_cache_preserves_storage_error_semantics() { + use std::sync::atomic::{AtomicBool, Ordering}; - let output = assemble_range(&(2..5), parts).unwrap(); - assert_eq!(Bytes::from_static(b"234"), output); + use common_error::ext::ErrorExt; + use object_store::layers::mock; + use object_store::services::Memory; + use store_api::storage::{FileId, RegionId}; + + use crate::cache::CacheManager; + + struct FailingReader { + inner: mock::oio::Reader, + fail: Arc, + kind: object_store::ErrorKind, + temporary: bool, + } + + impl FailingReader { + fn error(&self) -> object_store::Error { + let error = object_store::Error::new(self.kind, "page read failed"); + if self.temporary { + error.set_temporary() + } else { + error + } + } + } + + impl mock::oio::Read for FailingReader { + async fn open( + &self, + range: mock::BytesRange, + ) -> object_store::Result<(mock::RpRead, Box)> + { + if self.fail.load(Ordering::Relaxed) { + return Err(self.error()); + } + self.inner.open(range).await + } + + async fn read( + &self, + range: mock::BytesRange, + ) -> object_store::Result<(mock::RpRead, object_store::Buffer)> { + if self.fail.load(Ordering::Relaxed) { + return Err(self.error()); + } + self.inner.read(range).await + } + } + + for (kind, temporary) in [ + (object_store::ErrorKind::PermissionDenied, false), + (object_store::ErrorKind::NotFound, false), + (object_store::ErrorKind::Unexpected, true), + ] { + let mut baseline = None; + for cache_enabled in [false, true] { + let fail = Arc::new(AtomicBool::new(true)); + let layer = mock::MockLayerBuilder::default() + .reader_factory(Arc::new({ + let fail = fail.clone(); + move |_, _, inner| { + Box::new(FailingReader { + inner, + fail: fail.clone(), + kind, + temporary, + }) + } + })) + .build() + .unwrap(); + let store = ObjectStore::new(Memory::default()).unwrap().layer(layer); + store.write("bytes.parquet", vec![7; 100]).await.unwrap(); + let strategy = if cache_enabled { + CacheStrategy::EnableAll(Arc::new( + CacheManager::builder().page_cache_size(1024 * 1024).build(), + )) + } else { + CacheStrategy::Disabled + }; + let fetcher = SstParquetRangeFetcher::new( + RegionFileId::new(RegionId::new(1, 1), FileId::random()), + "bytes.parquet".into(), + store, + strategy, + None, + ) + .with_file_size(100); + let range = 0..4; + let ranges = std::slice::from_ref(&range); + let error = fetcher.fetch_bytes_with_cache(ranges).await.unwrap_err(); + let semantics = ( + error.status_code(), + error.retry_hint(), + error.is_object_not_found(), + ); + if cache_enabled { + assert!(matches!(error, Error::LoadParquetPage { .. })); + assert_eq!(Some(semantics), baseline); + } else { + assert!(matches!(error, Error::OpenDal { .. })); + baseline = Some(semantics); + } + // Failed shared loads must not poison the cache. + fail.store(false, Ordering::Relaxed); + assert_eq!( + fetcher.fetch_bytes_with_cache(ranges).await.unwrap(), + vec![Bytes::from_static(&[7; 4])] + ); + } + } + } + + #[tokio::test] + async fn test_cached_and_disabled_fetch_return_identical_ranges() { + use crate::cache::CacheManager; + use crate::cache::parquet_page::PAGE_SIZE; + use object_store::services::Memory; + use std::sync::Arc; + use store_api::storage::{FileId, RegionId}; + + let store = ObjectStore::new(Memory::default()).unwrap(); + let size = PAGE_SIZE + 10; + store + .write( + "bytes.parquet", + (0..size).map(|n| n as u8).collect::>(), + ) + .await + .unwrap(); + let id = RegionFileId::new(RegionId::new(1, 1), FileId::random()); + let manager = Arc::new(CacheManager::builder().page_cache_size(1024 * 1024).build()); + let cached = SstParquetRangeFetcher::new( + id, + "bytes.parquet".into(), + store.clone(), + CacheStrategy::EnableAll(manager.clone()), + None, + ); + let disabled = SstParquetRangeFetcher::new( + id, + "bytes.parquet".into(), + store.clone(), + CacheStrategy::Disabled, + None, + ); + let compact = SstParquetRangeFetcher::new( + id, + "bytes.parquet".into(), + store, + CacheStrategy::Compaction(manager), + None, + ); + let ranges = vec![size - 3..size, 1..6, PAGE_SIZE - 2..PAGE_SIZE + 3, 4..8]; + let expected = disabled.fetch_bytes_with_cache(&ranges).await.unwrap(); + assert_eq!( + cached.fetch_bytes_with_cache(&ranges).await.unwrap(), + expected + ); + assert_eq!( + cached.fetch_bytes_with_cache(&ranges).await.unwrap(), + expected + ); + assert_eq!( + compact.fetch_bytes_with_cache(&ranges).await.unwrap(), + expected + ); } } diff --git a/src/mito2/src/sst/parquet/reader.rs b/src/mito2/src/sst/parquet/reader.rs index 2ddb66a53f8..3f89d80bb91 100644 --- a/src/mito2/src/sst/parquet/reader.rs +++ b/src/mito2/src/sst/parquet/reader.rs @@ -2029,9 +2029,9 @@ impl RowGroupReaderBuilder { self.file_path.clone(), self.object_store.clone(), self.cache_strategy.clone(), - row_group_idx, fetch_metrics.cloned(), ) + .with_file_size(self.file_handle.meta_ref().file_size) .with_memory_tracker(self.memory_tracker.clone()); build_sst_parquet_record_batch_stream( @@ -3227,7 +3227,6 @@ mod tests { file_path.clone(), object_store, CacheStrategy::Disabled, - 0, None, ); let mut stream = build_sst_parquet_record_batch_stream( From 6491700fe641c67fe182c4c27b68811f2ceca4a6 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Mon, 2 Nov 2026 11:42:56 +0000 Subject: [PATCH 155/239] fix(meta): forward flushed Parquet bytes in region statistics (#10100) Signed-off-by: WenyXu --- src/common/meta/src/datanode.rs | 56 ++++++++++++++++++- src/meta-srv/src/gc/handler.rs | 1 + src/meta-srv/src/gc/mock.rs | 1 + .../handler/collect_leader_region_handler.rs | 1 + src/meta-srv/src/handler/failure_handler.rs | 1 + .../src/handler/persist_stats_handler.rs | 1 + .../src/handler/region_lease_handler.rs | 1 + src/meta-srv/src/selector/weight_compute.rs | 3 + src/standalone/src/information_extension.rs | 1 + 9 files changed, 65 insertions(+), 1 deletion(-) diff --git a/src/common/meta/src/datanode.rs b/src/common/meta/src/datanode.rs index 81135aefe0e..899063cf84a 100644 --- a/src/common/meta/src/datanode.rs +++ b/src/common/meta/src/datanode.rs @@ -102,6 +102,11 @@ pub struct RegionStat { pub region_manifest: RegionManifestInfo, /// The total bytes written of the region since region opened. pub written_bytes: u64, + /// Total Parquet bytes committed by flushes since the region opened. + /// + /// Excludes index files, compaction and imported SSTs. Resets on region reopen. + #[serde(default)] + pub flushed_parquet_bytes: u64, /// The total query CPU time of the region since region opened. /// /// Unit: nanoseconds. @@ -333,6 +338,7 @@ impl From<&api::v1::meta::RegionStat> for RegionStat { index_size: region_stat.index_size, region_manifest: region_stat.manifest.into(), written_bytes: region_stat.written_bytes, + flushed_parquet_bytes: region_stat.flushed_parquet_bytes, query_cpu_time: region_stat.query_cpu_time, query_scanned_bytes: region_stat.query_scanned_bytes, data_topic_latest_entry_id: region_stat.data_topic_latest_entry_id, @@ -576,7 +582,7 @@ mod tests { } #[test] - fn test_stat_val_deserializes_without_query_stats() { + fn test_stat_val_deserializes_without_optional_counters() { let stat = Stat { region_stats: vec![RegionStat { id: RegionId::new(1024, 1), @@ -597,6 +603,7 @@ mod tests { file_removed_cnt: 0, }, written_bytes: 0, + flushed_parquet_bytes: 4096, query_cpu_time: 10, query_scanned_bytes: 20, data_topic_latest_entry_id: 0, @@ -609,12 +616,14 @@ mod tests { let stat_val = DatanodeStatValue { stats: vec![stat] }; let mut value = serde_json::to_value(stat_val).unwrap(); let region_stat = value[0]["region_stats"][0].as_object_mut().unwrap(); + region_stat.remove("flushed_parquet_bytes"); region_stat.remove("query_cpu_time"); region_stat.remove("query_scanned_bytes"); let stat_val: DatanodeStatValue = serde_json::from_value(value).unwrap(); let region_stat = &stat_val.stats[0].region_stats[0]; + assert_eq!(region_stat.flushed_parquet_bytes, 0); assert_eq!(region_stat.query_cpu_time, 0); assert_eq!(region_stat.query_scanned_bytes, 0); } @@ -703,6 +712,51 @@ mod tests { assert_eq!(stat.region_stats[0].role, RegionRole::StagingLeader); } + #[test] + fn test_flushed_parquet_bytes_from_heartbeat() { + for bytes in [0, 4096, u64::MAX] { + let statistic = RegionStatistic { + flushed_parquet_bytes: bytes, + ..Default::default() + }; + let mut region = api::v1::meta::RegionStat::default(); + region.extensions.insert( + REGION_STATISTIC_KEY.to_string(), + serde_json::to_vec(&statistic).unwrap(), + ); + let request = HeartbeatRequest { + header: Some(RequestHeader::default()), + peer: Some(api::v1::meta::Peer { + id: 1, + addr: "127.0.0.1:3001".to_string(), + }), + region_stats: vec![region.clone()], + ..Default::default() + }; + let stat = Stat::try_from(&request).unwrap(); + let stat = &stat.region_stats[0]; + assert_eq!(stat.flushed_parquet_bytes, bytes); + let encoded = serde_json::to_vec(stat).unwrap(); + let decoded: RegionStat = serde_json::from_slice(&encoded).unwrap(); + assert_eq!(&decoded, stat); + + let mut legacy = serde_json::to_value(&statistic).unwrap(); + legacy + .as_object_mut() + .unwrap() + .remove("flushed_parquet_bytes"); + region.extensions.insert( + REGION_STATISTIC_KEY.to_string(), + serde_json::to_vec(&legacy).unwrap(), + ); + assert_eq!(RegionStat::from(®ion).flushed_parquet_bytes, 0); + } + assert_eq!( + RegionStat::from(&api::v1::meta::RegionStat::default()).flushed_parquet_bytes, + 0 + ); + } + #[test] fn test_env_vars_round_trip() { let mut vars = HashMap::new(); diff --git a/src/meta-srv/src/gc/handler.rs b/src/meta-srv/src/gc/handler.rs index 960a05eebce..1e490ebab58 100644 --- a/src/meta-srv/src/gc/handler.rs +++ b/src/meta-srv/src/gc/handler.rs @@ -495,6 +495,7 @@ fn dropped_region_stat(region_id: RegionId) -> RegionStat { file_removed_cnt: 0, }, written_bytes: 0, + flushed_parquet_bytes: 0, query_cpu_time: 0, query_scanned_bytes: 0, data_topic_latest_entry_id: 0, diff --git a/src/meta-srv/src/gc/mock.rs b/src/meta-srv/src/gc/mock.rs index e0409de6e40..c164c538917 100644 --- a/src/meta-srv/src/gc/mock.rs +++ b/src/meta-srv/src/gc/mock.rs @@ -453,6 +453,7 @@ fn mock_region_stat( min_timestamp: None, max_timestamp: None, written_bytes: 0, + flushed_parquet_bytes: 0, query_cpu_time: 0, query_scanned_bytes: 0, } diff --git a/src/meta-srv/src/handler/collect_leader_region_handler.rs b/src/meta-srv/src/handler/collect_leader_region_handler.rs index 6ae753e4808..0b886a489bf 100644 --- a/src/meta-srv/src/handler/collect_leader_region_handler.rs +++ b/src/meta-srv/src/handler/collect_leader_region_handler.rs @@ -91,6 +91,7 @@ mod tests { min_timestamp: None, max_timestamp: None, written_bytes: 0, + flushed_parquet_bytes: 0, query_cpu_time: 0, query_scanned_bytes: 0, } diff --git a/src/meta-srv/src/handler/failure_handler.rs b/src/meta-srv/src/handler/failure_handler.rs index d1d287d13ff..243f09caa23 100644 --- a/src/meta-srv/src/handler/failure_handler.rs +++ b/src/meta-srv/src/handler/failure_handler.rs @@ -100,6 +100,7 @@ mod tests { data_topic_latest_entry_id: 0, metadata_topic_latest_entry_id: 0, written_bytes: 0, + flushed_parquet_bytes: 0, query_cpu_time: 0, query_scanned_bytes: 0, min_timestamp: None, diff --git a/src/meta-srv/src/handler/persist_stats_handler.rs b/src/meta-srv/src/handler/persist_stats_handler.rs index 472072aa42b..b96af08c12f 100644 --- a/src/meta-srv/src/handler/persist_stats_handler.rs +++ b/src/meta-srv/src/handler/persist_stats_handler.rs @@ -298,6 +298,7 @@ mod tests { file_removed_cnt: 0, }, written_bytes, + flushed_parquet_bytes: 0, query_cpu_time: 0, query_scanned_bytes: 0, data_topic_latest_entry_id: 200, diff --git a/src/meta-srv/src/handler/region_lease_handler.rs b/src/meta-srv/src/handler/region_lease_handler.rs index 98ddd39cc86..e4338720216 100644 --- a/src/meta-srv/src/handler/region_lease_handler.rs +++ b/src/meta-srv/src/handler/region_lease_handler.rs @@ -180,6 +180,7 @@ mod test { data_topic_latest_entry_id: 0, metadata_topic_latest_entry_id: 0, written_bytes: 0, + flushed_parquet_bytes: 0, query_cpu_time: 0, query_scanned_bytes: 0, min_timestamp: None, diff --git a/src/meta-srv/src/selector/weight_compute.rs b/src/meta-srv/src/selector/weight_compute.rs index 27dcdf99b97..d39fda3b8e5 100644 --- a/src/meta-srv/src/selector/weight_compute.rs +++ b/src/meta-srv/src/selector/weight_compute.rs @@ -200,6 +200,7 @@ mod tests { data_topic_latest_entry_id: 0, metadata_topic_latest_entry_id: 0, written_bytes: 0, + flushed_parquet_bytes: 0, query_cpu_time: 0, query_scanned_bytes: 0, min_timestamp: None, @@ -234,6 +235,7 @@ mod tests { data_topic_latest_entry_id: 0, metadata_topic_latest_entry_id: 0, written_bytes: 0, + flushed_parquet_bytes: 0, query_cpu_time: 0, query_scanned_bytes: 0, min_timestamp: None, @@ -268,6 +270,7 @@ mod tests { data_topic_latest_entry_id: 0, metadata_topic_latest_entry_id: 0, written_bytes: 0, + flushed_parquet_bytes: 0, query_cpu_time: 0, query_scanned_bytes: 0, min_timestamp: None, diff --git a/src/standalone/src/information_extension.rs b/src/standalone/src/information_extension.rs index f42d7a84f7f..44fc266c5be 100644 --- a/src/standalone/src/information_extension.rs +++ b/src/standalone/src/information_extension.rs @@ -138,6 +138,7 @@ impl InformationExtension for StandaloneInformationExtension { data_topic_latest_entry_id: region_stat.data_topic_latest_entry_id, metadata_topic_latest_entry_id: region_stat.metadata_topic_latest_entry_id, written_bytes: region_stat.written_bytes, + flushed_parquet_bytes: region_stat.flushed_parquet_bytes, query_cpu_time: region_stat.query_cpu_time, query_scanned_bytes: region_stat.query_scanned_bytes, min_timestamp: region_stat.min_timestamp, From eef2d58a486fcdb304c1bfc08c4a6e01e12b368b Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Mon, 2 Nov 2026 18:19:28 +0000 Subject: [PATCH 156/239] fix(ci): skip unsupported stacked PRs in nightly merge (#10063) Signed-off-by: WenyXu --- .github/nightly-merge.md | 4 +++ .github/scripts/tests/nightly-merge.test.cjs | 32 ++++++++++++++++++-- .github/workflows/nightly-merge.yml | 13 +++++++- 3 files changed, 46 insertions(+), 3 deletions(-) diff --git a/.github/nightly-merge.md b/.github/nightly-merge.md index 08b523dbab2..cc9419ec652 100644 --- a/.github/nightly-merge.md +++ b/.github/nightly-merge.md @@ -13,6 +13,10 @@ GitHub may delay scheduled runs. The label alone controls nightly opt-in; the enqueue job does not check the PR body. PRs must be approved and clean, and have all review threads resolved. Drafts are skipped. +If GitHub rejects enqueueing because a PR belongs to a stack, the workflow skips +it and continues without failing the run or counting it toward the five-PR limit. +Other API errors still fail the run. Dry runs do not call the enqueue API, so +they cannot detect this stack rejection. Already queued PRs and manually enabled auto-merge requests are left alone; Maintainers can use **Merge when ready** to proceed earlier. Unready PRs wait for a later scheduled run; this workflow never enables auto-merge on their behalf. diff --git a/.github/scripts/tests/nightly-merge.test.cjs b/.github/scripts/tests/nightly-merge.test.cjs index 753b458ebce..d6c8792f7fe 100644 --- a/.github/scripts/tests/nightly-merge.test.cjs +++ b/.github/scripts/tests/nightly-merge.test.cjs @@ -26,7 +26,7 @@ const checked = '- [x] Merge during the nightly window '; async function run({ sync = false, body = checked, labels = ['merge nightly'], detail = {}, current = {}, threads = [[true]], count = 1, dryRun = false, - failFirst = false } = {}) { + failFirst = false, enqueueError = Error('head changed') } = {}) { const calls = []; let page = 0; const github = { @@ -43,7 +43,7 @@ async function run({ sync = false, body = checked, labels = ['merge nightly'], graphql: async (query, args) => { if (query.includes('mutation')) { calls.push(['enqueue', args]); - if (failFirst && calls.filter(call => call[0] === 'enqueue').length === 1) throw Error('head changed'); + if (failFirst && calls.filter(call => call[0] === 'enqueue').length === 1) throw enqueueError; return {}; } if (query.includes('reviewThreads')) { @@ -123,3 +123,31 @@ test('label-only opt-in survives PRs without a template and does not require a c // The label is authoritative for admission, even before checkbox synchronization. assert.equal((await run({ body: checked.replace('[x]', '[ ]') })).some(call => call[0] === 'enqueue'), true); }); + +const stackMessage = 'This pull request is part of a stack and must be enqueued using the asynchronous merge REST API. For more information, see https://docs.github.com/rest/pulls/pulls#merge-a-pull-request-asynchronously'; + +test('stacked PR rejection is skipped without consuming the successful enqueue cap', async () => { + const calls = await run({ count: 7, failFirst: true, + enqueueError: Object.assign(Error('GraphQL request failed'), { errors: [{ message: stackMessage }] }) }); + assert.equal(calls.some(c => c[0] === 'failed'), false); + assert.deepEqual(calls.find(c => c[0] === 'summary')[1], [ + ['PR', 'Result'], ['1', 'Skipped: stacked PR'], + ...[2, 3, 4, 5, 6].map(number => [String(number), 'Enqueued']) + ]); + assert.deepEqual(calls.filter(c => c[0] === 'enqueue').map(c => c[1].id), + ['PR_1', 'PR_2', 'PR_3', 'PR_4', 'PR_5', 'PR_6']); +}); + +test('mixed or unstructured errors are not mistaken for a stacked PR rejection', async () => { + for (const enqueueError of [ + Error(stackMessage), + Object.assign(Error('GraphQL request failed'), { errors: [] }), + Object.assign(Error('GraphQL request failed'), { errors: [{ message: stackMessage }, { message: 'Forbidden' }] }) + ]) { + const calls = await run({ count: 2, failFirst: true, enqueueError }); + assert.equal(calls.filter(c => c[0] === 'failed').length, 1); + assert.deepEqual(calls.find(c => c[0] === 'summary')[1], [ + ['PR', 'Result'], ['1', 'Failed: see workflow logs'], ['2', 'Enqueued'] + ]); + } +}); diff --git a/.github/workflows/nightly-merge.yml b/.github/workflows/nightly-merge.yml index 6188d12cd3f..05f5154358b 100644 --- a/.github/workflows/nightly-merge.yml +++ b/.github/workflows/nightly-merge.yml @@ -122,12 +122,23 @@ jobs: } if (!dryRun) { // GitHub enforces branch rules; never bypass or enable auto-merge. - await github.graphql(` + try { + await github.graphql(` mutation($id: ID!, $sha: GitObjectID!) { enqueuePullRequest(input: {pullRequestId: $id, expectedHeadOid: $sha}) { mergeQueueEntry { id } } }`, { id: pr.id, sha: pr.headRefOid }); + } catch (error) { + // Stacked PRs require a different API and are not handled by this workflow. + const stackError = 'This pull request is part of a stack and must be enqueued using the asynchronous merge REST API.'; + if (error.errors?.length && error.errors.every(item => item.message?.startsWith(stackError))) { + core.warning(`PR #${number}: skipped stacked PR`); + rows.push([String(number), 'Skipped: stacked PR']); + continue; + } + throw error; + } } enqueued++; rows.push([String(number), dryRun ? 'Would enqueue' : 'Enqueued']); From 5cecc99cbe847c30fd7cb3e53531afc6fffb33d9 Mon Sep 17 00:00:00 2001 From: Aum Date: Tue, 3 Nov 2026 03:16:10 +0000 Subject: [PATCH 157/239] fix(mito-codec): return the first field from dense decode_leftmost (#10090) DensePrimaryKeyCodec::decode_leftmost decoded the whole key and popped the last value, so codecs with more than one field returned the rightmost primary key column. Decode only the first field instead, which also removes the full-key decoding noted in the TODO. Current callers build single-field codecs and are unaffected. Fixes #10089 Signed-off-by: Aumnertic Co-authored-by: Claude Opus 5.5 --- src/mito-codec/src/row_converter/dense.rs | 46 +++++++++++++++++++++-- 1 file changed, 43 insertions(+), 3 deletions(-) diff --git a/src/mito-codec/src/row_converter/dense.rs b/src/mito-codec/src/row_converter/dense.rs index ad85f4ac8f1..4cf38880f2b 100644 --- a/src/mito-codec/src/row_converter/dense.rs +++ b/src/mito-codec/src/row_converter/dense.rs @@ -730,9 +730,13 @@ impl PrimaryKeyCodec for DensePrimaryKeyCodec { } fn decode_leftmost(&self, bytes: &[u8]) -> Result> { - // TODO(weny, yinwen): avoid decoding the whole primary key. - let mut values = self.decode_dense(bytes)?; - Ok(values.pop().map(|(_, v)| v)) + // Dense keys store fields in primary key order, so only the first + // field needs to be decoded. + let Some((_, field)) = self.ordered_primary_key_columns.first() else { + return Ok(None); + }; + let mut deserializer = Deserializer::new(bytes); + field.deserialize(&mut deserializer).map(Some) } } @@ -1003,6 +1007,42 @@ mod tests { ) } + #[test] + fn test_decode_leftmost_returns_first_field() { + let encoder = DensePrimaryKeyCodec::with_fields(vec![ + (0, SortField::new(ConcreteDataType::string_datatype())), + (1, SortField::new(ConcreteDataType::int32_datatype())), + (2, SortField::new(ConcreteDataType::string_datatype())), + ]); + let row = [ + Value::String("first".into()), + Value::Int32(42), + Value::String("last".into()), + ]; + let encoded = encoder + .encode(row.iter().map(|v| v.as_value_ref())) + .unwrap(); + assert_eq!( + encoder.decode_leftmost(&encoded).unwrap(), + Some(Value::String("first".into())) + ); + + let single = DensePrimaryKeyCodec::with_fields(vec![( + 0, + SortField::new(ConcreteDataType::int32_datatype()), + )]); + let encoded = single + .encode([Value::Int32(7).as_value_ref()].into_iter()) + .unwrap(); + assert_eq!( + single.decode_leftmost(&encoded).unwrap(), + Some(Value::Int32(7)) + ); + + let empty = DensePrimaryKeyCodec::with_fields(vec![]); + assert_eq!(empty.decode_leftmost(&[]).unwrap(), None); + } + #[test] fn test_encoded_value_at() { let data_types = [ From f58d598ca79a9a8461d68ed29f82545d326e6843 Mon Sep 17 00:00:00 2001 From: liyang Date: Tue, 3 Nov 2026 10:42:50 +0000 Subject: [PATCH 158/239] chore(ci): use standard AWS secrets in CI workflows (#10116) --- .github/workflows/dev-build.yml | 8 ++++---- .../workflows/metric-snapshot-s3-bench.yml | 20 +++++++++---------- .github/workflows/nightly-build.yml | 8 ++++---- .github/workflows/nightly-ci.yml | 12 +++++------ .github/workflows/oswal-s3-bench.yml | 16 +++++++-------- .github/workflows/release.yml | 8 ++++---- 6 files changed, 36 insertions(+), 36 deletions(-) diff --git a/.github/workflows/dev-build.yml b/.github/workflows/dev-build.yml index 563435b708e..da514f27def 100644 --- a/.github/workflows/dev-build.yml +++ b/.github/workflows/dev-build.yml @@ -185,8 +185,8 @@ jobs: large-page-size: ${{ inputs.large-page-size }} aws-ci-test-bucket: ${{ vars.AWS_CI_TEST_BUCKET }} aws-region: ${{ vars.AWS_CI_TEST_BUCKET_REGION }} - aws-access-key-id: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} - aws-secret-access-key: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} + aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }} build-linux-arm64-artifacts: name: Build linux-arm64 artifacts @@ -223,8 +223,8 @@ jobs: large-page-size: ${{ inputs.large-page-size }} aws-ci-test-bucket: ${{ vars.AWS_CI_TEST_BUCKET }} aws-region: ${{ vars.AWS_CI_TEST_BUCKET_REGION }} - aws-access-key-id: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} - aws-secret-access-key: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} + aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }} release-images-to-dockerhub: name: Build and push images to DockerHub diff --git a/.github/workflows/metric-snapshot-s3-bench.yml b/.github/workflows/metric-snapshot-s3-bench.yml index e4a12dde69f..724108e4065 100644 --- a/.github/workflows/metric-snapshot-s3-bench.yml +++ b/.github/workflows/metric-snapshot-s3-bench.yml @@ -66,8 +66,8 @@ jobs: - uses: astral-sh/setup-uv@v6 - name: Check S3 access before allocating EC2 env: - AWS_ACCESS_KEY_ID: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} - AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} + AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }} + AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }} run: uv run greptimedb/.github/scripts/metric-snapshot-s3-ops.py probe - name: Resolve the fixed benchmark runner id: runner @@ -148,20 +148,20 @@ jobs: findmnt -no SOURCE,FSTYPE -T "$DISK_ROOT" - name: Probe the real same-region S3 bucket env: - AWS_ACCESS_KEY_ID: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} - AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} + AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }} + AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }} run: uv run greptimedb/.github/scripts/metric-snapshot-s3-ops.py probe - name: Run the 500-table smoke, then one full packed export and restore env: - AWS_ACCESS_KEY_ID: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} - AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} + AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }} + AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }} run: python3 greptimedb/.github/scripts/metric-snapshot-s3-bench.py run - name: Check results for credentials id: scrub if: always() env: - AWS_ACCESS_KEY_ID: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} - AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} + AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }} + AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }} run: python3 greptimedb/.github/scripts/metric-snapshot-s3-ops.py scrub - name: Publish the readable result if: ${{ always() && steps.scrub.outcome == 'success' }} @@ -220,8 +220,8 @@ jobs: - name: Remove and verify only this run's S3 prefix if: always() env: - AWS_ACCESS_KEY_ID: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} - AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} + AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }} + AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }} run: uv run greptimedb/.github/scripts/metric-snapshot-s3-ops.py cleanup - name: Estimate the cumulative cloud cost if: ${{ always() && needs.provision.outputs.started_at != '' }} diff --git a/.github/workflows/nightly-build.yml b/.github/workflows/nightly-build.yml index c78ca4dc5b9..7b173c9a2be 100644 --- a/.github/workflows/nightly-build.yml +++ b/.github/workflows/nightly-build.yml @@ -153,8 +153,8 @@ jobs: image-namespace: ${{ vars.ECR_IMAGE_NAMESPACE }} aws-ci-test-bucket: ${{ vars.AWS_CI_TEST_BUCKET }} aws-region: ${{ vars.AWS_CI_TEST_BUCKET_REGION }} - aws-access-key-id: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} - aws-secret-access-key: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} + aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }} build-linux-arm64-artifacts: name: Build linux-arm64 artifacts @@ -179,8 +179,8 @@ jobs: image-namespace: ${{ vars.ECR_IMAGE_NAMESPACE }} aws-ci-test-bucket: ${{ vars.AWS_CI_TEST_BUCKET }} aws-region: ${{ vars.AWS_CI_TEST_BUCKET_REGION }} - aws-access-key-id: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} - aws-secret-access-key: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} + aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }} release-images-to-dockerhub: name: Build and push images to DockerHub diff --git a/.github/workflows/nightly-ci.yml b/.github/workflows/nightly-ci.yml index ffd1b01d477..2b6862bb827 100644 --- a/.github/workflows/nightly-ci.yml +++ b/.github/workflows/nightly-ci.yml @@ -30,8 +30,8 @@ jobs: data-root: sqlness-test aws-ci-test-bucket: ${{ vars.AWS_CI_TEST_BUCKET }} aws-region: ${{ vars.AWS_CI_TEST_BUCKET_REGION }} - aws-access-key-id: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} - aws-secret-access-key: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} + aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }} - name: Upload sqlness logs if: failure() uses: actions/upload-artifact@v4 @@ -113,8 +113,8 @@ jobs: RUST_MIN_STACK: 8388608 # 8MB CARGO_INCREMENTAL: 0 GT_S3_BUCKET: ${{ vars.AWS_CI_TEST_BUCKET }} - GT_S3_ACCESS_KEY_ID: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} - GT_S3_ACCESS_KEY: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} + GT_S3_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }} + GT_S3_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }} GT_S3_REGION: ${{ vars.AWS_CI_TEST_BUCKET_REGION }} GT_ETCD_TLS_ENDPOINTS: https://127.0.0.1:2378 GT_ETCD_ENDPOINTS: http://127.0.0.1:2379 @@ -167,8 +167,8 @@ jobs: RUST_MIN_STACK: 8388608 # 8 MiB, aligned with Rust CI and release tests. CARGO_INCREMENTAL: 0 GT_S3_BUCKET: ${{ vars.AWS_CI_TEST_BUCKET }} - GT_S3_ACCESS_KEY_ID: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} - GT_S3_ACCESS_KEY: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} + GT_S3_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }} + GT_S3_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }} GT_S3_REGION: ${{ vars.AWS_CI_TEST_BUCKET_REGION }} UNITTEST_LOG_DIR: "__unittest_logs" diff --git a/.github/workflows/oswal-s3-bench.yml b/.github/workflows/oswal-s3-bench.yml index 92243c89e58..81fdc4eefb1 100644 --- a/.github/workflows/oswal-s3-bench.yml +++ b/.github/workflows/oswal-s3-bench.yml @@ -235,8 +235,8 @@ jobs: --runtime-image "$RUNTIME_IMAGE" --instance-type "$INSTANCE_TYPE" --region "$S3_REGION" --prefix "$RUN_PREFIX" - name: Probe the bucket env: - AWS_ACCESS_KEY_ID: ${{ inputs.smoke && 'superpower_ci_user' || secrets.AWS_CI_TEST_ACCESS_KEY_ID }} - AWS_SECRET_ACCESS_KEY: ${{ inputs.smoke && 'superpower_password' || secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} + AWS_ACCESS_KEY_ID: ${{ inputs.smoke && 'superpower_ci_user' || secrets.AWS_ACCESS_KEY_ID }} + AWS_SECRET_ACCESS_KEY: ${{ inputs.smoke && 'superpower_password' || secrets.AWS_SECRET_ACCESS_KEY }} run: uv run greptimedb/.github/scripts/oswal-s3-bench.py probe --bucket "$S3_BUCKET" --region "$S3_REGION" --endpoint "$S3_ENDPOINT" --prefix "$RUN_PREFIX" - name: Put the Raft Engine WAL on the instance store run: bash greptimedb/.github/scripts/oswal-s3-bench-run.sh wal-disk "$WAL_ROOT" @@ -244,8 +244,8 @@ jobs: run: bash greptimedb/.github/scripts/oswal-s3-bench-run.sh generate - name: Load every cell env: - AWS_ACCESS_KEY_ID: ${{ inputs.smoke && 'superpower_ci_user' || secrets.AWS_CI_TEST_ACCESS_KEY_ID }} - AWS_SECRET_ACCESS_KEY: ${{ inputs.smoke && 'superpower_password' || secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} + AWS_ACCESS_KEY_ID: ${{ inputs.smoke && 'superpower_ci_user' || secrets.AWS_ACCESS_KEY_ID }} + AWS_SECRET_ACCESS_KEY: ${{ inputs.smoke && 'superpower_password' || secrets.AWS_SECRET_ACCESS_KEY }} run: bash greptimedb/.github/scripts/oswal-s3-bench-run.sh cells - name: Summary if: always() @@ -253,14 +253,14 @@ jobs: - name: Delete the run's S3 prefix if: always() env: - AWS_ACCESS_KEY_ID: ${{ inputs.smoke && 'superpower_ci_user' || secrets.AWS_CI_TEST_ACCESS_KEY_ID }} - AWS_SECRET_ACCESS_KEY: ${{ inputs.smoke && 'superpower_password' || secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} + AWS_ACCESS_KEY_ID: ${{ inputs.smoke && 'superpower_ci_user' || secrets.AWS_ACCESS_KEY_ID }} + AWS_SECRET_ACCESS_KEY: ${{ inputs.smoke && 'superpower_password' || secrets.AWS_SECRET_ACCESS_KEY }} run: uv run greptimedb/.github/scripts/oswal-s3-bench.py cleanup --bucket "$S3_BUCKET" --region "$S3_REGION" --endpoint "$S3_ENDPOINT" --prefix "$RUN_PREFIX" - name: Check the results hold no credentials if: always() env: - AWS_ACCESS_KEY_ID: ${{ inputs.smoke && 'superpower_ci_user' || secrets.AWS_CI_TEST_ACCESS_KEY_ID }} - AWS_SECRET_ACCESS_KEY: ${{ inputs.smoke && 'superpower_password' || secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} + AWS_ACCESS_KEY_ID: ${{ inputs.smoke && 'superpower_ci_user' || secrets.AWS_ACCESS_KEY_ID }} + AWS_SECRET_ACCESS_KEY: ${{ inputs.smoke && 'superpower_password' || secrets.AWS_SECRET_ACCESS_KEY }} run: bash greptimedb/.github/scripts/oswal-s3-bench-run.sh scrub - name: Upload results (no corpus) if: always() diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index b292b586453..d3c71199bc2 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -255,8 +255,8 @@ jobs: image-namespace: ${{ vars.ECR_IMAGE_NAMESPACE }} aws-ci-test-bucket: ${{ vars.AWS_CI_TEST_BUCKET }} aws-region: ${{ vars.AWS_CI_TEST_BUCKET_REGION }} - aws-access-key-id: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} - aws-secret-access-key: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} + aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }} build-linux-arm64-artifacts: name: Build linux-arm64 artifacts @@ -282,8 +282,8 @@ jobs: large-page-size: true aws-ci-test-bucket: ${{ vars.AWS_CI_TEST_BUCKET }} aws-region: ${{ vars.AWS_CI_TEST_BUCKET_REGION }} - aws-access-key-id: ${{ secrets.AWS_CI_TEST_ACCESS_KEY_ID }} - aws-secret-access-key: ${{ secrets.AWS_CI_TEST_SECRET_ACCESS_KEY }} + aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }} build-linux-riscv64-artifacts: name: Build linux-riscv64 artifacts From 7339a56abdfe6279a448995e3f939d48a37a9b9f Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Tue, 3 Nov 2026 11:33:34 +0000 Subject: [PATCH 159/239] feat: share streaming range results across scans (#9794) * feat: share streaming range results across scans Share post-merge range batches across concurrent readers while preserving stream boundaries, cancellation and retry behavior. Encapsulate range ownership, compaction and publication in RangeResultCache, backed by the common SyncCache primitive. Retain SST retirement fencing and memory accounting. Track live shared-buffer bytes and worker events alongside existing cache metrics. Signed-off-by: WenyXu * fix(mito2): clean up shared range registrations on release Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- src/common/cache/Cargo.toml | 2 +- src/common/cache/src/lib.rs | 2 + src/common/cache/src/sync_cache.rs | 111 + src/mito2/benches/bench_cache_stream.rs | 6 +- src/mito2/src/cache.rs | 109 +- src/mito2/src/cache/range_cache.rs | 1210 ++++++----- .../src/cache/range_cache/shared_range.rs | 413 ++++ src/mito2/src/engine/scan_test.rs | 32 +- src/mito2/src/error.rs | 29 +- src/mito2/src/metrics.rs | 12 + src/mito2/src/read/pruner.rs | 16 +- src/mito2/src/read/range_cache.rs | 1809 +++++++++-------- src/mito2/src/read/scan_util.rs | 68 +- src/mito2/src/read/seq_scan.rs | 15 +- src/mito2/src/read/series_candidate.rs | 179 +- src/mito2/src/read/series_reader.rs | 2 +- src/mito2/src/sst/file_purger.rs | 5 +- src/mito2/src/sst/parquet/file_range.rs | 2 +- 18 files changed, 2600 insertions(+), 1422 deletions(-) create mode 100644 src/common/cache/src/sync_cache.rs create mode 100644 src/mito2/src/cache/range_cache/shared_range.rs diff --git a/src/common/cache/Cargo.toml b/src/common/cache/Cargo.toml index afa9a8f3623..4a2118a6e6b 100644 --- a/src/common/cache/Cargo.toml +++ b/src/common/cache/Cargo.toml @@ -11,7 +11,7 @@ workspace = true bytes.workspace = true common-base.workspace = true futures.workspace = true -moka = { workspace = true, features = ["future"] } +moka = { workspace = true, features = ["future", "sync"] } prometheus.workspace = true [dev-dependencies] diff --git a/src/common/cache/src/lib.rs b/src/common/cache/src/lib.rs index 228f0ca3cf9..7cd68b2c81b 100644 --- a/src/common/cache/src/lib.rs +++ b/src/common/cache/src/lib.rs @@ -15,9 +15,11 @@ //! Shared memory cache primitives with native asynchronous load deduplication. mod cache; +mod sync_cache; pub mod metrics; pub mod page; pub mod paged; pub use cache::{Cache, RemovalCause}; +pub use sync_cache::SyncCache; diff --git a/src/common/cache/src/sync_cache.rs b/src/common/cache/src/sync_cache.rs new file mode 100644 index 00000000000..253a2861b91 --- /dev/null +++ b/src/common/cache/src/sync_cache.rs @@ -0,0 +1,111 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::borrow::Borrow; +use std::hash::Hash; +use std::sync::Arc; + +use crate::RemovalCause; + +/// A weighted completed-value cache for synchronous callers. +/// +/// Uses the same backend and removal causes as [`crate::Cache`]. Callers can +/// serialize publication and invalidation with their own lock without awaiting +/// under that lock. Values should be inexpensive to clone, such as `Arc`. +#[derive(Clone)] +pub struct SyncCache { + inner: moka::sync::Cache, +} + +impl SyncCache +where + K: Eq + Hash + Send + Sync + 'static, + V: Clone + Send + Sync + 'static, +{ + /// Creates a cache with the caller's weight and eviction accounting. + pub fn new(name: &str, capacity: u64, weigher: W, listener: L) -> Self + where + W: Fn(&K, &V) -> u32 + Send + Sync + 'static, + L: Fn(Arc, V, RemovalCause) + Send + Sync + 'static, + { + Self { + inner: moka::sync::Cache::builder() + .name(name) + .max_capacity(capacity) + .weigher(weigher) + .eviction_listener(move |key, value, cause| listener(key, value, cause.into())) + .build(), + } + } + + /// Gets a completed value without starting a load. + pub fn get(&self, key: &Q) -> Option + where + K: Borrow, + Q: Hash + Eq + ?Sized, + { + self.inner.get(key) + } + + /// Submits a completed value to the cache's admission and eviction policy. + /// Insertion does not guarantee retention. + pub fn insert(&self, key: K, value: V) { + self.inner.insert(key, value); + } + + /// Applies pending admission, eviction, and accounting work. + pub fn run_pending_tasks(&self) { + self.inner.run_pending_tasks(); + } + + /// Invalidates a completed entry. Concurrent producers are fenced by the caller. + pub fn invalidate(&self, key: &Q) + where + K: Borrow, + Q: Hash + Eq + ?Sized, + { + self.inner.invalidate(key); + } +} + +#[cfg(test)] +mod tests { + use std::sync::Mutex; + + use super::*; + + #[test] + fn test_borrowed_lookup_and_invalidation() { + let removed = Arc::new(Mutex::new(Vec::new())); + let events = removed.clone(); + let cache = SyncCache::new( + "test", + 1024, + |_: &String, _: &Arc| 1, + move |key, value, cause| events.lock().unwrap().push((key, value, cause)), + ); + let value = Arc::new(7); + cache.insert("a".to_string(), value.clone()); + cache.insert("b".to_string(), Arc::new(9)); + assert!(Arc::ptr_eq(&cache.get("a").unwrap(), &value)); + cache.invalidate("a"); + assert!(cache.get("a").is_none()); + assert_eq!(*cache.get("b").unwrap(), 9); + cache.run_pending_tasks(); + assert_eq!( + *removed.lock().unwrap(), + vec![(Arc::new("a".to_string()), value, RemovalCause::Explicit)] + ); + } +} diff --git a/src/mito2/benches/bench_cache_stream.rs b/src/mito2/benches/bench_cache_stream.rs index f183d10211a..70ea1f32bf4 100644 --- a/src/mito2/benches/bench_cache_stream.rs +++ b/src/mito2/benches/bench_cache_stream.rs @@ -12,10 +12,10 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! Benchmarks for `cache_flat_range_stream` overhead. +//! Benchmarks for the production range-result loader. //! -//! Compares consuming batches from a plain stream vs through the caching wrapper -//! that clones batches for the range cache. +//! Compares a plain stream with cold-cache construction through the loader. +//! Warm reuse, concurrent sharing, and remote I/O are measured separately. //! //! Run with: //! ```sh diff --git a/src/mito2/src/cache.rs b/src/mito2/src/cache.rs index 7161e27b5c2..8339bd5ce08 100644 --- a/src/mito2/src/cache.rs +++ b/src/mito2/src/cache.rs @@ -59,6 +59,7 @@ use crate::cache::cache_size::parquet_meta_size; use crate::cache::file_cache::{FileType, IndexKey}; use crate::cache::index::inverted_index::{InvertedIndexCache, InvertedIndexCacheRef}; use crate::cache::parquet_page::ParquetPageCache; +use crate::cache::range_cache::RangeScanCacheValue; use crate::cache::write_cache::WriteCacheRef; use crate::error::{ CompressObjectSnafu, DecompressObjectSnafu, InvalidMetadataSnafu, InvalidParquetSnafu, @@ -67,7 +68,7 @@ use crate::error::{ use crate::memtable::record_batch_estimated_size; use crate::metrics::{CACHE_BYTES, CACHE_EVICTION, CACHE_HIT, CACHE_MISS}; use crate::read::Batch; -use crate::read::range_cache::{RangeScanCacheKey, RangeScanCacheValue}; +use crate::read::range_cache::RangeScanCacheKey; use crate::read::read_columns::JsonTargetTypes; use crate::sst::file::{RegionFileId, RegionIndexId}; use crate::sst::parquet::PARQUET_METADATA_KEY; @@ -174,27 +175,6 @@ impl RangeResultMemoryLimiter { .ok() } - /// Transfers a load's reservation to its completed result without a gap in - /// accounting. Compaction may shrink buffers or add metadata; growing the - /// reservation must remain nonblocking just like loading more input. - pub(crate) fn try_resize( - &self, - permit: Option, - bytes: usize, - ) -> Option { - let Some(mut permit) = permit else { - return self.try_acquire(bytes); - }; - let required = bytes.div_ceil(self.permit_bytes).max(1); - let reserved = permit.num_permits(); - if required > reserved { - permit.merge(self.try_acquire((required - reserved).checked_mul(self.permit_bytes)?)?); - } else { - drop(permit.split(reserved - required)); - } - Some(permit) - } - pub(crate) async fn acquire(&self, bytes: usize) -> Result> { let permits = bytes.div_ceil(self.permit_bytes).max(1); if permits > self.total_permits { @@ -1022,7 +1002,7 @@ impl CacheStrategy { /// Calls [CacheManager::put_range_result()]. /// It does nothing if the strategy isn't [CacheStrategy::EnableAll]. - #[cfg(any(test, feature = "test"))] + #[cfg(test)] pub(crate) fn put_range_result( &self, key: RangeScanCacheKey, @@ -1033,7 +1013,7 @@ impl CacheStrategy { } } - pub(crate) fn range_cache_worker(&self) -> Option<&Arc> { + pub(crate) fn range_result_cache(&self) -> Option<&RangeResultCacheRef> { match self { Self::EnableAll(manager) => manager.range_result_cache.as_ref(), Self::Compaction(_) | Self::Disabled => None, @@ -1048,22 +1028,9 @@ impl CacheStrategy { } } - pub(crate) fn range_result_memory_limiter(&self) -> Option<&Arc> { - match self { - CacheStrategy::EnableAll(cache_manager) => { - Some(cache_manager.range_result_memory_limiter()) - } - CacheStrategy::Compaction(_) | CacheStrategy::Disabled => None, - } - } - - pub(crate) fn range_result_cache_size(&self) -> Option { - match self { - CacheStrategy::EnableAll(cache_manager) => { - Some(cache_manager.range_result_cache_size()) - } - CacheStrategy::Compaction(_) | CacheStrategy::Disabled => None, - } + #[cfg(test)] + pub(crate) fn range_concat_memory(&self) -> Option<&Arc> { + self.range_result_cache().map(|cache| cache.concat_memory()) } /// Calls [CacheManager::write_cache()]. @@ -1157,11 +1124,10 @@ pub struct CacheManager { /// Cache for time series selectors. selector_result_cache: Option, /// Cache for range scan outputs in flat format. - range_result_cache: Option, + range_result_cache: Option, /// Configured capacity for range scan outputs in flat format. + #[cfg(test)] range_result_cache_size: u64, - /// Shared memory limiter for async range-result cache tasks. - range_result_memory_limiter: Arc, /// Cache for index result. index_result_cache: Option, /// Cache for prefilter result. @@ -1540,7 +1506,7 @@ impl CacheManager { } /// Puts range scan result into cache. - #[cfg(any(test, feature = "test"))] + #[cfg(test)] pub(crate) fn put_range_result( &self, key: RangeScanCacheKey, @@ -1556,10 +1522,7 @@ impl CacheManager { self.range_result_cache.is_some() } - pub(crate) fn range_result_memory_limiter(&self) -> &Arc { - &self.range_result_memory_limiter - } - + #[cfg(test)] pub(crate) fn range_result_cache_size(&self) -> usize { self.range_result_cache_size as usize } @@ -1788,7 +1751,7 @@ impl CacheManagerBuilder { .build() }); let range_result_cache = (self.range_result_cache_size != 0) - .then(|| range_cache::RangeCacheWorker::new(self.range_result_cache_size)); + .then(|| range_cache::RangeResultCache::new(self.range_result_cache_size)); CacheManager { sst_meta_cache, sst_decoded_meta_cache, @@ -1800,11 +1763,8 @@ impl CacheManagerBuilder { puffin_metadata_cache: Some(Arc::new(puffin_metadata_cache)), selector_result_cache, range_result_cache, + #[cfg(test)] range_result_cache_size: self.range_result_cache_size, - range_result_memory_limiter: Arc::new(RangeResultMemoryLimiter::new( - self.range_result_cache_size as usize, - RANGE_RESULT_CONCAT_MEMORY_PERMIT.as_bytes() as usize, - )), index_result_cache, prefilter_result_cache, } @@ -2021,7 +1981,7 @@ type VectorCache = Cache<(ConcreteDataType, Value), VectorRef>; /// Maps (file id, row group id, time series row selector) to [SelectorResultValue]. type SelectorResultCache = Cache>; /// Maps partition-range scan key to cached flat batches. -type RangeResultCache = Arc; +type RangeResultCacheRef = Arc; #[cfg(test)] mod tests { @@ -2042,9 +2002,7 @@ mod tests { use crate::cache::test_util::{ parquet_meta, sst_parquet_meta, sst_parquet_meta_with_region_metadata, }; - use crate::read::range_cache::{ - RangeScanCacheKey, RangeScanCacheValue, ScanRequestFingerprintBuilder, - }; + use crate::read::range_cache::{RangeScanCacheKey, ScanRequestFingerprintBuilder}; use crate::read::read_columns::ReadColumns; use crate::sst::parquet::row_selection::RowGroupSelection; @@ -3014,7 +2972,7 @@ mod tests { } .build(), }; - let value = Arc::new(RangeScanCacheValue::new(Vec::new(), 0)); + let value = Arc::new(RangeScanCacheValue::empty()); assert!(cache.get_range_result(&key).is_none()); cache.put_range_result(key.clone(), value.clone()); @@ -3043,17 +3001,27 @@ mod tests { assert_eq!(cache.range_result_cache_size(), cache_size as usize); assert_eq!( - cache.range_result_memory_limiter().permit_bytes(), + cache + .range_result_cache + .as_ref() + .unwrap() + .concat_memory() + .permit_bytes(), RANGE_RESULT_CONCAT_MEMORY_PERMIT.as_bytes() as usize ); assert_eq!( - cache.range_result_memory_limiter().available_permits(), + cache + .range_result_cache + .as_ref() + .unwrap() + .concat_memory() + .available_permits(), (cache_size as usize).div_ceil(RANGE_RESULT_CONCAT_MEMORY_PERMIT.as_bytes() as usize) ); } #[tokio::test] - async fn range_result_memory_limiter_rejects_oversized_request() { + async fn test_range_concat_memory_rejects_oversized_request() { let limiter = RangeResultMemoryLimiter::new(2 * 1024, 1024); assert_eq!(limiter.available_permits(), 2); @@ -3066,7 +3034,7 @@ mod tests { } #[tokio::test] - async fn range_result_memory_limiter_allows_request_up_to_capacity() { + async fn test_range_concat_memory_allows_request_up_to_capacity() { let limiter = RangeResultMemoryLimiter::new(2 * 1024, 1024); let permit = limiter.acquire(2 * 1024).await.unwrap(); assert_eq!(limiter.available_permits(), 0); @@ -3075,23 +3043,18 @@ mod tests { } #[rstest::rstest] - #[case::empty_result(None, 0, Some(1))] - #[case::new_result(None, 1025, Some(2))] - #[case::unchanged(Some(1024), 1024, Some(1))] - #[case::grow(Some(1024), 2048, Some(2))] - #[case::shrink(Some(3072), 1024, Some(1))] - #[case::exhausted(Some(1024), 4096, None)] - #[case::oversized(Some(1024), usize::MAX, None)] + #[case::empty_result(0, Some(1))] + #[case::rounded_result(1025, Some(2))] + #[case::exhausted(4096, None)] + #[case::oversized(usize::MAX, None)] #[test] - fn range_result_memory_reservation_resizes_without_waiting( - #[case] input_size: Option, + fn test_range_result_memory_reservation_is_nonblocking( #[case] result_size: usize, #[case] expected_permits: Option, ) { let limiter = RangeResultMemoryLimiter::new(4096, 1024); let blocker = limiter.try_acquire(1024).unwrap(); - let input_permit = input_size.map(|size| limiter.try_acquire(size).unwrap()); - let result_permit = limiter.try_resize(input_permit, result_size); + let result_permit = limiter.try_acquire(result_size); assert_eq!( result_permit.as_ref().map(|permit| permit.num_permits()), expected_permits diff --git a/src/mito2/src/cache/range_cache.rs b/src/mito2/src/cache/range_cache.rs index 060af974764..7c58c18cc99 100644 --- a/src/mito2/src/cache/range_cache.rs +++ b/src/mito2/src/cache/range_cache.rs @@ -15,37 +15,478 @@ //! Coordinates range-result cache loads and SST retirement. No I/O or decoding //! runs under the coordinator lock; loaders use the shared query runtime. +mod shared_range; + use std::collections::HashMap; use std::future::Future; +use std::mem; +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex, Weak, mpsc}; -use moka::sync::Cache; +use async_stream::try_stream; +use common_cache::SyncCache; +use common_telemetry::warn; +use datafusion::execution::memory_pool::MemoryPool; +use datatypes::arrow::compute::concat_batches; +use datatypes::arrow::record_batch::RecordBatch; +use snafu::{IntoError, ResultExt}; use store_api::storage::FileId; -use tokio::sync::{OwnedSemaphorePermit, watch}; +use tokio::sync::{OnceCell, OwnedSemaphorePermit}; -use crate::cache::{ - RANGE_RESULT_TYPE, RangeResultMemoryLimiter, range_result_cache_weight, removal_cause_str, -}; +use crate::cache::range_cache::shared_range::{RangeRead, SharedRange}; +use crate::cache::{RANGE_RESULT_TYPE, RangeResultMemoryLimiter, range_result_cache_weight}; +use crate::error::{ComputeArrowSnafu, Result}; use crate::metrics::{CACHE_BYTES, CACHE_EVICTION, RANGE_CACHE_WORKER_EVENTS}; -use crate::read::range_cache::{RangeScanCacheKey, RangeScanCacheValue}; +use crate::read::BoxedRecordBatchStream; +use crate::read::range_cache::RangeScanCacheKey; use crate::sst::file::FileHandle; +use crate::sst::parquet::DEFAULT_READ_BATCH_SIZE; -/// Admission, not a queue: excess distinct misses keep using ordinary scans. -const MAX_IN_FLIGHT_RANGES: usize = 32; +pub(crate) const RANGE_CACHE_COMPACT_THRESHOLD_BYTES: usize = 8 * 1024 * 1024; -type SharedResult = Option>; +/// Separates active scans with different query memory reservations. +/// Completed results are reusable across scopes. The default scope is for +/// sources without a query-specific memory pool. +#[derive(Clone, Copy, Debug, Default, Eq, Hash, PartialEq)] +pub(crate) struct RangeCacheScope(usize); + +impl RangeCacheScope { + /// Uses the identity of the pool retained by the source scan. + pub(crate) fn for_memory_pool(pool: &Arc) -> Self { + Self(Arc::as_ptr(pool) as *const () as usize) + } +} + +/// Reports cache activity to the query that owns the reader or publication. +pub(crate) trait RangeCacheMetricsReport: Send + Sync { + /// Records a completed-result hit for this reader. + fn inc_range_cache_hit(&self); + /// Records a reader without a deliverable completed result. + fn inc_range_cache_miss(&self); + /// Records bytes published by the reader that finishes the shared scan. + fn inc_range_cache_size(&self, bytes: usize); +} + +/// Compacted batches with their original stream boundaries and memory weight. +pub(crate) struct RangeScanCacheValue { + cached_batches: Vec, + estimated_batch_size: usize, + /// Precomputed size of all compacted batches. + estimated_batches_size: usize, +} + +impl RangeScanCacheValue { + #[cfg(test)] + pub(crate) fn empty() -> Self { + Self::new(Vec::new(), 0) + } + + #[cfg(test)] + pub(crate) fn cached_batches(&self) -> impl Iterator { + self.cached_batches + .iter() + .map(|batch| (&batch.batch, batch.slice_lengths.as_slice())) + } + + #[cfg(test)] + pub(crate) fn estimated_batches_size(&self) -> usize { + self.estimated_batches_size + } + + fn new(cached_batches: Vec, estimated_batches_size: usize) -> Self { + Self { + cached_batches, + estimated_batch_size: DEFAULT_READ_BATCH_SIZE, + estimated_batches_size, + } + } + + pub(crate) fn estimated_size(&self) -> usize { + mem::size_of::() + + self.cached_batches.capacity() * mem::size_of::() + + self + .cached_batches + .iter() + .map(CachedBatchSlice::metadata_size) + .sum::() + + self.estimated_batches_size + } +} + +/// Stores completed range results and coordinates reader-owned shared scans. +pub(crate) struct RangeResultCache { + cache: SyncCache>, + state: Mutex, + /// Eviction callbacks and the last scan owner can run while `state` is locked. + /// Queue their cleanup instead of reentering the coordinator lock. + cleanup_tx: mpsc::Sender, + cleanup_rx: Mutex>, + /// Covers loads and completed deliveries, separately from cache residency + /// and the temporary concat-memory limiter. Never wait for this budget. + retention_memory: Arc, + /// Bounds temporary allocations while compacting a completed result. + concat_memory: Arc, +} + +impl RangeResultCache { + #[cfg(test)] + pub(crate) fn retention_memory(&self) -> &Arc { + &self.retention_memory + } + + #[cfg(test)] + pub(crate) fn concat_memory(&self) -> &Arc { + &self.concat_memory + } + + /// Replays a cached result or subscribes to one reader-owned streaming scan. + #[tracing::instrument(name = "range_result_lookup", skip_all)] + pub(crate) async fn read( + self: &Arc, + key: RangeScanCacheKey, + part_metrics: Arc, + files: Vec, + scope: RangeCacheScope, + build: F, + ) -> Result<(BoxedRecordBatchStream, usize)> + where + F: FnOnce() -> Fut, + Fut: Future>, + { + let cache = self; + let load = match cache.get_or_register(key, files, scope) { + RangeCacheLookup::Hit(value) => { + part_metrics.inc_range_cache_hit(); + let batch_size = value.value().estimated_batch_size; + return Ok((cached_flat_range_stream(value), batch_size)); + } + RangeCacheLookup::Loading(load) => load, + RangeCacheLookup::Bypass => { + part_metrics.inc_range_cache_miss(); + return build().await; + } + }; + part_metrics.inc_range_cache_miss(); + let cache = cache.clone(); + let limiter = cache.retention_memory.clone(); + let mut build = Some(build); + let mut owner_cursor = None; + let (shared, batch_size) = load + .shared_range + .get_or_try_init(|| async { + let builder = build.take().ok_or_else(|| { + crate::error::UnexpectedSnafu { + reason: "range scan builder already consumed", + } + .build() + })?; + let (source, batch_size) = Box::pin(builder()).await?; + let shared = SharedRange::new(source, limiter.clone()); + owner_cursor = shared.subscribe().await; + Ok::<_, crate::error::Error>((shared, batch_size)) + }) + .await?; + let shared = shared.clone(); + let batch_size = *batch_size; + let cursor = match owner_cursor { + Some(cursor) => Some(cursor), + None => shared.subscribe().await, + }; + let Some(mut cursor) = cursor else { + // An oversized range has closed admission. This caller has not emitted + // any rows, so its original ordinary stream remains a safe fallback. + let builder = build.ok_or_else(|| { + crate::error::UnexpectedSnafu { + reason: "range scan owner lost its initial cursor", + } + .build() + })?; + return builder().await; + }; + let metrics = part_metrics; + let stream = Box::pin(try_stream! { + // Keep the weak registry entry alive until this reader finishes. + let load = load; + loop { + match shared.read(&mut cursor).await { + RangeRead::Batch(batch) => yield batch, + RangeRead::Failed(error) => { + Err(crate::error::LoadRangeResultSnafu.into_error(error))?; + } + RangeRead::Finished => { + if let Some(batches) = shared.take_complete_batches().await { + // Finalization owns the registration until publication, so + // a new caller can still join during background compaction. + common_runtime::spawn_query(cache.finish_scan(load, batches, batch_size, metrics)); + } + break; + } + } + } + }); + Ok((stream, batch_size)) + } + + /// Creates independent retention and compaction budgets with the cache capacity. + pub(crate) fn new(capacity: u64) -> Arc { + let (tx, rx) = mpsc::channel(); + let eviction_tx = tx.clone(); + let cache = SyncCache::new( + RANGE_RESULT_TYPE, + capacity, + |k: &RangeScanCacheKey, v: &Arc| range_result_cache_weight(k, &v.value), + move |key, entry: Arc, cause| { + CACHE_BYTES + .with_label_values(&[RANGE_RESULT_TYPE]) + .sub(range_result_cache_weight(&key, &entry.value).into()); + CACHE_EVICTION + .with_label_values(&[RANGE_RESULT_TYPE, cause.as_str()]) + .inc(); + let _ = eviction_tx.send(RangeCacheCleanup::Evicted { + key, + generation: entry.generation, + }); + }, + ); + Arc::new(Self { + cache, + state: Mutex::new(State::default()), + cleanup_tx: tx, + cleanup_rx: Mutex::new(rx), + retention_memory: Arc::new(RangeResultMemoryLimiter::new(capacity as usize, 1024)), + concat_memory: Arc::new(RangeResultMemoryLimiter::new( + capacity as usize, + crate::cache::RANGE_RESULT_CONCAT_MEMORY_PERMIT.as_bytes() as usize, + )), + }) + } + + /// Compacts a completed scan, then publishes it under the retirement fence. + async fn finish_scan( + self: Arc, + load: Arc, + batches: Vec, + batch_size: usize, + metrics: Arc, + ) { + let bytes = batches + .iter() + .map(|cached| cached.batch.get_array_memory_size()) + .sum(); + let value = RangeScanCacheValue::new(batches, bytes); + let mut compacted = CacheConcatState::default(); + if let Err(error) = compacted.compact_range(&value, &self.concat_memory).await { + warn!(error; "Failed to compact shared range result"); + return; + } + drop(value); + let mut value = compacted.finish(); + value.estimated_batch_size = batch_size; + let bytes = load.key.0.estimated_size() + value.estimated_size(); + if self.complete(&load, Arc::new(value)) { + metrics.inc_range_cache_size(bytes); + } + } + + fn drain_cleanup(&self, state: &mut State) { + for cleanup in self.cleanup_rx.lock().unwrap().try_iter() { + match cleanup { + RangeCacheCleanup::Evicted { key, generation } => state.forget(&key, generation), + RangeCacheCleanup::Released { key, generation } => { + if state + .flights + .get(&key) + .is_some_and(|entry| entry.generation == generation) + { + state.flights.remove(&key); + } + } + } + } + } + + #[cfg(test)] + pub(crate) fn get(&self, key: &RangeScanCacheKey) -> Option> { + self.cache.get(key).map(|entry| entry.value.clone()) + } + + fn get_for_delivery(&self, key: &RangeScanCacheKey) -> Option { + let entry = self.cache.get(key)?; + // A hit without delivery budget must bypass, not start another load. + Some(match entry.deliver(&self.retention_memory) { + Some(delivery) => RangeCacheLookup::Hit(delivery), + None => { + RANGE_CACHE_WORKER_EVENTS + .with_label_values(&["bypass_memory"]) + .inc(); + RangeCacheLookup::Bypass + } + }) + } + + /// For test/benchmark producers that do not hold live SST handles. + #[cfg(test)] + pub(crate) fn insert(&self, key: RangeScanCacheKey, value: Arc) { + let mut state = self.state.lock().unwrap(); + self.drain_cleanup(&mut state); + self.publish(&mut state, key, value); + } + + fn publish(&self, state: &mut State, key: RangeScanCacheKey, value: Arc) { + let generation = state.next_generation(); + let indexed_key = Arc::new(key.clone()); + for (file, _) in &key.row_groups { + state + .files + .entry(*file) + .or_default() + .insert(indexed_key.clone(), generation); + } + CACHE_BYTES + .with_label_values(&[RANGE_RESULT_TYPE]) + .add(range_result_cache_weight(&key, &value).into()); + self.cache.insert( + key, + Arc::new(Entry { + generation, + value, + delivery: Mutex::new(Weak::new()), + }), + ); + self.drain_cleanup(state); + } + + /// Called at logical retirement, before the last snapshot releases the SST. + /// Existing readers may finish, but their results must not refill this cache. + pub(crate) fn invalidate_file(&self, file: FileId) { + RANGE_CACHE_WORKER_EVENTS + .with_label_values(&["invalidate"]) + .inc(); + let mut state = self.state.lock().unwrap(); + self.drain_cleanup(&mut state); + if let Some(keys) = state.files.remove(&file) { + for (key, generation) in keys { + self.cache.invalidate(key.as_ref()); + state.forget(&key, generation); + } + } + state.flights.retain(|(key, _), flight| { + let Some(flight) = flight.load.upgrade() else { + return false; + }; + if key.row_groups.iter().any(|(id, _)| *id == file) { + flight.invalidated.store(true, Ordering::Relaxed); + } + true + }); + self.drain_cleanup(&mut state); + } + + /// Looks up a retained result or registers one reader-owned streaming scan. + fn get_or_register( + self: &Arc, + key: RangeScanCacheKey, + files: Vec, + scope: RangeCacheScope, + ) -> RangeCacheLookup { + let mut state = self.state.lock().unwrap(); + self.drain_cleanup(&mut state); + if files.iter().any(FileHandle::is_deleted) { + RANGE_CACHE_WORKER_EVENTS + .with_label_values(&["bypass_retired"]) + .inc(); + return RangeCacheLookup::Bypass; + } + if let Some(lookup) = + crate::cache::update_hit_miss(self.get_for_delivery(&key), RANGE_RESULT_TYPE) + { + return lookup; + } + let scoped_key = (key, scope); + if let Some(load) = state + .flights + .get(&scoped_key) + .and_then(|entry| entry.load.upgrade()) + { + RANGE_CACHE_WORKER_EVENTS + .with_label_values(&["coalesced"]) + .inc(); + return RangeCacheLookup::Loading(load); + } + let load = Arc::new(RangeLoad { + shared_range: OnceCell::new(), + key: scoped_key.clone(), + generation: state.next_generation(), + files, + invalidated: AtomicBool::new(false), + cleanup_tx: self.cleanup_tx.clone(), + }); + state.flights.insert( + scoped_key, + RangeRegistration { + generation: load.generation, + load: Arc::downgrade(&load), + }, + ); + RANGE_CACHE_WORKER_EVENTS + .with_label_values(&["started"]) + .inc(); + RangeCacheLookup::Loading(load) + } + + /// Publishes only the current, successful scan of live files. Existing readers + /// may still finish an invalidated scan, but it cannot refill the cache. + fn complete(&self, load: &RangeLoad, value: Arc) -> bool { + // Streaming readers already own their buffer reservations. Cache-only + // residency needs no delivery reservation; cache hits acquire one later. + let mut state = self.state.lock().unwrap(); + self.drain_cleanup(&mut state); + let current = state + .flights + .get(&load.key) + .and_then(|entry| entry.load.upgrade()); + if !current.is_some_and(|current| current.generation == load.generation) { + return false; + } + if load.invalidated.load(Ordering::Relaxed) || load.files.iter().any(FileHandle::is_deleted) + { + RANGE_CACHE_WORKER_EVENTS + .with_label_values(&["late_fill_skipped"]) + .inc(); + return false; + } + self.publish(&mut state, load.key.0.clone(), value); + RANGE_CACHE_WORKER_EVENTS + .with_label_values(&["published"]) + .inc(); + true + } +} + +/// Cached result for one range scan. +#[derive(Debug, Clone)] +struct CachedBatchSlice { + batch: RecordBatch, + slice_lengths: Vec, +} + +impl CachedBatchSlice { + fn metadata_size(&self) -> usize { + self.slice_lengths.capacity() * mem::size_of::() + } +} /// Keeps completed results charged while subscribers or replay streams own them. /// Cache entries retain only a weak reference: after the last reader leaves, /// the resident value is covered solely by the cache's capacity accounting. -pub(crate) struct RangeCacheDelivery { +struct RangeCacheDelivery { value: Arc, _permit: OwnedSemaphorePermit, } impl RangeCacheDelivery { /// Shares a decoded result together with its memory reservation. - pub(crate) fn new(value: Arc, permit: OwnedSemaphorePermit) -> Arc { + fn new(value: Arc, permit: OwnedSemaphorePermit) -> Arc { Arc::new(Self { value, _permit: permit, @@ -53,14 +494,14 @@ impl RangeCacheDelivery { } /// Borrows the result without detaching it from the delivery reservation. - pub(crate) fn value(&self) -> &RangeScanCacheValue { + fn value(&self) -> &RangeScanCacheValue { &self.value } } -pub(crate) enum RangeCacheLookup { +enum RangeCacheLookup { Hit(Arc), - Loading(watch::Receiver), + Loading(Arc), Bypass, } @@ -71,7 +512,7 @@ struct Entry { } impl Entry { - fn deliver(&self, memory: &RangeResultMemoryLimiter) -> SharedResult { + fn deliver(&self, memory: &RangeResultMemoryLimiter) -> Option> { // Serialize cache-hit reservations so readers of the same result share // one charge, including when the last previous reader just left. let mut shared_delivery = self @@ -88,16 +529,48 @@ impl Entry { } } -struct Flight { +type FlightKey = (RangeScanCacheKey, RangeCacheScope); + +/// Readers own the scan; the coordinator only retains a weak registration. +/// Query-pool identity isolates active reservations, not completed cache values. +struct RangeLoad { + shared_range: OnceCell<(Arc, usize)>, + key: FlightKey, generation: u64, - sender: watch::Sender, - invalidated: bool, + files: Vec, + invalidated: AtomicBool, + cleanup_tx: mpsc::Sender, +} + +impl Drop for RangeLoad { + fn drop(&mut self) { + let _ = self.cleanup_tx.send(RangeCacheCleanup::Released { + key: self.key.clone(), + generation: self.generation, + }); + } +} + +struct RangeRegistration { + generation: u64, + load: Weak, +} + +enum RangeCacheCleanup { + Evicted { + key: Arc, + generation: u64, + }, + Released { + key: FlightKey, + generation: u64, + }, } #[derive(Default)] struct State { generation: u64, - flights: HashMap, + flights: HashMap, /// File IDs are globally unique. Index by physical file, not query region: /// a copied region can read an SST whose original region differs from its own. files: HashMap, u64>>, @@ -124,267 +597,105 @@ impl State { } } -pub(crate) struct RangeCacheWorker { - cache: Cache>, - state: Mutex, - /// Moka can call listeners during a cache operation. Queue only cleanup - /// metadata instead of reentering `state` from its eviction listener. - evictions: Mutex, u64)>>, - /// Covers loads and completed deliveries, separately from cache residency - /// and the temporary concat-memory limiter. Never wait for this budget. - pub(crate) load_memory: RangeResultMemoryLimiter, +/// Returns a stream that replays cached record batches. +fn cached_flat_range_stream(delivery: Arc) -> BoxedRecordBatchStream { + Box::pin(try_stream! { + // Borrow through the delivery guard so it survives every yield. Batches + // handed to downstream operators are then owned by the query, not by + // this completed-result backlog. + for cached_batch in &delivery.value().cached_batches { + let mut offset = 0; + for &len in &cached_batch.slice_lengths { + yield cached_batch.batch.slice(offset, len); + offset += len; + } + } + }) } -impl RangeCacheWorker { - pub(crate) fn new(capacity: u64) -> Arc { - let (tx, rx) = mpsc::channel(); - let cache = Cache::builder() - .max_capacity(capacity) - .weigher(|k: &RangeScanCacheKey, v: &Arc| range_result_cache_weight(k, &v.value)) - .eviction_listener(move |key, entry: Arc, cause| { - CACHE_BYTES - .with_label_values(&[RANGE_RESULT_TYPE]) - .sub(range_result_cache_weight(&key, &entry.value).into()); - CACHE_EVICTION - .with_label_values(&[RANGE_RESULT_TYPE, removal_cause_str(cause)]) - .inc(); - let _ = tx.send((key, entry.generation)); - }) - .build(); - Arc::new(Self { - cache, - state: Mutex::new(State::default()), - evictions: Mutex::new(rx), - load_memory: RangeResultMemoryLimiter::new(capacity as usize, 1024), - }) - } +#[derive(Default)] +struct CacheConcatState { + cached_batches: Vec, + estimated_size: usize, +} - fn drain_evictions(&self, state: &mut State) { - for (key, generation) in self.evictions.lock().unwrap().try_iter() { - state.forget(&key, generation); +impl CacheConcatState { + async fn compact( + &mut self, + batches: Vec, + limiter: &crate::cache::RangeResultMemoryLimiter, + ) -> Result<()> { + if batches.is_empty() { + return Ok(()); } + + let input_size = batches + .iter() + .map(RecordBatch::get_array_memory_size) + .sum::(); + let _permit = limiter.acquire(input_size).await?; + + let compacted = compact_record_batches(batches)?; + self.estimated_size += compacted.batch.get_array_memory_size(); + self.cached_batches.push(compacted); + Ok(()) } - #[cfg(test)] - pub(crate) fn get(&self, key: &RangeScanCacheKey) -> Option> { - self.cache.get(key).map(|entry| entry.value.clone()) - } - - fn get_for_delivery(&self, key: &RangeScanCacheKey) -> Option { - let entry = self.cache.get(key)?; - // A hit without delivery budget must bypass, not start another load. - Some(match entry.deliver(&self.load_memory) { - Some(delivery) => RangeCacheLookup::Hit(delivery), - None => { - RANGE_CACHE_WORKER_EVENTS - .with_label_values(&["bypass_memory"]) - .inc(); - RangeCacheLookup::Bypass + async fn compact_range( + &mut self, + value: &RangeScanCacheValue, + limiter: &crate::cache::RangeResultMemoryLimiter, + ) -> Result<()> { + let mut pending = Vec::new(); + let mut rows = 0; + let mut bytes = 0; + for cached in &value.cached_batches { + if cached.slice_lengths.len() > 1 { + self.compact(mem::take(&mut pending), limiter).await?; + self.estimated_size += cached.batch.get_array_memory_size(); + self.cached_batches.push(cached.clone()); + rows = 0; + bytes = 0; + continue; } - }) - } - - /// For test/benchmark producers that do not hold live SST handles. - #[cfg(any(test, feature = "test"))] - pub(crate) fn insert(&self, key: RangeScanCacheKey, value: Arc) { - let Some(permit) = self.load_memory.try_acquire(value.estimated_size()) else { - return; - }; - let delivery = RangeCacheDelivery::new(value, permit); - let mut state = self.state.lock().unwrap(); - self.drain_evictions(&mut state); - self.publish(&mut state, key, &delivery); - } - - fn publish( - &self, - state: &mut State, - key: RangeScanCacheKey, - delivery: &Arc, - ) { - let generation = state.next_generation(); - let indexed_key = Arc::new(key.clone()); - for (file, _) in &key.row_groups { - state - .files - .entry(*file) - .or_default() - .insert(indexed_key.clone(), generation); - } - CACHE_BYTES - .with_label_values(&[RANGE_RESULT_TYPE]) - .add(range_result_cache_weight(&key, &delivery.value).into()); - self.cache.insert( - key, - Arc::new(Entry { - generation, - value: delivery.value.clone(), - delivery: Mutex::new(Arc::downgrade(delivery)), - }), - ); - self.drain_evictions(state); - } - - /// Called at logical retirement, before the last snapshot releases the SST. - /// Existing readers may finish, but their results must not refill this cache. - pub(crate) fn invalidate_file(&self, file: FileId) { - RANGE_CACHE_WORKER_EVENTS - .with_label_values(&["invalidate"]) - .inc(); - let mut state = self.state.lock().unwrap(); - self.drain_evictions(&mut state); - if let Some(keys) = state.files.remove(&file) { - for (key, generation) in keys { - self.cache.invalidate(key.as_ref()); - state.forget(&key, generation); + rows += cached.batch.num_rows(); + bytes += cached.batch.get_array_memory_size(); + pending.push(cached.batch.clone()); + if should_compact_batches(pending.len(), rows, bytes) { + self.compact(mem::take(&mut pending), limiter).await?; + rows = 0; + bytes = 0; } } - for (key, flight) in &mut state.flights { - if key.row_groups.iter().any(|(id, _)| *id == file) { - flight.invalidated = true; - } - } - self.drain_evictions(&mut state); + self.compact(pending, limiter).await } - pub(crate) fn get_or_load( - self: &Arc, - key: RangeScanCacheKey, - files: Vec, - load: F, - ) -> RangeCacheLookup - where - F: Future + Send + 'static, - { - if files.iter().any(FileHandle::is_deleted) { - RANGE_CACHE_WORKER_EVENTS - .with_label_values(&["bypass_retired"]) - .inc(); - return RangeCacheLookup::Bypass; - } - if let Some(lookup) = - crate::cache::update_hit_miss(self.get_for_delivery(&key), RANGE_RESULT_TYPE) - { - return lookup; - } - let mut state = self.state.lock().unwrap(); - self.drain_evictions(&mut state); - // Close the miss/register race with completion and invalidation. - if files.iter().any(FileHandle::is_deleted) { - RANGE_CACHE_WORKER_EVENTS - .with_label_values(&["bypass_retired"]) - .inc(); - return RangeCacheLookup::Bypass; - } - if let Some(lookup) = self.get_for_delivery(&key) { - return lookup; - } - if let Some(flight) = state.flights.get(&key) { - RANGE_CACHE_WORKER_EVENTS - .with_label_values(&["coalesced"]) - .inc(); - return RangeCacheLookup::Loading(flight.sender.subscribe()); - } - if state.flights.len() >= MAX_IN_FLIGHT_RANGES { - RANGE_CACHE_WORKER_EVENTS - .with_label_values(&["bypass_capacity"]) - .inc(); - return RangeCacheLookup::Bypass; - } - let generation = state.next_generation(); - let (sender, receiver) = watch::channel(None); - state.flights.insert( - key.clone(), - Flight { - generation, - sender: sender.clone(), - invalidated: false, - }, - ); - drop(state); - - RANGE_CACHE_WORKER_EVENTS - .with_label_values(&["started"]) - .inc(); - let guard = LoadGuard { - worker: self.clone(), - key, - generation, - }; - common_runtime::spawn_query(async move { - // Losing the first requester does not cancel other subscribers. - // Losing every requester drops the loader and unregisters the flight. - let value = tokio::select! { - biased; - _ = sender.closed() => return, - value = load => value, - }; - guard.complete(value, &files); - }); - RangeCacheLookup::Loading(receiver) + fn finish(self) -> RangeScanCacheValue { + RangeScanCacheValue::new(self.cached_batches, self.estimated_size) } } -/// Also unregisters a cancelled or panicking load. Generation matching keeps a -/// stale task from unregistering a later retry of the same key. -struct LoadGuard { - worker: Arc, - key: RangeScanCacheKey, - generation: u64, +fn should_compact_batches(batches: usize, rows: usize, bytes: usize) -> bool { + batches > 1 && (rows > DEFAULT_READ_BATCH_SIZE || bytes > RANGE_CACHE_COMPACT_THRESHOLD_BYTES) } -impl LoadGuard { - fn complete(&self, value: SharedResult, files: &[FileHandle]) { - let mut state = self.worker.state.lock().unwrap(); - self.worker.drain_evictions(&mut state); - let Some(flight) = state.flights.get(&self.key) else { - return; - }; - if flight.generation != self.generation { - return; - } - let flight = state.flights.remove(&self.key).unwrap(); - let event = if let Some(value) = &value { - if !flight.invalidated && !files.iter().any(FileHandle::is_deleted) { - self.worker.publish(&mut state, self.key.clone(), value); - "published" - } else { - "late_fill_skipped" - } - } else { - "load_skipped" - }; - RANGE_CACHE_WORKER_EVENTS.with_label_values(&[event]).inc(); - // Fan out independently of Moka admission, including to old snapshots. - // Do not hold the coordinator lock while notifying readers. - drop(state); - let _ = flight.sender.send(value); - } -} - -impl Drop for LoadGuard { - fn drop(&mut self) { - let mut state = self.worker.state.lock().unwrap(); - if state - .flights - .get(&self.key) - .is_some_and(|flight| flight.generation == self.generation) - { - state.flights.remove(&self.key); - RANGE_CACHE_WORKER_EVENTS - .with_label_values(&["abandoned"]) - .inc(); - } +fn compact_record_batches(mut batches: Vec) -> Result { + debug_assert!(!batches.is_empty()); + let slice_lengths = batches.iter().map(RecordBatch::num_rows).collect(); + if batches.len() == 1 { + // Arrow's singleton concat aliases its input; a second empty batch lets + // the concat kernel copy sliced data without adding a replay boundary. + batches.push(batches[0].slice(0, 0)); } + let batch = concat_batches(&batches[0].schema(), &batches).context(ComputeArrowSnafu)?; + Ok(CachedBatchSlice { + batch, + slice_lengths, + }) } #[cfg(test)] mod tests { - use std::sync::atomic::{AtomicUsize, Ordering}; - use std::time::Duration; - - use tokio::sync::oneshot; - use super::*; use crate::cache::test_util::range_cache_key; use crate::sst::file::FileMeta; @@ -394,70 +705,16 @@ mod tests { Arc::new(RangeScanCacheValue::new(vec![], 0)) } - fn empty_delivery() -> Arc { - let value = empty_result(); - let memory = RangeResultMemoryLimiter::new(1024, 1024); - let permit = memory.try_acquire(value.estimated_size()).unwrap(); - RangeCacheDelivery::new(value, permit) - } - - fn receiver(lookup: RangeCacheLookup) -> watch::Receiver { - match lookup { - RangeCacheLookup::Loading(receiver) => receiver, - _ => panic!("expected a shared load"), - } - } - - async fn loaded(mut receiver: watch::Receiver) -> SharedResult { - tokio::time::timeout(Duration::from_secs(5), async { - receiver - .wait_for(|value| value.is_some()) - .await - .ok() - .and_then(|value| value.clone()) - }) - .await - .expect("load must finish or cancel") - } - - #[tokio::test] - async fn same_key_misses_share_one_load_even_after_first_reader_leaves() { - let worker = RangeCacheWorker::new(1024 * 1024); - let key = range_cache_key(&[FileId::random()]); - let count = Arc::new(AtomicUsize::new(0)); - let (release, gate) = oneshot::channel(); - let first_count = count.clone(); - let first = receiver(worker.get_or_load(key.clone(), vec![], async move { - first_count.fetch_add(1, Ordering::Relaxed); - gate.await.unwrap(); - Some(empty_delivery()) - })); - let mut followers = Vec::new(); - for _ in 0..16 { - let count = count.clone(); - followers.push(receiver(worker.get_or_load( - key.clone(), - vec![], - async move { - count.fetch_add(1, Ordering::Relaxed); - Some(empty_delivery()) - }, - ))); - } - drop(first); - release.send(()).unwrap(); - let mut results = Vec::new(); - for follower in followers { - results.push(loaded(follower).await.unwrap()); - } - assert_eq!(count.load(Ordering::Relaxed), 1); - assert!(results.iter().all(|value| Arc::ptr_eq(value, &results[0]))); - assert!(Arc::ptr_eq(&worker.get(&key).unwrap(), &results[0].value)); + fn loading(lookup: RangeCacheLookup) -> Arc { + let RangeCacheLookup::Loading(load) = lookup else { + panic!("expected active scan") + }; + load } #[test] - fn retirement_removes_all_dependent_keys_but_preserves_other_files() { - let worker = RangeCacheWorker::new(1024 * 1024); + fn test_retirement_removes_all_dependent_keys_but_preserves_other_files() { + let worker = RangeResultCache::new(1024 * 1024); let a = FileId::random(); let b = FileId::random(); let combined = range_cache_key(&[a, b, a]); @@ -478,8 +735,8 @@ mod tests { } #[test] - fn delayed_replacement_eviction_cannot_hide_a_new_entry_from_invalidation() { - let worker = RangeCacheWorker::new(1024 * 1024); + fn test_delayed_replacement_eviction_cannot_hide_a_new_entry_from_invalidation() { + let worker = RangeResultCache::new(1024 * 1024); let file = FileId::random(); let key = range_cache_key(&[file]); worker.insert(key.clone(), empty_result()); @@ -492,37 +749,166 @@ mod tests { assert!(worker.get(&key).is_none()); } - #[tokio::test] - async fn retired_inflight_result_serves_old_readers_without_refilling() { - let worker = RangeCacheWorker::new(1024 * 1024); + #[test] + fn test_cache_hits_share_delivery_budget_until_the_last_reader_leaves() { + let worker = RangeResultCache::new(1024); let file = FileId::random(); let key = range_cache_key(&[file]); - let (release, gate) = oneshot::channel(); - let pending = receiver(worker.get_or_load(key.clone(), vec![], async move { - gate.await.unwrap(); - Some(empty_delivery()) - })); + worker.insert(key.clone(), empty_result()); + assert_eq!(worker.retention_memory.available_permits(), 1); + + let blocker = worker.retention_memory.try_acquire(1024).unwrap(); + assert!(matches!( + worker.get_or_register(key.clone(), vec![], RangeCacheScope(0)), + RangeCacheLookup::Bypass + )); + assert!(worker.get(&key).is_some()); + drop(blocker); + + let RangeCacheLookup::Hit(first) = + worker.get_or_register(key.clone(), vec![], RangeCacheScope(0)) + else { + panic!("expected an accounted cache hit"); + }; + let RangeCacheLookup::Hit(second) = + worker.get_or_register(key.clone(), vec![], RangeCacheScope(0)) + else { + panic!("readers of the same value must share its reservation"); + }; + assert!(Arc::ptr_eq(&first, &second)); + assert_eq!(worker.retention_memory.available_permits(), 0); + worker.invalidate_file(file); - release.send(()).unwrap(); - assert!(loaded(pending).await.is_some()); + worker.cache.run_pending_tasks(); assert!(worker.get(&key).is_none()); + drop(first); + assert_eq!(worker.retention_memory.available_permits(), 0); + drop(second); + assert_eq!(worker.retention_memory.available_permits(), 1); } #[test] - fn old_snapshot_cannot_start_a_fill_after_retirement() { - let worker = RangeCacheWorker::new(1024 * 1024); - let file = FileHandle::new(FileMeta::default(), Arc::new(NoopFilePurger)); - file.mark_deleted(); - let key = range_cache_key(&[file.file_id().file_id()]); + fn test_streaming_registration_is_scoped_and_reader_owned() { + let worker = RangeResultCache::new(1024 * 1024); + let key = range_cache_key(&[FileId::random()]); + let first = loading(worker.get_or_register(key.clone(), vec![], RangeCacheScope(1))); + let second = loading(worker.get_or_register(key.clone(), vec![], RangeCacheScope(1))); + assert!(Arc::ptr_eq(&first, &second)); + let other_pool = loading(worker.get_or_register(key.clone(), vec![], RangeCacheScope(2))); + assert!(!Arc::ptr_eq(&first, &other_pool)); + let weak = Arc::downgrade(&first); + drop(first); + assert!(weak.upgrade().is_some()); + drop(second); + assert!(weak.upgrade().is_none()); + let retry = loading(worker.get_or_register(key.clone(), vec![], RangeCacheScope(1))); + // Publication must not reserve the active source's buffer twice. + let source_permit = worker.retention_memory.try_acquire(1024 * 1024).unwrap(); + assert!(worker.complete(&retry, empty_result())); + assert!(worker.get(&key).is_some()); assert!(matches!( - worker.get_or_load(key, vec![file], async { Some(empty_delivery()) }), + worker.get_or_register(key.clone(), vec![], RangeCacheScope(3)), RangeCacheLookup::Bypass )); + drop(source_permit); + assert!(matches!( + worker.get_or_register(key, vec![], RangeCacheScope(3)), + RangeCacheLookup::Hit(_) + )); + } + + #[test] + fn test_retirement_blocks_late_publication_and_old_snapshot_registration() { + for mark_handle in [false, true] { + let worker = RangeResultCache::new(1024 * 1024); + let file = FileHandle::new(FileMeta::default(), Arc::new(NoopFilePurger)); + let key = range_cache_key(&[file.file_id().file_id()]); + let load = loading(worker.get_or_register( + key.clone(), + vec![file.clone()], + RangeCacheScope(0), + )); + if mark_handle { + file.mark_deleted(); + } else { + worker.invalidate_file(file.file_id().file_id()); + } + worker.complete(&load, empty_result()); + assert!(worker.get(&key).is_none()); + if mark_handle { + assert!(matches!( + worker.get_or_register(key, vec![file], RangeCacheScope(0)), + RangeCacheLookup::Bypass + )); + } + } + } + + #[test] + fn test_scan_release_cleans_only_the_last_owner_under_the_state_lock() { + let cache = RangeResultCache::new(1024 * 1024); + let key = range_cache_key(&[FileId::random()]); + let first = loading(cache.get_or_register(key.clone(), vec![], RangeCacheScope(0))); + let second = loading(cache.get_or_register(key.clone(), vec![], RangeCacheScope(0))); + let weak = Arc::downgrade(&first); + let mut state = cache.state.lock().unwrap(); + drop(first); + cache.drain_cleanup(&mut state); + assert_eq!(state.flights.len(), 1); + // Last-owner release can happen while invalidation holds this lock. + drop(second); + assert!(weak.upgrade().is_none()); + cache.drain_cleanup(&mut state); + assert!(state.flights.is_empty()); + } + + #[test] + fn test_delayed_scan_release_keeps_a_replacement_registration() { + let cache = RangeResultCache::new(1024 * 1024); + let key = range_cache_key(&[FileId::random()]); + let old = loading(cache.get_or_register(key.clone(), vec![], RangeCacheScope(0))); + drop(old); + let delayed = cache.cleanup_rx.lock().unwrap().try_recv().unwrap(); + // Reproduce release arriving after another caller has replaced the dead Weak. + let new = loading(cache.get_or_register(key.clone(), vec![], RangeCacheScope(0))); + cache.cleanup_tx.send(delayed).unwrap(); + let shared = loading(cache.get_or_register(key, vec![], RangeCacheScope(0))); + assert!(Arc::ptr_eq(&new, &shared)); + drop(new); + drop(shared); + let mut state = cache.state.lock().unwrap(); + cache.drain_cleanup(&mut state); + assert!(state.flights.is_empty()); + } + + #[rstest::rstest] + fn test_distinct_active_scans_share_loads_and_prune_abandoned_registrations( + #[values(512, 2048, 8192)] count: usize, + ) { + let worker = RangeResultCache::new(1024 * 1024); + let keys: Vec<_> = (0..count) + .map(|i| range_cache_key(&[FileId::from_bytes((i as u128).to_le_bytes())])) + .collect(); + let loads: Vec<_> = keys + .iter() + .map(|key| loading(worker.get_or_register(key.clone(), vec![], RangeCacheScope(0)))) + .collect(); + for (key, load) in keys.iter().zip(&loads) { + let shared = loading(worker.get_or_register(key.clone(), vec![], RangeCacheScope(0))); + assert!(Arc::ptr_eq(load, &shared)); + } + let abandoned: Vec<_> = loads.iter().map(Arc::downgrade).collect(); + drop(loads); + assert!(abandoned.iter().all(|load| load.upgrade().is_none())); + let next = loading(worker.get_or_register(keys[0].clone(), vec![], RangeCacheScope(0))); + assert_eq!(worker.state.lock().unwrap().flights.len(), 1); + let shared = loading(worker.get_or_register(keys[0].clone(), vec![], RangeCacheScope(0))); + assert!(Arc::ptr_eq(&next, &shared)); } #[rstest::rstest] #[tokio::test] - async fn index_replacement_retirement_invalidates_and_blocks_old_reader_fills( + async fn test_index_replacement_retirement_invalidates_and_blocks_old_reader_fills( #[values(false, true)] gc_enabled: bool, ) { use common_base::readable_size::ReadableSize; @@ -570,14 +956,10 @@ mod tests { worker.insert(key.clone(), empty_result()); assert!(worker.get(&key).is_some()); let loading_key = key.clone().for_cross_range_merge(); - let (release, gate) = oneshot::channel(); - let pending = receiver(worker.get_or_load( + let pending = loading(worker.get_or_register( loading_key.clone(), vec![old_reader.clone()], - async move { - gate.await.unwrap(); - Some(empty_delivery()) - }, + RangeCacheScope(0), )); let updated = FileMeta { index_version: 1, @@ -587,147 +969,15 @@ mod tests { version.add_files(region.file_purger.clone(), [updated.clone()].into_iter()); version.remove_files([updated].into_iter()); - // R1-C1: neither a pre-existing flight nor a new request through the old + // Neither an existing scan nor a new request through the old // index handle may refill a range after the current handle is retired. assert!(old_reader.is_deleted()); assert!(worker.get(&key).is_none()); assert!(matches!( - worker.get_or_load(key, vec![old_reader], async { Some(empty_delivery()) }), + worker.get_or_register(key, vec![old_reader], RangeCacheScope(0)), RangeCacheLookup::Bypass )); - release.send(()).unwrap(); - assert!(loaded(pending).await.is_some()); + assert!(!worker.complete(&pending, empty_result())); assert!(worker.get(&loading_key).is_none()); } - - #[tokio::test] - async fn admission_rejection_does_not_discard_the_shared_result() { - let worker = RangeCacheWorker::new(1); - let key = range_cache_key(&[FileId::random()]); - let value = empty_result(); - let permit = worker - .load_memory - .try_acquire(value.estimated_size()) - .unwrap(); - let delivery = RangeCacheDelivery::new(value, permit); - let pending = - receiver(worker.get_or_load(key.clone(), vec![], async move { Some(delivery) })); - let delivery = loaded(pending).await.unwrap(); - worker.cache.run_pending_tasks(); - assert!(worker.get(&key).is_none()); - { - let mut state = worker.state.lock().unwrap(); - worker.drain_evictions(&mut state); - assert!(state.files.is_empty()); - } - assert_eq!(worker.load_memory.available_permits(), 0); - drop(delivery); - // Publishing wakes readers before the producer drops its watch sender. - // Its LoadGuard also needs the coordinator lock to finish cleaning up. - tokio::time::timeout(Duration::from_secs(5), async { - while worker.load_memory.available_permits() != 1 { - tokio::task::yield_now().await; - } - }) - .await - .expect("a rejected value must release its reservation after all owners leave"); - } - - #[test] - fn cache_hits_share_delivery_budget_until_the_last_reader_leaves() { - let worker = RangeCacheWorker::new(1024); - let file = FileId::random(); - let key = range_cache_key(&[file]); - worker.insert(key.clone(), empty_result()); - assert_eq!(worker.load_memory.available_permits(), 1); - - let blocker = worker.load_memory.try_acquire(1024).unwrap(); - assert!(matches!( - worker.get_or_load(key.clone(), vec![], async { - panic!("must not load on a hit") - }), - RangeCacheLookup::Bypass - )); - assert!(worker.get(&key).is_some()); - drop(blocker); - - let RangeCacheLookup::Hit(first) = - worker.get_or_load(key.clone(), vec![], async { panic!("must hit the cache") }) - else { - panic!("expected an accounted cache hit"); - }; - let RangeCacheLookup::Hit(second) = - worker.get_or_load(key.clone(), vec![], async { panic!("must reuse delivery") }) - else { - panic!("readers of the same value must share its reservation"); - }; - assert!(Arc::ptr_eq(&first, &second)); - assert_eq!(worker.load_memory.available_permits(), 0); - - worker.invalidate_file(file); - worker.cache.run_pending_tasks(); - assert!(worker.get(&key).is_none()); - drop(first); - assert_eq!(worker.load_memory.available_permits(), 0); - drop(second); - assert_eq!(worker.load_memory.available_permits(), 1); - } - - #[tokio::test] - async fn failed_load_releases_followers_and_allows_a_retry() { - let worker = RangeCacheWorker::new(1024 * 1024); - let key = range_cache_key(&[FileId::random()]); - let failed = receiver(worker.get_or_load(key.clone(), vec![], async { None })); - assert!(loaded(failed).await.is_none()); - let retry = receiver(worker.get_or_load(key, vec![], async { Some(empty_delivery()) })); - assert!(loaded(retry).await.is_some()); - } - - #[tokio::test] - async fn abandoning_all_subscribers_cancels_and_unregisters_the_load() { - let worker = RangeCacheWorker::new(1024 * 1024); - let key = range_cache_key(&[FileId::random()]); - let (started, start_rx) = oneshot::channel(); - let marker = Arc::new(()); - let weak_marker = Arc::downgrade(&marker); - let pending = receiver(worker.get_or_load(key.clone(), vec![], async move { - let _marker = marker; - started.send(()).unwrap(); - std::future::pending().await - })); - start_rx.await.unwrap(); - drop(pending); - tokio::time::timeout(Duration::from_secs(5), async { - while weak_marker.upgrade().is_some() - || !worker.state.lock().unwrap().flights.is_empty() - { - tokio::task::yield_now().await; - } - }) - .await - .expect("cancelled loaders must release resources and the flight key"); - assert!(worker.get(&key).is_none()); - let retry = receiver(worker.get_or_load(key, vec![], async { Some(empty_delivery()) })); - assert!(loaded(retry).await.is_some()); - } - - #[tokio::test] - async fn distinct_loads_are_bounded_without_queueing_misses() { - let worker = RangeCacheWorker::new(1024 * 1024); - let mut subscribers = Vec::new(); - for _ in 0..MAX_IN_FLIGHT_RANGES { - subscribers.push(receiver(worker.get_or_load( - range_cache_key(&[FileId::random()]), - vec![], - std::future::pending(), - ))); - } - assert!(matches!( - worker.get_or_load(range_cache_key(&[FileId::random()]), vec![], async { - Some(empty_delivery()) - }), - RangeCacheLookup::Bypass - )); - drop(subscribers); - } } diff --git a/src/mito2/src/cache/range_cache/shared_range.rs b/src/mito2/src/cache/range_cache/shared_range.rs new file mode 100644 index 00000000000..5a186f50235 --- /dev/null +++ b/src/mito2/src/cache/range_cache/shared_range.rs @@ -0,0 +1,413 @@ +// 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. + +//! A single range scan with independently advancing readers. + +use std::collections::VecDeque; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Weak}; + +use datatypes::arrow::record_batch::RecordBatch; +use futures::future::{BoxFuture, Shared}; +use futures::{FutureExt, TryStreamExt}; +use tokio::sync::{Mutex, OwnedSemaphorePermit}; + +use crate::cache::RangeResultMemoryLimiter; +use crate::cache::range_cache::{self, CachedBatchSlice}; +use crate::error::Error; +use crate::read::BoxedRecordBatchStream; + +type SharedBatch = Result, Arc>; +type NextBatch = Shared>; + +/// The registry must hold only a weak reference: readers own the scan lifetime. +/// The source is polled on demand, so an unused range does not run ahead. +/// Cache admission stops when retention reservations run out. Unread batches +/// remain owned by active readers; this is not a hard limit on their backlog. +pub(crate) struct SharedRange { + source: Arc>>, + buffer: Mutex, + limiter: Arc, +} + +/// A reader's cursor refers to original batch boundaries, including empty batches. +pub(crate) struct RangeCursor { + batch: Arc, +} + +pub(crate) enum RangeRead { + Batch(RecordBatch), + Finished, + Failed(Arc), +} + +impl SharedRange { + /// Wraps one source stream and its full-result retention budget. + pub(crate) fn new( + source: BoxedRecordBatchStream, + limiter: Arc, + ) -> Arc { + Arc::new(Self { + source: Arc::new(Mutex::new(Some(source))), + buffer: Mutex::new(RangeBuffer::default()), + limiter, + }) + } + + /// Register before consuming any data. Once full-result retention stops, + /// new queries use their own ordinary scan rather than joining mid-range. + pub(crate) async fn subscribe(&self) -> Option { + let mut buffer = self.buffer.lock().await; + if buffer.streaming_only || matches!(buffer.finished, Some(Err(_))) { + return None; + } + let batch = Arc::new(AtomicUsize::new(0)); + buffer.readers.push(Arc::downgrade(&batch)); + Some(RangeCursor { batch }) + } + + /// Exactly one reader hands the complete result to the existing compactor. + pub(crate) async fn take_complete_batches(&self) -> Option> { + let mut buffer = self.buffer.lock().await; + if buffer.streaming_only + || buffer.finalization_started + || !matches!(buffer.finished, Some(Ok(()))) + { + return None; + } + buffer.finalization_started = true; + let mut batches: Vec = Vec::new(); + let mut previous = None; + for entry in &buffer.batches { + let allocation = Arc::as_ptr(&entry.allocation); + if previous == Some(allocation) { + if let Some(last) = batches.last_mut() { + last.slice_lengths.push(entry.batch.num_rows()); + } + } else { + batches.push(CachedBatchSlice { + batch: entry.allocation.batch.clone(), + slice_lengths: vec![entry.batch.num_rows()], + }); + previous = Some(allocation); + } + } + Some(batches) + } + + /// Returns the next batch for this reader, polling the shared source as needed. + pub(crate) async fn read(&self, cursor: &mut RangeCursor) -> RangeRead { + loop { + let (index, next) = { + let mut buffer = self.buffer.lock().await; + buffer.release_consumed(); + let position = cursor.batch.load(Ordering::Relaxed); + if let Some(batch) = buffer.batches.get(position - buffer.first_batch) { + let batch = batch.batch.clone(); + cursor.batch.store(position + 1, Ordering::Relaxed); + buffer.release_consumed(); + return RangeRead::Batch(batch); + } + if let Some(result) = &buffer.finished { + return match result { + Ok(()) => RangeRead::Finished, + Err(error) => RangeRead::Failed(error.clone()), + }; + } + let index = buffer.first_batch + buffer.batches.len(); + let next = buffer + .next + .get_or_insert_with(|| { + let source = self.source.clone(); + async move { + let mut source = source.lock().await; + let result = match source.as_mut() { + Some(stream) => stream.try_next().await.map_err(Arc::new), + None => Ok(None), + }; + if !matches!(&result, Ok(Some(_))) { + // Release file readers as soon as the source terminates. + *source = None; + } + result + } + .boxed() + .shared() + }) + .clone(); + (index, next) + }; + + // No buffer lock is held across source I/O. A slower reader can still + // consume buffered batches while another reader awaits the next one. + let result = next.await; + let mut buffer = self.buffer.lock().await; + if index != buffer.first_batch + buffer.batches.len() || buffer.finished.is_some() { + continue; + } + buffer.next = None; + match result { + Ok(Some(batch)) => { + let bytes = batch.get_array_memory_size(); + let permit = if buffer.streaming_only { + None + } else { + self.limiter.try_acquire(bytes) + }; + buffer.retained_bytes += bytes; + crate::metrics::RANGE_CACHE_STAGING_BYTES.add(bytes as i64); + if buffer.streaming_only { + crate::metrics::RANGE_CACHE_STREAMING_BYTES.add(bytes as i64); + } + let missing_permit = permit.is_none(); + let allocation = Arc::new(BufferedAllocation { + batch: batch.clone(), + _permit: permit, + }); + buffer + .batches + .push_back(BufferedBatch { batch, allocation }); + if !buffer.streaming_only { + buffer.pending_batches += 1; + buffer.pending_rows += + buffer.batches.back().map_or(0, |b| b.batch.num_rows()); + buffer.pending_bytes += bytes; + let compact = missing_permit + || range_cache::should_compact_batches( + buffer.pending_batches, + buffer.pending_rows, + buffer.pending_bytes, + ); + let retained = !compact || buffer.compact_pending(&self.limiter); + if !retained { + crate::metrics::RANGE_CACHE_WORKER_EVENTS + .with_label_values(&["budget_fallback"]) + .inc(); + buffer.streaming_only = true; + crate::metrics::RANGE_CACHE_STREAMING_BYTES + .add(buffer.retained_bytes as i64); + } + } + buffer.release_consumed(); + } + Ok(None) => buffer.finished = Some(Ok(())), + Err(error) => buffer.finished = Some(Err(error)), + } + } + } +} + +#[derive(Default)] +struct RangeBuffer { + batches: VecDeque, + first_batch: usize, + pending_batches: usize, + pending_rows: usize, + pending_bytes: usize, + retained_bytes: usize, + next: Option, + finished: Option>>, + streaming_only: bool, + finalization_started: bool, + readers: Vec>, +} + +impl RangeBuffer { + /// Detaches the pending slices before deciding whether to stop cache admission. + /// Each compacted allocation owns one reservation, shared by its batch views. + fn compact_pending(&mut self, limiter: &RangeResultMemoryLimiter) -> bool { + let start = self.batches.len() - self.pending_batches; + let input = self + .batches + .iter() + .skip(start) + .map(|entry| entry.batch.clone()) + .collect(); + let compacted = match range_cache::compact_record_batches(input) { + Ok(compacted) => compacted, + Err(error) => { + common_telemetry::warn!(error; "Failed to compact shared range buffer"); + return false; + } + }; + let old_bytes = self.pending_bytes; + let bytes = compacted.batch.get_array_memory_size(); + // Release the old reservations before the nonblocking replacement attempt. + // Reader-owned batches retain their own Arrow references independently. + self.batches.truncate(start); + let permit = limiter.try_acquire(bytes); + let retained = permit.is_some(); + let allocation = Arc::new(BufferedAllocation { + batch: compacted.batch, + _permit: permit, + }); + let mut offset = 0; + for len in compacted.slice_lengths { + self.batches.push_back(BufferedBatch { + batch: allocation.batch.slice(offset, len), + allocation: allocation.clone(), + }); + offset += len; + } + self.retained_bytes = self.retained_bytes - old_bytes + bytes; + crate::metrics::RANGE_CACHE_STAGING_BYTES.add(bytes as i64 - old_bytes as i64); + self.pending_batches = 0; + self.pending_rows = 0; + self.pending_bytes = 0; + retained + } + + fn release_consumed(&mut self) { + if !self.streaming_only { + return; + } + let mut consumed = self.first_batch + self.batches.len(); + self.readers.retain(|reader| { + if let Some(position) = reader.upgrade() { + consumed = consumed.min(position.load(Ordering::Relaxed)); + true + } else { + false + } + }); + while self.first_batch < consumed { + if let Some(batch) = self.batches.pop_front() + && Arc::strong_count(&batch.allocation) == 1 + { + let bytes = batch.allocation.batch.get_array_memory_size(); + self.retained_bytes -= bytes; + crate::metrics::RANGE_CACHE_STAGING_BYTES.sub(bytes as i64); + crate::metrics::RANGE_CACHE_STREAMING_BYTES.sub(bytes as i64); + } + self.first_batch += 1; + } + } +} + +impl Drop for RangeBuffer { + fn drop(&mut self) { + crate::metrics::RANGE_CACHE_STAGING_BYTES.sub(self.retained_bytes as i64); + if self.streaming_only { + crate::metrics::RANGE_CACHE_STREAMING_BYTES.sub(self.retained_bytes as i64); + } + } +} + +struct BufferedBatch { + batch: RecordBatch, + allocation: Arc, +} + +struct BufferedAllocation { + batch: RecordBatch, + _permit: Option, +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use datatypes::arrow::array::{ArrayRef, Int64Array}; + use futures::stream; + use tokio::sync::Notify; + + use super::*; + + fn batch(value: i64) -> RecordBatch { + let values = Arc::new(Int64Array::from(vec![value])) as ArrayRef; + RecordBatch::try_from_iter([("value", values)]).unwrap() + } + + fn value(read: RangeRead) -> i64 { + match read { + RangeRead::Batch(batch) => batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + _ => panic!("expected a batch"), + } + } + + #[tokio::test] + async fn test_pending_source_does_not_block_buffered_readers() { + let produced = Arc::new(AtomicUsize::new(0)); + let counter = produced.clone(); + let gate = Arc::new(Notify::new()); + let ready = gate.clone(); + let source = async_stream::try_stream! { + counter.fetch_add(1, Ordering::Relaxed); + yield batch(1); + ready.notified().await; + counter.fetch_add(1, Ordering::Relaxed); + yield batch(2); + }; + let limiter = Arc::new(RangeResultMemoryLimiter::new(4096, 1)); + let shared = SharedRange::new(Box::pin(source), limiter.clone()); + let mut first = shared.subscribe().await.unwrap(); + let mut second = shared.subscribe().await.unwrap(); + assert_eq!(value(shared.read(&mut first).await), 1); + { + let next = shared.read(&mut first); + tokio::pin!(next); + assert!(futures::poll!(&mut next).is_pending()); + assert_eq!(value(shared.read(&mut second).await), 1); + } + drop(first); + gate.notify_one(); + assert_eq!(value(shared.read(&mut second).await), 2); + assert!(matches!( + shared.read(&mut second).await, + RangeRead::Finished + )); + assert_eq!(produced.load(Ordering::Relaxed), 2); + drop(shared); + assert_eq!(limiter.available_permits(), 4096); + } + + #[tokio::test] + async fn test_slow_reader_backlog_grows_and_is_reclaimed() { + let batches = (0..256).map(batch).collect::>(); + let expected_bytes = batches + .iter() + .map(RecordBatch::get_array_memory_size) + .sum::(); + let shared = SharedRange::new( + Box::pin(stream::iter(batches.into_iter().map(Ok))), + Arc::new(RangeResultMemoryLimiter::new(1, 1)), + ); + let mut fast = shared.subscribe().await.unwrap(); + let mut slow = shared.subscribe().await.unwrap(); + // A paused subscriber must not stop the other reader or discard its data. + for expected in 0..256 { + assert_eq!(value(shared.read(&mut fast).await), expected); + } + assert!(matches!(shared.read(&mut fast).await, RangeRead::Finished)); + { + let buffer = shared.buffer.lock().await; + assert!(buffer.streaming_only); + assert_eq!(buffer.retained_bytes, expected_bytes); + assert_eq!(buffer.batches.len(), 256); + } + assert!(shared.subscribe().await.is_none()); + for expected in 0..256 { + assert_eq!(value(shared.read(&mut slow).await), expected); + } + assert!(matches!(shared.read(&mut slow).await, RangeRead::Finished)); + let buffer = shared.buffer.lock().await; + assert!(buffer.batches.is_empty()); + assert_eq!(buffer.retained_bytes, 0); + } +} diff --git a/src/mito2/src/engine/scan_test.rs b/src/mito2/src/engine/scan_test.rs index 0651fd2dd4c..e495184b112 100644 --- a/src/mito2/src/engine/scan_test.rs +++ b/src/mito2/src/engine/scan_test.rs @@ -4916,18 +4916,21 @@ async fn test_two_phase_sparse_metric_exact_sequence_cases() { } } -/// A warmed two-phase exact scan must hit the existing range-result cache; a +/// A warmed exact series scan must hit the existing range-result cache; a /// second upper bound over the same files and an unbounded scan must not reuse /// its rows. +#[rstest::rstest] +#[case(true)] +#[case(false)] #[tokio::test] -async fn test_two_phase_sparse_metric_exact_sequence_cache_isolation() { +async fn test_sparse_metric_exact_sequence_cache_isolation(#[case] two_phase: bool) { let (_env, engine, region_id) = build_sparse_exact_metric_engine( - "test_two_phase_sparse_metric_exact_sequence_cache_isolation", - true, + "test_sparse_metric_exact_sequence_cache_isolation", + two_phase, ReadableSize::mb(64), ) .await; - // Make the fixture file-only so the two-phase candidate and data ranges are cacheable. + // Make the fixture file-only so both series scan modes can cache their ranges. test_util::flush_region(&engine, region_id, None).await; let restrictive = ScanRequest { @@ -4943,7 +4946,10 @@ async fn test_two_phase_sparse_metric_exact_sequence_cache_isolation() { let Scanner::Series(scanner) = scanner else { panic!("sparse metric scan must use SeriesScan"); }; - assert_eq!("two_phase", scanner.mode()); + assert_eq!( + if two_phase { "two_phase" } else { "legacy" }, + scanner.mode() + ); assert_eq!(1, scanner.properties().num_partitions()); let batches = RecordBatches::try_collect(scanner.build_stream().await.unwrap()) .await @@ -4956,18 +4962,20 @@ async fn test_two_phase_sparse_metric_exact_sequence_cache_isolation() { let (cold, _) = scan(restrictive.clone()).await; // Cache insertion is performed by the range stream's async concat task. - // Wait for the SeriesReader's data-stage entry, rather than treating a - // candidate-stage cache hit as evidence that the reader replayed its data. + // The data reader uses partition 0 in two-phase mode and partition 1 in legacy mode. let warm = tokio::time::timeout(Duration::from_secs(5), async { loop { let (rows, verbose) = scan(restrictive.clone()).await; - // The single output partition is 0; the candidate distributor uses - // synthetic partition 1. Only a partition-0 hit proves data replay. + // Inspect the data reader rather than candidate-stage cache hits. let data_metrics = verbose - .split("\"partition\":0, \"metrics\":") + .split(if two_phase { + "\"partition\":0, \"metrics\":" + } else { + "\"partition\":1, \"metrics\":" + }) .nth(1) .and_then(|metrics| metrics.split("\"partition\":").next()) - .expect("missing two-phase data-reader metrics"); + .expect("missing output-partition metrics"); if data_metrics.contains("\"range_cache_hit\":") { return rows; } diff --git a/src/mito2/src/error.rs b/src/mito2/src/error.rs index d9340883d2e..155978aa5b2 100644 --- a/src/mito2/src/error.rs +++ b/src/mito2/src/error.rs @@ -198,6 +198,13 @@ pub enum Error { location: Location, }, + #[snafu(display("Failed to load range result"))] + LoadRangeResult { + source: Arc, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Failed to load SST metadata"))] LoadSstMeta { source: Arc, @@ -1469,9 +1476,9 @@ impl Error { match self { Error::OpenDal { error, .. } => error.kind() == ErrorKind::NotFound, Error::ManifestDeltaNotFound { .. } => true, - Error::LoadSstMeta { source, .. } | Error::LoadParquetPage { source, .. } => { - source.is_object_not_found() - } + Error::LoadSstMeta { source, .. } + | Error::LoadParquetPage { source, .. } + | Error::LoadRangeResult { source, .. } => source.is_object_not_found(), _ => false, } } @@ -1596,7 +1603,8 @@ impl ErrorExt for Error { WriteParquet { .. } => StatusCode::StorageUnavailable, WriteGroup { source, .. } | LoadSstMeta { source, .. } - | LoadParquetPage { source, .. } => source.status_code(), + | LoadParquetPage { source, .. } + | LoadRangeResult { source, .. } => source.status_code(), InvalidBatch { .. } => StatusCode::InvalidArguments, InvalidRecordBatch { .. } => StatusCode::InvalidArguments, ConvertVector { source, .. } => source.status_code(), @@ -1753,6 +1761,7 @@ impl ErrorExt for Error { OpenRegion { source, .. } | LoadSstMeta { source, .. } | LoadParquetPage { source, .. } + | LoadRangeResult { source, .. } | WriteGroup { source, .. } | FlushRegion { source, .. } | BuildIndexAsync { source, .. } @@ -1900,10 +1909,14 @@ mod tests { let expected_status = source.status_code(); let expected_retry = source.retry_hint(); let source = Arc::new(source); - let wrapped = LoadSstMetaSnafu.into_error(source); - assert_eq!(expected_not_found, wrapped.is_object_not_found()); - assert_eq!(expected_status, wrapped.status_code()); - assert_eq!(expected_retry, wrapped.retry_hint()); + for wrapped in [ + LoadSstMetaSnafu.into_error(source.clone()), + LoadRangeResultSnafu.into_error(source), + ] { + assert_eq!(expected_not_found, wrapped.is_object_not_found()); + assert_eq!(expected_status, wrapped.status_code()); + assert_eq!(expected_retry, wrapped.retry_hint()); + } } } diff --git a/src/mito2/src/metrics.rs b/src/mito2/src/metrics.rs index b88b7e5231c..ea27baa3955 100644 --- a/src/mito2/src/metrics.rs +++ b/src/mito2/src/metrics.rs @@ -292,6 +292,18 @@ pub(crate) fn cache_loading_metrics( } // Cache metrics. +lazy_static! { + pub static ref RANGE_CACHE_STAGING_BYTES: IntGauge = register_int_gauge!( + "greptime_mito_range_cache_staging_bytes", + "Referenced batch bytes retained by active shared range buffers, including unread streaming batches" + ) + .unwrap(); + pub static ref RANGE_CACHE_STREAMING_BYTES: IntGauge = register_int_gauge!( + "greptime_mito_range_cache_streaming_bytes", + "Referenced bytes retained for active readers after full-result cache admission stops" + ).unwrap(); +} + lazy_static! { /// Native cache request and actual loader lifecycle; request is not a waiter count. pub static ref CACHE_NATIVE_LOAD_EVENTS: IntCounterVec = register_int_counter_vec!( diff --git a/src/mito2/src/read/pruner.rs b/src/mito2/src/read/pruner.rs index 47689db161c..e4d5bc129e2 100644 --- a/src/mito2/src/read/pruner.rs +++ b/src/mito2/src/read/pruner.rs @@ -152,6 +152,15 @@ impl PartitionPruner { pruned } + /// Returns the pre-filter mode for a file index in the scan input. + /// Defaults to [`PreFilterMode::SkipFields`] when no mode was recorded. + pub(crate) fn pre_filter_mode(&self, file_index: usize) -> PreFilterMode { + self.pre_filter_modes + .get(file_index) + .copied() + .unwrap_or(PreFilterMode::SkipFields) + } + /// Pre-fetches upcoming files starting from the given position. fn prefetch_upcoming_files(&self, current_pos: usize, partition_metrics: &PartitionMetrics) { let start = current_pos + 1; @@ -168,13 +177,6 @@ impl PartitionPruner { } } - fn pre_filter_mode(&self, file_index: usize) -> PreFilterMode { - self.pre_filter_modes - .get(file_index) - .copied() - .unwrap_or(PreFilterMode::SkipFields) - } - fn file_index(&self, index: RowGroupIndex) -> Option { self.pruner .inner diff --git a/src/mito2/src/read/range_cache.rs b/src/mito2/src/read/range_cache.rs index a3b857fc24d..ded0230f4b9 100644 --- a/src/mito2/src/read/range_cache.rs +++ b/src/mito2/src/read/range_cache.rs @@ -18,36 +18,39 @@ use std::future::Future; use std::mem; use std::sync::Arc; +#[cfg(any(test, feature = "test"))] use async_stream::try_stream; -use common_telemetry::warn; use common_time::Timestamp; use common_time::range::TimestampRange; use common_time::timestamp::TimeUnit; +use datafusion::execution::memory_pool::MemoryPool; use datafusion_expr::expr::Expr; use datafusion_expr::{Between, BinaryExpr, Operator}; -use datatypes::arrow::compute::concat_batches; +#[cfg(test)] use datatypes::arrow::record_batch::RecordBatch; use datatypes::prelude::ConcreteDataType; use datatypes::value::scalar_value_to_timestamp; +#[cfg(any(test, feature = "test"))] use futures::TryStreamExt; -use snafu::ResultExt; use store_api::region_engine::PartitionRange; use store_api::storage::{FileId, RegionId, SequenceRange, TimeSeriesRowSelector}; use table::predicate::is_string_timestamp_literal; -use tokio::sync::{mpsc, oneshot}; use crate::cache::CacheStrategy; -use crate::cache::range_cache::{RangeCacheDelivery, RangeCacheLookup, RangeCacheWorker}; -use crate::error::{ComputeArrowSnafu, Result}; +use crate::cache::range_cache::RangeCacheScope; +#[cfg(test)] +use crate::cache::range_cache::{RANGE_CACHE_COMPACT_THRESHOLD_BYTES, RangeScanCacheValue}; +use crate::error::Result; use crate::read::BoxedRecordBatchStream; +use crate::read::pruner::PartitionPruner; use crate::read::read_columns::ReadColumns; use crate::read::scan_region::StreamContext; use crate::read::scan_util::PartitionMetrics; use crate::read::series_reader::SeriesRange; use crate::region::options::MergeMode; +#[cfg(any(test, feature = "test"))] use crate::sst::parquet::DEFAULT_READ_BATCH_SIZE; - -const RANGE_CACHE_COMPACT_THRESHOLD_BYTES: usize = 8 * 1024 * 1024; +use crate::sst::parquet::file_range::PreFilterMode; /// Fingerprint of the scan request fields that affect partition range cache reuse. /// @@ -70,6 +73,8 @@ pub(crate) struct ScanRequestFingerprint { /// range results are never reused for scans with a different (or no) range. sequence_range: Option, stage: RangeScanStage, + /// File interpretation and snapshot guards, aligned with the key's row groups. + file_reads: Option>>, /// We keep the partition expr version to ensure we won't reuse the fingerprint after we change the partition expr. /// We store the version instead of the whole partition expr or partition expr filters. partition_expr_version: u64, @@ -127,11 +132,23 @@ impl ScanRequestFingerprintBuilder { merge_mode, sequence_range, stage: RangeScanStage::Data, + file_reads: None, partition_expr_version, } } } +/// Inputs affecting a file read's output. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct FileReadFingerprint { + target_schema_version: u64, + source_region: RegionId, + file_sequence: Option, + preserve_row_sequence: bool, + snapshot_sequence: Option, + pre_filter_mode: PreFilterMode, +} + /// Non-copiable struct of the fingerprint. #[derive(Debug, PartialEq, Eq, Hash)] struct SharedScanRequestFingerprint { @@ -177,6 +194,7 @@ impl ScanRequestFingerprint { filter_deleted: self.filter_deleted, merge_mode: self.merge_mode, sequence_range: self.sequence_range, + file_reads: self.file_reads.clone(), stage: self.stage, partition_expr_version: self.partition_expr_version, } @@ -191,6 +209,7 @@ impl ScanRequestFingerprint { filter_deleted: self.filter_deleted, merge_mode: self.merge_mode, sequence_range: self.sequence_range, + file_reads: self.file_reads.clone(), stage: RangeScanStage::CandidateSeries, partition_expr_version: self.partition_expr_version, } @@ -205,6 +224,7 @@ impl ScanRequestFingerprint { filter_deleted: self.filter_deleted, merge_mode: self.merge_mode, sequence_range: self.sequence_range, + file_reads: self.file_reads.clone(), stage: RangeScanStage::SeriesData(range), partition_expr_version: self.partition_expr_version, } @@ -212,6 +232,10 @@ impl ScanRequestFingerprint { pub(crate) fn estimated_size(&self) -> usize { mem::size_of::() + + self.file_reads.as_ref().map_or(0, |files| { + mem::size_of::>() + + files.capacity() * mem::size_of::() + }) + self.inner.read_columns.estimated_size() + self.inner.read_column_types.capacity() * mem::size_of::>() + self.inner.filters.capacity() * mem::size_of::() @@ -236,7 +260,7 @@ impl ScanRequestFingerprint { #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub(crate) struct RangeScanCacheKey { pub(crate) region_id: RegionId, - /// Sorted (file_id, row_group_index) pairs that uniquely identify the data this range covers. + /// Sorted file/row-group identities for merged results. pub(crate) row_groups: Vec<(FileId, i64)>, pub(crate) scan: ScanRequestFingerprint, } @@ -255,81 +279,6 @@ impl RangeScanCacheKey { } } -/// Cached result for one range scan. -#[derive(Debug)] -pub(crate) struct CachedBatchSlice { - batch: RecordBatch, - slice_lengths: Vec, -} - -impl CachedBatchSlice { - fn metadata_size(&self) -> usize { - self.slice_lengths.capacity() * mem::size_of::() - } -} - -pub(crate) struct RangeScanCacheValue { - cached_batches: Vec, - /// Precomputed size of all compacted batches. - estimated_batches_size: usize, -} - -impl RangeScanCacheValue { - pub(crate) fn new( - cached_batches: Vec, - estimated_batches_size: usize, - ) -> Self { - Self { - cached_batches, - estimated_batches_size, - } - } - - pub(crate) fn estimated_size(&self) -> usize { - mem::size_of::() - + self.cached_batches.capacity() * mem::size_of::() - + self - .cached_batches - .iter() - .map(CachedBatchSlice::metadata_size) - .sum::() - + self.estimated_batches_size - } -} - -/// Row groups and whether all sources are file-only for a partition range. -pub(crate) struct PartitionRangeRowGroups { - /// Sorted (file_id, row_group_index) pairs. - pub(crate) row_groups: Vec<(FileId, i64)>, - pub(crate) only_file_sources: bool, -} - -/// Collects (file_id, row_group_index) pairs from a partition range's row group indices. -pub(crate) fn collect_partition_range_row_groups( - stream_ctx: &StreamContext, - part_range: &PartitionRange, -) -> PartitionRangeRowGroups { - let range_meta = &stream_ctx.ranges[part_range.identifier]; - let mut row_groups = Vec::new(); - let mut only_file_sources = true; - - for index in &range_meta.row_group_indices { - if stream_ctx.is_file_range_index(*index) { - let file_id = stream_ctx.input.file_from_index(*index).file_id().file_id(); - row_groups.push((file_id, index.row_group_index)); - } else { - only_file_sources = false; - } - } - - row_groups.sort_unstable_by(|a, b| a.0.as_bytes().cmp(b.0.as_bytes()).then(a.1.cmp(&b.1))); - - PartitionRangeRowGroups { - row_groups, - only_file_sources, - } -} - /// Returns the timestamp range where all time-only predicates are guaranteed true. /// /// Returns `None` for empty input because there is no time-filter implication @@ -510,19 +459,22 @@ fn implied_from_between( pub(crate) fn build_range_cache_key( stream_ctx: &StreamContext, part_range: &PartitionRange, + partition_pruner: Option<&PartitionPruner>, ) -> Option { - build_range_cache_key_inner(stream_ctx, part_range, None) + build_range_cache_key_inner(stream_ctx, part_range, None, partition_pruner) } /// Builds a cache key for a candidate-series partition-range result. pub(crate) fn build_candidate_range_cache_key( stream_ctx: &StreamContext, part_range: &PartitionRange, + partition_pruner: &PartitionPruner, ) -> Option { build_range_cache_key_inner( stream_ctx, part_range, Some(RangeScanStage::CandidateSeries), + Some(partition_pruner), ) } @@ -531,11 +483,13 @@ pub(crate) fn build_series_range_cache_key( stream_ctx: &StreamContext, part_range: &PartitionRange, range: SeriesRange, + partition_pruner: &PartitionPruner, ) -> Option { build_range_cache_key_inner( stream_ctx, part_range, Some(RangeScanStage::SeriesData(range)), + Some(partition_pruner), ) } @@ -543,6 +497,7 @@ fn build_range_cache_key_inner( stream_ctx: &StreamContext, part_range: &PartitionRange, stage: Option, + partition_pruner: Option<&PartitionPruner>, ) -> Option { if !stream_ctx.input.cache_strategy.has_range_result_cache() { return None; @@ -560,10 +515,39 @@ fn build_range_cache_key_inner( return None; } - let rg = collect_partition_range_row_groups(stream_ctx, part_range); - if !rg.only_file_sources || rg.row_groups.is_empty() { + let range_meta = &stream_ctx.ranges[part_range.identifier]; + let mut files = Vec::with_capacity(range_meta.row_group_indices.len()); + for index in &range_meta.row_group_indices { + if !stream_ctx.is_file_range_index(*index) { + return None; + } + let handle = stream_ctx.input.file_from_index(*index); + let file = handle.meta_ref(); + files.push(( + (handle.file_id().file_id(), index.row_group_index), + FileReadFingerprint { + target_schema_version: stream_ctx.input.region_metadata().schema_version, + source_region: handle.region_id(), + file_sequence: file.sequence.map(|sequence| sequence.get()), + preserve_row_sequence: file.preserve_row_sequence, + snapshot_sequence: stream_ctx.input.snapshot_sequence, + pre_filter_mode: partition_pruner.map_or_else( + || stream_ctx.range_pre_filter_mode(part_range), + |pruner| pruner.pre_filter_mode(index.index - stream_ctx.input.num_memtables()), + ), + }, + )); + } + if files.is_empty() { return None; } + files.sort_unstable_by(|a, b| { + a.0.0 + .as_bytes() + .cmp(b.0.0.as_bytes()) + .then(a.0.1.cmp(&b.0.1)) + }); + let (row_groups, file_reads): (Vec<_>, Vec<_>) = files.into_iter().unzip(); // If the implied range covers this partition's `FileTimeRange`, drop // time-only predicates from the cache key so that queries with different @@ -571,7 +555,6 @@ fn build_range_cache_key_inner( // analyzable time-only predicate or some predicate had an unsupported shape // (e.g. `OR`), so we keep the fingerprint unchanged. When `time_filters` is // already empty, cloning the fingerprint is equivalent to stripping them. - let range_meta = &stream_ctx.ranges[part_range.identifier]; let (file_min, file_max) = range_meta.time_range; let covers = match stream_ctx.input.implied_time_range() { // An empty implied range can never cover a non-empty file range, so @@ -591,7 +574,7 @@ fn build_range_cache_key_inner( } else { fingerprint.clone() }; - let scan = match stage { + let mut scan = match stage { Some(RangeScanStage::CandidateSeries) => scan.for_candidate_series(), Some(RangeScanStage::SeriesData(assignment)) => scan.for_series_data(assignment), Some(RangeScanStage::DataForCrossRangeMerge) => { @@ -602,34 +585,15 @@ fn build_range_cache_key_inner( Some(RangeScanStage::Data) | None => scan, }; + scan.file_reads = Some(Arc::new(file_reads)); Some(RangeScanCacheKey { region_id: stream_ctx.input.region_metadata().region_id, - row_groups: rg.row_groups, + row_groups, scan, }) } -/// Returns a stream that replays cached record batches. -pub(crate) fn cached_flat_range_stream( - delivery: Arc, -) -> BoxedRecordBatchStream { - Box::pin(try_stream! { - // Borrow through the delivery guard so it survives every yield. Batches - // handed to downstream operators are then owned by the query, not by - // this completed-result backlog. - for cached_batch in &delivery.value().cached_batches { - let mut offset = 0; - for &len in &cached_batch.slice_lengths { - yield cached_batch.batch.slice(offset, len); - offset += len; - } - } - }) -} - -/// Builds only one decoded result per full cache key. A load must be driven -/// independently of any caller: merge readers can wait for another stream with -/// the same key before polling their first stream to completion. +/// Cache access for one range, including the SST handles that guard publication. pub(crate) struct CachedRangeReader<'a> { pub(crate) stream_ctx: &'a StreamContext, pub(crate) part_range: &'a PartitionRange, @@ -640,359 +604,125 @@ pub(crate) struct CachedRangeReader<'a> { impl CachedRangeReader<'_> { pub(crate) async fn build(self, build: F) -> Result<(BoxedRecordBatchStream, usize)> where - F: FnOnce() -> Fut + Clone + Send + 'static, - Fut: Future> + Send + 'static, + F: FnOnce() -> Fut, + Fut: Future>, { - let Self { - stream_ctx, - part_range, - key, - metrics, - } = self; - let strategy = &stream_ctx.input.cache_strategy; - let (Some(key), Some(worker)) = (key, strategy.range_cache_worker()) else { + self.build_in_scope(RangeCacheScope::default(), build).await + } + + /// Shares active scans only with readers using the same query memory pool. + pub(crate) async fn build_with_memory_pool( + self, + memory_pool: &Arc, + build: F, + ) -> Result<(BoxedRecordBatchStream, usize)> + where + F: FnOnce() -> Fut, + Fut: Future>, + { + self.build_in_scope(RangeCacheScope::for_memory_pool(memory_pool), build) + .await + } + + async fn build_in_scope( + self, + scope: RangeCacheScope, + build: F, + ) -> Result<(BoxedRecordBatchStream, usize)> + where + F: FnOnce() -> Fut, + Fut: Future>, + { + let Some(key) = self.key else { return build().await; }; - let files = stream_ctx.ranges[part_range.identifier] + let files = self.stream_ctx.ranges[self.part_range.identifier] .row_group_indices .iter() - .map(|index| stream_ctx.input.file_from_index(*index).clone()) + .map(|index| self.stream_ctx.input.file_from_index(*index).clone()) .collect(); - let loader = build.clone(); - let load_worker = worker.clone(); - let load_strategy = strategy.clone(); - let load_metrics = metrics.clone(); - let key_size = key.estimated_size(); - match worker.get_or_load(key, files, async move { - let result = async { - let (stream, _) = loader().await?; - collect_range_result(stream, &load_worker, &load_strategy).await - } - .await; - match result { - Ok(Some(value)) => { - load_metrics.inc_range_cache_size(key_size + value.value().estimated_size()); - Some(value) - } - Ok(None) => None, - Err(error) => { - warn!(error; "Failed to load range cache; falling back to an uncached scan"); - None - } - } - }) { - RangeCacheLookup::Hit(value) => { - metrics.inc_range_cache_hit(); - Ok((cached_flat_range_stream(value), DEFAULT_READ_BATCH_SIZE)) - } - RangeCacheLookup::Loading(receiver) => { - metrics.inc_range_cache_miss(); - Ok(( - shared_range_stream(receiver, build), - DEFAULT_READ_BATCH_SIZE, - )) - } - RangeCacheLookup::Bypass => { - metrics.inc_range_cache_miss(); - build().await - } - } + load_shared_range( + &self.stream_ctx.input.cache_strategy, + key, + self.metrics, + files, + scope, + build, + ) + .await } } -/// Waiting is lazy, but the producer is independent of every consumer. -fn shared_range_stream( - mut receiver: tokio::sync::watch::Receiver>>, - build: F, -) -> BoxedRecordBatchStream -where - F: FnOnce() -> Fut + Send + 'static, - Fut: Future> + Send + 'static, -{ - Box::pin(try_stream! { - let cached = receiver.wait_for(|cached| cached.is_some()).await - .ok().and_then(|cached| cached.clone()); - drop(receiver); - let mut stream = if let Some(cached) = cached { - cached_flat_range_stream(cached) - } else { - // A failed/oversized/cancelled fill is never a query result. - // Retry once through the ordinary streaming path, not the worker. - build().await?.0 - }; - while let Some(batch) = stream.try_next().await? { - yield batch; - } - }) -} - -/// The load budget covers retained input buffers and completed deliveries. Acquisition -/// must be nonblocking: waiting with a partially accumulated result deadlocks -/// when several loads consume all permits. Oversized ranges retain streaming -/// query semantics by falling back to an uncached scan. -async fn collect_range_result( - mut stream: BoxedRecordBatchStream, - worker: &RangeCacheWorker, +/// Joins an active scan while retaining independently advancing stream cursors. +#[cfg(any(test, feature = "test"))] +pub(crate) async fn get_or_load_range_result( strategy: &CacheStrategy, -) -> Result>> { - let mut permit: Option = None; - let mut buffer = CacheBatchBuffer::new(strategy); - while let Some(batch) = stream.try_next().await? { - let Some(next_permit) = worker - .load_memory - .try_acquire(batch.get_array_memory_size()) - else { - return Ok(None); - }; - if let Some(permit) = &mut permit { - permit.merge(next_permit); - } else { - permit = Some(next_permit); - } - buffer.push(batch)?; - } - let Some(value) = buffer.finish().await.transpose()? else { - return Ok(None); - }; - let Some(permit) = worker - .load_memory - .try_resize(permit, value.estimated_size()) - else { - return Ok(None); - }; - Ok(Some(RangeCacheDelivery::new(value, permit))) -} - -enum CacheConcatCommand { - Compact(Vec), - Finish { - pending: Vec, - result_tx: oneshot::Sender>>, - }, -} - -#[derive(Default)] -struct CacheConcatState { - cached_batches: Vec, - estimated_size: usize, -} - -impl CacheConcatState { - async fn compact( - &mut self, - batches: Vec, - limiter: &crate::cache::RangeResultMemoryLimiter, - ) -> Result<()> { - if batches.is_empty() { - return Ok(()); - } - - let input_size = batches - .iter() - .map(RecordBatch::get_array_memory_size) - .sum::(); - let _permit = limiter.acquire(input_size).await?; - - let compacted = compact_record_batches(batches)?; - self.estimated_size += compacted.batch.get_array_memory_size(); - self.cached_batches.push(compacted); - Ok(()) - } - - fn finish(self) -> RangeScanCacheValue { - RangeScanCacheValue::new(self.cached_batches, self.estimated_size) - } -} - -fn compact_record_batches(batches: Vec) -> Result { - debug_assert!(!batches.is_empty()); - - let slice_lengths = batches.iter().map(RecordBatch::num_rows).collect(); - build_cached_batch_slice(batches, slice_lengths) -} - -fn build_cached_batch_slice( - batches: Vec, - slice_lengths: Vec, -) -> Result { - let batch = if batches.len() == 1 { - batches.into_iter().next().unwrap() - } else { - let schema = batches[0].schema(); - concat_batches(&schema, &batches).context(ComputeArrowSnafu)? - }; - - Ok(CachedBatchSlice { - batch, - slice_lengths, - }) -} - -async fn run_cache_concat_task( - mut rx: mpsc::UnboundedReceiver, - limiter: Arc, - skip_threshold_bytes: usize, -) { - let mut state = CacheConcatState::default(); - - while let Some(cmd) = rx.recv().await { - match cmd { - CacheConcatCommand::Compact(batches) => { - if let Err(err) = state.compact(batches, &limiter).await { - warn!(err; "Failed to compact range cache batches"); - return; - } - // Close the channel to stop further work as soon as the cached - // size exceeds the configured cache budget. - if state.estimated_size > skip_threshold_bytes { - return; - } - } - CacheConcatCommand::Finish { pending, result_tx } => { - let compact_result = state - .compact(pending, &limiter) - .await - .map(|()| state.finish()); - let result = match compact_result { - Ok(v) => { - let value = Arc::new(v); - Ok(value) - } - Err(e) => { - warn!(e; "Failed to finalize range cache batches"); - - Err(e) - } - }; - - let _ = result_tx.send(result); - - break; - } - } - } -} - -struct CacheBatchBuffer { - buffered_batches: Vec, - buffered_rows: usize, - buffered_size: usize, - sender: Option>, - /// Kept until concat finishes so abandoning the shared load cancels concat too. - concat_task: Option>, -} - -impl CacheBatchBuffer { - fn new(cache_strategy: &CacheStrategy) -> Self { - let (sender, concat_task) = cache_strategy - .range_result_memory_limiter() - .map(|limiter| { - let skip_threshold_bytes = cache_strategy.range_result_cache_size().unwrap_or(0); - let (tx, rx) = mpsc::unbounded_channel(); - let task = common_runtime::spawn_query(run_cache_concat_task( - rx, - limiter.clone(), - skip_threshold_bytes, - )); - (tx, task) - }) - .unzip(); - - Self { - buffered_batches: Vec::new(), - buffered_rows: 0, - buffered_size: 0, - sender, - concat_task, - } - } - - fn push(&mut self, batch: RecordBatch) -> Result<()> { - if self.sender.is_none() { - return Ok(()); - } - - self.buffered_rows += batch.num_rows(); - self.buffered_size += batch.get_array_memory_size(); - self.buffered_batches.push(batch); - - if self.buffered_batches.len() > 1 - && (self.buffered_rows > DEFAULT_READ_BATCH_SIZE - || self.buffered_size > RANGE_CACHE_COMPACT_THRESHOLD_BYTES) - { - self.notify_compact(); - } - - Ok(()) - } - - fn notify_compact(&mut self) { - if self.buffered_batches.is_empty() || self.sender.is_none() { - return; - } - - let batches = mem::take(&mut self.buffered_batches); - self.buffered_rows = 0; - self.buffered_size = 0; - - let Some(sender) = &self.sender else { - return; - }; - if sender.send(CacheConcatCommand::Compact(batches)).is_err() { - self.sender = None; - } - } - - async fn finish(mut self) -> Option>> { - let sender = self.sender.take()?; - let (result_tx, result_rx) = oneshot::channel(); - sender - .send(CacheConcatCommand::Finish { - pending: mem::take(&mut self.buffered_batches), - result_tx, - }) - .ok()?; - // Keep ownership until the result arrives so cancellation also aborts - // concat while it waits for its temporary memory permits. - let result = result_rx.await.ok(); - self.concat_task = None; - result - } -} - -impl Drop for CacheBatchBuffer { - fn drop(&mut self) { - // Aborting releases queued batches even while concat waits for a permit. - if let Some(task) = &self.concat_task { - task.abort(); - } - } -} - -/// Wraps a stream to cache its output for future range cache hits. -#[cfg(feature = "test")] -pub(crate) fn cache_flat_range_stream( - mut stream: BoxedRecordBatchStream, - cache_strategy: CacheStrategy, key: RangeScanCacheKey, - part_metrics: PartitionMetrics, -) -> BoxedRecordBatchStream { - Box::pin(try_stream! { - let mut buffer = CacheBatchBuffer::new(&cache_strategy); - while let Some(batch) = stream.try_next().await? { - buffer.push(batch.clone())?; - yield batch; - } - - if let Some(Ok(value)) = buffer.finish().await { - part_metrics.inc_range_cache_size(key.estimated_size() + value.estimated_size()); - cache_strategy.put_range_result(key, value); - } - }) + part_metrics: &PartitionMetrics, + build: F, +) -> Result<(BoxedRecordBatchStream, usize)> +where + F: FnOnce() -> Fut, + Fut: Future>, +{ + get_or_load_range_result_in_scope( + strategy, + key, + part_metrics, + RangeCacheScope::default(), + build, + ) + .await } -/// Creates a `cache_flat_range_stream` with dummy internals for benchmarking. +/// Isolates active sources that reserve memory from a query-specific pool. +/// Pool identity affects active sharing only, never completed-result reuse. +#[cfg(any(test, feature = "test"))] +pub(crate) async fn get_or_load_range_result_in_scope( + strategy: &CacheStrategy, + key: RangeScanCacheKey, + part_metrics: &PartitionMetrics, + scope: RangeCacheScope, + build: F, +) -> Result<(BoxedRecordBatchStream, usize)> +where + F: FnOnce() -> Fut, + Fut: Future>, +{ + load_shared_range(strategy, key, part_metrics, Vec::new(), scope, build).await +} + +async fn load_shared_range( + strategy: &CacheStrategy, + key: RangeScanCacheKey, + part_metrics: &PartitionMetrics, + files: Vec, + scope: RangeCacheScope, + build: F, +) -> Result<(BoxedRecordBatchStream, usize)> +where + F: FnOnce() -> Fut, + Fut: Future>, +{ + let Some(cache) = strategy.range_result_cache() else { + return build().await; + }; + cache + .read( + key, + part_metrics.range_cache_metrics_reporter(), + files, + scope, + build, + ) + .await +} + +/// Benchmarks the production range-result loader with a fresh cache per call. /// +/// A zero capacity disables caching. A nonzero capacity measures cold construction, +/// including staging and consumption; background compaction may outlive the stream. +/// It does not measure warm reuse or concurrent load sharing. /// This avoids exposing `RangeScanCacheKey`, `ScanRequestFingerprint`, and /// `PartitionMetrics` publicly. #[cfg(feature = "test")] @@ -1038,7 +768,17 @@ pub fn bench_cache_flat_range_stream( let part_metrics = PartitionMetrics::new(region_id, 0, "bench", Instant::now(), false, &metrics_set); - cache_flat_range_stream(stream, cache_strategy, key, part_metrics) + Box::pin(try_stream! { + let (mut stream, _) = get_or_load_range_result( + &cache_strategy, + key, + &part_metrics, + || async { Ok((stream, DEFAULT_READ_BATCH_SIZE)) }, + ).await?; + while let Some(batch) = stream.try_next().await? { + yield batch; + } + }) } #[cfg(test)] @@ -1054,14 +794,13 @@ mod tests { use datafusion_expr::{Expr, col, lit}; use smallvec::smallvec; use store_api::storage::{FileId, RegionId}; - use tokio::sync::watch; use super::*; use crate::cache::CacheManager; - use crate::cache::test_util::range_cache_key; use crate::read::flat_projection::FlatProjectionMapper; + use crate::read::pruner::Pruner; use crate::read::range::{RangeMeta, RowGroupIndex, SourceIndex}; - use crate::read::scan_region::{PredicateGroup, ScanInput}; + use crate::read::scan_region::{PredicateGroup, ScanInput, StreamContext}; use crate::sst::file::FileTimeRange; use crate::test_util::memtable_util::metadata_with_primary_key; use crate::test_util::scheduler_util::SchedulerEnv; @@ -1161,7 +900,7 @@ mod tests { #[case(1024 * 1024, true)] #[case(1, false)] #[tokio::test] - async fn shared_range_read_preserves_batches_or_falls_back_on_budget_exhaustion( + async fn test_shared_range_read_preserves_batches_when_admission_budget_is_exhausted( #[case] capacity: u64, #[case] cached: bool, ) { @@ -1174,7 +913,7 @@ mod tests { .range_result_cache_size(capacity) .build(), )); - let key = build_range_cache_key(&ctx, &range).unwrap(); + let key = build_range_cache_key(&ctx, &range, None).unwrap(); let (_, metrics) = test_cache_context(&ctx.input.cache_strategy); let calls = Arc::new(AtomicUsize::new(0)); let expected = vec![make_batch(&[1, 2]), make_batch(&[3, 4, 5])]; @@ -1198,210 +937,61 @@ mod tests { .unwrap(); let actual: Vec<_> = stream.try_collect().await.unwrap(); assert_eq!(actual, expected); - assert_eq!(calls.load(Ordering::Relaxed), if cached { 1 } else { 2 }); + assert_eq!(calls.load(Ordering::Relaxed), 1); + if cached { + wait_cached(&ctx.input.cache_strategy, &key).await; + } assert_eq!( ctx.input.cache_strategy.get_range_result(&key).is_some(), cached ); - assert_eq!( - ctx.input - .cache_strategy - .range_cache_worker() - .unwrap() - .load_memory - .available_permits(), - (capacity as usize).div_ceil(1024) - ); - } - - fn start_shared_load( - strategy: &CacheStrategy, - key: RangeScanCacheKey, - batches: Vec, - ) -> watch::Receiver>> { - let worker = strategy.range_cache_worker().unwrap(); - let loader_worker = worker.clone(); - let loader_strategy = strategy.clone(); - let lookup = worker.get_or_load(key, vec![], async move { - collect_range_result( - Box::pin(futures::stream::iter(batches.into_iter().map(Ok))), - &loader_worker, - &loader_strategy, - ) - .await + let memory = &ctx + .input + .cache_strategy + .range_result_cache() .unwrap() - }); - let RangeCacheLookup::Loading(receiver) = lookup else { - panic!("expected a shared load"); - }; - receiver - } - - async fn wait_for_shared_load( - receiver: &mut watch::Receiver>>, - ) -> Option> { + .retention_memory(); tokio::time::timeout(std::time::Duration::from_secs(5), async { - receiver - .wait_for(|value| value.is_some()) - .await - .ok() - .and_then(|value| value.clone()) - }) - .await - .expect("a shared load must finish or fall back without waiting for consumers") - } - - #[tokio::test] - async fn completed_load_keeps_memory_until_all_subscribers_release() { - let strategy = test_cache_strategy(); - let worker = strategy.range_cache_worker().unwrap(); - let capacity = worker.load_memory.available_permits(); - let (key, _) = test_cache_context(&strategy); - let mut first = start_shared_load(&strategy, key.clone(), vec![make_batch(&[1, 2, 3])]); - let second = first.clone(); - drop(wait_for_shared_load(&mut first).await.unwrap()); - - // A completed flight no longer counts toward the flight limit, but - // unpolled subscribers must still reserve its decoded result memory. - assert!(worker.load_memory.available_permits() < capacity); - drop(first); - assert!(worker.load_memory.available_permits() < capacity); - drop(second); - tokio::time::timeout(std::time::Duration::from_secs(5), async { - while worker.load_memory.available_permits() != capacity { + while memory.available_permits() != (capacity as usize).div_ceil(1024) { tokio::task::yield_now().await; } }) .await - .expect("the last subscriber must release the delivery reservation"); - assert!( - worker.get(&key).is_some(), - "the cache still owns the result" - ); + .expect("finalization did not release the source reservation"); } - #[rstest::rstest] #[tokio::test] - async fn successive_loads_with_slow_consumers_exhaust_delivery_budget( - #[values(false, true)] invalidate: bool, - ) { - let strategy = CacheStrategy::EnableAll(Arc::new( - CacheManager::builder() - .range_result_cache_size(8192) - .build(), - )); - let worker = strategy.range_cache_worker().unwrap(); - let capacity = worker.load_memory.available_permits(); - let mut held_streams = Vec::new(); - let mut fallbacks = 0; - - // Each flight finishes before the next starts, exceeding the 32-flight - // limit cumulatively. Only the delivery budget can bound this backlog. - for _ in 0..40 { - let file = FileId::random(); - let key = range_cache_key(&[file]); - let expected = vec![make_large_binary_batch(1, 2048)]; - let mut receiver = start_shared_load(&strategy, key, expected.clone()); - let shared = wait_for_shared_load(&mut receiver).await.is_some(); - if invalidate { - worker.invalidate_file(file); - } - let fallback_batches = expected.clone(); - let stream = shared_range_stream(receiver, move || async move { + async fn test_cancelling_a_partial_load_releases_its_reservation() { + let strategy = test_cache_strategy(); + let worker = strategy.range_result_cache().unwrap().clone(); + let capacity = worker.retention_memory().available_permits(); + let (key, metrics) = test_cache_context(&strategy); + let (mut stream, _) = + get_or_load_range_result(&strategy, key.clone(), &metrics, || async { Ok(( - Box::pin(futures::stream::iter(fallback_batches.into_iter().map(Ok))) - as BoxedRecordBatchStream, + Box::pin(try_stream! { + yield make_batch(&[1, 2, 3]); + std::future::pending::<()>().await; + }) as BoxedRecordBatchStream, DEFAULT_READ_BATCH_SIZE, )) - }); - if shared { - held_streams.push(stream); - } else { - fallbacks += 1; - let actual: Vec<_> = - tokio::time::timeout(std::time::Duration::from_secs(5), stream.try_collect()) - .await - .expect("fallback must not wait for an unpolled range") - .unwrap(); - assert_eq!(actual, expected); - } - } - - assert!(!held_streams.is_empty()); - assert!(held_streams.len() < 32); - assert!( - fallbacks > 0, - "completed results must not escape the budget" - ); - assert!(worker.load_memory.available_permits() < capacity); - - // Transfer ownership out of watch receivers into actively polled replay - // streams, then leave them suspended at their first yield. - for stream in &mut held_streams { - assert!(stream.try_next().await.unwrap().is_some()); - } - assert!(worker.load_memory.available_permits() < capacity); - drop(held_streams); - tokio::time::timeout(std::time::Duration::from_secs(5), async { - while worker.load_memory.available_permits() != capacity { - tokio::task::yield_now().await; - } - }) - .await - .expect("dropping slow readers must unblock new shared loads"); - - let mut retry = start_shared_load( - &strategy, - range_cache_key(&[FileId::random()]), - vec![make_large_binary_batch(1, 2048)], - ); - assert!(wait_for_shared_load(&mut retry).await.is_some()); - } - - #[tokio::test] - async fn cancelling_a_partial_load_releases_its_reservation() { - let strategy = test_cache_strategy(); - let worker = strategy.range_cache_worker().unwrap().clone(); - let capacity = worker.load_memory.available_permits(); - let (key, _) = test_cache_context(&strategy); - let loader_worker = worker.clone(); - let lookup = worker.get_or_load(key.clone(), vec![], async move { - let stream = Box::pin(try_stream! { - yield make_batch(&[1, 2, 3]); - std::future::pending::<()>().await; - }); - collect_range_result(stream, &loader_worker, &strategy) - .await - .unwrap() - }); - let RangeCacheLookup::Loading(receiver) = lookup else { - panic!("expected a shared load"); - }; - tokio::time::timeout(std::time::Duration::from_secs(5), async { - while worker.load_memory.available_permits() == capacity { - tokio::task::yield_now().await; - } - }) - .await - .expect("the partial result must acquire its input reservation"); - drop(receiver); - tokio::time::timeout(std::time::Duration::from_secs(5), async { - while worker.load_memory.available_permits() != capacity { - tokio::task::yield_now().await; - } - }) - .await - .expect("cancellation must release the partial result reservation"); + }) + .await + .unwrap(); + assert!(stream.try_next().await.unwrap().is_some()); + assert!(worker.retention_memory().available_permits() < capacity); + drop(stream); + assert_eq!(worker.retention_memory().available_permits(), capacity); assert!(worker.get(&key).is_none()); } #[tokio::test] - async fn follower_can_finish_without_polling_the_first_range_stream() { + async fn test_follower_can_finish_without_polling_the_first_range_stream() { use std::sync::atomic::{AtomicUsize, Ordering}; let (ctx, range) = new_stream_context(vec![col("k0").eq(lit("a"))], None, FileTimeRange::default()).await; - let key = build_range_cache_key(&ctx, &range).unwrap(); + let key = build_range_cache_key(&ctx, &range, None).unwrap(); let (_, metrics) = test_cache_context(&ctx.input.cache_strategy); let gate = Arc::new(tokio::sync::Semaphore::new(0)); let calls = Arc::new(AtomicUsize::new(0)); @@ -1409,10 +999,11 @@ mod tests { let loader_gate = gate.clone(); let build = move || async move { count.fetch_add(1, Ordering::Relaxed); - let _permit = loader_gate.acquire().await.unwrap(); Ok(( - Box::pin(futures::stream::iter([Ok(make_batch(&[1, 2, 3]))])) - as BoxedRecordBatchStream, + Box::pin(try_stream! { + let _permit = loader_gate.acquire().await.unwrap(); + yield make_batch(&[1, 2, 3]); + }) as BoxedRecordBatchStream, DEFAULT_READ_BATCH_SIZE, )) }; @@ -1470,7 +1061,7 @@ mod tests { let (stream_ctx, part_range) = new_stream_context(filters, query_time_range, partition_time_range).await; - let key = build_range_cache_key(&stream_ctx, &part_range).unwrap(); + let key = build_range_cache_key(&stream_ctx, &part_range, None).unwrap(); assert_eq!( key.scan.filters(), @@ -1483,7 +1074,7 @@ mod tests { } #[tokio::test] - async fn range_cache_time_filter_key_cases() { + async fn test_range_cache_time_filter_key_cases() { let partition = ( Timestamp::new_millisecond(1000), Timestamp::new_millisecond(2000), @@ -1569,7 +1160,7 @@ mod tests { } #[tokio::test] - async fn two_distinct_queries_share_cache_key_when_both_cover() { + async fn test_two_distinct_queries_share_cache_key_when_both_cover() { let partition_range = ( Timestamp::new_millisecond(1000), Timestamp::new_millisecond(2000), @@ -1596,14 +1187,14 @@ mod tests { ) .await; - let key_a = build_range_cache_key(&ctx_a, &part_a).unwrap(); - let key_b = build_range_cache_key(&ctx_b, &part_b).unwrap(); + let key_a = build_range_cache_key(&ctx_a, &part_a, None).unwrap(); + let key_b = build_range_cache_key(&ctx_b, &part_b, None).unwrap(); assert_eq!(key_a.scan, key_b.scan); assert!(key_a.scan.time_filters().is_empty()); } #[tokio::test] - async fn scan_stages_have_separate_cache_keys() { + async fn test_scan_stages_have_separate_cache_keys() { let partition_range = ( Timestamp::new_millisecond(1000), Timestamp::new_millisecond(2000), @@ -1611,12 +1202,17 @@ mod tests { let (ctx, part_range) = new_stream_context(vec![col("k0").eq(lit("foo"))], None, partition_range).await; - let data_key = build_range_cache_key(&ctx, &part_range).unwrap(); - let candidate_key = build_candidate_range_cache_key(&ctx, &part_range).unwrap(); + let ctx = Arc::new(ctx); + let pruner = Arc::new(crate::read::pruner::Pruner::new(ctx.clone(), 0)); + let partition = PartitionPruner::new(pruner, std::slice::from_ref(&part_range)); + let data_key = build_range_cache_key(&ctx, &part_range, None).unwrap(); + let candidate_key = build_candidate_range_cache_key(&ctx, &part_range, &partition).unwrap(); let range_0 = SeriesRange::new(0, 2).unwrap(); let range_1 = SeriesRange::new(1, 2).unwrap(); - let series_key_0 = build_series_range_cache_key(&ctx, &part_range, range_0).unwrap(); - let series_key_1 = build_series_range_cache_key(&ctx, &part_range, range_1).unwrap(); + let series_key_0 = + build_series_range_cache_key(&ctx, &part_range, range_0, &partition).unwrap(); + let series_key_1 = + build_series_range_cache_key(&ctx, &part_range, range_1, &partition).unwrap(); let merge_key = data_key.clone().for_cross_range_merge(); assert_ne!(merge_key, data_key); @@ -1625,15 +1221,15 @@ mod tests { assert_eq!(merge_key.row_groups, data_key.row_groups); assert_eq!(merge_key, data_key.clone().for_cross_range_merge()); assert_ne!(data_key.scan, candidate_key.scan); - assert_ne!(data_key.scan, series_key_0.scan); assert_ne!(candidate_key.scan, series_key_0.scan); - assert_ne!(series_key_0.scan, series_key_1.scan); assert_eq!(data_key.row_groups, candidate_key.row_groups); + assert_ne!(data_key.scan, series_key_0.scan); + assert_ne!(series_key_0.scan, series_key_1.scan); assert_eq!(data_key.row_groups, series_key_0.row_groups); } #[tokio::test] - async fn disables_optimization_on_or_clause() { + async fn test_disables_optimization_on_or_clause() { let partition_range = ( Timestamp::new_millisecond(1000), Timestamp::new_millisecond(2000), @@ -1656,8 +1252,8 @@ mod tests { .await; assert!(ctx_a.input.implied_time_range().is_none()); - let key_a = build_range_cache_key(&ctx_a, &part_a).unwrap(); - let key_b = build_range_cache_key(&ctx_b, &part_b).unwrap(); + let key_a = build_range_cache_key(&ctx_a, &part_a, None).unwrap(); + let key_b = build_range_cache_key(&ctx_b, &part_b, None).unwrap(); assert_ne!(key_a.scan, key_b.scan); assert_eq!( key_a.scan.time_filters(), @@ -1666,7 +1262,7 @@ mod tests { } #[tokio::test] - async fn empty_implied_range_does_not_panic_on_non_ms_file_range() { + async fn test_empty_implied_range_does_not_panic_on_non_ms_file_range() { // Contradictory time predicates make the implied range empty. The // empty range's sentinel timestamps use `Timestamp::default()` (ms), // so without the `is_empty()` short-circuit the unit asserts would @@ -1693,13 +1289,13 @@ mod tests { Timestamp::new(2_000_000_000, TimeUnit::Nanosecond), ); - let key = build_range_cache_key(&ctx, &part_range).unwrap(); + let key = build_range_cache_key(&ctx, &part_range, None).unwrap(); // Empty implied range cannot cover, so time filters stay in the key. assert!(!key.scan.time_filters().is_empty()); } #[tokio::test] - async fn mixed_unit_file_range_is_compared_by_instant() { + async fn test_mixed_unit_file_range_is_compared_by_instant() { // After the time index unit was widened, pre-alter files keep the old // unit while `scan_implied_time_range` uses the current column unit. // The coverage decision must compare instants across units instead of @@ -1726,7 +1322,7 @@ mod tests { Timestamp::new_microsecond(1_000_000), Timestamp::new_microsecond(2_000_000), ); - let key = build_range_cache_key(&ctx, &part_range).unwrap(); + let key = build_range_cache_key(&ctx, &part_range, None).unwrap(); assert!(key.scan.time_filters().is_empty()); // Not covering: the implied range [1500, 3000)ms excludes the file's @@ -1745,7 +1341,7 @@ mod tests { Timestamp::new_microsecond(1_000_000), Timestamp::new_microsecond(2_000_000), ); - let key = build_range_cache_key(&ctx, &part_range).unwrap(); + let key = build_range_cache_key(&ctx, &part_range, None).unwrap(); assert!(!key.scan.time_filters().is_empty()); } @@ -1758,7 +1354,7 @@ mod tests { } #[test] - fn implied_time_range_supported_exprs() { + fn test_implied_time_range_supported_exprs() { for (expr, expected) in [ ( col("ts").gt_eq(ts_lit(1000)), @@ -1815,7 +1411,7 @@ mod tests { } #[test] - fn implied_time_range_unsupported_exprs() { + fn test_implied_time_range_unsupported_exprs() { let not_between = Expr::Between(Between { expr: Box::new(col("ts")), negated: true, @@ -1836,7 +1432,7 @@ mod tests { } #[test] - fn implied_time_range_unit_conversion() { + fn test_implied_time_range_unit_conversion() { let second_1 = lit(ScalarValue::TimestampSecond(Some(1), None)); let ns_1500 = lit(ScalarValue::TimestampNanosecond(Some(1_500_000_000), None)); let ns_1500_5 = lit(ScalarValue::TimestampNanosecond(Some(1_500_500_000), None)); @@ -1877,7 +1473,7 @@ mod tests { } #[test] - fn selector_after_merge_changes_fingerprint() { + fn test_selector_after_merge_changes_fingerprint() { let ordinary = test_scan_fingerprint( vec!["k0 = 'foo'".to_string()], vec![], @@ -1897,7 +1493,7 @@ mod tests { } #[test] - fn true_selector_after_merge_is_preserved_by_fingerprint_transforms() { + fn test_true_selector_after_merge_is_preserved_by_fingerprint_transforms() { let normalized = test_scan_fingerprint(vec!["k0 = 'foo'".to_string()], vec![], None, true, 0); @@ -1970,166 +1566,789 @@ mod tests { ArrowDataType::Binary, false, )])); - let payload = vec![b'x'; bytes_per_row]; - let values = (0..rows).map(|_| payload.as_slice()).collect::>(); + let payloads: Vec<_> = (0..rows) + .map(|i| vec![(i % 256) as u8; bytes_per_row]) + .collect(); + let values = payloads.iter().map(Vec::as_slice).collect::>(); RecordBatch::try_new(schema, vec![Arc::new(BinaryArray::from_vec(values))]).unwrap() } - #[test] - fn compact_record_batches_keeps_original_boundaries() { - let batches = vec![make_batch(&[1, 2]), make_batch(&[3]), make_batch(&[4, 5])]; + fn source_stream(batches: Vec) -> BoxedRecordBatchStream { + Box::pin(futures::stream::iter(batches.into_iter().map(Ok))) + } - let compacted = compact_record_batches(batches).unwrap(); - - assert_eq!(compacted.batch.num_rows(), 5); - assert_eq!(compacted.slice_lengths, vec![2, 1, 2]); + async fn wait_cached( + strategy: &CacheStrategy, + key: &RangeScanCacheKey, + ) -> Arc { + tokio::time::timeout(std::time::Duration::from_secs(5), async { + loop { + if let Some(value) = strategy.get_range_result(key) { + return value; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("range result was not published") } #[tokio::test] - async fn cached_flat_range_stream_replays_original_batches() { - let value = Arc::new(RangeScanCacheValue::new( - vec![CachedBatchSlice { - batch: make_batch(&[1, 2, 3]), - slice_lengths: vec![2, 1], - }], - make_batch(&[1, 2, 3]).get_array_memory_size(), - )); + async fn test_merged_key_covers_files_and_snapshot_independently_of_order() { + let (context, range) = new_stream_context( + vec![col("k0").eq(lit("a"))], + None, + ( + Timestamp::new_millisecond(0), + Timestamp::new_millisecond(10), + ), + ) + .await; + let mut context = Arc::new(context); + let first = build_range_cache_key(&context, &range, None).unwrap(); + let file = sst_file_handle_with_file_id(FileId::random(), 0, 10); + Arc::get_mut(&mut context).unwrap().input.files.push(file); + Arc::get_mut(&mut context).unwrap().ranges[0] + .indices + .push(SourceIndex { + index: 1, + num_row_groups: 1, + }); + Arc::get_mut(&mut context).unwrap().ranges[0] + .row_group_indices + .push(RowGroupIndex { + index: 1, + row_group_index: -1, + }); + let whole = build_range_cache_key(&context, &range, None).unwrap(); + assert_eq!(whole.row_groups.len(), 2); + assert!( + whole + .row_groups + .iter() + .any(|(_, row_group)| *row_group == -1) + ); + assert_ne!(first, whole); + Arc::get_mut(&mut context).unwrap().ranges[0] + .row_group_indices + .swap(0, 1); + let reordered = build_range_cache_key(&context, &range, None).unwrap(); + assert_eq!(whole, reordered); + Arc::get_mut(&mut context).unwrap().input.snapshot_sequence = Some(42); + let snapshot = build_range_cache_key(&context, &range, None).unwrap(); + assert_ne!(snapshot, reordered); + assert_eq!(snapshot.scan.stage, RangeScanStage::Data); + let sources = snapshot.scan.file_reads.as_ref().unwrap(); + assert_eq!(sources.len(), 2); + assert!( + sources + .iter() + .all(|source| source.snapshot_sequence == Some(42)) + ); + } - let worker = RangeCacheWorker::new(1024 * 1024); - let permit = worker - .load_memory - .try_acquire(value.estimated_size()) - .unwrap(); - let delivery = RangeCacheDelivery::new(value, permit); - let replayed = cached_flat_range_stream(delivery) - .try_collect::>() + fn new_partition_pruner( + context: &Arc, + range: &PartitionRange, + ) -> Arc { + let pruner = Arc::new(Pruner::new(context.clone(), 1)); + pruner.add_partition_ranges(std::slice::from_ref(range)); + Arc::new(PartitionPruner::new(pruner, std::slice::from_ref(range))) + } + + async fn sst_range_stream( + context: &Arc, + range: &PartitionRange, + metrics: &PartitionMetrics, + partition: Arc, + semaphore: Option>, + ) -> BoxedRecordBatchStream { + crate::read::seq_scan::SeqScan::build_flat_partition_range_read( + context, + range, + crate::read::seq_scan::FlatReadMode::Sequential, + metrics, + partition, + semaphore.clone(), + semaphore, + ) + .await + .unwrap() + .0 + } + + fn u64_values(batches: &[RecordBatch], name: &str) -> Vec> { + batches + .iter() + .flat_map(|batch| { + batch + .column_by_name(name) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .iter() + }) + .collect() + } + + async fn write_overlapping_sst_files( + env: &SchedulerEnv, + ) -> ( + store_api::metadata::RegionMetadataRef, + Vec, + ) { + use datatypes::arrow::array::UInt64Array; + use parquet::arrow::ArrowWriter; + use parquet::file::metadata::KeyValue; + use parquet::file::properties::WriterProperties; + use store_api::region_request::PathType; + + use crate::sst::file::{FileHandle, FileMeta}; + use crate::sst::parquet::PARQUET_METADATA_KEY; + use crate::test_util::sst_util::{ + new_record_batch_with_custom_sequence, sst_region_metadata, + }; + + let metadata = Arc::new(sst_region_metadata()); + let mut files = Vec::new(); + for sequence in [1, 2] { + let props = WriterProperties::builder() + .set_key_value_metadata(Some(vec![KeyValue::new( + PARQUET_METADATA_KEY.to_string(), + metadata.to_json().unwrap(), + )])) + .build(); + let first = new_record_batch_with_custom_sequence(&["a", "x"], 0, 2, sequence); + let mut bytes = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut bytes, first.schema(), Some(props)).unwrap(); + for start in [0, 2] { + let batch = + new_record_batch_with_custom_sequence(&["a", "x"], start, start + 2, sequence); + let mut columns = batch.columns().to_vec(); + columns[2] = Arc::new(UInt64Array::from(vec![ + sequence * 100 + start as u64, + sequence * 100 + start as u64 + 1, + ])); + writer + .write(&RecordBatch::try_new(batch.schema(), columns).unwrap()) + .unwrap(); + writer.flush().unwrap(); + } + writer.close().unwrap(); + let file = FileHandle::new( + FileMeta { + region_id: metadata.region_id, + file_id: FileId::random(), + time_range: (Timestamp::new_millisecond(0), Timestamp::new_millisecond(3)), + file_size: bytes.len() as u64, + num_rows: 4, + num_row_groups: 2, + preserve_row_sequence: true, + ..Default::default() + }, + crate::test_util::new_noop_file_purger(), + ); + env.access_layer + .object_store() + .write(&file.file_path("", PathType::Bare), bytes) + .await + .unwrap(); + files.push(file); + } + (metadata, files) + } + + #[tokio::test] + async fn test_range_cache_keeps_cross_sst_dedup_and_task_membership() { + use store_api::storage::consts::SEQUENCE_COLUMN_NAME; + + for capacity in [0, 1024, 1024 * 1024] { + let env = SchedulerEnv::new().await; + let (metadata, files) = write_overlapping_sst_files(&env).await; + let strategy = CacheStrategy::EnableAll(Arc::new( + CacheManager::builder() + .range_result_cache_size(capacity) + .build(), + )); + for (case, (max_sequence, last_row)) in [ + (1, false), + (2, false), + (1, false), + (2, true), + (2, false), + (1, true), + ] + .into_iter() + .enumerate() + { + let mapper = + FlatProjectionMapper::new(&metadata, 0..metadata.column_metadatas.len()) + .unwrap(); + let input = ScanInput::builder(env.access_layer.clone(), mapper) + .with_files(files.clone()) + .with_cache(strategy.clone()) + .with_append_mode(false) + .with_predicate( + PredicateGroup::new(&metadata, &[col("tag_0").eq(lit("a"))]).unwrap(), + ) + .with_batch_size(1) + .with_snapshot_sequence(Some(max_sequence)) + .with_series_row_selector( + last_row.then_some(TimeSeriesRowSelector::LastRow { after_merge: true }), + ) + .with_sequence_range(Some(SequenceRange::GtLtEq { + min: 0, + max: max_sequence, + })) + .build(); + let context = Arc::new(StreamContext::seq_scan_ctx(input)); + assert_eq!(context.ranges.len(), 1); + assert_eq!(context.ranges[0].row_group_indices.len(), 2); + let range = context.ranges[0].new_partition_range(0); + let (_, metrics) = test_cache_context(&strategy); + for run in 0..2 { + let partition = new_partition_pruner(&context, &range); + let semaphore = Arc::new(tokio::sync::Semaphore::new(1)); + let batches = tokio::time::timeout(std::time::Duration::from_secs(5), async { + let mut stream = sst_range_stream( + &context, + &range, + &metrics, + partition, + Some(semaphore.clone()), + ) + .await; + // The first scan at snapshot 2 has no completed cache entry. + if capacity > 0 && case == 1 && run == 0 { + assert!(build_range_cache_key(&context, &range, None).is_some()); + let (_, follower_metrics) = test_cache_context(&strategy); + let follower = sst_range_stream( + &context, + &range, + &follower_metrics, + new_partition_pruner(&context, &range), + Some(semaphore.clone()), + ) + .await; + let first = stream.try_next().await.unwrap().unwrap(); + assert!(first.num_rows() < 4); + drop(stream); + let batches = follower.try_collect::>().await.unwrap(); + assert!(batches.len() > 1); + assert_eq!(batches.first().unwrap(), &first); + assert_eq!(semaphore.available_permits(), 1); + let owner: serde_json::Value = + serde_json::from_str(&format!("{metrics:?}")).unwrap(); + let follower: serde_json::Value = + serde_json::from_str(&format!("{follower_metrics:?}")).unwrap(); + assert!(owner["metrics"]["num_file_ranges"].as_u64().unwrap() > 0); + assert_eq!(follower["metrics"]["num_file_ranges"], 0); + batches + } else { + stream.try_collect::>().await.unwrap() + } + }) + .await + .unwrap(); + assert!( + u64_values(&batches, SEQUENCE_COLUMN_NAME) + .iter() + .all(|sequence| *sequence == Some(max_sequence)) + ); + let values = u64_values(&batches, "field_0") + .into_iter() + .map(Option::unwrap) + .collect::>(); + let expected = if last_row { + vec![max_sequence * 100 + 3] + } else { + (0..4).map(|offset| max_sequence * 100 + offset).collect() + }; + assert_eq!(values, expected); + } + if capacity > 1024 { + let key = build_range_cache_key(&context, &range, None).unwrap(); + let value = wait_cached(&strategy, &key).await; + assert!(value.cached_batches().next().is_some()); + } + } + } + } + + fn range_cache_strategy(capacity: u64) -> CacheStrategy { + CacheStrategy::EnableAll(Arc::new( + CacheManager::builder() + .range_result_cache_size(capacity) + .build(), + )) + } + + #[tokio::test] + async fn test_shared_slices_compact_before_budget_fallback() { + for (capacity, count) in [(64 * 1024, 16), (16 * 1024 * 1024, 128)] { + let strategy = range_cache_strategy(capacity); + let (key, metrics) = test_cache_context(&strategy); + let backing = make_large_binary_batch(1024, 256); + let mut batches: Vec<_> = (0..count).map(|i| backing.slice(i, 1)).collect(); + batches.insert(count / 2, backing.slice(count / 2, 0)); + let input = batches.clone(); + let (mut owner, _) = + get_or_load_range_result(&strategy, key.clone(), &metrics, || async { + Ok((source_stream(input), 1)) + }) + .await + .unwrap(); + let (slow, _) = get_or_load_range_result(&strategy, key.clone(), &metrics, || async { + panic!("reader must share sliced source") + }) .await .unwrap(); - - assert_eq!(replayed.len(), 2); - assert_eq!(replayed[0].num_rows(), 2); - assert_eq!(replayed[1].num_rows(), 1); + for expected in &batches[..count * 3 / 4] { + assert_eq!(&owner.try_next().await.unwrap().unwrap(), expected); + } + let (late, _) = get_or_load_range_result(&strategy, key.clone(), &metrics, || async { + panic!("compacted source must still admit readers") + }) + .await + .unwrap(); + drop(owner); + assert_eq!(slow.try_collect::>().await.unwrap(), batches); + assert_eq!(late.try_collect::>().await.unwrap(), batches); + let value = wait_cached(&strategy, &key).await; + assert!(value.estimated_size() < 64 * 1024); + assert_eq!( + get_or_load_range_result(&strategy, key.clone(), &metrics, || async { + panic!("cached result") + }) + .await + .unwrap() + .0 + .try_collect::>() + .await + .unwrap(), + batches + ); + } } #[tokio::test] - async fn dropped_cache_buffer_releases_batches_while_waiting_for_memory() { + async fn test_failed_shared_stream_does_not_poison_a_later_request() { + let pool_a: Arc = + Arc::new(datafusion::execution::memory_pool::UnboundedMemoryPool::default()); + let pool_b: Arc = + Arc::new(datafusion::execution::memory_pool::UnboundedMemoryPool::default()); + let scope_a = RangeCacheScope::for_memory_pool(&pool_a); + let scope_b = RangeCacheScope::for_memory_pool(&pool_b); let strategy = test_cache_strategy(); - let limiter = strategy.range_result_memory_limiter().unwrap(); - // Take every permit so the concat task blocks before it can compact. - let _permit = limiter + let (key, metrics) = test_cache_context(&strategy); + let (mut failed, _) = get_or_load_range_result_in_scope( + &strategy, + key.clone(), + &metrics, + scope_a, + || async { + let source = futures::stream::iter(vec![ + Ok(make_batch(&[7])), + Err(crate::error::UnexpectedSnafu { + reason: "transient source error", + } + .build()), + ]); + Ok((Box::pin(source) as BoxedRecordBatchStream, 3)) + }, + ) + .await + .unwrap(); + // Keep a second reader alive so the failed generation remains registered. + let (mut existing, _) = get_or_load_range_result_in_scope( + &strategy, + key.clone(), + &metrics, + scope_a, + || async { panic!("existing reader must share the source") }, + ) + .await + .unwrap(); + let (other_pool, _) = get_or_load_range_result_in_scope( + &strategy, + key.clone(), + &metrics, + scope_b, + || async { Ok((source_stream(vec![make_batch(&[8])]), 1)) }, + ) + .await + .unwrap(); + assert_eq!(failed.try_next().await.unwrap().unwrap(), make_batch(&[7])); + let error = failed.try_next().await.unwrap_err(); + let crate::error::Error::LoadRangeResult { source, .. } = error else { + panic!("expected shared source error, got {error:?}"); + }; + assert!( + matches!(source.as_ref(), crate::error::Error::Unexpected { reason, .. } + if reason == "transient source error") + ); + let expected = make_batch(&[1, 2, 3]); + let (mut retry, _) = get_or_load_range_result_in_scope( + &strategy, + key.clone(), + &metrics, + scope_a, + || async { + Ok(( + Box::pin(futures::stream::iter([Ok(expected.clone())])) + as BoxedRecordBatchStream, + 3, + )) + }, + ) + .await + .unwrap(); + assert_eq!(retry.try_next().await.unwrap().unwrap(), expected); + assert_eq!( + existing.try_next().await.unwrap().unwrap(), + make_batch(&[7]) + ); + let error = existing.try_next().await.unwrap_err(); + let crate::error::Error::LoadRangeResult { source, .. } = error else { + panic!("expected shared source error, got {error:?}"); + }; + assert!( + matches!(source.as_ref(), crate::error::Error::Unexpected { reason, .. } + if reason == "transient source error") + ); + assert_eq!( + other_pool.try_collect::>().await.unwrap(), + vec![make_batch(&[8])] + ); + wait_cached(&strategy, &key).await; + let (cached, _) = + get_or_load_range_result_in_scope(&strategy, key, &metrics, scope_a, || async { + panic!("completed result must be reusable across pools") + }) + .await + .unwrap(); + assert_eq!( + cached.try_collect::>().await.unwrap(), + vec![make_batch(&[8])] + ); + } + + #[tokio::test] + async fn test_shared_ranges_build_once_before_consumers_poll() { + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::time::Duration; + + for enabled in [false, true] { + let strategy = range_cache_strategy(if enabled { 1024 * 1024 } else { 0 }); + let (key, metrics) = test_cache_context(&strategy); + let builds = AtomicUsize::new(0); + let gate = Arc::new(tokio::sync::Semaphore::new(0)); + let produced = Arc::new(AtomicUsize::new(0)); + let expected: Vec<_> = (0..32).map(|n| make_batch(&[n])).collect(); + let mut requests = (0..8) + .map(|_| { + Box::pin(get_or_load_range_result( + &strategy, + key.clone(), + &metrics, + || async { + builds.fetch_add(1, Ordering::Relaxed); + let produced = produced.clone(); + let gate = gate.clone(); + let stream = Box::pin(try_stream! { + gate.acquire().await.unwrap().forget(); + for n in 0..32 { + tokio::task::yield_now().await; + produced.fetch_add(1, Ordering::Relaxed); + yield make_batch(&[n]); + } + }); + Ok((stream as BoxedRecordBatchStream, 3)) + }, + )) + }) + .collect::>(); + // Polling all builders to completion must not require consuming any + // returned stream, matching the series distributor's construction order. + let streams = tokio::time::timeout( + Duration::from_secs(5), + futures::future::join_all(&mut requests), + ) + .await + .unwrap(); + assert_eq!(builds.load(Ordering::Relaxed), if enabled { 1 } else { 8 }); + gate.add_permits(8); + let results = futures::future::join_all( + streams + .into_iter() + .map(|stream| stream.unwrap().0.try_collect::>()), + ) + .await; + for result in results { + assert_eq!(result.unwrap(), expected); + } + assert_eq!( + produced.load(Ordering::Relaxed), + if enabled { 32 } else { 256 } + ); + } + } + + #[tokio::test] + async fn test_oversized_shared_range_keeps_existing_readers_and_declines_new_readers() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + let reads = Arc::new(AtomicUsize::new(0)); + let count = reads.clone(); + let strategy = range_cache_strategy(1024); + let (key, metrics) = test_cache_context(&strategy); + let expected = vec![ + make_batch(&[1]), + make_batch(&[2, 3]), + make_batch(&vec![5; 256]), + make_batch(&[4]), + ]; + let input = expected.clone(); + let (mut owner, _) = get_or_load_range_result(&strategy, key.clone(), &metrics, || async { + let source = futures::stream::iter(input.into_iter().map(move |batch| { + count.fetch_add(1, Ordering::Relaxed); + Ok(batch) + })); + Ok((Box::pin(source) as BoxedRecordBatchStream, 1)) + }) + .await + .unwrap(); + assert_eq!(reads.load(Ordering::Relaxed), 0); + let limiter = strategy.range_result_cache().unwrap().retention_memory(); + assert_eq!(limiter.available_permits(), 1); + assert!(strategy.get_range_result(&key).is_none()); + let (mut follower, _) = + get_or_load_range_result(&strategy, key.clone(), &metrics, || async { + panic!("existing follower must share") + }) + .await + .unwrap(); + assert_eq!(owner.try_next().await.unwrap().unwrap(), expected[0]); + assert_eq!(owner.try_next().await.unwrap().unwrap(), expected[1]); + assert_eq!(owner.try_next().await.unwrap().unwrap(), expected[2]); + let (late, _) = get_or_load_range_result(&strategy, key.clone(), &metrics, || async { + Ok((source_stream(vec![make_batch(&[9])]), 1)) + }) + .await + .unwrap(); + assert_eq!( + late.try_collect::>().await.unwrap(), + vec![make_batch(&[9])] + ); + // The first two batch views share one compacted allocation reservation. + assert_eq!(limiter.available_permits(), 0); + assert_eq!(follower.try_next().await.unwrap().unwrap(), expected[0]); + assert_eq!(limiter.available_permits(), 0); + assert_eq!(follower.try_next().await.unwrap().unwrap(), expected[1]); + assert_eq!(limiter.available_permits(), 1); + assert_eq!( + follower.try_collect::>().await.unwrap(), + expected[2..] + ); + assert_eq!( + owner.try_collect::>().await.unwrap(), + vec![make_batch(&[4])] + ); + assert!(strategy.get_range_result(&key).is_none()); + assert_eq!(reads.load(Ordering::Relaxed), 4); + assert_eq!(limiter.available_permits(), 1); + let (retry, _) = get_or_load_range_result(&strategy, key.clone(), &metrics, || async { + Ok((source_stream(vec![]), 1)) + }) + .await + .unwrap(); + assert!(retry.try_collect::>().await.unwrap().is_empty()); + wait_cached(&strategy, &key).await; + } + + #[tokio::test] + async fn test_dropping_reader_releases_unpolled_or_pending_source() { + for polled in [false, true] { + let strategy = test_cache_strategy(); + let (key, metrics) = test_cache_context(&strategy); + let source_gate = Arc::new(tokio::sync::Semaphore::new(1)); + let permit = source_gate.clone().acquire_owned().await.unwrap(); + let (mut stream, _) = + get_or_load_range_result(&strategy, key.clone(), &metrics, || async { + let source = Box::pin(try_stream! { + let _permit = permit; + std::future::pending::<()>().await; + yield make_batch(&[1]); + }); + Ok((source as BoxedRecordBatchStream, 1)) + }) + .await + .unwrap(); + assert_eq!(source_gate.available_permits(), 0); + if polled { + let next = stream.try_next(); + tokio::pin!(next); + assert!(futures::poll!(&mut next).is_pending()); + } + assert_eq!(source_gate.available_permits(), 0); + drop(stream); + assert_eq!(source_gate.available_permits(), 1); + assert!(strategy.get_range_result(&key).is_none()); + let (retry, _) = get_or_load_range_result(&strategy, key, &metrics, || async { + Ok((source_stream(vec![make_batch(&[2])]), 1)) + }) + .await + .unwrap(); + assert_eq!( + retry.try_collect::>().await.unwrap(), + vec![make_batch(&[2])] + ); + } + } + + #[tokio::test] + async fn test_cancelling_initial_reader_preserves_followers_and_releases_memory() { + let strategy = test_cache_strategy(); + let (key, metrics) = test_cache_context(&strategy); + let limiter = strategy.range_result_cache().unwrap().retention_memory(); + let initial = limiter.available_permits(); + let gate = Arc::new(tokio::sync::Notify::new()); + let producer_gate = gate.clone(); + let (mut owner, _) = get_or_load_range_result(&strategy, key.clone(), &metrics, || async { + let source = Box::pin(try_stream! { + yield make_batch(&[1]); + producer_gate.notified().await; + yield make_batch(&[2]); + }); + Ok((source as BoxedRecordBatchStream, 1)) + }) + .await + .unwrap(); + let (mut follower, _) = + get_or_load_range_result(&strategy, key.clone(), &metrics, || async { + panic!("follower must share") + }) + .await + .unwrap(); + assert_eq!(limiter.available_permits(), initial); + assert_eq!(owner.try_next().await.unwrap().unwrap(), make_batch(&[1])); + assert!(limiter.available_permits() < initial); + drop(owner); + assert_eq!( + follower.try_next().await.unwrap().unwrap(), + make_batch(&[1]) + ); + gate.notify_one(); + assert_eq!( + follower.try_next().await.unwrap().unwrap(), + make_batch(&[2]) + ); + // Stop before EOF: unfinished results must not be cached. + drop(follower); + assert_eq!(limiter.available_permits(), initial); + assert!(strategy.get_range_result(&key).is_none()); + } + + #[tokio::test] + async fn test_completed_stream_does_not_wait_for_cache_compaction() { + let strategy = test_cache_strategy(); + let (key, metrics) = test_cache_context(&strategy); + let limiter = strategy.range_concat_memory().unwrap(); + let permit = limiter .acquire(limiter.available_permits() * limiter.permit_bytes()) .await .unwrap(); - let batch = make_batch(&vec![1; DEFAULT_READ_BATCH_SIZE / 2 + 1]); - let weak = Arc::downgrade(batch.column(0)); - let mut buffer = CacheBatchBuffer::new(&strategy); - buffer.push(batch.clone()).unwrap(); - buffer.push(batch).unwrap(); - // Both batches were handed to the task instead of staying in the buffer. - assert!(buffer.buffered_batches.is_empty()); - - // Cancel only once the task is parked on the permit, otherwise the abort - // could land on a task that was never polled. + let expected = vec![make_batch(&[1]), make_batch(&[2])]; + let input = expected.clone(); + let (stream, _) = get_or_load_range_result(&strategy, key.clone(), &metrics, || async { + Ok((source_stream(input), 1)) + }) + .await + .unwrap(); + assert_eq!(stream.try_collect::>().await.unwrap(), expected); tokio::time::timeout(std::time::Duration::from_secs(5), async { while limiter.waited_acquires() == 0 { tokio::task::yield_now().await; } }) .await - .expect("concat task must wait for a memory permit"); - drop(buffer); - - tokio::time::timeout(std::time::Duration::from_secs(5), async { - while weak.upgrade().is_some() { - tokio::task::yield_now().await; - } + .expect("background compaction did not reach its memory wait"); + assert!(strategy.get_range_result(&key).is_none()); + let (joined, _) = get_or_load_range_result(&strategy, key.clone(), &metrics, || async { + panic!("finalizing scan must remain discoverable") }) .await - .expect("cancelled cache work must release input without waiting for a permit"); + .unwrap(); + assert_eq!(joined.try_collect::>().await.unwrap(), expected); + drop(permit); + wait_cached(&strategy, &key).await; } #[tokio::test] - async fn cache_batch_buffer_finishes_pending_batches() { - let strategy = test_cache_strategy(); - let batch = make_batch(&[1, 2, 3]); - let expected_size = batch.get_array_memory_size(); - - let mut buffer = CacheBatchBuffer::new(&strategy); - buffer.push(batch).unwrap(); - - let value = buffer.finish().await.unwrap().unwrap(); - assert_eq!(value.cached_batches.len(), 1); - assert_eq!(value.cached_batches[0].slice_lengths, vec![3]); - assert_eq!(value.estimated_batches_size, expected_size); - } - - #[tokio::test] - async fn cache_batch_buffer_compacts_when_rows_exceed_default_batch_size() { - let strategy = test_cache_strategy(); - let batch = make_batch(&vec![1; DEFAULT_READ_BATCH_SIZE / 2 + 1]); - - let mut buffer = CacheBatchBuffer::new(&strategy); - buffer.push(batch.clone()).unwrap(); - buffer.push(batch).unwrap(); - - assert_eq!(buffer.buffered_rows, 0); - assert!(buffer.buffered_batches.is_empty()); - - let value = buffer.finish().await.unwrap().unwrap(); - assert_eq!(value.cached_batches.len(), 1); - assert_eq!( - value.cached_batches[0].slice_lengths, - vec![ - DEFAULT_READ_BATCH_SIZE / 2 + 1, - DEFAULT_READ_BATCH_SIZE / 2 + 1 - ] - ); - } - - #[tokio::test] - async fn cache_batch_buffer_compacts_when_buffered_size_exceeds_threshold() { - let large_batch = make_large_binary_batch(DEFAULT_READ_BATCH_SIZE, 4096); - let strategy = CacheStrategy::EnableAll(Arc::new( - CacheManager::builder() - .range_result_cache_size((large_batch.get_array_memory_size() * 3) as u64) - .build(), - )); - - let mut buffer = CacheBatchBuffer::new(&strategy); - buffer.push(large_batch.clone()).unwrap(); - - assert_eq!(buffer.buffered_rows, large_batch.num_rows()); - assert_eq!(buffer.buffered_batches.len(), 1); - - buffer.push(large_batch.clone()).unwrap(); - - assert_eq!(buffer.buffered_rows, 0); - assert!(buffer.buffered_batches.is_empty()); - - let value = buffer.finish().await.unwrap().unwrap(); - assert_eq!(value.cached_batches.len(), 1); - assert_eq!( - value.cached_batches[0].slice_lengths, - vec![large_batch.num_rows(), large_batch.num_rows()] - ); - } - - #[tokio::test] - async fn cache_batch_buffer_skips_cache_when_compacted_size_exceeds_limit() { - let large_batch = make_large_binary_batch(DEFAULT_READ_BATCH_SIZE / 2 + 1, 4096); - // Budget only fits two large batches. - let budget = (large_batch.get_array_memory_size() as u64) * 2 + 1; - let strategy = CacheStrategy::EnableAll(Arc::new( - CacheManager::builder() - .range_result_cache_size(budget) - .build(), - )); - - let mut buffer = CacheBatchBuffer::new(&strategy); - for _ in 0..4 { - buffer.push(large_batch.clone()).unwrap(); + async fn test_range_load_compacts_at_thresholds_and_keeps_tail() { + let row_batch = make_batch(&vec![1; DEFAULT_READ_BATCH_SIZE / 2 + 1]); + // Two batches exceed the byte threshold without reaching the row threshold. + let byte_batch = make_large_binary_batch(1024, 4097); + assert!(byte_batch.num_rows() * 2 <= DEFAULT_READ_BATCH_SIZE); + assert!(byte_batch.get_array_memory_size() * 2 > RANGE_CACHE_COMPACT_THRESHOLD_BYTES); + assert!(row_batch.get_array_memory_size() * 2 < RANGE_CACHE_COMPACT_THRESHOLD_BYTES); + for batches in [ + vec![], + vec![make_batch(&[1, 2]), make_batch(&[3]), make_batch(&[4, 5])], + vec![make_batch(&[1, 2, 3])], + vec![row_batch.clone(), row_batch.clone(), row_batch], + vec![byte_batch.clone(), byte_batch.clone(), byte_batch], + ] { + let input_size = batches + .iter() + .map(RecordBatch::get_array_memory_size) + .sum::(); + let strategy = range_cache_strategy((input_size * 2).max(1024 * 1024) as u64); + let (key, metrics) = test_cache_context(&strategy); + let input = batches.clone(); + let expected_batch_size = 7; + let (stream, batch_size) = + get_or_load_range_result(&strategy, key.clone(), &metrics, || async move { + Ok((source_stream(input), expected_batch_size)) + }) + .await + .unwrap(); + assert_eq!(batch_size, expected_batch_size); + assert_eq!(stream.try_collect::>().await.unwrap(), batches); + let value = wait_cached(&strategy, &key).await; + let expected_lengths = if batches.is_empty() { + vec![] + } else if batches.iter().map(RecordBatch::num_rows).sum::() <= 5 { + vec![batches.iter().map(RecordBatch::num_rows).collect()] + } else { + vec![ + vec![batches[0].num_rows(), batches[1].num_rows()], + vec![batches[2].num_rows()], + ] + }; + assert_eq!( + value + .cached_batches() + .map(|(_, lengths)| lengths.to_vec()) + .collect::>(), + expected_lengths + ); + assert_eq!( + value.estimated_batches_size(), + value + .cached_batches() + .map(|(batch, _)| batch.get_array_memory_size()) + .sum::() + ); + let (hit, batch_size) = get_or_load_range_result(&strategy, key, &metrics, || async { + panic!("completed range must be reused without rebuilding"); + }) + .await + .unwrap(); + assert_eq!(batch_size, expected_batch_size); + assert_eq!(hit.try_collect::>().await.unwrap(), batches); } - assert!(!matches!(buffer.finish().await, Some(Ok(_)))); } } diff --git a/src/mito2/src/read/scan_util.rs b/src/mito2/src/read/scan_util.rs index 5f82743f92d..945dafb55fe 100644 --- a/src/mito2/src/read/scan_util.rs +++ b/src/mito2/src/read/scan_util.rs @@ -34,6 +34,7 @@ use smallvec::SmallVec; use snafu::ResultExt; use store_api::storage::{RegionId, SequenceRange, TimeSeriesRowSelector}; +use crate::cache::range_cache::RangeCacheMetricsReport; use crate::error::{ComputeArrowSnafu, Result}; use crate::memtable::MemScanMetrics; use crate::metrics::{ @@ -915,6 +916,26 @@ impl DedupMetricsReport for PartitionMetricsInner { } } +impl RangeCacheMetricsReport for PartitionMetricsInner { + /// Increments the total bytes added to the range cache. + fn inc_range_cache_size(&self, size: usize) { + let mut metrics = self.metrics.lock().unwrap(); + metrics.range_cache_size += size; + } + + /// Increments the range cache hit counter. + fn inc_range_cache_hit(&self) { + let mut metrics = self.metrics.lock().unwrap(); + metrics.range_cache_hit += 1; + } + + /// Increments the range cache miss counter. + fn inc_range_cache_miss(&self) { + let mut metrics = self.metrics.lock().unwrap(); + metrics.range_cache_miss += 1; + } +} + impl Drop for PartitionMetricsInner { fn drop(&mut self) { self.on_finish(false); @@ -1107,22 +1128,9 @@ impl PartitionMetrics { self.0.clone() } - /// Increments the total bytes added to the range cache. - pub(crate) fn inc_range_cache_size(&self, size: usize) { - let mut metrics = self.0.metrics.lock().unwrap(); - metrics.range_cache_size += size; - } - - /// Increments the range cache hit counter. - pub(crate) fn inc_range_cache_hit(&self) { - let mut metrics = self.0.metrics.lock().unwrap(); - metrics.range_cache_hit += 1; - } - - /// Increments the range cache miss counter. - pub(crate) fn inc_range_cache_miss(&self) { - let mut metrics = self.0.metrics.lock().unwrap(); - metrics.range_cache_miss += 1; + /// Returns the range-cache reporter while preserving query metric ownership. + pub(crate) fn range_cache_metrics_reporter(&self) -> Arc { + self.0.clone() } } @@ -1812,6 +1820,34 @@ mod tests { use crate::test_util::memtable_util::metadata_for_test; use crate::test_util::scheduler_util::SchedulerEnv; + #[test] + fn test_range_cache_reporter_preserves_metrics_and_lifetime() { + let metrics = PartitionMetrics::new( + RegionId::new(1, 1), + 0, + "test_range_cache_reporter", + Instant::now(), + false, + &ExecutionPlanMetricsSet::new(), + ); + let weak = Arc::downgrade(&metrics.0); + let reporter = metrics.range_cache_metrics_reporter(); + reporter.inc_range_cache_hit(); + reporter.inc_range_cache_miss(); + drop(metrics); + // Background publication can still report after the query drops its handle. + reporter.inc_range_cache_size(128); + { + let inner = weak.upgrade().unwrap(); + let metrics = inner.metrics.lock().unwrap(); + assert_eq!(metrics.range_cache_hit, 1); + assert_eq!(metrics.range_cache_miss, 1); + assert_eq!(metrics.range_cache_size, 128); + } + drop(reporter); + assert!(weak.upgrade().is_none()); + } + struct EmptyIterBuilder; impl IterBuilder for EmptyIterBuilder { diff --git a/src/mito2/src/read/seq_scan.rs b/src/mito2/src/read/seq_scan.rs index d8759ffaace..a0513300910 100644 --- a/src/mito2/src/read/seq_scan.rs +++ b/src/mito2/src/read/seq_scan.rs @@ -302,13 +302,14 @@ impl SeqScan { file_scan_semaphore: Option>, merge_semaphore: Option>, ) -> Result<(BoxedRecordBatchStream, usize)> { - let cache_key = build_range_cache_key(stream_ctx, part_range).map(|key| { - if read_mode == FlatReadMode::CrossRangeMerge { - key.for_cross_range_merge() - } else { - key - } - }); + let cache_key = + build_range_cache_key(stream_ctx, part_range, Some(&partition_pruner)).map(|key| { + if read_mode == FlatReadMode::CrossRangeMerge { + key.for_cross_range_merge() + } else { + key + } + }); let ctx = stream_ctx.clone(); let range = *part_range; let metrics = part_metrics.clone(); diff --git a/src/mito2/src/read/series_candidate.rs b/src/mito2/src/read/series_candidate.rs index 185d3ea0dff..00e0c556d81 100644 --- a/src/mito2/src/read/series_candidate.rs +++ b/src/mito2/src/read/series_candidate.rs @@ -156,17 +156,6 @@ impl SeriesCandidateScanner { for (range_idx, part_range) in all_ranges.into_iter().enumerate() { let range_builder = range_builder.clone(); tasks.push(common_runtime::spawn_query(async move { - let _permit = range_builder - .range_semaphore - .clone() - .acquire_owned() - .await - .map_err(|error| { - UnexpectedSnafu { - reason: format!("failed to acquire candidate range permit: {error}"), - } - .build() - })?; range_builder .build_range_stream(part_range, range_idx) .await @@ -367,7 +356,7 @@ impl SeriesCandidateRangeBuilder { let cache_key = if replaced_sources { None } else { - build_candidate_range_cache_key(&self.stream_ctx, &part_range) + build_candidate_range_cache_key(&self.stream_ctx, &part_range, &self.partition_pruner) }; let builder = self.clone(); let (stream, _) = CachedRangeReader { @@ -376,7 +365,7 @@ impl SeriesCandidateRangeBuilder { key: cache_key, metrics: &self.part_metrics, } - .build(move || async move { + .build_with_memory_pool(&self.memory_pool, move || async move { let stream = builder .build_uncached_range_stream(part_range, merge_partition) .await?; @@ -391,6 +380,14 @@ impl SeriesCandidateRangeBuilder { part_range: PartitionRange, merge_partition: usize, ) -> Result { + // Release the build permit before the cache loader drives parallel sources, + // which acquire permits from this same semaphore. + let _permit = self.range_semaphore.acquire().await.map_err(|error| { + UnexpectedSnafu { + reason: format!("failed to acquire candidate range permit: {error}"), + } + .build() + })?; let range_meta = &self.stream_ctx.ranges[part_range.identifier]; let mut sources = Vec::with_capacity(range_meta.row_group_indices.len()); for index in &range_meta.row_group_indices { @@ -775,6 +772,152 @@ mod tests { use crate::test_util::scheduler_util::SchedulerEnv; use crate::test_util::sst_util::sst_region_metadata_with_encoding; + #[tokio::test] + async fn shared_candidate_ranges_isolate_query_memory_pools() { + use datafusion::execution::memory_pool::GreedyMemoryPool; + use datatypes::arrow::array::BinaryDictionaryBuilder; + use parquet::arrow::ArrowWriter; + use parquet::file::metadata::KeyValue; + use parquet::file::properties::WriterProperties; + use store_api::region_request::PathType; + + use crate::read::scan_region::PredicateGroup; + use crate::sst::parquet::PARQUET_METADATA_KEY; + use crate::sst::parquet::flat_format::primary_key_column_index; + use crate::sst::{FlatSchemaOptions, to_flat_sst_arrow_schema}; + use crate::test_util::sst_util::{new_record_batch_by_range, new_sparse_primary_key}; + + let env = SchedulerEnv::new().await; + let metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let dense = new_record_batch_by_range(&["a", "x"], 0, 2); + let mut columns: Vec = vec![ + Arc::new(UInt32Array::from(vec![1, 1])), + Arc::new(UInt64Array::from(vec![10, 20])), + ]; + columns.extend_from_slice(dense.columns()); + let mut keys = BinaryDictionaryBuilder::::new(); + for tsid in [10, 20] { + keys.append(new_sparse_primary_key(&["a", "x"], &metadata, 1, tsid)) + .unwrap(); + } + let pk_index = primary_key_column_index(columns.len()); + columns[pk_index] = Arc::new(keys.finish()); + let batch = RecordBatch::try_new( + to_flat_sst_arrow_schema(&metadata, &FlatSchemaOptions::default()), + columns, + ) + .unwrap(); + let properties = WriterProperties::builder() + .set_key_value_metadata(Some(vec![KeyValue::new( + PARQUET_METADATA_KEY.to_string(), + metadata.to_json().unwrap(), + )])) + .build(); + let mut bytes = Vec::new(); + let mut writer = + ArrowWriter::try_new(&mut bytes, batch.schema(), Some(properties)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + let mut files = Vec::new(); + for _ in 0..2 { + let file = FileHandle::new( + FileMeta { + region_id: metadata.region_id, + file_id: FileId::random(), + time_range: (Timestamp::new_millisecond(0), Timestamp::new_millisecond(1)), + file_size: bytes.len() as u64, + num_rows: 2, + num_row_groups: 1, + preserve_row_sequence: true, + ..Default::default() + }, + new_noop_file_purger(), + ); + env.access_layer + .object_store() + .write(&file.file_path("", PathType::Bare), bytes.clone()) + .await + .unwrap(); + files.push(file); + } + let strategy = CacheStrategy::EnableAll(Arc::new( + CacheManager::builder() + .range_result_cache_size(1024 * 1024) + .build(), + )); + let mapper = + FlatProjectionMapper::new(&metadata, 0..metadata.column_metadatas.len()).unwrap(); + let input = ScanInput::builder(env.access_layer.clone(), mapper) + .with_files(files) + .with_cache(strategy) + .with_predicate( + PredicateGroup::new(&metadata, &[col("__table_id").eq(lit(1_u32))]).unwrap(), + ) + .build(); + let context = Arc::new(StreamContext::seq_scan_ctx(input)); + let ranges = context.partition_ranges(); + assert_eq!(ranges.len(), 1); + assert_eq!(context.ranges[0].row_group_indices.len(), 2); + let make_scanner = |pool: Arc| { + let pruner = Arc::new(Pruner::new_with_options( + context.clone(), + 1, + PrunerOptions { + retain_builders: true, + enable_predicate_prefilter: false, + }, + )); + let metrics = ExecutionPlanMetricsSet::new(); + let partition_metrics = PartitionMetrics::new( + metadata.region_id, + 1, + "pool-test", + Instant::now(), + false, + &metrics, + ); + SeriesCandidateScanner::try_new( + context.clone(), + vec![ranges.clone()], + pruner, + Arc::new(Semaphore::new(1)), + pool, + metrics, + partition_metrics, + ) + .unwrap() + }; + let low_pool = Arc::new(GreedyMemoryPool::new(0)); + let high_pool = Arc::new(UnboundedMemoryPool::default()); + let low = make_scanner(low_pool.clone()); + let high = make_scanner(high_pool.clone()); + tokio::time::timeout(std::time::Duration::from_secs(5), async { + // Register both scans before driving either source to its first batch. + let low_stream = low.build_stream().await.unwrap(); + let high_stream = high.build_stream().await.unwrap(); + let error = low_stream.try_collect::>().await.unwrap_err(); + assert!( + format!("{error:?}").contains("ResourcesExhausted"), + "{error:?}" + ); + let groups = high_stream.try_collect::>().await.unwrap(); + assert_eq!( + groups.into_iter().flatten().collect::>(), + [10, 20] + .map(|tsid| MetricSeriesId { table_id: 1, tsid }) + .to_vec(), + ); + }) + .await + .expect("candidate scans stalled with different query pools"); + drop(low); + drop(high); + assert_eq!(low_pool.reserved(), 0); + assert_eq!(high_pool.reserved(), 0); + } + /// Uses absent SST objects so any covered-SST read fails the test. async fn indexed_scanner() -> (SchedulerEnv, SeriesCandidateScanner, Arc) { let env = SchedulerEnv::new().await; @@ -916,7 +1059,14 @@ mod tests { .partitions .iter() .flatten() - .map(|range| build_candidate_range_cache_key(&scanner.stream_ctx, range).unwrap()) + .map(|range| { + build_candidate_range_cache_key( + &scanner.stream_ctx, + range, + &scanner.candidate_pruner, + ) + .unwrap() + }) .collect(); let groups = scanner .build_stream() @@ -957,6 +1107,7 @@ mod tests { // Index-backed ranges must not publish their empty residual streams as // complete candidate results for a future scan without the index. assert!(!format!("{:?}", scanner.part_metrics).contains("range_cache_miss")); + assert!(!format!("{:?}", scanner.part_metrics).contains("range_cache_hit")); for key in keys { assert!( scanner diff --git a/src/mito2/src/read/series_reader.rs b/src/mito2/src/read/series_reader.rs index 021928f2c01..bddac735d11 100644 --- a/src/mito2/src/read/series_reader.rs +++ b/src/mito2/src/read/series_reader.rs @@ -373,7 +373,7 @@ async fn build_series_partition_range( let range = filter.range; let cache_key = filter .enable_range_cache - .then(|| build_series_range_cache_key(&stream_ctx, &part_range, range)) + .then(|| build_series_range_cache_key(&stream_ctx, &part_range, range, &partition_pruner)) .flatten(); let ctx = stream_ctx.clone(); let metrics = part_metrics.clone(); diff --git a/src/mito2/src/sst/file_purger.rs b/src/mito2/src/sst/file_purger.rs index ad44c5bcc05..83147088e15 100644 --- a/src/mito2/src/sst/file_purger.rs +++ b/src/mito2/src/sst/file_purger.rs @@ -432,10 +432,7 @@ mod tests { let cache_key = crate::cache::test_util::range_cache_key(&[sst_file_id.file_id()]); cache.put_range_result( cache_key.clone(), - Arc::new(crate::read::range_cache::RangeScanCacheValue::new( - vec![], - 0, - )), + Arc::new(crate::cache::range_cache::RangeScanCacheValue::empty()), ); let file_purger = create_file_purger( gc_enabled, diff --git a/src/mito2/src/sst/parquet/file_range.rs b/src/mito2/src/sst/parquet/file_range.rs index f79569d6d4e..37bae5f1d34 100644 --- a/src/mito2/src/sst/parquet/file_range.rs +++ b/src/mito2/src/sst/parquet/file_range.rs @@ -589,7 +589,7 @@ impl FileRangeContext { } /// Mode to pre-filter columns in a range. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub enum PreFilterMode { /// Filters all columns. All, From 5fe9bdc814fa0dd81f84514f74b5f63f3cc6fdde Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Tue, 3 Nov 2026 11:43:33 +0000 Subject: [PATCH 160/239] feat(prometheus): store Remote Write v1 and v2 exemplars (#9922) * feat(prometheus): persist Remote Write exemplars Prometheus Remote Write v1 and v2 now persist valid exemplars in the source schema's metrics_exemplars table. A shared row contract keeps inherited series labels separate from exemplar labels, creates or validates an ordinary sparse append table, and preserves compatible pre-created timestamp precision. Malformed observations are dropped and counted without discarding valid samples or sibling exemplars. Permissions and admission precede writes, and v2 written-count headers use confirmed storage progress. Includes decoder, HTTP orchestration and physical-table routing coverage. The JSON2 row-schema helper becomes fallible; its three existing trace callers propagate that error in this layer. Signed-off-by: shuiyisong * feat(exemplar): persist source metric identities Signed-off-by: shuiyisong * fix(prometheus): bound exemplar write memory Signed-off-by: shuiyisong * fix(prometheus): correct batching with exemplars Signed-off-by: shuiyisong * fix(exemplar): accept mixed-case sparse encoding Signed-off-by: shuiyisong * perf(exemplar): reduce write preparation overhead Signed-off-by: shuiyisong * perf(prom): reuse exemplar source targets Signed-off-by: shuiyisong * fix: clippy Signed-off-by: shuiyisong * fix(prometheus): defer exemplar series expansion Signed-off-by: shuiyisong --------- Signed-off-by: shuiyisong --- Cargo.lock | 1 + .../meta/src/ddl/create_table/template.rs | 81 +- src/common/query/src/prelude.rs | 6 + src/frontend/src/instance/prom_store.rs | 69 ++ src/metric-engine/src/row_modifier.rs | 68 +- src/operator/Cargo.toml | 1 + src/operator/src/insert.rs | 471 +++++++++- src/servers/src/batcher/logical_table.rs | 106 +-- .../src/batcher/logical_table/test_util.rs | 96 +++ src/servers/src/exemplar.rs | 310 +++++++ src/servers/src/http/prom_store.rs | 472 +++++++++- src/servers/src/http/prom_store/exemplar.rs | 811 ++++++++++++++++++ src/servers/src/lib.rs | 1 + src/servers/src/metrics.rs | 10 + src/servers/src/otlp/trace/v2.rs | 6 +- src/servers/src/prom_remote_write/decode.rs | 275 +++++- src/servers/src/prom_remote_write/mod.rs | 14 +- .../src/prom_remote_write/row_builder.rs | 2 + src/servers/src/prom_remote_write/v2.rs | 311 +++++-- src/servers/src/query_handler.rs | 59 ++ src/servers/src/row_writer.rs | 22 +- tests-integration/src/tests/instance_test.rs | 566 +++++++++++- .../cases/exemplar_sparse/case.toml | 8 + .../cases/exemplar_sparse/setup.sql | 15 + .../cases/exemplar_sparse/verify.result | 42 + .../cases/exemplar_sparse/verify.sql | 5 + 26 files changed, 3614 insertions(+), 214 deletions(-) create mode 100644 src/servers/src/exemplar.rs create mode 100644 src/servers/src/http/prom_store/exemplar.rs create mode 100644 tests/compatibility/cases/exemplar_sparse/case.toml create mode 100644 tests/compatibility/cases/exemplar_sparse/setup.sql create mode 100644 tests/compatibility/cases/exemplar_sparse/verify.result create mode 100644 tests/compatibility/cases/exemplar_sparse/verify.sql diff --git a/Cargo.lock b/Cargo.lock index a5e8e050a72..67de9d76bab 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10180,6 +10180,7 @@ dependencies = [ "ahash 0.8.12", "api", "arrow 59.2.0", + "arrow-schema 59.2.0", "async-stream", "async-trait", "axum 0.8.4", diff --git a/src/common/meta/src/ddl/create_table/template.rs b/src/common/meta/src/ddl/create_table/template.rs index 45fcd1a0148..766a32780ba 100644 --- a/src/common/meta/src/ddl/create_table/template.rs +++ b/src/common/meta/src/ddl/create_table/template.rs @@ -17,11 +17,16 @@ use std::collections::HashMap; use api::v1::column_def::try_as_column_def; use api::v1::meta::Partition; use api::v1::region::{CreateRequest, RegionColumnDef}; -use api::v1::{ColumnDef, CreateTableExpr, SemanticType}; +use api::v1::{ColumnDataType, ColumnDef, CreateTableExpr, SemanticType}; +use common_catalog::consts::{MITO_ENGINE, MITO2_ENGINE}; use common_telemetry::warn; use snafu::{OptionExt, ResultExt}; -use store_api::metric_engine_consts::LOGICAL_TABLE_METADATA_KEY; +use store_api::metric_engine_consts::{ + DATA_SCHEMA_TABLE_ID_COLUMN_NAME, DATA_SCHEMA_TSID_COLUMN_NAME, LOGICAL_TABLE_METADATA_KEY, + PRIMARY_KEY_ENCODING, +}; use store_api::region_request::RegionRequirements; +use store_api::storage::consts::ReservedColumnId; use store_api::storage::{RegionId, RegionNumber}; use table::metadata::{TableId, TableInfo}; @@ -135,6 +140,28 @@ pub fn build_template_from_raw_table_info_for_physical_table( } pub(crate) fn build_template(create_table_expr: &CreateTableExpr) -> Result { + // Explicit sparse identity columns use the same IDs as metric physical regions. + let reserved_prefix = matches!( + create_table_expr.engine.as_str(), + MITO_ENGINE | MITO2_ENGINE + ) && create_table_expr + .table_options + .get(PRIMARY_KEY_ENCODING) + .is_some_and(|value| value.eq_ignore_ascii_case("sparse")) + && create_table_expr.primary_keys.starts_with(&[ + DATA_SCHEMA_TABLE_ID_COLUMN_NAME.into(), + DATA_SCHEMA_TSID_COLUMN_NAME.into(), + ]) + && [ + (DATA_SCHEMA_TABLE_ID_COLUMN_NAME, ColumnDataType::Uint32), + (DATA_SCHEMA_TSID_COLUMN_NAME, ColumnDataType::Uint64), + ] + .into_iter() + .all(|(name, data_type)| { + create_table_expr.column_defs.iter().any(|column| { + column.name == name && column.data_type == data_type as i32 && !column.is_nullable + }) + }); let column_defs = create_table_expr .column_defs .iter() @@ -159,7 +186,13 @@ pub(crate) fn build_template(create_table_expr: &CreateTableExpr) -> Result>(); @@ -280,6 +313,48 @@ mod tests { use super::*; use crate::key::test_utils; + #[test] + fn test_create_sparse_reserved_column_ids() { + let mut expr = CreateTableExpr { + engine: MITO_ENGINE.into(), + table_options: HashMap::from([(PRIMARY_KEY_ENCODING.into(), "sparse".into())]), + primary_keys: vec![ + DATA_SCHEMA_TABLE_ID_COLUMN_NAME.into(), + DATA_SCHEMA_TSID_COLUMN_NAME.into(), + "job".into(), + ], + time_index: "ts".into(), + column_defs: [ + ("ts", ColumnDataType::TimestampMillisecond), + (DATA_SCHEMA_TABLE_ID_COLUMN_NAME, ColumnDataType::Uint32), + (DATA_SCHEMA_TSID_COLUMN_NAME, ColumnDataType::Uint64), + ("job", ColumnDataType::String), + ] + .into_iter() + .map(|(name, data_type)| ColumnDef { + name: name.into(), + data_type: data_type as i32, + is_nullable: false, + ..Default::default() + }) + .collect(), + ..Default::default() + }; + let template = build_template(&expr).unwrap(); + assert_eq!( + template.primary_key, + vec![ReservedColumnId::table_id(), ReservedColumnId::tsid(), 3] + ); + assert_eq!( + template.column_defs[1].column_id, + ReservedColumnId::table_id() + ); + assert_eq!(template.column_defs[2].column_id, ReservedColumnId::tsid()); + // The names alone must not change ordinary dense-table column allocation. + expr.table_options.clear(); + assert_eq!(build_template(&expr).unwrap().primary_key, vec![1, 2, 3]); + } + #[test] fn test_build_one_sets_partition_expr_per_region() { // minimal template diff --git a/src/common/query/src/prelude.rs b/src/common/query/src/prelude.rs index ca70bf5520f..8bb486a2d30 100644 --- a/src/common/query/src/prelude.rs +++ b/src/common/query/src/prelude.rs @@ -91,6 +91,12 @@ pub fn greptime_native_histogram() -> &'static str { const GREPTIME_TIMESTAMP: &str = "greptime_timestamp"; /// Default value column name constant for backward compatibility. const GREPTIME_VALUE: &str = "greptime_value"; +/// Shared exemplar destination within each source metric's schema. +pub const METRICS_EXEMPLARS_TABLE: &str = "metrics_exemplars"; +/// Query-visible source metric name stored in the exemplar table. +pub const EXEMPLAR_METRIC_COLUMN: &str = "greptime_metric"; +/// JSON2 object containing exemplar-owned labels. +pub const EXEMPLAR_LABELS_COLUMN: &str = "exemplar_labels"; /// Stored series label that opts ordinary float samples into raw-delta math. pub const OTLP_AGGREGATION_TEMPORALITY_LABEL: &str = "otlp_aggregation_temporality"; /// Authoritative raw-delta value for [`OTLP_AGGREGATION_TEMPORALITY_LABEL`]. diff --git a/src/frontend/src/instance/prom_store.rs b/src/frontend/src/instance/prom_store.rs index 94687238ed0..a5e4c30b832 100644 --- a/src/frontend/src/instance/prom_store.rs +++ b/src/frontend/src/instance/prom_store.rs @@ -514,6 +514,75 @@ impl Instance { #[async_trait] impl PromStoreProtocolHandler for Instance { + async fn prepare_metric_write( + &self, + request: RowInsertRequests, + ctx: QueryContextRef, + ) -> ServerResult { + let physical_table = ctx + .extension(PHYSICAL_TABLE_PARAM) + .unwrap_or(GREPTIME_PHYSICAL_TABLE) + .to_string(); + self.inserter + .prepare_metric_row_inserts(request, ctx, &self.statement_executor, physical_table) + .await + .map_err(BoxedError::new) + .context(error::ExecuteGrpcQuerySnafu) + } + + async fn write_prepared_metric( + &self, + request: operator::insert::PreparedMetricInserts, + ctx: QueryContextRef, + ) -> ServerResult { + self.inserter + .handle_prepared_metric_inserts(request, ctx) + .await + .map_err(BoxedError::new) + .context(error::ExecuteGrpcQuerySnafu) + } + + async fn check_exemplar_sources( + &self, + sources: &std::collections::BTreeSet, + ctx: QueryContextRef, + ) -> ServerResult<()> { + let targets = sources + .iter() + .map(|name| { + PermissionTableTarget::new(ctx.current_catalog(), ctx.current_schema(), name) + }) + .collect(); + self.check_table_permission( + &ctx, + PermissionReq::Action(PROM_STORE_WRITE), + PermissionTableTargets::resolved(targets), + ) + .context(AuthSnafu)?; + Ok(()) + } + + async fn write_exemplars_prepared( + &self, + request: RowInsertRequests, + ctx: QueryContextRef, + ) -> ServerResult { + self.write_exemplars_with_sources(request, ctx, &[]).await + } + + async fn write_exemplars_with_sources( + &self, + request: RowInsertRequests, + ctx: QueryContextRef, + sources: &[table::metadata::TableInfoRef], + ) -> ServerResult { + self.inserter + .handle_exemplar_inserts(request, ctx, &self.statement_executor, sources) + .await + .map_err(BoxedError::new) + .context(error::ExecuteGrpcQuerySnafu) + } + async fn pre_write( &self, request: &RowInsertRequests, diff --git a/src/metric-engine/src/row_modifier.rs b/src/metric-engine/src/row_modifier.rs index 7d68f94aa0f..04f82a199fd 100644 --- a/src/metric-engine/src/row_modifier.rs +++ b/src/metric-engine/src/row_modifier.rs @@ -196,25 +196,12 @@ impl RowModifier { } ts_id_gen.finish() } else { - // Slow path: row contains null, recompute label hash - let mut hasher = TsidGenerator::default(); - // 1. Find out label names with non-null values and get the hash. - for (name, value) in iter.primary_keys_with_name() { - // The type is checked before. So only null is ignored. - if let Some(ValueData::StringValue(_)) = &value.value_data { - hasher.write_str(name); - } - } - let label_name_hash = hasher.finish(); - - // 2. Use label name hash as seed and continue with label values. - let mut final_hasher = TsidGenerator::new(label_name_hash); - for (_, value) in iter.primary_keys_with_name() { - if let Some(ValueData::StringValue(value)) = &value.value_data { - final_hasher.write_str(value); - } - } - final_hasher.finish() + TsidGenerator::from_sorted_tags(iter.primary_keys_with_name().filter_map( + |(name, value)| match &value.value_data { + Some(ValueData::StringValue(value)) => Some((name.as_str(), value.as_str())), + _ => None, + }, + )) }; ( @@ -231,6 +218,19 @@ pub struct TsidGenerator { } impl TsidGenerator { + /// Computes the metric TSID from non-null tags sorted by name, retaining empty values. + pub fn from_sorted_tags<'a>(tags: impl Iterator + Clone) -> u64 { + let mut names = Self::default(); + for (name, _) in tags.clone() { + names.write_str(name); + } + let mut values = Self::new(names.finish()); + for (_, value) in tags { + values.write_str(value); + } + values.finish() + } + pub fn new(label_name_hash: u64) -> Self { let mut hasher = FxHasher::default(); hasher.write_u64(label_name_hash); @@ -376,7 +376,7 @@ pub struct RowIter<'a> { impl RowIter<'_> { /// Returns the primary keys with their names. - fn primary_keys_with_name(&self) -> impl Iterator { + fn primary_keys_with_name(&self) -> impl Iterator + Clone { self.index.indices[..self.index.num_primary_key_column] .iter() .map(|idx| { @@ -790,6 +790,34 @@ mod tests { ); } + #[test] + fn test_sorted_tags_match_metric_rows() { + let names = ["b", "a", "absent"]; + for values in [ + [Some("B"), Some(""), Some("C")], + [Some("B"), Some(""), None], + [None, Some("A"), None], + [None, None, None], + ] { + let tags = names + .into_iter() + .zip(values) + .filter_map(|(name, value)| value.map(|value| (name, value))) + .collect::>(); + let tsid = + TsidGenerator::from_sorted_tags(tags.iter().map(|(name, value)| (*name, *value))); + assert_eq!( + tsid, + extract_tsid( + create_multi_label_schema(&names), + create_row_with_nulls(&values), + &create_name_to_column_id(&names), + 42 + ) + ); + } + } + #[test] fn test_tsid_same_with_multiple_null_labels() { // Test with multiple null labels diff --git a/src/operator/Cargo.toml b/src/operator/Cargo.toml index b89bf4b13d3..56c18a15f11 100644 --- a/src/operator/Cargo.toml +++ b/src/operator/Cargo.toml @@ -22,6 +22,7 @@ workspace = true ahash.workspace = true api.workspace = true arrow.workspace = true +arrow-schema.workspace = true async-stream.workspace = true async-trait.workspace = true backon.workspace = true diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index 5ecc01457a1..6e921577b7c 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -34,6 +34,7 @@ use api::v1::{ RowInsertRequest, RowInsertRequests, Rows, SemanticType, }; use arrow::datatypes::{DataType as ArrowDataType, Schema as ArrowSchema}; +use arrow_schema::extension::ExtensionType; use catalog::CatalogManagerRef; use client::{OutputData, OutputMeta}; use common_base::readable_size::ReadableSize; @@ -52,12 +53,17 @@ use common_meta::peer::Peer; use common_meta::rpc::ddl::TriggerReason; use common_query::Output; use common_query::native_histogram::{is_native_histogram_value_type, native_histogram_value_type}; -use common_query::prelude::{greptime_timestamp, greptime_value}; +use common_query::prelude::{ + EXEMPLAR_LABELS_COLUMN, EXEMPLAR_METRIC_COLUMN, METRICS_EXEMPLARS_TABLE, greptime_timestamp, + greptime_value, +}; use common_telemetry::tracing_context::TracingContext; use common_telemetry::{debug, error, warn}; use common_time::Timestamp; use common_time::timestamp::TimeUnit; -use datatypes::schema::{ColumnEncoding, CompressionOptions, SkippingIndexOptions}; +use datatypes::schema::{ + ColumnDefaultConstraint, ColumnEncoding, CompressionOptions, SkippingIndexOptions, +}; use futures_util::future; use meter_core::data::MeterRecord; use meter_macros::write_meter; @@ -69,12 +75,14 @@ use sql::partition::partition_rule_for_hexstring; use sql::statements::create::Partitions; use sql::statements::insert::Insert; use store_api::metric_engine_consts::{ - LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME, PHYSICAL_TABLE_METADATA_KEY, + DATA_SCHEMA_TABLE_ID_COLUMN_NAME, DATA_SCHEMA_TSID_COLUMN_NAME, LOGICAL_TABLE_METADATA_KEY, + METRIC_ENGINE_NAME, PHYSICAL_TABLE_METADATA_KEY, }; use store_api::mito_engine_options::{ APPEND_MODE_KEY, COMPACTION_OPTION_KEYS, COMPACTION_TWCS_PREFIX, COMPACTION_TYPE, COMPACTION_TYPE_TWCS, MERGE_MODE_KEY, TTL_KEY, }; +use store_api::storage::consts::ReservedColumnId; use store_api::storage::{RegionId, TableId}; use table::TableRef; use table::metadata::{TableInfo, TableInfoRef}; @@ -128,6 +136,8 @@ pub enum AutoCreateTableType { Logical(String), /// A physical table. Physical, + /// The ordinary sparse append table for metric exemplars. + Exemplar, /// A log table which is append-only. Log, /// A table that merges rows by `last_non_null` strategy. @@ -141,6 +151,7 @@ impl AutoCreateTableType { match self { AutoCreateTableType::Logical(_) => "logical", AutoCreateTableType::Physical => "physical", + AutoCreateTableType::Exemplar => "exemplar", AutoCreateTableType::Log => "log", AutoCreateTableType::LastNonNull => "last_non_null", AutoCreateTableType::Trace { .. } => "trace", @@ -163,7 +174,7 @@ impl AutoCreateTableType { /// and normal requests are requests with ttl set to other values. /// /// This is used to split requests for different processing. -#[derive(Clone)] +#[derive(Clone, Default)] pub struct InstantAndNormalInsertRequests { /// Requests with normal ttl. pub normal_requests: RegionInsertRequests, @@ -172,6 +183,20 @@ pub struct InstantAndNormalInsertRequests { pub instant_requests: RegionInsertRequests, } +/// Metric writes bound to the same resolved tables used by their exemplars. +#[derive(Default)] +pub struct PreparedMetricInserts { + requests: InstantAndNormalInsertRequests, + table_infos: HashMap, +} + +impl PreparedMetricInserts { + /// Returns the source tables captured during routing, including auto-created tables. + pub fn source_tables(&self) -> impl Iterator { + self.table_infos.values() + } +} + impl Inserter { /// Checks the assumptions of the logical bulk path without changing tables. /// Unsupported requests retain ordinary insertion, including schema policy, @@ -369,6 +394,119 @@ impl Inserter { .await } + /// Inserts exemplars using ordinary Mito storage and nanosecond auto-create defaults. + pub async fn handle_exemplar_inserts( + &self, + mut requests: RowInsertRequests, + ctx: QueryContextRef, + statement_executor: &StatementExecutor, + sources: &[TableInfoRef], + ) -> Result { + ensure!( + requests + .inserts + .iter() + .all(|r| r.table_name == METRICS_EXEMPLARS_TABLE), + InvalidInsertRequestSnafu { + reason: "invalid exemplar destination" + } + ); + self.bind_exemplar_table_ids(&mut requests, &ctx, sources) + .await?; + preprocess_row_insert_requests(&mut requests.inserts)?; + self.handle_row_inserts_with_create_type( + requests, + ctx, + statement_executor, + AutoCreateTableType::Exemplar, + false, + false, + ) + .await + } + + async fn bind_exemplar_table_ids( + &self, + requests: &mut RowInsertRequests, + ctx: &QueryContextRef, + sources: &[TableInfoRef], + ) -> Result<()> { + validate_column_count_match(requests)?; + let mut tables = HashMap::new(); + for info in sources.iter().filter(|info| { + info.catalog_name == ctx.current_catalog() && info.schema_name == ctx.current_schema() + }) { + if let Some(previous) = tables.insert(info.name.clone(), info.clone()) { + ensure!( + previous.table_id() == info.table_id(), + InvalidInsertRequestSnafu { + reason: format!( + "exemplar source {} changed during write preparation", + info.name + ), + } + ); + } + } + for request in &mut requests.inserts { + let Some(rows) = request.rows.as_mut() else { + continue; + }; + ensure!( + !rows + .schema + .iter() + .any(|c| c.column_name == DATA_SCHEMA_TABLE_ID_COLUMN_NAME), + InvalidInsertRequestSnafu { + reason: "exemplar table IDs must be resolved from their source tables", + } + ); + let metric_index = rows + .schema + .iter() + .position(|c| c.column_name == EXEMPLAR_METRIC_COLUMN) + .context(InvalidInsertRequestSnafu { + reason: "missing exemplar source metric", + })?; + for row in &mut rows.rows { + let Some(ValueData::StringValue(name)) = &row.values[metric_index].value_data + else { + return InvalidInsertRequestSnafu { + reason: "invalid exemplar source metric", + } + .fail(); + }; + if !tables.contains_key(name) { + let table = self.get_table(ctx.current_catalog(), &ctx.current_schema(), name).await? + .with_context(|| InvalidInsertRequestSnafu { + reason: format!("exemplar source {name} must be an existing logical metric-engine table"), + })?; + tables.insert(name.clone(), table.table_info()); + } + let info = &tables[name]; + ensure!( + info.meta.engine == METRIC_ENGINE_NAME + && info + .meta + .options + .extra_options + .contains_key(LOGICAL_TABLE_METADATA_KEY), + InvalidInsertRequestSnafu { + reason: format!( + "exemplar source {name} must be an existing logical metric-engine table" + ), + } + ); + row.values.push(ValueData::U32Value(info.table_id()).into()); + } + rows.schema.push(api::v1::helper::tag_column_schema( + DATA_SCHEMA_TABLE_ID_COLUMN_NAME, + ColumnDataType::Uint32, + )); + } + Ok(()) + } + /// Handles row inserts request and creates a log table on demand. pub async fn handle_log_inserts( &self, @@ -437,6 +575,8 @@ impl Inserter { is_single_value: bool, ) -> Result { let skip_wal = ctx.skip_wal(); + let align_time_index_unit = + matches!(create_type, AutoCreateTableType::Exemplar).then_some(TimeUnit::Nanosecond); let batcher = self .pending_rows_batcher @@ -463,7 +603,7 @@ impl Inserter { statement_executor, accommodate_existing_schema, is_single_value, - None, + align_time_index_unit, ) .await?; @@ -610,11 +750,25 @@ impl Inserter { /// Existing targets retain their backing tables and time-index units. pub async fn handle_metric_row_inserts( &self, - mut requests: RowInsertRequests, + requests: RowInsertRequests, ctx: QueryContextRef, statement_executor: &StatementExecutor, physical_table: String, ) -> Result { + let prepared = self + .prepare_metric_row_inserts(requests, ctx.clone(), statement_executor, physical_table) + .await?; + self.handle_prepared_metric_inserts(prepared, ctx).await + } + + /// Resolves schemas and routes metric rows before dependent exemplar writes start. + pub async fn prepare_metric_row_inserts( + &self, + mut requests: RowInsertRequests, + ctx: QueryContextRef, + statement_executor: &StatementExecutor, + physical_table: String, + ) -> Result { let skip_wal = ctx.skip_wal(); // remove empty requests @@ -672,7 +826,20 @@ impl Inserter { .convert(requests, skip_wal) .await?; - self.do_request(inserts, &table_infos, &ctx).await + Ok(PreparedMetricInserts { + requests: inserts, + table_infos, + }) + } + + /// Dispatches metric rows using the source table definitions captured during preparation. + pub async fn handle_prepared_metric_inserts( + &self, + prepared: PreparedMetricInserts, + ctx: QueryContextRef, + ) -> Result { + self.do_request(prepared.requests, &prepared.table_infos, &ctx) + .await } fn table_batcher( @@ -779,6 +946,7 @@ pub async fn admit_row_insert_batches( for (ctx, requests) in batches.iter() { let catalog = ctx.current_catalog(); let schema = ctx.current_schema(); + if ctx.write_rows_to_admit(catalog, &schema, 1) == 0 { continue; } @@ -1123,6 +1291,7 @@ impl Inserter { is_single_value: bool, align_time_index_unit: Option, ) -> Result { + let is_exemplar = matches!(&auto_create_table_type, AutoCreateTableType::Exemplar); let _timer = crate::metrics::CREATE_ALTER_ON_DEMAND .with_label_values(&[auto_create_table_type.as_str()]) .start_timer(); @@ -1162,16 +1331,19 @@ impl Inserter { .fail(); } }; - // Metric path: an existing destination table keeps its own - // time index unit (it may be bound to another physical - // table than the one selected by this request). - if matches!(auto_create_table_type, AutoCreateTableType::Logical(_)) - && let Some(rows) = req.rows.as_mut() + // Existing metric and exemplar destinations keep their time index unit. + if matches!( + auto_create_table_type, + AutoCreateTableType::Logical(_) | AutoCreateTableType::Exemplar + ) && let Some(rows) = req.rows.as_mut() && let Some(target_unit) = table_time_index_unit(&table) { convert_rows_time_unit(rows, target_unit)?; } let table_info = table.table_info(); + if matches!(auto_create_table_type, AutoCreateTableType::Exemplar) { + validate_exemplar_table(&table_info)?; + } if matches!(auto_create_table_type, AutoCreateTableType::Trace { .. }) { validate_trace_table_model(&table_info, ctx)?; } @@ -1200,6 +1372,9 @@ impl Inserter { match self.get_table(catalog, &schema, &req.table_name).await? { Some(table) => { let table_info = table.table_info(); + if matches!(auto_create_table_type, AutoCreateTableType::Exemplar) { + validate_exemplar_table(&table_info)?; + } if matches!(auto_create_table_type, AutoCreateTableType::Trace { .. }) { validate_trace_table_model(&table_info, ctx)?; } @@ -1306,6 +1481,7 @@ impl Inserter { } } AutoCreateTableType::Physical + | AutoCreateTableType::Exemplar | AutoCreateTableType::Log | AutoCreateTableType::LastNonNull => { // note that auto create table shouldn't be ttl instant table @@ -1314,6 +1490,33 @@ impl Inserter { let table = self .create_physical_table(create_table, None, ctx, statement_executor) .await?; + // A concurrent CREATE can return our proposed schema with the winner's ID. + let table = if is_exemplar { + let name = table.table_info().name.clone(); + self.get_table(catalog, &schema, &name) + .await? + .context(TableNotFoundSnafu { table_name: name })? + } else { + table + }; + if matches!(auto_create_table_type, AutoCreateTableType::Exemplar) { + validate_exemplar_table(&table.table_info())?; + if let Some(req) = requests + .inserts + .iter_mut() + .find(|r| r.table_name == table.table_info().name) + && let Some(alter) = self.get_alter_table_expr_on_demand( + req, &table, ctx, false, false, true, + )? + { + alter_tables.push(alter); + need_refresh_table_infos.insert(( + catalog.to_string(), + schema.clone(), + req.table_name.clone(), + )); + } + } let table_info = table.table_info(); if table_info.is_ttl_instant_table() { instant_table_ids.insert(table_info.table_id()); @@ -1484,9 +1687,32 @@ impl Inserter { ), })?; let table_info = table.table_info(); + if is_exemplar { + validate_exemplar_table(&table_info)?; + } table_infos.insert(table_info.table_id(), table.table_info()); } + if is_exemplar { + for info in table_infos.values() { + if let Some(unit) = info + .meta + .schema + .timestamp_column() + .and_then(|column| column.data_type.as_timestamp()) + .map(|timestamp| timestamp.unit()) + { + for request in &mut requests.inserts { + if request.table_name == info.name + && let Some(rows) = request.rows.as_mut() + { + convert_rows_time_unit(rows, unit)?; + } + } + } + } + } + Ok(CreateAlterTableResult { instant_table_ids, table_infos, @@ -1589,12 +1815,14 @@ impl Inserter { let schema = ctx.current_schema(); let mut table_options = std::collections::HashMap::with_capacity(4); fill_table_options_for_create(&mut table_options, create_type, ctx); - apply_per_table_semantic_options( - &mut table_options, - semantic_index, - ctx.current_schema().as_str(), - &req.table_name, - ); + if !matches!(create_type, AutoCreateTableType::Exemplar) { + apply_per_table_semantic_options( + &mut table_options, + semantic_index, + ctx.current_schema().as_str(), + &req.table_name, + ); + } let engine_name = if let AutoCreateTableType::Logical(_) = create_type { // engine should be metric engine when creating logical tables. @@ -1609,6 +1837,26 @@ impl Inserter { let mut create_table_expr = build_create_table_expr(&table_ref, request_schema, engine_name)?; + if matches!(create_type, AutoCreateTableType::Exemplar) { + // Preserve the stored schema order without shifting every incoming row. + create_table_expr + .column_defs + .sort_by_key(|column| column.name != DATA_SCHEMA_TABLE_ID_COLUMN_NAME); + create_table_expr + .primary_keys + .sort_by_key(|name| name != DATA_SCHEMA_TABLE_ID_COLUMN_NAME); + for column in &mut create_table_expr.column_defs { + if [ + DATA_SCHEMA_TABLE_ID_COLUMN_NAME, + DATA_SCHEMA_TSID_COLUMN_NAME, + ] + .contains(&column.name.as_str()) + { + column.is_nullable = false; + } + } + } + // extension set by the Splunk HEC handler for identity path if ctx.extension(SPLUNK_PK_METADATA_ORDER_KEY).is_some() { reorder_splunk_primary_keys(&mut create_table_expr.primary_keys); @@ -1947,6 +2195,145 @@ fn validate_column_count_match(requests: &RowInsertRequests) -> Result<()> { Ok(()) } +/// Validates the fixed exemplar contract before writing or scanning a shared table. +pub fn validate_exemplar_table(table: &TableInfo) -> Result<()> { + let meta = &table.meta; + let options = &meta.options.extra_options; + ensure!( + meta.engine == default_engine() + && options + .get(APPEND_MODE_KEY) + .is_some_and(|v| v.eq_ignore_ascii_case("true")) + && options + .get(store_api::metric_engine_consts::PRIMARY_KEY_ENCODING) + .is_some_and(|v| v.eq_ignore_ascii_case("sparse")), + InvalidInsertRequestSnafu { + reason: "metrics_exemplars requires Mito, append_mode=true and primary_key_encoding=sparse" + } + ); + let columns = meta.schema.column_schemas(); + let lookup = |name: &str| { + columns + .iter() + .enumerate() + .find(|(_, c)| c.name == name) + .with_context(|| InvalidInsertRequestSnafu { + reason: format!("missing exemplar column {name}"), + }) + }; + let (_, timestamp) = lookup(greptime_timestamp())?; + for (position, name, id, data_type) in [ + ( + 0, + DATA_SCHEMA_TABLE_ID_COLUMN_NAME, + ReservedColumnId::table_id(), + datatypes::prelude::ConcreteDataType::uint32_datatype(), + ), + ( + 1, + DATA_SCHEMA_TSID_COLUMN_NAME, + ReservedColumnId::tsid(), + datatypes::prelude::ConcreteDataType::uint64_datatype(), + ), + ] { + let (index, column) = lookup(name)?; + ensure!( + meta.primary_key_indices.get(position) == Some(&index) + && meta.column_ids.get(index) == Some(&id) + && column.data_type == data_type + && !column.is_nullable() + && column.default_constraint().is_none(), + InvalidInsertRequestSnafu { + reason: format!("invalid exemplar source identity column {name}"), + } + ); + } + ensure!( + timestamp.is_time_index() && timestamp.data_type.is_timestamp(), + InvalidInsertRequestSnafu { + reason: "invalid exemplar time index" + } + ); + let (metric_index, metric) = lookup(EXEMPLAR_METRIC_COLUMN)?; + ensure!( + metric.data_type.is_string() && meta.primary_key_indices.contains(&metric_index), + InvalidInsertRequestSnafu { + reason: "exemplar metric name must be a string tag" + } + ); + let (value_index, value) = lookup(greptime_value())?; + ensure!( + value.data_type == datatypes::prelude::ConcreteDataType::float64_datatype() + && !meta.primary_key_indices.contains(&value_index), + InvalidInsertRequestSnafu { + reason: "exemplar value must be a DOUBLE field" + } + ); + let (labels_index, labels) = lookup(EXEMPLAR_LABELS_COLUMN)?; + ensure!( + labels.data_type.is_json2() && !meta.primary_key_indices.contains(&labels_index), + InvalidInsertRequestSnafu { + reason: "exemplar_labels must be a JSON2 field" + } + ); + if let Some(extension) = labels + .extension_type::() + .map_err(|error| { + InvalidInsertRequestSnafu { + reason: error.to_string(), + } + .build() + })? + { + ensure!( + extension + .metadata() + .json_settings() + .type_hints() + .iter() + .all(|hint| hint.path.len() == 1 && hint.data_type.is_string()), + InvalidInsertRequestSnafu { + reason: "exemplar JSON2 hints must preserve string labels" + } + ); + } + for (index, column) in columns.iter().enumerate() { + if [ + greptime_timestamp(), + greptime_value(), + EXEMPLAR_METRIC_COLUMN, + EXEMPLAR_LABELS_COLUMN, + DATA_SCHEMA_TABLE_ID_COLUMN_NAME, + DATA_SCHEMA_TSID_COLUMN_NAME, + ] + .contains(&column.name.as_str()) + { + continue; + } + ensure!( + meta.primary_key_indices.contains(&index) + && column.data_type.is_string() + && column.is_nullable(), + InvalidInsertRequestSnafu { + reason: "inherited exemplar labels must be nullable string tags" + } + ); + // An absent inherited label must remain absent from the source series identity. + ensure!( + matches!( + column.default_constraint(), + None | Some(ColumnDefaultConstraint::Value( + datatypes::value::Value::Null + )) + ), + InvalidInsertRequestSnafu { + reason: "inherited exemplar labels must not have non-NULL defaults" + } + ); + } + Ok(()) +} + /// Rejects writes from a different built-in trace model before schema mutation. /// Unstamped, explicitly created tables remain subject to normal schema validation. pub fn validate_trace_table_model(table_info: &TableInfo, ctx: &QueryContextRef) -> Result<()> { @@ -2010,6 +2397,21 @@ pub fn fill_table_options_for_create( physical_table.clone(), ); } + AutoCreateTableType::Exemplar => { + table_options.retain(|key, _| { + !is_semantic_option_key(key) + && key != SEMANTIC_PER_TABLE_INDEX_KEY + && key != LOGICAL_TABLE_METADATA_KEY + && key != PHYSICAL_TABLE_METADATA_KEY + && key != MERGE_MODE_KEY + && key != "otlp_metric_compat" + }); + table_options.insert(APPEND_MODE_KEY.to_string(), "true".to_string()); + table_options.insert( + store_api::metric_engine_consts::PRIMARY_KEY_ENCODING.to_string(), + "sparse".to_string(), + ); + } AutoCreateTableType::Physical => { if let Some(append_mode) = ctx.extension(APPEND_MODE_KEY) { table_options.insert(APPEND_MODE_KEY.to_string(), append_mode.to_string()); @@ -3454,6 +3856,39 @@ mod tests { assert!(!table_options.contains_key("compaction.unknown")); } + #[test] + fn test_exemplar_create_options_isolate_metric_context() { + let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME); + for (key, value) in [ + (APPEND_MODE_KEY, "false"), + (MERGE_MODE_KEY, "last_non_null"), + (PHYSICAL_TABLE_METADATA_KEY, "true"), + (LOGICAL_TABLE_METADATA_KEY, "physical_metrics"), + ("otlp_metric_compat", "prometheus"), + (table::requests::SEMANTIC_SIGNAL_TYPE, "metric"), + (SEMANTIC_PER_TABLE_INDEX_KEY, "{}"), + ("ttl", "7d"), + ("compaction.twcs.time_window", "1h"), + ] { + ctx.set_extension(key, value); + } + let mut options = Default::default(); + fill_table_options_for_create(&mut options, &AutoCreateTableType::Exemplar, &Arc::new(ctx)); + assert_eq!( + options, + std::collections::HashMap::from([ + (APPEND_MODE_KEY.into(), "true".into()), + ( + store_api::metric_engine_consts::PRIMARY_KEY_ENCODING.into(), + "sparse".into() + ), + ("ttl".into(), "7d".into()), + (COMPACTION_TYPE.into(), "twcs".into()), + ("compaction.twcs.time_window".into(), "1h".into()), + ]) + ); + } + #[test] fn test_skip_wal_does_not_change_table_options() { check_skip_wal_does_not_change_table_options(false); diff --git a/src/servers/src/batcher/logical_table.rs b/src/servers/src/batcher/logical_table.rs index 26277c43406..707afe2ea0f 100644 --- a/src/servers/src/batcher/logical_table.rs +++ b/src/servers/src/batcher/logical_table.rs @@ -19,7 +19,7 @@ mod pending_worker; mod region_write; mod tables; #[cfg(test)] -mod test_util; +pub(crate) mod test_util; use std::collections::{HashMap, HashSet}; use std::future::{Future, ready}; @@ -173,6 +173,10 @@ impl LogicalTablePendingRowsBatcher { } impl LogicalTablePendingRowsBatcher { + pub(crate) fn is_synchronous(&self) -> bool { + self.pending_rows_batch_sync + } + pub async fn submit(&self, requests: RowInsertRequests, ctx: QueryContextRef) -> Result { self.submit_with(requests, ctx, |_| ready(Ok(()))) .await @@ -210,9 +214,9 @@ impl LogicalTablePendingRowsBatcher { /// table. Time index units need no check here — the bulk encode converts /// each request to its destination's unit. Destinations are resolved /// once per distinct (schema, table). - pub(crate) async fn accepts_bulk_destinations( + pub(crate) async fn accepts_bulk_destinations<'a>( &self, - batches: impl Iterator, + batches: impl Iterator, ) -> bool { // One request can select different physical tables per batch (e.g. // per-series physical-table labels), so the dedupe key includes the @@ -469,9 +473,7 @@ impl Drop for LogicalTablePendingRowsBatcher { #[cfg(test)] mod tests { use std::future::ready; - use std::num::NonZeroUsize; use std::sync::Arc; - use std::time::Duration; use api::v1::value::ValueData; use api::v1::{ @@ -481,99 +483,27 @@ mod tests { use arrow::array::{StringArray, TimestampMillisecondArray}; use arrow::datatypes::{DataType as ArrowDataType, Field, Schema as ArrowSchema}; use arrow::record_batch::RecordBatch; - use catalog::memory::MemoryCatalogManager; use common_error::ext::{BoxedError, ErrorExt}; use common_error::status_code::StatusCode; - use common_meta::test_util::MockDatanodeManager; use common_query::WriteProgress; use common_query::prelude::greptime_timestamp; - use datatypes::prelude::ConcreteDataType; - use datatypes::schema::{ColumnSchema as DtColumnSchema, Schema as DtSchema}; - use operator::test_util::{create_partition_rule_manager, prepare_mocked_backend}; - use operator::write_dispatcher::WriteDispatcher; - use session::context::{QueryContext, QueryContextRef}; - use table::dist_table::DistTable; - use table::test_util::table_info::test_table_info; + use session::context::QueryContext; use crate::batcher::logical_table::batch_convert::{RecordBatchWithTsIdx, TableBatch}; use crate::batcher::logical_table::flow_notifier::extract_timestamps; use crate::batcher::logical_table::pending_worker::WorkerCommand; - use crate::batcher::logical_table::test_util::{mock_aligned_tag_batch, mock_tag_batch}; - use crate::batcher::logical_table::{ - LogicalTablePendingRowsBatcher, PendingRowsSchemaAlterer, batch_key_from_ctx, + use crate::batcher::logical_table::test_util::{ + mock_aligned_tag_batch, mock_tag_batch, submission_batcher, }; - use crate::batcher::test_util::mock_table_flownode_cache; use crate::error; - struct UnexpectedSchemaAlterer; - - #[async_trait::async_trait] - impl PendingRowsSchemaAlterer for UnexpectedSchemaAlterer { - async fn create_tables_if_missing_batch( - &self, - _: &str, - _: &str, - _: &[(&str, &[ColumnSchema])], - _: bool, - _: QueryContextRef, - ) -> error::Result<()> { - unreachable!() - } - - async fn add_missing_prom_tag_columns_batch( - &self, - _: &str, - _: &str, - _: &[(&str, &[String])], - _: QueryContextRef, - ) -> error::Result<()> { - unreachable!() - } - } - - async fn submission_batcher(sync: bool) -> Arc { - let schema = DtSchema::try_new(vec![ - DtColumnSchema::new( - "ts", - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - DtColumnSchema::new("value", ConcreteDataType::float64_datatype(), true), - ]) - .unwrap(); - let table = test_table_info(1, "cpu", "public", "greptime", Arc::new(schema)); - let catalog = MemoryCatalogManager::new_with_table(DistTable::table(Arc::new(table))); - let partitions = create_partition_rule_manager(prepare_mocked_backend().await).await; - let nodes = Arc::new(MockDatanodeManager::new(())); - let dispatcher = Arc::new(WriteDispatcher::new(nodes.clone(), partitions.clone())); - let mut batcher = LogicalTablePendingRowsBatcher::try_new( - partitions, - nodes, - dispatcher, - catalog, - mock_table_flownode_cache(1, vec![]).await, - true, - Arc::new(UnexpectedSchemaAlterer), - Duration::from_secs(1), - 2, - 1, - 1, - 1, - NonZeroUsize::new(1).unwrap(), - ) - .unwrap(); - Arc::get_mut(&mut batcher).unwrap().pending_rows_batch_sync = sync; - batcher - } - fn submission_requests() -> RowInsertRequests { RowInsertRequests { inserts: vec![RowInsertRequest { table_name: "cpu".to_string(), rows: Some(Rows { schema: vec![ColumnSchema { - column_name: "ts".to_string(), + column_name: "greptime_timestamp".to_string(), datatype: ColumnDataType::TimestampMillisecond as i32, semantic_type: SemanticType::Timestamp as i32, ..Default::default() @@ -610,18 +540,13 @@ mod tests { for confirmed in [false, true] { for sync in [false, true] { for affected_rows in [0, 1, 2] { - let batcher = submission_batcher(sync).await; + let (batcher, mut receiver) = submission_batcher(sync).await; let ctx = QueryContext::arc(); let ctx = if confirmed { Arc::new(ctx.with_confirmed_writes()) } else { ctx }; - let (_, receiver) = batcher - .workers - .get_or_create(batch_key_from_ctx(&ctx), 1) - .await; - let mut receiver = receiver.unwrap(); let submit = batcher.submit_with(submission_requests(), ctx, |_| ready(Ok(7))); tokio::pin!(submit); let mut early_result = None; @@ -663,13 +588,8 @@ mod tests { #[tokio::test] async fn test_submit_output_with_preserves_confirmed_partial_progress() { for sync in [false, true] { - let batcher = submission_batcher(sync).await; + let (batcher, mut receiver) = submission_batcher(sync).await; let ctx = Arc::new(QueryContext::arc().with_confirmed_writes()); - let (_, receiver) = batcher - .workers - .get_or_create(batch_key_from_ctx(&ctx), 1) - .await; - let mut receiver = receiver.unwrap(); let (output, prepared) = batcher .submit_output_with(submission_requests(), ctx, |_| ready(Ok(7))) .await diff --git a/src/servers/src/batcher/logical_table/test_util.rs b/src/servers/src/batcher/logical_table/test_util.rs index 491bd6ff2b3..364cf0c8f5a 100644 --- a/src/servers/src/batcher/logical_table/test_util.rs +++ b/src/servers/src/batcher/logical_table/test_util.rs @@ -13,8 +13,11 @@ // limitations under the License. use std::collections::HashMap; +use std::num::NonZeroUsize; use std::sync::Arc; +use std::time::Duration; +use api::v1::ColumnSchema; use api::v1::flow::{DirtyWindowRequests, FlowRequest, FlowResponse}; use api::v1::meta::Peer; use api::v1::region::InsertRequests; @@ -22,13 +25,106 @@ use arrow::array::{StringArray, TimestampMillisecondArray}; use arrow::datatypes::{DataType as ArrowDataType, Field, Schema as ArrowSchema}; use arrow::record_batch::RecordBatch; use async_trait::async_trait; +use catalog::memory::MemoryCatalogManager; use common_meta::error::Result as MetaResult; use common_meta::node_manager::{ DatanodeManager, DatanodeRef, Flownode, FlownodeManager, FlownodeRef, }; +use common_meta::test_util::MockDatanodeManager; +use common_query::prelude::GREPTIME_PHYSICAL_TABLE; +use datatypes::prelude::ConcreteDataType; +use datatypes::schema::{ColumnSchema as DtColumnSchema, Schema as DtSchema}; +use operator::test_util::{create_partition_rule_manager, prepare_mocked_backend}; +use operator::write_dispatcher::WriteDispatcher; +use session::context::{QueryContext, QueryContextRef}; +use store_api::metric_engine_consts::{LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME}; +use table::dist_table::DistTable; +use table::test_util::table_info::test_table_info; use tokio::sync::mpsc; use crate::batcher::logical_table::batch_convert::RecordBatchWithTsIdx; +pub(crate) use crate::batcher::logical_table::pending_worker::WorkerCommand; +use crate::batcher::logical_table::{ + LogicalTablePendingRowsBatcher, PendingRowsSchemaAlterer, batch_key_from_ctx, +}; +use crate::batcher::test_util::mock_table_flownode_cache; +use crate::error; + +struct UnexpectedSchemaAlterer; + +#[async_trait] +impl PendingRowsSchemaAlterer for UnexpectedSchemaAlterer { + async fn create_tables_if_missing_batch( + &self, + _: &str, + _: &str, + _: &[(&str, &[ColumnSchema])], + _: bool, + _: QueryContextRef, + ) -> error::Result<()> { + unreachable!() + } + + async fn add_missing_prom_tag_columns_batch( + &self, + _: &str, + _: &str, + _: &[(&str, &[String])], + _: QueryContextRef, + ) -> error::Result<()> { + unreachable!() + } +} + +pub(crate) async fn submission_batcher( + sync: bool, +) -> ( + Arc, + mpsc::Receiver, +) { + let schema = DtSchema::try_new(vec![ + DtColumnSchema::new( + "greptime_timestamp", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + DtColumnSchema::new("greptime_value", ConcreteDataType::float64_datatype(), true), + ]) + .unwrap(); + let mut table = test_table_info(1, "cpu", "public", "greptime", Arc::new(schema)); + table.meta.engine = METRIC_ENGINE_NAME.into(); + table.meta.options.extra_options.insert( + LOGICAL_TABLE_METADATA_KEY.into(), + GREPTIME_PHYSICAL_TABLE.into(), + ); + let catalog = MemoryCatalogManager::new_with_table(DistTable::table(Arc::new(table))); + let partitions = create_partition_rule_manager(prepare_mocked_backend().await).await; + let nodes = Arc::new(MockDatanodeManager::new(())); + let dispatcher = Arc::new(WriteDispatcher::new(nodes.clone(), partitions.clone())); + let mut batcher = LogicalTablePendingRowsBatcher::try_new( + partitions, + nodes, + dispatcher, + catalog, + mock_table_flownode_cache(1, vec![]).await, + true, + Arc::new(UnexpectedSchemaAlterer), + Duration::from_secs(1), + 2, + 1, + 1, + 1, + NonZeroUsize::new(1).unwrap(), + ) + .unwrap(); + Arc::get_mut(&mut batcher).unwrap().pending_rows_batch_sync = sync; + let (_, receiver) = batcher + .workers + .get_or_create(batch_key_from_ctx(&QueryContext::arc()), 1) + .await; + (batcher, receiver.unwrap()) +} pub(in crate::batcher::logical_table) fn mock_tag_batch( tag_name: &str, diff --git a/src/servers/src/exemplar.rs b/src/servers/src/exemplar.rs new file mode 100644 index 00000000000..40599f14ab0 --- /dev/null +++ b/src/servers/src/exemplar.rs @@ -0,0 +1,310 @@ +// 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. + +//! Normalized metric exemplars shared by Remote Write and OTLP conversion. + +use std::collections::{BTreeMap, BTreeSet}; +use std::sync::Arc; + +use api::v1::value::ValueData; +use api::v1::{ColumnDataType, Row, RowInsertRequests, Value}; +use common_grpc::precision::Precision; +use common_query::prelude::{ + EXEMPLAR_LABELS_COLUMN, EXEMPLAR_METRIC_COLUMN, METRICS_EXEMPLARS_TABLE, greptime_timestamp, + greptime_value, +}; +use datatypes::json; +use metric_engine::row_modifier::TsidGenerator; +use snafu::{OptionExt, ensure}; +use store_api::metric_engine_consts::{ + DATA_SCHEMA_TABLE_ID_COLUMN_NAME, DATA_SCHEMA_TSID_COLUMN_NAME, +}; +use store_api::storage::consts::is_internal_column; + +use crate::error::{InvalidPromRemoteRequestSnafu, Result}; +use crate::row_writer::{self, MultiTableData}; + +/// One observation belonging to a query-visible metric series. +/// Metric names and inherited labels are shared until row construction. +#[derive(Debug, Clone, PartialEq)] +pub struct Exemplar { + pub metric: Arc, + pub series_labels: Arc>, + pub timestamp: i64, + pub value: f64, + pub labels: BTreeMap, +} + +/// Exemplars grouped by the source series' optional database routing control. +pub type ExemplarRequests = BTreeMap, Vec>; + +/// Bounds row capacity after each database's metrics share one exemplar schema. +pub(crate) fn estimated_row_bytes(requests: &ExemplarRequests) -> Option { + requests.values().try_fold(0usize, |total, exemplars| { + let mut tags = BTreeSet::new(); + let mut payload = 0usize; + // Shared series data is copied into each output row. + for exemplar in exemplars { + tags.extend(exemplar.series_labels.keys()); + payload = payload.checked_add(exemplar.metric.len().checked_mul(4)?)?; + for (name, value) in exemplar.series_labels.iter().chain(&exemplar.labels) { + // Include owned strings, label-map entries and JSON2 encoding scratch. + payload = payload.checked_add( + name.len() + .checked_add(value.len())? + .checked_mul(4)? + .checked_add(256)?, + )?; + } + } + // Include the source table ID attached after conversion, and Vec growth. + let columns = tags.len().checked_add(6)?; + let values = exemplars + .len() + .checked_mul(columns)? + .checked_mul(2 * std::mem::size_of::())?; + let rows = exemplars + .len() + .checked_mul(2 * std::mem::size_of::() + std::mem::size_of::())?; + total + .checked_add(values)? + .checked_add(rows)? + .checked_add(columns.checked_mul(512)?)? + .checked_add(payload) + }) +} + +/// Validates labels with the protocol's name policy, preserving duplicates as errors. +pub fn labels( + labels: I, + exemplar_owned: bool, + validate_name: impl Fn(&[u8]) -> bool, +) -> Result> +where + I: IntoIterator, +{ + let mut result = BTreeMap::new(); + let mut length = 0usize; + for (name, value) in labels { + ensure!( + validate_name(name.as_bytes()), + InvalidPromRemoteRequestSnafu { + msg: format!("invalid exemplar label name: {name}"), + } + ); + if exemplar_owned { + ensure!( + name != json::JSON2_REMAINDER_FIELD_NAME, + InvalidPromRemoteRequestSnafu { + msg: "exemplar label name is reserved by JSON2", + } + ); + } else { + ensure!( + ![ + greptime_timestamp(), + greptime_value(), + EXEMPLAR_METRIC_COLUMN, + EXEMPLAR_LABELS_COLUMN, + DATA_SCHEMA_TABLE_ID_COLUMN_NAME, + DATA_SCHEMA_TSID_COLUMN_NAME, + ] + .contains(&name.as_str()) + && !is_internal_column(&name), + InvalidPromRemoteRequestSnafu { + msg: format!("inherited exemplar label collides with reserved column: {name}"), + } + ); + } + length = length + .checked_add(name.chars().count()) + .and_then(|n| n.checked_add(value.chars().count())) + .context(InvalidPromRemoteRequestSnafu { + msg: "exemplar label length overflow", + })?; + ensure!( + !exemplar_owned || length <= 128, + InvalidPromRemoteRequestSnafu { + msg: "exemplar label names and values exceed 128 characters", + } + ); + ensure!( + result.insert(name.clone(), value).is_none(), + InvalidPromRemoteRequestSnafu { + msg: format!("duplicate exemplar label: {name}"), + } + ); + } + Ok(result) +} + +/// Converts Remote Write milliseconds without losing overflow information. +pub fn timestamp_from_millis(timestamp: i64) -> Result { + timestamp + .checked_mul(1_000_000) + .context(InvalidPromRemoteRequestSnafu { + msg: "exemplar timestamp is outside the signed nanosecond range", + }) +} + +/// Builds ordinary-table rows and retains source targets for authorization. +pub fn into_rows(exemplars: Vec) -> Result<(RowInsertRequests, BTreeSet)> { + let mut tables = MultiTableData::default(); + let mut sources = BTreeSet::new(); + for exemplar in exemplars { + ensure!( + !exemplar.metric.is_empty() && exemplar.metric.as_ref() != METRICS_EXEMPLARS_TABLE, + InvalidPromRemoteRequestSnafu { + msg: "invalid or reserved exemplar source metric name" + } + ); + sources.insert(exemplar.metric.to_string()); + let table = tables.get_or_default_table_data( + METRICS_EXEMPLARS_TABLE, + exemplar.series_labels.len() + 6, + 1, + ); + let mut row = table.alloc_one_row(); + let tsid = TsidGenerator::from_sorted_tags( + exemplar + .series_labels + .iter() + .map(|(name, value)| (name.as_str(), value.as_str())), + ); + row_writer::write_by_schema( + table, + std::iter::once(( + api::v1::helper::tag_column_schema( + DATA_SCHEMA_TSID_COLUMN_NAME, + ColumnDataType::Uint64, + ), + Some(ValueData::U64Value(tsid)), + )), + &mut row, + )?; + row_writer::write_tag(table, EXEMPLAR_METRIC_COLUMN, &exemplar.metric, &mut row)?; + row_writer::write_tags( + table, + exemplar + .series_labels + .iter() + .map(|(name, value)| (name.as_str(), value.clone())), + &mut row, + )?; + row_writer::write_fields( + table, + std::iter::once(( + greptime_value().to_string(), + ColumnDataType::Float64, + Some(ValueData::F64Value(exemplar.value)), + )), + &mut row, + )?; + row_writer::write_ts_to_nanos( + table, + greptime_timestamp(), + Some(exemplar.timestamp), + Precision::Nanosecond, + &mut row, + )?; + row_writer::write_by_schema( + table, + std::iter::once(( + row_writer::build_json2_column_schema(EXEMPLAR_LABELS_COLUMN)?, + Some(row_writer::encode_json2(exemplar.labels)?), + )), + &mut row, + )?; + table.add_row(row); + } + Ok((tables.into_row_insert_requests().0, sources)) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::prom_remote_write::validation::validate_label_name; + + #[test] + fn test_exemplar_labels_and_timestamp_validation() { + assert!( + labels( + [(String::from("traceID"), String::from("abc"))], + true, + validate_label_name + ) + .is_ok() + ); + assert!( + labels( + [ + (String::from("job"), String::new()), + (String::from("job"), String::new()) + ], + false, + validate_label_name, + ) + .is_err() + ); + assert!( + labels( + [(EXEMPLAR_METRIC_COLUMN.into(), "x".into())], + false, + validate_label_name + ) + .is_err() + ); + assert!(labels([("id".into(), "x".repeat(127))], true, validate_label_name).is_err()); + for name in ["", "request.method", "请求"] { + assert!(labels([(name.into(), "x".into())], true, validate_label_name).is_err()); + } + for name in [ + DATA_SCHEMA_TABLE_ID_COLUMN_NAME, + DATA_SCHEMA_TSID_COLUMN_NAME, + ] { + assert!(labels([(name.into(), "x".into())], false, validate_label_name).is_err()); + } + assert_eq!(timestamp_from_millis(-1).unwrap(), -1_000_000); + assert!(timestamp_from_millis(i64::MAX).is_err()); + } + + #[test] + fn test_exemplar_tsid_uses_only_source_labels() { + let exemplar = Exemplar { + metric: "cpu".into(), + series_labels: BTreeMap::from([("job".into(), "api".into())]).into(), + timestamp: 1, + value: 2.0, + labels: BTreeMap::from([("trace_id".into(), "a".into())]), + }; + let mut other = exemplar.clone(); + other.metric = "memory".into(); + other.labels.insert("trace_id".into(), "b".into()); + let (requests, _) = into_rows(vec![exemplar, other]).unwrap(); + let rows = requests.inserts[0].rows.as_ref().unwrap(); + let index = rows + .schema + .iter() + .position(|c| c.column_name == DATA_SCHEMA_TSID_COLUMN_NAME) + .unwrap(); + let expected = ValueData::U64Value(TsidGenerator::from_sorted_tags( + [("job", "api")].into_iter(), + )); + assert!( + rows.rows + .iter() + .all(|row| row.values[index].value_data.as_ref() == Some(&expected)) + ); + } +} diff --git a/src/servers/src/http/prom_store.rs b/src/servers/src/http/prom_store.rs index 5a1be58502e..eb7f0898411 100644 --- a/src/servers/src/http/prom_store.rs +++ b/src/servers/src/http/prom_store.rs @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +mod exemplar; + use std::sync::Arc; use api::prom_store::remote::ReadRequest; @@ -33,6 +35,7 @@ use prometheus::HistogramTimer; use prost::Message; use serde::{Deserialize, Serialize}; use session::context::{Channel, QueryContext, QueryContextRef}; +use snafu::IntoError; use snafu::prelude::*; use table::requests::{ METADATA_QUALITY_INFERRED, SEMANTIC_METRIC_METADATA_QUALITY, SEMANTIC_PER_TABLE_INDEX_KEY, @@ -51,7 +54,7 @@ use crate::prom_remote_write::decode::PromSeriesProcessor; use crate::prom_remote_write::v2::decode_remote_write_v2; use crate::prom_remote_write::validation::PromValidationMode; use crate::prom_remote_write::{ - REMOTE_WRITE_V1_VERSION, REMOTE_WRITE_V2_VERSION, decode_remote_write_request, + self, REMOTE_WRITE_V1_VERSION, REMOTE_WRITE_V2_VERSION, decode_remote_write_request, }; use crate::prom_store::{MAX_DECOMPRESSED_REQUEST_SIZE, snappy_decompress_limited}; use crate::query_handler::{PipelineHandlerRef, PromStoreProtocolHandlerRef, PromStoreResponse}; @@ -117,20 +120,29 @@ pub async fn remote_write( body: Bytes, ) -> Result { let is_zstd = content_encoding.contains(VM_ENCODING); + // Read after Bytes has collected every frame, including chunked-body reservations. + let body_reserved_bytes = extensions + .get::() + .map_or(0, BodyMemoryAccounting::reserved_bytes); match remote_write_proto(content_type) { RemoteWriteProto::V1 => { - remote_write_v1(state, params, query_ctx, pipeline_info, is_zstd, body).await + remote_write_v1( + state, + params, + query_ctx, + pipeline_info, + is_zstd, + body, + body_reserved_bytes, + ) + .await } RemoteWriteProto::V2 => { if let Some(response) = unsupported_remote_write_v2_encoding_response(&content_encoding) { return Ok(response); } - // Read after Bytes has collected every frame, including chunked-body reservations. - let body_reserved_bytes = extensions - .get::() - .map_or(0, BodyMemoryAccounting::reserved_bytes); remote_write_v2( state, params, @@ -157,6 +169,7 @@ async fn remote_write_v1( pipeline_info: PipelineInfo, is_zstd: bool, body: Bytes, + body_reserved_bytes: u64, ) -> Result { let PromStoreState { prom_store_handler, @@ -208,6 +221,7 @@ async fn remote_write_v1( // borrowing `col_indexes` dies inside `as_insert_requests`. Nothing below // references the buffer, so it need not span the write — the same goes for // the decode guards charged against the aggregate memory quota. + let exemplars = std::mem::take(&mut decoded.exemplars); let req = if processor.use_pipeline { drop(decoded); processor.exec_pipeline().await? @@ -217,7 +231,52 @@ async fn remote_write_v1( req }; drop(decode_guards); - let batches = into_prom_write_batches(req, query_ctx); + let query_ctx = if exemplars.is_empty() { + query_ctx + } else { + // Reserve for row expansion before copying shared metric names and series labels. + let bytes = crate::exemplar::estimated_row_bytes(&exemplars) + .and_then(|bytes| u64::try_from(bytes).ok()) + .ok_or_else(|| { + error::MemoryLimitExceededSnafu.into_error( + common_memory_manager::Error::MemoryLimitExceeded { + requested_bytes: u64::MAX, + limit_bytes: memory_limiter.limit_bytes(), + }, + ) + })?; + memory_limiter.check_request_fits(bytes, body_reserved_bytes)?; + let reservation = Arc::new( + memory_limiter + .acquire(bytes) + .await + .map_err(prom_remote_write::map_memory_acquire_timeout)?, + ); + let mut ctx = query_ctx.fork(); + ctx.set_write_progress(Arc::new( + session::write_progress::WriteProgress::with_reservation(reservation), + )); + Arc::new(ctx) + }; + let batches = into_prom_write_batches(req, query_ctx.clone()); + if !exemplars.is_empty() { + let outcome = exemplar::write( + prom_store_handler, + pending_rows_batcher, + prom_store_with_metric_engine, + batches, + Vec::new(), + exemplars, + query_ctx, + ) + .await; + record_remote_write_samples(&db, REMOTE_WRITE_V1_VERSION, outcome.written[0]); + record_remote_write_exemplars(&db, REMOTE_WRITE_V1_VERSION, outcome.written[2]); + if let Some(error) = outcome.error { + return Err(error); + } + return Ok((StatusCode::NO_CONTENT, write_cost_header_map(outcome.cost)).into_response()); + } let outcome = match write_prometheus_rows_with_progress( prom_store_handler, @@ -266,7 +325,7 @@ async fn remote_write_v2( // Pipeline processing is not supported for remote write v2 yet. Ignore the // optional pipeline parameter and ingest samples directly. - let _ = pipeline_info; + let has_pipeline = pipeline_info.pipeline_name.is_some(); let (db, mut query_ctx, _timer) = prepare_remote_write_context(¶ms, query_ctx, REMOTE_WRITE_V2_VERSION); @@ -276,15 +335,34 @@ async fn remote_write_v2( Ok(req) => req, Err(error) => return Ok(remote_write_v2_error_response(error, 0, 0, 0)), }; + if has_pipeline && req.exemplar_count > 0 { + return Ok(remote_write_v2_error_response( + error::NotSupportedSnafu { + feat: "pipelines with exemplars", + } + .build(), + 0, + 0, + 0, + )); + } // The v2 per-series metadata upgrades the written tables' semantic options // (metric type/unit, declared quality) at auto-create time. if let Some(index) = req.semantic_index.encode(&query_ctx.current_schema()) { query_ctx.set_extension(SEMANTIC_PER_TABLE_INDEX_KEY, index); } + if req.exemplar_count > 0 { + let counter = crate::metrics::PROM_STORE_REMOTE_WRITE_EXEMPLARS + .with_label_values(&[&db, REMOTE_WRITE_V2_VERSION]); + query_ctx.set_write_progress(Arc::new( + session::write_progress::WriteProgress::with_reservation(req.memory_guard.clone()) + .with_observer(move |rows| counter.inc_by(rows)), + )); + } let query_ctx = Arc::new(query_ctx); let (sample_count, histogram_count) = (req.sample_count, req.histogram_count); let mut sample_batches = into_prom_write_batches(req.samples, query_ctx.clone()); - let mut histogram_batches = into_prom_write_batches(req.histograms, query_ctx); + let mut histogram_batches = into_prom_write_batches(req.histograms, query_ctx.clone()); for (batches, histogram) in [(&mut sample_batches, false), (&mut histogram_batches, true)] { let counter = if histogram { crate::metrics::PROM_STORE_REMOTE_WRITE_HISTOGRAMS @@ -303,6 +381,39 @@ async fn remote_write_v2( *ctx = Arc::new(fork); } } + if req.exemplar_count > 0 { + let outcome = exemplar::write( + prom_store_handler, + pending_rows_batcher, + prom_store_with_metric_engine, + sample_batches, + histogram_batches, + req.exemplars, + query_ctx, + ) + .await; + if let Some(error) = outcome.error { + return Ok(remote_write_v2_error_response( + error, + outcome.written[0], + outcome.written[1], + outcome.written[2], + )); + } + let mut headers = write_cost_header_map(outcome.cost); + append_remote_write_v2_written_headers( + &mut headers, + outcome.written[0], + outcome.written[1], + outcome.written[2], + ); + let status = if outcome.accepted { + StatusCode::ACCEPTED + } else { + StatusCode::NO_CONTENT + }; + return Ok((status, headers).into_response()); + } let outcome = match write_prometheus_v2_rows_with_progress( prom_store_handler, pending_rows_batcher, @@ -462,7 +573,10 @@ async fn write_prometheus_rows_with_progress( // ordinary insert path, which routes per destination; time index // units need no check — the bulk encode converts each request to // its destination's unit. - if batcher.accepts_bulk_destinations(batches.iter()).await { + if batcher + .accepts_bulk_destinations(batches.iter().map(|(ctx, request)| (ctx, request))) + .await + { let mut rows_written = 0; for (temp_ctx, reqs) in batches { let rows = @@ -570,7 +684,9 @@ async fn write_prometheus_v2_rows_with_progress( })?; let batcher = if prom_store_with_metric_engine && let Some(batcher) = pending_rows_batcher.as_ref() - && batcher.accepts_bulk_destinations(batches.iter()).await + && batcher + .accepts_bulk_destinations(batches.iter().map(|(ctx, request)| (ctx, request))) + .await { Some(batcher.as_ref() as &dyn PromWriteBatcher) } else { @@ -699,6 +815,14 @@ fn record_remote_write_samples(db: &str, version: &str, rows: u64) { .inc_by(rows); } +fn record_remote_write_exemplars(db: &str, version: &str, rows: u64) { + if rows != 0 { + crate::metrics::PROM_STORE_REMOTE_WRITE_EXEMPLARS + .with_label_values(&[db, version]) + .inc_by(rows); + } +} + fn remote_write_v2_error_response( error: error::Error, samples: u64, @@ -1206,6 +1330,284 @@ mod tests { } } + struct ExemplarMemoryHandler { + limiter: ServerMemoryLimiter, + progress: Mutex>>, + } + + #[async_trait] + impl PromStoreProtocolHandler for ExemplarMemoryHandler { + async fn pre_write(&self, request: &RowInsertRequests, _: QueryContextRef) -> Result<()> { + let capacity: usize = request + .inserts + .iter() + .flat_map(|request| &request.rows.as_ref().unwrap().rows) + .map(|row| row.values.capacity() * std::mem::size_of::()) + .sum(); + assert!(self.limiter.used_bytes() >= capacity as u64); + Ok(()) + } + + async fn check_exemplar_sources( + &self, + _: &std::collections::BTreeSet, + _: QueryContextRef, + ) -> Result<()> { + Ok(()) + } + + async fn write_exemplars_prepared( + &self, + request: RowInsertRequests, + ctx: QueryContextRef, + ) -> Result { + *self.progress.lock().unwrap() = Some(ctx.write_progress().unwrap().clone()); + Ok(Output::new_with_affected_rows( + prom_write_row_count(&request) as usize, + )) + } + + async fn write_prepared( + &self, + _: RowInsertRequests, + _: QueryContextRef, + _: bool, + ) -> Result { + unreachable!() + } + + async fn write(&self, _: RowInsertRequests, _: QueryContextRef, _: bool) -> Result { + unreachable!() + } + + async fn write_all(&self, _: Vec, _: bool) -> Result>> { + unreachable!() + } + + async fn read(&self, _: ReadRequest, _: QueryContextRef) -> Result { + unreachable!() + } + } + + #[tokio::test] + async fn test_v1_exemplar_shared_series_expansion_is_charged() { + use api::prom_store::remote; + + use crate::http::{HttpOptions, HttpServerBuilder}; + + for large_field in ["__name__", "host"] { + for count in [128, 1] { + let request = remote::WriteRequest { + timeseries: vec![remote::TimeSeries { + labels: ["__name__", "host"] + .into_iter() + .map(|name| remote::Label { + name: name.into(), + value: if name == large_field { + "x".repeat(64 * 1024) + } else { + "metric".into() + }, + }) + .collect(), + exemplars: vec![remote::Exemplar::default(); count], + ..Default::default() + }], + ..Default::default() + }; + let limiter = ServerMemoryLimiter::new(1024 * 1024, OnExhaustedPolicy::Fail); + let handler = Arc::new(ExemplarMemoryHandler { + limiter: limiter.clone(), + progress: Mutex::new(None), + }); + let server = HttpServerBuilder::new(HttpOptions::default()) + .with_memory_limiter(limiter.clone()) + .with_prom_handler( + handler.clone(), + None, + true, + PromValidationMode::Strict, + None, + ) + .build(); + let response = server + .build(server.make_app()) + .unwrap() + .oneshot( + Request::post("/v1/prometheus/write") + .header(header::CONTENT_TYPE, "application/x-protobuf") + .header(header::CONTENT_ENCODING, DEFAULT_ENCODING) + .body(Body::from( + crate::prom_store::snappy_compress(&request.encode_to_vec()) + .unwrap(), + )) + .unwrap(), + ) + .await + .unwrap(); + let progress = handler.progress.lock().unwrap().take(); + if count == 1 { + assert_eq!(response.status(), StatusCode::NO_CONTENT); + assert!(progress.is_some()); + assert!(limiter.used_bytes() >= 64 * 1024); + } else { + assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS); + assert!(progress.is_none()); + } + drop(progress); + assert_eq!(limiter.used_bytes(), 0); + } + } + } + + #[tokio::test(start_paused = true)] + async fn test_exemplar_row_memory_is_charged_and_retained() { + use api::greptime_proto::io::prometheus::write::v2; + use api::prom_store::remote; + + use crate::http::{HttpOptions, HttpServerBuilder}; + + for version in [1, 2] { + for policy in [OnExhaustedPolicy::Fail, OnExhaustedPolicy::default()] { + for (distinct_tags, contended, expire) in [ + (true, false, false), + (false, false, false), + (false, true, false), + (false, true, true), + ] { + let mut v1 = remote::WriteRequest::default(); + let mut v2 = v2::Request { + symbols: vec![String::new()], + timeseries: vec![], + }; + for index in 0..1024 { + let labels = vec![ + remote::Label { + name: "__name__".into(), + value: format!("metric_{index}"), + }, + remote::Label { + name: format!("tag_{}", if distinct_tags { index } else { 0 }), + value: "value".into(), + }, + ]; + let mut labels_refs = Vec::new(); + for label in &labels { + labels_refs + .extend([v2.symbols.len() as u32, v2.symbols.len() as u32 + 1]); + v2.symbols.extend([label.name.clone(), label.value.clone()]); + } + v1.timeseries.push(remote::TimeSeries { + labels, + samples: if expire { + vec![remote::Sample::default()] + } else { + vec![] + }, + exemplars: vec![remote::Exemplar::default()], + ..Default::default() + }); + v2.timeseries.push(v2::TimeSeries { + labels_refs, + samples: if expire { + vec![v2::Sample::default()] + } else { + vec![] + }, + exemplars: vec![v2::Exemplar::default()], + ..Default::default() + }); + } + let (content_type, wire) = if version == 1 { + ("application/x-protobuf", v1.encode_to_vec()) + } else { + ( + "application/x-protobuf;proto=io.prometheus.write.v2.Request", + v2.encode_to_vec(), + ) + }; + let limiter = ServerMemoryLimiter::new(16 * 1024 * 1024, policy); + let mut occupied = contended + .then(|| limiter.try_acquire(16 * 1024 * 1024 - 256 * 1024).unwrap()); + let handler = Arc::new(ExemplarMemoryHandler { + limiter: limiter.clone(), + progress: Mutex::new(None), + }); + let server = HttpServerBuilder::new(HttpOptions::default()) + .with_memory_limiter(limiter.clone()) + .with_prom_handler( + handler.clone(), + None, + true, + PromValidationMode::Strict, + None, + ) + .build(); + let response = server.build(server.make_app()).unwrap().oneshot( + Request::post("/v1/prometheus/write") + .header(header::CONTENT_TYPE, content_type) + .header(header::CONTENT_ENCODING, DEFAULT_ENCODING) + .body(Body::from( + crate::prom_store::snappy_compress(&wire).unwrap(), + )) + .unwrap(), + ); + let started = tokio::time::Instant::now(); + tokio::pin!(response); + if contended && matches!(policy, OnExhaustedPolicy::Wait { .. }) { + assert!(futures::poll!(response.as_mut()).is_pending()); + if !expire { + drop(occupied.take()); + } + } + let response = response.await.unwrap(); + let timed_out = expire && matches!(policy, OnExhaustedPolicy::Wait { .. }); + let rejected = distinct_tags + || timed_out + || (contended && matches!(policy, OnExhaustedPolicy::Fail)); + assert_eq!( + response.status(), + if timed_out { + StatusCode::INTERNAL_SERVER_ERROR + } else if rejected { + StatusCode::TOO_MANY_REQUESTS + } else { + StatusCode::NO_CONTENT + }, + "version={version}, distinct={distinct_tags}, contended={contended}, expire={expire}, policy={policy:?}" + ); + if let OnExhaustedPolicy::Wait { timeout } = policy + && timed_out + { + assert_eq!(started.elapsed(), timeout); + } else { + assert_eq!(started.elapsed(), Duration::ZERO); + } + let progress = handler.progress.lock().unwrap().take(); + if !rejected { + assert!(progress.is_some()); + assert!(limiter.used_bytes() > 0); + if version == 2 { + assert_eq!( + response.headers()[REMOTE_WRITE_V2_EXEMPLARS_WRITTEN_HEADER], + "1024" + ); + } + } else { + assert!(progress.is_none()); + } + drop(progress); + assert_eq!( + limiter.used_bytes(), + occupied.as_ref().map_or(0, |guard| guard.granted_bytes()) + ); + drop(occupied); + assert_eq!(limiter.used_bytes(), 0); + } + } + } + } + fn memory_limited_prom_write_app(limiter: ServerMemoryLimiter) -> Router { let mut state = test_state(); state.memory_limiter = limiter.clone(); @@ -1219,6 +1621,54 @@ mod tests { )) } + #[tokio::test(start_paused = true)] + async fn test_exemplar_v1_reservation_accounts_for_chunked_body() { + use api::prom_store::remote; + + let wire = remote::WriteRequest { + timeseries: vec![remote::TimeSeries { + labels: vec![ + remote::Label { + name: "__name__".into(), + value: "metric".into(), + }, + remote::Label { + name: "job".into(), + value: "api".into(), + }, + ], + exemplars: vec![remote::Exemplar::default()], + ..Default::default() + }], + ..Default::default() + } + .encode_to_vec(); + let body = crate::prom_store::snappy_compress(&wire).unwrap(); + // Each frame retains a 1 KiB body permit. The rows fit alone, but not alongside both. + let frames = body + .chunks(body.len().div_ceil(2)) + .map(Bytes::copy_from_slice) + .collect::>(); + assert_eq!(frames.len(), 2); + let limiter = ServerMemoryLimiter::new(5 * 1024, OnExhaustedPolicy::default()); + let started = tokio::time::Instant::now(); + let response = memory_limited_prom_write_app(limiter.clone()) + .oneshot( + Request::post("/write") + .header(header::CONTENT_TYPE, "application/x-protobuf") + .header(header::CONTENT_ENCODING, DEFAULT_ENCODING) + .body(Body::from_stream(futures::stream::iter( + frames.into_iter().map(Ok::<_, std::io::Error>), + ))) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS); + assert_eq!(started.elapsed(), Duration::ZERO); + assert_eq!(limiter.used_bytes(), 0); + } + #[tokio::test(start_paused = true)] async fn test_remote_write_v2_accounts_for_body_when_rejecting_or_waiting() { let wire = api::greptime_proto::io::prometheus::write::v2::Request { diff --git a/src/servers/src/http/prom_store/exemplar.rs b/src/servers/src/http/prom_store/exemplar.rs new file mode 100644 index 00000000000..d6524e9ce71 --- /dev/null +++ b/src/servers/src/http/prom_store/exemplar.rs @@ -0,0 +1,811 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::HashSet; +use std::sync::Arc; + +use api::v1::RowInsertRequests; +use common_error::ext::BoxedError; +use futures::future::ready; +use operator::insert::PreparedMetricInserts; +use session::context::QueryContextRef; +use snafu::ResultExt; + +use crate::batcher::logical_table::LogicalTablePendingRowsBatcher; +use crate::error::{self, Result}; +use crate::exemplar::{ExemplarRequests, into_rows}; +use crate::http::prom_store::PromWriteBatch; +use crate::query_handler::PromStoreProtocolHandlerRef; + +#[derive(Default)] +pub(super) struct Outcome { + pub cost: usize, + pub accepted: bool, + /// Accepted rows on success; confirmed storage writes on error. + pub written: [u64; 3], + pub error: Option, +} + +enum PreparedWrite { + Metric(PreparedMetricInserts), + Rows(RowInsertRequests), +} + +/// Preflights all destinations together, retaining per-kind write progress. +pub(super) async fn write( + handler: PromStoreProtocolHandlerRef, + batcher: Option>, + metric_engine: bool, + samples: Vec, + histograms: Vec, + exemplars: ExemplarRequests, + ctx: QueryContextRef, +) -> Outcome { + match prepare(&handler, samples, histograms, exemplars, ctx).await { + Ok((writes, targets)) => { + let mut sources = Vec::new(); + let mut prepared = Vec::new(); + let mut source_preparation_failed = false; + for (kind, ctx, request) in writes { + if kind != 2 && metric_engine { + let (pinned, other): (Vec<_>, Vec<_>) = + request.inserts.into_iter().partition(|request| { + targets.contains(&( + ctx.current_catalog().to_string(), + ctx.current_schema(), + request.table_name.clone(), + )) + }); + if !other.is_empty() { + prepared.push(( + kind, + ctx.clone(), + Ok(PreparedWrite::Rows(RowInsertRequests { inserts: other })), + )); + } + if !pinned.is_empty() { + let result = handler + .prepare_metric_write( + RowInsertRequests { inserts: pinned }, + ctx.clone(), + ) + .await; + if let Ok(request) = &result { + sources.extend(request.source_tables().cloned()); + } else { + source_preparation_failed = true; + } + prepared.push((kind, ctx, result.map(PreparedWrite::Metric))); + } + } else { + prepared.push((kind, ctx, Ok(PreparedWrite::Rows(request)))); + } + } + // A failed source preparation must not fall back to a newer table with the same name. + if source_preparation_failed { + prepared.retain(|(kind, _, _)| *kind != 2); + } + let sample_batches = prepared.iter().filter_map(|(kind, ctx, request)| { + if *kind == 0 + && let Ok(PreparedWrite::Rows(request)) = request + { + Some((ctx, request)) + } else { + None + } + }); + let batcher = if metric_engine + && let Some(batcher) = batcher.as_ref() + && batcher.accepts_bulk_destinations(sample_batches).await + { + Some(batcher) + } else { + None + }; + let mut outputs = Vec::with_capacity(prepared.len()); + // Serialize schema preparation across batches that resolve to the same table. + for (kind, ctx, request) in prepared { + let mut deferred = false; + let output = match request { + Err(error) => Err(error), + Ok(PreparedWrite::Metric(request)) => { + handler.write_prepared_metric(request, ctx).await + } + Ok(PreparedWrite::Rows(request)) => { + if kind == 2 && !metric_engine { + error::NotSupportedSnafu { + feat: "exemplars require metric-engine ingestion", + } + .fail() + } else if kind == 2 { + handler + .write_exemplars_with_sources(request, ctx, &sources) + .await + } else if kind == 0 + && let Some(batcher) = batcher + { + deferred = !batcher.is_synchronous(); + batcher + .submit_output_with(request, ctx, |_| ready(Ok(()))) + .await + .map(|(output, ())| output) + } else { + handler.write_prepared(request, ctx, metric_engine).await + } + } + }; + outputs.push((kind, output, deferred)); + } + let mut outcome = Outcome::default(); + let mut pending_samples = Vec::new(); + for (kind, output, deferred) in outputs { + match output { + Ok(output) => { + outcome.cost += output.meta.cost; + if deferred && !output.meta.write_completions.is_empty() { + pending_samples.push(output); + continue; + } + let progress = output.write_progress().await; + outcome.written[kind] += progress.affected_rows as u64; + if let Err(error) = + progress.into_result().context(error::ExecuteGrpcQuerySnafu) + { + outcome.error.get_or_insert(error); + } + } + Err(error) => { + outcome.error.get_or_insert(error); + } + } + } + if outcome.error.is_none() { + outcome.accepted = !pending_samples.is_empty(); + outcome.written[0] += pending_samples + .iter() + .map(|output| output.extract_rows_and_cost().0 as u64) + .sum::(); + } else { + // Split submissions may share a counter; each completion owns its row count. + for output in pending_samples { + outcome.written[0] += output.write_progress().await.affected_rows as u64; + } + } + outcome + } + Err(error) => Outcome { + error: Some(error), + ..Default::default() + }, + } +} + +async fn prepare( + handler: &PromStoreProtocolHandlerRef, + samples: Vec, + histograms: Vec, + exemplars: ExemplarRequests, + ctx: QueryContextRef, +) -> Result<( + Vec<(usize, QueryContextRef, RowInsertRequests)>, + HashSet<(String, String, String)>, +)> { + let mut kinds = Vec::new(); + let mut prepared = Vec::new(); + let mut targets = HashSet::new(); + for (kind, batches) in [samples, histograms].into_iter().enumerate() { + for (ctx, request) in batches { + handler.pre_write(&request, ctx.clone()).await?; + kinds.push(kind); + // Keep per-batch routing options when admitting database-wide totals. + prepared.push((Arc::new(ctx.with_confirmed_writes()), request)); + } + } + for (schema, exemplars) in exemplars { + let target = ctx.fork(); + if let Some(schema) = schema { + target.set_current_schema(&schema); + } + let (request, sources) = into_rows(exemplars)?; + let target = Arc::new(target); + handler.pre_write(&request, target.clone()).await?; + let target = Arc::new(target.with_confirmed_writes()); + if !sources.is_empty() { + handler + .check_exemplar_sources(&sources, target.clone()) + .await?; + } + targets.extend(sources.into_iter().map(|name| { + ( + target.current_catalog().to_string(), + target.current_schema(), + name, + ) + })); + kinds.push(2); + prepared.push((target, request)); + } + operator::insert::admit_row_insert_batches(&mut prepared) + .await + .map_err(BoxedError::new) + .context(error::ExecuteGrpcQuerySnafu)?; + let writes = kinds + .into_iter() + .zip(prepared) + .map(|(kind, (ctx, request))| (kind, ctx, request)) + .collect(); + Ok((writes, targets)) +} + +#[cfg(test)] +mod tests { + use std::collections::{BTreeMap, BTreeSet}; + use std::sync::Mutex; + use std::sync::atomic::{AtomicBool, Ordering}; + + use api::prom_store::remote::ReadRequest; + use api::v1::{Row, RowInsertRequest, Rows}; + use async_trait::async_trait; + use common_query::{Output, OutputData, OutputMeta, WriteCompletion, WriteProgress}; + use session::context::QueryContext; + use tokio::sync::Notify; + + use super::*; + use crate::exemplar::Exemplar; + use crate::query_handler::{PromStoreProtocolHandler, PromStoreResponse}; + + #[derive(Default)] + struct Handler { + events: Mutex>, + exemplar_done: Arc, + deny_schema: bool, + fail_preparation: bool, + successful_writes: bool, + check_serial: bool, + writing: AtomicBool, + } + + #[async_trait] + impl PromStoreProtocolHandler for Handler { + async fn prepare_metric_write( + &self, + _: RowInsertRequests, + _: QueryContextRef, + ) -> Result { + if self.fail_preparation { + return error::InvalidPromRemoteRequestSnafu { + msg: "source preparation failed", + } + .fail(); + } + Ok(PreparedMetricInserts::default()) + } + + async fn write_prepared_metric( + &self, + _: PreparedMetricInserts, + ctx: QueryContextRef, + ) -> Result { + self.write_prepared(RowInsertRequests::default(), ctx, true) + .await + } + + async fn pre_write(&self, request: &RowInsertRequests, ctx: QueryContextRef) -> Result<()> { + self.events + .lock() + .unwrap() + .push(format!("pre:{}", ctx.current_schema())); + assert!(!request.inserts.is_empty()); + if self.deny_schema && ctx.current_schema() == "other" { + return error::InvalidPromRemoteRequestSnafu { + msg: "denied destination", + } + .fail(); + } + if let Some(schema) = ctx.extension("resolved_schema") { + ctx.set_current_schema(schema); + } + Ok(()) + } + + async fn check_exemplar_sources( + &self, + sources: &BTreeSet, + ctx: QueryContextRef, + ) -> Result<()> { + assert_eq!(sources, &BTreeSet::from(["metric".into()])); + self.events + .lock() + .unwrap() + .push(format!("source:{}", ctx.current_schema())); + Ok(()) + } + + async fn write_prepared( + &self, + _request: RowInsertRequests, + ctx: QueryContextRef, + _: bool, + ) -> Result { + if self.check_serial { + assert!(!self.writing.swap(true, Ordering::SeqCst)); + tokio::task::yield_now().await; + self.writing.store(false, Ordering::SeqCst); + } + assert!(ctx.confirmed_writes()); + assert_eq!( + ctx.write_rows_to_admit(ctx.current_catalog(), &ctx.current_schema(), 3), + 0 + ); + self.events.lock().unwrap().push("metric".into()); + if self.successful_writes { + if let Some(counter) = ctx.write_progress() { + counter.add_written(2); + } + return Ok(Output::new_with_affected_rows(2)); + } + let done = self.exemplar_done.clone(); + let counter = ctx.write_progress().cloned(); + let completion = WriteCompletion::with_progress(async move { + done.notified().await; + if let Some(counter) = counter { + counter.add_written(1); + } + WriteProgress { + affected_rows: 1, + error: Some(Arc::new(BoxedError::new( + error::InternalSnafu { + err_msg: "late metric failure", + } + .build(), + ))), + } + }); + Ok(Output::new( + OutputData::AffectedRows(2), + OutputMeta { + cost: 2, + write_completions: vec![completion], + ..Default::default() + }, + )) + } + + async fn write_exemplars_prepared( + &self, + _: RowInsertRequests, + ctx: QueryContextRef, + ) -> Result { + assert!(ctx.confirmed_writes()); + assert_eq!( + ctx.write_rows_to_admit(ctx.current_catalog(), &ctx.current_schema(), 1), + 0 + ); + self.events.lock().unwrap().push("exemplar".into()); + if let Some(counter) = ctx.write_progress() { + counter.add_written(1); + } + self.exemplar_done.notify_one(); + Ok(Output::new( + OutputData::AffectedRows(1), + OutputMeta::new_with_cost(1), + )) + } + + async fn write(&self, _: RowInsertRequests, _: QueryContextRef, _: bool) -> Result { + unreachable!() + } + async fn write_all( + &self, + _: Vec<(QueryContextRef, RowInsertRequests)>, + _: bool, + ) -> Result>> { + unreachable!() + } + async fn read(&self, _: ReadRequest, _: QueryContextRef) -> Result { + unreachable!() + } + } + + #[tokio::test] + async fn exemplar_requires_successful_metric_engine_source_preparation() { + for metric_engine in [false, true] { + let handler = Arc::new(Handler { + fail_preparation: true, + ..Default::default() + }); + let ctx = QueryContext::arc(); + let samples = if metric_engine { + vec![( + ctx.clone(), + RowInsertRequests { + inserts: vec![RowInsertRequest { + table_name: "metric".into(), + rows: Some(Rows { + schema: vec![], + rows: vec![Row { values: vec![] }], + }), + }], + }, + )] + } else { + vec![] + }; + let exemplars = ExemplarRequests::from([( + None, + vec![Exemplar { + metric: "metric".into(), + series_labels: BTreeMap::new().into(), + labels: BTreeMap::new(), + timestamp: 1, + value: 1.0, + }], + )]); + let outcome = write( + handler.clone(), + None, + metric_engine, + samples, + vec![], + exemplars, + ctx, + ) + .await; + assert!(outcome.error.is_some()); + assert_eq!(outcome.written, [0, 0, 0]); + assert!( + !handler + .events + .lock() + .unwrap() + .iter() + .any(|e| e == "metric" || e == "exemplar") + ); + } + } + + #[tokio::test] + async fn exemplar_preserves_batch_routing_after_preflight() { + use crate::http::prom_store::PHYSICAL_TABLE_PARAM; + + let handler: PromStoreProtocolHandlerRef = Arc::new(Handler::default()); + let mut ctx = QueryContext::with("greptime", "public"); + ctx.set_extension("resolved_schema", "tenant"); + let ctx = Arc::new(ctx); + let batch = |physical: Option<&str>| { + let mut target = ctx.fork(); + if let Some(physical) = physical { + target.set_extension(PHYSICAL_TABLE_PARAM, physical); + } + ( + Arc::new(target), + RowInsertRequests { + inserts: vec![RowInsertRequest { + table_name: "metric".into(), + rows: Some(Rows { + schema: vec![], + rows: vec![Row { values: vec![] }], + }), + }], + }, + ) + }; + let exemplars = ExemplarRequests::from([( + None, + vec![Exemplar { + metric: "metric".into(), + series_labels: BTreeMap::new().into(), + labels: BTreeMap::new(), + timestamp: 1, + value: 1.0, + }], + )]); + let (writes, targets) = prepare( + &handler, + vec![ + batch(Some("physical_a")), + batch(Some("physical_b")), + batch(None), + ], + vec![batch(Some("histogram_physical"))], + exemplars, + ctx, + ) + .await + .unwrap(); + assert_eq!( + targets, + HashSet::from([("greptime".into(), "tenant".into(), "metric".into())]) + ); + let expected = [ + (0, Some("physical_a")), + (0, Some("physical_b")), + (0, None), + (1, Some("histogram_physical")), + (2, None), + ]; + assert_eq!(writes.len(), expected.len()); + for ((kind, ctx, request), (expected_kind, physical)) in writes.into_iter().zip(expected) { + assert_eq!(kind, expected_kind); + assert_eq!(ctx.extension(PHYSICAL_TABLE_PARAM), physical); + assert_eq!(ctx.current_schema(), "tenant"); + assert!(ctx.confirmed_writes()); + assert_eq!(ctx.write_rows_to_admit("greptime", "tenant", 1), 0); + assert_eq!(request.inserts.len(), 1); + } + } + + #[tokio::test] + async fn exemplar_submits_metric_batches_sequentially() { + let handler = Arc::new(Handler { + successful_writes: true, + check_serial: true, + ..Default::default() + }); + let ctx = QueryContext::arc(); + let samples = ["metric", "other", "other"] + .into_iter() + .map(|name| { + ( + Arc::new(ctx.fork()), + RowInsertRequests { + inserts: vec![RowInsertRequest { + table_name: name.into(), + rows: Some(Rows { + schema: vec![], + rows: vec![Row { values: vec![] }; 2], + }), + }], + }, + ) + }) + .collect(); + let exemplars = ExemplarRequests::from([( + None, + vec![Exemplar { + metric: "metric".into(), + series_labels: BTreeMap::new().into(), + labels: BTreeMap::new(), + timestamp: 1, + value: 1.0, + }], + )]); + let outcome = write(handler, None, true, samples, vec![], exemplars, ctx).await; + assert!(outcome.error.is_none()); + assert_eq!(outcome.written, [6, 0, 1]); + } + + #[tokio::test] + async fn exemplar_v2_batching_acknowledgement_and_partial_counts() { + use api::greptime_proto::io::prometheus::write::v2; + use common_memory_manager::OnExhaustedPolicy; + use prost::Message; + + use crate::batcher::logical_table::test_util::{WorkerCommand, submission_batcher}; + use crate::http::extractor::PipelineInfo; + use crate::http::prom_store::{PromStoreState, RemoteWriteQuery, remote_write_v2}; + use crate::prom_remote_write::REMOTE_WRITE_V2_VERSION; + use crate::prom_remote_write::validation::PromValidationMode; + use crate::request_memory_limiter::ServerMemoryLimiter; + + for (sync, successful_writes) in + [(false, false), (false, true), (true, false), (true, true)] + { + let (batcher, mut receiver) = submission_batcher(sync).await; + let limiter = ServerMemoryLimiter::new(1024 * 1024, OnExhaustedPolicy::Fail); + let db = "exemplar_confirmed_counts"; + let counters = [ + &*crate::metrics::PROM_STORE_REMOTE_WRITE_SAMPLES, + &*crate::metrics::PROM_STORE_REMOTE_WRITE_EXEMPLARS, + ] + .map(|counter| counter.with_label_values(&[db, REMOTE_WRITE_V2_VERSION])); + let before = counters.each_ref().map(|counter| counter.get()); + let request = v2::Request { + symbols: vec![ + String::new(), + "__name__".into(), + "metric".into(), + "cpu".into(), + ], + timeseries: vec![ + v2::TimeSeries { + labels_refs: vec![1, 2], + samples: vec![ + v2::Sample { + timestamp: 1, + value: 1.0, + ..Default::default() + }; + 2 + ], + exemplars: vec![v2::Exemplar { + labels_refs: vec![], + value: 0.25, + timestamp: 1, + }], + ..Default::default() + }, + v2::TimeSeries { + labels_refs: vec![1, 3], + samples: vec![ + v2::Sample { + timestamp: 1, + value: 1.0, + ..Default::default() + }; + 2 + ], + ..Default::default() + }, + ], + }; + let response = remote_write_v2( + PromStoreState { + prom_store_handler: Arc::new(Handler { + successful_writes, + ..Default::default() + }), + pipeline_handler: None, + prom_store_with_metric_engine: true, + prom_validation_mode: PromValidationMode::Strict, + pending_rows_batcher: Some(batcher), + memory_limiter: limiter.clone(), + }, + RemoteWriteQuery { + db: Some(db.into()), + ..Default::default() + }, + QueryContext::with("greptime", "public"), + PipelineInfo { + pipeline_name: None, + pipeline_version: None, + pipeline_params: Default::default(), + }, + false, + crate::prom_store::snappy_compress(&request.encode_to_vec()) + .unwrap() + .into(), + 0, + ); + tokio::pin!(response); + let mut early_response = None; + let command = tokio::select! { + biased; + result = &mut response => { + early_response = Some(result.unwrap()); + receiver.try_recv().unwrap() + } + command = receiver.recv() => command.unwrap(), + }; + assert_eq!(early_response.is_some(), !sync && successful_writes); + assert!(limiter.used_bytes() > 0); + let WorkerCommand::Submit { + ctx, + response_tx, + _completion, + .. + } = command + else { + panic!("expected sample submission") + }; + let affected_rows = if successful_writes { 2 } else { 1 }; + ctx.write_progress().unwrap().add_written(affected_rows); + let progress = WriteProgress { + affected_rows, + error: (!successful_writes).then(|| { + Arc::new(BoxedError::new( + error::InternalSnafu { + err_msg: "partial queued write", + } + .build(), + )) + }), + }; + assert_eq!( + response_tx.send(progress).is_ok(), + sync || !successful_writes + ); + drop(_completion); + let response = match early_response { + Some(response) => response, + None => response.await.unwrap(), + }; + if successful_writes { + assert_eq!(response.status(), if sync { 204 } else { 202 }); + } else { + assert!(response.status().is_server_error()); + } + let samples = if successful_writes { "4" } else { "2" }; + for (kind, count) in [ + ("samples", samples), + ("histograms", "0"), + ("exemplars", "1"), + ] { + assert_eq!( + response.headers()[format!("x-prometheus-remote-write-{kind}-written")], + count + ); + } + for ((counter, before), count) in counters + .into_iter() + .zip(before) + .zip([if successful_writes { 4 } else { 2 }, 1]) + { + assert_eq!(counter.get() - before, count); + } + drop(ctx); + assert_eq!(limiter.used_bytes(), 0); + } + } + + #[tokio::test] + async fn exemplar_preflights_all_databases_and_retains_delayed_partial_progress() { + for denied in [false, true] { + let handler = Arc::new(Handler { + deny_schema: denied, + ..Default::default() + }); + let mut exemplars = ExemplarRequests::new(); + for schema in if denied { + vec![None, Some("other".into())] + } else { + vec![None] + } { + exemplars.insert( + schema, + vec![Exemplar { + metric: "metric".into(), + series_labels: BTreeMap::new().into(), + labels: BTreeMap::new(), + timestamp: 1, + value: 1.0, + }], + ); + } + let ctx = QueryContext::arc(); + let samples = vec![( + ctx.clone(), + RowInsertRequests { + inserts: vec![RowInsertRequest { + table_name: "metric".into(), + rows: Some(Rows { + schema: vec![], + rows: vec![Row { values: vec![] }; 2], + }), + }], + }, + )]; + let outcome = write(handler.clone(), None, true, samples, vec![], exemplars, ctx).await; + assert!(outcome.error.is_some()); + let events = handler.events.lock().unwrap(); + if denied { + assert_eq!(outcome.written, [0, 0, 0]); + assert!(!events.iter().any(|e| e == "metric" || e == "exemplar")); + } else { + assert_eq!(outcome.written, [1, 0, 1]); + assert_eq!(outcome.cost, 3); + assert_eq!( + *events, + [ + "pre:public", + "pre:public", + "source:public", + "metric", + "exemplar" + ] + ); + } + } + } +} diff --git a/src/servers/src/lib.rs b/src/servers/src/lib.rs index 2fed3d43de0..2b0d0515a5b 100644 --- a/src/servers/src/lib.rs +++ b/src/servers/src/lib.rs @@ -23,6 +23,7 @@ pub mod batcher; pub mod configurator; pub(crate) mod elasticsearch; pub mod error; +pub mod exemplar; pub mod grpc; mod hint_headers; diff --git a/src/servers/src/metrics.rs b/src/servers/src/metrics.rs index 963ab51b1e2..ef5f30ba0b5 100644 --- a/src/servers/src/metrics.rs +++ b/src/servers/src/metrics.rs @@ -123,6 +123,16 @@ lazy_static! { &[METRIC_DB_LABEL, METRIC_VERSION_LABEL] ) .unwrap(); + pub static ref PROM_STORE_REMOTE_WRITE_EXEMPLARS: IntCounterVec = register_int_counter_vec!( + "greptime_servers_prometheus_remote_write_exemplars", + "Acknowledged Prometheus remote write exemplars", + &[METRIC_DB_LABEL, METRIC_VERSION_LABEL] + ).unwrap(); + pub static ref EXEMPLARS_DROPPED_TOTAL: IntCounterVec = register_int_counter_vec!( + "greptime_servers_exemplars_dropped_total", + "Exemplars dropped during input validation", + &["protocol"] + ).unwrap(); /// The native histograms count of Prometheus remote write per protocol version. pub static ref PROM_STORE_REMOTE_WRITE_HISTOGRAMS: IntCounterVec = register_int_counter_vec!( "greptime_servers_prometheus_remote_write_histograms", diff --git a/src/servers/src/otlp/trace/v2.rs b/src/servers/src/otlp/trace/v2.rs index dd43037822f..185beb43183 100644 --- a/src/servers/src/otlp/trace/v2.rs +++ b/src/servers/src/otlp/trace/v2.rs @@ -91,13 +91,13 @@ impl FixedTraceColumnIndexes { })?, span_attributes: writer.ensure_column(row_writer::build_json2_column_schema( SPAN_ATTRIBUTES_COLUMN, - ))?, + )?)?, scope_attributes: writer.ensure_column(row_writer::build_json2_column_schema( SCOPE_ATTRIBUTES_COLUMN, - ))?, + )?)?, resource_attributes: writer.ensure_column(row_writer::build_json2_column_schema( RESOURCE_ATTRIBUTES_COLUMN, - ))?, + )?)?, span_events: writer .ensure_column(row_writer::build_json_column_schema(SPAN_EVENTS_COLUMN))?, span_links: writer diff --git a/src/servers/src/prom_remote_write/decode.rs b/src/servers/src/prom_remote_write/decode.rs index f7324b4be6e..580fc5c57a3 100644 --- a/src/servers/src/prom_remote_write/decode.rs +++ b/src/servers/src/prom_remote_write/decode.rs @@ -14,10 +14,11 @@ //! Decoding of Prometheus remote write protobuf payloads. -use std::collections::BTreeMap; use std::collections::btree_map::Entry; +use std::collections::{BTreeMap, HashSet}; +use std::sync::Arc; -use api::prom_store::remote::Sample; +use api::prom_store::remote::{Exemplar, Sample}; use bytes::Buf; use common_query::prelude::{greptime_timestamp, greptime_value}; use pipeline::{ContextReq, GreptimePipelineParams, PipelineContext, PipelineDefinition}; @@ -33,8 +34,8 @@ use crate::error::InternalSnafu; use crate::http::event::PipelineIngestRequest; use crate::pipeline::run_pipeline; use crate::prom_remote_write::row_builder::{PromCtx, TablesBuilder}; -use crate::prom_remote_write::types::PromLabel; -use crate::prom_remote_write::validation::PromValidationMode; +use crate::prom_remote_write::types::{PromLabel, RawBytes}; +use crate::prom_remote_write::validation::{PromValidationMode, validate_label_name}; #[allow(deprecated)] use crate::prom_store::{ DATABASE_LABEL_ALT_BYTES, DATABASE_LABEL_BYTES, METRIC_NAME_LABEL_BYTES, @@ -51,6 +52,9 @@ pub(crate) struct PromTimeSeries { pub(crate) labels: RepeatedField, pub(crate) samples: RepeatedField, + exemplars: Vec, + label_names: HashSet, + duplicate_labels: bool, } impl Clear for PromTimeSeries { @@ -64,6 +68,9 @@ impl Clear for PromTimeSeries { } self.labels.clear(); self.samples.clear(); + self.exemplars.clear(); + self.label_names.clear(); + self.duplicate_labels = false; } } @@ -98,6 +105,7 @@ impl PromTimeSeries { return Err(DecodeError::new("delimited length exceeded")); } + self.duplicate_labels |= !self.label_names.insert(label.name); #[allow(deprecated)] let is_special_label = match label.name { METRIC_NAME_LABEL_BYTES => { @@ -138,7 +146,12 @@ impl PromTimeSeries { )?; Ok(()) } - 3u32 => prost::encoding::skip_field(wire_type, tag, buf, Default::default()), + 3u32 => { + let mut exemplar = Exemplar::default(); + merge(wire_type, &mut exemplar, buf, Default::default())?; + self.exemplars.push(exemplar); + Ok(()) + } 4u32 => Err(DecodeError::new( "remote write v1 native histogram ingestion is unsupported; use remote write v2", )), @@ -154,6 +167,66 @@ impl PromTimeSeries { let label_num = self.labels.len(); let row_num = self.samples.len(); + if !self.exemplars.is_empty() { + if self.duplicate_labels || self.table_name.is_empty() { + return Err(DecodeError::new( + "exemplar series requires a metric name and unique labels", + )); + } + let inherited = self + .labels + .iter() + .map(|label| { + Ok(( + prom_validation_mode + .decode_label_name(label.name)? + .to_string(), + prom_validation_mode.decode_string(label.value)?, + )) + }) + .collect::, DecodeError>>()?; + let inherited = Arc::new( + crate::exemplar::labels(inherited, false, validate_label_name) + .map_err(|error| DecodeError::new(error.to_string()))?, + ); + let metric: Arc = self.table_name.as_str().into(); + for exemplar in self.exemplars.drain(..) { + let observation: crate::error::Result<_> = (|| { + let labels = crate::exemplar::labels( + exemplar + .labels + .into_iter() + .map(|label| (label.name, label.value)), + true, + validate_label_name, + )?; + Ok(crate::exemplar::Exemplar { + metric: metric.clone(), + series_labels: inherited.clone(), + timestamp: crate::exemplar::timestamp_from_millis(exemplar.timestamp)?, + value: exemplar.value, + labels, + }) + })(); + let observation = match observation { + Ok(observation) => observation, + Err(_) => { + crate::metrics::EXEMPLARS_DROPPED_TOTAL + .with_label_values(&["prometheus_v1"]) + .inc(); + continue; + } + }; + table_builders + .exemplars + .entry(self.schema.clone()) + .or_default() + .push(observation); + } + } + if row_num == 0 { + return Ok(()); + } let prom_ctx = PromCtx { schema: self.schema.take(), physical_table: self.physical_table.take(), @@ -237,6 +310,11 @@ impl<'a> PromWriteRequest<'a> { return Err(DecodeError::new("delimited length exceeded")); } + if processor.use_pipeline && !self.series.exemplars.is_empty() { + return Err(DecodeError::new( + "Remote Write pipelines with exemplars are unsupported", + )); + } if processor.use_pipeline { processor.consume_series_to_pipeline_map( &mut self.series, @@ -469,6 +547,193 @@ mod tests { } } + #[test] + fn test_exemplar_decoder_shares_series_data_before_row_expansion() { + let metric = "m".repeat(64 * 1024); + let label = "x".repeat(64 * 1024); + let count = 128; + let wire = WriteRequest { + timeseries: vec![TimeSeries { + labels: vec![ + Label { + name: "__name__".into(), + value: metric.clone(), + }, + Label { + name: "host".into(), + value: label.clone(), + }, + ], + exemplars: vec![Exemplar::default(); count], + ..Default::default() + }], + ..Default::default() + } + .encode_to_vec(); + assert!(wire.len() < 256 * 1024); + let mut decoder = PromWriteRequest::default(); + decoder + .decode( + wire, + PromValidationMode::Strict, + &mut PromSeriesProcessor::default_processor(), + ) + .unwrap(); + let requests = std::mem::take(&mut decoder.table_data.exemplars); + drop(decoder); + + let observations = &requests[&None]; + assert_eq!(observations.len(), count); + let first = &observations[0]; + for observation in observations { + assert_eq!(&*observation.metric, metric.as_str()); + assert_eq!(observation.series_labels["host"], label); + // Decoding retains one copy of series strings regardless of exemplar count. + assert_eq!(observation.metric.as_ptr(), first.metric.as_ptr()); + assert_eq!( + observation.series_labels["host"].as_ptr(), + first.series_labels["host"].as_ptr() + ); + } + assert!( + crate::exemplar::estimated_row_bytes(&requests).unwrap() + >= count * (metric.len() + label.len()) + ); + } + + #[test] + fn test_exemplar_decoder_preserves_data_across_requests() { + let valid = Exemplar { + timestamp: -1, + value: 0.5, + labels: vec![Label { + name: "traceID".into(), + value: "abc".into(), + }], + }; + let series = TimeSeries { + labels: vec![Label { + name: "__name__".into(), + value: "metric".into(), + }], + exemplars: vec![valid.clone()], + ..Default::default() + }; + let encode = |series| { + WriteRequest { + timeseries: vec![series], + ..Default::default() + } + .encode_to_vec() + }; + let sample = Sample { + timestamp: 2, + value: 1.0, + }; + let mut decoder = PromWriteRequest::default(); + let mut processor = PromSeriesProcessor::default_processor(); + for invalid in [ + Exemplar { + labels: vec![valid.labels[0].clone(); 2], + ..valid.clone() + }, + Exemplar { + labels: vec![Label { + name: "traceID".into(), + value: "x".repeat(128), + }], + ..valid.clone() + }, + Exemplar { + timestamp: i64::MAX, + ..valid.clone() + }, + ] { + for (exemplars, samples, expected) in [ + (vec![valid.clone()], vec![], 1), + (vec![invalid.clone()], vec![], 0), + ( + vec![invalid.clone(), valid.clone(), invalid, valid.clone()], + vec![sample], + 2, + ), + ] { + let sample_count = samples.len(); + decoder + .decode( + encode(TimeSeries { + exemplars, + samples, + ..series.clone() + }), + PromValidationMode::Strict, + &mut processor, + ) + .unwrap(); + let observations = decoder + .table_data + .exemplars + .get(&None) + .map(Vec::as_slice) + .unwrap_or_default(); + assert_eq!(observations.len(), expected); + for observation in observations { + assert_eq!(observation.timestamp, -1_000_000); + assert_eq!(observation.value, 0.5); + assert_eq!(observation.labels["traceID"], "abc"); + } + assert_eq!( + decoder + .as_row_insert_requests() + .all_req() + .map(|r| r.rows.unwrap().rows.len()) + .sum::(), + sample_count + ); + } + for name in ["__name__", "job"] { + let mut invalid = series.clone(); + invalid.labels.extend([ + Label { + name: name.into(), + value: "a".into(), + }, + Label { + name: name.into(), + value: "b".into(), + }, + ]); + assert!( + decoder + .decode(encode(invalid), PromValidationMode::Strict, &mut processor) + .is_err() + ); + } + // A pooled decoder must not leak rejected exemplars into its next request. + decoder + .decode( + encode(TimeSeries { + samples: vec![sample], + exemplars: vec![], + ..series.clone() + }), + PromValidationMode::Strict, + &mut processor, + ) + .unwrap(); + assert!(decoder.table_data.exemplars.is_empty()); + assert_eq!(decoder.as_row_insert_requests().all_req().count(), 1); + } + processor.use_pipeline = true; + assert!( + decoder + .decode(encode(series), PromValidationMode::Strict, &mut processor) + .unwrap_err() + .to_string() + .contains("pipelines with exemplars") + ); + } + #[test] fn test_decode_rejects_remote_write_v1_native_histograms() { for samples in [ diff --git a/src/servers/src/prom_remote_write/mod.rs b/src/servers/src/prom_remote_write/mod.rs index a54b50ad55b..2b67304b620 100644 --- a/src/servers/src/prom_remote_write/mod.rs +++ b/src/servers/src/prom_remote_write/mod.rs @@ -27,10 +27,11 @@ pub(crate) mod v2; pub mod validation; use bytes::Bytes; +use common_error::ext::BoxedError; use common_memory_manager::MemoryGuard; use lazy_static::lazy_static; use object_pool::Pool; -use snafu::ResultExt; +use snafu::{IntoError, ResultExt}; use crate::error; use crate::prom_remote_write::decode::{PromSeriesProcessor, PromWriteRequest}; @@ -137,3 +138,14 @@ pub async fn decode_remote_write_request( let tables = std::mem::take(&mut request.table_data); Ok((tables, guards)) } + +/// Maps memory wait timeouts to 5xx responses so Remote Write senders retry. +pub(crate) fn map_memory_acquire_timeout(error: error::Error) -> error::Error { + match error { + error::Error::MemoryLimitExceeded { + source: source @ common_memory_manager::Error::MemoryAcquireTimeout { .. }, + .. + } => error::OtherSnafu.into_error(BoxedError::new(source)), + error => error, + } +} diff --git a/src/servers/src/prom_remote_write/row_builder.rs b/src/servers/src/prom_remote_write/row_builder.rs index 08f5da61326..a87041b2eb0 100644 --- a/src/servers/src/prom_remote_write/row_builder.rs +++ b/src/servers/src/prom_remote_write/row_builder.rs @@ -37,12 +37,14 @@ pub struct PromCtx { pub struct TablesBuilder<'a> { pub tables: HashMap>>, pub(crate) raw_data: Vec, + pub(crate) exemplars: crate::exemplar::ExemplarRequests, } impl<'a> Clear for TablesBuilder<'a> { fn clear(&mut self) { self.tables.clear(); self.raw_data.clear(); + self.exemplars.clear(); } } diff --git a/src/servers/src/prom_remote_write/v2.rs b/src/servers/src/prom_remote_write/v2.rs index 3c36370fa2a..a8e11fc65e6 100644 --- a/src/servers/src/prom_remote_write/v2.rs +++ b/src/servers/src/prom_remote_write/v2.rs @@ -30,7 +30,6 @@ use api::v1::ColumnSchema; use api::v1::value::ValueData; use api::v1::{ColumnDataType, RowInsertRequest, Rows, SemanticType, Value}; use bytes::{Buf, Bytes}; -use common_error::ext::BoxedError; use common_memory_manager::{MemoryGuard, OnExhaustedPolicy}; #[cfg(test)] use common_query::native_histogram::*; @@ -51,7 +50,7 @@ use table::requests::{ use crate::error::{self, Result}; use crate::prom_remote_write::row_builder::PromCtx; -use crate::prom_remote_write::{REMOTE_WRITE_V2_VERSION, decompress_remote_write_body}; +use crate::prom_remote_write::{self, REMOTE_WRITE_V2_VERSION, decompress_remote_write_body}; #[allow(deprecated)] use crate::prom_store::{ DATABASE_LABEL, DATABASE_LABEL_ALT, METRIC_NAME_LABEL, PHYSICAL_TABLE_LABEL, @@ -125,6 +124,8 @@ pub(crate) struct RemoteWriteV2WriteRequests { pub histograms: ContextReq, pub sample_count: u64, pub histogram_count: u64, + pub exemplars: crate::exemplar::ExemplarRequests, + pub exemplar_count: u64, /// Per-table semantic metadata from the series' inline `Metadata`, folded /// into table options at auto-create time. pub semantic_index: SemanticIndexes, @@ -237,16 +238,7 @@ pub(crate) async fn decode_remote_write_v2( } OnExhaustedPolicy::Fail => decode.await, }; - result.map_err(|error| match error { - error::Error::MemoryLimitExceeded { - source: source @ common_memory_manager::Error::MemoryAcquireTimeout { .. }, - .. - } => { - // Prometheus retries 5xx; 429 may discard a request delayed by other writes. - error::OtherSnafu.into_error(BoxedError::new(source)) - } - error => error, - }) + result.map_err(prom_remote_write::map_memory_acquire_timeout) } async fn decode_remote_write_v2_once( @@ -335,6 +327,8 @@ fn convert_remote_write_v2( let mut label_names = HashSet::new(); let mut sample_count_total = 0; let mut histogram_count_total = 0; + let mut exemplar_count_total = 0; + let mut exemplars = crate::exemplar::ExemplarRequests::new(); let mut labels_refs = Vec::new(); let mut metadata = Metadata::default(); let mut scratch = LeafScratch::default(); @@ -343,8 +337,13 @@ fn convert_remote_write_v2( for series in request.timeseries { let counts = scan_series(series, &mut labels_refs, &mut metadata) .context(error::DecodePromRemoteRequestSnafu)?; + budget.reserve( + counts + .exemplars + .checked_mul(2 * std::mem::size_of::()), + )?; - if counts.samples == 0 && counts.histograms == 0 { + if counts.samples == 0 && counts.histograms == 0 && counts.exemplars == 0 { decode_series_leaves(series, None, Vec::new(), 0, &mut scratch)?; continue; } @@ -359,39 +358,52 @@ fn convert_remote_write_v2( }))?; let (prom_ctx, table_name, tags) = resolve_series_labels(&request.symbols, &labels_refs, &mut label_names)?; - ensure_no_internal_histogram_labels(&tags)?; - budget.reserve( - request - .symbols - .get(metadata.unit_ref as usize) - .map_or(Some(0), |unit| unit.len().checked_mul(4)), - )?; - record_series_metadata( - &mut semantic_index, - &request.symbols, - &metadata, - &prom_ctx, - &table_name, - )?; - let has_other_value_type = if counts.samples > 0 { - counts.histograms > 0 - || histogram_tables + let inherited = if counts.exemplars > 0 { + crate::exemplar::labels( + tags.iter() + .map(|(key, value)| ((*key).to_string(), value.clone())), + false, + |name| !name.is_empty(), + )? + } else { + Default::default() + }; + let exemplar_metric = table_name.clone(); + let exemplar_schema = prom_ctx.schema.clone(); + if counts.samples > 0 || counts.histograms > 0 { + ensure_no_internal_histogram_labels(&tags)?; + budget.reserve( + request + .symbols + .get(metadata.unit_ref as usize) + .map_or(Some(0), |unit| unit.len().checked_mul(4)), + )?; + record_series_metadata( + &mut semantic_index, + &request.symbols, + &metadata, + &prom_ctx, + &table_name, + )?; + let has_other_value_type = if counts.samples > 0 { + counts.histograms > 0 + || histogram_tables + .get(&prom_ctx) + .is_some_and(|tables| tables.contains_key(&table_name)) + } else { + sample_tables .get(&prom_ctx) .is_some_and(|tables| tables.contains_key(&table_name)) - } else { - sample_tables - .get(&prom_ctx) - .is_some_and(|tables| tables.contains_key(&table_name)) - }; - ensure!( - !has_other_value_type, - error::InvalidPromRemoteRequestSnafu { - msg: format!( - "remote write v2 metric `{table_name}` contains both samples and native histograms" - ), - } - ); - + }; + ensure!( + !has_other_value_type, + error::InvalidPromRemoteRequestSnafu { + msg: format!( + "remote write v2 metric `{table_name}` contains both samples and native histograms" + ), + } + ); + } let column_count = tags.len() .checked_add(2) @@ -450,39 +462,114 @@ fn convert_remote_write_v2( })())?; let (writer, row_count) = if counts.samples > 0 { ( - SeriesWriter::Samples(get_or_create_table_data( + Some(SeriesWriter::Samples(get_or_create_table_data( &mut sample_tables, prom_ctx, table_name, column_count, counts.samples, - )), + ))), counts.samples, ) - } else { + } else if counts.histograms > 0 { ( - SeriesWriter::Histograms(get_or_create_table_data( + Some(SeriesWriter::Histograms(get_or_create_table_data( &mut histogram_tables, prom_ctx, table_name, column_count, counts.histograms, - )), + ))), counts.histograms, ) + } else { + (None, 0) }; - decode_series_leaves(series, Some(writer), tags, row_count, &mut scratch)?; + decode_series_leaves(series, writer, tags, row_count, &mut scratch)?; + if !scratch.exemplars.is_empty() { + let inherited = Arc::new(inherited); + let exemplar_metric: Arc = exemplar_metric.into(); + for exemplar in scratch.exemplars.drain(..) { + ensure!( + exemplar.labels_refs.len() % 2 == 0, + error::InvalidPromRemoteRequestSnafu { + msg: "exemplar label references must contain name/value pairs", + } + ); + let exemplar_label_bytes = + exemplar + .labels_refs + .iter() + .try_fold(0usize, |bytes, index| { + bytes.checked_add( + request.symbols.get(*index as usize).map_or(0, |s| s.len()), + ) + }); + // Symbol references can expand into repeated owned labels for every exemplar. + budget.reserve((|| { + label_bytes? + .checked_add(exemplar_label_bytes?)? + .checked_mul(4)? + .checked_add( + labels_refs + .len() + .checked_add(exemplar.labels_refs.len())? + .checked_mul(128)?, + )? + .checked_add(2 * std::mem::size_of::()) + })())?; + let labels = exemplar + .labels_refs + .as_chunks::<2>() + .0 + .iter() + .map(|pair| { + Ok(( + symbol_ref(&request.symbols, pair[0], "exemplar label")?.to_string(), + symbol_ref(&request.symbols, pair[1], "exemplar value")?.to_string(), + )) + }) + .collect::>>()?; + let observation: Result<_> = (|| { + Ok(crate::exemplar::Exemplar { + metric: exemplar_metric.clone(), + series_labels: inherited.clone(), + timestamp: crate::exemplar::timestamp_from_millis(exemplar.timestamp)?, + value: exemplar.value, + labels: crate::exemplar::labels(labels, true, |name| !name.is_empty())?, + }) + })(); + let observation = match observation { + Ok(observation) => observation, + Err(_) => { + crate::metrics::EXEMPLARS_DROPPED_TOTAL + .with_label_values(&["prometheus_v2"]) + .inc(); + continue; + } + }; + exemplars + .entry(exemplar_schema.clone()) + .or_default() + .push(observation); + exemplar_count_total = checked_total(exemplar_count_total, 1, "exemplar")?; + } + } sample_count_total = checked_total(sample_count_total, counts.samples, "sample")?; histogram_count_total = checked_total(histogram_count_total, counts.histograms, "histogram")?; } + budget.reserve(crate::exemplar::estimated_row_bytes(&exemplars))?; + Ok(RemoteWriteV2WriteRequests { samples: into_context_req(sample_tables), histograms: into_context_req(histogram_tables), sample_count: sample_count_total, histogram_count: histogram_count_total, + exemplar_count: exemplar_count_total, + exemplars, semantic_index, memory_guard: Arc::new(budget.guard), }) @@ -555,6 +642,7 @@ fn metric_type_value(wire_type: i32) -> Option<&'static str> { #[derive(Default)] struct SeriesCounts { + exemplars: usize, samples: usize, histograms: usize, } @@ -598,6 +686,10 @@ fn scan_series( } TIME_SERIES_EXEMPLARS_TAG => { take_length_delimited(wire_type, &mut buf)?; + counts.exemplars = counts + .exemplars + .checked_add(1) + .ok_or_else(|| DecodeError::new("exemplar count overflow"))?; } TIME_SERIES_METADATA_TAG => { message::merge(wire_type, metadata, &mut buf, DecodeContext::default()).map_err( @@ -619,6 +711,7 @@ struct LeafScratch { sample: Sample, histogram: Histogram, exemplar: Exemplar, + exemplars: Vec, } enum SeriesWriter<'a> { @@ -635,6 +728,7 @@ fn decode_series_leaves( ) -> Result<()> { let mut sample_row_template = None; let mut histogram_row_template = None; + scratch.exemplars.clear(); while buf.has_remaining() { let (tag, wire_type) = decode_key(&mut buf).context(error::DecodePromRemoteRequestSnafu)?; @@ -781,6 +875,9 @@ fn decode_series_leaves( error }) .context(error::DecodePromRemoteRequestSnafu)?; + scratch + .exemplars + .push(std::mem::take(&mut scratch.exemplar)); } TIME_SERIES_METADATA_TAG => { take_length_delimited(wire_type, &mut buf) @@ -1357,7 +1454,7 @@ mod tests { } #[test] - fn test_fused_decoder_ignores_exemplar_symbol_refs() { + fn test_fused_decoder_rejects_invalid_exemplar_symbol_refs() { let request = Request { symbols: vec![ String::new(), @@ -1375,7 +1472,66 @@ mod tests { }], }; - assert_eq!(decode_test_request(request).unwrap().sample_count, 1); + assert!(decode_test_request(request).is_err()); + } + + #[test] + fn test_invalid_exemplars_preserve_samples_and_counts() { + let valid = Exemplar { + labels_refs: vec![3, 4], + timestamp: 1, + value: 0.5, + }; + for invalid in [ + Exemplar { + labels_refs: vec![3, 4, 3, 4], + ..valid.clone() + }, + Exemplar { + labels_refs: vec![3, 5], + ..valid.clone() + }, + Exemplar { + timestamp: i64::MAX, + ..valid.clone() + }, + Exemplar { + labels_refs: vec![6, 4], + ..valid.clone() + }, + ] { + for retain_valid in [false, true] { + let mut observations = vec![invalid.clone()]; + if retain_valid { + observations.extend([valid.clone(), invalid.clone(), valid.clone()]); + } + let request = Request { + symbols: vec![ + "".into(), + METRIC_NAME_LABEL.into(), + "metric".into(), + "trace_id".into(), + "abc".into(), + "x".repeat(128), + datatypes::json::JSON2_REMAINDER_FIELD_NAME.into(), + ], + timeseries: vec![TimeSeries { + labels_refs: vec![1, 2], + samples: vec![Sample::default()], + exemplars: observations, + ..Default::default() + }], + }; + let converted = decode_test_request(request).unwrap(); + assert_eq!(converted.sample_count, 1); + let expected = if retain_valid { 2 } else { 0 }; + assert_eq!(converted.exemplar_count, expected); + assert_eq!( + converted.exemplars.values().map(Vec::len).sum::(), + expected as usize + ); + } + } } #[test] @@ -1634,6 +1790,38 @@ mod tests { assert!(budget.reserve(usize::MAX.checked_mul(2)).is_err()); } + #[tokio::test] + async fn test_exemplar_conversion_reserves_repeated_labels() { + let limiter = ServerMemoryLimiter::new(128 * 1024, OnExhaustedPolicy::Fail); + let label = "x".repeat(4096); + let mut request = + request_with_sample(vec![(METRIC_NAME_LABEL, "metric"), ("host", &label)]); + request.timeseries[0].samples.clear(); + request.timeseries[0].exemplars = vec![Exemplar::default(); 128]; + let body = + Bytes::from(crate::prom_store::snappy_compress(&request.encode_to_vec()).unwrap()); + assert!(matches!( + decode_remote_write_v2(false, body, &limiter, 0).await, + Err(error::Error::MemoryLimitExceeded { .. }) + )); + assert_eq!(limiter.used_bytes(), 0); + + request.timeseries[0].exemplars.truncate(1); + let body = + Bytes::from(crate::prom_store::snappy_compress(&request.encode_to_vec()).unwrap()); + let decoded = decode_remote_write_v2(false, body, &limiter, 0) + .await + .unwrap(); + assert_eq!(decoded.exemplar_count, 1); + assert!(limiter.used_bytes() >= label.len() as u64); + let progress = + session::write_progress::WriteProgress::with_reservation(decoded.memory_guard.clone()); + drop(decoded); + assert!(limiter.used_bytes() >= label.len() as u64); + drop(progress); + assert_eq!(limiter.used_bytes(), 0); + } + #[tokio::test] async fn test_large_symbols_do_not_reserve_body_multiples_or_retain_scratch() { let limiter = ServerMemoryLimiter::new(64 * 1024 * 1024, OnExhaustedPolicy::Fail); @@ -1757,7 +1945,17 @@ mod tests { ("区域", "西"), ("has-dash", "value"), ]; - let sample = request_with_sample(labels.clone()); + let mut sample = request_with_sample(labels.clone()); + let name = push_test_symbol(&mut sample.symbols, "trace.id"); + let value = push_test_symbol(&mut sample.symbols, "abc"); + sample.timeseries[0].exemplars = vec![Exemplar { + labels_refs: vec![name, value], + ..Default::default() + }]; + let decoded = decode_test_request(sample.clone()).unwrap(); + assert_eq!(decoded.exemplar_count, 1); + assert_eq!(decoded.exemplars[&None][0].series_labels["区域"], "西"); + assert_eq!(decoded.exemplars[&None][0].labels["trace.id"], "abc"); let (requests, _, _, _) = test_util::write_requests(sample).unwrap(); let sample_rows = requests[0].rows.as_ref().unwrap(); assert_eq!(sample_rows.schema[3].column_name, "区域"); @@ -2225,7 +2423,7 @@ mod tests { } #[test] - fn test_into_context_req_converts_histograms_and_ignores_exemplars() { + fn test_into_context_req_preserves_histogram_exemplars() { let request = Request { symbols: vec![ "".to_string(), @@ -2256,6 +2454,11 @@ mod tests { assert_eq!(ctx_req.sample_count, 1); assert_eq!(ctx_req.histogram_count, 1); + assert_eq!(ctx_req.exemplar_count, 1); + assert_eq!( + ctx_req.exemplars[&None][0].metric.as_ref(), + "histogram_metric" + ); assert_eq!(ctx_req.samples.all_req().count(), 1); assert_eq!(ctx_req.histograms.all_req().count(), 1); } diff --git a/src/servers/src/query_handler.rs b/src/servers/src/query_handler.rs index 57f55855ac1..c0d1d7c9b0b 100644 --- a/src/servers/src/query_handler.rs +++ b/src/servers/src/query_handler.rs @@ -39,6 +39,7 @@ use headers::HeaderValue; use log_query::LogQuery; use opentelemetry_proto::tonic::collector::logs::v1::ExportLogsServiceRequest; use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest; +use operator::insert::PreparedMetricInserts; use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest; use pipeline::{GreptimePipelineParams, Pipeline, PipelineInfo, PipelineVersion, PipelineWay}; use serde_json::Value; @@ -109,6 +110,64 @@ pub struct PromStoreResponse { #[async_trait] pub trait PromStoreProtocolHandler { + /// Resolves metric destinations before source identities are attached to exemplars. + /// Requires pre-write checks and admission to have succeeded for the entire request. + async fn prepare_metric_write( + &self, + _request: RowInsertRequests, + _ctx: QueryContextRef, + ) -> Result { + crate::error::NotSupportedSnafu { + feat: "metric write preparation", + } + .fail() + } + + /// Executes the same routed metric writes used to resolve exemplar source identities. + async fn write_prepared_metric( + &self, + _request: PreparedMetricInserts, + _ctx: QueryContextRef, + ) -> Result { + crate::error::NotSupportedSnafu { + feat: "prepared metric ingestion", + } + .fail() + } + + /// Writes exemplars using table identities captured while preparing the source metrics. + async fn write_exemplars_with_sources( + &self, + request: RowInsertRequests, + ctx: QueryContextRef, + _sources: &[TableInfoRef], + ) -> Result { + self.write_exemplars_prepared(request, ctx).await + } + + /// Authorizes source series, including inputs with no metric samples. + async fn check_exemplar_sources( + &self, + _sources: &std::collections::BTreeSet, + _ctx: QueryContextRef, + ) -> Result<()> { + crate::error::NotSupportedSnafu { + feat: "exemplar source authorization", + } + .fail() + } + + async fn write_exemplars_prepared( + &self, + _request: RowInsertRequests, + _ctx: QueryContextRef, + ) -> Result { + crate::error::NotSupportedSnafu { + feat: "exemplar ingestion", + } + .fail() + } + /// Runs pre-write checks/hooks for prometheus remote write requests. async fn pre_write(&self, request: &RowInsertRequests, ctx: QueryContextRef) -> Result<()>; diff --git a/src/servers/src/row_writer.rs b/src/servers/src/row_writer.rs index 74ef6cedc0d..95a6f34260a 100644 --- a/src/servers/src/row_writer.rs +++ b/src/servers/src/row_writer.rs @@ -16,7 +16,7 @@ use std::collections::HashMap; use std::sync::Arc; use ahash::{HashMap as AHashMap, HashMapExt}; -use api::helper::encode_json_value; +use api::helper::{ColumnDataTypeWrapper, encode_json_value}; use api::v1::column_data_type_extension::TypeExt; use api::v1::helper::time_index_column_schema; use api::v1::value::ValueData; @@ -33,6 +33,8 @@ use common_time::timestamp::TimeUnit; use common_time::timestamp::TimeUnit::Nanosecond; use datatypes::extension::json::{Json2ExtensionType, JsonMetadata}; use datatypes::json::JsonSettings; +use datatypes::prelude::ConcreteDataType; +use datatypes::types::json_type::JsonNativeType; use datatypes::value::Value as DataValue; use serde::Serialize; use snafu::{OptionExt, ResultExt, ensure}; @@ -356,7 +358,7 @@ pub fn write_json( ) } -pub(crate) fn build_json2_column_schema(name: impl ToString) -> ColumnSchema { +pub(crate) fn build_json2_column_schema(name: impl ToString) -> Result { let extension = Json2ExtensionType::new(Arc::new(JsonMetadata::new(JsonSettings::new_v2()))); let mut options = ColumnOptions::default(); options.options.insert( @@ -369,13 +371,23 @@ pub(crate) fn build_json2_column_schema(name: impl ToString) -> ColumnSchema { .insert(EXTENSION_TYPE_METADATA_KEY.to_string(), metadata); } - ColumnSchema { + Ok(ColumnSchema { column_name: name.to_string(), datatype: ColumnDataType::Json as i32, semantic_type: SemanticType::Field as i32, + datatype_extension: ColumnDataTypeWrapper::try_from(ConcreteDataType::json2( + JsonNativeType::Null, + )) + .map_err(|error| { + InternalSnafu { + err_msg: error.to_string(), + } + .build() + })? + .into_parts() + .1, options: Some(options), - ..Default::default() - } + }) } /// Encodes a JSON2 field without constructing its column schema. diff --git a/tests-integration/src/tests/instance_test.rs b/tests-integration/src/tests/instance_test.rs index c1f886c6ed7..86bd1409c1a 100644 --- a/tests-integration/src/tests/instance_test.rs +++ b/tests-integration/src/tests/instance_test.rs @@ -50,7 +50,7 @@ async fn test_sparse_primary_key_schema_validation(instance: Arc { + add(&mut request, &[("__name__", "unpinned")]); + request.timeseries[1].samples.clear(); + request.timeseries[1].histograms.push(v2::Histogram { + count: Some(v2::histogram::Count::CountInt(0)), + timestamp: 1000, + ..Default::default() + }); + } + "schema_alias" => { + add( + &mut request, + &[("__name__", "unpinned"), ("left_tag", "left")], + ); + add( + &mut request, + &[ + ("__name__", "unpinned"), + ("__schema__", "public"), + ("right_tag", "right"), + ], + ); + } + "physical_conflict" => { + for (physical, job) in [("physical_a", "one"), ("physical_b", "two")] { + execute_sql(frontend, &format!( + "CREATE TABLE {physical} (greptime_timestamp TIMESTAMP TIME INDEX, greptime_value DOUBLE, job STRING PRIMARY KEY) ENGINE=metric WITH ('physical_metric_table'='')" + )).await; + add( + &mut request, + &[ + ("__name__", "unpinned"), + ("__physical_table__", physical), + ("job", job), + ], + ); + } + } + _ => unreachable!(), + } + let response = app + .oneshot( + Request::post("/v1/prometheus/write") + .header("Content-Encoding", "snappy") + .header( + "Content-Type", + "application/x-protobuf;proto=io.prometheus.write.v2.Request", + ) + .body(Body::from( + snappy_compress(&request.encode_to_vec()).unwrap(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!( + response.status(), + 204, + "{:?}", + to_bytes(response.into_body(), usize::MAX).await.unwrap() + ); + assert_eq!( + response.headers()["x-prometheus-remote-write-samples-written"], + if kind == "histogram" { "1" } else { "3" } + ); + assert_eq!( + response.headers()["x-prometheus-remote-write-histograms-written"], + if kind == "histogram" { "1" } else { "0" } + ); + assert_eq!( + response.headers()["x-prometheus-remote-write-exemplars-written"], + "1" + ); + let filter = match kind { + "histogram" => "greptime_native_histogram IS NOT NULL", + "schema_alias" => "left_tag = 'left' OR right_tag = 'right'", + "physical_conflict" => "job IN ('one', 'two')", + _ => unreachable!(), + }; + let output = execute_sql( + frontend, + &format!("SELECT count(*) FROM unpinned WHERE {filter}"), + ) + .await; + check_output_stream( + output.data, + &format!( + "+----------+\n| count(*) |\n+----------+\n| {} |\n+----------+", + if kind == "histogram" { 1 } else { 2 } + ), + ) + .await; +} + +#[rstest] +#[case::lowercase("sparse", true, true, 9)] +#[case::uppercase("SPARSE", true, true, 9)] +#[case::mixed_case("SpArSe", true, true, 9)] +#[case::dense("dense", true, false, 9)] +#[case::not_append_only("sparse", false, false, 9)] +#[case::seconds("sparse", true, true, 0)] +#[case::milliseconds("sparse", true, true, 3)] +#[case::microseconds("sparse", true, true, 6)] +#[tokio::test(flavor = "multi_thread")] +async fn test_exemplar_precreated_table_options( + #[case] encoding: &str, + #[case] append_mode: bool, + #[case] accepted: bool, + #[case] precision: u32, +) { + use std::collections::BTreeMap; + + use servers::exemplar::{Exemplar, into_rows}; + use servers::query_handler::PromStoreProtocolHandler; + + use crate::standalone::GreptimeDbStandaloneBuilder; + + let instance = GreptimeDbStandaloneBuilder::new("exemplar_precreated_table_options") + .build() + .await; + let frontend = instance.fe_instance(); + execute_sql(frontend, "CREATE TABLE source_physical (ts TIMESTAMP TIME INDEX, val DOUBLE) ENGINE=metric WITH ('physical_metric_table'='')").await; + execute_sql(frontend, "CREATE TABLE metric_source (job STRING PRIMARY KEY, ts TIMESTAMP TIME INDEX, val DOUBLE) ENGINE=metric WITH ('on_physical_table'='source_physical')").await; + execute_sql(frontend, &format!( + "CREATE TABLE metrics_exemplars (__table_id INT UNSIGNED NOT NULL, __tsid BIGINT UNSIGNED NOT NULL, greptime_metric STRING, greptime_timestamp TIMESTAMP({precision}) TIME INDEX, greptime_value DOUBLE, exemplar_labels JSON2, PRIMARY KEY (__table_id, __tsid, greptime_metric)) ENGINE=mito WITH ('append_mode'='{append_mode}', 'primary_key_encoding'='{encoding}', 'compaction.type'='twcs')" + )).await; + let (request, _) = into_rows(vec![Exemplar { + metric: "metric_source".into(), + series_labels: BTreeMap::from([("job".into(), "api".into())]).into(), + timestamp: 1_234_567_890, + value: 2.0, + labels: BTreeMap::new(), + }]) + .unwrap(); + let result = frontend + .write_exemplars_prepared(request, QueryContext::arc()) + .await; + if accepted { + assert_eq!(result.unwrap().extract_rows_and_cost().0, 1); + } else { + let error = result.unwrap_err(); + assert_eq!(error.status_code(), StatusCode::InvalidArguments); + assert!( + error.output_msg().contains( + "metrics_exemplars requires Mito, append_mode=true and primary_key_encoding=sparse" + ), + "{error:?}" + ); + } + let table = frontend + .catalog_manager() + .table("greptime", "public", "metrics_exemplars", None) + .await + .unwrap() + .unwrap(); + // Invalid destinations must fail before auto-ALTER adds the inherited label. + assert_eq!( + table.schema().column_schema_by_name("job").is_some(), + accepted + ); + let sql = if accepted { + let timestamp = 1_234_567_890 / 10_i64.pow(9 - precision); + format!( + "SELECT count(*) FROM metrics_exemplars WHERE job = 'api' AND greptime_value = 2.0 AND CAST(greptime_timestamp AS BIGINT) = {timestamp}" + ) + } else { + "SELECT count(*) FROM metrics_exemplars".into() + }; + let output = execute_sql(frontend, &sql).await; + check_output_stream( + output.data, + &format!( + "+----------+\n| count(*) |\n+----------+\n| {} |\n+----------+", + usize::from(accepted) + ), + ) + .await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn test_exemplar_source_identity_policy() { + use std::collections::BTreeMap; + + use servers::exemplar::{Exemplar, into_rows}; + use servers::query_handler::PromStoreProtocolHandler; + + use crate::standalone::GreptimeDbStandaloneBuilder; + + let instance = GreptimeDbStandaloneBuilder::new("exemplar_source_identity") + .build() + .await; + let frontend = instance.fe_instance(); + let request = |name: &str| { + into_rows(vec![Exemplar { + metric: name.into(), + series_labels: BTreeMap::from([("job".into(), "api".into())]).into(), + timestamp: 1, + value: 2.0, + labels: BTreeMap::new(), + }]) + .unwrap() + .0 + }; + execute_sql( + frontend, + "CREATE TABLE ordinary_source (ts TIMESTAMP TIME INDEX, val DOUBLE)", + ) + .await; + for name in ["missing_source", "ordinary_source"] { + let error = frontend + .write_exemplars_prepared(request(name), QueryContext::arc()) + .await + .unwrap_err(); + assert!( + error + .output_msg() + .contains("existing logical metric-engine table"), + "{error:?}" + ); + } + assert!( + frontend + .catalog_manager() + .table("greptime", "public", "metrics_exemplars", None) + .await + .unwrap() + .is_none() + ); + execute_sql(frontend, "CREATE TABLE source_physical (ts TIMESTAMP TIME INDEX, val DOUBLE) ENGINE=metric WITH ('physical_metric_table'='')").await; + let create = "CREATE TABLE metric_source (job STRING PRIMARY KEY, ts TIMESTAMP TIME INDEX, val DOUBLE) ENGINE=metric WITH ('on_physical_table'='source_physical')"; + execute_sql(frontend, create).await; + let source = frontend + .catalog_manager() + .table("greptime", "public", "metric_source", None) + .await + .unwrap() + .unwrap() + .table_info(); + // A metric write's captured table ID must not be replaced by a later name lookup. + execute_sql(frontend, "DROP TABLE metric_source").await; + execute_sql(frontend, create).await; + let replacement = frontend + .catalog_manager() + .table("greptime", "public", "metric_source", None) + .await + .unwrap() + .unwrap() + .table_info(); + assert_ne!(source.table_id(), replacement.table_id()); + frontend + .write_exemplars_with_sources( + request("metric_source"), + QueryContext::arc(), + std::slice::from_ref(&source), + ) + .await + .unwrap(); + frontend + .write_exemplars_prepared(request("metric_source"), QueryContext::arc()) + .await + .unwrap(); + let table = frontend + .catalog_manager() + .table("greptime", "public", "metrics_exemplars", None) + .await + .unwrap() + .unwrap(); + let info = table.table_info(); + assert_eq!( + info.meta + .schema + .column_schemas() + .iter() + .map(|column| column.name.as_str()) + .collect::>(), + [ + "__table_id", + "__tsid", + "greptime_metric", + "job", + "greptime_value", + "greptime_timestamp", + "exemplar_labels", + ] + ); + assert_eq!(info.meta.primary_key_indices, [0, 1, 2, 3]); + let output = execute_sql( + frontend, + "SELECT __table_id FROM metrics_exemplars ORDER BY __table_id", + ) + .await; + let OutputData::Stream(stream) = output.data else { + panic!("expected stream") + }; + let batches = util::collect_batches(stream).await.unwrap(); + let ids = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_primitive::() + .values() + .to_vec() + }) + .collect::>(); + assert_eq!(ids, vec![source.table_id(), replacement.table_id()]); +} diff --git a/tests/compatibility/cases/exemplar_sparse/case.toml b/tests/compatibility/cases/exemplar_sparse/case.toml new file mode 100644 index 00000000000..d245bc00b20 --- /dev/null +++ b/tests/compatibility/cases/exemplar_sparse/case.toml @@ -0,0 +1,8 @@ +name = "exemplar_sparse" +reason = "Preserve reserved sparse source identity keys, append duplicates and nanosecond timestamps across restart." +introduced_by = "persistent_metric_exemplars" +topologies = ["standalone", "distributed"] +from_range = [">=v1.4.0-alpha.0"] +to_range = [">=v1.4.0-alpha.0"] +features = ["sparse", "table", "read"] +owner = "mito2" diff --git a/tests/compatibility/cases/exemplar_sparse/setup.sql b/tests/compatibility/cases/exemplar_sparse/setup.sql new file mode 100644 index 00000000000..dc1a09b4b48 --- /dev/null +++ b/tests/compatibility/cases/exemplar_sparse/setup.sql @@ -0,0 +1,15 @@ +CREATE TABLE exemplar_source_ids ( + __table_id INT UNSIGNED NOT NULL, + __tsid BIGINT UNSIGNED NOT NULL, + greptime_metric STRING, + greptime_timestamp TIMESTAMP(9) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY (__table_id, __tsid, greptime_metric) +) WITH ('append_mode'='true', 'primary_key_encoding'='sparse'); +INSERT INTO exemplar_source_ids VALUES + (42, 18446744073709551615, 'cpu', 20, 1.0), + (42, 18446744073709551615, 'cpu', 20, 1.0), + (43, 18446744073709551615, 'memory', 10, 2.0); +ADMIN FLUSH_TABLE('exemplar_source_ids'); +ALTER TABLE exemplar_source_ids ADD COLUMN job STRING PRIMARY KEY; +INSERT INTO exemplar_source_ids VALUES (42, 7, 'cpu', 30, 3.0, 'api'); diff --git a/tests/compatibility/cases/exemplar_sparse/verify.result b/tests/compatibility/cases/exemplar_sparse/verify.result new file mode 100644 index 00000000000..67e016ea127 --- /dev/null +++ b/tests/compatibility/cases/exemplar_sparse/verify.result @@ -0,0 +1,42 @@ +SELECT * FROM exemplar_source_ids ORDER BY __tsid, greptime_timestamp; + ++------------+----------------------+-----------------+-------------------------------+----------------+-----+ +| __table_id | __tsid | greptime_metric | greptime_timestamp | greptime_value | job | ++------------+----------------------+-----------------+-------------------------------+----------------+-----+ +| 42 | 7 | cpu | 1970-01-01T00:00:00.000000030 | 3.0 | api | +| 43 | 18446744073709551615 | memory | 1970-01-01T00:00:00.000000010 | 2.0 | | +| 42 | 18446744073709551615 | cpu | 1970-01-01T00:00:00.000000020 | 1.0 | | +| 42 | 18446744073709551615 | cpu | 1970-01-01T00:00:00.000000020 | 1.0 | | ++------------+----------------------+-----------------+-------------------------------+----------------+-----+ + +ADMIN FLUSH_TABLE('exemplar_source_ids'); + ++------------------------------------------+ +| ADMIN FLUSH_TABLE('exemplar_source_ids') | ++------------------------------------------+ +| 0 | ++------------------------------------------+ + +ADMIN COMPACT_TABLE('exemplar_source_ids', 'swcs', '86400'); + ++-------------------------------------------------------------+ +| ADMIN COMPACT_TABLE('exemplar_source_ids', 'swcs', '86400') | ++-------------------------------------------------------------+ +| 0 | ++-------------------------------------------------------------+ + +SELECT COUNT(*) FROM exemplar_source_ids WHERE __table_id = 42 AND __tsid = 18446744073709551615 AND job IS NULL; + ++----------+ +| count(*) | ++----------+ +| 2 | ++----------+ + +SELECT COUNT(*) FROM exemplar_source_ids WHERE __table_id = 43 AND __tsid = 18446744073709551615; + ++----------+ +| count(*) | ++----------+ +| 1 | ++----------+ diff --git a/tests/compatibility/cases/exemplar_sparse/verify.sql b/tests/compatibility/cases/exemplar_sparse/verify.sql new file mode 100644 index 00000000000..71f1112bd9d --- /dev/null +++ b/tests/compatibility/cases/exemplar_sparse/verify.sql @@ -0,0 +1,5 @@ +SELECT * FROM exemplar_source_ids ORDER BY __tsid, greptime_timestamp; +ADMIN FLUSH_TABLE('exemplar_source_ids'); +ADMIN COMPACT_TABLE('exemplar_source_ids', 'swcs', '86400'); +SELECT COUNT(*) FROM exemplar_source_ids WHERE __table_id = 42 AND __tsid = 18446744073709551615 AND job IS NULL; +SELECT COUNT(*) FROM exemplar_source_ids WHERE __table_id = 43 AND __tsid = 18446744073709551615; From c17a5e09f74069a7e62b9a512d6d95229344c219 Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Tue, 3 Nov 2026 12:52:05 +0000 Subject: [PATCH 161/239] fix(auth): preserve physical metric table targets (#10084) * fix(auth): preserve physical metric table targets Signed-off-by: shuiyisong * fix: nit Signed-off-by: shuiyisong --------- Signed-off-by: shuiyisong --- src/auth/src/permission.rs | 76 +++- src/frontend/src/instance.rs | 156 +++++++- src/frontend/src/instance/grpc.rs | 34 +- src/servers/src/http/prometheus.rs | 529 ++++++++++++++------------ src/servers/src/prometheus_handler.rs | 9 +- 5 files changed, 526 insertions(+), 278 deletions(-) diff --git a/src/auth/src/permission.rs b/src/auth/src/permission.rs index 901b935d6a2..5342076f3f5 100644 --- a/src/auth/src/permission.rs +++ b/src/auth/src/permission.rs @@ -34,6 +34,9 @@ pub struct PermissionTableTarget { pub catalog: String, pub schema: String, pub table: String, + /// Set from catalog metadata, never from a client-supplied table name. + /// Physical metric tables contain data belonging to multiple logical tables. + pub is_physical_metric_table: bool, } impl PermissionTableTarget { @@ -47,8 +50,15 @@ impl PermissionTableTarget { catalog: catalog.into(), schema: schema.into(), table: table.into(), + is_physical_metric_table: false, } } + + /// Sets whether catalog metadata identifies this target as a physical metric table. + pub fn with_physical_metric_table(mut self, is_physical_metric_table: bool) -> Self { + self.is_physical_metric_table = is_physical_metric_table; + self + } } /// The result of resolving every table target in a permission request. @@ -183,6 +193,11 @@ pub enum PermissionReq<'a> { GrpcRequest(&'a Request), SqlStatement(&'a Statement), Action(PermissionAction), + /// Enumerates a schema's metadata without accessing any table data. + SchemaRead { + catalog: &'a str, + schema: &'a str, + }, BulkInsert { catalog: &'a str, schema: &'a str, @@ -199,6 +214,7 @@ impl<'a> PermissionReq<'a> { } PermissionReq::SqlStatement(stmt) => stmt.is_readonly(), PermissionReq::Action(action) => action.access_mode() == AccessMode::Read, + PermissionReq::SchemaRead { .. } => true, PermissionReq::GrpcRequest(_) | PermissionReq::BulkInsert { .. } => false, } @@ -241,6 +257,17 @@ pub trait PermissionChecker: Send + Sync { true } + /// Extracts user-visible gRPC targets before the frontend classifies them + /// using catalog metadata. Unrecognized requests must fail closed. + fn grpc_table_targets( + &self, + _request: &Request, + _catalog: &str, + _schema: &str, + ) -> PermissionTableTargets { + PermissionTableTargets::Unresolved + } + /// Checks the operation privilege and its resolved user-visible table targets. /// /// [`PermissionTableTargets::Resolved`] with an empty vector still requires @@ -296,6 +323,18 @@ impl PermissionChecker for Option<&PermissionCheckerRef> { } } + fn grpc_table_targets( + &self, + request: &Request, + catalog: &str, + schema: &str, + ) -> PermissionTableTargets { + match self { + Some(checker) => checker.grpc_table_targets(request, catalog, schema), + None => PermissionTableTargets::Unresolved, + } + } + fn check_permission_with_table_targets( &self, user_info: UserInfoRef, @@ -386,6 +425,20 @@ mod tests { Ok(PermissionResp::Reject) } + fn grpc_table_targets( + &self, + request: &Request, + catalog: &str, + schema: &str, + ) -> PermissionTableTargets { + match request { + Request::RowInserts(requests) => { + PermissionTableTargets::from_row_insert_requests(catalog, schema, requests) + } + _ => PermissionTableTargets::Unresolved, + } + } + fn check_permission_with_table_targets( &self, _user_info: UserInfoRef, @@ -600,8 +653,8 @@ mod tests { #[test] fn test_table_target_permission_forwarding() { - let checker: PermissionCheckerRef = Arc::new(TargetAwarePermissionChecker); - let checker = Some(&checker); + let inner_checker: PermissionCheckerRef = Arc::new(TargetAwarePermissionChecker); + let checker = Some(&inner_checker); assert!(checker.uses_table_targets()); let allowed = checker @@ -658,6 +711,25 @@ mod tests { let no_checker: Option<&PermissionCheckerRef> = None; assert!(!no_checker.uses_table_targets()); + for request in [ + Request::RowInserts(RowInsertRequests { + inserts: vec![api::v1::RowInsertRequest { + table_name: "allowed".to_string(), + ..Default::default() + }], + }), + Request::RowInserts(RowInsertRequests::default()), + Request::Query(api::v1::QueryRequest::default()), + ] { + assert_eq!( + inner_checker.grpc_table_targets(&request, "catalog", "schema"), + checker.grpc_table_targets(&request, "catalog", "schema"), + ); + assert_eq!( + PermissionTableTargets::Unresolved, + no_checker.grpc_table_targets(&request, "catalog", "schema"), + ); + } let allowed = no_checker .check_permission_with_table_targets( crate::userinfo_by_name(None), diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index 554a9677af7..4580b611352 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -1238,29 +1238,16 @@ impl Instance { let mut seen = HashSet::with_capacity(targets.len()); targets.retain(|target| seen.insert(target.clone())); } - if let [target] = targets.as_slice() { - return if self - .is_physical_query_permission_target(target, query_ctx) - .await? - { - Ok(PermissionTableTargets::Unresolved) - } else { - Ok(PermissionTableTargets::resolved(targets)) - }; - } - // Bound catalog load and inspect results in target order to preserve serial semantics. - for chunk in targets.chunks(CONCURRENCY) { + for chunk in targets.chunks_mut(CONCURRENCY) { let results = future::join_all( chunk .iter() .map(|target| self.is_physical_query_permission_target(target, query_ctx)), ) .await; - for result in results { - if result? { - return Ok(PermissionTableTargets::Unresolved); - } + for (target, result) in chunk.iter_mut().zip(results) { + target.is_physical_metric_table = result?; } } @@ -1494,6 +1481,35 @@ impl PrometheusHandler for Instance { Ok(allowed) } + async fn filter_metadata_schema_names( + &self, + schema_names: Vec, + query_ctx: &QueryContextRef, + ) -> server_error::Result> { + let checker = self.plugins.get::(); + let mut allowed = Vec::with_capacity(schema_names.len()); + for schema in schema_names { + match checker + .as_ref() + .check_permission( + query_ctx.current_user(), + PermissionReq::SchemaRead { + catalog: query_ctx.current_catalog(), + schema: &schema, + }, + ) + .context(AuthSnafu) + { + Ok(_) => allowed.push(schema), + Err(error) + if error.status_code() + == common_error::status_code::StatusCode::PermissionDenied => {} + Err(error) => return Err(error), + } + } + Ok(allowed) + } + async fn query_metric_names( &self, matchers: Vec, @@ -2121,9 +2137,9 @@ mod tests { ) -> auth::error::Result { let reject = match targets { PermissionTableTargets::Unresolved => true, - PermissionTableTargets::Resolved(targets) => { - targets.iter().any(|target| target.table == "denied") - } + PermissionTableTargets::Resolved(targets) => targets + .iter() + .any(|target| target.table == "denied" || target.is_physical_metric_table), }; Ok(if reject { PermissionResp::Reject @@ -3553,6 +3569,103 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_physical_target_permission_reaches_checker() -> TestResult<()> { + struct PhysicalChecker { + allow: bool, + } + impl PermissionChecker for PhysicalChecker { + fn check_permission( + &self, + _: UserInfoRef, + _: PermissionReq, + ) -> auth::error::Result { + Ok(PermissionResp::Allow) + } + + fn grpc_table_targets( + &self, + request: &api::v1::greptime_request::Request, + catalog: &str, + schema: &str, + ) -> PermissionTableTargets { + let api::v1::greptime_request::Request::Ddl(ddl) = request else { + panic!("expected DDL") + }; + let Some(api::v1::ddl_request::Expr::CommentOn(comment)) = &ddl.expr else { + panic!("expected comment") + }; + PermissionTableTargets::resolved(vec![PermissionTableTarget::new( + catalog, + schema, + &comment.object_name, + )]) + } + + fn check_permission_with_table_targets( + &self, + _: UserInfoRef, + _: PermissionReq, + targets: PermissionTableTargets, + ) -> auth::error::Result { + let PermissionTableTargets::Resolved(targets) = targets else { + panic!("physical target must retain its identity") + }; + assert_eq!(targets.len(), 1); + assert_eq!(targets[0].catalog, "greptime"); + assert_eq!(targets[0].schema, "public"); + assert_eq!(targets[0].table, "physical_metric"); + assert!(targets[0].is_physical_metric_table); + Ok(if self.allow { + PermissionResp::Allow + } else { + PermissionResp::Reject + }) + } + } + for allow in [true, false] { + let plugins = Plugins::new(); + plugins.insert::(Arc::new(PhysicalChecker { allow })); + let instance = test_instance_with_plugins( + test_physical_table(1024, "physical_metric")?, + test_table(1025, "target")?, + plugins, + ) + .await?; + let ctx = test_query_ctx(1); + for sql in [ + "SHOW CREATE TABLE physical_metric", + "SELECT * FROM physical_metric", + ] { + let result = instance + .check_sql_permission(&parse_one_sql(sql), &ctx) + .await; + assert_eq!(allow, result.is_ok()); + if let Err(error) = result { + assert_eq!(StatusCode::PermissionDenied, error.status_code()); + } + } + if !allow { + let request = api::v1::greptime_request::Request::Ddl(api::v1::DdlRequest { + expr: Some(api::v1::ddl_request::Expr::CommentOn( + api::v1::CommentOnExpr { + object_name: "physical_metric".to_string(), + object_type: api::v1::CommentObjectType::Table as i32, + ..Default::default() + }, + )), + }); + let err = servers::query_handler::grpc::GrpcQueryHandler::do_query( + &instance, request, ctx, + ) + .await + .unwrap_err(); + assert_eq!(StatusCode::PermissionDenied, err.status_code()); + } + } + Ok(()) + } + #[tokio::test] async fn test_physical_query_targets_fail_closed() -> TestResult<()> { let physical_table = "physical_metric"; @@ -3577,9 +3690,10 @@ mod tests { .await .unwrap() ); - let physical_target = PermissionTableTarget::new("greptime", "public", physical_table); + let physical_target = PermissionTableTarget::new("greptime", "public", physical_table) + .with_physical_metric_table(true); assert_eq!( - PermissionTableTargets::Unresolved, + PermissionTableTargets::Resolved(vec![logical_target.clone(), physical_target.clone()]), instance .resolve_query_permission_targets( PermissionTableTargets::resolved( diff --git a/src/frontend/src/instance/grpc.rs b/src/frontend/src/instance/grpc.rs index 7dfd97cd69d..4d1ff1e8d2b 100644 --- a/src/frontend/src/instance/grpc.rs +++ b/src/frontend/src/instance/grpc.rs @@ -27,7 +27,8 @@ use api::v1::{ use async_stream::try_stream; use async_trait::async_trait; use auth::{ - PermissionChecker, PermissionCheckerRef, PermissionReq, PermissionResp, PermissionTableTargets, + PermissionChecker, PermissionCheckerRef, PermissionReq, PermissionResp, PermissionTableTarget, + PermissionTableTargets, }; use common_error::ext::BoxedError; use common_grpc::flight::do_put::DoPutResponse; @@ -85,6 +86,20 @@ impl GrpcQueryHandler for Instance { Some(&ctx.current_schema()), ) .context(PermissionSnafu)?; + + if let Some(checker) = self.plugins.get::() + && checker.uses_table_targets() + { + let targets = checker.grpc_table_targets( + &request, + ctx.current_catalog(), + &ctx.current_schema(), + ); + let targets = self.resolve_query_permission_targets(targets, &ctx) + .await.map_err(BoxedError::new).context(ExternalSnafu)?; + self.check_table_permission(&ctx, PermissionReq::GrpcRequest(&request), targets) + .context(PermissionSnafu)?; + } } let output = match request { @@ -439,6 +454,23 @@ impl Instance { table_name: table_name.to_string(), })?; + let target = PermissionTableTarget::new( + &table_name.catalog_name, + &table_name.schema_name, + &table_name.table_name, + ) + .with_physical_metric_table(table.table_info().is_physical_table()); + plugins.get::().as_ref() + .check_permission_with_table_targets( + ctx.current_user(), + PermissionReq::BulkInsert { + catalog: &table_name.catalog_name, + schema: &table_name.schema_name, + table: &table_name.table_name, + }, + PermissionTableTargets::resolved(vec![target]), + ).context(PermissionSnafu)?; + // Check permissions for the table let interceptor_ref = plugins.get::>(); let interceptor = interceptor_ref.as_ref(); diff --git a/src/servers/src/http/prometheus.rs b/src/servers/src/http/prometheus.rs index e417044ef43..82bbc665202 100644 --- a/src/servers/src/http/prometheus.rs +++ b/src/servers/src/http/prometheus.rs @@ -60,7 +60,7 @@ use serde::de::{self, MapAccess, Visitor}; use serde::{Deserialize, Serialize}; use serde_json::Value; use session::context::{QueryContext, QueryContextRef}; -use snafu::{Location, OptionExt, ResultExt}; +use snafu::{Location, OptionExt, ResultExt, ensure}; use store_api::metric_engine_consts::{ DATA_SCHEMA_TABLE_ID_COLUMN_NAME, DATA_SCHEMA_TSID_COLUMN_NAME, LOGICAL_TABLE_METADATA_KEY, }; @@ -1828,61 +1828,17 @@ pub async fn label_values_query( return PrometheusJsonResponse::success(PrometheusResponse::LabelValues(table_names)); } else if label_name == FIELD_NAME_LABEL { try_call_return_response!(handler.check_query_permission(&[], &query_ctx).await); - let enumerate_all = matches.is_empty(); - let metric_names = matches - .iter() - .map(|selector| retrieve_exact_metric_name_from_promql(selector)) - .collect::>(); - if metric_names.iter().any(Option::is_none) { - try_call_return_response!( - handler - .check_query_target_permission(PermissionTableTargets::Unresolved, &query_ctx,) - .await - ); - } - let metric_names = metric_names.into_iter().flatten().collect::>(); - - try_call_return_response!( - handler - .check_query_target_permission( - current_schema_metric_targets(&query_ctx, &metric_names), - &query_ctx, - ) - .await - ); - - let (field_columns, table_names) = if !enumerate_all && metric_names.is_empty() { - (HashSet::new(), Vec::new()) - } else { - handle_schema_err!( - retrieve_field_names(&query_ctx, handler.catalog_manager(), metric_names).await - ) - .unwrap_or_default() - }; - try_call_return_response!( - handler - .check_query_target_permission( - current_schema_metric_targets(&query_ctx, &table_names), - &query_ctx, - ) - .await - ); + let field_columns = + handle_schema_err!(retrieve_field_names(&query_ctx, &handler, matches).await) + .unwrap_or_default(); let mut field_columns = field_columns.into_iter().collect::>(); field_columns.sort_unstable(); truncate_results(&mut field_columns, params.limit); return PrometheusJsonResponse::success(PrometheusResponse::LabelValues(field_columns)); } else if is_database_selection_label(&label_name) { try_call_return_response!(handler.check_query_permission(&[], &query_ctx).await); - let catalog_manager = handler.catalog_manager(); - - let (mut schema_names, targets) = try_call_return_response!( - retrieve_schema_names(&query_ctx, catalog_manager, matches).await - ); - try_call_return_response!( - handler - .check_query_target_permission(targets, &query_ctx) - .await - ); + let mut schema_names = + try_call_return_response!(retrieve_schema_names(&query_ctx, &handler, matches).await); truncate_results(&mut schema_names, params.limit); return PrometheusJsonResponse::success(PrometheusResponse::LabelValues(schema_names)); } @@ -2300,100 +2256,132 @@ fn table_has_native_histogram_value(table_info: &TableInfo) -> bool { .any(|column| is_native_histogram_value_type(&column.data_type)) } -async fn retrieve_field_names( - query_ctx: &QueryContext, - manager: CatalogManagerRef, - matches: Vec, -) -> Result<(HashSet, Vec)> { - let mut field_columns = HashSet::new(); - let mut table_names = Vec::new(); - let catalog = query_ctx.current_catalog(); - let schema = query_ctx.current_schema(); +fn metadata_name_selectors(matches: &[String]) -> Result>> { + matches + .iter() + .map(|query| { + let expr = promql_parser::parser::parse(query) + .map_err(|reason| InvalidQuerySnafu { reason }.build())?; + let PromqlExpr::VectorSelector(mut selector) = expr else { + return InvalidQuerySnafu { + reason: "expected vector selector".to_string(), + } + .fail(); + }; + ensure!( + selector.matchers.or_matchers.is_empty(), + InvalidQuerySnafu { + reason: "OR matchers are not supported in schema and field metadata queries" + .to_string(), + } + ); + Ok(take_metric_name_matchers(&mut selector)) + }) + .collect() +} - if matches.is_empty() { - // query all tables if no matcher is provided - let mut tables = manager.tables(catalog, &schema, Some(query_ctx)); +async fn readable_metadata_tables( + query_ctx: &QueryContextRef, + handler: &PrometheusHandlerRef, + schema: &str, + selectors: &[Vec], +) -> Result> { + let manager = handler.catalog_manager(); + let exact_names = selectors + .iter() + .map(|matchers| { + matchers + .iter() + .find(|matcher| matcher.op == MatchOp::Equal) + .map(|matcher| &matcher.value) + }) + .collect::>>(); + let mut candidates = Vec::new(); + if let Some(names) = exact_names.filter(|names| !names.is_empty()) { + for name in names { + if let Some(table) = manager + .table(query_ctx.current_catalog(), schema, name, Some(query_ctx)) + .await? + { + candidates.push(table); + } + } + } else { + let mut tables = manager.tables(query_ctx.current_catalog(), schema, Some(query_ctx)); while let Some(table) = tables.next().await { - let table = table?; - if is_internal_physical_metric_table(&table) { - continue; - } - table_names.push(table.table_info().name.clone()); - for column in table.field_columns() { - field_columns.insert(column.name); - } - } - return Ok((field_columns, table_names)); - } - - for table_name in matches { - let table = manager - .table(catalog, &schema, &table_name, Some(query_ctx)) - .await? - .with_context(|| TableNotFoundSnafu { - catalog: catalog.to_string(), - schema: schema.clone(), - table: table_name.clone(), - })?; - - if is_internal_physical_metric_table(&table) { - continue; - } - table_names.push(table.table_info().name.clone()); - for column in table.field_columns() { - field_columns.insert(column.name); + candidates.push(table?); } } - Ok((field_columns, table_names)) + candidates.retain(|table| { + !is_internal_physical_metric_table(table) + && (selectors.is_empty() + || selectors + .iter() + .any(|matchers| metric_name_matches(&table.table_info().name, matchers))) + }); + let names = candidates + .iter() + .map(|table| table.table_info().name.clone()) + .collect(); + let allowed: HashSet<_> = handler + .filter_metadata_metric_names(names, schema, query_ctx) + .await? + .into_iter() + .collect(); + candidates.retain(|table| allowed.contains(&table.table_info().name)); + Ok(candidates) +} + +async fn retrieve_field_names( + query_ctx: &QueryContextRef, + handler: &PrometheusHandlerRef, + matches: Vec, +) -> Result> { + let mut selectors = metadata_name_selectors(&matches)?; + // Only an omitted match[] enumerates all fields; label-only selectors add no candidates. + selectors.retain(|matchers| !matchers.is_empty()); + if !matches.is_empty() && selectors.is_empty() { + return Ok(HashSet::new()); + } + let tables = + readable_metadata_tables(query_ctx, handler, &query_ctx.current_schema(), &selectors) + .await?; + Ok(tables + .iter() + .flat_map(|table| table.field_columns().map(|column| column.name)) + .collect()) } async fn retrieve_schema_names( - query_ctx: &QueryContext, - catalog_manager: CatalogManagerRef, + query_ctx: &QueryContextRef, + handler: &PrometheusHandlerRef, matches: Vec, -) -> Result<(Vec, PermissionTableTargets)> { - let mut schemas = Vec::new(); - let mut targets = Vec::new(); - let catalog = query_ctx.current_catalog(); - let metric_names = matches - .iter() - .map(|match_item| retrieve_exact_metric_name_from_promql(match_item)) - .collect::>(); - let unresolved = metric_names.is_empty() || metric_names.iter().any(Option::is_none); - - let candidate_schemas = catalog_manager - .schema_names(catalog, Some(query_ctx)) +) -> Result> { + let selectors = metadata_name_selectors(&matches)?; + let candidate_schemas = handler + .catalog_manager() + .schema_names(query_ctx.current_catalog(), Some(query_ctx)) .await?; - + let candidate_schemas = handler + .filter_metadata_schema_names(candidate_schemas, query_ctx) + .await?; + let mut schemas = Vec::new(); for schema in candidate_schemas { - let mut found = true; - for table_name in metric_names.iter().flatten() { - let table = catalog_manager - .table(catalog, &schema, table_name, Some(query_ctx)) - .await?; - let Some(table) = table else { - found = false; - continue; - }; - targets.push(PermissionTableTarget::new(catalog, &schema, table_name)); - if is_internal_physical_metric_table(&table) { - found = false; - } + if selectors.is_empty() { + schemas.push(schema); + continue; } - - if found { + let tables = readable_metadata_tables(query_ctx, handler, &schema, &selectors).await?; + if selectors.iter().all(|matchers| { + tables + .iter() + .any(|table| metric_name_matches(&table.table_info().name, matchers)) + }) { schemas.push(schema); } } - schemas.sort_unstable(); - - let targets = if unresolved { - PermissionTableTargets::Unresolved - } else { - PermissionTableTargets::resolved(targets) - }; - Ok((schemas, targets)) + Ok(schemas) } fn retrieve_exact_metric_name_from_promql(query: &str) -> Option { @@ -2676,13 +2664,13 @@ mod tests { targets: PermissionTableTargets, _: &QueryContextRef, ) -> Result<()> { - if matches!( - targets, - PermissionTableTargets::Resolved(targets) - if self.denied_table.is_some_and(|denied| { - targets.iter().any(|target| target.table == denied) - }) - ) { + let denied = match targets { + PermissionTableTargets::Unresolved => true, + PermissionTableTargets::Resolved(targets) => self + .denied_table + .is_some_and(|denied| targets.iter().any(|target| target.table == denied)), + }; + if denied { return auth::error::PermissionDeniedSnafu .fail() .context(crate::error::AuthSnafu); @@ -2702,6 +2690,17 @@ mod tests { .collect()) } + async fn filter_metadata_schema_names( + &self, + schema_names: Vec, + _: &QueryContextRef, + ) -> Result> { + Ok(schema_names + .into_iter() + .filter(|schema| schema != "denied_schema") + .collect()) + } + async fn query_metric_names( &self, _: Vec, @@ -2926,64 +2925,127 @@ mod tests { assert_eq!(range_count + 1, range_histogram.get_sample_count()); } - #[tokio::test] - async fn test_field_name_enumeration_fails_closed() { - use axum::response::IntoResponse; + fn metadata_test_handler(manager: CatalogManagerRef) -> PrometheusHandlerRef { + Arc::new(TestPrometheusHandler { + catalog_manager: manager, + deny_operation: false, + denied_table: Some("denied"), + metric_names: Vec::new(), + label_metric_names: Vec::new(), + label_lookups: Mutex::new(Vec::new()), + queries: Mutex::new(Vec::new()), + ordered_outputs: Mutex::new(Vec::new()), + }) + } - let mut allowed = test_table_info( - 1024, - "allowed", - DEFAULT_SCHEMA_NAME, + #[tokio::test] + async fn test_metadata_enumeration_filters_inaccessible_candidates() { + let manager = MemoryCatalogManager::with_default_setup(); + for schema in ["empty", "denied_schema"] { + manager + .register_schema_sync(RegisterSchemaRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: schema.to_string(), + }) + .unwrap(); + } + for (id, schema, name, field) in [ + (1024, DEFAULT_SCHEMA_NAME, "allowed", "allowed_field"), + (1025, DEFAULT_SCHEMA_NAME, "denied", "secret_field"), + (1026, "denied_schema", "allowed", "hidden_field"), + ] { + let info = test_table_info( + id, + name, + schema, + DEFAULT_CATALOG_NAME, + Arc::new(Schema::new(vec![ColumnSchema::new( + field, + ConcreteDataType::float64_datatype(), + true, + )])), + ); + manager + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: schema.to_string(), + table_name: name.to_string(), + table_id: id, + table: EmptyTable::from_table_info(&info), + }) + .unwrap(); + } + let handler = metadata_test_handler(manager); + let query_ctx = Arc::new(QueryContext::with( DEFAULT_CATALOG_NAME, - Arc::new(Schema::new(vec![])), - ); - allowed.meta.options.extra_options.insert( - LOGICAL_TABLE_METADATA_KEY.to_string(), - "physical_metrics".to_string(), - ); - let manager = MemoryCatalogManager::new_with_table(EmptyTable::from_table_info(&allowed)); - let mut denied = allowed.clone(); - denied.ident.table_id = 1025; - denied.name = "denied".to_string(); - manager - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: denied.name.clone(), - table_id: denied.table_id(), - table: EmptyTable::from_table_info(&denied), - }) - .unwrap(); - let query_ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME); - assert_eq!( + DEFAULT_SCHEMA_NAME, + )); + for selectors in [ + vec![], + vec![r#"{__name__=~"allowed|denied"}"#.to_string()], vec!["allowed".to_string(), "denied".to_string()], - retrieve_table_names(&query_ctx, manager.clone(), Vec::new()) + vec!["allowed".to_string(), r#"{job="nonexistent"}"#.to_string()], + vec![r#"allowed{__schema__="other"}"#.to_string()], + ] { + let fields = retrieve_field_names(&query_ctx, &handler, selectors) + .await + .unwrap(); + assert_eq!(fields, HashSet::from(["allowed_field".to_string()])); + } + for selectors in [ + vec![r#"{job="nonexistent"}"#.to_string()], + vec![r#"{__database__="nonexistent"}"#.to_string()], + vec![r#"{__field__="nonexistent"}"#.to_string()], + vec!["missing".to_string(), r#"{job="nonexistent"}"#.to_string()], + ] { + assert!( + retrieve_field_names(&query_ctx, &handler, selectors) + .await + .unwrap() + .is_empty() + ); + } + assert!( + retrieve_field_names(&query_ctx, &handler, vec!["denied".to_string()]) .await .unwrap() + .is_empty() ); - + assert!( + retrieve_field_names(&query_ctx, &handler, vec!["{".to_string()]) + .await + .is_err() + ); + for selector in ["allowed", r#"{__name__=~"allowed|denied"}"#] { + assert_eq!( + retrieve_schema_names(&query_ctx, &handler, vec![selector.to_string()]) + .await + .unwrap(), + vec![DEFAULT_SCHEMA_NAME] + ); + } + assert!( + retrieve_schema_names(&query_ctx, &handler, vec!["denied".to_string()]) + .await + .unwrap() + .is_empty() + ); + let schemas = retrieve_schema_names(&query_ctx, &handler, vec![]) + .await + .unwrap(); + assert!(schemas.contains(&"empty".to_string())); + assert!(!schemas.contains(&"denied_schema".to_string())); let response = label_values_query( - State(Arc::new(TestPrometheusHandler { - catalog_manager: manager, - deny_operation: false, - denied_table: Some("denied"), - metric_names: Vec::new(), - label_metric_names: Vec::new(), - label_lookups: Mutex::new(Vec::new()), - queries: Mutex::new(Vec::new()), - ordered_outputs: Mutex::new(Vec::new()), - })), + State(handler), Path(FIELD_NAME_LABEL.to_string()), - Extension(query_ctx), + Extension(QueryContext::with( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + )), Query(LabelValueQuery::default()), ) .await; - - assert_eq!(Some(StatusCode::PermissionDenied), response.status_code); - assert_eq!( - axum::http::StatusCode::FORBIDDEN, - response.into_response().status() - ); + assert_eq!(None, response.status_code); } /// A handler over the logical metric tables `cpu_user` and `cpu_system`, @@ -3649,69 +3711,36 @@ mod tests { ); assert_eq!(table_names, vec!["cpu_usage"]); - let query_ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME); - let (schemas, targets) = - retrieve_schema_names(&query_ctx, manager.clone(), vec!["cpu_usage".to_string()]) + let query_ctx = Arc::new(QueryContext::with( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + )); + let handler = metadata_test_handler(manager.clone()); + for selector in ["cpu_usage", r#"{__name__=~"cpu.*"}"#] { + assert_eq!( + retrieve_schema_names(&query_ctx, &handler, vec![selector.to_string()]) + .await + .unwrap(), + vec![DEFAULT_SCHEMA_NAME] + ); + } + assert!( + retrieve_schema_names(&query_ctx, &handler, vec![]) .await - .unwrap(); - assert_eq!(schemas, vec![DEFAULT_SCHEMA_NAME]); - assert_eq!( - targets, - PermissionTableTargets::Resolved(vec![PermissionTableTarget::new( - DEFAULT_CATALOG_NAME, - DEFAULT_SCHEMA_NAME, - "cpu_usage", - )]) + .unwrap() + .contains(&DEFAULT_SCHEMA_NAME.to_string()) ); - - let (_, targets) = retrieve_schema_names( - &query_ctx, - manager.clone(), - vec![r#"{__name__=~"cpu.*"}"#.to_string()], - ) - .await - .unwrap(); - assert_eq!(targets, PermissionTableTargets::Unresolved); - - let (schemas, targets) = retrieve_schema_names(&query_ctx, manager.clone(), vec![]) - .await - .unwrap(); - assert!(schemas.contains(&DEFAULT_SCHEMA_NAME.to_string())); - assert_eq!(targets, PermissionTableTargets::Unresolved); - - let (schemas, targets) = retrieve_schema_names( - &query_ctx, - manager.clone(), + for selectors in [ vec!["physical_metrics".to_string()], - ) - .await - .unwrap(); - assert!(schemas.is_empty()); - assert_eq!( - targets, - PermissionTableTargets::Resolved(vec![PermissionTableTarget::new( - DEFAULT_CATALOG_NAME, - DEFAULT_SCHEMA_NAME, - "physical_metrics", - )]) - ); - - let (schemas, targets) = retrieve_schema_names( - &query_ctx, - manager.clone(), vec!["missing".to_string(), "physical_metrics".to_string()], - ) - .await - .unwrap(); - assert!(schemas.is_empty()); - assert_eq!( - targets, - PermissionTableTargets::Resolved(vec![PermissionTableTarget::new( - DEFAULT_CATALOG_NAME, - DEFAULT_SCHEMA_NAME, - "physical_metrics", - )]) - ); + ] { + assert!( + retrieve_schema_names(&query_ctx, &handler, selectors) + .await + .unwrap() + .is_empty() + ); + } assert_eq!( vec!["cpu_usage".to_string()], @@ -3733,21 +3762,15 @@ mod tests { .is_err() ); - let (fields, table_names) = retrieve_field_names( - &query_ctx, - manager.clone(), - vec!["physical_metrics".to_string()], - ) - .await - .unwrap(); + let fields = + retrieve_field_names(&query_ctx, &handler, vec!["physical_metrics".to_string()]) + .await + .unwrap(); assert!(fields.is_empty()); - assert!(table_names.is_empty()); - - let (fields, table_names) = retrieve_field_names(&query_ctx, manager, vec![]) + let fields = retrieve_field_names(&query_ctx, &handler, vec![]) .await .unwrap(); assert_eq!(fields, HashSet::from(["value".to_string()])); - assert_eq!(table_names, vec!["cpu_usage"]); } #[tokio::test] diff --git a/src/servers/src/prometheus_handler.rs b/src/servers/src/prometheus_handler.rs index 1e211b501ab..2c7aec1fb61 100644 --- a/src/servers/src/prometheus_handler.rs +++ b/src/servers/src/prometheus_handler.rs @@ -168,7 +168,7 @@ pub trait PrometheusHandler { query_ctx: &QueryContextRef, ) -> Result<()>; - /// Removes inaccessible metric names from logical-table metadata enumeration results. + /// Removes inaccessible names from non-physical table metadata enumeration results. async fn filter_metadata_metric_names( &self, metric_names: Vec, @@ -176,6 +176,13 @@ pub trait PrometheusHandler { query_ctx: &QueryContextRef, ) -> Result>; + /// Removes schemas the user cannot access from metadata enumeration results. + async fn filter_metadata_schema_names( + &self, + schema_names: Vec, + query_ctx: &QueryContextRef, + ) -> Result>; + /// Query metric table names by the `__name__` matchers. async fn query_metric_names( &self, From 8deb3872b3283b01ebc629a5d1642b7053b65fbb Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Tue, 3 Nov 2026 13:25:14 +0000 Subject: [PATCH 162/239] fix: include peer context in distributed query errors (#10118) Signed-off-by: WenyXu --- src/catalog/src/error.rs | 3 ++- src/catalog/src/information_extension.rs | 4 +++- src/client/src/error.rs | 2 +- 3 files changed, 6 insertions(+), 3 deletions(-) diff --git a/src/catalog/src/error.rs b/src/catalog/src/error.rs index 58c5ef96b96..24bf61a6b0e 100644 --- a/src/catalog/src/error.rs +++ b/src/catalog/src/error.rs @@ -291,8 +291,9 @@ pub enum Error { location: Location, }, - #[snafu(display("Failed to handle query"))] + #[snafu(display("Failed to handle query, addr: {}", addr))] HandleQuery { + addr: String, source: common_meta::error::Error, #[snafu(implicit)] location: Location, diff --git a/src/catalog/src/information_extension.rs b/src/catalog/src/information_extension.rs index f22c169df13..0c572eac562 100644 --- a/src/catalog/src/information_extension.rs +++ b/src/catalog/src/information_extension.rs @@ -84,7 +84,9 @@ impl DistributedInformationExtension { header: None, }) .await - .context(crate::error::HandleQuerySnafu) + .context(crate::error::HandleQuerySnafu { + addr: node.peer.addr, + }) } })) .await?; diff --git a/src/client/src/error.rs b/src/client/src/error.rs index d9fa8c96050..afa834655e5 100644 --- a/src/client/src/error.rs +++ b/src/client/src/error.rs @@ -42,7 +42,7 @@ pub enum Error { location: Location, }, - #[snafu(display("Failed to do Flight get, code: {}", tonic_code))] + #[snafu(display("Failed to do Flight get, addr: {}, code: {}", addr, tonic_code))] FlightGet { addr: String, tonic_code: Code, From adf29740b28147cfb9035f53e1c41515b5b519b3 Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Wed, 4 Nov 2026 03:39:51 +0000 Subject: [PATCH 163/239] fix(auth): redact inline user-provider credentials (#10115) * fix(auth): redact inline user-provider credentials Signed-off-by: shuiyisong * fix(auth): close user-provider redaction gaps Signed-off-by: shuiyisong --------- Signed-off-by: shuiyisong --- Cargo.lock | 1 + src/auth/src/common.rs | 129 +++++++++++++++++- src/auth/src/lib.rs | 5 +- .../src/user_provider/static_user_provider.rs | 7 +- src/cmd/src/frontend.rs | 28 ++-- src/cmd/src/standalone.rs | 31 +++-- src/cmd/tests/load_config_test.rs | 42 +++++- src/frontend/src/frontend.rs | 60 +++++++- src/plugins/src/frontend.rs | 5 +- src/standalone/Cargo.toml | 1 + src/standalone/src/options.rs | 22 ++- 11 files changed, 290 insertions(+), 41 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 67de9d76bab..cf522479457 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -14724,6 +14724,7 @@ version = "1.4.0-alpha.0" dependencies = [ "api", "async-trait", + "auth", "catalog", "client", "common-base", diff --git a/src/auth/src/common.rs b/src/auth/src/common.rs index 4b88cc60fc4..a008832490a 100644 --- a/src/auth/src/common.rs +++ b/src/auth/src/common.rs @@ -14,10 +14,11 @@ use std::sync::{Arc, LazyLock}; -use common_base::secrets::SecretString; +use common_base::secrets::{ExposeSecret, SecretString}; use digest::Digest; use hmac::{Hmac, Mac}; use pbkdf2::pbkdf2_hmac; +use serde::Serializer; use sha1::Sha1; use sha2::Sha256; use snafu::{OptionExt, ensure}; @@ -39,6 +40,9 @@ pub const MAX_PBKDF2_SHA256_ITERATIONS: u32 = 1_000_000; pub const MAX_PBKDF2_SHA256_SALT_LEN: usize = 1024; pub const PG_SCRAM_SHA256_KEY_LEN: usize = 32; +/// Placeholder used when rendering potentially secret user-provider options. +pub(crate) const REDACTED_USER_PROVIDER_OPTION: &str = "[REDACTED]"; + type HmacSha256 = Hmac; /// Process-wide secret used to derive a stable mock salt for unknown users, so @@ -52,8 +56,8 @@ pub fn userinfo_by_name(username: Option) -> UserInfoRef { } pub fn user_provider_from_option(opt: &str) -> Result { - let (name, content) = opt.split_once(':').with_context(|| InvalidConfigSnafu { - value: opt.to_string(), + let (name, content) = opt.split_once(':').context(InvalidConfigSnafu { + value: REDACTED_USER_PROVIDER_OPTION, msg: "UserProviderOption must be in format `